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

Add online current-tail LoRA benchmark

Browse Source
This commit is contained in:
Xin Wang
2026-08-23 15:10:16 +08:00
parent 1069a0beee
commit dc932cc165
2 changed files with 236 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
omni_cached_av.py
View File

@@ -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:

238
train_omni_turn_lora.py
View File

@@ -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,9 +635,28 @@ 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.
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
@@ -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")