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
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}"
)