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

Add transition-sparse Omni AV caching

Browse Source
This commit is contained in:
Xin Wang
2026-08-23 18:28:07 +08:00
parent 8b373cd19b
commit cb4abfdf82
2 changed files with 333 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+44 -8
View File
@@ -903,17 +903,20 @@ 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
## Thinker-LoRA sampling and cache policy
This is a planned second-stage experiment, not a reported benchmark result.
This is an implemented second-stage experiment, not yet 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.
0.115 GiB, which extrapolates to about 17.25 GiB per media hour. In the latest
controlled single-L20 test without gradient checkpointing, 200 updates took
57.718 s (3.465 update/s) with every partial cached, 130.151 s (1.537
update/s) with current partials recomputed using processor defaults, and
111.146 s (1.799 update/s) after CUDA log-Mel and bounded-audio-padding
optimization. The optimized online run spent 53.408 s in 163 frozen-AV calls;
the cached run therefore remains the preferred training path when storage is
available.
The proposed training sampler is therefore:
@@ -934,6 +937,39 @@ 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.
`cache_omni_av_pieces.py --anchor-policy transition-sparse` implements this
policy without changing the trainer cache schema. A transition is any change
in the three-class EasyCom target vector. The default two-second dense band is
defined explicitly as one second before through one second after that target
row. Every complete two-second source block is still cached once as reusable
history. Ordinary selected anchors and transition-band anchors cache only
their exact current causal phase. Label information selects training examples
offline but never changes the media endpoint supplied to Qwen.
Inspect the selection counts for the 24-second pilot without loading media or
the model:
```bash
python cache_omni_av_pieces.py \
--model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
--easycom-session-dir /data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed/Session_7 \
--target /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame/Session_7.npy \
--output-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_7_24s \
--end 24 \
--anchor-policy transition-sparse \
--ordinary-anchor-hz 1 \
--transition-before-seconds 1 \
--transition-after-seconds 1 \
--selection-only
```
After auditing the counts, remove `--selection-only` and add the normal AV
extraction options. Use `--end 0` only after the 24-second cache and replay
verification pass; zero selects the complete target timeline. The resulting
manifest contains only selected training anchors, so `train_omni_turn_lora.py`
uses the sparse dataset automatically and must be run without
`--online-current-tail`.
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
cache_omni_av_pieces.py
+289 -35
View File
@@ -1,10 +1,12 @@
#!/usr/bin/env python3
"""Cache exact Qwen2.5-Omni AV pieces for Thinker LoRA training.
For every source-aligned 2-second block this script caches ten causal phases:
0.2, 0.4, ..., 2.0 seconds. The 2.0-second phase is reused as completed
history; phases 0.2--1.8 are used only for the current partial block. No
future samples enter a phase.
Dense mode caches all ten causal phases (0.2, 0.4, ..., 2.0 seconds) for every
source-aligned 2-second block. Transition-sparse mode instead caches every
complete 2.0-second history block, a low-rate grid of ordinary current
anchors, and dense 5-Hz current anchors near target-state transitions. No
future media samples enter a phase; target transitions affect training-sample
selection only.
The optional verification compares a full-sequence reconstruction from the
cache with the existing open-stream KV replay. Run this 24-second pilot before
@@ -59,7 +61,10 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
"--end",
type=float,
default=24.0,
help="Dense session endpoint in seconds; use 24 for the first pilot",
help=(
"Session endpoint in seconds; use 24 for a pilot or 0 for the "
"complete target timeline"
),
)
parser.add_argument("--step-seconds", type=float, default=0.2)
parser.add_argument("--native-block-seconds", type=float, default=2.0)
@@ -68,6 +73,38 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
parser.add_argument("--frame-cache-size", type=int, default=64)
parser.add_argument(
"--anchor-policy",
choices=("dense", "transition-sparse"),
default="dense",
help=(
"dense caches every 5-Hz partial; transition-sparse caches a "
"low-rate ordinary grid plus dense anchors near label transitions"
),
)
parser.add_argument(
"--ordinary-anchor-hz",
type=float,
default=1.0,
help="Ordinary-anchor rate used by transition-sparse mode",
)
parser.add_argument(
"--transition-before-seconds",
type=float,
default=1.0,
help="Dense 5-Hz sampling duration before each target transition",
)
parser.add_argument(
"--transition-after-seconds",
type=float,
default=1.0,
help="Dense 5-Hz sampling duration after each target transition",
)
parser.add_argument(
"--selection-only",
action="store_true",
help="Print anchor/variant counts without loading media or the model",
)
parser.add_argument(
"--av-batch-size",
type=int,
@@ -145,6 +182,193 @@ def _load_targets(path: Path) -> np.ndarray:
return targets
def _transition_step_indices(targets: np.ndarray, num_steps: int) -> list[int]:
"""Return target-row indices whose three-class state changed."""
timeline_steps = min(int(num_steps), int(len(targets)))
if timeline_steps <= 1:
return []
current = np.asarray(targets[1:timeline_steps])
previous = np.asarray(targets[: timeline_steps - 1])
changed = np.any(current != previous, axis=1)
return (np.flatnonzero(changed) + 1).astype(np.int64).tolist()
def _build_cache_plan(
*,
targets: np.ndarray,
num_steps: int,
steps_per_block: int,
memory_blocks: int,
step_seconds: float,
anchor_policy: str,
ordinary_anchor_hz: float,
transition_before_seconds: float,
transition_after_seconds: float,
) -> tuple[list[dict[str, Any]], dict[int, set[int]], dict[str, Any]]:
"""Select train anchors and the minimum frozen-AV phase variants they need."""
if num_steps <= 0:
raise ValueError("num_steps must be positive")
if anchor_policy not in {"dense", "transition-sparse"}:
raise ValueError(f"Unknown anchor policy: {anchor_policy}")
num_blocks = math.ceil(num_steps / steps_per_block)
phase_counts = {
block_index: min(
steps_per_block,
num_steps - block_index * steps_per_block,
)
for block_index in range(num_blocks)
}
valid_anchor_count = max(
0,
min(num_steps, int(len(targets)) - NUM_HORIZONS),
)
transition_indices = _transition_step_indices(targets, num_steps)
if anchor_policy == "dense":
selected_indices = list(range(num_steps))
ordinary_indices = set(selected_indices)
transition_band_indices: set[int] = set()
required_variants = {
block_index: set(range(1, phase_count + 1))
for block_index, phase_count in phase_counts.items()
}
ordinary_stride_steps = 1
else:
if ordinary_anchor_hz <= 0:
raise ValueError("--ordinary-anchor-hz must be positive")
ordinary_stride_steps = _integer_ratio(
1.0 / ordinary_anchor_hz,
step_seconds,
"ordinary anchor period / step",
)
if ordinary_stride_steps <= 0:
raise ValueError("Ordinary anchor stride must be positive")
if transition_before_seconds < 0 or transition_after_seconds < 0:
raise ValueError("Transition sampling durations must be non-negative")
before_steps = _integer_ratio(
transition_before_seconds,
step_seconds,
"transition-before / step",
)
after_steps = _integer_ratio(
transition_after_seconds,
step_seconds,
"transition-after / step",
)
ordinary_indices = {
step_index
for step_index in range(valid_anchor_count)
if (step_index + 1) % ordinary_stride_steps == 0
}
transition_band_indices = set()
for transition_index in transition_indices:
band_start = max(0, transition_index - before_steps)
band_end = min(valid_anchor_count - 1, transition_index + after_steps)
if band_end >= band_start:
transition_band_indices.update(range(band_start, band_end + 1))
selected_indices = sorted(ordinary_indices | transition_band_indices)
if valid_anchor_count and not selected_indices:
raise ValueError("Sparse anchor policy selected no valid training anchors")
required_variants: dict[int, set[int]] = {}
# Every complete source block can be reused as retained causal history.
for block_index, phase_count in phase_counts.items():
if phase_count == steps_per_block:
required_variants.setdefault(block_index, set()).add(steps_per_block)
# A selected current anchor additionally needs its exact causal phase.
for step_index in selected_indices:
block_index = step_index // steps_per_block
tail_steps = step_index % steps_per_block + 1
required_variants.setdefault(block_index, set()).add(tail_steps)
anchors: list[dict[str, Any]] = []
for step_index in selected_indices:
anchor = make_anchor(
step_index=step_index,
steps_per_block=steps_per_block,
memory_blocks=memory_blocks,
step_seconds=step_seconds,
num_target_steps=len(targets),
num_horizons=NUM_HORIZONS,
)
if anchor_policy == "transition-sparse":
reasons = []
if step_index in ordinary_indices:
reasons.append("ordinary_grid")
if step_index in transition_band_indices:
reasons.append("transition_band")
anchor["selection_reasons"] = reasons
anchors.append(anchor)
valid_selected_anchor_count = sum(
bool(anchor["has_future_target"]) for anchor in anchors
)
complete_history_variants = sum(
steps_per_block in variants for variants in required_variants.values()
)
stats = {
"policy": anchor_policy,
"dense_timeline_steps": int(num_steps),
"dense_valid_training_anchors": int(valid_anchor_count),
"manifest_anchors": len(anchors),
"selected_training_anchors": int(valid_selected_anchor_count),
"selected_fraction_of_valid_dense": (
valid_selected_anchor_count / valid_anchor_count
if valid_anchor_count
else 0.0
),
"ordinary_anchor_hz": ordinary_anchor_hz,
"ordinary_stride_steps": ordinary_stride_steps,
"ordinary_grid_anchors": len(ordinary_indices),
"target_transition_count": len(transition_indices),
"transition_before_seconds": transition_before_seconds,
"transition_after_seconds": transition_after_seconds,
"transition_band_anchors": len(transition_band_indices),
"ordinary_transition_overlap": len(
ordinary_indices & transition_band_indices
),
"cached_phase_variants": sum(
len(variants) for variants in required_variants.values()
),
"dense_phase_variants": int(num_steps),
"complete_history_variants": int(complete_history_variants),
"transition_definition": "any_change_in_target_three_class_vector",
"selection_uses_labels_only_for_offline_training_sampling": True,
}
return anchors, required_variants, stats
def _print_cache_plan(stats: dict[str, Any]) -> None:
print("Anchor/cache selection plan")
print(f" policy: {stats['policy']}")
print(
" train anchors: "
f"{stats['selected_training_anchors']} / "
f"{stats['dense_valid_training_anchors']} "
f"({100.0 * stats['selected_fraction_of_valid_dense']:.2f}%)"
)
if stats["policy"] == "transition-sparse":
print(
" ordinary/transition/overlap: "
f"{stats['ordinary_grid_anchors']} / "
f"{stats['transition_band_anchors']} / "
f"{stats['ordinary_transition_overlap']}"
)
print(f" detected target transitions: {stats['target_transition_count']}")
print(
" cached AV phase variants: "
f"{stats['cached_phase_variants']} / {stats['dense_phase_variants']} "
f"({100.0 * stats['cached_phase_variants'] / max(1, stats['dense_phase_variants']):.2f}%)"
)
print(
" reusable complete 2s variants: "
f"{stats['complete_history_variants']}"
)
def _safe_output_path(path: Path, overwrite: bool) -> None:
if path.exists() and not overwrite:
raise FileExistsError(f"Refusing to overwrite {path}; pass --overwrite")
@@ -298,6 +522,7 @@ def _cache_blocks(
output_dir: Path,
num_steps: int,
steps_per_block: int,
required_variants: dict[int, set[int]],
step_seconds: float,
block_seconds: float,
av_batch_size: int,
@@ -305,18 +530,20 @@ def _cache_blocks(
audio_right_padding_samples: int | None,
overwrite: bool,
) -> list[dict[str, Any]]:
num_blocks = math.ceil(num_steps / steps_per_block)
block_plan = sorted(
block_index for block_index, variants in required_variants.items() if variants
)
if not block_plan:
raise ValueError("Cache plan contains no AV phase variants")
records: list[dict[str, Any]] = []
input_embeddings = model.get_input_embeddings()
progress = tqdm(
total=num_blocks,
total=len(block_plan),
desc=f"cache AV blocks (batch={av_batch_size})",
unit="block",
)
for tile_start in range(0, num_blocks, av_batch_size):
block_indices = list(
range(tile_start, min(tile_start + av_batch_size, num_blocks))
)
for tile_start in range(0, len(block_plan), av_batch_size):
block_indices = block_plan[tile_start : tile_start + av_batch_size]
variants_by_block: dict[int, dict[str, dict[str, torch.Tensor]]] = {
index: {} for index in block_indices
}
@@ -331,6 +558,16 @@ def _cache_blocks(
for index in block_indices
}
for block_index in block_indices:
invalid = sorted(
tail_steps
for tail_steps in required_variants[block_index]
if tail_steps < 1 or tail_steps > phase_counts[block_index]
)
if invalid:
raise ValueError(
f"Block {block_index} requests unavailable phases {invalid}; "
f"phase_count={phase_counts[block_index]}"
)
_safe_output_path(
output_dir / cache_block_filename(block_index), overwrite
)
@@ -339,7 +576,7 @@ def _cache_blocks(
active_indices = [
index
for index in block_indices
if tail_steps <= phase_counts[index]
if tail_steps in required_variants[index]
]
if not active_indices:
continue
@@ -405,6 +642,7 @@ def _cache_blocks(
"block_index": block_index,
"path": path.name,
"phase_count": phase_counts[block_index],
"cached_tail_steps": sorted(required_variants[block_index]),
"token_lengths_without_query": lengths_by_block[block_index],
"size_bytes": path.stat().st_size,
}
@@ -497,8 +735,8 @@ def _verify_cache(
def main(argv: Sequence[str] | None = None) -> None:
args = parse_args(argv)
core._configure_qwen_logging()
if args.end <= 0:
raise ValueError("--end must be positive")
if args.end < 0:
raise ValueError("--end must be non-negative; 0 selects the full target")
if args.av_batch_size <= 0:
raise ValueError("--av-batch-size must be positive")
if (
@@ -520,12 +758,31 @@ def main(argv: Sequence[str] | None = None) -> None:
memory_blocks = _integer_ratio(
args.memory_seconds, args.native_block_seconds, "memory / native block"
)
num_steps = _integer_ratio(args.end, args.step_seconds, "end / step")
targets = _load_targets(args.target)
if args.end == 0:
num_steps = int(len(targets))
end_seconds = num_steps * args.step_seconds
else:
num_steps = _integer_ratio(args.end, args.step_seconds, "end / step")
end_seconds = args.end
if num_steps > len(targets):
raise ValueError(
f"Requested {num_steps} cached steps but target has only {len(targets)}"
)
anchors, required_variants, selection_stats = _build_cache_plan(
targets=targets,
num_steps=num_steps,
steps_per_block=steps_per_block,
memory_blocks=memory_blocks,
step_seconds=args.step_seconds,
anchor_policy=args.anchor_policy,
ordinary_anchor_hz=args.ordinary_anchor_hz,
transition_before_seconds=args.transition_before_seconds,
transition_after_seconds=args.transition_after_seconds,
)
_print_cache_plan(selection_stats)
if args.selection_only:
return
output_dir = args.output_dir.expanduser().resolve()
output_dir.mkdir(parents=True, exist_ok=True)
@@ -547,7 +804,7 @@ def main(argv: Sequence[str] | None = None) -> None:
media = CausalSourceAV(
specs,
output_fps=args.video_fps,
required_duration=args.end,
required_duration=end_seconds,
frame_cache_size=args.frame_cache_size,
decoder_backend=args.decoder,
decoder_gpu_id=args.decoder_gpu_id,
@@ -555,7 +812,7 @@ def main(argv: Sequence[str] | None = None) -> None:
max_open_readers=args.max_open_readers,
)
print(
f"Direct media ready: duration={args.end:.1f}s, "
f"Direct media ready: duration={end_seconds:.1f}s, "
f"grid={media.expected_grid_frames} frames, audio={len(media.audio)} samples"
)
@@ -611,6 +868,7 @@ def main(argv: Sequence[str] | None = None) -> None:
output_dir=output_dir,
num_steps=num_steps,
steps_per_block=steps_per_block,
required_variants=required_variants,
step_seconds=args.step_seconds,
block_seconds=args.native_block_seconds,
av_batch_size=args.av_batch_size,
@@ -622,19 +880,8 @@ def main(argv: Sequence[str] | None = None) -> None:
print(
f"AV cache generation: {cache_seconds:.3f}s, "
f"{len(block_records) / cache_seconds:.3f} block/s, "
f"{args.end / cache_seconds:.3f}x media realtime"
f"{end_seconds / cache_seconds:.3f}x media realtime"
)
anchors = [
make_anchor(
step_index=step_index,
steps_per_block=steps_per_block,
memory_blocks=memory_blocks,
step_seconds=args.step_seconds,
num_target_steps=len(targets),
num_horizons=NUM_HORIZONS,
)
for step_index in range(num_steps)
]
checks = _verify_cache(
model=model,
extractor=extractor,
@@ -647,6 +894,11 @@ def main(argv: Sequence[str] | None = None) -> None:
step_seconds=args.step_seconds,
)
size = cache_size_bytes(
[static_path, *[output_dir / record["path"] for record in block_records]]
)
selection_stats["cache_tensor_size_bytes"] = int(size)
selection_stats["cache_tensor_size_gib"] = size / (1024 ** 3)
manifest = {
"schema_version": SCHEMA_VERSION,
"cache_kind": "qwen2.5-omni-frozen-av-open-stream-phases",
@@ -655,7 +907,7 @@ def main(argv: Sequence[str] | None = None) -> None:
"target_path": str(args.target.expanduser().resolve()),
"target_steps": int(len(targets)),
"cached_steps": num_steps,
"cached_end_seconds": args.end,
"cached_end_seconds": end_seconds,
"hidden_size": int(query_embeddings.shape[-1]),
"floating_dtype": "bfloat16",
"step_seconds": args.step_seconds,
@@ -669,7 +921,12 @@ def main(argv: Sequence[str] | None = None) -> None:
"video_max_pixels": args.video_max_pixels,
"use_audio_in_video": True,
"interval_convention": "half_open_[start,end)",
"phase_policy": "cache_0.2_to_2.0s;_2.0s_is_completed_history",
"phase_policy": (
"cache_0.2_to_2.0s;_2.0s_is_completed_history"
if args.anchor_policy == "dense"
else "cache_selected_current_phases;_all_complete_2.0s_history_blocks"
),
"anchor_selection": selection_stats,
"av_batch_size": args.av_batch_size,
"audio_feature_device": (
None if audio_feature_device is None else str(audio_feature_device)
@@ -679,7 +936,7 @@ def main(argv: Sequence[str] | None = None) -> None:
"processor-default-device-and-model-maximum-padding"
),
"cache_generation_seconds": cache_seconds,
"cache_media_realtime_factor": args.end / cache_seconds,
"cache_media_realtime_factor": end_seconds / cache_seconds,
"cache_blocks_per_second": len(block_records) / cache_seconds,
"label_policy": "target[i+1:i+11]",
"static_path": static_path.name,
@@ -693,9 +950,6 @@ def main(argv: Sequence[str] | None = None) -> None:
},
}
save_json(manifest_path, manifest)
size = cache_size_bytes(
[static_path, *[output_dir / record["path"] for record in block_records]]
)
print(
f"Saved {len(block_records)} blocks / {len(anchors)} anchors to {output_dir}"
)