wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add two-GPU AV producer for Omni LoRA

Browse Source
This commit is contained in:
Xin Wang
2026-08-23 15:29:26 +08:00
parent dc932cc165
commit c2162e1170
2 changed files with 273 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+103 -2
View File
@@ -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
progress.
EasyCom. The first Ego4D 8-second versus 20-second context ablation has also
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
+170 -10
View File
@@ -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__":