diff --git a/README.md b/README.md index 4e49ed4..856e238 100644 --- a/README.md +++ b/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 -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 diff --git a/cache_omni_av_pieces.py b/cache_omni_av_pieces.py index 4382e87..e0ca0e8 100644 --- a/cache_omni_av_pieces.py +++ b/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: -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}" )