diff --git a/omni_cached_av.py b/omni_cached_av.py index 5cb79c7..a7df11d 100644 --- a/omni_cached_av.py +++ b/omni_cached_av.py @@ -192,6 +192,7 @@ def reconstruct_open_stream_sequence( store: CachedBlockStore, static: dict[str, Any], anchor: dict[str, Any], + current_piece: dict[str, torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.LongTensor]: """Materialize one differentiable Thinker input from frozen AV embeddings. @@ -219,7 +220,13 @@ def reconstruct_open_stream_sequence( last_shift = 0 for sequence_index, (block_index, tail_steps) in enumerate(selected): - piece = store.variant(block_index, tail_steps) + is_current = sequence_index == len(selected) - 1 + piece = ( + current_piece + if is_current and current_piece is not None + else store.variant(block_index, tail_steps) + ) + validate_piece(piece) 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: diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py index 8581231..1dd950f 100644 --- a/train_omni_turn_lora.py +++ b/train_omni_turn_lora.py @@ -19,6 +19,7 @@ import importlib.metadata import json import math import random +import time from pathlib import Path from typing import Any, Sequence @@ -42,6 +43,7 @@ from omni_cached_av import ( collate_cached_sequences, forward_last_state, load_json, + pieces_to_cache, reconstruct_open_stream_sequence, save_json, torch_load, @@ -82,6 +84,24 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: 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("--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("--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) @@ -144,11 +164,13 @@ class OmniTurnCacheDataset(Dataset[dict[str, Any]]): 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 @@ -224,11 +246,6 @@ class OmniTurnCacheDataset(Dataset[dict[str, Any]]): 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( @@ -243,12 +260,139 @@ class OmniTurnCacheDataset(Dataset[dict[str, Any]]): raise RuntimeError( f"Unexpected future target shape {tuple(targets.shape)} at {target_index}" ) - return { - "inputs_embeds": embeddings, - "position_ids": positions, + 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, + ) -> 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"]) + 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, + ) + 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 def _load_initial_head( @@ -350,6 +494,7 @@ def _evaluate_loss( device: torch.device, model_dtype: torch.dtype, max_batches: int, + materialize_batch: Any | None = None, ) -> float: language_was_training = language_model.training head_was_training = head.training @@ -359,6 +504,8 @@ def _evaluate_loss( 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 ) @@ -398,6 +545,17 @@ def main(argv: Sequence[str] | None = None) -> None: 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 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.num_workers != 0: print( "Warning: each worker owns an independent block-file LRU; use 0 for " @@ -427,6 +585,12 @@ def main(argv: Sequence[str] | None = None) -> None: 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) train_loader = DataLoader( @@ -436,7 +600,7 @@ def main(argv: Sequence[str] | None = None) -> None: num_workers=args.num_workers, pin_memory=True, persistent_workers=args.num_workers > 0, - collate_fn=collate_cached_sequences, + collate_fn=collate_fn, generator=generator, ) eval_loader = DataLoader( @@ -445,14 +609,14 @@ def main(argv: Sequence[str] | None = None) -> None: shuffle=False, num_workers=0, pin_memory=True, - collate_fn=collate_cached_sequences, + collate_fn=collate_fn, ) 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( + base_model, processor = core.load_qwen_omni( args.model_id, dtype=args.dtype, device_map=args.device, @@ -471,11 +635,30 @@ def main(argv: Sequence[str] | None = None) -> None: 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() + materialize_batch: OnlineCurrentTailMaterializer | None = None + if args.online_current_tail: + assert args.easycom_session_dir is not None + materialize_batch = OnlineCurrentTailMaterializer( + dataset=dataset, + model=base_model, + processor=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, + ) + print( + "Cache policy benchmark: complete 2s history cached; " + "current 0.2--1.8s tail computed online" + ) + 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, @@ -507,8 +690,11 @@ def main(argv: Sequence[str] | None = None) -> None: device=device, model_dtype=model_dtype, max_batches=args.eval_max_batches, + materialize_batch=materialize_batch, ) 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) @@ -520,6 +706,7 @@ def main(argv: Sequence[str] | None = None) -> None: unit="update", disable=args.no_progress, ) + training_started = time.perf_counter() for update_step in range(1, args.max_steps + 1): accumulated_loss = 0.0 for _ in range(args.gradient_accumulation_steps): @@ -528,6 +715,8 @@ def main(argv: Sequence[str] | None = None) -> None: except StopIteration: iterator = iter(train_loader) batch = next(iterator) + if materialize_batch is not None: + batch = materialize_batch(batch) inputs_embeds = batch["inputs_embeds"].to( device=device, dtype=model_dtype, non_blocking=True ) @@ -575,6 +764,15 @@ def main(argv: Sequence[str] | None = None) -> None: flush=True, ) progress.close() + training_seconds = time.perf_counter() - training_started + 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, + } final_loss = _evaluate_loss( language_model, @@ -583,6 +781,7 @@ def main(argv: Sequence[str] | None = None) -> None: device=device, model_dtype=model_dtype, max_batches=args.eval_max_batches, + materialize_batch=materialize_batch, ) adapter_dir = output_dir / "adapter" language_model.save_pretrained(adapter_dir) @@ -610,6 +809,9 @@ def main(argv: Sequence[str] | None = None) -> None: "num_examples": len(dataset), "initial_loss": initial_loss, "final_loss": final_loss, + "training_seconds": training_seconds, + "updates_per_second": args.max_steps / training_seconds, + "online_tail_timing": online_timing, "loss_reduction_ratio": ( (initial_loss - final_loss) / initial_loss if initial_loss else None ), @@ -625,6 +827,16 @@ def main(argv: Sequence[str] | None = None) -> None: } 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" + ) + 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" + ) 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")