File size: 9,311 Bytes
5f0e4a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 | """Read-only capture of exact v2 trajectories and selected AR context KV."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Callable, Mapping
import torch
from .schema import LATENT_HEIGHT, LATENT_WIDTH, NUM_STEPS
from .v2_schema import CONTEXT_BLOCK_IDS, block_key
def _clone_cpu(tensor: torch.Tensor) -> torch.Tensor:
return tensor.detach().to(device="cpu").contiguous()
@dataclass
class _ActiveStep:
chunk_id: int
step_id: int
tensors: dict[str, torch.Tensor]
shared: dict[str, torch.Tensor]
class PredictorV2TeacherCapture:
"""Capture Full-DiT targets plus step-invariant selected-layer AR memory."""
def __init__(
self,
transformer: torch.nn.Module,
*,
on_chunk: Callable[
[int, dict[str, torch.Tensor], dict[int, dict[str, torch.Tensor]]], None
],
num_steps: int = NUM_STEPS,
context_block_ids: tuple[int, ...] = CONTEXT_BLOCK_IDS,
) -> None:
self.transformer = transformer
self.on_chunk = on_chunk
self.num_steps = num_steps
self.context_block_ids = tuple(context_block_ids)
self.call_index = 0
self.active: _ActiveStep | None = None
self.chunk_steps: list[_ActiveStep] = []
self.image_condition_latent: torch.Tensor | None = None
self.text_context: dict[int, dict[str, torch.Tensor]] | None = None
self.chunk_context: dict[int, dict[str, torch.Tensor]] | None = None
self._handles: list[Any] = []
def __enter__(self) -> "PredictorV2TeacherCapture":
self._handles.append(
self.transformer.register_forward_pre_hook(self._transformer_pre, with_kwargs=True)
)
self._handles.append(
self.transformer.register_forward_hook(self._transformer_post, with_kwargs=True)
)
self._handles.append(
self.transformer.final_layer.register_forward_pre_hook(self._final_pre, with_kwargs=True)
)
return self
def __exit__(self, exc_type, exc, traceback) -> bool:
for handle in self._handles:
handle.remove()
self._handles.clear()
self.active = None
return False
def _capture_context(
self,
kv_cache: list[Mapping[str, torch.Tensor | None]],
) -> dict[int, dict[str, torch.Tensor]]:
result: dict[int, dict[str, torch.Tensor]] = {}
for block_id in self.context_block_ids:
cache = kv_cache[block_id]
k_txt = cache.get("k_txt")
v_txt = cache.get("v_txt")
if k_txt is None or v_txt is None:
raise RuntimeError(f"Block {block_id} text KV is unavailable")
k_vision = cache.get("k_vision")
v_vision = cache.get("v_vision")
if (k_vision is None) != (v_vision is None):
raise RuntimeError(f"Block {block_id} has incomplete vision KV")
if k_vision is None:
empty_shape = (2, int(k_txt.shape[1]), 0, int(k_txt.shape[3]))
k_vision_cpu = torch.empty(empty_shape, dtype=k_txt.dtype, device="cpu")
v_vision_cpu = torch.empty(empty_shape, dtype=v_txt.dtype, device="cpu")
else:
k_vision_cpu = _clone_cpu(k_vision)
v_vision_cpu = _clone_cpu(v_vision)
result[block_id] = {
"k_vision": k_vision_cpu,
"v_vision": v_vision_cpu,
}
if self.text_context is None:
result[block_id]["k_txt"] = _clone_cpu(k_txt)
result[block_id]["v_txt"] = _clone_cpu(v_txt)
if self.text_context is None:
self.text_context = {
block_id: {
"k_txt": result[block_id].pop("k_txt"),
"v_txt": result[block_id].pop("v_txt"),
}
for block_id in self.context_block_ids
}
return result
def _transformer_pre(self, module, args, kwargs) -> None:
is_denoise = (
kwargs.get("ar_vision_inference", False)
and not kwargs.get("cache_vision", False)
)
if not is_denoise:
return
if self.active is not None:
raise RuntimeError("Nested denoising capture is not supported")
chunk_id, step_id = divmod(self.call_index, self.num_steps)
model_input = kwargs["hidden_states"]
if model_input.shape[1] != 65:
raise ValueError(f"Teacher denoising input must have 65 channels, got {model_input.shape}")
if self.image_condition_latent is None:
self.image_condition_latent = _clone_cpu(model_input[:, 32:64, 0:1])
mask = model_input[:, 64:65]
if not torch.all(mask[:, :, 0] == 1) or not torch.all(mask[:, :, 1:] == 0):
raise ValueError("Unexpected I2V condition mask in first chunk")
if step_id == 0:
if self.chunk_context is not None:
raise RuntimeError("Previous chunk context was not flushed")
self.chunk_context = self._capture_context(kwargs["kv_cache"])
timestep = kwargs["timestep"].reshape(-1)[0:1]
shared = {
"action_labels": _clone_cpu(kwargs["action"].reshape(1, -1).round().long()),
"target_viewmats": _clone_cpu(kwargs["viewmats"]),
"target_Ks": _clone_cpu(kwargs["Ks"]),
"rope_temporal_size": torch.tensor([int(kwargs["rope_temporal_size"])], dtype=torch.int64),
"start_rope_start_idx": torch.tensor(
[int(kwargs["start_rope_start_idx"])], dtype=torch.int64
),
}
self.active = _ActiveStep(
chunk_id=chunk_id,
step_id=step_id,
tensors={
"timestep": _clone_cpu(timestep.float()),
"noisy_sample": _clone_cpu(model_input[:, :32]),
},
shared=shared,
)
def _final_pre(self, module, args, kwargs) -> None:
if self.active is None:
return
hidden, condition = args[0], args[1]
batch, tokens, hidden_size = hidden.shape
spatial_tokens = LATENT_HEIGHT * LATENT_WIDTH
if tokens % spatial_tokens:
raise ValueError(f"Final hidden token count {tokens} is not frame-aligned")
frames = tokens // spatial_tokens
compact = condition.reshape(batch, frames, spatial_tokens, hidden_size)[:, :, 0]
expanded = compact[:, :, None].expand(batch, frames, spatial_tokens, hidden_size)
if not torch.equal(
expanded.reshape(batch, tokens, hidden_size),
condition.reshape(batch, tokens, hidden_size),
):
raise ValueError("Final-layer condition varies inside a latent frame")
self.active.tensors["frame_condition"] = _clone_cpu(compact)
self.active.tensors["final_hidden"] = _clone_cpu(hidden)
def _transformer_post(self, module, args, kwargs, output) -> None:
if self.active is None:
return
velocity = output[0] if isinstance(output, tuple) else output
self.active.tensors["velocity"] = _clone_cpu(velocity)
required = {"timestep", "noisy_sample", "frame_condition", "final_hidden", "velocity"}
missing = required.difference(self.active.tensors)
if missing:
raise RuntimeError(f"Incomplete v2 teacher step capture: {sorted(missing)}")
self.chunk_steps.append(self.active)
completed_step = self.active.step_id
self.active = None
self.call_index += 1
if completed_step == self.num_steps - 1:
self._flush_chunk()
def _flush_chunk(self) -> None:
if len(self.chunk_steps) != self.num_steps:
raise RuntimeError(f"Expected {self.num_steps} steps, got {len(self.chunk_steps)}")
if self.chunk_context is None:
raise RuntimeError("Chunk context was not captured")
chunk_id = self.chunk_steps[0].chunk_id
if any(step.chunk_id != chunk_id for step in self.chunk_steps):
raise RuntimeError("Captured steps cross a chunk boundary")
tensors = dict(self.chunk_steps[0].shared)
for step in self.chunk_steps:
for name, tensor in step.tensors.items():
tensors[f"step_{step.step_id}_{name}"] = tensor
self.on_chunk(chunk_id, tensors, self.chunk_context)
self.chunk_steps.clear()
self.chunk_context = None
def case_tensors(self) -> dict[str, torch.Tensor]:
if self.image_condition_latent is None or self.text_context is None:
raise RuntimeError("Case condition or text context has not been captured")
result = {"image_condition_latent": self.image_condition_latent}
for block_id, cache in self.text_context.items():
result[block_key(block_id, "k_txt")] = cache["k_txt"]
result[block_key(block_id, "v_txt")] = cache["v_txt"]
return result
@property
def text_token_count(self) -> int:
if self.text_context is None:
raise RuntimeError("Text context has not been captured")
first = self.text_context[self.context_block_ids[0]]["k_txt"]
return int(first.shape[2])
|