icassp2027-experiment/train_omni_turn_lora.py

#!/usr/bin/env python3
"""Train a frozen or LoRA Qwen2.5-Omni Thinker + 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 inspect
import json
import math
import os
import queue
import random
import re
import threading
import time
from contextlib import ExitStack
from dataclasses import dataclass
from datetime import timedelta
from pathlib import Path
from typing import Any, Sequence

import numpy as np
import torch
import torch.distributed as dist
from torch import nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
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,
    pieces_to_cache,
    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",
)

_TRANSFORMER_LAYER_PATTERN = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)")


@dataclass(frozen=True)
class DistributedRun:
    enabled: bool
    rank: int
    local_rank: int
    world_size: int
    device: torch.device

    @property
    def is_main(self) -> bool:
        return self.rank == 0


def _initialize_distributed(args: argparse.Namespace) -> DistributedRun:
    """Initialize torchrun/NCCL when WORLD_SIZE is greater than one."""

    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    if world_size <= 1:
        device = torch.device(args.device)
        if device.type != "cuda":
            raise ValueError("This LoRA trainer requires a CUDA device")
        torch.cuda.set_device(device)
        return DistributedRun(
            enabled=False,
            rank=0,
            local_rank=device.index or 0,
            world_size=1,
            device=device,
        )

    if not dist.is_available():
        raise RuntimeError("torch.distributed is unavailable in this PyTorch build")
    if "LOCAL_RANK" not in os.environ or "RANK" not in os.environ:
        raise RuntimeError("Multi-GPU training must be launched with torchrun")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    dist.init_process_group(
        backend="nccl",
        timeout=timedelta(minutes=float(args.ddp_timeout_minutes)),
    )
    return DistributedRun(
        enabled=True,
        rank=dist.get_rank(),
        local_rank=local_rank,
        world_size=dist.get_world_size(),
        device=torch.device("cuda", local_rank),
    )


def _distributed_barrier(run: DistributedRun) -> None:
    if run.enabled:
        dist.barrier()


def _distributed_average(value: float, run: DistributedRun) -> float:
    if not run.enabled:
        return float(value)
    tensor = torch.tensor(value, dtype=torch.float64, device=run.device)
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    return float((tensor / run.world_size).item())


def _synchronize_trainable_parameters(
    modules: Sequence[torch.nn.Module], run: DistributedRun
) -> None:
    """Broadcast only LoRA/head tensors, avoiding a second 3B-model broadcast."""

    if not run.enabled:
        return
    for module in modules:
        for parameter in module.parameters():
            if parameter.requires_grad:
                dist.broadcast(parameter.data, src=0)


def _wrap_ddp(module: torch.nn.Module, run: DistributedRun) -> torch.nn.Module:
    if not run.enabled:
        return module
    kwargs: dict[str, Any] = {
        "device_ids": [run.local_rank],
        "output_device": run.local_rank,
        "broadcast_buffers": False,
        "find_unused_parameters": False,
    }
    # PyTorch >= 2.7 can skip the expensive initial broadcast of all frozen
    # base-model parameters.  The checkpoint is identical on every rank and
    # trainable LoRA/head tensors were synchronized explicitly above.
    if "init_sync" in inspect.signature(DDP.__init__).parameters:
        kwargs["init_sync"] = False
    return DDP(module, **kwargs)


def _load_base_model(
    args: argparse.Namespace, run: DistributedRun
) -> tuple[torch.nn.Module, Any]:
    """Load one rank at a time so four shard readers do not hit storage together."""

    if not run.enabled:
        return core.load_qwen_omni(
            args.model_id,
            dtype=args.dtype,
            device_map=str(run.device),
            use_flash_attention_2=args.flash_attention_2,
        )

    loaded: tuple[torch.nn.Module, Any] | None = None
    for loading_rank in range(run.world_size):
        if run.rank == loading_rank:
            print(
                f"[rank {run.rank}/{run.world_size}] loading model on {run.device}",
                flush=True,
            )
            loaded = core.load_qwen_omni(
                args.model_id,
                dtype=args.dtype,
                device_map=str(run.device),
                use_flash_attention_2=args.flash_attention_2,
            )
        dist.barrier()
    if loaded is None:
        raise RuntimeError(f"Rank {run.rank} did not load a model")
    return loaded


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(
        "--adapter-checkpoint",
        type=Path,
        help=(
            "Existing PEFT adapter directory to continue from; optimizer and "
            "scheduler are reinitialized"
        ),
    )
    parser.add_argument(
        "--freeze-thinker",
        action="store_true",
        help=(
            "Do not create or update LoRA parameters; run the same cached "
            "Thinker sequences under no_grad and train only the turn head"
        ),
    )
    parser.add_argument(
        "--freeze-head",
        action="store_true",
        help=(
            "Keep the initialized turn head fixed and update only Thinker "
            "LoRA parameters; requires --head-checkpoint"
        ),
    )
    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(
        "--online-current-tail",
        action="store_true",
        help=(
            "Benchmark complete-2s-block caching: read completed history from "
            "cache but recompute the current 0.2--1.8s AV tail from source media"
        ),
    )
    parser.add_argument(
        "--easycom-session-dir",
        type=Path,
        help="Original Session_N media; required by --online-current-tail",
    )
    parser.add_argument(
        "--producer-device",
        help=(
            "Optional second CUDA device for asynchronous frozen-AV production, "
            "for example cuda:1 while --device is cuda:0"
        ),
    )
    parser.add_argument(
        "--prefetch-depth",
        type=int,
        default=4,
        help="Maximum producer-ready CPU batches queued ahead of training",
    )
    parser.add_argument("--frame-cache-size", type=int, default=64)
    parser.add_argument(
        "--audio-feature-device",
        default=None,
        help=(
            "Device used by the native Whisper log-Mel extractor for online "
            "current tails, for example model, cuda:0, or cpu; omitted "
            "preserves processor defaults"
        ),
    )
    parser.add_argument(
        "--audio-right-padding-samples",
        type=int,
        default=None,
        help=(
            "For online current tails, replace model-maximum audio padding "
            "with actual samples plus this many zeros (hop-aligned); omitted "
            "preserves the original maximum-length padding"
        ),
    )
    parser.add_argument(
        "--decoder",
        choices=("decord", "torchcodec", "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("--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(
        "--lora-last-n-layers",
        type=int,
        default=0,
        help=(
            "Apply a newly created LoRA adapter only to the final N Thinker "
            "language layers; 0 keeps all matched layers"
        ),
    )
    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(
        "--save-every-steps",
        type=int,
        default=0,
        help=(
            "Save an evaluation-ready adapter/head checkpoint every N optimizer "
            "updates under checkpoint-step-N; 0 disables intermediate saves"
        ),
    )
    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",
        help="Single-GPU device; torchrun overrides this with LOCAL_RANK",
    )
    parser.add_argument(
        "--ddp-timeout-minutes",
        type=float,
        default=60.0,
        help="NCCL process-group timeout used by torchrun multi-GPU training",
    )
    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,
        defer_current_tail: bool = False,
    ) -> 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]]] = []
        self.defer_current_tail = bool(defer_current_tail)
        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]
        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}"
            )
        item: dict[str, Any] = {
            "targets": targets,
            "step_index": target_index,
        }
        if self.defer_current_tail:
            item["source_index"] = source_index
            item["anchor"] = anchor
        else:
            embeddings, positions = reconstruct_open_stream_sequence(
                store=source["store"],
                static=source["static"],
                anchor=anchor,
            )
            item["inputs_embeds"] = embeddings
            item["position_ids"] = positions
        return item


def _collate_deferred_tail(batch: Sequence[dict[str, Any]]) -> dict[str, Any]:
    if len(batch) != 1:
        raise ValueError("--online-current-tail benchmark currently requires batch size 1")
    return batch[0]


class OnlineCurrentTailMaterializer:
    """Recompute only the current partial block while history stays cached."""

    def __init__(
        self,
        *,
        dataset: OmniTurnCacheDataset,
        model: torch.nn.Module,
        processor: Any,
        session_dir: Path,
        frame_cache_size: int,
        decoder: str,
        decoder_gpu_id: int,
        decoder_threads: int,
        max_open_readers: int,
        audio_feature_device: str | torch.device | None,
        audio_right_padding_samples: int | None,
    ) -> None:
        if len(dataset.sources) != 1:
            raise ValueError("Online-tail benchmark supports exactly one cache directory")
        from streaming_turn_model_directly import (
            CausalSourceAV,
            _easycom_directory_specs,
        )

        self.dataset = dataset
        self.model = model
        source = dataset.sources[0]
        manifest = source["manifest"]
        config = core.StreamConfig(
            step_seconds=float(manifest["step_seconds"]),
            native_block_seconds=float(manifest["native_block_seconds"]),
            memory_seconds=float(manifest["memory_seconds"]),
            video_fps=float(manifest["video_fps"]),
            video_min_pixels=int(manifest["video_min_pixels"]),
            video_max_pixels=int(manifest["video_max_pixels"]),
        )
        media = CausalSourceAV(
            _easycom_directory_specs(session_dir),
            output_fps=config.video_fps,
            required_duration=float(manifest["cached_end_seconds"]),
            frame_cache_size=frame_cache_size,
            decoder_backend=decoder,
            decoder_gpu_id=decoder_gpu_id,
            decoder_threads=decoder_threads,
            max_open_readers=max_open_readers,
        )
        self.extractor = core.OpenStreamReplayExtractor(
            model, processor, config, fixed_grid_av=media
        )
        self.step_seconds = config.step_seconds
        self.block_seconds = config.native_block_seconds
        self.steps_per_block = int(source["static"]["steps_per_block"])
        if audio_feature_device == "model":
            try:
                audio_feature_device = next(model.parameters()).device
            except StopIteration as exc:
                raise RuntimeError("Online-tail AV model has no parameters") from exc
        self.audio_feature_device = audio_feature_device
        self.audio_right_padding_samples = audio_right_padding_samples
        self.calls = 0
        self.online_av_calls = 0
        self.total_seconds = 0.0
        self.online_av_seconds = 0.0

    def reset_timing(self) -> None:
        self.calls = 0
        self.online_av_calls = 0
        self.total_seconds = 0.0
        self.online_av_seconds = 0.0

    @torch.no_grad()
    def __call__(self, item: dict[str, Any]) -> dict[str, Any]:
        started = time.perf_counter()
        source_index = int(item["source_index"])
        source = self.dataset.sources[source_index]
        anchor = item["anchor"]
        tail_steps = int(anchor["tail_steps"])
        current_piece = None
        if tail_steps < self.steps_per_block:
            av_started = time.perf_counter()
            block_start = int(anchor["block_index"]) * self.block_seconds
            interval_end = block_start + tail_steps * self.step_seconds
            pieces = self.extractor._prepare_pieces(
                video_path=None,
                audio_path=None,
                interval=(block_start, interval_end),
                use_audio_in_video=True,
                audio_feature_device=self.audio_feature_device,
                audio_right_padding_samples=self.audio_right_padding_samples,
            )
            closure_embeddings = self.model.get_input_embeddings()(
                pieces.closure_ids
            )
            current_piece = pieces_to_cache(
                pieces,
                closure_embeddings=closure_embeddings,
                floating_dtype=torch.bfloat16,
            )
            self.online_av_seconds += time.perf_counter() - av_started
            self.online_av_calls += 1
        embeddings, positions = reconstruct_open_stream_sequence(
            store=source["store"],
            static=source["static"],
            anchor=anchor,
            current_piece=current_piece,
        )
        batch = collate_cached_sequences(
            [
                {
                    "inputs_embeds": embeddings,
                    "position_ids": positions,
                    "targets": item["targets"],
                    "step_index": item["step_index"],
                }
            ]
        )
        self.total_seconds += time.perf_counter() - started
        self.calls += 1
        return batch


class AsyncTailBatchProducer:
    """Prepare online-tail batches on a second GPU behind a bounded queue."""

    def __init__(
        self,
        *,
        loader: DataLoader,
        materializer: OnlineCurrentTailMaterializer,
        total_batches: int,
        prefetch_depth: int,
    ) -> None:
        if total_batches <= 0:
            raise ValueError("total_batches must be positive")
        if prefetch_depth <= 0:
            raise ValueError("--prefetch-depth must be positive")
        self.loader = loader
        self.materializer = materializer
        self.total_batches = int(total_batches)
        self.queue: queue.Queue[Any] = queue.Queue(maxsize=int(prefetch_depth))
        self.stop_event = threading.Event()
        self.thread = threading.Thread(
            target=self._run,
            name="omni-av-tail-producer",
            daemon=True,
        )
        self.produced = 0
        self.wait_seconds = 0.0

    def _put(self, value: Any) -> bool:
        while not self.stop_event.is_set():
            try:
                self.queue.put(value, timeout=0.1)
                return True
            except queue.Full:
                continue
        return False

    def _run(self) -> None:
        try:
            iterator = iter(self.loader)
            for _ in range(self.total_batches):
                if self.stop_event.is_set():
                    return
                try:
                    deferred = next(iterator)
                except StopIteration:
                    iterator = iter(self.loader)
                    deferred = next(iterator)
                prepared = self.materializer(deferred)
                if not self._put(prepared):
                    return
                self.produced += 1
        except BaseException as error:
            self._put(error)

    def start(self) -> None:
        self.thread.start()

    def next_batch(self) -> dict[str, Any]:
        started = time.perf_counter()
        value = self.queue.get()
        self.wait_seconds += time.perf_counter() - started
        if isinstance(value, BaseException):
            raise RuntimeError("Asynchronous AV producer failed") from value
        if not isinstance(value, dict):
            raise TypeError(f"Producer returned {type(value)!r}, expected dict")
        return value

    def close(self) -> None:
        self.stop_event.set()
        self.thread.join(timeout=30.0)
        if self.thread.is_alive():
            raise RuntimeError("Asynchronous AV producer did not stop")


def _load_initial_head(
    checkpoint: Path | None,
    *,
    hidden_size: int,
    dropout: float,
    verbose: bool = True,
) -> 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(**config)
    head.load_state_dict(payload.get("head", payload))
    if verbose:
        print(f"Initialized classification head from {checkpoint}")
    return head


def _prepare_language_lora(
    model: torch.nn.Module,
    *,
    adapter_checkpoint: Path | None,
    rank: int,
    alpha: int,
    dropout: float,
    target_modules: Sequence[str],
    last_n_layers: int,
    gradient_checkpointing: bool,
    verbose: bool = True,
) -> torch.nn.Module:
    try:
        from peft import LoraConfig, PeftModel, 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
    if adapter_checkpoint is not None:
        adapter_checkpoint = adapter_checkpoint.expanduser().resolve()
        config_path = adapter_checkpoint / "adapter_config.json"
        if not config_path.is_file():
            raise FileNotFoundError(
                f"Missing PEFT adapter config: {config_path}"
            )
        language_model = PeftModel.from_pretrained(
            language_model,
            str(adapter_checkpoint),
            is_trainable=True,
        )
        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()
        if verbose:
            language_model.print_trainable_parameters()
            trainable_lora = sum(
                1
                for name, parameter in language_model.named_parameters()
                if parameter.requires_grad and "lora_" in name
            )
            print(
                f"Continued trainable LoRA from {adapter_checkpoint}; "
                f"trainable LoRA tensors={trainable_lora}"
            )
        return language_model

    matched_all = [
        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_all:
        raise RuntimeError(
            "No language-model linear layers matched LoRA targets "
            f"{list(target_modules)}"
        )
    layer_indices = sorted(
        {
            int(match.group(1))
            for name in matched_all
            if (match := _TRANSFORMER_LAYER_PATTERN.search(name)) is not None
        }
    )
    selected_layer_indices: list[int] | None = None
    matched = matched_all
    if last_n_layers:
        if not layer_indices:
            raise RuntimeError(
                "Could not identify Thinker layer indices in matched LoRA modules"
            )
        if last_n_layers > len(layer_indices):
            raise ValueError(
                f"--lora-last-n-layers={last_n_layers} exceeds the "
                f"{len(layer_indices)} matched Thinker layers"
            )
        selected_layer_indices = layer_indices[-last_n_layers:]
        selected = set(selected_layer_indices)
        matched = [
            name
            for name in matched_all
            if (
                (match := _TRANSFORMER_LAYER_PATTERN.search(name)) is not None
                and int(match.group(1)) in selected
            )
        ]
        if not matched:
            raise RuntimeError(
                "Final-layer selection removed every matched LoRA module"
            )
    # PEFT's layers_pattern matcher requires additional nesting before the
    # ModuleList name in some versions.  The extracted Omni Thinker exposes
    # modules directly as ``layers.28...``, so layers_to_transform can reject
    # every otherwise valid target.  Explicit full module paths use PEFT's
    # exact/suffix-list matching and are stable for this submodel hierarchy.
    peft_target_modules = matched if selected_layer_indices is not None else list(
        target_modules
    )
    config = LoraConfig(
        r=rank,
        lora_alpha=alpha,
        lora_dropout=dropout,
        target_modules=peft_target_modules,
        bias="none",
    )
    language_model = get_peft_model(language_model, config)
    targeted_module_names = getattr(
        getattr(language_model, "base_model", None), "targeted_module_names", None
    )
    if targeted_module_names is not None:
        actual = list(targeted_module_names)
        missing = [
            expected
            for expected in matched
            if not any(
                name == expected or name.endswith(f".{expected}") for name in actual
            )
        ]
        if missing:
            raise RuntimeError(
                "PEFT did not inject all selected LoRA modules; first missing="
                f"{missing[0]}, missing={len(missing)}, targeted={len(actual)}"
            )
    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()
    if verbose:
        language_model.print_trainable_parameters()
        print(
            f"LoRA matched {len(matched)} language linear modules; first={matched[0]}, "
            f"last={matched[-1]}"
        )
        if selected_layer_indices is not None:
            print(
                "LoRA layer restriction: "
                f"{selected_layer_indices[0]}--{selected_layer_indices[-1]} "
                f"({len(selected_layer_indices)} of {len(layer_indices)} layers)"
            )
    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,
    distributed_run: DistributedRun,
    materialize_batch: Any | None = None,
) -> float:
    language_was_training = language_model.training
    head_was_training = head.training
    language_model.eval()
    head.eval()
    loss_sum = 0.0
    example_count = 0
    for batch_index, batch in enumerate(loader):
        if max_batches > 0 and batch_index >= max_batches:
            break
        if materialize_batch is not None:
            batch = materialize_batch(batch)
        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())
        batch_size = int(targets.shape[0])
        loss_sum += float(egospeak_multihorizon_loss(logits, targets).item()) * batch_size
        example_count += batch_size
    language_model.train(language_was_training)
    head.train(head_was_training)
    totals = torch.tensor(
        [loss_sum, float(example_count)], dtype=torch.float64, device=device
    )
    if distributed_run.enabled:
        dist.all_reduce(totals, op=dist.ReduceOp.SUM)
    if float(totals[1].item()) <= 0:
        raise RuntimeError("Evaluation loader produced no examples")
    return float((totals[0] / totals[1]).item())


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 _save_weight_checkpoint(
    checkpoint_dir: Path,
    *,
    raw_language_model: torch.nn.Module,
    raw_head: TurnPredictionHead,
    args: argparse.Namespace,
    hidden_size: int,
    training_step: int,
) -> None:
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    if not args.freeze_thinker:
        raw_language_model.save_pretrained(checkpoint_dir / "adapter")
    torch.save(
        {
            "head": raw_head.state_dict(),
            "head_config": {
                "hidden_size": hidden_size,
                "num_horizons": NUM_HORIZONS,
                "num_classes": NUM_CLASSES,
                "dropout": raw_head.dropout,
                "num_input_layers": raw_head.num_input_layers,
                "layer_fusion": raw_head.layer_fusion,
                "learnable_layer_mixture": raw_head.learnable_layer_mixture,
                "transition_auxiliary": raw_head.transition_auxiliary,
            },
            "initial_head_checkpoint": (
                str(args.head_checkpoint.resolve()) if args.head_checkpoint else None
            ),
            "head_frozen": bool(args.freeze_head),
            "training_step": training_step,
        },
        checkpoint_dir / "turn_head.pt",
    )


def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
    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.save_every_steps < 0:
        raise ValueError("--save-every-steps must be non-negative")
    if args.online_current_tail:
        if run.enabled:
            raise ValueError(
                "DDP currently supports fully cached AV training only; remove "
                "--online-current-tail when launching with torchrun"
            )
        if args.easycom_session_dir is None:
            raise ValueError(
                "--easycom-session-dir is required by --online-current-tail"
            )
        if len(args.cache_dir) != 1:
            raise ValueError("Online-tail benchmark requires exactly one --cache-dir")
        if args.batch_size != 1:
            raise ValueError("Online-tail benchmark currently requires --batch-size 1")
        if args.num_workers != 0:
            raise ValueError("Online-tail benchmark currently requires --num-workers 0")
    if args.producer_device is not None:
        if not args.online_current_tail:
            raise ValueError("--producer-device requires --online-current-tail")
        if args.producer_device == args.device:
            raise ValueError("Producer and Thinker trainer must use different devices")
        if not args.producer_device.startswith("cuda:"):
            raise ValueError("--producer-device must be a CUDA device such as cuda:1")
    if args.prefetch_depth <= 0:
        raise ValueError("--prefetch-depth must be positive")
    if args.ddp_timeout_minutes <= 0:
        raise ValueError("--ddp-timeout-minutes must be positive")
    if args.adapter_checkpoint is not None and args.head_checkpoint is None:
        if run.is_main:
            print(
                "Warning: --adapter-checkpoint was provided without "
                "--head-checkpoint; the turn head will be reinitialized."
            )
    if args.freeze_thinker and args.adapter_checkpoint is not None:
        raise ValueError(
            "--freeze-thinker cannot be combined with --adapter-checkpoint"
        )
    if args.freeze_thinker and args.freeze_head:
        raise ValueError(
            "--freeze-thinker and --freeze-head would leave no trainable parameters"
        )
    if args.freeze_head and args.head_checkpoint is None:
        raise ValueError("--freeze-head requires --head-checkpoint")
    if args.lora_last_n_layers < 0:
        raise ValueError("--lora-last-n-layers must be non-negative")
    if args.adapter_checkpoint is not None and args.lora_last_n_layers:
        raise ValueError(
            "Layer selection is stored in an existing adapter config; omit "
            "--lora-last-n-layers when using --adapter-checkpoint"
        )
    if args.freeze_thinker and args.lora_last_n_layers:
        raise ValueError(
            "--lora-last-n-layers has no effect with --freeze-thinker"
        )
    if args.freeze_thinker and args.gradient_checkpointing:
        raise ValueError(
            "--gradient-checkpointing is unnecessary with --freeze-thinker"
        )
    if (
        args.audio_right_padding_samples is not None
        and args.audio_right_padding_samples < 0
    ):
        raise ValueError("--audio-right-padding-samples must be non-negative")
    if not args.online_current_tail and (
        args.audio_feature_device is not None
        or args.audio_right_padding_samples is not None
    ):
        raise ValueError(
            "Audio preprocessing overrides apply only with --online-current-tail"
        )
    if args.num_workers != 0 and run.is_main:
        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" / "adapter_config.json",
        output_dir / "turn_head.pt",
        output_dir / "train_summary.json",
    )
    intermediate_outputs = list(output_dir.glob("checkpoint-step-*"))
    if not args.overwrite_output and (
        any(path.exists() for path in known_outputs) or intermediate_outputs
    ):
        raise FileExistsError(
            f"Training output already exists in {output_dir}; pass --overwrite-output"
        )
    if run.is_main:
        output_dir.mkdir(parents=True, exist_ok=True)
    _distributed_barrier(run)
    _seed_everything(args.seed)
    core._configure_qwen_logging()
    device = run.device

    dataset = OmniTurnCacheDataset(
        args.cache_dir,
        target_override=args.target,
        max_examples=args.max_examples,
        seed=args.seed,
        block_cache_size=args.block_cache_size,
        defer_current_tail=args.online_current_tail,
    )
    collate_fn = (
        _collate_deferred_tail
        if args.online_current_tail
        else collate_cached_sequences
    )
    generator = torch.Generator().manual_seed(args.seed + run.rank)
    train_sampler: DistributedSampler | None = None
    eval_sampler: DistributedSampler | None = None
    if run.enabled:
        train_sampler = DistributedSampler(
            dataset,
            num_replicas=run.world_size,
            rank=run.rank,
            shuffle=True,
            seed=args.seed,
            drop_last=False,
        )
        eval_sampler = DistributedSampler(
            dataset,
            num_replicas=run.world_size,
            rank=run.rank,
            shuffle=False,
            drop_last=False,
        )
    train_loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=train_sampler is None,
        sampler=train_sampler,
        num_workers=args.num_workers,
        pin_memory=True,
        persistent_workers=args.num_workers > 0,
        collate_fn=collate_fn,
        generator=generator,
    )
    eval_loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        sampler=eval_sampler,
        num_workers=0,
        pin_memory=True,
        collate_fn=collate_fn,
    )
    if run.is_main:
        print(
            f"Cached training examples: {len(dataset)}, hidden_size={dataset.hidden_size}, "
            f"world={run.world_size}, per_gpu_batch={args.batch_size}, "
            f"accum={args.gradient_accumulation_steps}, "
            "global_batch="
            f"{run.world_size * args.batch_size * args.gradient_accumulation_steps}"
        )

    base_model, processor = _load_base_model(args, run)
    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")
    if args.freeze_thinker:
        base_model.requires_grad_(False)
        language_model = base_model.model
        language_model.config.use_cache = False
        language_model.eval()
        if run.is_main:
            print(
                "Frozen-Thinker control: no LoRA adapter; only the turn head "
                "is trainable"
            )
    else:
        language_model = _prepare_language_lora(
            base_model,
            adapter_checkpoint=args.adapter_checkpoint,
            rank=args.lora_rank,
            alpha=args.lora_alpha,
            dropout=args.lora_dropout,
            target_modules=target_modules,
            last_n_layers=args.lora_last_n_layers,
            gradient_checkpointing=args.gradient_checkpointing,
            verbose=run.is_main,
        )
    materialize_batch: OnlineCurrentTailMaterializer | None = None
    producer_model: torch.nn.Module | None = None
    if args.online_current_tail:
        assert args.easycom_session_dir is not None
        materializer_model = base_model
        materializer_processor = processor
        if args.producer_device is not None:
            producer_model, materializer_processor = core.load_qwen_omni(
                args.model_id,
                dtype=args.dtype,
                device_map=args.producer_device,
                use_flash_attention_2=args.flash_attention_2,
            )
            producer_model.requires_grad_(False)
            materializer_model = producer_model
        materialize_batch = OnlineCurrentTailMaterializer(
            dataset=dataset,
            model=materializer_model,
            processor=materializer_processor,
            session_dir=args.easycom_session_dir,
            frame_cache_size=args.frame_cache_size,
            decoder=args.decoder,
            decoder_gpu_id=args.decoder_gpu_id,
            decoder_threads=args.decoder_threads,
            max_open_readers=args.max_open_readers,
            audio_feature_device=args.audio_feature_device,
            audio_right_padding_samples=args.audio_right_padding_samples,
        )
        print(
            "Cache policy benchmark: complete 2s history cached; "
            "current 0.2--1.8s tail computed online"
        )
        print(
            "Online-tail audio preprocessing: "
            f"feature_device={args.audio_feature_device or 'processor default'}, "
            "padding="
            + (
                "model maximum"
                if args.audio_right_padding_samples is None
                else f"actual + >= {args.audio_right_padding_samples} samples"
            )
        )
        if args.producer_device is not None:
            print(
                f"Producer-consumer enabled: frozen AV={args.producer_device}, "
                f"Thinker LoRA={device}, queue={args.prefetch_depth}"
            )
            # The consumer retains the same language object through its PEFT
            # wrapper, so the outer multimodal container and towers can go.
            del base_model
            del processor
            gc.collect()
            torch.cuda.empty_cache()
    else:
        # All AV phases are cached, so the modality towers can be discarded.
        del base_model
        del processor
        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,
        verbose=run.is_main,
    ).to(device=device, dtype=torch.float32)
    if args.freeze_head:
        head.requires_grad_(False)
        head.eval()
        if run.is_main:
            print("Frozen-head control: only Thinker LoRA parameters are trainable")

    raw_language_model = language_model
    raw_head = head
    synchronization_modules = tuple(
        module
        for module, trainable in (
            (raw_language_model, not args.freeze_thinker),
            (raw_head, not args.freeze_head),
        )
        if trainable
    )
    _synchronize_trainable_parameters(synchronization_modules, run)
    if not args.freeze_thinker:
        language_model = _wrap_ddp(raw_language_model, run)
    if not args.freeze_head:
        head = _wrap_ddp(raw_head, run)
    # Model initialization is synchronized above; rank-specific RNG streams
    # avoid identical dropout masks for different distributed samples.
    _seed_everything(args.seed + run.rank)

    lora_parameters = [
        parameter
        for parameter in raw_language_model.parameters()
        if parameter.requires_grad
    ]
    if not args.freeze_thinker and not lora_parameters:
        raise RuntimeError("LoRA preparation produced no trainable parameters")
    if args.freeze_thinker:
        optimizer = torch.optim.AdamW(
            [{"params": raw_head.parameters(), "lr": args.head_learning_rate}],
            weight_decay=args.weight_decay,
        )
        head_group_index = 0
    elif args.freeze_head:
        optimizer = torch.optim.AdamW(
            [{"params": lora_parameters, "lr": args.learning_rate}],
            weight_decay=args.weight_decay,
        )
        head_group_index = None
    else:
        optimizer = torch.optim.AdamW(
            [
                {"params": lora_parameters, "lr": args.learning_rate},
                {"params": raw_head.parameters(), "lr": args.head_learning_rate},
            ],
            weight_decay=args.weight_decay,
        )
        head_group_index = 1
    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,
        distributed_run=run,
        materialize_batch=materialize_batch,
    )
    if run.is_main:
        print(f"Initial cached loss: {initial_loss:.6f}")
    if materialize_batch is not None:
        materialize_batch.reset_timing()
    if args.freeze_thinker:
        language_model.eval()
    else:
        language_model.train()
    if args.freeze_head:
        head.eval()
    else:
        head.train()
    optimizer.zero_grad(set_to_none=True)
    _distributed_barrier(run)
    torch.cuda.synchronize(device)
    training_started = time.perf_counter()
    batch_producer: AsyncTailBatchProducer | None = None
    if materialize_batch is not None and args.producer_device is not None:
        batch_producer = AsyncTailBatchProducer(
            loader=train_loader,
            materializer=materialize_batch,
            total_batches=args.max_steps * args.gradient_accumulation_steps,
            prefetch_depth=args.prefetch_depth,
        )
        batch_producer.start()
    sampler_epoch = 0
    if train_sampler is not None:
        train_sampler.set_epoch(sampler_epoch)
    iterator = iter(train_loader) if batch_producer is None else None
    history: list[dict[str, float | int]] = []
    saved_checkpoint_steps: list[int] = []
    progress = tqdm(
        total=args.max_steps,
        desc=("Frozen Thinker head" if args.freeze_thinker else "Thinker LoRA"),
        unit="update",
        disable=args.no_progress or not run.is_main,
    )
    for update_step in range(1, args.max_steps + 1):
        accumulated_loss = 0.0
        for accumulation_index in range(args.gradient_accumulation_steps):
            if batch_producer is not None:
                batch = batch_producer.next_batch()
            else:
                assert iterator is not None
                try:
                    batch = next(iterator)
                except StopIteration:
                    sampler_epoch += 1
                    if train_sampler is not None:
                        train_sampler.set_epoch(sampler_epoch)
                    iterator = iter(train_loader)
                    batch = next(iterator)
            if materialize_batch is not None and batch_producer is None:
                batch = materialize_batch(batch)
            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)
            synchronize_gradients = (
                accumulation_index == args.gradient_accumulation_steps - 1
            )
            with ExitStack() as stack:
                if run.enabled and not synchronize_gradients:
                    if not args.freeze_thinker:
                        assert isinstance(language_model, DDP)
                        stack.enter_context(language_model.no_sync())
                    if not args.freeze_head:
                        assert isinstance(head, DDP)
                        stack.enter_context(head.no_sync())
                if args.freeze_thinker:
                    with torch.no_grad():
                        states = forward_last_state(
                            language_model,
                            inputs_embeds=inputs_embeds,
                            position_ids=position_ids,
                            attention_mask=attention_mask,
                            readout_indices=readout_indices,
                        )
                else:
                    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 = (
            list(raw_head.parameters())
            if args.freeze_thinker
            else (
                lora_parameters
                if args.freeze_head
                else [*lora_parameters, *raw_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)
        reported_loss = _distributed_average(accumulated_loss, run)
        lora_lr = (
            0.0 if args.freeze_thinker else float(optimizer.param_groups[0]["lr"])
        )
        head_lr = (
            0.0
            if head_group_index is None
            else float(optimizer.param_groups[head_group_index]["lr"])
        )
        record = {
            "step": update_step,
            "loss": reported_loss,
            "grad_norm": float(grad_norm),
            "lora_lr": lora_lr,
            "head_lr": head_lr,
        }
        if run.is_main:
            history.append(record)
            progress.update(1)
            progress.set_postfix(
                loss=f"{reported_loss:.4f}",
                lr=f"{head_lr:.2g}" if args.freeze_thinker else f"{lora_lr:.2g}",
            )
        if run.is_main and (update_step == 1 or update_step % args.log_every == 0):
            print(
                f"step={update_step} loss={reported_loss:.6f} "
                f"grad_norm={float(grad_norm):.4f} "
                + (
                    f"head_lr={head_lr:.3g}"
                    if args.freeze_thinker
                    else f"lora_lr={lora_lr:.3g} head_lr={head_lr:.3g}"
                ),
                flush=True,
            )
        if args.save_every_steps and update_step % args.save_every_steps == 0:
            torch.cuda.synchronize(device)
            _distributed_barrier(run)
            if run.is_main:
                checkpoint_dir = output_dir / f"checkpoint-step-{update_step}"
                _save_weight_checkpoint(
                    checkpoint_dir,
                    raw_language_model=raw_language_model,
                    raw_head=raw_head,
                    args=args,
                    hidden_size=dataset.hidden_size,
                    training_step=update_step,
                )
                saved_checkpoint_steps.append(update_step)
                print(f"Saved intermediate checkpoint to {checkpoint_dir}", flush=True)
            _distributed_barrier(run)
    progress.close()
    torch.cuda.synchronize(device)
    _distributed_barrier(run)
    training_seconds = time.perf_counter() - training_started
    if batch_producer is not None:
        batch_producer.close()
    online_timing = None
    if materialize_batch is not None:
        online_timing = {
            "materialize_calls": materialize_batch.calls,
            "materialize_seconds": materialize_batch.total_seconds,
            "online_av_calls": materialize_batch.online_av_calls,
            "online_av_seconds": materialize_batch.online_av_seconds,
            "consumer_queue_wait_seconds": (
                batch_producer.wait_seconds if batch_producer is not None else None
            ),
            "producer_batches": (
                batch_producer.produced if batch_producer is not None else None
            ),
        }

    final_loss = _evaluate_loss(
        language_model,
        head,
        eval_loader,
        device=device,
        model_dtype=model_dtype,
        max_batches=args.eval_max_batches,
        distributed_run=run,
        materialize_batch=materialize_batch,
    )
    local_peak_gib = torch.cuda.max_memory_allocated(device) / (1024 ** 3)
    peak_tensor = torch.tensor(local_peak_gib, dtype=torch.float64, device=device)
    if run.enabled:
        gathered_peaks = [torch.zeros_like(peak_tensor) for _ in range(run.world_size)]
        dist.all_gather(gathered_peaks, peak_tensor)
        per_rank_peak_gib = [float(value.item()) for value in gathered_peaks]
    else:
        per_rank_peak_gib = [local_peak_gib]

    if run.is_main:
        _save_weight_checkpoint(
            output_dir,
            raw_language_model=raw_language_model,
            raw_head=raw_head,
            args=args,
            hidden_size=dataset.hidden_size,
            training_step=args.max_steps,
        )
    global_batch_size = (
        run.world_size * args.batch_size * args.gradient_accumulation_steps
    )
    summary = {
        "schema_version": 1,
        "status": "pilot" if args.max_examples else "training",
        "model_id": args.model_id,
        "model_variant": (
            "frozen_base_thinker_head_only"
            if args.freeze_thinker
            else (
                "thinker_lora_frozen_head"
                if args.freeze_head
                else "thinker_lora"
            )
        ),
        "thinker_frozen": bool(args.freeze_thinker),
        "head_frozen": bool(args.freeze_head),
        "lora_last_n_layers": int(args.lora_last_n_layers),
        "initial_adapter_checkpoint": (
            str(args.adapter_checkpoint.expanduser().resolve())
            if args.adapter_checkpoint is not None
            else None
        ),
        "optimizer_state_resumed": False,
        "cache_dirs": [str(path.expanduser().resolve()) for path in args.cache_dir],
        "num_examples": len(dataset),
        "initial_loss": initial_loss,
        "final_loss": final_loss,
        "training_seconds": training_seconds,
        "updates_per_second": args.max_steps / training_seconds,
        "estimated_examples_per_second": (
            args.max_steps * global_batch_size / training_seconds
        ),
        "distributed": {
            "enabled": run.enabled,
            "world_size": run.world_size,
            "per_gpu_batch_size": args.batch_size,
            "gradient_accumulation_steps": args.gradient_accumulation_steps,
            "global_batch_size": global_batch_size,
            "per_rank_peak_cuda_allocated_gib": per_rank_peak_gib,
        },
        "online_tail_timing": online_timing,
        "loss_reduction_ratio": (
            (initial_loss - final_loss) / initial_loss if initial_loss else None
        ),
        "peak_cuda_allocated_gib": max(per_rank_peak_gib),
        "producer_peak_cuda_allocated_gib": (
            torch.cuda.max_memory_allocated(torch.device(args.producer_device))
            / (1024 ** 3)
            if args.producer_device is not None
            else None
        ),
        "packages": {
            "torch": torch.__version__,
            "transformers": _package_version("transformers"),
            "peft": _package_version("peft"),
            "ms-swift": _package_version("ms-swift"),
        },
        "args": _jsonable_args(args),
        "history": history,
        "intermediate_checkpoint_steps": saved_checkpoint_steps,
    }
    if run.is_main:
        save_json(output_dir / "train_summary.json", summary)
        print(f"Final cached loss: {final_loss:.6f} (initial {initial_loss:.6f})")
        print(
            f"Training speed: {training_seconds:.3f}s / {args.max_steps} updates = "
            f"{args.max_steps / training_seconds:.3f} update/s; "
            f"estimated {summary['estimated_examples_per_second']:.2f} examples/s"
        )
        if online_timing is not None:
            print(
                "Online current-tail AV: "
                f"{online_timing['online_av_seconds']:.3f}s / "
                f"{online_timing['online_av_calls']} calls"
            )
            if online_timing["consumer_queue_wait_seconds"] is not None:
                print(
                    "Consumer queue wait: "
                    f"{online_timing['consumer_queue_wait_seconds']:.3f}s; "
                    f"producer batches={online_timing['producer_batches']}"
                )
        if args.freeze_thinker:
            print("Frozen Thinker: no LoRA adapter was created")
        else:
            print(f"Saved LoRA adapter to {output_dir / 'adapter'}")
        print(f"Saved turn head to {output_dir / 'turn_head.pt'}")
        print(
            "Peak CUDA allocated across ranks: "
            f"{summary['peak_cuda_allocated_gib']:.2f} GiB"
        )
        if summary["producer_peak_cuda_allocated_gib"] is not None:
            print(
                "Producer peak CUDA allocated: "
                f"{summary['producer_peak_cuda_allocated_gib']:.2f} GiB"
            )
    _distributed_barrier(run)


def main(argv: Sequence[str] | None = None) -> None:
    args = parse_args(argv)
    run = _initialize_distributed(args)
    try:
        _run_training(args, run)
    finally:
        if run.enabled and dist.is_initialized():
            dist.destroy_process_group()


if __name__ == "__main__":
    main()