Add online current-tail LoRA benchmark

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")