Add cached AV pipeline for Omni turn LoRA

cache_omni_av_pieces.py

@@ -0,0 +1,426 @@
+#!/usr/bin/env python3
+"""Cache exact Qwen2.5-Omni AV pieces for Thinker LoRA training.
+
+For every source-aligned 2-second block this script caches ten causal phases:
+0.2, 0.4, ..., 2.0 seconds.  The 2.0-second phase is reused as completed
+history; phases 0.2--1.8 are used only for the current partial block.  No
+future samples enter a phase.
+
+The optional verification compares a full-sequence reconstruction from the
+cache with the existing open-stream KV replay.  Run this 24-second pilot before
+caching a full EasyCom session.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import platform
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from tqdm.auto import tqdm
+
+import streaming_turn_model as core
+from easycom_head import NUM_HORIZONS
+from omni_cached_av import (
+    SCHEMA_VERSION,
+    CachedBlockStore,
+    anchor_intervals,
+    cache_block_filename,
+    cache_size_bytes,
+    forward_last_state,
+    make_anchor,
+    pieces_to_cache,
+    reconstruct_open_stream_sequence,
+    save_json,
+)
+from streaming_turn_model_directly import CausalSourceAV, _easycom_directory_specs
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-id", required=True)
+    parser.add_argument("--easycom-session-dir", type=Path, required=True)
+    parser.add_argument("--target", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument(
+        "--end",
+        type=float,
+        default=24.0,
+        help="Dense session endpoint in seconds; use 24 for the first pilot",
+    )
+    parser.add_argument("--step-seconds", type=float, default=0.2)
+    parser.add_argument("--native-block-seconds", type=float, default=2.0)
+    parser.add_argument("--memory-seconds", type=float, default=8.0)
+    parser.add_argument("--video-fps", type=float, default=2.0)
+    parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
+    parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
+    parser.add_argument("--frame-cache-size", type=int, default=64)
+    parser.add_argument("--decoder", choices=("decord", "pynvcodec"), default="decord")
+    parser.add_argument("--decoder-gpu-id", type=int, default=0)
+    parser.add_argument("--decoder-threads", type=int, default=1)
+    parser.add_argument("--max-open-readers", type=int, default=1)
+    parser.add_argument(
+        "--dtype",
+        choices=("auto", "bfloat16", "float16", "float32"),
+        default="bfloat16",
+    )
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--flash-attention-2", action="store_true")
+    parser.add_argument(
+        "--verify-every",
+        type=int,
+        default=10,
+        help="Compare every Nth 5-Hz state with exact KV replay; 0 disables",
+    )
+    parser.add_argument(
+        "--minimum-cosine",
+        type=float,
+        default=0.999,
+        help="Fail if cached full-forward/reference KV cosine is below this",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Overwrite this cache's known block/static/manifest files",
+    )
+    return parser.parse_args(argv)
+
+
+def _integer_ratio(numerator: float, denominator: float, name: str) -> int:
+    ratio = numerator / denominator
+    rounded = round(ratio)
+    if not math.isclose(ratio, rounded, rel_tol=0.0, abs_tol=1e-6):
+        raise ValueError(f"{name} must be integral, got {ratio}")
+    return int(rounded)
+
+
+def _load_targets(path: Path) -> np.ndarray:
+    targets = np.load(path, mmap_mode="r")
+    if targets.ndim != 2 or targets.shape[1] != 3:
+        raise ValueError(f"Expected target [T,3], got {targets.shape}: {path}")
+    return targets
+
+
+def _safe_output_path(path: Path, overwrite: bool) -> None:
+    if path.exists() and not overwrite:
+        raise FileExistsError(f"Refusing to overwrite {path}; pass --overwrite")
+
+
+def _model_device(model: torch.nn.Module) -> torch.device:
+    try:
+        return next(model.parameters()).device
+    except StopIteration as exc:
+        raise RuntimeError("Model has no parameters") from exc
+
+
+@torch.inference_mode()
+def _cache_blocks(
+    *,
+    extractor: core.OpenStreamReplayExtractor,
+    model: torch.nn.Module,
+    output_dir: Path,
+    num_steps: int,
+    steps_per_block: int,
+    step_seconds: float,
+    block_seconds: float,
+    overwrite: bool,
+) -> list[dict[str, Any]]:
+    num_blocks = math.ceil(num_steps / steps_per_block)
+    records: list[dict[str, Any]] = []
+    input_embeddings = model.get_input_embeddings()
+    for block_index in tqdm(range(num_blocks), desc="cache AV blocks", unit="block"):
+        first_step = block_index * steps_per_block
+        phase_count = min(steps_per_block, num_steps - first_step)
+        block_start = block_index * block_seconds
+        variants: dict[str, dict[str, torch.Tensor]] = {}
+        token_lengths: dict[str, int] = {}
+        for tail_steps in range(1, phase_count + 1):
+            interval_end = block_start + tail_steps * step_seconds
+            pieces = extractor._prepare_pieces(
+                video_path=None,
+                audio_path=None,
+                interval=(block_start, interval_end),
+                use_audio_in_video=True,
+            )
+            closure_embeddings = input_embeddings(pieces.closure_ids)
+            cached = pieces_to_cache(
+                pieces,
+                closure_embeddings=closure_embeddings,
+                floating_dtype=torch.bfloat16,
+            )
+            variants[str(tail_steps)] = cached
+            token_lengths[str(tail_steps)] = int(
+                cached["prefix_embeddings"].shape[1]
+                + cached["payload_embeddings"].shape[1]
+                + cached["closure_embeddings"].shape[1]
+            )
+        path = output_dir / cache_block_filename(block_index)
+        _safe_output_path(path, overwrite)
+        torch.save(
+            {
+                "schema_version": SCHEMA_VERSION,
+                "block_index": block_index,
+                "start_seconds": block_start,
+                "end_seconds": min(
+                    block_start + block_seconds, num_steps * step_seconds
+                ),
+                "variants": variants,
+            },
+            path,
+        )
+        records.append(
+            {
+                "block_index": block_index,
+                "path": path.name,
+                "phase_count": phase_count,
+                "token_lengths_without_query": token_lengths,
+                "size_bytes": path.stat().st_size,
+            }
+        )
+    return records
+
+
+@torch.inference_mode()
+def _verify_cache(
+    *,
+    model: torch.nn.Module,
+    extractor: core.OpenStreamReplayExtractor,
+    cache_dir: Path,
+    static: dict[str, Any],
+    anchors: list[dict[str, Any]],
+    verify_every: int,
+    minimum_cosine: float,
+    block_seconds: float,
+    step_seconds: float,
+) -> list[dict[str, float]]:
+    if verify_every <= 0:
+        return []
+    selected = [
+        anchor
+        for anchor in anchors
+        if (int(anchor["step_index"]) + 1) % verify_every == 0
+    ]
+    if anchors and anchors[-1] not in selected:
+        selected.append(anchors[-1])
+    store = CachedBlockStore(cache_dir, max_blocks=max(2, int(static["memory_blocks"])))
+    device = _model_device(model)
+    checks: list[dict[str, float]] = []
+    for anchor in tqdm(selected, desc="verify cached replay", unit="state"):
+        embeddings, positions = reconstruct_open_stream_sequence(
+            store=store, static=static, anchor=anchor
+        )
+        embeddings = embeddings.unsqueeze(0).to(device=device, dtype=model.dtype)
+        positions = positions.unsqueeze(1).to(device=device)
+        attention_mask = torch.ones(
+            1, embeddings.shape[1], dtype=torch.long, device=device
+        )
+        cached_state = forward_last_state(
+            model.model,
+            inputs_embeds=embeddings,
+            position_ids=positions,
+            attention_mask=attention_mask,
+            readout_indices=torch.tensor(
+                [embeddings.shape[1] - 1], dtype=torch.long, device=device
+            ),
+        )
+        exact_state = extractor.state_from_scratch(
+            video_path=None,
+            audio_path=None,
+            intervals=anchor_intervals(
+                anchor,
+                block_seconds=block_seconds,
+                step_seconds=step_seconds,
+            ),
+            use_audio_in_video=True,
+        )
+        cached_float = cached_state.float()
+        exact_float = exact_state.float()
+        check = {
+            "time_seconds": float(anchor["time_seconds"]),
+            "cosine": float(
+                F.cosine_similarity(cached_float, exact_float, dim=-1).item()
+            ),
+            "mean_abs": float((cached_float - exact_float).abs().mean().item()),
+            "max_abs": float((cached_float - exact_float).abs().max().item()),
+        }
+        checks.append(check)
+        print(
+            f"verify t={check['time_seconds']:.1f}s "
+            f"cosine={check['cosine']:.8f} "
+            f"mean_abs={check['mean_abs']:.6f} "
+            f"max_abs={check['max_abs']:.6f}",
+            flush=True,
+        )
+        if check["cosine"] < minimum_cosine:
+            raise RuntimeError(
+                "Cached full-sequence reconstruction disagrees with open-stream "
+                f"KV replay at t={check['time_seconds']:.1f}s: "
+                f"cosine={check['cosine']:.8f} < {minimum_cosine:.8f}"
+            )
+    return checks
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    args = parse_args(argv)
+    core._configure_qwen_logging()
+    if args.end <= 0:
+        raise ValueError("--end must be positive")
+    steps_per_block = _integer_ratio(
+        args.native_block_seconds, args.step_seconds, "native block / step"
+    )
+    memory_blocks = _integer_ratio(
+        args.memory_seconds, args.native_block_seconds, "memory / native block"
+    )
+    num_steps = _integer_ratio(args.end, args.step_seconds, "end / step")
+    targets = _load_targets(args.target)
+    if num_steps > len(targets):
+        raise ValueError(
+            f"Requested {num_steps} cached steps but target has only {len(targets)}"
+        )
+
+    output_dir = args.output_dir.expanduser().resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    manifest_path = output_dir / "manifest.json"
+    static_path = output_dir / "static.pt"
+    _safe_output_path(manifest_path, args.overwrite)
+    _safe_output_path(static_path, args.overwrite)
+
+    config = core.StreamConfig(
+        step_seconds=args.step_seconds,
+        native_block_seconds=args.native_block_seconds,
+        memory_seconds=args.memory_seconds,
+        video_fps=args.video_fps,
+        video_min_pixels=args.video_min_pixels,
+        video_max_pixels=args.video_max_pixels,
+    )
+    config.validate()
+    specs = _easycom_directory_specs(args.easycom_session_dir)
+    media = CausalSourceAV(
+        specs,
+        output_fps=args.video_fps,
+        required_duration=args.end,
+        frame_cache_size=args.frame_cache_size,
+        decoder_backend=args.decoder,
+        decoder_gpu_id=args.decoder_gpu_id,
+        decoder_threads=args.decoder_threads,
+        max_open_readers=args.max_open_readers,
+    )
+    print(
+        f"Direct media ready: duration={args.end:.1f}s, "
+        f"grid={media.expected_grid_frames} frames, audio={len(media.audio)} samples"
+    )
+
+    model, processor = core.load_qwen_omni(
+        args.model_id,
+        dtype=args.dtype,
+        device_map=args.device,
+        use_flash_attention_2=args.flash_attention_2,
+    )
+    model.requires_grad_(False)
+    extractor = core.OpenStreamReplayExtractor(
+        model, processor, config, fixed_grid_av=media
+    )
+    with torch.inference_mode():
+        query_ids = extractor.query_ids
+        query_embeddings = model.get_input_embeddings()(query_ids)
+    positions_per_chunk = core._tmrope_positions_per_chunk(model)
+    static = {
+        "schema_version": SCHEMA_VERSION,
+        "steps_per_block": steps_per_block,
+        "memory_blocks": memory_blocks,
+        "positions_per_chunk": positions_per_chunk,
+        "query_ids": query_ids.detach().cpu().long(),
+        "query_embeddings": query_embeddings.detach().cpu().to(torch.bfloat16),
+    }
+    torch.save(static, static_path)
+
+    block_records = _cache_blocks(
+        extractor=extractor,
+        model=model,
+        output_dir=output_dir,
+        num_steps=num_steps,
+        steps_per_block=steps_per_block,
+        step_seconds=args.step_seconds,
+        block_seconds=args.native_block_seconds,
+        overwrite=args.overwrite,
+    )
+    anchors = [
+        make_anchor(
+            step_index=step_index,
+            steps_per_block=steps_per_block,
+            memory_blocks=memory_blocks,
+            step_seconds=args.step_seconds,
+            num_target_steps=len(targets),
+            num_horizons=NUM_HORIZONS,
+        )
+        for step_index in range(num_steps)
+    ]
+    checks = _verify_cache(
+        model=model,
+        extractor=extractor,
+        cache_dir=output_dir,
+        static=static,
+        anchors=anchors,
+        verify_every=args.verify_every,
+        minimum_cosine=args.minimum_cosine,
+        block_seconds=args.native_block_seconds,
+        step_seconds=args.step_seconds,
+    )
+
+    manifest = {
+        "schema_version": SCHEMA_VERSION,
+        "cache_kind": "qwen2.5-omni-frozen-av-open-stream-phases",
+        "session": args.easycom_session_dir.resolve().name,
+        "model_id": args.model_id,
+        "target_path": str(args.target.expanduser().resolve()),
+        "target_steps": int(len(targets)),
+        "cached_steps": num_steps,
+        "cached_end_seconds": args.end,
+        "hidden_size": int(query_embeddings.shape[-1]),
+        "floating_dtype": "bfloat16",
+        "step_seconds": args.step_seconds,
+        "native_block_seconds": args.native_block_seconds,
+        "memory_seconds": args.memory_seconds,
+        "steps_per_block": steps_per_block,
+        "memory_blocks": memory_blocks,
+        "positions_per_chunk": positions_per_chunk,
+        "video_fps": args.video_fps,
+        "video_min_pixels": args.video_min_pixels,
+        "video_max_pixels": args.video_max_pixels,
+        "use_audio_in_video": True,
+        "interval_convention": "half_open_[start,end)",
+        "phase_policy": "cache_0.2_to_2.0s;_2.0s_is_completed_history",
+        "label_policy": "target[i+1:i+11]",
+        "static_path": static_path.name,
+        "blocks": block_records,
+        "anchors": anchors,
+        "verification": checks,
+        "media": media.metadata(),
+        "software": {
+            "python": platform.python_version(),
+            "torch": torch.__version__,
+        },
+    }
+    save_json(manifest_path, manifest)
+    size = cache_size_bytes(
+        [static_path, *[output_dir / record["path"] for record in block_records]]
+    )
+    print(
+        f"Saved {len(block_records)} blocks / {len(anchors)} anchors to {output_dir}"
+    )
+    print(f"Cache tensor size: {size / (1024 ** 3):.3f} GiB")
+    if checks:
+        print(
+            "Verification: "
+            f"lowest cosine={min(item['cosine'] for item in checks):.8f}, "
+            f"worst max_abs={max(item['max_abs'] for item in checks):.6f}"
+        )
+
+
+if __name__ == "__main__":
+    main()

omni_cached_av.py

@@ -0,0 +1,365 @@
+"""Cached open-stream AV pieces for Qwen2.5-Omni turn-taking LoRA.
+
+The modality towers are frozen for the first fine-tuning experiment.  This
+module serializes their exact output embeddings and Qwen's local TMRoPE
+coordinates, then reconstructs the same still-open media turn without using a
+KV cache.  The resulting sequence can be passed through the Thinker with
+``use_cache=False`` so gradients reach LoRA adapters in every retained block.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+from collections import OrderedDict
+from pathlib import Path
+from typing import Any, Iterable, Sequence
+
+import torch
+
+
+SCHEMA_VERSION = 1
+PIECE_TENSOR_NAMES = (
+    "prefix_embeddings",
+    "prefix_position_ids",
+    "payload_embeddings",
+    "payload_position_ids",
+    "closure_ids",
+    "closure_position_ids",
+    "closure_embeddings",
+)
+
+
+def load_json(path: str | Path) -> dict[str, Any]:
+    return json.loads(Path(path).read_text())
+
+
+def save_json(path: str | Path, payload: dict[str, Any]) -> None:
+    Path(path).write_text(json.dumps(payload, indent=2) + "\n")
+
+
+def torch_load(path: str | Path) -> Any:
+    """Load tensor-only cache files across old and new PyTorch defaults."""
+
+    try:
+        return torch.load(Path(path), map_location="cpu", weights_only=True)
+    except TypeError:
+        return torch.load(Path(path), map_location="cpu")
+
+
+def cache_block_filename(block_index: int) -> str:
+    return f"block_{block_index:06d}.pt"
+
+
+def pieces_to_cache(
+    pieces: Any,
+    *,
+    closure_embeddings: torch.Tensor,
+    floating_dtype: torch.dtype = torch.bfloat16,
+) -> dict[str, torch.Tensor]:
+    """Detach one ``OpenMediaPieces`` object into a portable CPU payload."""
+
+    output: dict[str, torch.Tensor] = {}
+    for name in PIECE_TENSOR_NAMES:
+        value = closure_embeddings if name == "closure_embeddings" else getattr(pieces, name)
+        value = value.detach().cpu().contiguous()
+        if value.is_floating_point():
+            value = value.to(floating_dtype)
+        output[name] = value
+    validate_piece(output)
+    return output
+
+
+def validate_piece(piece: dict[str, torch.Tensor]) -> None:
+    missing = [name for name in PIECE_TENSOR_NAMES if name not in piece]
+    if missing:
+        raise ValueError(f"Cached AV piece is missing {missing}")
+    prefix = piece["prefix_embeddings"]
+    payload = piece["payload_embeddings"]
+    closure = piece["closure_embeddings"]
+    if prefix.ndim != 3 or prefix.shape[0] != 1:
+        raise ValueError(f"Invalid prefix embeddings: {tuple(prefix.shape)}")
+    if payload.ndim != 3 or payload.shape[0] != 1:
+        raise ValueError(f"Invalid payload embeddings: {tuple(payload.shape)}")
+    if closure.ndim != 3 or closure.shape[0] != 1:
+        raise ValueError(f"Invalid closure embeddings: {tuple(closure.shape)}")
+    hidden_sizes = {int(prefix.shape[-1]), int(payload.shape[-1]), int(closure.shape[-1])}
+    if len(hidden_sizes) != 1:
+        raise ValueError(f"Cached piece hidden sizes disagree: {sorted(hidden_sizes)}")
+    expected_positions = (
+        ("prefix_position_ids", prefix.shape[1]),
+        ("payload_position_ids", payload.shape[1]),
+        ("closure_position_ids", closure.shape[1]),
+    )
+    for name, length in expected_positions:
+        value = piece[name]
+        if value.shape != (3, 1, length):
+            raise ValueError(
+                f"{name} must have shape {(3, 1, length)}, got {tuple(value.shape)}"
+            )
+    if piece["closure_ids"].shape != (1, closure.shape[1]):
+        raise ValueError(
+            "closure_ids/closure_embeddings mismatch: "
+            f"{tuple(piece['closure_ids'].shape)} vs {tuple(closure.shape)}"
+        )
+
+
+def make_anchor(
+    *,
+    step_index: int,
+    steps_per_block: int,
+    memory_blocks: int,
+    step_seconds: float,
+    num_target_steps: int,
+    num_horizons: int,
+) -> dict[str, Any]:
+    """Describe the exact completed-history/current-tail selection at time t."""
+
+    if step_index < 0:
+        raise ValueError("step_index must be non-negative")
+    block_index = step_index // steps_per_block
+    tail_steps = step_index % steps_per_block + 1
+    first_history = max(0, block_index - (memory_blocks - 1))
+    history_blocks = list(range(first_history, block_index))
+    return {
+        "step_index": int(step_index),
+        "time_seconds": float((step_index + 1) * step_seconds),
+        "block_index": int(block_index),
+        "tail_steps": int(tail_steps),
+        "history_blocks": history_blocks,
+        "target_index": int(step_index),
+        "has_future_target": bool(step_index + num_horizons < num_target_steps),
+    }
+
+
+class CachedBlockStore:
+    """Small CPU LRU for block shard files shared by adjacent samples."""
+
+    def __init__(self, cache_dir: str | Path, max_blocks: int = 16) -> None:
+        if max_blocks <= 0:
+            raise ValueError("max_blocks must be positive")
+        self.cache_dir = Path(cache_dir)
+        self.max_blocks = int(max_blocks)
+        self._blocks: OrderedDict[int, dict[str, Any]] = OrderedDict()
+
+    def get(self, block_index: int) -> dict[str, Any]:
+        cached = self._blocks.get(block_index)
+        if cached is not None:
+            self._blocks.move_to_end(block_index)
+            return cached
+        path = self.cache_dir / cache_block_filename(block_index)
+        if not path.is_file():
+            raise FileNotFoundError(path)
+        payload = torch_load(path)
+        if int(payload.get("block_index", -1)) != block_index:
+            raise ValueError(f"Block index mismatch in {path}")
+        variants = payload.get("variants")
+        if not isinstance(variants, dict) or not variants:
+            raise ValueError(f"No variants in {path}")
+        self._blocks[block_index] = payload
+        self._blocks.move_to_end(block_index)
+        while len(self._blocks) > self.max_blocks:
+            self._blocks.popitem(last=False)
+        return payload
+
+    def variant(self, block_index: int, tail_steps: int) -> dict[str, torch.Tensor]:
+        block = self.get(block_index)
+        variants = block["variants"]
+        piece = variants.get(str(tail_steps), variants.get(tail_steps))
+        if piece is None:
+            raise KeyError(
+                f"Block {block_index} has no tail_steps={tail_steps}; "
+                f"available={sorted(str(key) for key in variants)}"
+            )
+        validate_piece(piece)
+        return piece
+
+
+def _strip_embedding_batch(value: torch.Tensor) -> torch.Tensor:
+    if value.ndim != 3 or value.shape[0] != 1:
+        raise ValueError(f"Expected [1,L,D], got {tuple(value.shape)}")
+    return value[0]
+
+
+def _strip_position_batch(value: torch.Tensor) -> torch.LongTensor:
+    if value.ndim != 3 or value.shape[:2] != (3, 1):
+        raise ValueError(f"Expected [3,1,L], got {tuple(value.shape)}")
+    return value[:, 0, :].long()
+
+
+def reconstruct_open_stream_sequence(
+    *,
+    store: CachedBlockStore,
+    static: dict[str, Any],
+    anchor: dict[str, Any],
+) -> tuple[torch.Tensor, torch.LongTensor]:
+    """Materialize one differentiable Thinker input from frozen AV embeddings.
+
+    Complete history blocks use their 2.0-second variant.  Only the current
+    block uses its 0.2--2.0-second phase.  The first retained block keeps the
+    chat/media prefix; later blocks append payload only.  The final block's
+    media closure and the artificial assistant readout query are appended once.
+    """
+
+    steps_per_block = int(static["steps_per_block"])
+    positions_per_chunk = int(static["positions_per_chunk"])
+    history_blocks = [int(value) for value in anchor["history_blocks"]]
+    current_block = int(anchor["block_index"])
+    selected: list[tuple[int, int]] = [
+        (block_index, steps_per_block) for block_index in history_blocks
+    ]
+    selected.append((current_block, int(anchor["tail_steps"])))
+    if not selected:
+        raise ValueError("An anchor must select at least its current block")
+
+    embedding_parts: list[torch.Tensor] = []
+    position_parts: list[torch.LongTensor] = []
+    first_payload_base: int | None = None
+    last_piece: dict[str, torch.Tensor] | None = None
+    last_shift = 0
+
+    for sequence_index, (block_index, tail_steps) in enumerate(selected):
+        piece = store.variant(block_index, tail_steps)
+        local_payload_positions = _strip_position_batch(piece["payload_position_ids"])
+        local_payload_base = int(local_payload_positions.min().item())
+        if first_payload_base is None:
+            first_payload_base = local_payload_base
+        desired_payload_base = first_payload_base + sequence_index * positions_per_chunk
+        position_shift = desired_payload_base - local_payload_base
+
+        if sequence_index == 0:
+            embedding_parts.append(_strip_embedding_batch(piece["prefix_embeddings"]))
+            position_parts.append(_strip_position_batch(piece["prefix_position_ids"]))
+        embedding_parts.append(_strip_embedding_batch(piece["payload_embeddings"]))
+        position_parts.append(local_payload_positions + position_shift)
+        last_piece = piece
+        last_shift = position_shift
+
+    if last_piece is None:
+        raise RuntimeError("No cached AV piece was selected")
+    closure_embeddings = _strip_embedding_batch(last_piece["closure_embeddings"])
+    closure_positions = (
+        _strip_position_batch(last_piece["closure_position_ids"]) + last_shift
+    )
+    embedding_parts.append(closure_embeddings)
+    position_parts.append(closure_positions)
+
+    query_embeddings = static["query_embeddings"]
+    if query_embeddings.ndim != 3 or query_embeddings.shape[0] != 1:
+        raise ValueError(
+            f"static query_embeddings must be [1,Q,D], got {tuple(query_embeddings.shape)}"
+        )
+    query_embeddings = query_embeddings[0]
+    query_start = int(closure_positions.max().item()) + 1
+    query_scalar = torch.arange(
+        query_start,
+        query_start + query_embeddings.shape[0],
+        dtype=torch.long,
+    )
+    query_positions = query_scalar.view(1, -1).expand(3, -1)
+    embedding_parts.append(query_embeddings)
+    position_parts.append(query_positions)
+
+    embeddings = torch.cat(embedding_parts, dim=0).contiguous()
+    position_ids = torch.cat(position_parts, dim=1).contiguous()
+    if position_ids.shape != (3, embeddings.shape[0]):
+        raise RuntimeError(
+            "Reconstructed position/embedding mismatch: "
+            f"{tuple(position_ids.shape)} vs {tuple(embeddings.shape)}"
+        )
+    return embeddings, position_ids
+
+
+def collate_cached_sequences(
+    batch: Sequence[dict[str, Any]],
+) -> dict[str, Any]:
+    """Right-pad variable token sequences while preserving each readout index."""
+
+    if not batch:
+        raise ValueError("Cannot collate an empty batch")
+    hidden_size = int(batch[0]["inputs_embeds"].shape[-1])
+    max_length = max(int(item["inputs_embeds"].shape[0]) for item in batch)
+    dtype = batch[0]["inputs_embeds"].dtype
+    inputs_embeds = torch.zeros(
+        len(batch), max_length, hidden_size, dtype=dtype
+    )
+    position_ids = torch.zeros(3, len(batch), max_length, dtype=torch.long)
+    attention_mask = torch.zeros(len(batch), max_length, dtype=torch.long)
+    readout_indices = torch.empty(len(batch), dtype=torch.long)
+    targets = torch.stack([item["targets"] for item in batch])
+
+    for batch_index, item in enumerate(batch):
+        embeddings = item["inputs_embeds"]
+        positions = item["position_ids"]
+        length = int(embeddings.shape[0])
+        if embeddings.shape[-1] != hidden_size:
+            raise ValueError("Mixed hidden sizes in one batch")
+        if positions.shape != (3, length):
+            raise ValueError(
+                f"Invalid position_ids {tuple(positions.shape)} for length {length}"
+            )
+        inputs_embeds[batch_index, :length] = embeddings
+        position_ids[:, batch_index, :length] = positions
+        attention_mask[batch_index, :length] = 1
+        readout_indices[batch_index] = length - 1
+
+    return {
+        "inputs_embeds": inputs_embeds,
+        "position_ids": position_ids,
+        "attention_mask": attention_mask,
+        "readout_indices": readout_indices,
+        "targets": targets,
+        "step_indices": torch.tensor(
+            [int(item["step_index"]) for item in batch], dtype=torch.long
+        ),
+    }
+
+
+def anchor_intervals(
+    anchor: dict[str, Any],
+    *,
+    block_seconds: float,
+    step_seconds: float,
+) -> list[tuple[float, float]]:
+    """Return source intervals used by the reference KV replay for an anchor."""
+
+    intervals = [
+        (index * block_seconds, (index + 1) * block_seconds)
+        for index in anchor["history_blocks"]
+    ]
+    block_index = int(anchor["block_index"])
+    current_start = block_index * block_seconds
+    current_end = current_start + int(anchor["tail_steps"]) * step_seconds
+    intervals.append((current_start, current_end))
+    return intervals
+
+
+def forward_last_state(
+    language_model: torch.nn.Module,
+    *,
+    inputs_embeds: torch.Tensor,
+    position_ids: torch.LongTensor,
+    attention_mask: torch.LongTensor,
+    readout_indices: torch.LongTensor,
+) -> torch.Tensor:
+    """Run full-sequence causal Thinker attention and gather each query state."""
+
+    outputs = language_model(
+        inputs_embeds=inputs_embeds,
+        position_ids=position_ids,
+        attention_mask=attention_mask,
+        use_cache=False,
+        output_hidden_states=False,
+        return_dict=True,
+        cache_position=torch.arange(
+            inputs_embeds.shape[1], device=inputs_embeds.device
+        ),
+    )
+    batch_indices = torch.arange(
+        inputs_embeds.shape[0], device=inputs_embeds.device
+    )
+    return outputs.last_hidden_state[batch_indices, readout_indices]
+
+
+def cache_size_bytes(paths: Iterable[Path]) -> int:
+    return sum(path.stat().st_size for path in paths if path.is_file())

requirements.txt

@@ -8,5 +8,8 @@ Pillow
 tqdm
 scikit-learn
 
+# Used by train_omni_turn_lora.py; already included in the ms-swift image.
+peft
+
 # Optional on compatible NVIDIA GPUs: flash-attn
 # Optional NVDEC backend: PyNvVideoCodec

train_omni_turn_lora.py

@@ -0,0 +1,634 @@
+#!/usr/bin/env python3
+"""Train Qwen2.5-Omni Thinker LoRA + EgoSpeak 10x3 head from AV caches.
+
+This is a custom task runner intended to execute inside the ms-swift image.
+It uses the PEFT library shipped with ms-swift because Swift's normal SFT CLI
+optimizes generated text tokens, whereas this experiment reads one continuous
+Thinker hidden state and applies a multi-horizon classification loss.
+
+The cached audio and vision towers remain frozen.  The reconstructed sequence
+is forwarded with ``use_cache=False`` so all retained 2-second blocks
+participate in autograd through the Thinker LoRA adapters.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import importlib.metadata
+import json
+import math
+import random
+from pathlib import Path
+from typing import Any, Sequence
+
+import numpy as np
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+from tqdm.auto import tqdm
+
+import streaming_turn_model as core
+from easycom_head import (
+    NUM_CLASSES,
+    NUM_HORIZONS,
+    TurnPredictionHead,
+    checkpoint_head_config,
+    egospeak_multihorizon_loss,
+)
+from omni_cached_av import (
+    SCHEMA_VERSION,
+    CachedBlockStore,
+    collate_cached_sequences,
+    forward_last_state,
+    load_json,
+    reconstruct_open_stream_sequence,
+    save_json,
+    torch_load,
+)
+
+
+DEFAULT_LORA_TARGETS = (
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "down_proj",
+)
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-id", required=True)
+    parser.add_argument(
+        "--cache-dir",
+        type=Path,
+        action="append",
+        required=True,
+        help="Cache directory from cache_omni_av_pieces.py; repeat for sessions",
+    )
+    parser.add_argument(
+        "--target",
+        type=Path,
+        help="Override target path for a single --cache-dir",
+    )
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--head-checkpoint", type=Path)
+    parser.add_argument("--max-examples", type=int, default=0)
+    parser.add_argument("--max-steps", type=int, default=200)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
+    parser.add_argument("--num-workers", type=int, default=0)
+    parser.add_argument("--block-cache-size", type=int, default=16)
+    parser.add_argument("--learning-rate", type=float, default=2e-5)
+    parser.add_argument("--head-learning-rate", type=float, default=1e-4)
+    parser.add_argument("--weight-decay", type=float, default=0.01)
+    parser.add_argument("--warmup-steps", type=int, default=10)
+    parser.add_argument("--max-grad-norm", type=float, default=1.0)
+    parser.add_argument("--lora-rank", type=int, default=8)
+    parser.add_argument("--lora-alpha", type=int, default=16)
+    parser.add_argument("--lora-dropout", type=float, default=0.05)
+    parser.add_argument(
+        "--lora-target-modules",
+        default=",".join(DEFAULT_LORA_TARGETS),
+    )
+    parser.add_argument("--head-dropout", type=float, default=0.1)
+    parser.add_argument("--gradient-checkpointing", action="store_true")
+    parser.add_argument("--seed", type=int, default=1337)
+    parser.add_argument("--log-every", type=int, default=5)
+    parser.add_argument("--eval-max-batches", type=int, default=4)
+    parser.add_argument(
+        "--dtype",
+        choices=("bfloat16", "float16"),
+        default="bfloat16",
+    )
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--flash-attention-2", action="store_true")
+    parser.add_argument("--overwrite-output", action="store_true")
+    parser.add_argument("--no-progress", action="store_true")
+    return parser.parse_args(argv)
+
+
+def _seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def _package_version(name: str) -> str | None:
+    try:
+        return importlib.metadata.version(name)
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def _load_targets(path: Path) -> np.ndarray:
+    targets = np.load(path, mmap_mode="r")
+    if targets.ndim != 2 or targets.shape[1] != NUM_CLASSES:
+        raise ValueError(f"Expected [T,{NUM_CLASSES}] target, got {targets.shape}: {path}")
+    return targets
+
+
+class OmniTurnCacheDataset(Dataset[dict[str, Any]]):
+    """One sample per 200-ms anchor, reconstructed lazily on CPU."""
+
+    def __init__(
+        self,
+        cache_dirs: Sequence[Path],
+        *,
+        target_override: Path | None,
+        max_examples: int,
+        seed: int,
+        block_cache_size: int,
+    ) -> None:
+        if target_override is not None and len(cache_dirs) != 1:
+            raise ValueError("--target can override only one --cache-dir")
+        self.sources: list[dict[str, Any]] = []
+        self.examples: list[tuple[int, dict[str, Any]]] = []
+        hidden_size: int | None = None
+        reference_contract: tuple[Any, ...] | None = None
+
+        for source_index, raw_cache_dir in enumerate(cache_dirs):
+            cache_dir = raw_cache_dir.expanduser().resolve()
+            manifest_path = cache_dir / "manifest.json"
+            manifest = load_json(manifest_path)
+            if int(manifest.get("schema_version", -1)) != SCHEMA_VERSION:
+                raise ValueError(
+                    f"Unsupported cache schema in {manifest_path}: "
+                    f"{manifest.get('schema_version')}"
+                )
+            static = torch_load(cache_dir / manifest["static_path"])
+            target_path = (
+                target_override.expanduser().resolve()
+                if target_override is not None
+                else Path(manifest["target_path"])
+            )
+            targets = _load_targets(target_path)
+            if len(targets) != int(manifest["target_steps"]):
+                raise ValueError(
+                    f"Target length changed for {manifest_path}: "
+                    f"{len(targets)} vs {manifest['target_steps']}"
+                )
+            source_hidden = int(manifest["hidden_size"])
+            hidden_size = source_hidden if hidden_size is None else hidden_size
+            if source_hidden != hidden_size:
+                raise ValueError("Cache directories have different hidden sizes")
+            contract = (
+                manifest["model_id"],
+                manifest["step_seconds"],
+                manifest["native_block_seconds"],
+                manifest["memory_seconds"],
+                manifest["video_fps"],
+                manifest["positions_per_chunk"],
+            )
+            reference_contract = contract if reference_contract is None else reference_contract
+            if contract != reference_contract:
+                raise ValueError(
+                    "Cache directories do not share the same model/timing contract: "
+                    f"{manifest_path}"
+                )
+            self.sources.append(
+                {
+                    "cache_dir": cache_dir,
+                    "manifest": manifest,
+                    "static": static,
+                    "targets": targets,
+                    "store": CachedBlockStore(cache_dir, max_blocks=block_cache_size),
+                }
+            )
+            for anchor in manifest["anchors"]:
+                target_index = int(anchor["target_index"])
+                if not bool(anchor.get("has_future_target", True)):
+                    continue
+                if target_index + NUM_HORIZONS >= len(targets):
+                    continue
+                self.examples.append((source_index, anchor))
+
+        if not self.examples:
+            raise ValueError("No cache anchors have ten future target steps")
+        if max_examples < 0:
+            raise ValueError("--max-examples must be non-negative")
+        if max_examples and max_examples < len(self.examples):
+            generator = random.Random(seed)
+            selected = generator.sample(range(len(self.examples)), max_examples)
+            self.examples = [self.examples[index] for index in sorted(selected)]
+        self.hidden_size = int(hidden_size or 0)
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        source_index, anchor = self.examples[index]
+        source = self.sources[source_index]
+        embeddings, positions = reconstruct_open_stream_sequence(
+            store=source["store"],
+            static=source["static"],
+            anchor=anchor,
+        )
+        target_index = int(anchor["target_index"])
+        targets = torch.from_numpy(
+            np.array(
+                source["targets"][
+                    target_index + 1 : target_index + NUM_HORIZONS + 1
+                ],
+                dtype=np.float32,
+                copy=True,
+            )
+        )
+        if targets.shape != (NUM_HORIZONS, NUM_CLASSES):
+            raise RuntimeError(
+                f"Unexpected future target shape {tuple(targets.shape)} at {target_index}"
+            )
+        return {
+            "inputs_embeds": embeddings,
+            "position_ids": positions,
+            "targets": targets,
+            "step_index": target_index,
+        }
+
+
+def _load_initial_head(
+    checkpoint: Path | None,
+    *,
+    hidden_size: int,
+    dropout: float,
+) -> TurnPredictionHead:
+    if checkpoint is None:
+        return TurnPredictionHead(hidden_size=hidden_size, dropout=dropout)
+    try:
+        payload = torch.load(checkpoint, map_location="cpu", weights_only=True)
+    except (TypeError, RuntimeError):
+        payload = torch.load(checkpoint, map_location="cpu")
+    if not isinstance(payload, dict):
+        raise ValueError(f"Expected dict head checkpoint: {checkpoint}")
+    config = checkpoint_head_config(payload)
+    if int(config["hidden_size"]) != hidden_size:
+        raise ValueError(
+            f"Head hidden size {config['hidden_size']} != cache hidden size {hidden_size}"
+        )
+    head = TurnPredictionHead(
+        hidden_size=hidden_size,
+        num_horizons=int(config["num_horizons"]),
+        num_classes=int(config["num_classes"]),
+        dropout=float(config["dropout"]),
+    )
+    head.load_state_dict(payload.get("head", payload))
+    print(f"Initialized classification head from {checkpoint}")
+    return head
+
+
+def _prepare_language_lora(
+    model: torch.nn.Module,
+    *,
+    rank: int,
+    alpha: int,
+    dropout: float,
+    target_modules: Sequence[str],
+    gradient_checkpointing: bool,
+) -> torch.nn.Module:
+    try:
+        from peft import LoraConfig, get_peft_model
+    except ImportError as exc:
+        raise RuntimeError(
+            "PEFT is required. The ms-swift 3.12.5 image should include it."
+        ) from exc
+
+    model.requires_grad_(False)
+    language_model = model.model
+    matched = [
+        name
+        for name, module in language_model.named_modules()
+        if isinstance(module, nn.Linear) and name.rsplit(".", 1)[-1] in target_modules
+    ]
+    if not matched:
+        raise RuntimeError(
+            "No language-model linear layers matched LoRA targets "
+            f"{list(target_modules)}"
+        )
+    config = LoraConfig(
+        r=rank,
+        lora_alpha=alpha,
+        lora_dropout=dropout,
+        target_modules=list(target_modules),
+        bias="none",
+    )
+    language_model = get_peft_model(language_model, config)
+    language_model.config.use_cache = False
+    if gradient_checkpointing:
+        try:
+            language_model.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs={"use_reentrant": False}
+            )
+        except TypeError:
+            language_model.gradient_checkpointing_enable()
+    language_model.print_trainable_parameters()
+    print(
+        f"LoRA matched {len(matched)} language layers; first={matched[0]}, "
+        f"last={matched[-1]}"
+    )
+    return language_model
+
+
+def _lr_lambda(step: int, *, total_steps: int, warmup_steps: int) -> float:
+    if warmup_steps > 0 and step < warmup_steps:
+        return float(step + 1) / float(warmup_steps)
+    remaining = max(1, total_steps - warmup_steps)
+    progress = min(1.0, max(0.0, (step - warmup_steps) / remaining))
+    return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+
+@torch.inference_mode()
+def _evaluate_loss(
+    language_model: torch.nn.Module,
+    head: TurnPredictionHead,
+    loader: DataLoader,
+    *,
+    device: torch.device,
+    model_dtype: torch.dtype,
+    max_batches: int,
+) -> float:
+    language_was_training = language_model.training
+    head_was_training = head.training
+    language_model.eval()
+    head.eval()
+    losses: list[float] = []
+    for batch_index, batch in enumerate(loader):
+        if max_batches > 0 and batch_index >= max_batches:
+            break
+        inputs_embeds = batch["inputs_embeds"].to(
+            device=device, dtype=model_dtype, non_blocking=True
+        )
+        position_ids = batch["position_ids"].to(device=device, non_blocking=True)
+        attention_mask = batch["attention_mask"].to(device=device, non_blocking=True)
+        readout_indices = batch["readout_indices"].to(device=device, non_blocking=True)
+        targets = batch["targets"].to(device=device, dtype=torch.float32)
+        states = forward_last_state(
+            language_model,
+            inputs_embeds=inputs_embeds,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            readout_indices=readout_indices,
+        )
+        logits = head(states.float())
+        losses.append(float(egospeak_multihorizon_loss(logits, targets).item()))
+    language_model.train(language_was_training)
+    head.train(head_was_training)
+    return float(np.mean(losses))
+
+
+def _jsonable_args(args: argparse.Namespace) -> dict[str, Any]:
+    output: dict[str, Any] = {}
+    for key, value in vars(args).items():
+        if isinstance(value, Path):
+            output[key] = str(value)
+        elif isinstance(value, list) and value and isinstance(value[0], Path):
+            output[key] = [str(item) for item in value]
+        else:
+            output[key] = value
+    return output
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    args = parse_args(argv)
+    if args.max_steps <= 0:
+        raise ValueError("--max-steps must be positive")
+    if args.batch_size <= 0 or args.gradient_accumulation_steps <= 0:
+        raise ValueError("Batch size and gradient accumulation must be positive")
+    if args.num_workers != 0:
+        print(
+            "Warning: each worker owns an independent block-file LRU; use 0 for "
+            "the 16-example pilot and increase only after profiling."
+        )
+    output_dir = args.output_dir.expanduser().resolve()
+    known_outputs = (
+        output_dir / "adapter_config.json",
+        output_dir / "adapter",
+        output_dir / "turn_head.pt",
+        output_dir / "train_summary.json",
+    )
+    if not args.overwrite_output and any(path.exists() for path in known_outputs):
+        raise FileExistsError(
+            f"Training output already exists in {output_dir}; pass --overwrite-output"
+        )
+    output_dir.mkdir(parents=True, exist_ok=True)
+    _seed_everything(args.seed)
+    core._configure_qwen_logging()
+    device = torch.device(args.device)
+    if device.type != "cuda":
+        raise ValueError("This LoRA pilot is intended for one CUDA GPU")
+
+    dataset = OmniTurnCacheDataset(
+        args.cache_dir,
+        target_override=args.target,
+        max_examples=args.max_examples,
+        seed=args.seed,
+        block_cache_size=args.block_cache_size,
+    )
+    generator = torch.Generator().manual_seed(args.seed)
+    train_loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        persistent_workers=args.num_workers > 0,
+        collate_fn=collate_cached_sequences,
+        generator=generator,
+    )
+    eval_loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=True,
+        collate_fn=collate_cached_sequences,
+    )
+    print(
+        f"Cached training examples: {len(dataset)}, hidden_size={dataset.hidden_size}, "
+        f"batch={args.batch_size} x accum={args.gradient_accumulation_steps}"
+    )
+
+    base_model, _ = core.load_qwen_omni(
+        args.model_id,
+        dtype=args.dtype,
+        device_map=args.device,
+        use_flash_attention_2=args.flash_attention_2,
+    )
+    target_modules = tuple(
+        value.strip() for value in args.lora_target_modules.split(",") if value.strip()
+    )
+    if not target_modules:
+        raise ValueError("--lora-target-modules is empty")
+    language_model = _prepare_language_lora(
+        base_model,
+        rank=args.lora_rank,
+        alpha=args.lora_alpha,
+        dropout=args.lora_dropout,
+        target_modules=target_modules,
+        gradient_checkpointing=args.gradient_checkpointing,
+    )
+    # Cached embeddings make both modality towers unnecessary during training.
+    # Keeping only the language model materially reduces L20 memory use.
+    del base_model
+    gc.collect()
+    torch.cuda.empty_cache()
+    model_dtype = next(language_model.parameters()).dtype
+    head = _load_initial_head(
+        args.head_checkpoint,
+        hidden_size=dataset.hidden_size,
+        dropout=args.head_dropout,
+    ).to(device=device, dtype=torch.float32)
+
+    lora_parameters = [parameter for parameter in language_model.parameters() if parameter.requires_grad]
+    if not lora_parameters:
+        raise RuntimeError("LoRA preparation produced no trainable parameters")
+    optimizer = torch.optim.AdamW(
+        [
+            {"params": lora_parameters, "lr": args.learning_rate},
+            {"params": head.parameters(), "lr": args.head_learning_rate},
+        ],
+        weight_decay=args.weight_decay,
+    )
+    scheduler = torch.optim.lr_scheduler.LambdaLR(
+        optimizer,
+        lr_lambda=lambda step: _lr_lambda(
+            step, total_steps=args.max_steps, warmup_steps=args.warmup_steps
+        ),
+    )
+
+    initial_loss = _evaluate_loss(
+        language_model,
+        head,
+        eval_loader,
+        device=device,
+        model_dtype=model_dtype,
+        max_batches=args.eval_max_batches,
+    )
+    print(f"Initial cached loss: {initial_loss:.6f}")
+    language_model.train()
+    head.train()
+    optimizer.zero_grad(set_to_none=True)
+    iterator = iter(train_loader)
+    history: list[dict[str, float | int]] = []
+    progress = tqdm(
+        total=args.max_steps,
+        desc="Thinker LoRA",
+        unit="update",
+        disable=args.no_progress,
+    )
+    for update_step in range(1, args.max_steps + 1):
+        accumulated_loss = 0.0
+        for _ in range(args.gradient_accumulation_steps):
+            try:
+                batch = next(iterator)
+            except StopIteration:
+                iterator = iter(train_loader)
+                batch = next(iterator)
+            inputs_embeds = batch["inputs_embeds"].to(
+                device=device, dtype=model_dtype, non_blocking=True
+            )
+            position_ids = batch["position_ids"].to(device=device, non_blocking=True)
+            attention_mask = batch["attention_mask"].to(device=device, non_blocking=True)
+            readout_indices = batch["readout_indices"].to(device=device, non_blocking=True)
+            targets = batch["targets"].to(device=device, dtype=torch.float32)
+            states = forward_last_state(
+                language_model,
+                inputs_embeds=inputs_embeds,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                readout_indices=readout_indices,
+            )
+            logits = head(states.float())
+            loss = egospeak_multihorizon_loss(logits, targets)
+            (loss / args.gradient_accumulation_steps).backward()
+            accumulated_loss += float(loss.detach()) / args.gradient_accumulation_steps
+
+        trainable_parameters = [*lora_parameters, *head.parameters()]
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            trainable_parameters, args.max_grad_norm
+        )
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad(set_to_none=True)
+        record = {
+            "step": update_step,
+            "loss": accumulated_loss,
+            "grad_norm": float(grad_norm),
+            "lora_lr": float(optimizer.param_groups[0]["lr"]),
+            "head_lr": float(optimizer.param_groups[1]["lr"]),
+        }
+        history.append(record)
+        progress.update(1)
+        progress.set_postfix(
+            loss=f"{accumulated_loss:.4f}",
+            lr=f"{optimizer.param_groups[0]['lr']:.2g}",
+        )
+        if update_step == 1 or update_step % args.log_every == 0:
+            print(
+                f"step={update_step} loss={accumulated_loss:.6f} "
+                f"grad_norm={float(grad_norm):.4f} "
+                f"lora_lr={optimizer.param_groups[0]['lr']:.3g}",
+                flush=True,
+            )
+    progress.close()
+
+    final_loss = _evaluate_loss(
+        language_model,
+        head,
+        eval_loader,
+        device=device,
+        model_dtype=model_dtype,
+        max_batches=args.eval_max_batches,
+    )
+    adapter_dir = output_dir / "adapter"
+    language_model.save_pretrained(adapter_dir)
+    torch.save(
+        {
+            "head": head.state_dict(),
+            "head_config": {
+                "hidden_size": dataset.hidden_size,
+                "num_horizons": NUM_HORIZONS,
+                "num_classes": NUM_CLASSES,
+                "dropout": head.dropout,
+            },
+            "initial_head_checkpoint": (
+                str(args.head_checkpoint.resolve()) if args.head_checkpoint else None
+            ),
+            "training_step": args.max_steps,
+        },
+        output_dir / "turn_head.pt",
+    )
+    summary = {
+        "schema_version": 1,
+        "status": "pilot" if args.max_examples else "training",
+        "model_id": args.model_id,
+        "cache_dirs": [str(path.expanduser().resolve()) for path in args.cache_dir],
+        "num_examples": len(dataset),
+        "initial_loss": initial_loss,
+        "final_loss": final_loss,
+        "loss_reduction_ratio": (
+            (initial_loss - final_loss) / initial_loss if initial_loss else None
+        ),
+        "peak_cuda_allocated_gib": torch.cuda.max_memory_allocated(device) / (1024 ** 3),
+        "packages": {
+            "torch": torch.__version__,
+            "transformers": _package_version("transformers"),
+            "peft": _package_version("peft"),
+            "ms-swift": _package_version("ms-swift"),
+        },
+        "args": _jsonable_args(args),
+        "history": history,
+    }
+    save_json(output_dir / "train_summary.json", summary)
+    print(f"Final cached loss: {final_loss:.6f} (initial {initial_loss:.6f})")
+    print(f"Saved LoRA adapter to {adapter_dir}")
+    print(f"Saved turn head to {output_dir / 'turn_head.pt'}")
+    print(f"Peak CUDA allocated: {summary['peak_cuda_allocated_gib']:.2f} GiB")
+
+
+if __name__ == "__main__":
+    main()