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

Add Qwen tail and chunk ablations

Browse Source
This commit is contained in:
Xin Wang
2026-09-17 02:32:08 +08:00
parent a183fe5944
commit a9f748d6a6
7 changed files with 1219 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
PROJECT_CONTEXT.md
+6
View File
@@ -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
README.md
+5
View File
@@ -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
docs/qwen_tail_chunk_ablation.md
+162
View File
@@ -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.
run_smoke_easycom_direct.sh
+9 -5
View File
@@ -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 \
streaming_turn_model.py
+734 -12
View File
@@ -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(
streaming_turn_model_directly.py
+74 -15
View File
@@ -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(
test_streaming_chunk_ablation.py
+229
View File
@@ -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()