366 lines
14 KiB
Python
366 lines
14 KiB
Python
|
"""Cached open-stream AV pieces for Qwen2.5-Omni turn-taking LoRA.
|
||
|
|
|
||
|
|
The modality towers are frozen for the first fine-tuning experiment. This
|
||
|
|
module serializes their exact output embeddings and Qwen's local TMRoPE
|
||
|
|
coordinates, then reconstructs the same still-open media turn without using a
|
||
|
|
KV cache. The resulting sequence can be passed through the Thinker with
|
||
|
|
``use_cache=False`` so gradients reach LoRA adapters in every retained block.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import math
|
||
|
|
from collections import OrderedDict
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Iterable, Sequence
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
SCHEMA_VERSION = 1
|
||
|
|
PIECE_TENSOR_NAMES = (
|
||
|
|
"prefix_embeddings",
|
||
|
|
"prefix_position_ids",
|
||
|
|
"payload_embeddings",
|
||
|
|
"payload_position_ids",
|
||
|
|
"closure_ids",
|
||
|
|
"closure_position_ids",
|
||
|
|
"closure_embeddings",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def load_json(path: str | Path) -> dict[str, Any]:
|
||
|
|
return json.loads(Path(path).read_text())
|
||
|
|
|
||
|
|
|
||
|
|
def save_json(path: str | Path, payload: dict[str, Any]) -> None:
|
||
|
|
Path(path).write_text(json.dumps(payload, indent=2) + "\n")
|
||
|
|
|
||
|
|
|
||
|
|
def torch_load(path: str | Path) -> Any:
|
||
|
|
"""Load tensor-only cache files across old and new PyTorch defaults."""
|
||
|
|
|
||
|
|
try:
|
||
|
|
return torch.load(Path(path), map_location="cpu", weights_only=True)
|
||
|
|
except TypeError:
|
||
|
|
return torch.load(Path(path), map_location="cpu")
|
||
|
|
|
||
|
|
|
||
|
|
def cache_block_filename(block_index: int) -> str:
|
||
|
|
return f"block_{block_index:06d}.pt"
|
||
|
|
|
||
|
|
|
||
|
|
def pieces_to_cache(
|
||
|
|
pieces: Any,
|
||
|
|
*,
|
||
|
|
closure_embeddings: torch.Tensor,
|
||
|
|
floating_dtype: torch.dtype = torch.bfloat16,
|
||
|
|
) -> dict[str, torch.Tensor]:
|
||
|
|
"""Detach one ``OpenMediaPieces`` object into a portable CPU payload."""
|
||
|
|
|
||
|
|
output: dict[str, torch.Tensor] = {}
|
||
|
|
for name in PIECE_TENSOR_NAMES:
|
||
|
|
value = closure_embeddings if name == "closure_embeddings" else getattr(pieces, name)
|
||
|
|
value = value.detach().cpu().contiguous()
|
||
|
|
if value.is_floating_point():
|
||
|
|
value = value.to(floating_dtype)
|
||
|
|
output[name] = value
|
||
|
|
validate_piece(output)
|
||
|
|
return output
|
||
|
|
|
||
|
|
|
||
|
|
def validate_piece(piece: dict[str, torch.Tensor]) -> None:
|
||
|
|
missing = [name for name in PIECE_TENSOR_NAMES if name not in piece]
|
||
|
|
if missing:
|
||
|
|
raise ValueError(f"Cached AV piece is missing {missing}")
|
||
|
|
prefix = piece["prefix_embeddings"]
|
||
|
|
payload = piece["payload_embeddings"]
|
||
|
|
closure = piece["closure_embeddings"]
|
||
|
|
if prefix.ndim != 3 or prefix.shape[0] != 1:
|
||
|
|
raise ValueError(f"Invalid prefix embeddings: {tuple(prefix.shape)}")
|
||
|
|
if payload.ndim != 3 or payload.shape[0] != 1:
|
||
|
|
raise ValueError(f"Invalid payload embeddings: {tuple(payload.shape)}")
|
||
|
|
if closure.ndim != 3 or closure.shape[0] != 1:
|
||
|
|
raise ValueError(f"Invalid closure embeddings: {tuple(closure.shape)}")
|
||
|
|
hidden_sizes = {int(prefix.shape[-1]), int(payload.shape[-1]), int(closure.shape[-1])}
|
||
|
|
if len(hidden_sizes) != 1:
|
||
|
|
raise ValueError(f"Cached piece hidden sizes disagree: {sorted(hidden_sizes)}")
|
||
|
|
expected_positions = (
|
||
|
|
("prefix_position_ids", prefix.shape[1]),
|
||
|
|
("payload_position_ids", payload.shape[1]),
|
||
|
|
("closure_position_ids", closure.shape[1]),
|
||
|
|
)
|
||
|
|
for name, length in expected_positions:
|
||
|
|
value = piece[name]
|
||
|
|
if value.shape != (3, 1, length):
|
||
|
|
raise ValueError(
|
||
|
|
f"{name} must have shape {(3, 1, length)}, got {tuple(value.shape)}"
|
||
|
|
)
|
||
|
|
if piece["closure_ids"].shape != (1, closure.shape[1]):
|
||
|
|
raise ValueError(
|
||
|
|
"closure_ids/closure_embeddings mismatch: "
|
||
|
|
f"{tuple(piece['closure_ids'].shape)} vs {tuple(closure.shape)}"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def make_anchor(
|
||
|
|
*,
|
||
|
|
step_index: int,
|
||
|
|
steps_per_block: int,
|
||
|
|
memory_blocks: int,
|
||
|
|
step_seconds: float,
|
||
|
|
num_target_steps: int,
|
||
|
|
num_horizons: int,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
"""Describe the exact completed-history/current-tail selection at time t."""
|
||
|
|
|
||
|
|
if step_index < 0:
|
||
|
|
raise ValueError("step_index must be non-negative")
|
||
|
|
block_index = step_index // steps_per_block
|
||
|
|
tail_steps = step_index % steps_per_block + 1
|
||
|
|
first_history = max(0, block_index - (memory_blocks - 1))
|
||
|
|
history_blocks = list(range(first_history, block_index))
|
||
|
|
return {
|
||
|
|
"step_index": int(step_index),
|
||
|
|
"time_seconds": float((step_index + 1) * step_seconds),
|
||
|
|
"block_index": int(block_index),
|
||
|
|
"tail_steps": int(tail_steps),
|
||
|
|
"history_blocks": history_blocks,
|
||
|
|
"target_index": int(step_index),
|
||
|
|
"has_future_target": bool(step_index + num_horizons < num_target_steps),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class CachedBlockStore:
|
||
|
|
"""Small CPU LRU for block shard files shared by adjacent samples."""
|
||
|
|
|
||
|
|
def __init__(self, cache_dir: str | Path, max_blocks: int = 16) -> None:
|
||
|
|
if max_blocks <= 0:
|
||
|
|
raise ValueError("max_blocks must be positive")
|
||
|
|
self.cache_dir = Path(cache_dir)
|
||
|
|
self.max_blocks = int(max_blocks)
|
||
|
|
self._blocks: OrderedDict[int, dict[str, Any]] = OrderedDict()
|
||
|
|
|
||
|
|
def get(self, block_index: int) -> dict[str, Any]:
|
||
|
|
cached = self._blocks.get(block_index)
|
||
|
|
if cached is not None:
|
||
|
|
self._blocks.move_to_end(block_index)
|
||
|
|
return cached
|
||
|
|
path = self.cache_dir / cache_block_filename(block_index)
|
||
|
|
if not path.is_file():
|
||
|
|
raise FileNotFoundError(path)
|
||
|
|
payload = torch_load(path)
|
||
|
|
if int(payload.get("block_index", -1)) != block_index:
|
||
|
|
raise ValueError(f"Block index mismatch in {path}")
|
||
|
|
variants = payload.get("variants")
|
||
|
|
if not isinstance(variants, dict) or not variants:
|
||
|
|
raise ValueError(f"No variants in {path}")
|
||
|
|
self._blocks[block_index] = payload
|
||
|
|
self._blocks.move_to_end(block_index)
|
||
|
|
while len(self._blocks) > self.max_blocks:
|
||
|
|
self._blocks.popitem(last=False)
|
||
|
|
return payload
|
||
|
|
|
||
|
|
def variant(self, block_index: int, tail_steps: int) -> dict[str, torch.Tensor]:
|
||
|
|
block = self.get(block_index)
|
||
|
|
variants = block["variants"]
|
||
|
|
piece = variants.get(str(tail_steps), variants.get(tail_steps))
|
||
|
|
if piece is None:
|
||
|
|
raise KeyError(
|
||
|
|
f"Block {block_index} has no tail_steps={tail_steps}; "
|
||
|
|
f"available={sorted(str(key) for key in variants)}"
|
||
|
|
)
|
||
|
|
validate_piece(piece)
|
||
|
|
return piece
|
||
|
|
|
||
|
|
|
||
|
|
def _strip_embedding_batch(value: torch.Tensor) -> torch.Tensor:
|
||
|
|
if value.ndim != 3 or value.shape[0] != 1:
|
||
|
|
raise ValueError(f"Expected [1,L,D], got {tuple(value.shape)}")
|
||
|
|
return value[0]
|
||
|
|
|
||
|
|
|
||
|
|
def _strip_position_batch(value: torch.Tensor) -> torch.LongTensor:
|
||
|
|
if value.ndim != 3 or value.shape[:2] != (3, 1):
|
||
|
|
raise ValueError(f"Expected [3,1,L], got {tuple(value.shape)}")
|
||
|
|
return value[:, 0, :].long()
|
||
|
|
|
||
|
|
|
||
|
|
def reconstruct_open_stream_sequence(
|
||
|
|
*,
|
||
|
|
store: CachedBlockStore,
|
||
|
|
static: dict[str, Any],
|
||
|
|
anchor: dict[str, Any],
|
||
|
|
) -> tuple[torch.Tensor, torch.LongTensor]:
|
||
|
|
"""Materialize one differentiable Thinker input from frozen AV embeddings.
|
||
|
|
|
||
|
|
Complete history blocks use their 2.0-second variant. Only the current
|
||
|
|
block uses its 0.2--2.0-second phase. The first retained block keeps the
|
||
|
|
chat/media prefix; later blocks append payload only. The final block's
|
||
|
|
media closure and the artificial assistant readout query are appended once.
|
||
|
|
"""
|
||
|
|
|
||
|
|
steps_per_block = int(static["steps_per_block"])
|
||
|
|
positions_per_chunk = int(static["positions_per_chunk"])
|
||
|
|
history_blocks = [int(value) for value in anchor["history_blocks"]]
|
||
|
|
current_block = int(anchor["block_index"])
|
||
|
|
selected: list[tuple[int, int]] = [
|
||
|
|
(block_index, steps_per_block) for block_index in history_blocks
|
||
|
|
]
|
||
|
|
selected.append((current_block, int(anchor["tail_steps"])))
|
||
|
|
if not selected:
|
||
|
|
raise ValueError("An anchor must select at least its current block")
|
||
|
|
|
||
|
|
embedding_parts: list[torch.Tensor] = []
|
||
|
|
position_parts: list[torch.LongTensor] = []
|
||
|
|
first_payload_base: int | None = None
|
||
|
|
last_piece: dict[str, torch.Tensor] | None = None
|
||
|
|
last_shift = 0
|
||
|
|
|
||
|
|
for sequence_index, (block_index, tail_steps) in enumerate(selected):
|
||
|
|
piece = store.variant(block_index, tail_steps)
|
||
|
|
local_payload_positions = _strip_position_batch(piece["payload_position_ids"])
|
||
|
|
local_payload_base = int(local_payload_positions.min().item())
|
||
|
|
if first_payload_base is None:
|
||
|
|
first_payload_base = local_payload_base
|
||
|
|
desired_payload_base = first_payload_base + sequence_index * positions_per_chunk
|
||
|
|
position_shift = desired_payload_base - local_payload_base
|
||
|
|
|
||
|
|
if sequence_index == 0:
|
||
|
|
embedding_parts.append(_strip_embedding_batch(piece["prefix_embeddings"]))
|
||
|
|
position_parts.append(_strip_position_batch(piece["prefix_position_ids"]))
|
||
|
|
embedding_parts.append(_strip_embedding_batch(piece["payload_embeddings"]))
|
||
|
|
position_parts.append(local_payload_positions + position_shift)
|
||
|
|
last_piece = piece
|
||
|
|
last_shift = position_shift
|
||
|
|
|
||
|
|
if last_piece is None:
|
||
|
|
raise RuntimeError("No cached AV piece was selected")
|
||
|
|
closure_embeddings = _strip_embedding_batch(last_piece["closure_embeddings"])
|
||
|
|
closure_positions = (
|
||
|
|
_strip_position_batch(last_piece["closure_position_ids"]) + last_shift
|
||
|
|
)
|
||
|
|
embedding_parts.append(closure_embeddings)
|
||
|
|
position_parts.append(closure_positions)
|
||
|
|
|
||
|
|
query_embeddings = static["query_embeddings"]
|
||
|
|
if query_embeddings.ndim != 3 or query_embeddings.shape[0] != 1:
|
||
|
|
raise ValueError(
|
||
|
|
f"static query_embeddings must be [1,Q,D], got {tuple(query_embeddings.shape)}"
|
||
|
|
)
|
||
|
|
query_embeddings = query_embeddings[0]
|
||
|
|
query_start = int(closure_positions.max().item()) + 1
|
||
|
|
query_scalar = torch.arange(
|
||
|
|
query_start,
|
||
|
|
query_start + query_embeddings.shape[0],
|
||
|
|
dtype=torch.long,
|
||
|
|
)
|
||
|
|
query_positions = query_scalar.view(1, -1).expand(3, -1)
|
||
|
|
embedding_parts.append(query_embeddings)
|
||
|
|
position_parts.append(query_positions)
|
||
|
|
|
||
|
|
embeddings = torch.cat(embedding_parts, dim=0).contiguous()
|
||
|
|
position_ids = torch.cat(position_parts, dim=1).contiguous()
|
||
|
|
if position_ids.shape != (3, embeddings.shape[0]):
|
||
|
|
raise RuntimeError(
|
||
|
|
"Reconstructed position/embedding mismatch: "
|
||
|
|
f"{tuple(position_ids.shape)} vs {tuple(embeddings.shape)}"
|
||
|
|
)
|
||
|
|
return embeddings, position_ids
|
||
|
|
|
||
|
|
|
||
|
|
def collate_cached_sequences(
|
||
|
|
batch: Sequence[dict[str, Any]],
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
"""Right-pad variable token sequences while preserving each readout index."""
|
||
|
|
|
||
|
|
if not batch:
|
||
|
|
raise ValueError("Cannot collate an empty batch")
|
||
|
|
hidden_size = int(batch[0]["inputs_embeds"].shape[-1])
|
||
|
|
max_length = max(int(item["inputs_embeds"].shape[0]) for item in batch)
|
||
|
|
dtype = batch[0]["inputs_embeds"].dtype
|
||
|
|
inputs_embeds = torch.zeros(
|
||
|
|
len(batch), max_length, hidden_size, dtype=dtype
|
||
|
|
)
|
||
|
|
position_ids = torch.zeros(3, len(batch), max_length, dtype=torch.long)
|
||
|
|
attention_mask = torch.zeros(len(batch), max_length, dtype=torch.long)
|
||
|
|
readout_indices = torch.empty(len(batch), dtype=torch.long)
|
||
|
|
targets = torch.stack([item["targets"] for item in batch])
|
||
|
|
|
||
|
|
for batch_index, item in enumerate(batch):
|
||
|
|
embeddings = item["inputs_embeds"]
|
||
|
|
positions = item["position_ids"]
|
||
|
|
length = int(embeddings.shape[0])
|
||
|
|
if embeddings.shape[-1] != hidden_size:
|
||
|
|
raise ValueError("Mixed hidden sizes in one batch")
|
||
|
|
if positions.shape != (3, length):
|
||
|
|
raise ValueError(
|
||
|
|
f"Invalid position_ids {tuple(positions.shape)} for length {length}"
|
||
|
|
)
|
||
|
|
inputs_embeds[batch_index, :length] = embeddings
|
||
|
|
position_ids[:, batch_index, :length] = positions
|
||
|
|
attention_mask[batch_index, :length] = 1
|
||
|
|
readout_indices[batch_index] = length - 1
|
||
|
|
|
||
|
|
return {
|
||
|
|
"inputs_embeds": inputs_embeds,
|
||
|
|
"position_ids": position_ids,
|
||
|
|
"attention_mask": attention_mask,
|
||
|
|
"readout_indices": readout_indices,
|
||
|
|
"targets": targets,
|
||
|
|
"step_indices": torch.tensor(
|
||
|
|
[int(item["step_index"]) for item in batch], dtype=torch.long
|
||
|
|
),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def anchor_intervals(
|
||
|
|
anchor: dict[str, Any],
|
||
|
|
*,
|
||
|
|
block_seconds: float,
|
||
|
|
step_seconds: float,
|
||
|
|
) -> list[tuple[float, float]]:
|
||
|
|
"""Return source intervals used by the reference KV replay for an anchor."""
|
||
|
|
|
||
|
|
intervals = [
|
||
|
|
(index * block_seconds, (index + 1) * block_seconds)
|
||
|
|
for index in anchor["history_blocks"]
|
||
|
|
]
|
||
|
|
block_index = int(anchor["block_index"])
|
||
|
|
current_start = block_index * block_seconds
|
||
|
|
current_end = current_start + int(anchor["tail_steps"]) * step_seconds
|
||
|
|
intervals.append((current_start, current_end))
|
||
|
|
return intervals
|
||
|
|
|
||
|
|
|
||
|
|
def forward_last_state(
|
||
|
|
language_model: torch.nn.Module,
|
||
|
|
*,
|
||
|
|
inputs_embeds: torch.Tensor,
|
||
|
|
position_ids: torch.LongTensor,
|
||
|
|
attention_mask: torch.LongTensor,
|
||
|
|
readout_indices: torch.LongTensor,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""Run full-sequence causal Thinker attention and gather each query state."""
|
||
|
|
|
||
|
|
outputs = language_model(
|
||
|
|
inputs_embeds=inputs_embeds,
|
||
|
|
position_ids=position_ids,
|
||
|
|
attention_mask=attention_mask,
|
||
|
|
use_cache=False,
|
||
|
|
output_hidden_states=False,
|
||
|
|
return_dict=True,
|
||
|
|
cache_position=torch.arange(
|
||
|
|
inputs_embeds.shape[1], device=inputs_embeds.device
|
||
|
|
),
|
||
|
|
)
|
||
|
|
batch_indices = torch.arange(
|
||
|
|
inputs_embeds.shape[0], device=inputs_embeds.device
|
||
|
|
)
|
||
|
|
return outputs.last_hidden_state[batch_indices, readout_indices]
|
||
|
|
|
||
|
|
|
||
|
|
def cache_size_bytes(paths: Iterable[Path]) -> int:
|
||
|
|
return sum(path.stat().st_size for path in paths if path.is_file())
|