Add two-GPU AV producer for Omni LoRA

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,
-            processor=processor,
+            model=materializer_model,
+            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:
-                batch = next(iterator)
-            except StopIteration:
-                iterator = iter(train_loader)
-                batch = next(iterator)
-            if materialize_batch is not None:
+            if batch_producer is not None:
+                batch = batch_producer.next_batch()
+            else:
+                assert iterator is not None
+                try:
+                    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__":