Add transition-sparse Omni AV caching

README.md

@@ -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
+0.115 GiB, which extrapolates to about 17.25 GiB per media hour.  In the latest
-complete 2-second history blocks and recomputing every current partial online
+controlled single-L20 test without gradient checkpointing, 200 updates took
-reduced storage, but the single-L20 pilot slowed 200 LoRA updates from 83 s
+57.718 s (3.465 update/s) with every partial cached, 130.151 s (1.537
-(2.39 update/s) to 254.085 s (0.787 update/s).  Online AV preparation used
+update/s) with current partials recomputed using processor defaults, and
-98.654 s for 163 partials and peak allocation was 9.56 GiB, so the bottleneck
+111.146 s (1.799 update/s) after CUDA log-Mel and bounded-audio-padding
-is repeated preprocessing/modality encoding rather than memory capacity.
+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

@@ -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:
+Dense mode caches all ten causal phases (0.2, 0.4, ..., 2.0 seconds) for every
-0.2, 0.4, ..., 2.0 seconds.  The 2.0-second phase is reused as completed
+source-aligned 2-second block.  Transition-sparse mode instead caches every
-history; phases 0.2--1.8 are used only for the current partial block.  No
+complete 2.0-second history block, a low-rate grid of ordinary current
-future samples enter a phase.
+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):
+    for tile_start in range(0, len(block_plan), av_batch_size):
-        block_indices = list(
+        block_indices = block_plan[tile_start : tile_start + av_batch_size]
-            range(tile_start, min(tile_start + av_batch_size, num_blocks))
-        )
         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:
+    if args.end < 0:
-        raise ValueError("--end must be positive")
+        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}"
     )