Add two-GPU AV producer for Omni LoRA

README.md

@@ -40,8 +40,8 @@ classification head.
 The full Thinker input path, embedded audio, direct video decoding,
 FlashAttention-2, 5 Hz streaming extraction, cache replay, target alignment,
 head training, and EgoSpeak-compatible evaluation have now been exercised on
-EasyCom. Ego4D feature extraction and multi-seed EasyCom evaluation remain in
+EasyCom. The first Ego4D 8-second versus 20-second context ablation has also
-progress.
+been recorded. Multi-seed evaluation remains in progress.
 
 ## Current EasyCom results
 
@@ -64,6 +64,69 @@ The full per-horizon and per-class snapshot, machine-readable CSV, provenance,
 and follow-up checks are in
 [`results/easycom_open_stream_vs_causal_prefix_2026-08-19/`](results/easycom_open_stream_vs_causal_prefix_2026-08-19/README.md).
 
+### EasyCom 8-second versus 20-second context
+
+An additional single-run open-stream comparison recorded on 2026-08-23 gives
+**58.51937** full-timeline mean mAP with an 8-second context budget versus
+**58.316** with 20 seconds, a difference of **+0.20337** points. The aggregate
+difference is small and hides a clear horizon tradeoff: 8 seconds is stronger
+at 0.2--0.8 seconds, whereas 20 seconds is stronger from 1.0 through 2.0
+seconds. These runs must be repeated with identical additional seeds before
+choosing a context budget for the paper.
+
+The raw 8-second console output, machine-readable comparison, and detailed
+interpretation are archived in
+[`results/easycom_open_stream_8s_vs_20s_2026-08-23/`](results/easycom_open_stream_8s_vs_20s_2026-08-23/README.md).
+
+## Current Ego4D context-length results
+
+The following single-seed snapshot was recorded on 2026-08-22. Both variants
+use open-stream replay, 2 FPS Decord video, 5 Hz predictions, the same frozen
+Qwen2.5-Omni-3B Thinker and independently trained `10 x 3` heads. They differ
+only in the rolling causal-memory budget: 8 seconds versus 20 seconds. Both
+runs use the same EgoSpeak Ego4D train/test manifest and test timeline.
+
+| Context | EgoSpeak full-timeline mean mAP | Strict causal mean mAP |
+|---|---:|---:|
+| 8 seconds | **65.5943** | **65.6920** |
+| 20 seconds | 63.1408 | 63.2298 |
+| 20s minus 8s | -2.4535 | -2.4622 |
+
+The strict-causal per-horizon comparison is:
+
+| Horizon | 8s mAP | 20s mAP | 20s minus 8s |
+|---:|---:|---:|---:|
+| 0.2s | **74.2213** | 71.7403 | -2.4810 |
+| 0.4s | **71.0180** | 68.5945 | -2.4234 |
+| 0.6s | **68.3746** | 65.9287 | -2.4459 |
+| 0.8s | **66.3190** | 63.8606 | -2.4584 |
+| 1.0s | **64.7470** | 62.3329 | -2.4140 |
+| 1.2s | **63.6416** | 61.2290 | -2.4126 |
+| 1.4s | **62.8697** | 60.4035 | -2.4662 |
+| 1.6s | **62.3142** | 59.7870 | -2.5272 |
+| 1.8s | **61.8705** | 59.3678 | -2.5027 |
+| 2.0s | **61.5444** | 59.0541 | -2.4903 |
+| Mean | **65.6920** | 63.2298 | -2.4622 |
+
+For the exact full-timeline protocol, averaging class AP over the ten
+horizons gives:
+
+| Class | 8s mean AP | 20s mean AP | 20s minus 8s |
+|---|---:|---:|---:|
+| Background | **69.5063** | 67.2409 | -2.2653 |
+| Camera wearer speaking | **59.4005** | 56.5784 | -2.8222 |
+| Other people speaking | **67.8760** | 65.6030 | -2.2730 |
+
+The 8-second representation is better at every horizon and for every class.
+The approximately 2.4--2.5 point gap is nearly constant across anticipation
+horizons, while the largest class degradation is for camera-wearer speech.
+This pattern does not support a long-horizon benefit from the additional 12
+seconds of context. A plausible hypothesis is that excessive history dilutes
+locally predictive turn-taking cues in the final frozen Thinker state, but the
+single-seed comparison cannot establish that mechanism. Both context settings
+must be repeated with the same additional seeds before treating the difference
+as a final paper result.
+
 ## Why Omni-only instead of Omni + RNN
 
 EgoSpeak separately extracts visual and audio features and then uses GRU,
@@ -840,6 +903,44 @@ Recommended experimental procedure:
 Class-balanced/focal loss, a learned turn-query token, and LoRA are later
 ablations. They should not replace the frozen-Omni linear/MLP-head baseline.
 
+## Planned Thinker-LoRA sampling and cache policy
+
+This is a planned second-stage experiment, not a reported benchmark result.
+Dense 5 Hz partial-prefix caching is fast during training but does not fit the
+combined EasyCom and Ego4D storage budget.  The 24-second Session 7 pilot used
+0.115 GiB, which extrapolates to about 17.25 GiB per media hour.  Caching only
+complete 2-second history blocks and recomputing every current partial online
+reduced storage, but the single-L20 pilot slowed 200 LoRA updates from 83 s
+(2.39 update/s) to 254.085 s (0.787 update/s).  Online AV preparation used
+98.654 s for 163 partials and peak allocation was 9.56 GiB, so the bottleneck
+is repeated preprocessing/modality encoding rather than memory capacity.
+
+The proposed training sampler is therefore:
+
+1. retain a uniform low-rate set of ordinary anchors, initially 1 Hz;
+2. retain every 5 Hz anchor within a two-second band around annotated speaker
+   onset, offset, handover, and overlap-state transitions;
+3. optionally add high-loss or high-uncertainty anchors after a first training
+   pass, while selecting them only from the training split;
+4. cache exact partial AV embeddings only for these selected anchors, plus one
+   reusable complete embedding for every required 2-second history block;
+5. enforce a 240 GiB cache target and a hard 300 GB upper bound, leaving room
+   for manifests, checkpoints, and extracted evaluation states.
+
+Transition-centered selection is training-example sampling, not inference
+look-ahead: model inputs must still end at the current time.  Final extraction,
+head training, and Table 2 evaluation remain on the complete 5 Hz timeline.
+The paper must report the sampling rule and compare uniform-only sampling with
+transition-enhanced sampling so that any gain is not attributed solely to
+additional transition examples.
+
+For a speed-only alternative, `train_omni_turn_lora.py` supports a two-GPU
+producer-consumer benchmark.  The producer GPU runs the frozen AV towers and
+places current-tail embeddings in a bounded CPU queue while the consumer GPU
+runs Thinker LoRA.  This preserves exact online partial semantics; it changes
+only scheduling.  Compare its reported update rate and consumer queue-wait
+time with the synchronous 0.787 update/s pilot before adopting it.
+
 ## EgoSpeak-compatible evaluation
 
 Run inference on the official test sessions (`Session_3`, `Session_2`, and

train_omni_turn_lora.py

@@ -18,7 +18,9 @@ import gc
 import importlib.metadata
 import json
 import math
+import queue
 import random
+import threading
 import time
 from pathlib import Path
 from typing import Any, Sequence
@@ -97,6 +99,19 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         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("--decoder", choices=("decord", "pynvcodec"), default="decord")
     parser.add_argument("--decoder-gpu-id", type=int, default=0)
@@ -395,6 +410,81 @@ class OnlineCurrentTailMaterializer:
         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,
     *,
@@ -556,6 +646,15 @@ def main(argv: Sequence[str] | None = None) -> None:
             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.num_workers != 0:
         print(
             "Warning: each worker owns an independent block-file LRU; use 0 for "
@@ -636,12 +735,24 @@ def main(argv: Sequence[str] | None = None) -> None:
         gradient_checkpointing=args.gradient_checkpointing,
     )
     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=base_model,
+            model=materializer_model,
-            processor=processor,
+            processor=materializer_processor,
             session_dir=args.easycom_session_dir,
             frame_cache_size=args.frame_cache_size,
             decoder=args.decoder,
@@ -653,6 +764,17 @@ def main(argv: Sequence[str] | None = None) -> None:
             "Cache policy benchmark: complete 2s history cached; "
             "current 0.2--1.8s tail computed online"
         )
+        if args.producer_device is not None:
+            print(
+                f"Producer-consumer enabled: frozen AV={args.producer_device}, "
+                f"Thinker LoRA={args.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
@@ -698,7 +820,17 @@ def main(argv: Sequence[str] | None = None) -> None:
     language_model.train()
     head.train()
     optimizer.zero_grad(set_to_none=True)
-    iterator = iter(train_loader)
+    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()
+    iterator = iter(train_loader) if batch_producer is None else None
     history: list[dict[str, float | int]] = []
     progress = tqdm(
         total=args.max_steps,
@@ -706,16 +838,19 @@ 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):
-            try:
+            if batch_producer is not None:
-                batch = next(iterator)
+                batch = batch_producer.next_batch()
-            except StopIteration:
+            else:
-                iterator = iter(train_loader)
+                assert iterator is not None
-                batch = next(iterator)
+                try:
-            if materialize_batch is not None:
+                    batch = next(iterator)
+                except StopIteration:
+                    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
@@ -765,6 +900,8 @@ def main(argv: Sequence[str] | None = None) -> None:
             )
     progress.close()
     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 = {
@@ -772,6 +909,12 @@ def main(argv: Sequence[str] | None = None) -> None:
             "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(
@@ -816,6 +959,12 @@ def main(argv: Sequence[str] | None = None) -> None:
             (initial_loss - final_loss) / initial_loss if initial_loss else None
         ),
         "peak_cuda_allocated_gib": torch.cuda.max_memory_allocated(device) / (1024 ** 3),
+        "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"),
@@ -837,9 +986,20 @@ def main(argv: Sequence[str] | None = None) -> None:
             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(f"Peak CUDA allocated: {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"
+        )
 
 
 if __name__ == "__main__":