#!/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( 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)) 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, }, "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()