From a9f748d6a6cd5dac5bbbc160a3c77e376a345d16 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Thu, 17 Sep 2026 02:32:08 +0800 Subject: [PATCH] Add Qwen tail and chunk ablations --- PROJECT_CONTEXT.md | 6 + README.md | 5 + docs/qwen_tail_chunk_ablation.md | 162 +++++++ run_smoke_easycom_direct.sh | 14 +- streaming_turn_model.py | 746 ++++++++++++++++++++++++++++++- streaming_turn_model_directly.py | 89 +++- test_streaming_chunk_ablation.py | 229 ++++++++++ 7 files changed, 1219 insertions(+), 32 deletions(-) create mode 100644 docs/qwen_tail_chunk_ablation.md create mode 100644 test_streaming_chunk_ablation.py diff --git a/PROJECT_CONTEXT.md b/PROJECT_CONTEXT.md index 5593fd4..3fd55a7 100644 --- a/PROJECT_CONTEXT.md +++ b/PROJECT_CONTEXT.md @@ -81,6 +81,12 @@ All are pending in that log; Omni LoRA is deferred and no semantic gate is compl ### Implemented but not GPU-verified +- Qwen2.5-Omni now has two matched chunk-handling controls around the existing + 5 Hz partial-tail path: native completed 2-second blocks held between + boundaries, and non-native 200-ms payload commits. Both preserve dense causal + alignment and save state-age/processing diagnostics, but still require the + documented real-GPU smoke test before a full run. See + [`docs/qwen_tail_chunk_ablation.md`](docs/qwen_tail_chunk_ablation.md). - MiniCPM-o 4.5 input-only feature extraction has two unmeasured modes: a native 1 Hz state causally held on the 5 Hz grid, and a 5 Hz completed-cache plus temporary-partial-tail branch. The latter refuses the normal diff --git a/README.md b/README.md index b7de0ea..603605f 100644 --- a/README.md +++ b/README.md @@ -1036,6 +1036,11 @@ be included. ## Direct extraction without a prepared 2 FPS MP4 +The direct extractor also implements the matched Qwen chunk-handling controls +`completed-block-hold` and `micro-chunk-200ms`. Their exact definitions, +required smoke checks, full EasyCom launch commands, and processing diagnostics +are in [`docs/qwen_tail_chunk_ablation.md`](docs/qwen_tail_chunk_ablation.md). + `streaming_turn_model_directly.py` removes the offline video-transcoding requirement. It reads original frame PTS with Decord, constructs a global 2 FPS lattice, maps every lattice point to the latest source frame whose PTS diff --git a/docs/qwen_tail_chunk_ablation.md b/docs/qwen_tail_chunk_ablation.md new file mode 100644 index 0000000..6b67fd4 --- /dev/null +++ b/docs/qwen_tail_chunk_ablation.md @@ -0,0 +1,162 @@ +# Qwen partial-tail and chunk-size ablation + +This runbook freezes the three-way comparison used to test whether causal +5 Hz adaptation needs an explicit partial tail. All modes use the same frozen +Qwen2.5-Omni backbone, right-edged 5 Hz output grid, 2 FPS causal video grid, +8-second memory budget, downstream head, split, and optimization. + +| Mode | New media at output time `t` | 5 Hz freshness | Interpretation | +|---|---|---:|---| +| `open-stream-replay` | completed native 2 s blocks plus a temporary growing tail of 0.2--2.0 s | yes | proposed partial-tail representation | +| `completed-block-hold` | only completed native 2 s blocks; discard the incomplete tail and hold the last state | no | native-chunk stale control | +| `micro-chunk-200ms` | commit every new 0.2 s interval as an independent payload | yes | non-native micro-chunk stress control | + +The 2-second hold mode does not backfill future information. Its first nine +rows (0.2--1.8 s) are zero initialization because no native block has yet +completed; the first model state appears at 2.0 s. Saved +`initialization_mask`, `state_source_times`, `state_update_mask`, +`payload_commit_mask`, and `native_boundary_mask` make this +auditable. Report a sensitivity analysis excluding initialization rows if +they have a non-negligible effect. + +The micro-chunk mode is deliberately not described as native Qwen streaming. +It uses the model's physical-time TMRoPE rate to allocate five positions per +0.2-second step and refuses to run if ten micro steps do not tile one native +2-second/50-position chunk exactly. It tests whether freshness alone is +sufficient when the model's native chunk construction is removed. + +Do not infer compute cost only from chunk duration. For a 2-second source +interval, partial-tail extraction prepares 0.2+...+2.0 = 11 seconds of media, +the hold control prepares 2 seconds in one call, and the micro control prepares +2 seconds across ten calls. Wall time, peak memory, processor calls, and saved +`production_prepared_media_seconds` should all be reported. + +## Required short GPU smoke tests + +Run each mode on the same four seconds before launching complete sessions: + +```bash +cd /data1/yfzx/wangx/Code/icassp2027-experiment + +MODE=open-stream-replay \ +SESSION_ID=4 GPU=0 END=4 SHORT_MEMORY_SECONDS=8 \ +FEATURE_ROOT=features/qwen_chunk_ablation_smoke \ +LOG_ROOT=logs/qwen_chunk_ablation_smoke \ +bash run_smoke_easycom_direct.sh + +MODE=completed-block-hold \ +SESSION_ID=4 GPU=0 END=4 SHORT_MEMORY_SECONDS=8 \ +FEATURE_ROOT=features/qwen_chunk_ablation_smoke \ +LOG_ROOT=logs/qwen_chunk_ablation_smoke \ +bash run_smoke_easycom_direct.sh + +MODE=micro-chunk-200ms \ +SESSION_ID=4 GPU=0 END=4 SHORT_MEMORY_SECONDS=8 \ +FEATURE_ROOT=features/qwen_chunk_ablation_smoke \ +LOG_ROOT=logs/qwen_chunk_ablation_smoke \ +bash run_smoke_easycom_direct.sh +``` + +The expected feature shape for every mode is `(20, D)`. Inspect the schedule +and processing accounting: + +```bash +python - <<'PY' +from pathlib import Path +import torch + +root = Path("features/qwen_chunk_ablation_smoke") +for mode in ( + "open-stream-replay", + "completed-block-hold", + "micro-chunk-200ms", +): + path = root / f"Session_4_{mode}_gpu0.pt" + payload = torch.load(path, map_location="cpu", weights_only=False) + diagnostics = payload["extraction_diagnostics"] + print("\n", mode, tuple(payload["states"].shape)) + print("processor calls:", diagnostics["production_processor_calls"]) + print( + "prepared media seconds:", + diagnostics["production_prepared_media_seconds"], + ) + print("source times:", diagnostics["state_source_times"].tolist()) + print("initialization rows:", int(diagnostics["initialization_mask"].sum())) + print("state updates:", diagnostics["state_update_mask"].nonzero().flatten().tolist()) + print("payload commits:", diagnostics["payload_commit_mask"].nonzero().flatten().tolist()) + print("native boundaries:", diagnostics["native_boundary_mask"].nonzero().flatten().tolist()) + print("verification:", payload["verification"]) +PY +``` + +For four seconds, expected production accounting is: + +| Mode | Processor calls | Prepared media seconds | State update rows | Payload commit rows | +|---|---:|---:|---|---| +| `open-stream-replay` | 20 | 22.0 | every row | 9, 19 | +| `completed-block-hold` | 2 | 4.0 | 9, 19 | 9, 19 | +| `micro-chunk-200ms` | 20 | 4.0 | every row | every row | + +All three modes should report native 2-second boundaries at rows 9 and 19; +that timing fact is distinct from whether a mode updates or commits there. + +`cache_replay` checks should be numerically exact or effectively exact. +`micro_chunk_vs_continuous_semantics` is a diagnostic comparison, not an +equivalence requirement: the micro representation intentionally changes input +segmentation. + +## Extract all EasyCom sessions + +The stable multi-session launcher already forwards `MODE`. Run the controls +separately so their outputs and logs cannot overwrite each other. Adjust GPU +IDs and roots to the remote host if needed. + +```bash +cd /data1/yfzx/wangx/Code/icassp2027-experiment + +PROJECT_ROOT=$PWD \ +MODEL=/data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \ +VIDEO_ROOT=/data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed \ +TARGET_ROOT=/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame \ +FEATURE_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/qwen_completed_block_hold_2fps_8s \ +LOG_ROOT=$PWD/logs/qwen_completed_block_hold_2fps_8s \ +MODE=completed-block-hold MEMORY_SECONDS=8 GPUS="0 1 2 3" \ +bash run_extract_easycom_direct.sh + +PROJECT_ROOT=$PWD \ +MODEL=/data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \ +VIDEO_ROOT=/data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed \ +TARGET_ROOT=/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame \ +FEATURE_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/qwen_micro_chunk_200ms_2fps_8s \ +LOG_ROOT=$PWD/logs/qwen_micro_chunk_200ms_2fps_8s \ +MODE=micro-chunk-200ms MEMORY_SECONDS=8 GPUS="0 1 2 3" \ +bash run_extract_easycom_direct.sh +``` + +Each output is mode-tagged. Existing training code therefore fails closed if +the wrong directory is paired with `--expected-mode`. + +## Train matched frozen heads + +Use the same seeds and training hyperparameters as the partial-tail baseline. +The commands below show one matched run per control: + +```bash +python train_turn_head.py \ + --feature-root /data1/yfzx/wangx/Datasets/easycom/omni_features/qwen_completed_block_hold_2fps_8s \ + --target-root /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame \ + --output checkpoints/easycom_qwen_completed_block_hold_8s.pt \ + --expected-mode completed-block-hold \ + --epochs 50 --batch-size 512 --device cuda:0 + +python train_turn_head.py \ + --feature-root /data1/yfzx/wangx/Datasets/easycom/omni_features/qwen_micro_chunk_200ms_2fps_8s \ + --target-root /data1/yfzx/wangx/Datasets/easycom/easycom_target_frame \ + --output checkpoints/easycom_qwen_micro_chunk_200ms_8s.pt \ + --expected-mode micro-chunk-200ms \ + --epochs 50 --batch-size 512 --device cuda:0 +``` + +This ablation is interpretable only when the partial-tail baseline is +re-extracted or otherwise confirmed to use the same 8-second memory, decoder, +2 FPS sampling, readout layers, split, seeds, and head training recipe. diff --git a/run_smoke_easycom_direct.sh b/run_smoke_easycom_direct.sh index b140835..56d54f7 100644 --- a/run_smoke_easycom_direct.sh +++ b/run_smoke_easycom_direct.sh @@ -16,16 +16,20 @@ LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_direct_smoke}" SESSION_ID="${SESSION_ID:-1}" GPU="${GPU:-0}" +MODE="${MODE:-open-stream-replay}" # Short smoke: END=4. Full session smoke: END="" (empty). END="${END:-4}" MEMORY_SECONDS="${MEMORY_SECONDS:-20}" +# The historical short smoke used persistent KV (0). Override this to 8 for +# the matched bounded-memory chunk ablation. +SHORT_MEMORY_SECONDS="${SHORT_MEMORY_SECONDS:-0}" VIDEO_FPS="${VIDEO_FPS:-2}" session="Session_${SESSION_ID}" session_dir="${VIDEO_ROOT}/${session}" target="${TARGET_ROOT}/${session}.npy" -output="${FEATURE_ROOT}/${session}_gpu${GPU}.pt" -log="${LOG_ROOT}/${session}_gpu${GPU}.log" +output="${FEATURE_ROOT}/${session}_${MODE}_gpu${GPU}.pt" +log="${LOG_ROOT}/${session}_${MODE}_gpu${GPU}.log" mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}" cd "${PROJECT_ROOT}" @@ -36,13 +40,13 @@ extra_args=() if [[ -n "${END}" ]]; then # Bounded smoke: do not pass --target (length would conflict with --end). extra_args+=(--end "${END}") - MEMORY_SECONDS=0 + MEMORY_SECONDS="${SHORT_MEMORY_SECONDS}" else [[ -s "${target}" ]] || { echo "Missing: ${target}" >&2; exit 1; } extra_args+=(--target "${target}") fi -echo "GPU=${GPU} session=${session} end=${END:-full} memory=${MEMORY_SECONDS}" +echo "GPU=${GPU} session=${session} mode=${MODE} end=${END:-full} memory=${MEMORY_SECONDS}" echo "log -> ${log}" { @@ -52,7 +56,7 @@ echo "log -> ${log}" --model-id "${MODEL}" \ --easycom-session-dir "${session_dir}" \ --use-audio-in-video \ - --mode open-stream-replay \ + --mode "${MODE}" \ --memory-seconds "${MEMORY_SECONDS}" \ --video-fps "${VIDEO_FPS}" \ --dtype bfloat16 \ diff --git a/streaming_turn_model.py b/streaming_turn_model.py index c96a810..52893a9 100644 --- a/streaming_turn_model.py +++ b/streaming_turn_model.py @@ -1,12 +1,16 @@ """Omni-only causal turn-taking features at 5 Hz. -This module implements three inference paths for Qwen2.5-Omni-3B: +This module implements five inference paths for Qwen2.5-Omni-3B: 1. ``causal-prefix``: recompute a causal AV window ending at every 200 ms. This is deliberately slow, but is the simplest correctness baseline. 2. ``open-stream-replay``: cache completed AV blocks inside one still-open user/media turn; closure and readout tokens remain temporary. -3. ``overlap-working-memory``: a frozen ablation that serializes cached +3. ``completed-block-hold``: discard the incomplete native block and hold the + most recent completed-block state on the 5 Hz output grid. +4. ``micro-chunk-200ms``: commit each 200-ms interval independently as an + explicitly non-native chunk-size stress test. +5. ``overlap-working-memory``: a frozen ablation that serializes cached complete blocks and an overlapping ``[t-2,t]`` recent recap. Neither path uses an RNN. A small head maps one Thinker state to the same @@ -37,7 +41,7 @@ import logging import math import subprocess from collections import OrderedDict -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Any, Literal @@ -63,6 +67,8 @@ from transformers.cache_utils import DynamicCache Mode = Literal[ "causal-prefix", "open-stream-replay", + "completed-block-hold", + "micro-chunk-200ms", "overlap-working-memory", ] BackboneLoader = Literal["qwen", "roma"] @@ -389,6 +395,7 @@ class ExtractionResult: times: torch.Tensor # [N], seconds in the original session states: torch.Tensor # [N,D] or preserved query layers [N,L,D] verification: list[dict[str, Any]] + diagnostics: dict[str, Any] = field(default_factory=dict) @dataclass @@ -621,6 +628,29 @@ def _time_grid(start: float, end: float, step: float) -> list[float]: return [start + i * step for i in range(1, count + 1)] +def _periodic_boundary_mask( + times: list[float], + *, + origin: float, + period: float, +) -> torch.BoolTensor: + """Mark right edges aligned to a positive multiple of ``period``.""" + + values = [] + for current_time in times: + ratio = (current_time - origin) / period + values.append( + ratio > 0 + and math.isclose( + ratio, + round(ratio), + rel_tol=0.0, + abs_tol=1e-6, + ) + ) + return torch.tensor(values, dtype=torch.bool) + + def _media_content( *, video_path: str | Path | None, @@ -1076,6 +1106,31 @@ def _tmrope_positions_per_chunk( return int(rounded) +def _tmrope_positions_per_step( + model: Qwen2_5OmniThinkerForConditionalGeneration, + step_seconds: float, +) -> int: + """Return the native TMRoPE advance for one micro-chunk step.""" + + positions_per_second = float(model.config.position_id_per_seconds) + value = positions_per_second * step_seconds + rounded = round(value) + if not math.isclose(value, rounded, rel_tol=0.0, abs_tol=1e-6): + raise ValueError(f"Non-integral TMRoPE positions per step: {value}") + steps_per_chunk = _require_integer_ratio( + float(model.config.seconds_per_chunk), + step_seconds, + "model native chunk / micro step", + ) + positions_per_chunk = _tmrope_positions_per_chunk(model) + if rounded * steps_per_chunk != positions_per_chunk: + raise ValueError( + "Micro-step TMRoPE advances do not tile one native chunk: " + f"{rounded} * {steps_per_chunk} != {positions_per_chunk}" + ) + return int(rounded) + + def _cache_length(cache: Any | None) -> int: if cache is None: return 0 @@ -1906,10 +1961,634 @@ class OpenStreamReplayExtractor: global_step_offset=0, progress=progress, ) + steps_per_block = _require_integer_ratio( + self.config.native_block_seconds, + self.config.step_seconds, + "native block / step", + ) return ExtractionResult( times=torch.tensor(all_times, dtype=torch.float64), states=torch.cat(all_states, dim=0), verification=all_checks, + diagnostics={ + "production_processor_calls": len(all_times), + "production_prepared_media_seconds": float( + sum( + ((index % steps_per_block) + 1) + * self.config.step_seconds + for index in range(len(all_times)) + ) + ), + "state_source_times": torch.tensor( + all_times, dtype=torch.float64 + ), + "initialization_mask": torch.zeros( + len(all_times), dtype=torch.bool + ), + "state_update_mask": torch.ones( + len(all_times), dtype=torch.bool + ), + "payload_commit_mask": _periodic_boundary_mask( + all_times, + origin=start, + period=self.config.native_block_seconds, + ), + "native_boundary_mask": _periodic_boundary_mask( + all_times, + origin=start, + period=self.config.native_block_seconds, + ), + }, + ) + + +class CompletedBlockHoldExtractor(OpenStreamReplayExtractor): + """Commit native 2-s blocks and causally hold the last emitted state. + + This is the deliberately stale control for partial-tail processing. It + never constructs an incomplete AV interval and never backfills a state + into rows before that block was fully observed. + """ + + def _zero_state(self) -> torch.Tensor: + hidden_size = _hidden_size(self.model) + dtype = getattr(self.model, "dtype", torch.float32) + if self.config.readout_layer_fusion == "stack": + shape = (1, self.config.readout_last_n_layers, hidden_size) + else: + shape = (1, hidden_size) + return torch.zeros(shape, dtype=dtype, device="cpu") + + def extract( + self, + *, + video_path: str | Path | None, + audio_path: str | Path | None, + start: float, + end: float, + use_audio_in_video: bool, + verify_every: int = 0, + verify_continuous_every: int = 0, + show_progress: bool = True, + ) -> ExtractionResult: + if verify_every < 0 or verify_continuous_every < 0: + raise ValueError("verification intervals must be >= 0") + times = _time_grid(start, end, self.config.step_seconds) + committed_cache = None + block_start = start + history_start = start + payload_base = None + positions_per_chunk = _tmrope_positions_per_chunk(self.model) + latest = self._zero_state() + source_time = start + states: list[torch.Tensor] = [] + source_times: list[float] = [] + updates: list[bool] = [] + initialization: list[bool] = [] + checks: list[dict[str, Any]] = [] + processor_calls = 0 + prepared_seconds = 0.0 + replay_block_count = max( + 0, + _require_integer_ratio( + self.config.memory_seconds, + self.config.native_block_seconds, + "memory / native block", + ) + - 1, + ) + needs_rolling_rebuild = ( + self.config.memory_seconds < end - start - 1e-6 + ) + completed_blocks = CompletedBlockPieceCache( + origin=start, + block_seconds=self.config.native_block_seconds, + max_blocks=replay_block_count if needs_rolling_rebuild else 0, + ) + + with tqdm( + total=len(times), + desc="completed-block-hold", + unit="state", + dynamic_ncols=True, + mininterval=0.0, + miniters=1, + disable=not show_progress, + ) as progress: + for global_step, current_time in enumerate(times, start=1): + block_end = block_start + self.config.native_block_seconds + update = math.isclose( + current_time, + block_end, + rel_tol=0.0, + abs_tol=1e-6, + ) + if update: + if ( + committed_cache is not None + and block_start - history_start + >= self.config.memory_seconds - 1e-6 + ): + history_start = block_start - ( + self.config.memory_seconds + - self.config.native_block_seconds + ) + history_intervals = _block_intervals( + history_start, + block_start, + self.config.native_block_seconds, + ) + if history_intervals: + prepared_history = completed_blocks.get_many( + history_intervals + ) + ( + committed_cache, + _, + payload_base, + _, + _, + ) = self._cache_from_intervals( + video_path=video_path, + audio_path=audio_path, + intervals=history_intervals, + use_audio_in_video=use_audio_in_video, + prepared_blocks=prepared_history, + ) + else: + committed_cache = None + payload_base = None + + pieces = self._prepare_pieces( + video_path=video_path, + audio_path=audio_path, + interval=(block_start, block_end), + use_audio_in_video=use_audio_in_video, + ) + processor_calls += 1 + prepared_seconds += block_end - block_start + if payload_base is None: + payload_base = int( + pieces.payload_position_ids.min().item() + ) + block_index = round( + (block_start - history_start) + / self.config.native_block_seconds + ) + payload_offset = ( + payload_base + block_index * positions_per_chunk + ) + ( + committed_cache, + _, + closure_position_ids, + ) = _append_open_pieces( + self.model, + pieces, + past_key_values=committed_cache, + rope_offset=payload_offset, + include_prefix=committed_cache is None, + ) + latest = _read_open_state( + self.model, + committed_cache, + closure_ids=pieces.closure_ids, + closure_position_ids=closure_position_ids, + query_ids=self.query_ids, + last_n_layers=self.config.readout_last_n_layers, + layer_fusion=self.config.readout_layer_fusion, + ).cpu() + source_time = current_time + completed_blocks.put( + CachedOpenMediaBlock( + start=block_start, + end=block_end, + pieces=pieces, + ) + ) + + if verify_every > 0 and global_step % verify_every == 0: + exact = self.state_from_scratch( + video_path=video_path, + audio_path=audio_path, + intervals=_block_intervals( + history_start, + current_time, + self.config.native_block_seconds, + ), + use_audio_in_video=use_audio_in_video, + ) + checks.append( + _compare_states( + current_time, + latest, + exact, + kind="completed_hold_cache_replay", + ) + ) + if ( + verify_continuous_every > 0 + and global_step % verify_continuous_every == 0 + ): + continuous = self.continuous_reference.state_for_intervals( + video_path=video_path, + audio_path=audio_path, + intervals=[(history_start, current_time)], + use_audio_in_video=use_audio_in_video, + ) + checks.append( + _compare_states( + current_time, + latest, + continuous, + kind="completed_hold_continuous_semantics", + ) + ) + block_start = block_end + + states.append(latest) + source_times.append(source_time) + updates.append(update) + initialization.append(source_time <= start + 1e-6) + progress.set_postfix_str( + f"source={source_time:.1f}s", refresh=False + ) + progress.update(1) + + return ExtractionResult( + times=torch.tensor(times, dtype=torch.float64), + states=torch.cat(states, dim=0), + verification=checks, + diagnostics={ + "production_processor_calls": processor_calls, + "production_prepared_media_seconds": float(prepared_seconds), + "state_source_times": torch.tensor( + source_times, dtype=torch.float64 + ), + "initialization_mask": torch.tensor( + initialization, dtype=torch.bool + ), + "state_update_mask": torch.tensor(updates, dtype=torch.bool), + "payload_commit_mask": torch.tensor(updates, dtype=torch.bool), + "native_boundary_mask": torch.tensor(updates, dtype=torch.bool), + }, + ) + + +class MicroChunkReplayExtractor(OpenStreamReplayExtractor): + """Commit each new 200-ms payload as a non-native streaming stress test.""" + + def __init__( + self, + model: Qwen2_5OmniThinkerForConditionalGeneration, + processor: Qwen2_5OmniProcessor, + config: StreamConfig, + fixed_grid_av: FixedGridAV | None = None, + ) -> None: + super().__init__(model, processor, config, fixed_grid_av=fixed_grid_av) + self.positions_per_step = _tmrope_positions_per_step( + model, config.step_seconds + ) + + def _prepare_micro_pieces( + self, + *, + video_path: str | Path | None, + audio_path: str | Path | None, + interval: tuple[float, float], + use_audio_in_video: bool, + ) -> OpenMediaPieces: + duration = interval[1] - interval[0] + if not math.isclose( + duration, + self.config.step_seconds, + rel_tol=0.0, + abs_tol=1e-6, + ): + raise ValueError( + f"Micro interval is {duration:g}s, expected " + f"{self.config.step_seconds:g}s" + ) + pieces = self._prepare_pieces( + video_path=video_path, + audio_path=audio_path, + interval=interval, + use_audio_in_video=use_audio_in_video, + ) + temporal_positions = pieces.payload_position_ids[0] + temporal_span = int( + temporal_positions.max().item() + - temporal_positions.min().item() + ) + if temporal_span >= self.positions_per_step: + raise RuntimeError( + "A 200-ms payload does not fit its native TMRoPE step: " + f"interval={interval}, temporal_span={temporal_span}, " + f"step_positions={self.positions_per_step}" + ) + return pieces + + def _micro_cache_from_intervals( + self, + *, + video_path: str | Path | None, + audio_path: str | Path | None, + intervals: list[tuple[float, float]], + use_audio_in_video: bool, + prepared_blocks: list[CachedOpenMediaBlock] | None = None, + ) -> tuple[Any, int, torch.LongTensor, torch.LongTensor]: + if not intervals: + raise ValueError("At least one micro interval is required") + if prepared_blocks is not None and len(prepared_blocks) != len(intervals): + raise ValueError("prepared micro-block count does not match intervals") + media_cache = None + payload_base = None + closure_ids = None + closure_position_ids = None + for index, interval in enumerate(intervals): + if prepared_blocks is None: + pieces = self._prepare_micro_pieces( + video_path=video_path, + audio_path=audio_path, + interval=interval, + use_audio_in_video=use_audio_in_video, + ) + else: + prepared = prepared_blocks[index] + if not ( + math.isclose( + prepared.start, + interval[0], + rel_tol=0.0, + abs_tol=1e-6, + ) + and math.isclose( + prepared.end, + interval[1], + rel_tol=0.0, + abs_tol=1e-6, + ) + ): + raise ValueError("prepared micro-block interval mismatch") + pieces = prepared.pieces + if payload_base is None: + payload_base = int(pieces.payload_position_ids.min().item()) + payload_offset = payload_base + index * self.positions_per_step + ( + media_cache, + _, + closure_position_ids, + ) = _append_open_pieces( + self.model, + pieces, + past_key_values=media_cache, + rope_offset=payload_offset, + include_prefix=index == 0, + ) + closure_ids = pieces.closure_ids + if media_cache is None or payload_base is None: + raise RuntimeError("Micro-chunk replay produced no cache") + if closure_ids is None or closure_position_ids is None: + raise RuntimeError("Micro-chunk replay produced no closure") + return media_cache, payload_base, closure_ids, closure_position_ids + + def state_from_micro_scratch( + self, + *, + video_path: str | Path | None, + audio_path: str | Path | None, + intervals: list[tuple[float, float]], + use_audio_in_video: bool, + ) -> torch.Tensor: + cache, _, closure_ids, closure_positions = ( + self._micro_cache_from_intervals( + video_path=video_path, + audio_path=audio_path, + intervals=intervals, + use_audio_in_video=use_audio_in_video, + ) + ) + return _read_open_state( + self.model, + cache, + closure_ids=closure_ids, + closure_position_ids=closure_positions, + query_ids=self.query_ids, + last_n_layers=self.config.readout_last_n_layers, + layer_fusion=self.config.readout_layer_fusion, + ) + + def extract( + self, + *, + video_path: str | Path | None, + audio_path: str | Path | None, + start: float, + end: float, + use_audio_in_video: bool, + verify_every: int = 0, + verify_continuous_every: int = 0, + show_progress: bool = True, + ) -> ExtractionResult: + if verify_every < 0 or verify_continuous_every < 0: + raise ValueError("verification intervals must be >= 0") + times = _time_grid(start, end, self.config.step_seconds) + committed_cache = None + block_start = start + history_start = start + payload_base = None + states: list[torch.Tensor] = [] + checks: list[dict[str, Any]] = [] + prepared_seconds = 0.0 + needs_rolling_rebuild = ( + self.config.memory_seconds < end - start - 1e-6 + ) + replay_step_count = max( + 0, + _require_integer_ratio( + self.config.memory_seconds - self.config.native_block_seconds, + self.config.step_seconds, + "micro replay memory / step", + ), + ) + prepared_cache = CompletedBlockPieceCache( + origin=start, + block_seconds=self.config.step_seconds, + max_blocks=replay_step_count if needs_rolling_rebuild else 0, + ) + + with tqdm( + total=len(times), + desc="micro-chunk-200ms", + unit="state", + dynamic_ncols=True, + mininterval=0.0, + miniters=1, + disable=not show_progress, + ) as progress: + for global_step, current_time in enumerate(times, start=1): + while ( + current_time + > block_start + self.config.native_block_seconds + 1e-6 + ): + block_start += self.config.native_block_seconds + + if ( + committed_cache is not None + and block_start - history_start + >= self.config.memory_seconds - 1e-6 + ): + history_start = block_start - ( + self.config.memory_seconds + - self.config.native_block_seconds + ) + history_intervals = _block_intervals( + history_start, + block_start, + self.config.step_seconds, + ) + if history_intervals: + prepared_history = prepared_cache.get_many( + history_intervals + ) + ( + committed_cache, + payload_base, + _, + _, + ) = self._micro_cache_from_intervals( + video_path=video_path, + audio_path=audio_path, + intervals=history_intervals, + use_audio_in_video=use_audio_in_video, + prepared_blocks=prepared_history, + ) + else: + committed_cache = None + payload_base = None + + interval_start = current_time - self.config.step_seconds + pieces = self._prepare_micro_pieces( + video_path=video_path, + audio_path=audio_path, + interval=(interval_start, current_time), + use_audio_in_video=use_audio_in_video, + ) + prepared_seconds += current_time - interval_start + if payload_base is None: + payload_base = int(pieces.payload_position_ids.min().item()) + micro_index = round( + (interval_start - history_start) + / self.config.step_seconds + ) + payload_offset = ( + payload_base + micro_index * self.positions_per_step + ) + ( + committed_cache, + _, + closure_positions, + ) = _append_open_pieces( + self.model, + pieces, + past_key_values=committed_cache, + rope_offset=payload_offset, + include_prefix=committed_cache is None, + ) + state = _read_open_state( + self.model, + committed_cache, + closure_ids=pieces.closure_ids, + closure_position_ids=closure_positions, + query_ids=self.query_ids, + last_n_layers=self.config.readout_last_n_layers, + layer_fusion=self.config.readout_layer_fusion, + ) + states.append(state.cpu()) + prepared_cache.put( + CachedOpenMediaBlock( + start=interval_start, + end=current_time, + pieces=pieces, + ) + ) + + if verify_every > 0 and global_step % verify_every == 0: + intervals = _block_intervals( + history_start, + current_time, + self.config.step_seconds, + ) + exact = self.state_from_micro_scratch( + video_path=video_path, + audio_path=audio_path, + intervals=intervals, + use_audio_in_video=use_audio_in_video, + ) + checks.append( + _compare_states( + current_time, + state, + exact, + kind="micro_chunk_cache_replay", + ) + ) + if ( + verify_continuous_every > 0 + and global_step % verify_continuous_every == 0 + ): + continuous = self.continuous_reference.state_for_intervals( + video_path=video_path, + audio_path=audio_path, + intervals=[(history_start, current_time)], + use_audio_in_video=use_audio_in_video, + ) + checks.append( + _compare_states( + current_time, + state, + continuous, + kind="micro_chunk_vs_continuous_semantics", + ) + ) + + if math.isclose( + current_time, + block_start + self.config.native_block_seconds, + rel_tol=0.0, + abs_tol=1e-6, + ): + block_start = current_time + progress.set_postfix_str( + f"media={current_time:.1f}s", refresh=False + ) + progress.update(1) + + output_times = torch.tensor(times, dtype=torch.float64) + return ExtractionResult( + times=output_times, + states=torch.cat(states, dim=0), + verification=checks, + diagnostics={ + "production_processor_calls": len(times), + "production_prepared_media_seconds": float(prepared_seconds), + "state_source_times": output_times.clone(), + "initialization_mask": torch.zeros( + len(times), dtype=torch.bool + ), + "state_update_mask": torch.ones( + len(times), dtype=torch.bool + ), + "payload_commit_mask": torch.ones( + len(times), dtype=torch.bool + ), + "native_boundary_mask": _periodic_boundary_mask( + times, + origin=start, + period=self.config.native_block_seconds, + ), + "tmrope_positions_per_micro_step": self.positions_per_step, + }, ) @@ -2240,8 +2919,8 @@ def _compare_states( *, kind: str, ) -> dict[str, Any]: - cached_f = cached.float() - exact_f = exact.float() + cached_f = cached.detach().float() + exact_f = exact.detach().to(device=cached_f.device, dtype=torch.float32) cached_flat = cached_f.reshape(1, -1) exact_flat = exact_f.reshape(1, -1) return { @@ -2394,7 +3073,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--output", type=Path, required=True) parser.add_argument( "--mode", - choices=("causal-prefix", "open-stream-replay"), + choices=( + "causal-prefix", + "open-stream-replay", + "completed-block-hold", + "micro-chunk-200ms", + ), default="open-stream-replay", ) parser.add_argument("--model-id", default=MODEL_ID) @@ -2554,7 +3238,12 @@ def main() -> None: show_progress=not args.no_progress, ) else: - extractor = OpenStreamReplayExtractor( + extractor_class = { + "open-stream-replay": OpenStreamReplayExtractor, + "completed-block-hold": CompletedBlockHoldExtractor, + "micro-chunk-200ms": MicroChunkReplayExtractor, + }[args.mode] + extractor = extractor_class( model, processor, config, fixed_grid_av=fixed_grid_av ) result = extractor.extract( @@ -2573,11 +3262,16 @@ def main() -> None: "backbone_loader": args.backbone_loader, "backbone_metadata": backbone_metadata(model), "mode": args.mode, - "stream_representation": ( - "single_causal_window" - if args.mode == "causal-prefix" - else "single_open_media_turn_with_2s_payload_commits" - ), + "stream_representation": { + "causal-prefix": "single_causal_window", + "open-stream-replay": ( + "native_2s_completed_blocks_plus_temporary_partial_tail" + ), + "completed-block-hold": ( + "native_2s_completed_blocks_forward_held_to_5hz" + ), + "micro-chunk-200ms": "non_native_200ms_payload_commits", + }[args.mode], "config": asdict(config), "start": args.start, "end": end, @@ -2614,9 +3308,37 @@ def main() -> None: if args.mode == "open-stream-replay" else None ), + "chunk_ablation_policy": ( + { + "kind": "native_completed_block_hold", + "native_chunk_seconds": config.native_block_seconds, + "output_step_seconds": config.step_seconds, + "tail_policy": "discard_incomplete_tail_and_hold_last_state", + "initial_state": "zeros_until_first_complete_native_block", + "fresh_at_every_5hz_step": False, + "backbone": "frozen", + } + if args.mode == "completed-block-hold" + else ( + { + "kind": "non_native_micro_chunk_stress_test", + "micro_chunk_seconds": config.step_seconds, + "native_chunk_seconds": config.native_block_seconds, + "tmrope_positions_per_micro_step": ( + _tmrope_positions_per_step(model, config.step_seconds) + ), + "tail_policy": "commit_each_5hz_interval_as_a_payload", + "fresh_at_every_5hz_step": True, + "backbone": "frozen", + } + if args.mode == "micro-chunk-200ms" + else None + ) + ), "times": result.times, "states": result.states, "verification": result.verification, + "extraction_diagnostics": result.diagnostics, } if args.head_checkpoint is not None: head, head_loss_name = _load_head( diff --git a/streaming_turn_model_directly.py b/streaming_turn_model_directly.py index f2416c7..4f0ba10 100644 --- a/streaming_turn_model_directly.py +++ b/streaming_turn_model_directly.py @@ -1025,6 +1025,8 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: choices=( "causal-prefix", "open-stream-replay", + "completed-block-hold", + "micro-chunk-200ms", "overlap-working-memory", ), default="open-stream-replay", @@ -1414,8 +1416,17 @@ def main( use_audio_in_video=True, show_progress=not args.no_progress, ) - elif args.mode == "open-stream-replay": - extractor = core.OpenStreamReplayExtractor( + elif args.mode in { + "open-stream-replay", + "completed-block-hold", + "micro-chunk-200ms", + }: + extractor_class = { + "open-stream-replay": core.OpenStreamReplayExtractor, + "completed-block-hold": core.CompletedBlockHoldExtractor, + "micro-chunk-200ms": core.MicroChunkReplayExtractor, + }[args.mode] + extractor = extractor_class( model, processor, config, fixed_grid_av=reader ) result = extractor.extract( @@ -1459,25 +1470,33 @@ def main( ) persistent_memory = ( - args.mode == "open-stream-replay" and args.memory_seconds == 0 + args.mode + in { + "open-stream-replay", + "completed-block-hold", + "micro-chunk-200ms", + } + and args.memory_seconds == 0 ) payload: dict[str, Any] = { "model_id": args.model_id, "backbone_loader": args.backbone_loader, "backbone_metadata": core.backbone_metadata(model), "mode": args.mode, - "stream_representation": ( - "single_causal_prefix" - if args.mode == "causal-prefix" - else ( - "single_open_media_turn_with_2s_payload_commits" - if args.mode == "open-stream-replay" - else ( - "serialized_cached_complete_blocks_then_overlapping_" - "global_history_current_frame_recent_recap" - ) - ) - ), + "stream_representation": { + "causal-prefix": "single_causal_prefix", + "open-stream-replay": ( + "native_2s_completed_blocks_plus_temporary_partial_tail" + ), + "completed-block-hold": ( + "native_2s_completed_blocks_forward_held_to_5hz" + ), + "micro-chunk-200ms": "non_native_200ms_payload_commits", + "overlap-working-memory": ( + "serialized_cached_complete_blocks_then_overlapping_" + "global_history_current_frame_recent_recap" + ), + }[args.mode], "config": asdict(config), "requested_memory_seconds": args.memory_seconds, "persistent_thinker_kv": persistent_memory, @@ -1514,6 +1533,45 @@ def main( if args.mode == "open-stream-replay" else None ), + "chunk_ablation_policy": ( + { + "kind": "native_completed_block_hold", + "native_chunk_seconds": config.native_block_seconds, + "output_step_seconds": config.step_seconds, + "tail_policy": "discard_incomplete_tail_and_hold_last_state", + "initial_state": "zeros_until_first_complete_native_block", + "fresh_at_every_5hz_step": False, + "memory": ( + "persistent_full_extraction_thinker_kv" + if persistent_memory + else "rolling_native_completed_blocks" + ), + "backbone": "frozen", + } + if args.mode == "completed-block-hold" + else ( + { + "kind": "non_native_micro_chunk_stress_test", + "micro_chunk_seconds": config.step_seconds, + "native_chunk_seconds": config.native_block_seconds, + "tmrope_positions_per_micro_step": ( + core._tmrope_positions_per_step( + model, config.step_seconds + ) + ), + "tail_policy": "commit_each_5hz_interval_as_a_payload", + "fresh_at_every_5hz_step": True, + "memory": ( + "persistent_full_extraction_thinker_kv" + if persistent_memory + else "rolling_micro_payloads_on_native_boundaries" + ), + "backbone": "frozen", + } + if args.mode == "micro-chunk-200ms" + else None + ) + ), "overlap_working_memory_policy": ( { "total_presentation_seconds": config.memory_seconds, @@ -1556,6 +1614,7 @@ def main( "times": result.times, "states": result.states, "verification": result.verification, + "extraction_diagnostics": result.diagnostics, } if args.head_checkpoint is not None: head, head_loss_name = core._load_head( diff --git a/test_streaming_chunk_ablation.py b/test_streaming_chunk_ablation.py new file mode 100644 index 0000000..ff07e62 --- /dev/null +++ b/test_streaming_chunk_ablation.py @@ -0,0 +1,229 @@ +"""CPU-only schedule and feature-contract tests for Qwen chunk ablations.""" + +from __future__ import annotations + +import tempfile +import types +import unittest +from pathlib import Path +from unittest import mock + +import numpy as np +import torch + +import streaming_turn_model as core +from easycom_head import load_session_data +from local_turn_features import reference_alignment + + +class FakeModel: + def __init__(self) -> None: + self.config = types.SimpleNamespace( + seconds_per_chunk=2.0, + position_id_per_seconds=25.0, + text_config=types.SimpleNamespace(hidden_size=3), + ) + self.dtype = torch.float32 + self.device = torch.device("cpu") + + +class FakeCache: + def __init__(self, count: int) -> None: + self.count = count + + +def fake_pieces() -> core.OpenMediaPieces: + return core.OpenMediaPieces( + prefix_embeddings=torch.zeros(1, 1, 3), + prefix_position_ids=torch.zeros(3, 1, 1, dtype=torch.long), + payload_embeddings=torch.zeros(1, 1, 3), + payload_position_ids=torch.zeros(3, 1, 1, dtype=torch.long), + closure_ids=torch.zeros(1, 1, dtype=torch.long), + closure_position_ids=torch.ones(3, 1, 1, dtype=torch.long), + ) + + +class RecordingMixin: + def _prepare_pieces(self, **kwargs): # type: ignore[no-untyped-def] + self.prepared_intervals.append(tuple(kwargs["interval"])) + return fake_pieces() + + +class RecordingHold(RecordingMixin, core.CompletedBlockHoldExtractor): + def __init__(self, *args, **kwargs): # type: ignore[no-untyped-def] + super().__init__(*args, **kwargs) + self.prepared_intervals: list[tuple[float, float]] = [] + + +class RecordingMicro(RecordingMixin, core.MicroChunkReplayExtractor): + def __init__(self, *args, **kwargs): # type: ignore[no-untyped-def] + super().__init__(*args, **kwargs) + self.prepared_intervals: list[tuple[float, float]] = [] + + +def fake_append( # type: ignore[no-untyped-def] + model, + pieces, + *, + past_key_values, + rope_offset, + include_prefix, +): + count = 0 if past_key_values is None else past_key_values.count + return FakeCache(count + 1), rope_offset + 1, pieces.closure_position_ids + + +def fake_read(model, cache, **kwargs): # type: ignore[no-untyped-def] + return torch.full((1, 3), float(cache.count), dtype=torch.float32) + + +class QwenChunkAblationTests(unittest.TestCase): + def setUp(self) -> None: + self.model = FakeModel() + self.processor = object() + self.config = core.StreamConfig(memory_seconds=4.0) + self.patches = [ + mock.patch.object( + core, + "_readout_query_ids", + return_value=torch.zeros(1, 1, dtype=torch.long), + ), + mock.patch.object(core, "_append_open_pieces", side_effect=fake_append), + mock.patch.object(core, "_read_open_state", side_effect=fake_read), + ] + for patcher in self.patches: + patcher.start() + self.addCleanup(patcher.stop) + + def test_completed_native_blocks_are_held_without_future_backfill(self) -> None: + extractor = RecordingHold(self.model, self.processor, self.config) + result = extractor.extract( + video_path="unused.mp4", + audio_path=None, + start=0.0, + end=4.0, + use_audio_in_video=True, + show_progress=False, + ) + + self.assertEqual(tuple(result.states.shape), (20, 3)) + self.assertTrue(torch.equal(result.states[:9], torch.zeros(9, 3))) + self.assertTrue(torch.equal(result.states[9:19], torch.ones(10, 3))) + self.assertTrue(torch.equal(result.states[19], torch.full((3,), 2.0))) + self.assertEqual( + result.diagnostics["state_update_mask"].nonzero().flatten().tolist(), + [9, 19], + ) + self.assertTrue( + torch.equal( + result.diagnostics["state_update_mask"], + result.diagnostics["payload_commit_mask"], + ) + ) + self.assertTrue( + torch.equal( + result.diagnostics["state_update_mask"], + result.diagnostics["native_boundary_mask"], + ) + ) + self.assertTrue(result.diagnostics["initialization_mask"][:9].all()) + self.assertFalse(result.diagnostics["initialization_mask"][9:].any()) + self.assertEqual( + result.diagnostics["state_source_times"].tolist(), + [0.0] * 9 + [2.0] * 10 + [4.0], + ) + self.assertEqual(extractor.prepared_intervals, [(0.0, 2.0), (2.0, 4.0)]) + self.assertEqual(result.diagnostics["production_processor_calls"], 2) + self.assertAlmostEqual( + result.diagnostics["production_prepared_media_seconds"], 4.0 + ) + + def test_micro_chunks_update_every_200ms_without_growing_tail(self) -> None: + config = core.StreamConfig(memory_seconds=2.0) + extractor = RecordingMicro(self.model, self.processor, config) + result = extractor.extract( + video_path="unused.mp4", + audio_path=None, + start=0.0, + end=2.0, + use_audio_in_video=True, + show_progress=False, + ) + + expected = torch.arange(1, 11, dtype=torch.float32).view(-1, 1) + expected = expected.expand(-1, 3) + self.assertTrue(torch.equal(result.states, expected)) + self.assertTrue(result.diagnostics["state_update_mask"].all()) + self.assertTrue(result.diagnostics["payload_commit_mask"].all()) + self.assertEqual( + result.diagnostics["native_boundary_mask"] + .nonzero() + .flatten() + .tolist(), + [9], + ) + self.assertTrue( + torch.equal( + result.diagnostics["state_source_times"], result.times + ) + ) + self.assertEqual(len(extractor.prepared_intervals), 10) + for index, interval in enumerate(extractor.prepared_intervals): + self.assertAlmostEqual(interval[0], index * 0.2) + self.assertAlmostEqual(interval[1], (index + 1) * 0.2) + self.assertEqual(result.diagnostics["production_processor_calls"], 10) + self.assertAlmostEqual( + result.diagnostics["production_prepared_media_seconds"], 2.0 + ) + self.assertEqual(result.diagnostics["tmrope_positions_per_micro_step"], 5) + + def test_micro_tmrope_step_must_tile_the_native_chunk(self) -> None: + self.assertEqual(core._tmrope_positions_per_step(self.model, 0.2), 5) + with self.assertRaisesRegex(ValueError, "Non-integral TMRoPE"): + core._tmrope_positions_per_step(self.model, 0.15) + + def test_both_modes_load_through_head_and_local_alignment_contracts(self) -> None: + states = torch.randn(20, 8) + targets = np.tile( + np.asarray( + [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 1]], + dtype=np.float32, + ), + (5, 1), + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + feature_root = root / "features" + target_root = root / "targets" + feature_root.mkdir() + target_root.mkdir() + feature_path = feature_root / "Session_4.pt" + np.save(target_root / "Session_4.npy", targets) + + for mode in ("completed-block-hold", "micro-chunk-200ms"): + payload = { + "states": states, + "times": (torch.arange(20, dtype=torch.float64) + 1) / 5, + "label_indices": torch.arange(20), + "mode": mode, + "config": {"step_seconds": 0.2}, + "start": 0.0, + } + torch.save(payload, feature_path) + loaded = load_session_data( + feature_root, + target_root, + 4, + expected_mode=mode, + ) + self.assertTrue(torch.equal(loaded.states, states)) + times, indices = reference_alignment(payload) + self.assertTrue(torch.equal(times, payload["times"])) + self.assertTrue(torch.equal(indices, payload["label_indices"])) + + with self.assertRaisesRegex(ValueError, "expected 'causal-prefix'"): + load_session_data(feature_root, target_root, 4) + + +if __name__ == "__main__": + unittest.main()