diff --git a/cache_omni_av_pieces.py b/cache_omni_av_pieces.py new file mode 100644 index 0000000..1343b1e --- /dev/null +++ b/cache_omni_av_pieces.py @@ -0,0 +1,426 @@ +#!/usr/bin/env python3 +"""Cache exact Qwen2.5-Omni AV pieces for Thinker LoRA training. + +For every source-aligned 2-second block this script caches ten causal phases: +0.2, 0.4, ..., 2.0 seconds. The 2.0-second phase is reused as completed +history; phases 0.2--1.8 are used only for the current partial block. No +future samples enter a phase. + +The optional verification compares a full-sequence reconstruction from the +cache with the existing open-stream KV replay. Run this 24-second pilot before +caching a full EasyCom session. +""" + +from __future__ import annotations + +import argparse +import math +import platform +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +import torch +import torch.nn.functional as F +from tqdm.auto import tqdm + +import streaming_turn_model as core +from easycom_head import NUM_HORIZONS +from omni_cached_av import ( + SCHEMA_VERSION, + CachedBlockStore, + anchor_intervals, + cache_block_filename, + cache_size_bytes, + forward_last_state, + make_anchor, + pieces_to_cache, + reconstruct_open_stream_sequence, + save_json, +) +from streaming_turn_model_directly import CausalSourceAV, _easycom_directory_specs + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-id", required=True) + parser.add_argument("--easycom-session-dir", type=Path, required=True) + parser.add_argument("--target", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument( + "--end", + type=float, + default=24.0, + help="Dense session endpoint in seconds; use 24 for the first pilot", + ) + parser.add_argument("--step-seconds", type=float, default=0.2) + parser.add_argument("--native-block-seconds", type=float, default=2.0) + parser.add_argument("--memory-seconds", type=float, default=8.0) + parser.add_argument("--video-fps", type=float, default=2.0) + parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28) + parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28) + parser.add_argument("--frame-cache-size", type=int, default=64) + parser.add_argument("--decoder", choices=("decord", "pynvcodec"), default="decord") + parser.add_argument("--decoder-gpu-id", type=int, default=0) + parser.add_argument("--decoder-threads", type=int, default=1) + parser.add_argument("--max-open-readers", type=int, default=1) + parser.add_argument( + "--dtype", + choices=("auto", "bfloat16", "float16", "float32"), + default="bfloat16", + ) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--flash-attention-2", action="store_true") + parser.add_argument( + "--verify-every", + type=int, + default=10, + help="Compare every Nth 5-Hz state with exact KV replay; 0 disables", + ) + parser.add_argument( + "--minimum-cosine", + type=float, + default=0.999, + help="Fail if cached full-forward/reference KV cosine is below this", + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Overwrite this cache's known block/static/manifest files", + ) + return parser.parse_args(argv) + + +def _integer_ratio(numerator: float, denominator: float, name: str) -> int: + ratio = numerator / denominator + rounded = round(ratio) + if not math.isclose(ratio, rounded, rel_tol=0.0, abs_tol=1e-6): + raise ValueError(f"{name} must be integral, got {ratio}") + return int(rounded) + + +def _load_targets(path: Path) -> np.ndarray: + targets = np.load(path, mmap_mode="r") + if targets.ndim != 2 or targets.shape[1] != 3: + raise ValueError(f"Expected target [T,3], got {targets.shape}: {path}") + return targets + + +def _safe_output_path(path: Path, overwrite: bool) -> None: + if path.exists() and not overwrite: + raise FileExistsError(f"Refusing to overwrite {path}; pass --overwrite") + + +def _model_device(model: torch.nn.Module) -> torch.device: + try: + return next(model.parameters()).device + except StopIteration as exc: + raise RuntimeError("Model has no parameters") from exc + + +@torch.inference_mode() +def _cache_blocks( + *, + extractor: core.OpenStreamReplayExtractor, + model: torch.nn.Module, + output_dir: Path, + num_steps: int, + steps_per_block: int, + step_seconds: float, + block_seconds: float, + overwrite: bool, +) -> list[dict[str, Any]]: + num_blocks = math.ceil(num_steps / steps_per_block) + records: list[dict[str, Any]] = [] + input_embeddings = model.get_input_embeddings() + for block_index in tqdm(range(num_blocks), desc="cache AV blocks", unit="block"): + first_step = block_index * steps_per_block + phase_count = min(steps_per_block, num_steps - first_step) + block_start = block_index * block_seconds + variants: dict[str, dict[str, torch.Tensor]] = {} + token_lengths: dict[str, int] = {} + for tail_steps in range(1, phase_count + 1): + interval_end = block_start + tail_steps * step_seconds + pieces = extractor._prepare_pieces( + video_path=None, + audio_path=None, + interval=(block_start, interval_end), + use_audio_in_video=True, + ) + closure_embeddings = input_embeddings(pieces.closure_ids) + cached = pieces_to_cache( + pieces, + closure_embeddings=closure_embeddings, + floating_dtype=torch.bfloat16, + ) + variants[str(tail_steps)] = cached + token_lengths[str(tail_steps)] = int( + cached["prefix_embeddings"].shape[1] + + cached["payload_embeddings"].shape[1] + + cached["closure_embeddings"].shape[1] + ) + path = output_dir / cache_block_filename(block_index) + _safe_output_path(path, overwrite) + torch.save( + { + "schema_version": SCHEMA_VERSION, + "block_index": block_index, + "start_seconds": block_start, + "end_seconds": min( + block_start + block_seconds, num_steps * step_seconds + ), + "variants": variants, + }, + path, + ) + records.append( + { + "block_index": block_index, + "path": path.name, + "phase_count": phase_count, + "token_lengths_without_query": token_lengths, + "size_bytes": path.stat().st_size, + } + ) + return records + + +@torch.inference_mode() +def _verify_cache( + *, + model: torch.nn.Module, + extractor: core.OpenStreamReplayExtractor, + cache_dir: Path, + static: dict[str, Any], + anchors: list[dict[str, Any]], + verify_every: int, + minimum_cosine: float, + block_seconds: float, + step_seconds: float, +) -> list[dict[str, float]]: + if verify_every <= 0: + return [] + selected = [ + anchor + for anchor in anchors + if (int(anchor["step_index"]) + 1) % verify_every == 0 + ] + if anchors and anchors[-1] not in selected: + selected.append(anchors[-1]) + store = CachedBlockStore(cache_dir, max_blocks=max(2, int(static["memory_blocks"]))) + device = _model_device(model) + checks: list[dict[str, float]] = [] + for anchor in tqdm(selected, desc="verify cached replay", unit="state"): + embeddings, positions = reconstruct_open_stream_sequence( + store=store, static=static, anchor=anchor + ) + embeddings = embeddings.unsqueeze(0).to(device=device, dtype=model.dtype) + positions = positions.unsqueeze(1).to(device=device) + attention_mask = torch.ones( + 1, embeddings.shape[1], dtype=torch.long, device=device + ) + cached_state = forward_last_state( + model.model, + inputs_embeds=embeddings, + position_ids=positions, + attention_mask=attention_mask, + readout_indices=torch.tensor( + [embeddings.shape[1] - 1], dtype=torch.long, device=device + ), + ) + exact_state = extractor.state_from_scratch( + video_path=None, + audio_path=None, + intervals=anchor_intervals( + anchor, + block_seconds=block_seconds, + step_seconds=step_seconds, + ), + use_audio_in_video=True, + ) + cached_float = cached_state.float() + exact_float = exact_state.float() + check = { + "time_seconds": float(anchor["time_seconds"]), + "cosine": float( + F.cosine_similarity(cached_float, exact_float, dim=-1).item() + ), + "mean_abs": float((cached_float - exact_float).abs().mean().item()), + "max_abs": float((cached_float - exact_float).abs().max().item()), + } + checks.append(check) + print( + f"verify t={check['time_seconds']:.1f}s " + f"cosine={check['cosine']:.8f} " + f"mean_abs={check['mean_abs']:.6f} " + f"max_abs={check['max_abs']:.6f}", + flush=True, + ) + if check["cosine"] < minimum_cosine: + raise RuntimeError( + "Cached full-sequence reconstruction disagrees with open-stream " + f"KV replay at t={check['time_seconds']:.1f}s: " + f"cosine={check['cosine']:.8f} < {minimum_cosine:.8f}" + ) + return checks + + +def main(argv: Sequence[str] | None = None) -> None: + args = parse_args(argv) + core._configure_qwen_logging() + if args.end <= 0: + raise ValueError("--end must be positive") + steps_per_block = _integer_ratio( + args.native_block_seconds, args.step_seconds, "native block / step" + ) + memory_blocks = _integer_ratio( + args.memory_seconds, args.native_block_seconds, "memory / native block" + ) + num_steps = _integer_ratio(args.end, args.step_seconds, "end / step") + targets = _load_targets(args.target) + if num_steps > len(targets): + raise ValueError( + f"Requested {num_steps} cached steps but target has only {len(targets)}" + ) + + output_dir = args.output_dir.expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + manifest_path = output_dir / "manifest.json" + static_path = output_dir / "static.pt" + _safe_output_path(manifest_path, args.overwrite) + _safe_output_path(static_path, args.overwrite) + + config = core.StreamConfig( + step_seconds=args.step_seconds, + native_block_seconds=args.native_block_seconds, + memory_seconds=args.memory_seconds, + video_fps=args.video_fps, + video_min_pixels=args.video_min_pixels, + video_max_pixels=args.video_max_pixels, + ) + config.validate() + specs = _easycom_directory_specs(args.easycom_session_dir) + media = CausalSourceAV( + specs, + output_fps=args.video_fps, + required_duration=args.end, + frame_cache_size=args.frame_cache_size, + decoder_backend=args.decoder, + decoder_gpu_id=args.decoder_gpu_id, + decoder_threads=args.decoder_threads, + max_open_readers=args.max_open_readers, + ) + print( + f"Direct media ready: duration={args.end:.1f}s, " + f"grid={media.expected_grid_frames} frames, audio={len(media.audio)} samples" + ) + + model, processor = core.load_qwen_omni( + args.model_id, + dtype=args.dtype, + device_map=args.device, + use_flash_attention_2=args.flash_attention_2, + ) + model.requires_grad_(False) + extractor = core.OpenStreamReplayExtractor( + model, processor, config, fixed_grid_av=media + ) + with torch.inference_mode(): + query_ids = extractor.query_ids + query_embeddings = model.get_input_embeddings()(query_ids) + positions_per_chunk = core._tmrope_positions_per_chunk(model) + static = { + "schema_version": SCHEMA_VERSION, + "steps_per_block": steps_per_block, + "memory_blocks": memory_blocks, + "positions_per_chunk": positions_per_chunk, + "query_ids": query_ids.detach().cpu().long(), + "query_embeddings": query_embeddings.detach().cpu().to(torch.bfloat16), + } + torch.save(static, static_path) + + block_records = _cache_blocks( + extractor=extractor, + model=model, + output_dir=output_dir, + num_steps=num_steps, + steps_per_block=steps_per_block, + step_seconds=args.step_seconds, + block_seconds=args.native_block_seconds, + overwrite=args.overwrite, + ) + anchors = [ + make_anchor( + step_index=step_index, + steps_per_block=steps_per_block, + memory_blocks=memory_blocks, + step_seconds=args.step_seconds, + num_target_steps=len(targets), + num_horizons=NUM_HORIZONS, + ) + for step_index in range(num_steps) + ] + checks = _verify_cache( + model=model, + extractor=extractor, + cache_dir=output_dir, + static=static, + anchors=anchors, + verify_every=args.verify_every, + minimum_cosine=args.minimum_cosine, + block_seconds=args.native_block_seconds, + step_seconds=args.step_seconds, + ) + + manifest = { + "schema_version": SCHEMA_VERSION, + "cache_kind": "qwen2.5-omni-frozen-av-open-stream-phases", + "session": args.easycom_session_dir.resolve().name, + "model_id": args.model_id, + "target_path": str(args.target.expanduser().resolve()), + "target_steps": int(len(targets)), + "cached_steps": num_steps, + "cached_end_seconds": args.end, + "hidden_size": int(query_embeddings.shape[-1]), + "floating_dtype": "bfloat16", + "step_seconds": args.step_seconds, + "native_block_seconds": args.native_block_seconds, + "memory_seconds": args.memory_seconds, + "steps_per_block": steps_per_block, + "memory_blocks": memory_blocks, + "positions_per_chunk": positions_per_chunk, + "video_fps": args.video_fps, + "video_min_pixels": args.video_min_pixels, + "video_max_pixels": args.video_max_pixels, + "use_audio_in_video": True, + "interval_convention": "half_open_[start,end)", + "phase_policy": "cache_0.2_to_2.0s;_2.0s_is_completed_history", + "label_policy": "target[i+1:i+11]", + "static_path": static_path.name, + "blocks": block_records, + "anchors": anchors, + "verification": checks, + "media": media.metadata(), + "software": { + "python": platform.python_version(), + "torch": torch.__version__, + }, + } + save_json(manifest_path, manifest) + size = cache_size_bytes( + [static_path, *[output_dir / record["path"] for record in block_records]] + ) + print( + f"Saved {len(block_records)} blocks / {len(anchors)} anchors to {output_dir}" + ) + print(f"Cache tensor size: {size / (1024 ** 3):.3f} GiB") + if checks: + print( + "Verification: " + f"lowest cosine={min(item['cosine'] for item in checks):.8f}, " + f"worst max_abs={max(item['max_abs'] for item in checks):.6f}" + ) + + +if __name__ == "__main__": + main() diff --git a/omni_cached_av.py b/omni_cached_av.py new file mode 100644 index 0000000..5cb79c7 --- /dev/null +++ b/omni_cached_av.py @@ -0,0 +1,365 @@ +"""Cached open-stream AV pieces for Qwen2.5-Omni turn-taking LoRA. + +The modality towers are frozen for the first fine-tuning experiment. This +module serializes their exact output embeddings and Qwen's local TMRoPE +coordinates, then reconstructs the same still-open media turn without using a +KV cache. The resulting sequence can be passed through the Thinker with +``use_cache=False`` so gradients reach LoRA adapters in every retained block. +""" + +from __future__ import annotations + +import json +import math +from collections import OrderedDict +from pathlib import Path +from typing import Any, Iterable, Sequence + +import torch + + +SCHEMA_VERSION = 1 +PIECE_TENSOR_NAMES = ( + "prefix_embeddings", + "prefix_position_ids", + "payload_embeddings", + "payload_position_ids", + "closure_ids", + "closure_position_ids", + "closure_embeddings", +) + + +def load_json(path: str | Path) -> dict[str, Any]: + return json.loads(Path(path).read_text()) + + +def save_json(path: str | Path, payload: dict[str, Any]) -> None: + Path(path).write_text(json.dumps(payload, indent=2) + "\n") + + +def torch_load(path: str | Path) -> Any: + """Load tensor-only cache files across old and new PyTorch defaults.""" + + try: + return torch.load(Path(path), map_location="cpu", weights_only=True) + except TypeError: + return torch.load(Path(path), map_location="cpu") + + +def cache_block_filename(block_index: int) -> str: + return f"block_{block_index:06d}.pt" + + +def pieces_to_cache( + pieces: Any, + *, + closure_embeddings: torch.Tensor, + floating_dtype: torch.dtype = torch.bfloat16, +) -> dict[str, torch.Tensor]: + """Detach one ``OpenMediaPieces`` object into a portable CPU payload.""" + + output: dict[str, torch.Tensor] = {} + for name in PIECE_TENSOR_NAMES: + value = closure_embeddings if name == "closure_embeddings" else getattr(pieces, name) + value = value.detach().cpu().contiguous() + if value.is_floating_point(): + value = value.to(floating_dtype) + output[name] = value + validate_piece(output) + return output + + +def validate_piece(piece: dict[str, torch.Tensor]) -> None: + missing = [name for name in PIECE_TENSOR_NAMES if name not in piece] + if missing: + raise ValueError(f"Cached AV piece is missing {missing}") + prefix = piece["prefix_embeddings"] + payload = piece["payload_embeddings"] + closure = piece["closure_embeddings"] + if prefix.ndim != 3 or prefix.shape[0] != 1: + raise ValueError(f"Invalid prefix embeddings: {tuple(prefix.shape)}") + if payload.ndim != 3 or payload.shape[0] != 1: + raise ValueError(f"Invalid payload embeddings: {tuple(payload.shape)}") + if closure.ndim != 3 or closure.shape[0] != 1: + raise ValueError(f"Invalid closure embeddings: {tuple(closure.shape)}") + hidden_sizes = {int(prefix.shape[-1]), int(payload.shape[-1]), int(closure.shape[-1])} + if len(hidden_sizes) != 1: + raise ValueError(f"Cached piece hidden sizes disagree: {sorted(hidden_sizes)}") + expected_positions = ( + ("prefix_position_ids", prefix.shape[1]), + ("payload_position_ids", payload.shape[1]), + ("closure_position_ids", closure.shape[1]), + ) + for name, length in expected_positions: + value = piece[name] + if value.shape != (3, 1, length): + raise ValueError( + f"{name} must have shape {(3, 1, length)}, got {tuple(value.shape)}" + ) + if piece["closure_ids"].shape != (1, closure.shape[1]): + raise ValueError( + "closure_ids/closure_embeddings mismatch: " + f"{tuple(piece['closure_ids'].shape)} vs {tuple(closure.shape)}" + ) + + +def make_anchor( + *, + step_index: int, + steps_per_block: int, + memory_blocks: int, + step_seconds: float, + num_target_steps: int, + num_horizons: int, +) -> dict[str, Any]: + """Describe the exact completed-history/current-tail selection at time t.""" + + if step_index < 0: + raise ValueError("step_index must be non-negative") + block_index = step_index // steps_per_block + tail_steps = step_index % steps_per_block + 1 + first_history = max(0, block_index - (memory_blocks - 1)) + history_blocks = list(range(first_history, block_index)) + return { + "step_index": int(step_index), + "time_seconds": float((step_index + 1) * step_seconds), + "block_index": int(block_index), + "tail_steps": int(tail_steps), + "history_blocks": history_blocks, + "target_index": int(step_index), + "has_future_target": bool(step_index + num_horizons < num_target_steps), + } + + +class CachedBlockStore: + """Small CPU LRU for block shard files shared by adjacent samples.""" + + def __init__(self, cache_dir: str | Path, max_blocks: int = 16) -> None: + if max_blocks <= 0: + raise ValueError("max_blocks must be positive") + self.cache_dir = Path(cache_dir) + self.max_blocks = int(max_blocks) + self._blocks: OrderedDict[int, dict[str, Any]] = OrderedDict() + + def get(self, block_index: int) -> dict[str, Any]: + cached = self._blocks.get(block_index) + if cached is not None: + self._blocks.move_to_end(block_index) + return cached + path = self.cache_dir / cache_block_filename(block_index) + if not path.is_file(): + raise FileNotFoundError(path) + payload = torch_load(path) + if int(payload.get("block_index", -1)) != block_index: + raise ValueError(f"Block index mismatch in {path}") + variants = payload.get("variants") + if not isinstance(variants, dict) or not variants: + raise ValueError(f"No variants in {path}") + self._blocks[block_index] = payload + self._blocks.move_to_end(block_index) + while len(self._blocks) > self.max_blocks: + self._blocks.popitem(last=False) + return payload + + def variant(self, block_index: int, tail_steps: int) -> dict[str, torch.Tensor]: + block = self.get(block_index) + variants = block["variants"] + piece = variants.get(str(tail_steps), variants.get(tail_steps)) + if piece is None: + raise KeyError( + f"Block {block_index} has no tail_steps={tail_steps}; " + f"available={sorted(str(key) for key in variants)}" + ) + validate_piece(piece) + return piece + + +def _strip_embedding_batch(value: torch.Tensor) -> torch.Tensor: + if value.ndim != 3 or value.shape[0] != 1: + raise ValueError(f"Expected [1,L,D], got {tuple(value.shape)}") + return value[0] + + +def _strip_position_batch(value: torch.Tensor) -> torch.LongTensor: + if value.ndim != 3 or value.shape[:2] != (3, 1): + raise ValueError(f"Expected [3,1,L], got {tuple(value.shape)}") + return value[:, 0, :].long() + + +def reconstruct_open_stream_sequence( + *, + store: CachedBlockStore, + static: dict[str, Any], + anchor: dict[str, Any], +) -> tuple[torch.Tensor, torch.LongTensor]: + """Materialize one differentiable Thinker input from frozen AV embeddings. + + Complete history blocks use their 2.0-second variant. Only the current + block uses its 0.2--2.0-second phase. The first retained block keeps the + chat/media prefix; later blocks append payload only. The final block's + media closure and the artificial assistant readout query are appended once. + """ + + steps_per_block = int(static["steps_per_block"]) + positions_per_chunk = int(static["positions_per_chunk"]) + history_blocks = [int(value) for value in anchor["history_blocks"]] + current_block = int(anchor["block_index"]) + selected: list[tuple[int, int]] = [ + (block_index, steps_per_block) for block_index in history_blocks + ] + selected.append((current_block, int(anchor["tail_steps"]))) + if not selected: + raise ValueError("An anchor must select at least its current block") + + embedding_parts: list[torch.Tensor] = [] + position_parts: list[torch.LongTensor] = [] + first_payload_base: int | None = None + last_piece: dict[str, torch.Tensor] | None = None + last_shift = 0 + + for sequence_index, (block_index, tail_steps) in enumerate(selected): + piece = store.variant(block_index, tail_steps) + local_payload_positions = _strip_position_batch(piece["payload_position_ids"]) + local_payload_base = int(local_payload_positions.min().item()) + if first_payload_base is None: + first_payload_base = local_payload_base + desired_payload_base = first_payload_base + sequence_index * positions_per_chunk + position_shift = desired_payload_base - local_payload_base + + if sequence_index == 0: + embedding_parts.append(_strip_embedding_batch(piece["prefix_embeddings"])) + position_parts.append(_strip_position_batch(piece["prefix_position_ids"])) + embedding_parts.append(_strip_embedding_batch(piece["payload_embeddings"])) + position_parts.append(local_payload_positions + position_shift) + last_piece = piece + last_shift = position_shift + + if last_piece is None: + raise RuntimeError("No cached AV piece was selected") + closure_embeddings = _strip_embedding_batch(last_piece["closure_embeddings"]) + closure_positions = ( + _strip_position_batch(last_piece["closure_position_ids"]) + last_shift + ) + embedding_parts.append(closure_embeddings) + position_parts.append(closure_positions) + + query_embeddings = static["query_embeddings"] + if query_embeddings.ndim != 3 or query_embeddings.shape[0] != 1: + raise ValueError( + f"static query_embeddings must be [1,Q,D], got {tuple(query_embeddings.shape)}" + ) + query_embeddings = query_embeddings[0] + query_start = int(closure_positions.max().item()) + 1 + query_scalar = torch.arange( + query_start, + query_start + query_embeddings.shape[0], + dtype=torch.long, + ) + query_positions = query_scalar.view(1, -1).expand(3, -1) + embedding_parts.append(query_embeddings) + position_parts.append(query_positions) + + embeddings = torch.cat(embedding_parts, dim=0).contiguous() + position_ids = torch.cat(position_parts, dim=1).contiguous() + if position_ids.shape != (3, embeddings.shape[0]): + raise RuntimeError( + "Reconstructed position/embedding mismatch: " + f"{tuple(position_ids.shape)} vs {tuple(embeddings.shape)}" + ) + return embeddings, position_ids + + +def collate_cached_sequences( + batch: Sequence[dict[str, Any]], +) -> dict[str, Any]: + """Right-pad variable token sequences while preserving each readout index.""" + + if not batch: + raise ValueError("Cannot collate an empty batch") + hidden_size = int(batch[0]["inputs_embeds"].shape[-1]) + max_length = max(int(item["inputs_embeds"].shape[0]) for item in batch) + dtype = batch[0]["inputs_embeds"].dtype + inputs_embeds = torch.zeros( + len(batch), max_length, hidden_size, dtype=dtype + ) + position_ids = torch.zeros(3, len(batch), max_length, dtype=torch.long) + attention_mask = torch.zeros(len(batch), max_length, dtype=torch.long) + readout_indices = torch.empty(len(batch), dtype=torch.long) + targets = torch.stack([item["targets"] for item in batch]) + + for batch_index, item in enumerate(batch): + embeddings = item["inputs_embeds"] + positions = item["position_ids"] + length = int(embeddings.shape[0]) + if embeddings.shape[-1] != hidden_size: + raise ValueError("Mixed hidden sizes in one batch") + if positions.shape != (3, length): + raise ValueError( + f"Invalid position_ids {tuple(positions.shape)} for length {length}" + ) + inputs_embeds[batch_index, :length] = embeddings + position_ids[:, batch_index, :length] = positions + attention_mask[batch_index, :length] = 1 + readout_indices[batch_index] = length - 1 + + return { + "inputs_embeds": inputs_embeds, + "position_ids": position_ids, + "attention_mask": attention_mask, + "readout_indices": readout_indices, + "targets": targets, + "step_indices": torch.tensor( + [int(item["step_index"]) for item in batch], dtype=torch.long + ), + } + + +def anchor_intervals( + anchor: dict[str, Any], + *, + block_seconds: float, + step_seconds: float, +) -> list[tuple[float, float]]: + """Return source intervals used by the reference KV replay for an anchor.""" + + intervals = [ + (index * block_seconds, (index + 1) * block_seconds) + for index in anchor["history_blocks"] + ] + block_index = int(anchor["block_index"]) + current_start = block_index * block_seconds + current_end = current_start + int(anchor["tail_steps"]) * step_seconds + intervals.append((current_start, current_end)) + return intervals + + +def forward_last_state( + language_model: torch.nn.Module, + *, + inputs_embeds: torch.Tensor, + position_ids: torch.LongTensor, + attention_mask: torch.LongTensor, + readout_indices: torch.LongTensor, +) -> torch.Tensor: + """Run full-sequence causal Thinker attention and gather each query state.""" + + outputs = language_model( + inputs_embeds=inputs_embeds, + position_ids=position_ids, + attention_mask=attention_mask, + use_cache=False, + output_hidden_states=False, + return_dict=True, + cache_position=torch.arange( + inputs_embeds.shape[1], device=inputs_embeds.device + ), + ) + batch_indices = torch.arange( + inputs_embeds.shape[0], device=inputs_embeds.device + ) + return outputs.last_hidden_state[batch_indices, readout_indices] + + +def cache_size_bytes(paths: Iterable[Path]) -> int: + return sum(path.stat().st_size for path in paths if path.is_file()) diff --git a/requirements.txt b/requirements.txt index 2419e09..b6ea753 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,5 +8,8 @@ Pillow tqdm scikit-learn +# Used by train_omni_turn_lora.py; already included in the ms-swift image. +peft + # Optional on compatible NVIDIA GPUs: flash-attn # Optional NVDEC backend: PyNvVideoCodec diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py new file mode 100644 index 0000000..8581231 --- /dev/null +++ b/train_omni_turn_lora.py @@ -0,0 +1,634 @@ +#!/usr/bin/env python3 +"""Train Qwen2.5-Omni Thinker LoRA + EgoSpeak 10x3 head from AV caches. + +This is a custom task runner intended to execute inside the ms-swift image. +It uses the PEFT library shipped with ms-swift because Swift's normal SFT CLI +optimizes generated text tokens, whereas this experiment reads one continuous +Thinker hidden state and applies a multi-horizon classification loss. + +The cached audio and vision towers remain frozen. The reconstructed sequence +is forwarded with ``use_cache=False`` so all retained 2-second blocks +participate in autograd through the Thinker LoRA adapters. +""" + +from __future__ import annotations + +import argparse +import gc +import importlib.metadata +import json +import math +import random +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +import torch +from torch import nn +from torch.utils.data import DataLoader, Dataset +from tqdm.auto import tqdm + +import streaming_turn_model as core +from easycom_head import ( + NUM_CLASSES, + NUM_HORIZONS, + TurnPredictionHead, + checkpoint_head_config, + egospeak_multihorizon_loss, +) +from omni_cached_av import ( + SCHEMA_VERSION, + CachedBlockStore, + collate_cached_sequences, + forward_last_state, + load_json, + reconstruct_open_stream_sequence, + save_json, + torch_load, +) + + +DEFAULT_LORA_TARGETS = ( + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", +) + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-id", required=True) + parser.add_argument( + "--cache-dir", + type=Path, + action="append", + required=True, + help="Cache directory from cache_omni_av_pieces.py; repeat for sessions", + ) + parser.add_argument( + "--target", + type=Path, + help="Override target path for a single --cache-dir", + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--head-checkpoint", type=Path) + parser.add_argument("--max-examples", type=int, default=0) + parser.add_argument("--max-steps", type=int, default=200) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--gradient-accumulation-steps", type=int, default=1) + parser.add_argument("--num-workers", type=int, default=0) + parser.add_argument("--block-cache-size", type=int, default=16) + parser.add_argument("--learning-rate", type=float, default=2e-5) + parser.add_argument("--head-learning-rate", type=float, default=1e-4) + parser.add_argument("--weight-decay", type=float, default=0.01) + parser.add_argument("--warmup-steps", type=int, default=10) + parser.add_argument("--max-grad-norm", type=float, default=1.0) + parser.add_argument("--lora-rank", type=int, default=8) + parser.add_argument("--lora-alpha", type=int, default=16) + parser.add_argument("--lora-dropout", type=float, default=0.05) + parser.add_argument( + "--lora-target-modules", + default=",".join(DEFAULT_LORA_TARGETS), + ) + parser.add_argument("--head-dropout", type=float, default=0.1) + parser.add_argument("--gradient-checkpointing", action="store_true") + parser.add_argument("--seed", type=int, default=1337) + parser.add_argument("--log-every", type=int, default=5) + parser.add_argument("--eval-max-batches", type=int, default=4) + parser.add_argument( + "--dtype", + choices=("bfloat16", "float16"), + default="bfloat16", + ) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--flash-attention-2", action="store_true") + parser.add_argument("--overwrite-output", action="store_true") + parser.add_argument("--no-progress", action="store_true") + return parser.parse_args(argv) + + +def _seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def _package_version(name: str) -> str | None: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return None + + +def _load_targets(path: Path) -> np.ndarray: + targets = np.load(path, mmap_mode="r") + if targets.ndim != 2 or targets.shape[1] != NUM_CLASSES: + raise ValueError(f"Expected [T,{NUM_CLASSES}] target, got {targets.shape}: {path}") + return targets + + +class OmniTurnCacheDataset(Dataset[dict[str, Any]]): + """One sample per 200-ms anchor, reconstructed lazily on CPU.""" + + def __init__( + self, + cache_dirs: Sequence[Path], + *, + target_override: Path | None, + max_examples: int, + seed: int, + block_cache_size: int, + ) -> None: + if target_override is not None and len(cache_dirs) != 1: + raise ValueError("--target can override only one --cache-dir") + self.sources: list[dict[str, Any]] = [] + self.examples: list[tuple[int, dict[str, Any]]] = [] + hidden_size: int | None = None + reference_contract: tuple[Any, ...] | None = None + + for source_index, raw_cache_dir in enumerate(cache_dirs): + cache_dir = raw_cache_dir.expanduser().resolve() + manifest_path = cache_dir / "manifest.json" + manifest = load_json(manifest_path) + if int(manifest.get("schema_version", -1)) != SCHEMA_VERSION: + raise ValueError( + f"Unsupported cache schema in {manifest_path}: " + f"{manifest.get('schema_version')}" + ) + static = torch_load(cache_dir / manifest["static_path"]) + target_path = ( + target_override.expanduser().resolve() + if target_override is not None + else Path(manifest["target_path"]) + ) + targets = _load_targets(target_path) + if len(targets) != int(manifest["target_steps"]): + raise ValueError( + f"Target length changed for {manifest_path}: " + f"{len(targets)} vs {manifest['target_steps']}" + ) + source_hidden = int(manifest["hidden_size"]) + hidden_size = source_hidden if hidden_size is None else hidden_size + if source_hidden != hidden_size: + raise ValueError("Cache directories have different hidden sizes") + contract = ( + manifest["model_id"], + manifest["step_seconds"], + manifest["native_block_seconds"], + manifest["memory_seconds"], + manifest["video_fps"], + manifest["positions_per_chunk"], + ) + reference_contract = contract if reference_contract is None else reference_contract + if contract != reference_contract: + raise ValueError( + "Cache directories do not share the same model/timing contract: " + f"{manifest_path}" + ) + self.sources.append( + { + "cache_dir": cache_dir, + "manifest": manifest, + "static": static, + "targets": targets, + "store": CachedBlockStore(cache_dir, max_blocks=block_cache_size), + } + ) + for anchor in manifest["anchors"]: + target_index = int(anchor["target_index"]) + if not bool(anchor.get("has_future_target", True)): + continue + if target_index + NUM_HORIZONS >= len(targets): + continue + self.examples.append((source_index, anchor)) + + if not self.examples: + raise ValueError("No cache anchors have ten future target steps") + if max_examples < 0: + raise ValueError("--max-examples must be non-negative") + if max_examples and max_examples < len(self.examples): + generator = random.Random(seed) + selected = generator.sample(range(len(self.examples)), max_examples) + self.examples = [self.examples[index] for index in sorted(selected)] + self.hidden_size = int(hidden_size or 0) + + def __len__(self) -> int: + return len(self.examples) + + def __getitem__(self, index: int) -> dict[str, Any]: + source_index, anchor = self.examples[index] + source = self.sources[source_index] + embeddings, positions = reconstruct_open_stream_sequence( + store=source["store"], + static=source["static"], + anchor=anchor, + ) + target_index = int(anchor["target_index"]) + targets = torch.from_numpy( + np.array( + source["targets"][ + target_index + 1 : target_index + NUM_HORIZONS + 1 + ], + dtype=np.float32, + copy=True, + ) + ) + if targets.shape != (NUM_HORIZONS, NUM_CLASSES): + raise RuntimeError( + f"Unexpected future target shape {tuple(targets.shape)} at {target_index}" + ) + return { + "inputs_embeds": embeddings, + "position_ids": positions, + "targets": targets, + "step_index": target_index, + } + + +def _load_initial_head( + checkpoint: Path | None, + *, + hidden_size: int, + dropout: float, +) -> TurnPredictionHead: + if checkpoint is None: + return TurnPredictionHead(hidden_size=hidden_size, dropout=dropout) + try: + payload = torch.load(checkpoint, map_location="cpu", weights_only=True) + except (TypeError, RuntimeError): + payload = torch.load(checkpoint, map_location="cpu") + if not isinstance(payload, dict): + raise ValueError(f"Expected dict head checkpoint: {checkpoint}") + config = checkpoint_head_config(payload) + if int(config["hidden_size"]) != hidden_size: + raise ValueError( + f"Head hidden size {config['hidden_size']} != cache hidden size {hidden_size}" + ) + head = TurnPredictionHead( + hidden_size=hidden_size, + num_horizons=int(config["num_horizons"]), + num_classes=int(config["num_classes"]), + dropout=float(config["dropout"]), + ) + head.load_state_dict(payload.get("head", payload)) + print(f"Initialized classification head from {checkpoint}") + return head + + +def _prepare_language_lora( + model: torch.nn.Module, + *, + rank: int, + alpha: int, + dropout: float, + target_modules: Sequence[str], + gradient_checkpointing: bool, +) -> torch.nn.Module: + try: + from peft import LoraConfig, get_peft_model + except ImportError as exc: + raise RuntimeError( + "PEFT is required. The ms-swift 3.12.5 image should include it." + ) from exc + + model.requires_grad_(False) + language_model = model.model + matched = [ + name + for name, module in language_model.named_modules() + if isinstance(module, nn.Linear) and name.rsplit(".", 1)[-1] in target_modules + ] + if not matched: + raise RuntimeError( + "No language-model linear layers matched LoRA targets " + f"{list(target_modules)}" + ) + config = LoraConfig( + r=rank, + lora_alpha=alpha, + lora_dropout=dropout, + target_modules=list(target_modules), + bias="none", + ) + language_model = get_peft_model(language_model, config) + language_model.config.use_cache = False + if gradient_checkpointing: + try: + language_model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + except TypeError: + language_model.gradient_checkpointing_enable() + language_model.print_trainable_parameters() + print( + f"LoRA matched {len(matched)} language layers; first={matched[0]}, " + f"last={matched[-1]}" + ) + return language_model + + +def _lr_lambda(step: int, *, total_steps: int, warmup_steps: int) -> float: + if warmup_steps > 0 and step < warmup_steps: + return float(step + 1) / float(warmup_steps) + remaining = max(1, total_steps - warmup_steps) + progress = min(1.0, max(0.0, (step - warmup_steps) / remaining)) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + +@torch.inference_mode() +def _evaluate_loss( + language_model: torch.nn.Module, + head: TurnPredictionHead, + loader: DataLoader, + *, + device: torch.device, + model_dtype: torch.dtype, + max_batches: int, +) -> float: + language_was_training = language_model.training + head_was_training = head.training + language_model.eval() + head.eval() + losses: list[float] = [] + for batch_index, batch in enumerate(loader): + if max_batches > 0 and batch_index >= max_batches: + break + inputs_embeds = batch["inputs_embeds"].to( + device=device, dtype=model_dtype, non_blocking=True + ) + position_ids = batch["position_ids"].to(device=device, non_blocking=True) + attention_mask = batch["attention_mask"].to(device=device, non_blocking=True) + readout_indices = batch["readout_indices"].to(device=device, non_blocking=True) + targets = batch["targets"].to(device=device, dtype=torch.float32) + states = forward_last_state( + language_model, + inputs_embeds=inputs_embeds, + position_ids=position_ids, + attention_mask=attention_mask, + readout_indices=readout_indices, + ) + logits = head(states.float()) + losses.append(float(egospeak_multihorizon_loss(logits, targets).item())) + language_model.train(language_was_training) + head.train(head_was_training) + return float(np.mean(losses)) + + +def _jsonable_args(args: argparse.Namespace) -> dict[str, Any]: + output: dict[str, Any] = {} + for key, value in vars(args).items(): + if isinstance(value, Path): + output[key] = str(value) + elif isinstance(value, list) and value and isinstance(value[0], Path): + output[key] = [str(item) for item in value] + else: + output[key] = value + return output + + +def main(argv: Sequence[str] | None = None) -> None: + args = parse_args(argv) + if args.max_steps <= 0: + raise ValueError("--max-steps must be positive") + if args.batch_size <= 0 or args.gradient_accumulation_steps <= 0: + raise ValueError("Batch size and gradient accumulation must be positive") + if args.num_workers != 0: + print( + "Warning: each worker owns an independent block-file LRU; use 0 for " + "the 16-example pilot and increase only after profiling." + ) + output_dir = args.output_dir.expanduser().resolve() + known_outputs = ( + output_dir / "adapter_config.json", + output_dir / "adapter", + output_dir / "turn_head.pt", + output_dir / "train_summary.json", + ) + if not args.overwrite_output and any(path.exists() for path in known_outputs): + raise FileExistsError( + f"Training output already exists in {output_dir}; pass --overwrite-output" + ) + output_dir.mkdir(parents=True, exist_ok=True) + _seed_everything(args.seed) + core._configure_qwen_logging() + device = torch.device(args.device) + if device.type != "cuda": + raise ValueError("This LoRA pilot is intended for one CUDA GPU") + + dataset = OmniTurnCacheDataset( + args.cache_dir, + target_override=args.target, + max_examples=args.max_examples, + seed=args.seed, + block_cache_size=args.block_cache_size, + ) + generator = torch.Generator().manual_seed(args.seed) + train_loader = DataLoader( + dataset, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.num_workers, + pin_memory=True, + persistent_workers=args.num_workers > 0, + collate_fn=collate_cached_sequences, + generator=generator, + ) + eval_loader = DataLoader( + dataset, + batch_size=args.batch_size, + shuffle=False, + num_workers=0, + pin_memory=True, + collate_fn=collate_cached_sequences, + ) + print( + f"Cached training examples: {len(dataset)}, hidden_size={dataset.hidden_size}, " + f"batch={args.batch_size} x accum={args.gradient_accumulation_steps}" + ) + + base_model, _ = core.load_qwen_omni( + args.model_id, + dtype=args.dtype, + device_map=args.device, + use_flash_attention_2=args.flash_attention_2, + ) + target_modules = tuple( + value.strip() for value in args.lora_target_modules.split(",") if value.strip() + ) + if not target_modules: + raise ValueError("--lora-target-modules is empty") + language_model = _prepare_language_lora( + base_model, + rank=args.lora_rank, + alpha=args.lora_alpha, + dropout=args.lora_dropout, + target_modules=target_modules, + gradient_checkpointing=args.gradient_checkpointing, + ) + # Cached embeddings make both modality towers unnecessary during training. + # Keeping only the language model materially reduces L20 memory use. + del base_model + gc.collect() + torch.cuda.empty_cache() + model_dtype = next(language_model.parameters()).dtype + head = _load_initial_head( + args.head_checkpoint, + hidden_size=dataset.hidden_size, + dropout=args.head_dropout, + ).to(device=device, dtype=torch.float32) + + lora_parameters = [parameter for parameter in language_model.parameters() if parameter.requires_grad] + if not lora_parameters: + raise RuntimeError("LoRA preparation produced no trainable parameters") + optimizer = torch.optim.AdamW( + [ + {"params": lora_parameters, "lr": args.learning_rate}, + {"params": head.parameters(), "lr": args.head_learning_rate}, + ], + weight_decay=args.weight_decay, + ) + scheduler = torch.optim.lr_scheduler.LambdaLR( + optimizer, + lr_lambda=lambda step: _lr_lambda( + step, total_steps=args.max_steps, warmup_steps=args.warmup_steps + ), + ) + + initial_loss = _evaluate_loss( + language_model, + head, + eval_loader, + device=device, + model_dtype=model_dtype, + max_batches=args.eval_max_batches, + ) + print(f"Initial cached loss: {initial_loss:.6f}") + language_model.train() + head.train() + optimizer.zero_grad(set_to_none=True) + iterator = iter(train_loader) + history: list[dict[str, float | int]] = [] + progress = tqdm( + total=args.max_steps, + desc="Thinker LoRA", + unit="update", + disable=args.no_progress, + ) + for update_step in range(1, args.max_steps + 1): + accumulated_loss = 0.0 + for _ in range(args.gradient_accumulation_steps): + try: + batch = next(iterator) + except StopIteration: + iterator = iter(train_loader) + batch = next(iterator) + inputs_embeds = batch["inputs_embeds"].to( + device=device, dtype=model_dtype, non_blocking=True + ) + position_ids = batch["position_ids"].to(device=device, non_blocking=True) + attention_mask = batch["attention_mask"].to(device=device, non_blocking=True) + readout_indices = batch["readout_indices"].to(device=device, non_blocking=True) + targets = batch["targets"].to(device=device, dtype=torch.float32) + states = forward_last_state( + language_model, + inputs_embeds=inputs_embeds, + position_ids=position_ids, + attention_mask=attention_mask, + readout_indices=readout_indices, + ) + logits = head(states.float()) + loss = egospeak_multihorizon_loss(logits, targets) + (loss / args.gradient_accumulation_steps).backward() + accumulated_loss += float(loss.detach()) / args.gradient_accumulation_steps + + trainable_parameters = [*lora_parameters, *head.parameters()] + grad_norm = torch.nn.utils.clip_grad_norm_( + trainable_parameters, args.max_grad_norm + ) + optimizer.step() + scheduler.step() + optimizer.zero_grad(set_to_none=True) + record = { + "step": update_step, + "loss": accumulated_loss, + "grad_norm": float(grad_norm), + "lora_lr": float(optimizer.param_groups[0]["lr"]), + "head_lr": float(optimizer.param_groups[1]["lr"]), + } + history.append(record) + progress.update(1) + progress.set_postfix( + loss=f"{accumulated_loss:.4f}", + lr=f"{optimizer.param_groups[0]['lr']:.2g}", + ) + if update_step == 1 or update_step % args.log_every == 0: + print( + f"step={update_step} loss={accumulated_loss:.6f} " + f"grad_norm={float(grad_norm):.4f} " + f"lora_lr={optimizer.param_groups[0]['lr']:.3g}", + flush=True, + ) + progress.close() + + final_loss = _evaluate_loss( + language_model, + head, + eval_loader, + device=device, + model_dtype=model_dtype, + max_batches=args.eval_max_batches, + ) + adapter_dir = output_dir / "adapter" + language_model.save_pretrained(adapter_dir) + torch.save( + { + "head": head.state_dict(), + "head_config": { + "hidden_size": dataset.hidden_size, + "num_horizons": NUM_HORIZONS, + "num_classes": NUM_CLASSES, + "dropout": head.dropout, + }, + "initial_head_checkpoint": ( + str(args.head_checkpoint.resolve()) if args.head_checkpoint else None + ), + "training_step": args.max_steps, + }, + output_dir / "turn_head.pt", + ) + summary = { + "schema_version": 1, + "status": "pilot" if args.max_examples else "training", + "model_id": args.model_id, + "cache_dirs": [str(path.expanduser().resolve()) for path in args.cache_dir], + "num_examples": len(dataset), + "initial_loss": initial_loss, + "final_loss": final_loss, + "loss_reduction_ratio": ( + (initial_loss - final_loss) / initial_loss if initial_loss else None + ), + "peak_cuda_allocated_gib": torch.cuda.max_memory_allocated(device) / (1024 ** 3), + "packages": { + "torch": torch.__version__, + "transformers": _package_version("transformers"), + "peft": _package_version("peft"), + "ms-swift": _package_version("ms-swift"), + }, + "args": _jsonable_args(args), + "history": history, + } + save_json(output_dir / "train_summary.json", summary) + print(f"Final cached loss: {final_loss:.6f} (initial {initial_loss:.6f})") + print(f"Saved LoRA adapter to {adapter_dir}") + print(f"Saved turn head to {output_dir / 'turn_head.pt'}") + print(f"Peak CUDA allocated: {summary['peak_cuda_allocated_gib']:.2f} GiB") + + +if __name__ == "__main__": + main()