#!/usr/bin/env python3 """Train Qwen2.5-Omni Thinker LoRA + EgoSpeak 10x3 head from AV caches. This is a custom task runner intended to execute inside the ms-swift image. It uses the PEFT library shipped with ms-swift because Swift's normal SFT CLI optimizes generated text tokens, whereas this experiment reads one continuous Thinker hidden state and applies a multi-horizon classification loss. The cached audio and vision towers remain frozen. The reconstructed sequence is forwarded with ``use_cache=False`` so all retained 2-second blocks participate in autograd through the Thinker LoRA adapters. """ from __future__ import annotations import argparse import gc import importlib.metadata import inspect import json import math import os import queue import random 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", ) @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("--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("--head-dropout", type=float, default=0.1) parser.add_argument("--gradient-checkpointing", action="store_true") parser.add_argument("--seed", type=int, default=1337) parser.add_argument("--log-every", type=int, default=5) parser.add_argument("--eval-max-batches", type=int, default=4) parser.add_argument( "--dtype", choices=("bfloat16", "float16"), default="bfloat16", ) parser.add_argument( "--device", default="cuda:0", 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, *, rank: int, alpha: int, dropout: float, target_modules: Sequence[str], gradient_checkpointing: bool, verbose: bool = True, ) -> torch.nn.Module: try: from peft import LoraConfig, get_peft_model except ImportError as exc: raise RuntimeError( "PEFT is required. The ms-swift 3.12.5 image should include it." ) from exc model.requires_grad_(False) language_model = model.model matched = [ name for name, module in language_model.named_modules() if isinstance(module, nn.Linear) and name.rsplit(".", 1)[-1] in target_modules ] if not matched: raise RuntimeError( "No language-model linear layers matched LoRA targets " f"{list(target_modules)}" ) config = LoraConfig( r=rank, lora_alpha=alpha, lora_dropout=dropout, target_modules=list(target_modules), bias="none", ) language_model = get_peft_model(language_model, config) language_model.config.use_cache = False if gradient_checkpointing: try: language_model.gradient_checkpointing_enable( gradient_checkpointing_kwargs={"use_reentrant": False} ) except TypeError: language_model.gradient_checkpointing_enable() if verbose: language_model.print_trainable_parameters() print( f"LoRA matched {len(matched)} language layers; first={matched[0]}, " f"last={matched[-1]}" ) return language_model def _lr_lambda(step: int, *, total_steps: int, warmup_steps: int) -> float: if warmup_steps > 0 and step < warmup_steps: return float(step + 1) / float(warmup_steps) remaining = max(1, total_steps - warmup_steps) progress = min(1.0, max(0.0, (step - warmup_steps) / remaining)) return 0.5 * (1.0 + math.cos(math.pi * progress)) @torch.inference_mode() def _evaluate_loss( language_model: torch.nn.Module, head: TurnPredictionHead, loader: DataLoader, *, device: torch.device, model_dtype: torch.dtype, max_batches: int, 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 _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.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.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", ) if not args.overwrite_output and any(path.exists() for path in known_outputs): raise FileExistsError( f"Training output already exists in {output_dir}; pass --overwrite-output" ) 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") language_model = _prepare_language_lora( base_model, rank=args.lora_rank, alpha=args.lora_alpha, dropout=args.lora_dropout, target_modules=target_modules, gradient_checkpointing=args.gradient_checkpointing, 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) raw_language_model = language_model raw_head = head _synchronize_trainable_parameters((raw_language_model, raw_head), run) language_model = _wrap_ddp(raw_language_model, run) 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 lora_parameters: raise RuntimeError("LoRA preparation produced no trainable parameters") 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, ) 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() language_model.train() 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]] = [] progress = tqdm( total=args.max_steps, desc="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: assert isinstance(language_model, DDP) assert isinstance(head, DDP) stack.enter_context(language_model.no_sync()) stack.enter_context(head.no_sync()) states = forward_last_state( language_model, inputs_embeds=inputs_embeds, position_ids=position_ids, attention_mask=attention_mask, readout_indices=readout_indices, ) logits = head(states.float()) loss = egospeak_multihorizon_loss(logits, targets) (loss / args.gradient_accumulation_steps).backward() accumulated_loss += float(loss.detach()) / args.gradient_accumulation_steps trainable_parameters = [*lora_parameters, *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) record = { "step": update_step, "loss": reported_loss, "grad_norm": float(grad_norm), "lora_lr": float(optimizer.param_groups[0]["lr"]), "head_lr": float(optimizer.param_groups[1]["lr"]), } if run.is_main: history.append(record) progress.update(1) progress.set_postfix( loss=f"{reported_loss:.4f}", lr=f"{optimizer.param_groups[0]['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"lora_lr={optimizer.param_groups[0]['lr']:.3g}", flush=True, ) 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] adapter_dir = output_dir / "adapter" if run.is_main: raw_language_model.save_pretrained(adapter_dir) torch.save( { "head": raw_head.state_dict(), "head_config": { "hidden_size": dataset.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 ), "training_step": args.max_steps, }, output_dir / "turn_head.pt", ) 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, "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, } 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']}" ) print(f"Saved LoRA adapter to {adapter_dir}") 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()