Add files using upload-large-folder tool

backup_2026-09-14/code/baseline/SPEC.md

@@ -0,0 +1,88 @@
+# Baseline 动作臂契约（五臂公平消融）
+
+主实验：同一个 Wan2.2-TI2V-5B、同一份数据、同一套训练配方，**唯一自由变量是动作以什么方式进模型**：
+
+| 臂 | 代表作 | 机制（要按上游代码移植，不是凭印象） | 上游代码入口 |
+|---|---|---|---|
+| `arope` | 本项目 | 世界坐标 RoPE + known/new 掩码，0 参数 | `actionrope/arope.py` |
+| `linear` | ReactiveGWM | 每个 DiT block 一个线性层，把动作映射成 hidden 上的偏置 | `vender/ReactiveGWM/inference/models/dit.py`、`training/data/action_utils.py`、`inference/utils/actions.py` |
+| `xattn` | Matrix-Game 3 | 键盘/动作 token 与视频 token 沿时间轴做 cross-attention（前若干 block） | `vender/Matrix-Game/Matrix-Game-3/wan/modules/action_module.py`、`model.py`、`utils/conditions.py` |
+| `prompt` | Incantation | 动作走文本，逐 latent（逐 cell）的 prompt 控制 | `vender/Incantation/modules/causal_model.py`、`inference.py`、`prompts/` |
+| `adaln` | AlayaWorld | 动作嵌入进 adaLN 调制（与时间步嵌入同路） | `vender/AlayaWorld/alaya/control/action.py`、`alaya/model/components.py`、`ltx2/modules/model.py`、`timestep_embedding.py` |
+
+旧报告的口径（要对齐）：四臂全部从零起训、同数据同 seed；**带动作的臂收到完全相同的控制信号（逐 cell 累计偏移）**；
+文本一律剥除动作词（prompt 臂除外——文本就是它的通道）。Linear 移植自 ReactiveGWM 官方实现（30 个逐块线性偏置，+0.18M），
+Cross-Attn 复刻 Matrix-Game 3.0 键盘通路（前 15 块时间轴交叉注意力，+98.6M）。
+
+## 钩子接口（`actionrope/arope.py::arope_forward(..., arm=, action_inputs=)`）
+
+`arm` 是一个 `nn.Module`（`baseline/base.py::ActionArm` 的子类），`arope_forward` 在固定位置调用它；
+`arm=None` 时 `arope_forward` 与原实现逐位相同。调用顺序与形状：
+
+```python
+grid = dict(b, f=21, h=15, w=26, n_tok_per_frame=390)      # token 网格
+feats   = arm.encode(action_inputs, grid)                    # 一次；返回任意张量/tuple（建议 [b, f, d]），会作为张量参数穿过梯度检查点
+t_mod   = arm.modify_t_mod(t_mod, feats, grid)               # t_mod: [b, S, 6, dim]（逐 token，S = f·h·w；帧 0 的 t 已置 0）
+context = arm.extra_context(context, feats, grid)            # 文本 context: [b, L, dim] 或 [(b f), L, dim]（逐 cell）
+x       = arm.after_patch(x, feats, grid)                    # x: [b, S, dim]，patch embedding 之后、blocks 之前
+for i, block in enumerate(dit.blocks):                       # 每个 block 内部（在梯度检查点里）：
+    x = arm.block_pre(i, x, feats, grid)                     #   block 入口
+    ... self_attn ...
+    x = arm.block_mid(i, x, feats, grid)                     #   自注意力之后、文本 cross-attn 之前
+    ... cross_attn ...
+    x = arm.block_pre_ffn(i, x, feats, grid)                 #   文本 cross-attn 之后、FFN 之前（Matrix-Game 3 的注入位置）
+    ... ffn ...
+    x = arm.block_post(i, x, feats, grid)                    #   block 出口
+```
+
+token 排列是 `(f h w)`：token 序号 `s = f_idx * 390 + i_h * 26 + j_w`，所以 `[b, S, dim] → [b, f, h*w, dim]` 用 `x.view(b, f, -1, dim)`。
+
+`ActionArm` 基类（`baseline/base.py`）所有钩子默认恒等、`encode` 默认返回 None；子类只覆盖自己用到的。
+基类还提供：`install(dit)`（默认空；需要读 dit.dim / num_heads / 层数时用），`n_new_params()`，`zero_init_check()`。
+
+### 动作输入 `action_inputs`（dataset 统一提供，臂自己选用）
+
+```python
+{
+  "offset_px":  float32 [b, 21, 2],   # 逐 cell 累计玩家屏幕投影位移 (dx, dy) px，cell 0 = 0；与 ARoPE 同一个信号（sidecar）
+  "offset_tok": float32 [b, 21, 2],   # = offset_px / 32
+  "delta_tok":  float32 [b, 21, 2],   # 逐 cell 增量 = offset_tok[k] − offset_tok[k−1]，cell 0 = 0（速度）
+  "action_idx": int64   [b, 21],      # 9 类离散标签（0 idle，1 up … 按 prompts.ACTIONS），已翻回真实屏幕方向
+}
+```
+
+**主信号是 `offset_px` / `offset_tok`（累计偏移），与旧报告一致。** 上游代码用离散按键的（Matrix-Game 键盘）
+按其原样用 `action_idx`（one-hot / embedding），并在 docstring 里写明；若上游同时接受连续量（鼠标/相机），
+连续量用 `delta_tok`。每个臂在 `baseline/<name>.py` 顶部 docstring 写清：上游是怎么做的（文件:行）、这里怎么对应、
+用了哪个信号、新增参数量、初始化方式（**新增路径必须零初始化或等效地保证"装上瞬间输出与原模型相同"**）。
+
+### 文本
+
+`linear` / `xattn` / `adaln` / `arope`：`text_mode=scene`（动作词剥离，`text_table_actionrope.pt`）。
+`prompt`：`text_mode=scene_action`（动作词翻回真实方向，`text_table_eybx_mirror.pt`），逐 cell 文本已经是
+`arope_forward` 的原生能力（每帧 token 只 attend 自己 cell 的���本）。Incantation 若有别的逐 latent 做法
+（例如逐帧 prompt 掩码、专门的 attention），按上游移植并说明差异。
+
+### 训练 / 推理接入（integration 阶段做，臂的作者不要改这些文件）
+
+* `actionrope/train.py --arm {arope,plain,linear,xattn,prompt,adaln}`：`plain` 保留（= prompt 臂的旧名，行为不变）；
+  新臂：`install_arope(dit, mask_channel=False)`（普通 RoPE、无 mask、权重全 1）+ `arm = baseline.build_arm(name, dit)`，
+  arm 的参数与 DiT 一起训练（同 lr）；ckpt 导出时 arm 参数以 `arm.` 前缀一起写进 safetensors；config.json 记 arm 名与新增参数量。
+* `actionrope/infer.py --arm <name>`：按 ckpt 元数据/键集自动装 arm；动作脚本 → `action_inputs`（`offset_px` 同 ARoPE；
+  `action_idx` 由每 cell 速度方向得到，同 `velocity_to_label(mirror=False)`）。
+* `baseline/__init__.py::build_arm(name, dit) -> ActionArm`、`ARMS = {...}`。
+
+### 测试（每个臂 `tests/test_arm_<name>.py`，单卡，模型只加载一次）
+
+1. 装上后 **输出与原模型逐位相同**（零初始化）——`arope_forward(arm=arm)` vs `arope_forward()`。
+2. 给非零动作、把新增参数随机扰动后输出**有变化**且梯度能回到新增参数与 DiT（`use_gradient_checkpointing=True`）。
+3. 新增参数量打印并与上游/旧报告对照（linear ≈ 0.18M，xattn ≈ 98.6M；差得远要解释）。
+4. ckpt 键集：`arm.state_dict()` 键名无冲突、可 strict 加载。
+5. 单卡 2 步真实训练 smoke（`--limit 8 --max_steps 2`）在 integration 阶段跑。
+
+## 纪律
+
+只在 `/opt/dlami/nvme/zhiyangdeng/ActionRoPE` 内读写；新文件只放 `baseline/`、`tests/`；`vender/` 只读（不改、不删、不跑它们的脚本、不下载它们的权重）；
+不改 `actionrope/*.py`（钩子已加好；确需改动写进 report 交 integration 决定）；不动 data/、models/、DiffSynth-Studio/；
+禁止 rm -rf、sudo、pip 卸载/升降级、杀非自己启动的进程、任何网络上传/下载（vender 已克隆好）；不读凭据。
+GPU：正在跑 8 卡训练（每卡 24 GB 已用，143 GB 总），单卡前向+反向 smoke 用分配给你的那一张，跑完立刻释放。

backup_2026-09-14/code/baseline/__init__.py

@@ -0,0 +1,200 @@
+"""baseline 动作臂的注册表与装配入口（baseline/SPEC.md「训练 / 推理接入」）。
+
+    from baseline import ARMS, build_arm
+    arm = build_arm("linear", dit)          # 建参数 → install(dit) → 搬到 dit 的 device / dtype
+
+五臂里 `arope`（本项目）与 `plain`（= prompt 臂旧名）不走 ActionArm：arope 的动作在 RoPE 里，
+plain 与 prompt 行为完全相同（train.py 里 plain 仍按旧路径 arm=None 跑，prompt 走 PromptArm 的零钩子，
+两者前向逐位一致；保留 plain 是为了旧 ckpt / 旧脚本不改）。
+
+ckpt 约定：arm 参数以 `arm.` 前缀与 DiT 权重写在同一个 safetensors 里（train.py::export_dit_weights），
+safetensors metadata 里记 `arm=<name>`；`detect_arm_from_keys` 只看键集就能判断臂名与 xattn 的结构配置，
+metadata 缺失（手工拼的 ckpt）时推理仍能自动装对。
+"""
+from __future__ import annotations
+
+import json
+from typing import Iterable
+
+import torch
+
+from baseline.base import ActionArm
+
+# 各臂的类延迟导入：prompt.py 会把 data/code 加进 sys.path、xattn.py 依赖 einops，
+# 不用某个臂时不必付这些 import 的代价，也避免一个臂坏了拖垮别的臂。
+_ARM_MODULES = {
+    "linear": ("baseline.linear", "LinearArm"),
+    "xattn": ("baseline.xattn", "XAttnArm"),
+    "prompt": ("baseline.prompt", "PromptArm"),
+    "adaln": ("baseline.adaln", "AdaLNArm"),
+}
+# 不走 ActionArm 的两个臂（train.py / infer.py 的 --arm 也接受它们）
+NATIVE_ARMS = ("arope", "plain")
+# 全部合法臂名（CLI choices 用）
+ARM_NAMES = tuple(NATIVE_ARMS) + tuple(_ARM_MODULES)
+# 动作走文本的臂：text_mode=scene_action（动作词翻回真实方向）；其余剥动作从句
+TEXT_ARMS = ("plain", "prompt")
+# ckpt 里 arm 参数的前缀
+ARM_PREFIX = "arm."
+
+
+class _LazyArms(dict):
+    """ARMS[name] → 臂类；第一次取时才 import 对应模块。"""
+
+    def __missing__(self, name):
+        if name not in _ARM_MODULES:
+            raise KeyError(f"未知动作臂 {name!r}，可选 {list(_ARM_MODULES)}（arope / plain 不走 ActionArm）")
+        mod_name, cls_name = _ARM_MODULES[name]
+        import importlib
+
+        cls = getattr(importlib.import_module(mod_name), cls_name)
+        self[name] = cls
+        return cls
+
+    def __contains__(self, name):
+        return name in _ARM_MODULES
+
+    def __iter__(self):
+        return iter(_ARM_MODULES)
+
+    def __len__(self):
+        return len(_ARM_MODULES)
+
+
+ARMS = _LazyArms()
+
+
+def arm_text_mode(name: str) -> str:
+    """臂 → dataset 的 text_mode：动作走文本的臂保留动作从句，其余剥掉。"""
+    return "scene_action" if name in TEXT_ARMS else "scene"
+
+
+def parse_arm_kwargs(spec) -> dict:
+    """--arm_kwargs 的 JSON 串（或已是 dict / None）→ dict。"""
+    if spec is None or spec == "":
+        return {}
+    if isinstance(spec, dict):
+        return dict(spec)
+    kw = json.loads(spec)
+    if not isinstance(kw, dict):
+        raise ValueError(f"--arm_kwargs 应为 JSON 对象，得到 {spec!r}")
+    return kw
+
+
+def build_arm(name: str, dit, arm_kwargs: dict | str | None = None) -> ActionArm:
+    """按名字建臂、install(dit)、搬到 dit 的 device / dtype。arope / plain 没有 ActionArm ⇒ 返回 None。
+
+    install 由各臂自己按 dit.dim / 层数建参数并 .to(dit)；这里再统一 .to 一次是兜底：
+    某个臂的 install 若只建了参数没搬卡（例如 prompt 臂无参数），DeepSpeed 下混着 CPU 参数会炸。
+    """
+    if name in NATIVE_ARMS:
+        return None
+    cls = ARMS[name]
+    arm = cls(**parse_arm_kwargs(arm_kwargs))
+    arm.install(dit)
+    ref = dit.patch_embedding.weight
+    arm.to(device=ref.device, dtype=ref.dtype)
+    return arm
+
+
+# --------------------------------------------------------------------------
+# ckpt 键集 → 臂
+# --------------------------------------------------------------------------
+
+def split_arm_keys(sd: dict) -> tuple[dict, dict]:
+    """safetensors 的 state_dict → (DiT 部分, arm 部分（已去掉 'arm.' 前缀）)。"""
+    dit_sd, arm_sd = {}, {}
+    for k, v in sd.items():
+        if k.startswith(ARM_PREFIX):
+            arm_sd[k[len(ARM_PREFIX):]] = v
+        else:
+            dit_sd[k] = v
+    return dit_sd, arm_sd
+
+
+def detect_arm_from_keys(arm_keys: Iterable[str], shapes: dict | None = None) -> tuple[str | None, dict]:
+    """去掉前缀后的 arm 键集 → (臂名, 构造 kwargs)。没有 arm 键 ⇒ (None, {})：可能是 arope / plain / prompt，需看 metadata。
+
+    各臂的键名互不重叠（linear: action_embedders.*；xattn: action_modules.*；adaln: embedder.* / proj.*），
+    xattn 的结构配置能从形状反推：keyboard_attn_kv.weight 的 in_features = hidden_size(128)·window_frames，
+    有无 mouse_mlp.* = enable_mouse，有无 keyboard_embed.* = enable_keyboard，blocks = action_modules.<i> 的 i 集合。
+    """
+    keys = list(arm_keys)
+    if not keys:
+        return None, {}
+    heads = {k.split(".")[0] for k in keys}
+    if heads == {"action_embedders"}:
+        return "linear", {}
+    if heads == {"action_modules"}:
+        blocks = sorted({int(k.split(".")[1]) for k in keys})
+        kw: dict = {"blocks": blocks}
+        kw["enable_mouse"] = any(".mouse_mlp." in k for k in keys)
+        kw["enable_keyboard"] = any(".keyboard_embed." in k for k in keys)
+        if shapes is not None and kw["enable_keyboard"]:
+            kv_key = f"action_modules.{blocks[0]}.keyboard_attn_kv.weight"
+            emb_key = f"action_modules.{blocks[0]}.keyboard_embed.2.weight"
+            if kv_key in shapes and emb_key in shapes:
+                hidden = int(shapes[emb_key][0])
+                kw["window_frames"] = int(shapes[kv_key][1]) // hidden
+                kw["hidden_size"] = hidden
+        elif shapes is not None and kw["enable_mouse"]:
+            mlp_key = f"action_modules.{blocks[0]}.mouse_mlp.0.weight"
+            if mlp_key in shapes:
+                # in_features = 2·window_frames + dim；dim 从 proj_mouse 的 out_features 拿
+                proj_key = f"action_modules.{blocks[0]}.proj_mouse.weight"
+                dim = int(shapes[proj_key][0]) if proj_key in shapes else 3072
+                kw["window_frames"] = (int(shapes[mlp_key][1]) - dim) // 2
+        return "xattn", kw
+    if heads <= {"embedder", "proj"}:
+        kw = {}
+        if shapes is not None and "embedder.0.weight" in shapes:
+            n_in = int(shapes["embedder.0.weight"][1])
+            # n_in = n_axes · freq_dim_per_axis；默认 freq 32 ⇒ 4 轴 128 / 2 轴 64
+            kw["use_delta"] = n_in >= 4 * 32
+        return "adaln", kw
+    raise ValueError(f"无法从 arm 键集识别臂：顶层键 {sorted(heads)}")
+
+
+def detect_arm_from_ckpt(sd: dict, metadata: dict | None = None) -> tuple[str | None, dict]:
+    """(state_dict, safetensors metadata) → (臂名, 构造 kwargs)。
+
+    metadata 里的 arm 名优先（它能区分 arope / plain / prompt 这三个无 arm 键的臂）；有 arm 键时
+    以键集反推的结构配置为准，并核对与 metadata 一致。metadata 缺失且无 arm 键 ⇒ (None, {})，由调用方决定默认。
+    """
+    _, arm_sd = split_arm_keys(sd)
+    shapes = {k: tuple(v.shape) for k, v in arm_sd.items()}
+    name_keys, kw = detect_arm_from_keys(arm_sd.keys(), shapes)
+    name_meta = (metadata or {}).get("arm")
+    if name_meta and name_keys and name_meta != name_keys:
+        raise ValueError(f"ckpt metadata 说 arm={name_meta!r}，键集却是 {name_keys!r}")
+    if metadata and metadata.get("arm_kwargs"):
+        # 训练时的构造参数（JSON）比形状反推更完整（例如 heads_num）
+        try:
+            kw = {**parse_arm_kwargs(metadata["arm_kwargs"]), **kw}
+        except (ValueError, json.JSONDecodeError):
+            pass
+    return name_meta or name_keys, kw
+
+
+def make_action_inputs(offset_px: torch.Tensor, action_idx: torch.Tensor | None = None) -> dict:
+    """offset_px [b, 21, 2] (+ action_idx [b, 21]) → SPEC 的 action_inputs dict（offset_tok / delta_tok 由此派生）。
+
+    dataset.py 与 infer.py 共用，保证两边 delta_tok 的定义一致（cell 0 = 0，其余为相邻 cell 之差）。
+    """
+    from actionrope.arope import PX_PER_TOKEN
+
+    off_px = torch.as_tensor(offset_px, dtype=torch.float32)
+    off_tok = off_px / PX_PER_TOKEN
+    delta = torch.zeros_like(off_tok)
+    delta[..., 1:, :] = off_tok[..., 1:, :] - off_tok[..., :-1, :]
+    out = {"offset_px": off_px, "offset_tok": off_tok, "delta_tok": delta}
+    if action_idx is not None:
+        out["action_idx"] = torch.as_tensor(action_idx, dtype=torch.int64)
+    return out
+
+
+__all__ = [
+    "ARMS", "ARM_NAMES", "NATIVE_ARMS", "TEXT_ARMS", "ARM_PREFIX", "ActionArm",
+    "build_arm", "arm_text_mode", "parse_arm_kwargs", "split_arm_keys",
+    "detect_arm_from_keys", "detect_arm_from_ckpt", "make_action_inputs",
+]

backup_2026-09-14/code/baseline/adaln.py

@@ -0,0 +1,207 @@
+"""adaln 臂：动作嵌进 adaLN 调制（AlayaWorld 的 Action-AdaLN 移植到 Wan2.2-TI2V-5B）。
+
+上游机制（vender/AlayaWorld，基座是 LTX-2.3，不是 Wan）
+------------------------------------------------------
+alaya/model/loader.py:16,163 实际装的是 `ltx2/modules/model_ltx_2_3.py::LTX23Model`，
+`_runtime_transformer_overrides` (loader.py:206) 把 `enable_action_control=cfg.control.uses("action")` 传进去。
+（`ltx2/modules/model.py:969-997` 的 81 类离散 `DiscreteActionEmbedder` 是 LTX-2 旧路径，训练/推理都不走它。）
+
+1. 动作编码（连续 6 维、逐 latent 帧）
+   * `alaya/control/action.py:78-115 c2w_to_adaln_action_vectors`：相机 c2w 按 latent 步长（8 帧）抽样，
+     相邻两个 latent 的相对位姿 `inv(c2w[i-1]) @ c2w[i]` → 平移 xyz + 欧拉角 xyz，共 6 维；latent 0 没有前驱 ⇒ 0。
+     **这是逐 latent 的速度（增量），不是累计位移。** 每轴再除以 `control.action_scale`
+     （schema.py:210，默认 "0.14,0.075,0.22,0.17,0.70,0.16"）把量级归一到 ~[-1, 1]。
+   * `model_ltx_2_3.py:130-175 ActionAdaLNEmbedder`：每轴 `value * freq_scale(1000)` 做 32 维正弦编码
+     （schema.py:211-212 `action_freq_scale=1000, action_freq_dim_per_axis=32`；频率 10000^(-i/16)，[cos, sin]），
+     6 轴拼成 192 维 → `Linear(192, dim) → SiLU → Linear(dim, dim)` ⇒ `a` [B, T, dim]。
+2. 进 adaLN（与时间步嵌入同路，逐帧广播到该帧所有 token）
+   * `model_ltx_2_3.py:782-786`：`action_adaln_projection = Sequential(SiLU, Linear(dim, 6*dim))`，
+     与 `AdaLayerNormSingle.linear`（timestep_embedding.py:120-127：`linear(silu(embedded_timestep))`）同构。
+   * `model_ltx_2_3.py:1171-1176, 1252-1253`：`a0 = projection(a)`，`a`/`a0` 按 `tokens_per_frame` repeat_interleave
+     到逐 token，然后 `timestep_emb += a0`（[B, N, 6*dim]，所有 block 共用的 shift/scale/gate 六元组），
+     `embedded_timestep += a`（[B, N, dim]，只给输出头 `_process_output` 的 shift/scale 用，model_ltx_2_3.py:829-840）。
+     block 内 `get_ada_values` (547-555) = `scale_shift_table + timestep_emb` 切 6 段——与 Wan 的
+     `block.modulation + t_mod` 完全同构。条件（clean）帧的 token 也照加自己的动作（1153-1156 注释）。
+   * 不是每层一套：只有一套 embedder + 一套 projection，30 层共用（同时间步嵌入一样）。
+3. 初始化：`model_ltx_2_3.py:785-786` 最后一层 `Linear(dim, 6*dim)` 权重 normal(std=1e-3)、偏置 0（"近零"）；
+   embedder MLP 用 PyTorch 默认初始化。
+
+本实现的对应关系
+----------------
+Wan 的时间步路径（arope.py::arope_forward）：`t = time_embedding(sinusoid(t_tok))` [b, S, dim]，
+`t_mod = time_projection(t)` = `Linear(SiLU(t))` 展成 [b, S, 6, dim]，每个 block `modulation + t_mod` 切 6 段。
+所以上游 `timestep_emb += a0` 一一对应到钩子 `modify_t_mod`：
+
+    a  = embedder(actions)            [b, f, dim]        （同上游 ActionAdaLNEmbedder，轴数改为 4）
+    a0 = proj(a) = Linear(SiLU(a))    [b, f, 6, dim]     （同上游 action_adaln_projection）
+    t_mod[b, s, :, :] += a0[b, frame(s)]                  （帧 k 的 390 个 token 共用一份；token 排列 (f h w)）
+
+* 未移植的部分：上游还把 `a` 加进输出头的 `embedded_timestep`。Wan 的 `dit.head(x, t)` 用的 `t` 没有钩子可改
+  （SPEC 只给了 `modify_t_mod`），且那只是最后一个 LayerNorm 的 shift/scale，30 层 × 6 元组的主路径已完整，故省略；
+  若 integration 想补，只需在 arope_forward 里对 `t` 也加一条钩子。
+* 帧 0 是条件帧：数据集保证 `offset_tok[:, 0] = delta_tok[:, 0] = 0`，与上游 latent 0 的零向量一致；
+  和上游一样，零向量也照走 embedder（得到一个常量"静止"嵌入），不额外屏蔽——上游对 clean 条件 token 也加动作。
+* 帧数不等于 21 时按上游做线性插值到 f（`F.interpolate(mode='linear', align_corners=True)`）。
+
+信号
+----
+上游用的是逐 latent 速度（相对位姿），本项目的主信号是累计偏移 `offset_tok`（SPEC：与 ARoPE 同一个信号）。
+按 SPEC 的约定两者都给：动作向量 = `[offset_tok(dx, dy), delta_tok(dx, dy)]` 共 4 轴，每轴各做 32 维正弦编码
+（4 × 32 = 128 维进 MLP）。每轴的归一化尺度 `ACTION_SCALE_TOK` 相当于上游的 `action_scale`，
+取自 sidecar 实测（3000 个 valid clip）：|offset_tok| 的 p99 ≈ (9.5, 6.4)，|delta_tok| 的 p99 ≈ (0.59, 0.44)，
+取整为 (10, 6.5, 0.6, 0.45)，归一后 ~[-1, 1]，再乘 freq_scale=1000 落到与时间步同量级（正弦编码分辨得开）。
+上游把动作先转成模型 dtype（bf16）再编码，这里正弦编码用 float32/float64 算完再转 bf16 进 MLP——
+bf16 只有 3 位有效数字，累计偏移 10 token 会被量化到 0.06 token，没必要丢这个精度。
+
+参数量（dim=3072）
+-----------------
+embedder: Linear(128→3072) 396,288 + Linear(3072→3072) 9,440,256 = 9,836,544
+proj:     Linear(3072→18432) 56,641,536
+合计 66,478,080 ≈ 66.5M（上游 LTX-2.3 inner_dim=4096、6 轴时约 118M；旧报告 linear +0.18M、xattn +98.6M）。
+
+初始化
+------
+上游最后一层是 std=1e-3 的"近零"，装上会有 ~1e-3 量级的扰动。SPEC 要求装上瞬间输出与原模型逐位相同，
+所以这里把 `proj` 的最后一层 Linear 权重和偏置都置 0（`t_mod + 0` 在 bf16 下逐位不变）；
+embedder 保持 PyTorch 默认初始化，与上游一致。第一步只有 proj 的权重有非零梯度，之后 embedder 才跟着动。
+"""
+from __future__ import annotations
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from baseline.base import ActionArm
+
+# 动作向量的轴顺序：[offset_dx, offset_dy, delta_dx, delta_dy]，单位 token
+ACTION_AXES = ("offset_dx", "offset_dy", "delta_dx", "delta_dy")
+# 每轴归一化尺度（上游 action_scale 的对应物），sidecar 实测 p99 取整，见模块 docstring
+ACTION_SCALE_TOK = (10.0, 6.5, 0.6, 0.45)
+# 上游默认：schema.py:211-212
+FREQ_SCALE = 1000.0
+FREQ_DIM_PER_AXIS = 32
+
+
+def sinusoidal_per_axis(position: torch.Tensor, dim: int) -> torch.Tensor:
+    """照抄 model_ltx_2_3.py:151-160 `ActionAdaLNEmbedder._sinusoidal_embedding`：
+    [..., ] 位置 → [..., dim]，频率 10000^(-i/(dim//2))，输出顺序 [cos, sin]，float64 算完转回输入 dtype。"""
+    orig_shape = position.shape
+    pos_flat = position.reshape(-1).to(torch.float64)
+    freqs = torch.pow(
+        10000,
+        -torch.arange(dim // 2, dtype=torch.float64, device=position.device).div(max(dim // 2, 1)),
+    )
+    sinusoid = torch.outer(pos_flat, freqs)
+    emb = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=-1)
+    return emb.view(*orig_shape, dim).to(position.dtype)
+
+
+class AdaLNArm(ActionArm):
+    """AlayaWorld 式 Action-AdaLN：一套 embedder + 一套 6·dim 投影，加到逐 token 的 t_mod 上，30 层共用。"""
+
+    name = "adaln"
+
+    def __init__(
+        self,
+        freq_dim_per_axis: int = FREQ_DIM_PER_AXIS,
+        freq_scale: float = FREQ_SCALE,
+        action_scale: tuple[float, ...] = ACTION_SCALE_TOK,
+        use_delta: bool = True,
+    ):
+        super().__init__()
+        self.freq_dim_per_axis = int(freq_dim_per_axis)
+        self.freq_scale = float(freq_scale)
+        self.use_delta = bool(use_delta)
+        n_axes = 4 if self.use_delta else 2
+        assert len(action_scale) >= n_axes, f"action_scale 至少要 {n_axes} 个值"
+        # 用 buffer 存尺度：随模块 .to() 走，不进 state_dict（persistent=False），ckpt 键集只有权重
+        self.register_buffer("action_scale", torch.tensor(action_scale[:n_axes], dtype=torch.float32), persistent=False)
+        self.n_axes = n_axes
+        self.dim: int | None = None
+        # 真正的参数在 install(dit) 里按 dit.dim 建
+        self.embedder: nn.Sequential | None = None
+        self.proj: nn.Sequential | None = None
+
+    # ---- 建参数 ----
+    def install(self, dit) -> None:
+        """按 dit.dim 建 embedder / proj，放到 dit 的 device / dtype 上，proj 末层零初始化。"""
+        dim = int(dit.dim)
+        ref = dit.time_projection[-1].weight  # Wan 的 Linear(dim, 6*dim)，拿它的 device / dtype
+        self.dim = dim
+        # 同上游 ActionAdaLNEmbedder.mlp：Linear(n_axes*freq, dim) → SiLU → Linear(dim, dim)，默认初始化
+        self.embedder = nn.Sequential(
+            nn.Linear(self.n_axes * self.freq_dim_per_axis, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim),
+        )
+        # 同上游 action_adaln_projection：SiLU → Linear(dim, 6*dim)；上游 std=1e-3，这里置 0 保证装上不改输出
+        self.proj = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True))
+        nn.init.zeros_(self.proj[-1].weight)
+        nn.init.zeros_(self.proj[-1].bias)
+        self.to(device=ref.device, dtype=ref.dtype)
+
+    # ---- 动作 → 向量 ----
+    def action_vectors(self, action_inputs: dict, grid: dict) -> torch.Tensor:
+        """action_inputs → 归一化后的 [b, f, n_axes] float32（单位 token / 尺度）。
+
+        delta_tok 缺省时由 offset_tok 差分得到（定义同 SPEC：cell 0 = 0）。帧数不等于 grid['f'] 时按上游线性插值。
+        """
+        offset = torch.as_tensor(action_inputs["offset_tok"], dtype=torch.float32)
+        assert offset.dim() == 3 and offset.shape[-1] == 2, f"offset_tok 应为 [b, f, 2]，得到 {tuple(offset.shape)}"
+        parts = [offset]
+        if self.use_delta:
+            if action_inputs.get("delta_tok") is not None:
+                delta = torch.as_tensor(action_inputs["delta_tok"], dtype=torch.float32, device=offset.device)
+            else:
+                delta = torch.zeros_like(offset)
+                delta[:, 1:] = offset[:, 1:] - offset[:, :-1]
+            parts.append(delta)
+        vec = torch.cat(parts, dim=-1)  # [b, f_in, n_axes]
+        f = int(grid["f"])
+        if vec.shape[1] != f:
+            # 同 model_ltx_2_3.py:1164-1168：帧数不匹配就沿时间轴线性插值
+            vec = F.interpolate(vec.permute(0, 2, 1), size=f, mode="linear", align_corners=True).permute(0, 2, 1)
+        return vec / self.action_scale.to(vec.device)
+
+    # ---- 钩子 ----
+    def encode(self, action_inputs: dict | None, grid: dict):
+        """动作 → a0 [b, f, 6, dim]（模型 dtype），只算一次；没有动作输入 ⇒ None（钩子全部恒等）。"""
+        if action_inputs is None or self.proj is None:
+            return None
+        w = self.embedder[0].weight
+        vec = self.action_vectors(action_inputs, grid).to(w.device)  # [b, f, n_axes] float32
+        # 逐轴正弦编码（× freq_scale 落到时间步量级），拼成 [b, f, n_axes*freq_dim]
+        emb = torch.cat(
+            [sinusoidal_per_axis(vec[..., k] * self.freq_scale, self.freq_dim_per_axis) for k in range(self.n_axes)],
+            dim=-1,
+        )
+        a = self.embedder(emb.to(w.dtype))  # [b, f, dim]
+        a0 = self.proj(a)  # [b, f, 6*dim]
+        return a0.unflatten(-1, (6, self.dim))
+
+    def modify_t_mod(self, t_mod: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        """t_mod [b, S, 6, dim] += a0 按帧广播（帧 k 的 n_tok_per_frame 个 token 共用一份）。"""
+        if feats is None:
+            return t_mod
+        b, s = t_mod.shape[:2]
+        f, n = int(grid["f"]), int(grid["n_tok_per_frame"])
+        assert s == f * n, f"S={s} 与 f*n_tok_per_frame={f * n} 不符"
+        assert feats.shape[0] in (1, b) and feats.shape[1] == f, f"feats 形状 {tuple(feats.shape)} 与 grid 不符"
+        # token 排列 (f h w)：[b, f, 1, 6, dim] 广播到 [b, f, n, 6, dim] 再摊平成 S
+        add = feats.unsqueeze(2).expand(feats.shape[0], f, n, 6, self.dim).reshape(feats.shape[0], s, 6, self.dim)
+        return t_mod + add.to(dtype=t_mod.dtype, device=t_mod.device)
+
+    # ---- 工具 ----
+    def zero_init_check(self) -> bool:
+        """proj 末层全零 ⇒ a0 ≡ 0 ⇒ t_mod 不变。"""
+        if self.proj is None:
+            return False
+        lin = self.proj[-1]
+        return bool((lin.weight == 0).all() and (lin.bias == 0).all())
+
+    @staticmethod
+    def expected_n_params(dim: int, n_axes: int = 4, freq_dim_per_axis: int = FREQ_DIM_PER_AXIS) -> int:
+        """按结构算的参数量，测试用来对照 n_new_params()。"""
+        n_in = n_axes * freq_dim_per_axis
+        return (n_in * dim + dim) + (dim * dim + dim) + (dim * 6 * dim + 6 * dim)

backup_2026-09-14/code/baseline/base.py

@@ -0,0 +1,51 @@
+"""动作臂基类：baseline/SPEC.md 里钩子接口的默认实现（全部恒等）。
+
+各 baseline（linear / xattn / prompt / adaln）继承它，只覆盖自己用到的钩子。
+`arope_forward(arm=None)` 时一个钩子都不会被调用，所以基类本身不影响 ARoPE 臂。
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class ActionArm(nn.Module):
+    name = "base"
+
+    def install(self, dit) -> None:
+        """需要按 dit 的 dim / num_heads / 层数建参数时在这里做（build_arm 会调用一次）。"""
+
+    # ---- 钩子：签名固定，见 baseline/SPEC.md ----
+    def encode(self, action_inputs: dict | None, grid: dict):
+        """动作输入 → 特征（任意张量或 tuple），只算一次；None 表示这个臂不需要。"""
+        return None
+
+    def modify_t_mod(self, t_mod: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        return t_mod
+
+    def extra_context(self, context: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        return context
+
+    def after_patch(self, x: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        return x
+
+    def block_pre(self, idx: int, x: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        return x
+
+    def block_mid(self, idx: int, x: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        return x
+
+    def block_pre_ffn(self, idx: int, x: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        """文本 cross-attn 之后、FFN 之前。"""
+        return x
+
+    def block_post(self, idx: int, x: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        return x
+
+    # ---- 工具 ----
+    def n_new_params(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+    def zero_init_check(self) -> bool:
+        """"装上瞬间不改变输出"的粗检：至少有一条通向残差流的路径是零。子类按自己的结构覆盖。"""
+        return True

backup_2026-09-14/code/baseline/linear.py

@@ -0,0 +1,97 @@
+"""`linear` 臂：ReactiveGWM 的逐块线性动作偏置（30 个 bias-free Linear，加在每个 DiT block 入口）。
+
+上游机制（vender/ReactiveGWM，只读）
+-----------------------------------
+* 模型 `inference/models/dit.py::WanModelAction`（同一个 Wan2.2-TI2V-5B 基座，dim 3072 / 30 层）：
+  - L207-209 `self.action_embedders = nn.ModuleList([nn.Linear(num_buttons, dim, bias=False) for _ in range(num_layers)])`
+    —— 每个 block 一个 **无偏置** 线性层，把动作向量投到 hidden 维；没有门控（docstring："No gates"）。
+  - L222-226 `_bin_action`：原始逐帧键盘 one-hot `[B, T_raw, K]` 用 `adaptive_max_pool1d` 压到 latent 帧数 `[B, f, K]`
+    （每个 latent 帧 = 4 个视频帧的"按过即 1"）。
+  - L228-234 `_inject_action`：`emb = action_embedders[i](action[B,f,K].to(x.dtype))` → `[B, f, C]`，
+    沿 h·w 展开成逐 token 偏置 `[B, f·h·w, C]`，`x = x + bias`。
+  - L284-286 前向：`for i, block: x = _inject_action(x, action, i, f, h, w); x = block(x, ctx, t_mod, freqs)`
+    —— 偏置加在 **block 入口的残差流**（进 norm1/self-attn 之前），不是 adaLN、不是 cross-attn。
+* 动作向量 `training/data/action_utils.py` / `inference/utils/actions.py`：parquet 的 10 个按键列
+  （`inference/constants.py::SF_BUTTON_COLS`，UP/DOWN/LEFT/RIGHT + 6 攻击键）0/1 值，`hold_last_upsample`（10 帧窗）
+  补成逐视频帧 dense one-hot，不做归一化。
+* 初始化：`training/bidirectional/train.py` L17-18 / L134-137："ActionModule keys stay at their default (zero / xavier) init"，
+  其余权重从 Wan2.2-TI2V-5B 按形状拷贝（`_transfer_weights`）。训练时 action_embedders 与 DiT 全参数同 lr 训练；
+  推理有 `action_cfg_scale`（`inference/pipeline.py` L218-224，动作置零做无动作分支）。
+* 上游参数量：30 × 10 × 3072 = 921,600（10 键）。
+
+本实现的对应（baseline/SPEC.md 的钩子）
+--------------------------------------
+* `encode`：取 `action_inputs["offset_tok"]`（float32 `[b, 21, 2]`，逐 cell 累计屏幕位移 (dx, dy)，单位 token，
+  与 ARoPE 臂**完全相同的控制信号**——旧报告口径"带动作的臂收到相同的逐 cell 累计偏移"）。
+  上游的"逐帧 one-hot → adaptive_max_pool 到 f 个 latent 帧"这一步，在我们这边由 sidecar 的
+  `frames_to_cells`（cell 内逐帧偏移取均值）已经做完，dataset 直接给逐 cell 的 `[b, 21, 2]`，
+  所以这里不再池化，只做形状校验。不归一化（上游也不归一化；零初始化下量纲只影响学习动态，不影响装上瞬间的等价性）。
+  cell 0 的偏移恒为 0 ⇒ 无偏置 Linear 在条件帧上的偏置恒为 0，与"首帧是干净条件帧"一致。
+* `block_pre(i, x)`：`x + expand(action_embedders[i](offset_tok))`，与上游 `_inject_action` 同位置同算式。
+  小矩阵乘在 fp32 里算再转回 x.dtype：offset_tok 最大约 ±11 token，bf16 在该量级的分辨率是 1/16 token = 2 px，
+  fp32 保住亚像素信息；上游输入是 0/1 one-hot，bf16 本来就无损，所以它直接 `.to(x.dtype)`。
+* 参数量：30 × 2 × 3072 = **184,320 ≈ 0.18M**（与旧报告 +0.18M 一致：旧报告就是 2 维偏移输入）。
+* 初始化：**全零**（SPEC 要求"装上瞬间输出与原模型逐位相同"；上游 docstring 也把 zero 列为默认之一）。
+  零权重 ⇒ 偏置张量逐位为 0 ⇒ `x + 0 == x` 逐位相等。
+* ckpt 键：`action_embedders.{i}.weight`（与上游键布局同名；integration 阶段导出时加 `arm.` 前缀，与 DiT 键无冲突）。
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from baseline.base import ActionArm
+
+# 上游 SF_BUTTON_COLS 是 10 维 one-hot；这里换成 SPEC 的 (dx, dy) 累计偏移，2 维
+ACTION_DIM = 2
+
+
+class LinearArm(ActionArm):
+    name = "linear"
+
+    def __init__(self, dim: int = 3072, num_layers: int = 30, in_dim: int = ACTION_DIM):
+        super().__init__()
+        self.dim, self.num_layers, self.in_dim = dim, num_layers, in_dim
+        self.action_embedders = self._build(dim, num_layers, in_dim)
+
+    @staticmethod
+    def _build(dim: int, num_layers: int, in_dim: int) -> nn.ModuleList:
+        # 与上游 dit.py L207-209 同构：每层一个 bias-free Linear；权重零初始化（见模块 docstring）
+        embedders = nn.ModuleList([nn.Linear(in_dim, dim, bias=False) for _ in range(num_layers)])
+        for lin in embedders:
+            nn.init.zeros_(lin.weight)
+        return embedders
+
+    def install(self, dit) -> None:
+        """按 dit 的 dim / 层数建参数，并放到 dit 的 device / dtype（与 DiT 一起以 bf16 训练）。"""
+        dim, num_layers = int(dit.dim), len(dit.blocks)
+        if (dim, num_layers) != (self.dim, self.num_layers):
+            self.dim, self.num_layers = dim, num_layers
+            self.action_embedders = self._build(dim, num_layers, self.in_dim)
+        ref = dit.patch_embedding.weight
+        self.to(device=ref.device, dtype=ref.dtype)
+
+    # ---- 钩子 ----
+    def encode(self, action_inputs: dict | None, grid: dict):
+        """offset_tok [b, f, 2] → 同一张量（float32）。None ⇒ 不注入（对应上游 keyboard_action=None 的分支）。"""
+        if action_inputs is None:
+            return None
+        off = action_inputs["offset_tok"]
+        b, f = grid["b"], grid["f"]
+        assert off.shape == (b, f, self.in_dim), f"offset_tok 形状应为 [{b}, {f}, {self.in_dim}]，得到 {tuple(off.shape)}"
+        return off.to(torch.float32)
+
+    def block_pre(self, idx: int, x: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        """上游 `_inject_action`：逐帧线性偏置沿帧内 h·w 个 token 广播后加到残差流。"""
+        if feats is None:
+            return x
+        b, f, n = grid["b"], grid["f"], grid["n_tok_per_frame"]
+        weight = self.action_embedders[idx].weight
+        emb = F.linear(feats.to(device=x.device), weight.to(torch.float32)).to(x.dtype)  # [b, f, C]
+        bias = emb.unsqueeze(2).expand(b, f, n, self.dim).reshape(b, f * n, self.dim)
+        return x + bias
+
+    # ---- 工具 ----
+    def zero_init_check(self) -> bool:
+        return all(bool((lin.weight == 0).all()) for lin in self.action_embedders)

backup_2026-09-14/code/baseline/prompt.py

@@ -0,0 +1,267 @@
+"""Prompt 臂（Incantation 移植）：动作走文本，逐 latent 帧（逐 cell）的 prompt 控制。零新增参数、零钩子。
+
+上游机制（vender/Incantation，只读）
+--------------------------------------
+Incantation 的动作接口**纯粹是自然语言**，模型里没有任何动作专用模块：
+
+* 模型结构 `modules/causal_model.py`：`CausalWanModel` 与 Wan2.2-TI2V-5B 的 `WanModel` 参数名逐一对应
+  （:9 "Weight-compatible with WanModel"、:47-48、:238、:284-286、:386-389），子模块只有
+  patch_embedding / text_embedding / time_embedding / time_projection / blocks / head（:434-450），
+  没有 action embedding、action cross-attn、adaLN 动作支路。`modules/attention.py` 只有注意力后端、
+  RoPE、Norm，同样没有动作模块。它与教师的唯一差异是自注意力改成分块因果 + KV cache（:93-133、:135-224），
+  那是流式生成的效率设计，不是动作通路。
+* 动作 → 文本 `inference.py:76-159 PromptGenerator`：动作脚本（`"Move forward"` 或 `["Move forward", 秒]`）
+  按 `TIME_STEP = 0.25 s`（:85）展开成时间线，每个 latent 帧一条 prompt：
+  `_steps = max(1, round(duration / 0.25))`（:133-134），没给时长的用 `DURATION_MAP`（:87-131，缺省 1.0 s）；
+  两个实体的时间线取最长者、短的用最后一个动作补齐（:152-154）；模板
+  `f"person_1 performs {p}. boss_1 performs {b}."`（:157）。0.25 s = 16 fps 下 4 帧 = 恰好 1 个 latent 帧，
+  与本项目"cell k = 帧 4k−3…4k"的粒度完全一致。
+* 文本编码 `inference.py:247-260 encode_prompts`：每条**唯一**串过一次 UMT5 得 [512, 4096]（去重，:250-258），
+  再按帧展开成 list（:259）。
+* 逐帧对齐 `inference.py:318-353 generate`：第 fi 个 latent 帧用 `prompt_embeds_list[pi]` 做 context（:320-321），
+  每帧开始把 `crossattn_cache["is_init"]` 清掉（:323-324）重新算该帧文本的 k/v。也就是说**每帧 token 只
+  cross-attend 自己那帧的文本**——没有 attention mask 拼接，也没有跨帧的文本序列；逐帧对齐靠"一次只生成一帧"实现。
+* 逐帧文本掩码 `causal_model.py:340-344`（`cross_out * cross_attn_mask`）+ :673-686 / :760-772
+  （`frame_mask = (t > 0)`）：只有噪声帧（t > 0）收文本，干净帧（t = 0，含写 KV cache 的那一遍）不收，
+  注释说明是为了"文本信息不泄漏进自注意力 KV cache"（:760-764）。
+* 已发布的训练前向 `causal_model.py:589-712 _forward_train` 收的 context 是**每个样本一条**（:650-655），
+  并不逐帧；逐帧文本只在流式推理里出现。训练配方（Self-Forcing 式蒸馏、逐帧 rollout）未开源
+  （README "Release Status"）。
+
+本实现的对应关系
+----------------
+本仓库 `arope_forward` 原生支持逐 cell 文本（`actionrope/arope.py::_embed_context / _block_forward`：
+context `[b, f, L, 4096]` ⇒ q 重排成 `(b f) (h w) c`，k/v 为 `(b f) L c`，每帧 token 只 attend 自己 cell 的文本；
+唯一串只算一次 text_embedding，`context_ids` 去重），这正是 Incantation "逐 latent 帧一条 prompt、每帧只看自己
+文本" 的双向并行版本。所以这个臂**不需要任何钩子**（全部继承基类的恒等实现），只声明它要的数据契约：
+
+* `text_mode = "scene_action"`：数据集不剥动作从句，动作词按 `prompts.mirror_action` 翻回真实屏幕方向
+  （`text_table_eybx_mirror.pt`）。这就是它的动作通道。
+* `per_cell_context = True`：context 必须是 `[b, f, 512, 4096]`（dataset 默认如此）。`extra_context` 里做核对，
+  拿到整条 context 直接报错，避免"以为逐 cell 其实整条"的静默错误（`require_per_cell=False` 可放行）。
+* `uses_offset_px = False`：`arope_forward(offset_px=None)`（普通 RoPE）、无 mask 通道、loss 权重全 1。
+  `action_inputs` 里的 `offset_px / offset_tok / delta_tok / action_idx` 一个都不用——Incantation 也没有任何
+  非文本的动作输入；`action_idx` 的信息已经以动作从句的形式在文本里（infer 侧由 `velocity_to_label(mirror=False)`
+  得到同一个标签再拼进 prompt）。
+
+一致处 / 差异处（逐条）
+----------------------
+1. [一致] 动作只走文本，模型零改动、零新增参数（Incantation 的 student 与教师权重同构）。
+2. [一致] 粒度：1 条 prompt / 1 个 latent 帧（0.25 s = 4 视频帧）= 本项目的 1 个 cell。
+3. [一致] 逐帧对齐方式：每帧 token 只 cross-attend 自己那帧的文本 k/v，不是把各帧文本拼成长序列，
+   也不是 attention mask。Incantation 靠逐帧流式生成天然实现，这里靠 `(b f)` 分组的 cross-attn 并行实现；
+   所有 cell 文本相同时两者都退化成整条条件。
+4. [一致] 文本编码去重：唯一串只编码一次（`encode_prompts` ↔ `_embed_context` + `context_ids`）。
+5. [一致] 时间线展开规则：`round(秒 / 0.25)` 个 cell，缺省时长 1.0 s，脚本短于视频用最后一个动作补齐
+   （见下面 `script_to_timeline`）。
+6. [差异] prompt 模板：Incantation 是多实体槽位 `person_1 performs {p}. boss_1 performs {b}.`，不含场景
+   （场景由首帧给）；本数据集是单实体 + 场景从句 `In a 2.5D top down view, {scene}, the player is {action}.`
+   （`data/code/scenes.py::PROMPT_TMPL`）。不能换成 Incantation 模板：模型在本数据集模板上训练、推理时
+   查的是预计算文本表（不挂 T5），且转场语料要求场景从句逐 cell 变化。动作词表也不同
+   （Incantation 是 Elden Ring 的 11 个玩家动作 + 31 个 boss 动作；这里是 9 个方向标签）。
+7. [差异，未复刻] `cross_attn_mask = (t > 0)`：Incantation 的干净帧不收文本。本实现里帧 0（TI2V 条件帧，t=0）
+   照常 cross-attend cell 0 的文本。原因：(a) 它是因果 student 的 KV-cache 设计（防文本泄漏进 cache），
+   不是动作接口本身；本模型双向、无 cache，不存在泄漏问题；(b) 钩子接口里 cross-attn 输出没有单独的
+   挂点（`block_mid` 在 cross-attn 之前、`block_post` 在 ffn 之后），且 `cross_attn.o` 带 bias，无法靠构造
+   context 让输出恰为 0，精确复刻要改 `actionrope/arope.py`（加一个逐帧 cross-attn 掩码参数）；
+   (c) 复刻会让"装上瞬间输出与原模型相同"不成立，并且与其它四臂（帧 0 都 attend 场景文本）不对等。
+   若 integration 决定要它，见文末 `CROSS_ATTN_MASK_NOTE`。
+8. [差异] 自注意力：Incantation 是分块因果（帧 f 只看 ≤ f 的帧）+ KV cache 流式；本实现与其它四臂一样是
+   Wan 原生的双向全注意力，一次并行出 21 个 cell。这是生成范式的差异，五臂共享，不属于动作通路。
+9. [差异] Incantation 已发布的训练前向是整条 context；本实现训练即逐 cell（与其推理一致，比其发布代码更"逐帧"）。
+
+信号 / 参数量 / 初始化
+---------------------
+* 信号：文本（动作从句），不读 `action_inputs`。
+* 新增参数：0（Incantation 同样为 0：其 student 的 state_dict 键集与教师一致）。
+* 初始化：无新参数，装上瞬间与原模型**逐位**相同（`arope_forward(arm=PromptArm())` 只多了一次形状核对）。
+* 旧报告口径："Wan + Prompt" = `train.py --arm plain` 的行为，本臂与之完全相同，只是走 ActionArm 接口。
+
+推理时"动作脚本 → 逐 cell prompt"
+--------------------------------
+对照 `vender/Incantation/prompts/prompts.txt`（手工模式：一行一条 prompt、逐帧）与
+`examples/actions_example.json`（动作模式：`[["Move forward", 1.5], "Roll forward", ...]`），本臂建议：
+
+* 动作模式（推荐，与 `infer.py --actions "right:1.0:10,left:1.0:11"` 同一套速度/方向约定）：
+  `script_to_cell_prompts(scene, [["right", 2.5], "left", ["idle", 1.0]], trigger=None)` ⇒ 21 条
+  `PROMPT_TMPL` 串；`script_to_action_spec(script, mult)` ⇒ `--actions` 字符串，让 `offset_px`
+  （只用于实测位移对比、不进模型）与文本说的方向一致。动作词 = 真实屏幕方向（`mirror=False` 口径）。
+* 手工模式：一行一条、共 21 行（cell 0 是条件帧也要给），串必须能在 `text_table_eybx_mirror.pt` 查到，
+  否则回退加载 UMT5（`infer.py::TextEncoder`）。
+* CFG 负条件：逐 cell 裸动作串 `In a 2.5D top down view, the player is {action}.`
+  （训练 dropout 用的就是它，`infer.py::BARE_TMPL`），不是 `NULL_PROMPT`（不在 eybx 表里）。
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+import torch
+
+from baseline.base import ActionArm
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if os.path.join(ROOT, "data", "code") not in sys.path:
+    sys.path.insert(0, os.path.join(ROOT, "data", "code"))
+from scenes import ACTIONS, PROMPT_TMPL, SCENE_NL  # noqa: E402
+
+# ---- Incantation 的时间线常量（inference.py:85, :133-134, :140） ----
+TIME_STEP_S = 0.25          # 1 条 prompt 覆盖 0.25 s = 1 个 latent 帧 = 本项目 1 个 cell
+DEFAULT_DURATION_S = 1.0    # DURATION_MAP 查不到时的缺省时长（inference.py:140）
+CELLS = 21
+
+# 动作脚本里的方向名 → scenes.ACTIONS 下标（与 actionrope/infer.py::DIR_NAMES 一致；不 import infer 是因为
+# 它 import 时会带起 opencv / diffsynth 等推理依赖）
+LABEL_NAMES = ["idle", "up", "up-left", "left", "down-left", "down", "down-right", "right", "up-right"]
+DIR_NAMES = {
+    "idle": 0, "up": 1, "up-left": 2, "left": 3, "down-left": 4,
+    "down": 5, "down-right": 6, "right": 7, "up-right": 8,
+    "ul": 2, "dl": 4, "dr": 6, "ur": 8, "u": 1, "d": 5, "l": 3, "r": 7,
+}
+
+CROSS_ATTN_MASK_NOTE = (
+    "Incantation 的 (t > 0) 逐帧 cross-attn 掩码未复刻。要复刻需在 actionrope/arope.py::_block_forward 的 "
+    "cross_attn 之后乘一个 [b, S, 1] 的逐帧掩码（帧 0 为 0），并在 arope_forward 加参数透传；"
+    "这会改变帧 0 token 的表示（进而经自注意力影响全部帧），装上瞬间就不再与原模型逐位相同。"
+)
+
+
+class PromptArm(ActionArm):
+    """零钩子：动作全在逐 cell 文本里，模型本身一个参数都不加。"""
+
+    name = "prompt"
+    # ---- 数据 / 前向契约（integration 按这几个属性配 dataset 与 arope_forward）----
+    text_mode = "scene_action"     # 动作从句保留并翻回真实方向（text_table_eybx_mirror.pt）
+    per_cell_context = True        # context 必须是 [b, f, 512, 4096]
+    uses_offset_px = False         # arope_forward(offset_px=None)：普通 RoPE
+    mask_channel = False           # 无 known/new 掩码通道
+    new_weight = 1.0               # loss 权重全 1
+
+    def __init__(self, require_per_cell: bool = True):
+        super().__init__()
+        self.require_per_cell = bool(require_per_cell)
+
+    def install(self, dit) -> None:
+        """无参数可建；只记下几何信息供 extra_context 的核对用。"""
+        self._dim = int(getattr(dit, "dim", 0))
+
+    def encode(self, action_inputs, grid: dict):
+        """不消费 action_inputs（Incantation 没有非文本的动作输入）。返回 None ⇒ 钩子收到 feats=None。"""
+        return None
+
+    def extra_context(self, context: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        """恒等；只核对 context 真是逐 cell 的（[(b f), L, dim]），否则这个臂根本没有动作通道。"""
+        if self.require_per_cell and grid is not None:
+            b, f = int(grid["b"]), int(grid["f"])
+            if context.shape[0] != b * f:
+                raise ValueError(
+                    f"prompt 臂要求逐 cell 文本（context [b, f, 512, 4096] ⇒ 嵌入后 {b * f} 行），"
+                    f"得到 {context.shape[0]} 行。请给 dataset text_mode='scene_action' 的 [b, 21, 512, 4096] context；"
+                    f"确要整条 context 用 PromptArm(require_per_cell=False)。"
+                )
+        return context
+
+    # ---- 工具 ----
+    def n_new_params(self) -> int:
+        return 0
+
+    def zero_init_check(self) -> bool:
+        return True  # 没有任何通向残差流的新路径
+
+    def describe(self) -> dict:
+        return dict(name=self.name, text_mode=self.text_mode, per_cell_context=self.per_cell_context,
+                    uses_offset_px=self.uses_offset_px, mask_channel=self.mask_channel,
+                    new_weight=self.new_weight, n_new_params=0, upstream="Incantation (text-only action interface)")
+
+
+# --------------------------------------------------------------------------
+# 推理：动作脚本 → 逐 cell prompt（对照 inference.py::PromptGenerator）
+# --------------------------------------------------------------------------
+
+def _steps(duration_s: float) -> int:
+    """Incantation 的 `_steps`：max(1, round(秒 / 0.25))。"""
+    return max(1, int(round(float(duration_s) / TIME_STEP_S)))
+
+
+def script_to_timeline(script: list, n_cells: int = CELLS, default_duration_s: float = DEFAULT_DURATION_S) -> list[int]:
+    """动作脚本 → 长度 n_cells 的动作标签时间线（scenes.ACTIONS 下标）。
+
+    脚本项：`"right"`（缺省时长）或 `["right", 秒]`（None 秒 ⇒ 缺省），方向名见 DIR_NAMES，也可直接给标签下标。
+    展开规则照 Incantation：每项 round(秒/0.25) 个 cell；总长不足 n_cells 用最后一个动作补齐
+    （脚本为空 ⇒ idle，对应其 "Standing"）；超出则截断（本项目视频固定 21 个 cell，Incantation 是按脚本定长）。
+    """
+    tl: list[int] = []
+    for item in script:
+        if isinstance(item, (list, tuple)):
+            name, dur = item[0], (item[1] if len(item) > 1 and item[1] is not None else default_duration_s)
+        else:
+            name, dur = item, default_duration_s
+        if isinstance(name, str):
+            key = name.strip().lower()
+            if key not in DIR_NAMES:
+                raise ValueError(f"未知动作 {name!r}，可选 {sorted(DIR_NAMES)} 或标签下标 0..8")
+            label = DIR_NAMES[key]
+        else:
+            label = int(name)
+            if not 0 <= label < len(ACTIONS):
+                raise ValueError(f"动作标签下标越界: {label}")
+        tl += [label] * _steps(dur)
+    if not tl:
+        tl = [0]
+    tl += [tl[-1]] * (n_cells - len(tl))
+    return tl[:n_cells]
+
+
+def scene_clause(scene: str) -> str:
+    """scene key（tidal_flats）或完整从句（"in a ..."）。"""
+    if scene in SCENE_NL:
+        return SCENE_NL[scene]
+    if scene.lower().startswith("in a"):
+        return scene
+    raise ValueError(f"scene 应为 key {sorted(SCENE_NL)} 或 'in a ...' 从句，得到 {scene!r}")
+
+
+def timeline_to_cell_prompts(scene: str, timeline: list[int], trigger: str | None = None) -> list[str]:
+    """标签时间线 → 21 条数据集模板 prompt（动作词 = 真实屏幕方向，配 text_table_eybx_mirror.pt）。
+
+    trigger "burn:<B>:<start>:<n>" / "cut:<B>:<start>:<n>"：与 infer.py::compose_cell_prompts 同一约定，
+    cell [start, start+n) 用转场从句，之后切成纯 B 场景。
+    """
+    if len(timeline) != CELLS:
+        raise ValueError(f"时间线应有 {CELLS} 个 cell，得到 {len(timeline)}")
+    clauses = [scene_clause(scene)] * CELLS
+    if trigger:
+        kind, b, start, n = trigger.split(":")
+        start, n = int(start), int(n)
+        tmpl = {"burn": "{a} burning away into {b}", "cut": "{a} cutting to {b}"}
+        if kind not in tmpl:
+            raise ValueError(f"trigger 类型应为 burn/cut，得到 {kind!r}")
+        if not (0 <= start < CELLS and n >= 1):
+            raise ValueError(f"trigger 的 start/n 越界: {start}/{n}")
+        b_nl = scene_clause(b)
+        trig = tmpl[kind].format(a=clauses[0], b=b_nl[3:] if b_nl.startswith("in ") else b_nl)
+        for k in range(start, CELLS):
+            clauses[k] = trig if k < start + n else b_nl
+    return [PROMPT_TMPL.format(scene=c, action=ACTIONS[l]) for c, l in zip(clauses, timeline)]
+
+
+def script_to_cell_prompts(scene: str, script: list, trigger: str | None = None) -> list[str]:
+    """Incantation 风格动作脚本 → 21 条逐 cell prompt。例：
+    script_to_cell_prompts("tidal_flats", [["right", 2.5], "left", ["idle", 1.0]])
+    """
+    return timeline_to_cell_prompts(scene, script_to_timeline(script), trigger)
+
+
+def script_to_action_spec(script: list, mult: float = 1.0) -> str:
+    """同一脚本 → infer.py 的 `--actions` 串（"right:1.0:10,left:1.0:11"），让实测位移的参考指令与文本一致。
+
+    时间线里连续相同的标签合并成一段；idle 段速度为 0（"idle:0:n"）。
+    """
+    tl = script_to_timeline(script)
+    segs: list[tuple[int, int]] = []
+    for lab in tl:
+        if segs and segs[-1][0] == lab:
+            segs[-1] = (lab, segs[-1][1] + 1)
+        else:
+            segs.append((lab, 1))
+    return ",".join(f"{LABEL_NAMES[lab]}:{0.0 if lab == 0 else mult}:{n}" for lab, n in segs)

backup_2026-09-14/code/baseline/xattn.py

@@ -0,0 +1,288 @@
+"""xattn 臂：Matrix-Game 3 的动作模块（键盘 cross-attn + 鼠标时间轴 self-attn）移植到 Wan2.2-TI2V-5B。
+
+上游机制（vender/Matrix-Game/Matrix-Game-3，只读）
+------------------------------------------------
+* `wan/modules/action_module.py::ActionModule`（L39-333，"action module from GameFactory"）：
+  - 输入：`keyboard_condition [B, N_frames, C_key]`（逐视频帧、多热按键向量，`utils/conditions.py` L73-107
+    的 KEYBOARD_IDX = forward/back/left/right 各占一位，按下为 1；同时按两键 ⇒ 两位为 1）；
+    `mouse_condition [B, N_frames, 2]`（逐视频帧的连续相机增量，CAM_VALUE=0.1 量级）。
+    N_frames = 4·(T−1)+1，T = latent 帧数（L192-198）。
+  - 时间窗口（L207-214 / L268-281）：在帧序列前面补 `pad_t = 4·windows_size` 个帧 0 的拷贝，
+    latent 帧 i 取补齐后 `[4i, 4i+pad_t)` 这 12 帧（= 该 latent 之前 3 个 latent 跨度里的按键），
+    展平成一个 `windows·4·C` 向量 ⇒ 每个 latent 帧一条动作 token。
+  - 键盘分支（L267-332）：`keyboard_embed`（Linear(C_key,128)-SiLU-Linear(128,128)）逐帧嵌入后按窗口展平成 1536 维，
+    `keyboard_attn_kv: 1536 → 2·1024` 得到 k/v（每 latent 帧一条 token）；视频 token 经 `mouse_attn_q: dim → 1024`
+    得到 q；q/k 各做无仿射 RMSNorm（head_dim 64，eps 1e-6；`WanRMSNorm.weight` 是死参数，forward 没用，L16-28）；
+    q 重排成 `(B S) T H D`（S = h·w，同一空间位置的 T 帧作为一条序列），k/v 沿时间轴（T 条）与之做
+    **非因果** attention（`causal=False`），q/k 先加时间轴 RoPE（`rope_dim_list=[8,28,28]`，`theta=256`，
+    h/w 轴位置恒为 0，等价于只旋转每头前 8 维，L146-184 / L308-315）；`proj_keyboard: 1024 → dim`，残差相加。
+  - 鼠标分支（L201-265）：把视频 token 重排成 `(B S) T C`，与同窗口展平的鼠标量（12·2 = 24 维）在通道上拼接，
+    `mouse_mlp`（Linear(dim+24,1024)-GELU-Linear(1024,1024)-LayerNorm）后 `t_qkv: 1024 → 3·1024`，
+    在时间轴 T 上做 self-attention（同样的 RMSNorm + 时间 RoPE），`proj_mouse: 1024 → dim` 残差相加。
+    鼠标分支在键盘分支之前，键盘分支的 q 取的是加过鼠标输出之后的 hidden。
+* `wan/modules/model.py`：L533-536 只有 `block_idx ∈ action_config['blocks']` 的 block 才建 ActionModule；
+  L617-662 在 block 内 **文本 cross-attn 之后、FFN 之前** 调 `self.action_model(x, T, h, w, mouse_cond, keyboard_cond)`；
+  L1097-1103 `init_weights` 把 `proj_mouse` / `proj_keyboard` 的权重和偏置置零（其余 Linear xavier_uniform、偏置 0）。
+* 配置：Matrix-Game-3 仓库不带 config.json；`Matrix-Game-2/configs/distilled_model/universal/config.json` 给出
+  `blocks=[0..14]`（前 15 块）、`hidden_size=128, keyboard_hidden_dim=mouse_hidden_dim=1024, heads_num=16,
+  windows_size=3, vae_time_compression_ratio=4, rope_theta=256, rope_dim_list=[8,28,28], qk_norm=True, qkv_bias=False`，
+  与 ActionModule 的默认值一致；本实现沿用这组超参（`img_hidden_size` 换成本模型的 dim=3072）。
+
+本实现的对应关系
+----------------
+* 钩子位置：用 `block_pre_ffn`（文本 cross-attn 之后、FFN 之前），与上游 model.py 的注入位置一致
+  （2026-09-05 arope_forward 补了这个钩子；之前用的 block_mid 比上游早一个子层）。
+  上游在 `mouse_cond is not None` 时还有 `x = self.norm3(x)` 直接替换残差流的写法（model.py L600-601），
+  这会在零初始化时就改变预训练模型的输出，与"装上不改输出"矛盾，这里不复刻。
+* 信号：
+  - 键盘 ← `action_inputs["action_idx"]`（[b, 21] 的 9 类标签），按上游 KEYBOARD_IDX 的键位编码映射成 4 位多热
+    [up, down, left, right]（对角 = 两键同按，idle = 全 0；9 类 ↔ 9 个互异向量，与 one-hot 等价的信息量，
+    见 `KEY_TABLE`），即 `keyboard_dim_in=4`（Matrix-Game-2 配置的取值）。
+  - 鼠标 ← `action_inputs["delta_tok"]`（[b, 21, 2] 逐 cell 速度，token 单位，连续量），对应上游的相机增量。
+  - 数据是逐 cell 的，上游是逐视频帧的：按 SPEC 的 cell↔帧映射（帧 0 = cell 0，帧 4k−3…4k = cell k）把每个 cell 的
+    值复制到它的 4 帧上，再走上游一模一样的补齐 + 12 帧窗口（`_window_cell_index`）。cell 内 4 帧相同，
+    窗口里会出现重复列——这是上游结构的直接后果，保留不动（也是参数量 1536→2048 的来源）。
+    窗口长度用视频帧数 `window_frames` 表示（默认 4·windows_size = 12；`window_frames=1` 退化成每 cell 一条嵌入）。
+  - 动作特征在 `encode` 里只算一次：`[b, f, 4·windows_size, 6]`（前 4 列按键、后 2 列鼠标），作为一个张量穿过梯度检查点；
+    各 block 的 ActionModule 自己做 embed / kv 投影（上游每块一份参数，不共享）。
+* 上游 `k.repeat(S,1,1,1)` 与 `rearrange(q, "B (T S) ... -> (B S) T ...")` 的 batch 排列在 B>1 时对不上
+  （前者 s 外 b 内，后者 b 外 s 内），这里用 expand 按 (b, s) 正确对齐；B=1 时两者相同。
+* attention 用 `F.scaled_dot_product_attention`（上游是 flash_attn_func，数学等价）；RMSNorm 在 float32 里算再转回。
+
+新增参数量（dim=3072，前 15 块，每块一份）
+-----------------------------------------
+* 键盘分支每块 9,454,336：keyboard_embed 17,152 + mouse_attn_q 3072·1024 + keyboard_attn_kv 1536·2048 + proj_keyboard 1024·3072
+  ⇒ 15 块 141,815,040（≈141.8M）。
+* 鼠标分支每块 10,514,432：mouse_mlp 4,222,976（含 LayerNorm 2,048）+ t_qkv 1024·3072 + proj_mouse 1024·3072
+  ⇒ 15 块 157,716,480（≈157.7M）。
+* 默认（键盘 + 鼠标）合计 299,531,520（≈299.5M）。
+* 旧报告的 +98.6M 对应 `enable_mouse=False, window_frames=1`（每个 latent 帧只取一帧按键嵌入，kv 输入 128 → 2048，
+  每块 6,570,752，15 块 98,561,280）：即"只复刻键盘通路、不做 12 帧窗口"。上游默认的 4·windows_size=12 帧窗口把 kv
+  输入放大到 1536，仅键盘通路就是 141.8M。本实现按上游代码保留窗口与鼠标分支，二者都可用构造参数关掉。
+
+初始化
+------
+所有 Linear 权重 xavier_uniform、偏置 0（上游 `init_weights`），`proj_keyboard` / `proj_mouse` 权重置零且无偏置（qkv_bias=False），
+LayerNorm 取默认 (1, 0)。两条分支都只经 proj 进入残差流，所以装上瞬间输出与原模型逐位相同（`zero_init_check`）。
+"""
+from __future__ import annotations
+
+from typing import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from baseline.base import ActionArm
+
+# ActionModule 的时间压缩比：一个 latent 帧对应 4 个视频帧（与 SPEC 的 cell↔帧一致）
+VAE_TIME_RATIO = 4
+# 上游 rope_dim_list=[8,28,28] 里只有 t 轴那 8 维真的在旋转（h/w 位置恒 0）
+ROPE_T_DIM = 8
+ROPE_THETA = 256.0
+
+# actionrope.prompts.ACTIONS 的 9 类 → 上游 KEYBOARD_IDX 风格的多热按键 [up(forward), down(back), left, right]
+KEY_TABLE = torch.tensor([
+    [0, 0, 0, 0],  # 0 standing idle
+    [1, 0, 0, 0],  # 1 moving up
+    [1, 0, 1, 0],  # 2 moving up-left
+    [0, 0, 1, 0],  # 3 moving left
+    [0, 1, 1, 0],  # 4 moving down-left
+    [0, 1, 0, 0],  # 5 moving down
+    [0, 1, 0, 1],  # 6 moving down-right
+    [0, 0, 0, 1],  # 7 moving right
+    [1, 0, 0, 1],  # 8 moving up-right
+], dtype=torch.float32)
+KEYBOARD_DIM_IN = KEY_TABLE.shape[1]
+MOUSE_DIM_IN = 2
+
+
+def _rms_norm(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    """上游 WanRMSNorm（无仿射；它的 weight 在 forward 里没用到）。float32 里算，回到输入 dtype。"""
+    xf = x.float()
+    return (xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps)).to(x.dtype)
+
+
+def _temporal_rope(x: torch.Tensor) -> torch.Tensor:
+    """上游 get_rotary_pos_embed + apply_rotary_emb 在 rope_sizes=[T,1,1] 下的等价形式。
+
+    x: [N, T, H, D]。只有每头前 ROPE_T_DIM 维按时间位置 0..T-1 旋转（相邻两维成对，theta=256），
+    其余维位置为 0（cos=1, sin=0）原样保留。上游在 float32 里算再 type_as 回去，这里一样。
+    """
+    T, D = x.shape[1], x.shape[-1]
+    assert D >= ROPE_T_DIM
+    half = ROPE_T_DIM // 2
+    inv = 1.0 / (ROPE_THETA ** (torch.arange(0, ROPE_T_DIM, 2, device=x.device).float() / ROPE_T_DIM))  # [4]
+    ang = torch.arange(T, device=x.device).float()[:, None] * inv[None, :]  # [T, 4]
+    cos = ang.cos().view(1, T, 1, half, 1)
+    sin = ang.sin().view(1, T, 1, half, 1)
+    xr = x[..., :ROPE_T_DIM].float().reshape(x.shape[0], T, x.shape[2], half, 2)
+    x_re, x_im = xr.unbind(-1)
+    rot = torch.stack([x_re * cos[..., 0] - x_im * sin[..., 0],
+                       x_re * sin[..., 0] + x_im * cos[..., 0]], dim=-1).flatten(-2)
+    return torch.cat([rot.to(x.dtype), x[..., ROPE_T_DIM:]], dim=-1)
+
+
+def _window_cell_index(f: int, window_frames: int) -> torch.Tensor:
+    """上游"补 pad_t 帧 + 每 latent 取 [4i, 4i+pad_t)"的窗口（pad_t = 4·windows_size），折算成 [f, pad_t] 的 cell 下标。
+
+    补齐后第 p 个视频帧 = 原视频帧 max(p − pad_t, 0)，即 latent i 看的是原视频帧 [4i − pad_t, 4i)；
+    原视频帧 j 属于 cell (j+3)//4（帧 0 = cell 0，帧 4k−3…4k = cell k）。这样 feats 直接按 cell gather，
+    不必真的展开成 81 帧。pad_t 这里放宽成任意 window_frames（=1 时 latent i 只看帧 4i−1，即 cell i 自己）。
+    """
+    pad_t = window_frames
+    i = torch.arange(f).view(f, 1)
+    j = torch.arange(pad_t).view(1, pad_t)
+    frame = (VAE_TIME_RATIO * i + j - pad_t).clamp(min=0)
+    return (frame + VAE_TIME_RATIO - 1) // VAE_TIME_RATIO
+
+
+class MG3ActionModule(nn.Module):
+    """一个 block 的动作模块（上游 ActionModule 去掉 memory 分支，B>1 的 k/v 对齐修正）。"""
+
+    def __init__(self, img_hidden_size: int, *, hidden_size: int = 128, keyboard_hidden_dim: int = 1024,
+                 mouse_hidden_dim: int = 1024, heads_num: int = 16, window_frames: int = 12,
+                 enable_keyboard: bool = True, enable_mouse: bool = True, qkv_bias: bool = False):
+        super().__init__()
+        self.heads_num = heads_num
+        self.window_frames = window_frames
+        self.enable_keyboard = enable_keyboard
+        self.enable_mouse = enable_mouse
+        win = window_frames  # 上游 = vae_time_compression_ratio · windows_size = 12
+        if enable_mouse:
+            c = mouse_hidden_dim
+            self.mouse_mlp = nn.Sequential(
+                nn.Linear(MOUSE_DIM_IN * win + img_hidden_size, c, bias=True),
+                nn.GELU(approximate="tanh"),
+                nn.Linear(c, c),
+                nn.LayerNorm(c),
+            )
+            self.t_qkv = nn.Linear(c, c * 3, bias=qkv_bias)
+            self.proj_mouse = nn.Linear(c, img_hidden_size, bias=qkv_bias)
+        if enable_keyboard:
+            self.keyboard_embed = nn.Sequential(
+                nn.Linear(KEYBOARD_DIM_IN, hidden_size, bias=True), nn.SiLU(), nn.Linear(hidden_size, hidden_size, bias=True))
+            self.mouse_attn_q = nn.Linear(img_hidden_size, keyboard_hidden_dim, bias=qkv_bias)  # 名字照抄上游：其实是键盘分支的 q
+            self.keyboard_attn_kv = nn.Linear(hidden_size * win, keyboard_hidden_dim * 2, bias=qkv_bias)
+            self.proj_keyboard = nn.Linear(keyboard_hidden_dim, img_hidden_size, bias=qkv_bias)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """上游 init_weights：Linear xavier_uniform + 零偏置；proj_* 置零（进入残差流的唯一路径）。"""
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.LayerNorm):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+        for name in ("proj_mouse", "proj_keyboard"):
+            if hasattr(self, name):
+                nn.init.zeros_(getattr(self, name).weight)
+
+    def _attend(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+        """q/k/v: [N, T, H, D]；RMSNorm + 时间 RoPE + 非因果 SDPA，返回 [N, T, H·D]。"""
+        q = _temporal_rope(_rms_norm(q).to(v.dtype))
+        k = _temporal_rope(_rms_norm(k).to(v.dtype))
+        out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=False)
+        return out.transpose(1, 2).flatten(2)
+
+    def forward(self, x: torch.Tensor, feats: torch.Tensor, grid: dict) -> torch.Tensor:
+        """x: [b, (f h w), dim]；feats: [b, f, window_frames, 4+2]（按键多热 + 鼠标连续量，已按窗口展开）。"""
+        b, f, S = grid["b"], grid["f"], grid["n_tok_per_frame"]
+        assert x.shape[1] == f * S and feats.shape[:2] == (b, f)
+        H = self.heads_num
+        dtype = x.dtype
+        hidden = x
+
+        if self.enable_mouse:
+            # 上游 L201-265：同一空间位置的 T 帧作为序列，与窗口鼠标量拼接后做时间轴 self-attention
+            mouse_win = feats[..., KEYBOARD_DIM_IN:].reshape(b, f, -1).to(dtype)  # [b, f, 24]，(window d) 排列
+            h_t = rearrange(hidden, "b (f s) c -> (b s) f c", f=f)
+            m = mouse_win.unsqueeze(1).expand(b, S, f, mouse_win.shape[-1]).reshape(b * S, f, -1)
+            g = self.mouse_mlp(torch.cat([h_t, m], dim=-1))
+            q, k, v = rearrange(self.t_qkv(g), "n t (k h d) -> k n t h d", k=3, h=H)
+            attn = rearrange(self._attend(q, k, v), "(b s) f c -> b (f s) c", b=b)
+            hidden = hidden + self.proj_mouse(attn)
+
+        if self.enable_keyboard:
+            # 上游 L267-332：每 latent 帧一条键盘 token 作 k/v，视频 token 作 q，沿时间轴 cross-attention
+            key_win = feats[..., :KEYBOARD_DIM_IN].to(self.keyboard_embed[0].weight.dtype)  # [b, f, 12, 4]
+            key_tok = self.keyboard_embed(key_win).reshape(b, f, -1)  # [b, f, 12·128]
+            q = rearrange(self.mouse_attn_q(hidden), "b (f s) (h d) -> (b s) f h d", f=f, h=H)
+            k, v = rearrange(self.keyboard_attn_kv(key_tok), "b f (k h d) -> k b f h d", k=2, h=H)
+            # 上游 k.repeat(S) 在 B>1 时与 q 的 (b s) 排列错位，这里按 (b, s) 展开
+            k = k.unsqueeze(1).expand(b, S, *k.shape[1:]).reshape(b * S, *k.shape[1:])
+            v = v.unsqueeze(1).expand(b, S, *v.shape[1:]).reshape(b * S, *v.shape[1:])
+            attn = rearrange(self._attend(q, k, v), "(b s) f c -> b (f s) c", b=b)
+            hidden = hidden + self.proj_keyboard(attn)
+        return hidden
+
+
+class XAttnArm(ActionArm):
+    """Matrix-Game 3 动作模块臂：前 `blocks` 个 block 的 block_pre_ffn 钩子上各挂一个 MG3ActionModule。"""
+
+    name = "xattn"
+
+    def __init__(self, blocks: Iterable[int] = range(15), *, enable_keyboard: bool = True, enable_mouse: bool = True,
+                 windows_size: int = 3, window_frames: int | None = None, hidden_size: int = 128,
+                 keyboard_hidden_dim: int = 1024, mouse_hidden_dim: int = 1024, heads_num: int = 16):
+        """windows_size 是上游的单位（latent 帧数，窗口 = 4·windows_size 个视频帧）；window_frames 给了就直接按视频帧数
+        （旧报告口径 `enable_mouse=False, window_frames=1`）。"""
+        super().__init__()
+        self.blocks = sorted(int(i) for i in blocks)
+        if window_frames is None:
+            window_frames = VAE_TIME_RATIO * windows_size
+        assert window_frames >= 1
+        self.window_frames = window_frames
+        self.cfg = dict(enable_keyboard=enable_keyboard, enable_mouse=enable_mouse, window_frames=window_frames,
+                        hidden_size=hidden_size, keyboard_hidden_dim=keyboard_hidden_dim,
+                        mouse_hidden_dim=mouse_hidden_dim, heads_num=heads_num)
+        # ModuleDict 的键用 block 下标的字符串：state_dict 键形如 action_modules.3.proj_keyboard.weight
+        self.action_modules = nn.ModuleDict()
+
+    def install(self, dit) -> None:
+        assert len(self.action_modules) == 0, "XAttnArm 只能 install 一次"
+        assert max(self.blocks) < len(dit.blocks), f"blocks 超出 DiT 层数 {len(dit.blocks)}"
+        ref = dit.patch_embedding.weight
+        for i in self.blocks:
+            self.action_modules[str(i)] = MG3ActionModule(dit.dim, **self.cfg)
+        # 参数与 DiT 同 dtype / 同卡一起训练（ZeRO-2 下混合 dtype 会出问题）
+        self.action_modules.to(device=ref.device, dtype=ref.dtype)
+
+    # ---- 钩子 ----
+    def encode(self, action_inputs: dict | None, grid: dict):
+        """action_idx [b, f] + delta_tok [b, f, 2] → 窗口特征 float32 [b, f, window_frames, 6]。"""
+        if action_inputs is None:
+            return None
+        b, f = grid["b"], grid["f"]
+        idx = torch.as_tensor(action_inputs["action_idx"]).long()
+        assert idx.shape == (b, f), f"action_idx 应为 [b={b}, f={f}]，得到 {tuple(idx.shape)}"
+        device = idx.device
+        keys = KEY_TABLE.to(device)[idx]  # [b, f, 4]
+        if self.cfg["enable_mouse"]:
+            delta = torch.as_tensor(action_inputs["delta_tok"]).to(device=device, dtype=torch.float32)
+            assert delta.shape == (b, f, 2), f"delta_tok 应为 [b={b}, f={f}, 2]，得到 {tuple(delta.shape)}"
+        else:
+            delta = torch.zeros(b, f, 2, device=device)
+        cond = torch.cat([keys, delta], dim=-1)  # [b, f, 6]
+        win = _window_cell_index(f, self.window_frames).to(device)  # [f, window_frames]
+        return cond[:, win]  # [b, f, window_frames, 6]
+
+    def block_pre_ffn(self, idx: int, x: torch.Tensor, feats, grid: dict) -> torch.Tensor:
+        key = str(idx)
+        if feats is None or key not in self.action_modules:
+            return x
+        return self.action_modules[key](x, feats, grid)
+
+    # ---- 工具 ----
+    def zero_init_check(self) -> bool:
+        ok = len(self.action_modules) > 0
+        for mod in self.action_modules.values():
+            for name in ("proj_mouse", "proj_keyboard"):
+                if hasattr(mod, name):
+                    lin = getattr(mod, name)
+                    ok = ok and bool((lin.weight == 0).all()) and (lin.bias is None or bool((lin.bias == 0).all()))
+        return ok

backup_2026-09-14/code/configs/accelerate_zero2.yaml

@@ -0,0 +1,26 @@
+# 8 卡 DeepSpeed ZeRO-2，bf16，无 offload，梯度裁剪 1.0，梯度累积 1。
+# train_micro_batch_size_per_gpu 由 accelerate 填 "auto" ⇒ 取 train.py 里 DataLoader 的 batch_size（默认 1）。
+# 不用 deepspeed_config_file：accelerate 1.14 里它与 mixed_precision 等键互斥，yaml 原生字段更省事。
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: false
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

backup_2026-09-14/code/configs/accelerate_zero2_4gpu.yaml

@@ -0,0 +1,26 @@
+# 4 卡（自测用，配合 CUDA_VISIBLE_DEVICES=0,1,2,3） DeepSpeed ZeRO-2，bf16，无 offload，梯度裁剪 1.0，梯度累积 1。
+# train_micro_batch_size_per_gpu 由 accelerate 填 "auto" ⇒ 取 train.py 里 DataLoader 的 batch_size（默认 1）。
+# 不用 deepspeed_config_file：accelerate 1.14 里它与 mixed_precision 等键互斥，yaml 原生字段更省事。
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: false
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

backup_2026-09-14/code/configs/accelerate_zero2_ga2.yaml

@@ -0,0 +1,27 @@
+# 与 accelerate_zero2.yaml 相同，只把梯度累积改成 2（全局 batch 16）；train.py 传 --grad_accum 2
+# 8 卡 DeepSpeed ZeRO-2，bf16，无 offload，梯度裁剪 1.0，梯度累积 1。
+# train_micro_batch_size_per_gpu 由 accelerate 填 "auto" ⇒ 取 train.py 里 DataLoader 的 batch_size（默认 1）。
+# 不用 deepspeed_config_file：accelerate 1.14 里它与 mixed_precision 等键互斥，yaml 原生字段更省事。
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  gradient_accumulation_steps: 2
+  gradient_clipping: 1.0
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: false
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

backup_2026-09-14/code/demo/engine.py

@@ -0,0 +1,326 @@
+"""交互式 rollout 引擎：常驻模型 + 滑窗 KV + 流式因果 VAE，一次生成一个 cell。
+
+与 `distill/eval/rollout_eval.py:deploy_rollout_flowmap` 走**同一条**采样路径（同样的
+`frame_cond` / `fm.euler_step` / t=0 提交），区别只有三点，都是为了交互：
+
+  1. `off_ext` 是**增长的**，不预先铺满 —— 动作是玩家现给的。`frame_cond` 只看
+     [k−window, k] 这一段历史窗口，所以往后追加不影响任何已提交的 cell。
+  2. VAE 逐 cell 解码而不是最后整片解码。Wan 的 3D VAE 在时间上是因果的（decoder 带
+     `_feat_map` 缓存），把这个缓存跨调用留住就能无缝流式出帧 —— 见 StreamVAE。
+     整片 tiled 解码 161 帧要 3.6 s；流式下每 cell 只解 1 个 latent frame。
+  3. 不跑到头就停：无限续走。世界坐标是窗口锚定的（rope_cap == sink+kv_window），
+     所以走多远都不会越界 —— 这正是 ARoPE 相对编码的性质。
+
+一个 cell = 4 帧 = 16 fps 下的 250 ms。实测稳态 N=4 时 DiT 约 175 ms/cell，
+所以生成比播放快，**缓冲深度是唯一的延迟旋钮**：见 server.py 的 lead_cells。
+"""
+from __future__ import annotations
+
+import os
+import sys
+import threading
+import time
+from collections import deque
+
+import numpy as np
+import torch
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
+
+from actionrope import geometry as G          # noqa: E402
+from actionrope import infer                  # noqa: E402
+from actionrope.prompts import NULL_PROMPT    # noqa: E402
+from distill import flowmap as fm             # noqa: E402
+from safetensors import safe_open              # noqa: E402
+from safetensors.torch import load_file        # noqa: E402
+
+from distill.arope_dit import AropeDiT, unwrap  # noqa: E402
+from distill.eval.rollout_eval import _ckpt_metadata, frames_to_cells_n  # noqa: E402
+from distill.stage3_dmd.dmd_arope import frame_cond      # noqa: E402
+
+CELL_FRAMES = 4
+
+
+def safetensors_keys(path: str) -> list[str]:
+    """只读 header 拿键名 —— 不为了数个数就把 10 GB 读进内存。"""
+    with safe_open(path, "pt") as f:
+        return list(f.keys())
+
+
+class StreamVAE:
+    """持有 Wan VAE decoder 的因果时间缓存，逐 latent frame 解码。
+
+    `VideoVAE_.decode` / `VideoVAE38_.decode` 每次调用都 `clear_cache()` 再从头跑整段；我们要的
+    是跨调用保留，所以把循环体拆出来。`conv2` 是 CausalConv3d(kernel=1)、时间感受野为 1，可逐帧过。
+
+    ★ Wan2.2 的 VAE 是 `VideoVAE38_`，与 2.1 的 `VideoVAE_` 有两处不同，漏掉任何一处都直接报错
+      或出错帧（实测：漏 first_chunk ⇒ "size of tensor a (2) must match b (4)"）：
+        1. **整段的第一个** latent frame 要 `first_chunk=True` —— 它决定 avg_shortcut 出 1 帧还是 4 帧
+        2. 循环之后要 `unpatchify(out, patch_size=2)`
+      所以第 1 帧出 1 帧、之后每帧出 4 帧，正好是 (F−1)/4+1 的压缩率。
+    """
+
+    def __init__(self, vae, device):
+        self.vae = vae
+        self.m = vae.model
+        self.device = device
+        for attr in ("decoder", "conv2", "clear_cache"):
+            assert hasattr(self.m, attr), f"VAE 内核没有 {attr}，DiffSynth 版本变了"
+        self.is38 = type(self.m).__name__.endswith("38_")
+        self.first = True
+
+    def reset(self):
+        self.m.clear_cache()
+        self.first = True
+
+    @torch.no_grad()
+    def decode_one(self, z: torch.Tensor) -> np.ndarray:
+        """z [1,48,1,h,w] → uint8 [n,H,W,3]（整段首次 n=1，之后 n=4）。"""
+        scale = self.vae.scale
+        z = z.to(self.device)
+        if isinstance(scale[0], torch.Tensor):
+            s0 = scale[0].to(dtype=z.dtype, device=z.device).view(1, -1, 1, 1, 1)
+            s1 = scale[1].to(dtype=z.dtype, device=z.device).view(1, -1, 1, 1, 1)
+            z = z / s1 + s0
+        else:
+            z = z / scale[1] + scale[0]
+        x = self.m.conv2(z)
+        self.m._conv_idx = [0]
+        kw = {"first_chunk": True} if (self.is38 and self.first) else {}
+        out, self.m._feat_map, self.m._conv_idx = self.m.decoder(
+            x, feat_cache=self.m._feat_map, feat_idx=self.m._conv_idx, **kw)
+        self.first = False
+        if self.is38:
+            from diffsynth.models.wan_video_vae import unpatchify
+            out = unpatchify(out, patch_size=2)
+        out = out.clamp_(-1, 1).float()                       # [1,3,n,H,W]
+        arr = ((out[0].permute(1, 2, 3, 0) + 1) * 127.5).clamp_(0, 255)
+        return arr.to(torch.uint8).cpu().numpy()
+
+
+class TinyStreamVAE:
+    """TAEHV（madebyollin/taehv）的流式解码器 —— demo 的默认 VAE。
+
+    Wan2.2 5B 用 `taew2_2`；`_super` 是同一架构的大解码器，README 明说是给「实时世界模型、
+    TAEHV 输出即最终��出」的场景（我们正是），而这个分辨率下两者都是 3 ms/cell，所以默认取 super。
+
+    实测（480×832、单 H200、稳态）：全量 Wan VAE 流式 87 ms/cell → TAEHV 3 ms/cell，快 ~30×。
+    合上 DiT 的 167 ms，一个 cell（4 帧 = 16 fps 下 250 ms）总共 170 ms ⇒ 0.68× 实时。
+
+    两处容易搞错：
+      · 布局是 **NTCHW**，我们的 latent 是 NCTHW，要 permute
+      · TAEHV **不做** mean/std 反缩放（"decodes exactly what diffusion models use"），
+        所以喂的是 DiT 直接输出的 latent，不是 WanVideoVAE.decode 里那个 z/std+mean
+    """
+
+    def __init__(self, weights: str, device, dtype=torch.float16):
+        import sys as _s
+        tp = os.path.join(ROOT, "vendor", "taehv")
+        if tp not in _s.path:
+            _s.path.insert(0, tp)
+        from taehv import TAEHV, StreamingTAEHV
+        self.device, self.dtype = device, dtype
+        self.tae = TAEHV(weights).to(device, dtype).eval()
+        self.st = StreamingTAEHV(self.tae)
+
+    def reset(self):
+        self.st.reset()
+
+    @torch.no_grad()
+    def decode_one(self, z: torch.Tensor) -> np.ndarray:
+        x = z.to(self.device, self.dtype).permute(0, 2, 1, 3, 4)      # NCTHW -> NTCHW
+        out, frames = self.st.decode(x), []
+        while out is not None:
+            frames.append(out)
+            out = self.st.decode()
+        if not frames:
+            return np.zeros((0, 1, 1, 3), np.uint8)
+        v = torch.cat(frames, dim=1)[0]                                # [t,3,H,W]，值域 [0,1]
+        v = (v.permute(0, 2, 3, 1).clamp(0, 1) * 255).to(torch.uint8)
+        return v.cpu().numpy()
+
+
+class Engine:
+    """常驻 GPU 的交互式生成器。generate 循环跑在自己的线程里。"""
+
+    def __init__(self, ckpt: str, *, device: str = "cuda", steps: int = 4,
+                 sink_size: int = 3, kv_window: int = 17, rope_cap: int = 20,
+                 scene: str | None = None, seed: int = 0, vae: str = "taehv_super",
+                 log=print):
+        self.log = log
+        self.device = torch.device(device)
+        self.steps = steps
+        self.sink_size, self.kv_window, self.rope_cap = sink_size, kv_window, rope_cap
+        assert rope_cap == sink_size + kv_window, \
+            f"rope_cap({rope_cap}) 必须 == sink({sink_size})+window({kv_window})"
+        self.seed = seed
+
+        t0 = time.time()
+        # 832 键喂不进 load_models 的 strict 加载（它只认 827），所以：基座 → 装壳 → 灌权重。
+        # 顺序不能反：install_flowmap 会把读出层零初始化，装完再 load 才拿得到训出来的 delta。
+        # 与 rollout_eval:720-727 同一条路径。
+        meta = _ckpt_metadata(ckpt)
+        n_keys = len(safetensors_keys(ckpt))
+        if n_keys < 830:
+            raise SystemExit(f"{ckpt} 只有 {n_keys} 个键，不是 flow map 学生 —— demo 要 832 键的。")
+        self.pipe, dit, _ = infer.load_models(None, True, device, need_vae=True, arm="arope")
+        self.dit = AropeDiT(dit=dit, mask_channel=True, causal=True)
+        unwrap(self.dit).install_flowmap(form=meta.get("flowmap_form", "residual"),
+                                         gate=float(meta.get("delta_gate", 1.0)))
+        unwrap(self.dit).dit.load_state_dict(load_file(ckpt), strict=True)
+        self.dit.eval().requires_grad_(False)
+        unwrap(self.dit).move_rope_table(self.device)
+        if vae == "full":
+            self.svae = StreamVAE(self.pipe.vae, self.device)
+        else:
+            name = "taew2_2_super" if vae == "taehv_super" else "taew2_2"
+            w = os.path.join(ROOT, "vendor", "taehv", f"{name}.pth")
+            if not os.path.exists(w):
+                raise SystemExit(f"没有 {w} —— git clone https://github.com/madebyollin/taehv vendor/taehv")
+            self.svae = TinyStreamVAE(w, self.device)
+        self.vae_kind = vae
+
+        sc = torch.load(infer.SIDECAR, map_location="cpu", weights_only=False)
+        self.walk_speed = np.array(sc["meta"]["walk_speed_px_s"], np.float64)
+
+        # 文本走查表（与 rollout_eval:742 同一条）；表里没有的串会回退去加载 UMT5，
+        # 那要多占 ~11 GB，所以 demo 默认用 NULL_PROMPT —— 它一定在表里。
+        enc = infer.TextEncoder(infer.TEXT_TABLE["arope"], device)
+        prompts = [infer.strip_action(scene) if scene else NULL_PROMPT] * G.CELLS
+        ctx, ids = infer.build_context(enc, prompts, device, per_cell=True)
+        # ARoPE 臂的文本只含场景、动作已剥掉，所以 21 格必然逐位相同 —— 相同才能在无限
+        # rollout 里复用同一格（等价于 extend_context 的 "hold"）。不同就是口径出问题了。
+        assert torch.allclose(ctx[:, 0], ctx[:, -1]), "逐 cell 文本不一致，不能复用单格上下文"
+        self.ctx1, self.ids1 = ctx[:, :1].contiguous(), ids[:, :1].contiguous()
+        self.log(f"[engine] 模型 {time.time() - t0:.0f}s · {n_keys} 键 flow map · "
+                 f"N={steps} · sink={sink_size} window={kv_window} · vae={vae}")
+
+        self.sigmas = fm.anystep_sigmas(steps)
+        fm.assert_anystep_axis(torch.as_tensor(self.sigmas), terminal=True)
+
+        self._lock = threading.Lock()
+        self._vel = np.zeros(2, np.float64)      # 当前速度指令 px/s
+        self.out = deque()                        # 待推的帧
+        self._out_lock = threading.Lock()
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+        self.lead_cells = 1
+        self.stats = {"cells": 0, "gen_ms": 0.0, "dec_ms": 0.0}
+        self.cache = None
+
+    # ---- 会话 ----
+
+    def reset(self, first_frame: np.ndarray):
+        """用一张首帧开一局。会丢弃旧的 KV / VAE 缓存。"""
+        with self._lock:
+            bare = unwrap(self.dit)
+            if self.cache is not None:
+                bare.release_kv_cache(self.cache)
+            z0 = infer.encode_first_frame(self.pipe, first_frame, tiled=False)
+            self.z0 = z0
+            self.lat_h, self.lat_w = z0.shape[-2], z0.shape[-1]
+            self.cache = bare.init_kv_cache(kv_window=self.kv_window,
+                                            sink_size=self.sink_size, rope_cap=self.rope_cap)
+            self.k = 0
+            self.frame_off = np.zeros((1, 2), np.float64)     # 逐帧世界偏移，随 cell 增长
+            self._push_off()
+            ow, mk = frame_cond(self.off_ext, 0, self.sink_size, self.kv_window,
+                                self.lat_h, self.lat_w)
+            zeros = torch.zeros(1, device=self.device)
+            with torch.no_grad():
+                self._call(z0, zeros, 0, True, ow, mk)        # 真首帧 t=r=0 进缓存
+            self.svae.reset()
+            with self._out_lock:
+                self.out.clear()
+                for f in self.svae.decode_one(z0):
+                    self.out.append(f)
+            self.g = torch.Generator("cpu").manual_seed(self.seed)
+            self.stats = {"cells": 0, "gen_ms": 0.0, "dec_ms": 0.0}
+
+    def _push_off(self):
+        cells = self.k + 1
+        cell_off = frames_to_cells_n(self.frame_off, cells)
+        self.off_ext = torch.tensor(cell_off, dtype=torch.float32,
+                                    device=self.device).unsqueeze(0)
+
+    def set_velocity(self, vx: float, vy: float):
+        with self._lock:
+            self._vel = np.array([vx, vy], np.float64)
+
+    def velocity_from_keys(self, dx: float, dy: float, mult: float = 1.0) -> np.ndarray:
+        """方向单位向量 → 逐轴速度（横纵步速不同是透视造成的，与 infer.parse_actions 同口径）。"""
+        n = float(np.hypot(dx, dy))
+        if n < 1e-9:
+            return np.zeros(2)
+        return self.walk_speed * np.array([dx, dy]) / n * mult
+
+    # ---- 生成 ----
+
+    def _call(self, x, t, k, commit, ow, mk, r=None):
+        return self.dit(x, t, self.ctx1, context_ids=self.ids1,
+                        r_timestep=(t if r is None else r),
+                        first_frame_cond=False, frame_index=k, kv_cache=self.cache,
+                        commit=commit, offset_win=ow, mask_in_k=mk)
+
+    @torch.no_grad()
+    def step_once(self):
+        """按当前速度生成并解码一个 cell。"""
+        with self._lock:
+            vel = self._vel.copy()
+        t0 = time.time()
+        # 4 帧的逐帧偏移递推，与 action_offsets 同一条
+        last = self.frame_off[-1]
+        new = np.stack([last + vel * (i + 1) / G.FPS for i in range(CELL_FRAMES)], 0)
+        self.frame_off = np.concatenate([self.frame_off, new], 0)
+        self.k += 1
+        self._push_off()
+        k = self.k
+        ow, mk = frame_cond(self.off_ext, k, self.sink_size, self.kv_window,
+                            self.lat_h, self.lat_w)
+        x = torch.randn(1, self.z0.shape[1], 1, self.lat_h, self.lat_w,
+                        generator=self.g).to(dtype=self.z0.dtype, device=self.device)
+        for i in range(len(self.sigmas) - 1):
+            s_t, s_r = float(self.sigmas[i]), float(self.sigmas[i + 1])
+            tt = torch.full((1,), s_t * fm.NUM_TRAIN_TIMESTEPS, device=self.device)
+            rr = torch.full((1,), s_r * fm.NUM_TRAIN_TIMESTEPS, device=self.device)
+            x = fm.euler_step(x, self._call(x, tt, k, False, ow, mk, r=rr), s_t, s_r)
+        zeros = torch.zeros(1, device=self.device)
+        self._call(x.detach(), zeros, k, True, ow, mk)        # 每 cell 恰好提交一次，t=0
+        t1 = time.time()
+        frames = self.svae.decode_one(x)
+        t2 = time.time()
+        with self._out_lock:
+            for f in frames:
+                self.out.append(f)
+        self.stats = {"cells": k, "gen_ms": (t1 - t0) * 1e3, "dec_ms": (t2 - t1) * 1e3}
+
+    def _loop(self):
+        while not self._stop.is_set():
+            with self._out_lock:
+                buffered = len(self.out)
+            # 只保持 lead_cells 个 cell 的缓冲：多了就是白白增加操作延迟
+            if buffered >= self.lead_cells * CELL_FRAMES:
+                time.sleep(0.004)
+                continue
+            try:
+                self.step_once()
+            except Exception as e:                            # noqa: BLE001
+                self.log(f"[engine] 生成失败：{type(e).__name__}: {e}")
+                self._stop.set()
+
+    def start(self):
+        self._stop.clear()
+        self._thread = threading.Thread(target=self._loop, daemon=True, name="arope-engine")
+        self._thread.start()
+
+    def stop(self):
+        self._stop.set()
+        if self._thread is not None:
+            self._thread.join(timeout=5)
+
+    def pop_frame(self):
+        with self._out_lock:
+            return self.out.popleft() if self.out else None

backup_2026-09-14/code/demo/index.html

@@ -0,0 +1,96 @@
+<!doctype html>
+<meta charset="utf-8">
+<title>ActionRoPE — 交互 rollout</title>
+<style>
+  :root { color-scheme: dark; }
+  body { margin:0; background:#0b0d10; color:#d8dee9; font:14px/1.5 ui-monospace,SFMono-Regular,Menlo,monospace;
+         display:flex; flex-direction:column; align-items:center; gap:12px; padding:18px; }
+  video { width:min(95vw,832px); aspect-ratio:832/480; background:#000; border-radius:8px;
+          image-rendering:auto; outline:1px solid #222831; }
+  .row { display:flex; gap:14px; align-items:center; flex-wrap:wrap; justify-content:center; }
+  .k { display:inline-block; min-width:22px; padding:2px 6px; border:1px solid #3b4252; border-radius:4px;
+       text-align:center; background:#181b20; transition:background .08s; }
+  .k.on { background:#5e81ac; color:#fff; border-color:#81a1c1; }
+  select,button { background:#181b20; color:#d8dee9; border:1px solid #3b4252; border-radius:5px; padding:5px 9px;
+                  font:inherit; cursor:pointer; }
+  #stats { color:#7b8494; font-size:12px; white-space:pre; }
+  b { color:#a3be8c; font-weight:600; }
+</style>
+
+<h3 style="margin:4px 0;color:#88c0d0">ActionRoPE · any-step AR rollout</h3>
+<video id="v" autoplay playsinline muted></video>
+
+<div class="row">
+  <span>移动</span>
+  <span class="k" id="kw">W</span><span class="k" id="ka">A</span>
+  <span class="k" id="ks">S</span><span class="k" id="kd">D</span>
+  <span style="color:#616e88">（或方向键 · Shift 加速）</span>
+</div>
+
+<div class="row">
+  <select id="ff"></select>
+  <button id="rst">重开</button>
+  <span>缓冲</span>
+  <select id="lead"><option value="1" selected>1 cell（最低延迟）</option>
+    <option value="2">2 cell</option><option value="3">3 cell（最稳）</option></select>
+</div>
+<div id="stats">连接中…</div>
+
+<script>
+const V=document.getElementById('v'), S=document.getElementById('stats');
+let ws=null, keys={w:0,a:0,s:0,d:0}, shift=false, lastSent='';
+
+const KEYMAP={KeyW:'w',KeyA:'a',KeyS:'s',KeyD:'d',
+              ArrowUp:'w',ArrowLeft:'a',ArrowDown:'s',ArrowRight:'d'};
+
+function paint(){ for(const k of ['w','a','s','d'])
+  document.getElementById('k'+k).classList.toggle('on', !!keys[k]); }
+
+function send(){
+  if(!ws || ws.readyState!==1) return;
+  // 屏幕坐标：y 向下为正。按 W 是往「前/上」走，对应世界偏移 dy 为负。
+  const dx=(keys.d?1:0)-(keys.a?1:0), dy=(keys.s?1:0)-(keys.w?1:0);
+  const mult=shift?1.6:1.0, sig=`${dx},${dy},${mult}`;
+  if(sig===lastSent) return; lastSent=sig;
+  ws.send(JSON.stringify({type:'action',dx,dy,mult}));
+}
+addEventListener('keydown',e=>{ if(e.repeat) return;
+  if(e.key==='Shift'){shift=true;send();return;}
+  const k=KEYMAP[e.code]; if(!k) return; e.preventDefault(); keys[k]=1; paint(); send(); });
+addEventListener('keyup',e=>{
+  if(e.key==='Shift'){shift=false;send();return;}
+  const k=KEYMAP[e.code]; if(!k) return; e.preventDefault(); keys[k]=0; paint(); send(); });
+addEventListener('blur',()=>{ keys={w:0,a:0,s:0,d:0}; shift=false; paint(); send(); });
+
+async function start(){
+  const pc=new RTCPeerConnection({sdpSemantics:'unified-plan'});
+  pc.addTransceiver('video',{direction:'recvonly'});
+  pc.ontrack=e=>{ V.srcObject=e.streams[0]; };
+  await pc.setLocalDescription(await pc.createOffer());
+  // 等 ICE 收集完再发 offer：这是本机/隧道场景，没有 TURN，一次性发全比 trickle 简单可靠
+  await new Promise(r=>{ if(pc.iceGatheringState==='complete') return r();
+    pc.onicegatheringstatechange=()=>pc.iceGatheringState==='complete'&&r(); });
+  const a=await fetch('/offer',{method:'POST',headers:{'Content-Type':'application/json'},
+    body:JSON.stringify({sdp:pc.localDescription.sdp,type:pc.localDescription.type})}).then(r=>r.json());
+  await pc.setRemoteDescription(a);
+
+  ws=new WebSocket(`ws://${location.host}/ws`);
+  ws.onopen=()=>{ lastSent=''; send(); setInterval(()=>ws.readyState===1&&
+    ws.send(JSON.stringify({type:'stats'})), 500); };
+  ws.onmessage=ev=>{ const m=JSON.parse(ev.data);
+    if(m.type==='stats'){
+      const budget=250; // 一个 cell = 4 帧 @16fps
+      S.textContent=`cell ${m.cells}  ·  生成 ${m.gen_ms.toFixed(0)} ms  解码 ${m.dec_ms.toFixed(0)} ms`
+        +`  ·  ${((m.gen_ms+m.dec_ms)/budget).toFixed(2)}× 实时  ·  缓冲 ${m.buffered} 帧`;
+    } };
+}
+fetch('/frames').then(r=>r.json()).then(fs=>{
+  const sel=document.getElementById('ff');
+  for(const f of fs){ const o=document.createElement('option'); o.value=o.textContent=f; sel.append(o); }
+});
+document.getElementById('rst').onclick=()=>ws&&ws.send(JSON.stringify(
+  {type:'reset',first_frame:document.getElementById('ff').value}));
+document.getElementById('lead').onchange=e=>ws&&ws.send(JSON.stringify(
+  {type:'config',lead_cells:+e.target.value}));
+start().catch(e=>S.textContent='连接失败: '+e);
+</script>

backup_2026-09-14/code/demo/run.sh

@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+# ActionRoPE 交互 demo 启动器。默认只绑回环 —— 这是共用机器。
+#
+#   bash demo/run.sh                                  # 用 stage2-flowmap 的 7500
+#   CKPT=.../ema_step150.safetensors bash demo/run.sh  # 换个 ckpt
+#   STEPS=2 bash demo/run.sh                           # 更快、质量略降
+#
+# 远程玩：在你自己电脑上 ssh -N -L 8420:127.0.0.1:8420 <这台机器>，再开 localhost:8420
+set -euo pipefail
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+CKPT="${CKPT:-$ROOT/outputs/distill/stage2_flowmap/ema/ema_step7500.safetensors}"
+PORT="${PORT:-8420}"; HOST="${HOST:-127.0.0.1}"; STEPS="${STEPS:-4}"; GPU="${GPU:-}"
+[ -f "$CKPT" ] || { echo "ABORT: 没有 ckpt $CKPT"; exit 1; }
+
+# 挑一张空卡；显式给 GPU= 就用那张。demo 要约 17 GB。
+if [ -z "$GPU" ]; then
+  GPU="$(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits \
+         | awk -F', *' '$2<2000 {print $1; exit}')"
+  [ -n "$GPU" ] || { echo "ABORT: 没有空闲的卡（都 >2 GB 占用）。别抢别人的。"; exit 1; }
+fi
+echo "[demo] GPU $GPU · ckpt $(basename "$CKPT") · N=$STEPS 步 · http://$HOST:$PORT/"
+export PYTHONPATH="$ROOT" DIFFSYNTH_SKIP_DOWNLOAD=True CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES="$GPU"
+exec "$ROOT/.venv/bin/python" "$ROOT/demo/server.py" \
+  --ckpt "$CKPT" --host "$HOST" --port "$PORT" --steps "$STEPS" --vae "${VAE:-taehv_super}" "$@"

backup_2026-09-14/code/demo/server.py

@@ -0,0 +1,195 @@
+"""ActionRoPE 交互 demo：WebRTC 视频 + WebSocket 控制。
+
+  bash demo/run.sh                      # 默认绑 127.0.0.1:8420，只有本机能连
+  ssh -N -L 8420:127.0.0.1:8420 <这台机器>   # 在你自己电脑上开隧道，然后浏览器 localhost:8420
+
+★ 默认**只绑回环**。这是共用机器，绑 0.0.0.0 等于把一个能跑 5B 模型的接口暴露给整个内网，
+  所以要暴露必须显式 --host 0.0.0.0，且会打印一条警告。别图省事。
+
+延迟预算（单张 H200、N=4 步、实测）：
+    DiT 生成      ~175 ms / cell
+    流式 VAE 解码  ~25 ms / cell（整片 tiled 解码是 3.6 s/161 帧，流式后不再是瓶颈）
+    一个 cell     = 4 帧 = 16 fps 下 250 ms 的画面
+  所以生成比播放快，**延迟由缓冲深度决定，不由传输决定**：lead_cells=1 时，一次按键最坏
+  要等「当前 cell 播完 + 下一个 cell 生成」≈ 250 + 200 = 450 ms 才看得见效果。
+  WebRTC 相对 WS+JPEG 只省 20~50 ms —— 选它是因为它自带拥塞控制和帧率节拍，
+  不是因为它能解决这 450 ms。真要压延迟得动 cell 粒度（模型的时间压缩率是 4），那是模型的事。
+"""
+from __future__ import annotations
+
+import argparse
+import asyncio
+import fractions
+import json
+import os
+import sys
+import time
+
+import av
+import numpy as np
+from aiohttp import WSMsgType, web
+from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+ROOT = os.path.dirname(HERE)
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from demo.engine import CELL_FRAMES, Engine  # noqa: E402
+
+FPS = 16
+TIME_BASE = fractions.Fraction(1, FPS)
+
+
+class GameTrack(VideoStreamTrack):
+    """把引擎的帧队列按 16 fps 节拍推出去。
+
+    队列空时**重复上一帧**而不是阻塞：WebRTC 的节拍一旦被拖住，浏览器侧的 jitter buffer
+    会加深，之后每一帧都多欠一份延迟，比丢一帧糟得多。
+    """
+
+    kind = "video"
+
+    def __init__(self, engine: Engine):
+        super().__init__()
+        self.engine = engine
+        self.last = None
+        self.n = 0
+        self.repeats = 0
+
+    async def recv(self):
+        pts, _ = await self.next_timestamp()
+        f = self.engine.pop_frame()
+        if f is None:
+            f = self.last
+            self.repeats += 1
+            if f is None:                       # 还没 reset，先给一块黑
+                f = np.zeros((480, 832, 3), np.uint8)
+        self.last = f
+        self.n += 1
+        frame = av.VideoFrame.from_ndarray(f, format="rgb24")
+        frame.pts, frame.time_base = pts, TIME_BASE
+        return frame
+
+
+def build_app(engine: Engine, first_frames: list[str], log=print) -> web.Application:
+    app = web.Application()
+    app["pcs"] = set()
+
+    async def index(_):
+        return web.FileResponse(os.path.join(HERE, "index.html"))
+
+    async def frames(_):
+        return web.json_response([os.path.basename(p) for p in first_frames])
+
+    async def offer(request):
+        params = await request.json()
+        pc = RTCPeerConnection()
+        app["pcs"].add(pc)
+
+        @pc.on("connectionstatechange")
+        async def on_state():
+            log(f"[rtc] {pc.connectionState}")
+            if pc.connectionState in ("failed", "closed"):
+                await pc.close()
+                app["pcs"].discard(pc)
+
+        pc.addTrack(GameTrack(engine))
+        await pc.setRemoteDescription(RTCSessionDescription(params["sdp"], params["type"]))
+        answer = await pc.createAnswer()
+        await pc.setLocalDescription(answer)
+        return web.json_response({"sdp": pc.localDescription.sdp, "type": pc.localDescription.type})
+
+    async def ws(request):
+        sock = web.WebSocketResponse(heartbeat=10)
+        await sock.prepare(request)
+        loop = asyncio.get_running_loop()
+        async for msg in sock:
+            if msg.type is not WSMsgType.TEXT:
+                continue
+            try:
+                m = json.loads(msg.data)
+            except Exception:                                        # noqa: BLE001
+                continue
+            t = m.get("type")
+            if t == "action":
+                v = engine.velocity_from_keys(float(m.get("dx", 0)), float(m.get("dy", 0)),
+                                              float(m.get("mult", 1.0)))
+                engine.set_velocity(float(v[0]), float(v[1]))
+            elif t == "reset":
+                name = m.get("first_frame")
+                path = next((p for p in first_frames if os.path.basename(p) == name),
+                            first_frames[0])
+                # 必须走 infer.load_first_frame：它按 26:15 居中裁剪再缩到 832×480。
+                # 裸 imread 会把 bg50 的 1536×1024 原图直接喂进去 ⇒ latent 从 30×52 变 64×96，
+                # token 数 4 倍，实测生成从 ~190 ms 涨到 768 ms，而且步速尺度也错了。
+                from actionrope import infer as _inf
+                img = _inf.load_first_frame(path)
+                # reset 会重建 KV 与 VAE 缓存，是阻塞的 GPU 操作 —— 丢到线程池，
+                # 否则 aiohttp 的事件循环会卡住，连带 WebRTC 的心跳一起停。
+                engine.stop()
+                await loop.run_in_executor(None, engine.reset, img)
+                engine.start()
+                await sock.send_json({"type": "reset_ok", "first_frame": os.path.basename(path)})
+            elif t == "config":
+                if "lead_cells" in m:
+                    engine.lead_cells = max(1, min(4, int(m["lead_cells"])))
+            elif t == "stats":
+                await sock.send_json({"type": "stats", **engine.stats,
+                                      "buffered": len(engine.out),
+                                      "lead_cells": engine.lead_cells})
+        return sock
+
+    app.router.add_get("/", index)
+    app.router.add_get("/frames", frames)
+    app.router.add_post("/offer", offer)
+    app.router.add_get("/ws", ws)
+
+    async def on_shutdown(app):
+        for pc in list(app["pcs"]):
+            await pc.close()
+        engine.stop()
+    app.on_shutdown.append(on_shutdown)
+    return app
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(description="ActionRoPE 交互 demo")
+    ap.add_argument("--ckpt", required=True, help="832 键的 flow map 学生（stage2/stage3 的 ema）")
+    ap.add_argument("--host", default="127.0.0.1", help="默认只绑回环。共用机器，绑 0.0.0.0 前先想清楚")
+    ap.add_argument("--port", type=int, default=8420)
+    ap.add_argument("--steps", type=int, default=4, help="any-step 采样步数；2 更快、4 更稳")
+    ap.add_argument("--sink_size", type=int, default=3)
+    ap.add_argument("--kv_window", type=int, default=17)
+    ap.add_argument("--rope_cap", type=int, default=20)
+    ap.add_argument("--scene", default=None, help="场景描述串；不给就用 NULL_PROMPT（bg50 口径）")
+    ap.add_argument("--first_frame_dir", default=os.path.join(ROOT, "data/eval/eybx_bg_50"))
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--vae", choices=["taehv_super", "taehv", "full"], default="taehv_super",
+                    help="解码器。taehv* = TAEHV 流式（3 ms/cell），full = 全量 Wan VAE（87 ms/cell）")
+    args = ap.parse_args(argv)
+
+    import glob
+    firsts = sorted(glob.glob(os.path.join(args.first_frame_dir, "*.png")))
+    if not firsts:
+        raise SystemExit(f"{args.first_frame_dir} 里没有 png 首帧")
+
+    if args.host not in ("127.0.0.1", "localhost", "::1"):
+        print(f"!! 正在把 demo 绑到 {args.host} —— 这台机器是共用的，内网里任何人都能连上并占用这张卡。")
+        print("!! 想安全一点就用默认的 127.0.0.1 + ssh -N -L 隧道。")
+
+    engine = Engine(args.ckpt, steps=args.steps, sink_size=args.sink_size,
+                    kv_window=args.kv_window, rope_cap=args.rope_cap,
+                    scene=args.scene, seed=args.seed, vae=args.vae)
+    from actionrope import infer as _inf
+    t0 = time.time()
+    engine.reset(_inf.load_first_frame(firsts[0]))
+    engine.start()
+    print(f"[demo] 首局就绪 {time.time() - t0:.1f}s · {len(firsts)} 张首帧可选")
+    print(f"[demo] 打开 http://{args.host}:{args.port}/ "
+          f"（远程的话先 ssh -N -L {args.port}:127.0.0.1:{args.port} 这台机器）")
+    web.run_app(build_app(engine, firsts), host=args.host, port=args.port, print=None)
+
+
+if __name__ == "__main__":
+    main()

backup_2026-09-14/code/distill/AUDIT_CFPP_2026-09-10.md

@@ -0,0 +1,128 @@
+# CF++ Stage2 CD + Stage3 DMD 移植审计报告（终版）
+
+路径缩写：`$R=/opt/dlami/nvme/zhiyangdeng/ActionRoPE`，`$U=$R/hf/worldmind_ref/WorldMind/distill/code`，`$G=$R/hf/worldmind_ref/GraphGWM/train/code_distill`。未加前缀的 `distill/…` 均在 `$R` 下。
+
+## 1. 结论
+
+**CD 逻辑（代码+数学）正确。** `distill/stage2_cd/cd_arope.py` 对 `$U/stage2_cd/cd_gm.py:94-134` 是逐行忠实实现：σ/t/idx 配对、Euler 符号、cell 0 重钉、detach、cells 1: 上的 MSE、cond/clean_x 不变、v 空间 CFG、跨 rank 同步的 idx 流与按 micro 计数的快进、fp32 EMA 影子在真实优化器步之后更新一次、ema_start 硬拷贝窗口、ZeRO-2 下 zip 对齐——CPU 上 oracle 与上游转写位级一致（test_stage2_cd 9/9、test_stage23_upstream_diff 全过）。唯一实质弱点是结构性且已记录的：默认 4 节点部署轴上 σ=0.2516（部署终读点）从不是一致性 source，是固定点族的自由分量（杠杆为上游 σ=0.096 的 2.6×）；docstring 的"等于教师一步估计"只在 step 0 成立。**DMD 逻辑（代码+数学）正确**：x0 = x_t − σv、共用一份噪声与 t、(pred_fake − pred_real)/‖gen − pred_real‖ 归一化+clamp+nan_to_num、0.5·MSE 代理梯度恰为 +grad 且锚 cell 为 0、[20,980) 索引钳制、x0 空间 CFG、双向教师/critic 无 clean_x、流式重放与融合梯度精确相等（rel<1e-6）、5:1 步序与优化器常量与 `$U/stage3_dmd/dmd_gm.py:193-246, 681-706` 一致，FSDP 习语（prepare 前 fp32、use_orig_params 断言、裸 bf16 教师、每模型 clip、get_state_dict 导出、`_KVRef`）在 torch 2.11/accelerate 1.14 源码层面无误。**一个设计层面的 major**（DMD-1）：打分尾巴的 t=0 锚点是学生生成的 cell，且端口同时去掉了所有流形上参考（无真 cell 0、sink=0、VAE 重锚硬拒），其文档论证把 RoPE 坐标漂移与 latent 值漂移混为一谈，导致 DMD 目标对锚点携带的全局外观偏移无恢复力——这不是代码 bug，而是"已记录但论证不完整"的偏差。**另两个 major 是基础设施**：stage2 启动器无 8 卡分支（裸 `DISTILL_NGPU=8` 会以等效 batch 16 训练），以及部署评测脚本用 sink 3/w17 评估按 sink 0/w20 训练的学生。当前 live run（`outputs/distill/stage2_cd`，01:43:07 起，world 8/GA 1/train_batch_size 8，teacher=student=stage1_long_K3 step1000，start_step 0，~5.06 s/步）已由人工绕过了第一个 major，配方正确。
+
+## 2. 已确认发现（按复核后严重度排序）
+
+复核规则：3 位反驳者中 ≥2 未反驳即确认；表中级别为复核后共识。
+
+| ID | 级别 | 位置（端口 ↔ 上游） | 机理 / 数学与后果 | 精确修复 |
+|---|---|---|---|---|
+| **INFRA-1** (CDINFRA-1≡R1) | **major** | 端口 `distill/stage2_cd/run.sh:42-43`（`GRAD_ACCUM:-2`，`ACC_CFG` 硬编码 `accelerate_zero2_4gpu_ga${GRAD_ACCUM}.yaml`，无 `DISTILL_NGPU` 分支）、`:93`（abort 文案"只有 2/4 两份"）、`:132`（`--num_processes $DISTILL_NPROC`）；`distill/common.sh:36-44`；`cd_arope.py:513-519`（GA 门只比 yaml GA==`--grad_accum`，文案只列 4gpu）；`configs/accelerate_zero2_8gpu_ga1.yaml:4-5`（"等效 batch 仍是 8…别改 GA"）存在但从不被选；`stage2_cd/config.json:68-71`（4×1×2=8）。↔ 上游 `$U/../launch/stage2_v3_ext.sh:37-42`（8×1×1=8，"Reverted from an aborted accum=2 attempt"）、`$U/stage2_cd/cd_gm.py:274`（GA 默认 1）；本仓库正确范式 `distill/stage2_flowmap/run.sh:52-58`、`distill/stage1_long/run.sh:37-43`。 | eff = world×bs×GA。accelerate 1.14 `launch.py:1272-1277` 仅在 CLI 为 None 时采用 yaml 值 ⇒ `--num_processes 8` 覆盖 yaml 的 4，GA 仍 2 ⇒ **8×1×2=16**（文档/上游 8）；DeepSpeed train_batch_size 同为 16（accelerator.py:2223-2227）。`GRAD_ACCUM=1` ⇒ `4gpu_ga1.yaml` 不存在 ⇒ 误导性 abort。续训禁改 GA（`cd_arope.py:587-590`）⇒ step 0 即锁死。后果：裸启动=未记录配方（2× 样本/步、wall-clock 不降、EMA 每样本更新减半），只有 `run.sh:124` 一行"等效 batch=16"提示，无守卫。live run 已手工传 `GRAD_ACCUM=1 ACC_CFG=…8gpu_ga1.yaml` 绕过。 | `run.sh:42-43` 改为：`if [ "$DISTILL_NGPU" = "8" ]; then GRAD_ACCUM="${GRAD_ACCUM:-1}"; ACC_CFG="${ACC_CFG:-$ROOT/distill/configs/accelerate_zero2_8gpu_ga${GRAD_ACCUM}.yaml}"; else GRAD_ACCUM="${GRAD_ACCUM:-2}"; ACC_CFG="${ACC_CFG:-$ROOT/distill/configs/accelerate_zero2_4gpu_ga${GRAD_ACCUM}.yaml}"; fi`；`:93` 后加 `[ $((DISTILL_NPROC*GRAD_ACCUM)) = 8 ]` 断言（`EFF_BATCH_OVERRIDE=1` 放行）；`:93` 与 `cd_arope.py:515-518` 文案列出真实 4 份 yaml；`RUNBOOK.md:104` 附 8 卡形式；`config.json:69-71` 注 "4×GA2 或 8×GA1"；同样修 `distill/stage1_ar/run.sh:39-40`；修 `configs/accelerate_zero2_8gpu_ga2.yaml:4-5,11` 过期头注释（写 GA1/eff 8，实为 GA2/eff 16）。 |
+| **DMD-1** | **major** | 端口 `distill/stage3_dmd/dmd_arope.py:449-457`（`a=total-1-Keff; tail_sl=generated[:,:,a:]`��`first_frame_cond=True`）、`distill/flow_pf.py:328-331`（tail cell 0 不加噪 t=0）、`dmd_arope.py:209-211,222-233`（每 rollout 一个 exit_step；p=1/4 时锚点是 σ=1 的一步输出）、`:658-660`（`--reanchor_every` 硬拒）、`stage3_dmd/config.json:39,46`、`PORT_SPEC.md:375-382`、R10 `:556`、对抗评审 `:666-668`（[MAJOR] 未闭合）。↔ 上游 `$U/stage3_dmd/dmd_gm.py:365-367`（sink=3："frame0 is a REAL VAE-encoded latent → permanent on-manifold anchor"）、`:674`（tail=cat(real frame0, last-K)）、`:171-178`（decode→末帧→encode 重锚已实现）；`DMD_WINDOWING_SPEC.md:61-64,85`（ReactiveGWM 连续尾巴**配** `dmd_tail_reanchor: vae`）；`$G/stage3_dmd/dmd_gm.py:158-164`（重锚存在的原因："darken/wash root cause"）；`diag_pull.py:5-12`（frankenstein 变暗仅"假设 B"，无结果 ⇒ R10 前提未验证）；SFT `actionrope/train.py:234-235`（cell 0 恒为真 x0）。 | 端口打分条件 c=(x_a^θ, off−off_a)，x_a^θ~p_θ（a=N−20∈[1,7]），未加噪、不打分、detach。(i) 教师在其 cell-0 支撑外被查询（无高斯平滑）。(ii) 对整段窗口的全局外观变换 g（亮度/色偏）：教师与 critic 都以 g·x̃_a 为条件 ⇒ pred_fake−pred_real≈0 ⇒ 对 g 零梯度；上游真 frame0 钉住 g=id。sink=0+窗口锚定坐标（off_win=off[j]−off[a]）对 RoPE 精确，但与 latent 值离流形正交——`config.json:39` 的"抗漂移机制"论证是混淆。目标只惩罚窗内相对漂移**率**，对锚点携带的累积偏移无恢复力；dmd 带/critic/gnorm 全部看不到。复核者补充：训练时未约束段只有 cell 1..a≤7（真 cell 0 仍在缓存内），部署 sink=0 过 cell 20 后偏移为随机游走。 | (a) 文档：`config.json` `_reanchor_every/_scored_tail`、`PORT_SPEC` R3/R10 改为"sink=0+锚定仅解决 RoPE 坐标 OOD；尾锚是生成 cell（N−20∈[1,7]）；目标对锚点携带的全局外观偏移盲"，并显式闭合 `PORT_SPEC.md:666-668`。(b) **门**：把 `distill/eval/sweep_photodrift.sh`（41 cell 逐 cell std/mean vs stage2 基线，**按训练几何 sink 0/w20 跑**，见 GEO-1）列为 stage3 冒烟与 EMA 里程碑挑选的必做项，与 rollout_eval 位移一起。(c) 仅当出现窗恒定变暗/偏移：加 `--tail_reanchor vae`，仅对 tail cell a 做 `dmd_gm.py:171-178` 的 decode(6 latent)→末像素帧→encode，detach，不动学生缓存（修 (i)，不修 g-pin）；或按步 rank-sync 抽 a∈{0, N−K} 交替头/尾窗口（修 g-pin；代价 cell 21..27 后驱逐信号减半，一位复核者反对）。 |
+| **GEO-1** (rollout F1) | **major**（原 blocker） | 端口训练 `distill/stage3_dmd/run.sh:57-59`（KV_WINDOW 20/SINK 0/ROPE_CAP 20）、`dmd_arope.py:543-549`；`distill/eval/rollout_eval.py:588-590` 默认 0/20/20 **但** `distill/eval/sweep8.sh:50`、`photodrift.sh:16`、`paired50.sh:23` 硬编码 `--sink_size 3 --kv_window 17 --rope_cap 20`；`longhorizon/README.md:47`；本机 42 份 `outputs/distill/eval/stage1_long_K3/*.json` 全记 sink 3/w17/121 cell；`distill/tests/preflight.py:100,119-125`（锚到第一个保留 cell=cell 0 ⇒ 绝对坐标）；`rollout_eval.py:745` 仅为 flowmap 读元数据。↔ 上游 `$U/stage3_dmd/run.sh:46`（训练=部署 sink3/w17/cap20）、`DMD_WINDOWING_SPEC.md:48-50`（train/infer 同几何）；`PORT_SPEC.md:359`、`RUNBOOK.md:170-174`（两几何仅 k≤20 重合，只说"可比较"）；`test_world_rope_window.py:19-28,323-345`（sink>0 = OOD 负控制）。 | k≤20 保留集相同；k=21 差 {0,3}，k≥23 各差 3 cell。sink 3 下窗内跨度 357 px(k21)→426 px=1.14×包络(k25)→686 px=1.84×(k40) 线性增长（sink 0 恒 346 px）；当前 cell 掩码 known 比例 0.58→0.17(k40)→0(k≥48，=README 的 12 s)。时间位置两者同为 0..20，差异全在 h/w RoPE 与掩码通道——恰是 test c2/d2 证明 stage3 rollout 从不产生的输入。后果：stage3 EMA 之间及 vs stage2 的 paired50/photodrift（cell>20）度量的是 OOD 外推，ckpt 选择轴错误；训练本身与默认 rollout_eval 路径一致（故降级 major）。 | 三个脚本删掉 `--sink_size 3 --kv_window 17 --rope_cap 20`（默认即训练几何）；顺手替换死路径 `/tmp/claude-2011/…`（`sweep8.sh:17`、`photodrift.sh:12`、`paired50.sh:12`，本机无法运行）。`rollout_eval.py` 在 `_ckpt_metadata`（`:120-127`，`:745` 附近）读 do_save 写的 `kv_window/sink_size/rope_cap`（`dmd_arope.py:1090-1093`），存在且与 CLI 不等则拒绝，除非 `--allow_geometry_mismatch`（stage1/2 ckpt 无这些键，不受限）。sink3/w17 只作 RUNBOOK:174 的显式 A/B，不作唯一选择轴。stage1_long_K3 的 121-cell 基线需按 sink 0/w20 重跑。若产品几何确为 sink 3：改 `run.sh:57-58` 并重写 `PORT_SPEC.md:359`/`_geometry`——二选一。 |
+| **CD-1** | minor | 端口 `distill/stage2_cd/cd_arope.py:118-122`（docstring："学生在那里的行为等于 stage1 教师的一步估计"）、`:741-743,:758`（`idx∈[0,K-2]`）、`:138-146`（`build_cd_axis` 只向下追加���、`:299-301`/`run.sh:52`/`config.json _cd_tail_t`（默认 0，"不改用户拍板的轴"）；部署终读同节点 `distill/eval/rollout_eval.py:212-214`、`dmd_arope.py:228`。↔ 上游 `$U/stage2_cd/cd_gm.py:65-69`（N=48，自由节点 σ=0.0962）、`:435`；`$U/stage2_cd/run.sh:41`（部署 1000,250）；48 轴节点 45..47=[0.25,0.1786,0.0962]（CPU 验证）⇒ 上游终读 σ=0.2516 落在 SOURCE 节点旁。 | f(x,σ)=x−σv（c_skip=1, c_out=−σ）。平稳解族 {f(·,σ_i)=f(·,σ_3)∘Φ_2∘…∘Φ_i}，f(·,σ_3) 任意（含常数）；CM 边界条件只在 σ→0 由参数化强制，此处 c_out=−0.2516（上游 −0.096 的 2.6×）。共享参数下 pair 0-2 的梯度会顺带移动 v_θ(·,σ_3)，无恢复项。docstring 等式仅 step 0 成立（student==ema==teacher）。已验证 `CD_TAIL_T=100` ⇒ 轴 [1,0.75,0.5006,0.2516,0.1011]，4 对，部署节点不变。后果有界：`cm_err/tgt_err` 已按 idx 记录（`:249,:716,:799-800`；idx==2 行的 tgt_err 即 ‖f_EMA(·,0.2516)−x0‖²，live run 目前 ~0.04-0.05），stage3 在 exit_step=3 时直接训练该节点（`dmd_arope.py:219-228`）。 | 改 `:120-121` docstring 为"初始时等于教师一步估计，此后无约束项维持（上游 σ=0.096 节点同理，仅杠杆小 2.6×）"；RUNBOOK/PORT_SPEC §Stage2 注明哪个节点无约束。训练中盯 `train_log.csv` idx==2 行的 tgt_err 相对 step≤200 基线；仅当 cd 降、tgt_err 升时，以**新 OUT** 跑 `CD_TAIL_T=100` 消融（续训不能改轴；轴系用户拍板，不改默认；改则同步 `cd_arope.py:299-301`/`run.sh:52`/`config.json`/`test_stage2_cd.py::test_1`）。 |
+| **INFRA-2** (CDINFRA-2) | minor | 端口 `cd_arope.py:575-597`（续训只查 step/grad_accum/axis_sigmas；`:594` 重指 student/ema，`args.teacher_ckpt` 不动，`:617` 原样加载）、`:738-740`（progress.json 无 ckpt 路径；`runlib.py:224-229`）、`:449-452`（meta gate 只比 tf/windowed，不比身份）、`run.sh:15`（`AR=… # 换一个 stage1 检查点`）、`:38/:76`（固定 OUT + RESUME=auto）、`:123`（横幅"teacher = student init = $AR"续训时为假）；`_smoke_stage2_cd/progress.json` 实测键集无 ckpt 字段。↔ 上游 `$U/stage2_cd/cd_gm.py:326-337,472-473`（同洞，但 `stage2_v3_ext.sh:17-29,41-42` 硬编码 RD/TEACHER 并断言状态）；`PORT_SPEC.md:249`（P1 同源要求）。 | 目标 cm_tn=lat_tn−σ_tn·EMA(lat_tn)，lat_tn=lat_t+(σ_tn−σ_t)·T_B(lat_t)。student/EMA 若来自对 T_A 蒸馏的 run，链条锚在 f_A(·,0.25)∘Φ_{T_B}，loss 有限且下降，无报错。后果：一旦 progress.json 出现（step 500），>10 min 后换 `AR=` 进默认 OUT 的重启静默续训旧权重。live run 已从 step 0 起。 | do_save `:738-740` 的 `write_progress` 加 `student_ckpt=args.student_ckpt, teacher_ckpt=args.teacher_ckpt`（CLI 值，`student_init` 续训后是 dits 路径）；`:593` 后加 `for k in ("student_ckpt","teacher_ckpt"): old,new=prog.get(k),getattr(args,k); if old and new and realpath(old)!=realpath(new): raise SystemExit(...)`（dry_run 下两者为 None，需 None-safe）；旧 progress.json 回退读 dits header 同名键（`:722-725` 已写）；stage3 镜像 `dmd_arope.py:1112` 与 `:891-914`；修 `run.sh:123` 横幅；`RUNBOOK.md:180` 记"换 AR= ⇒ 换 OUT 或 RESUME=none"。**live run 结束后再改。** |
+| **FSDP-CAST** (S3-FSDP-1≡dit-wrapper F1) | minor | 端口 `dmd_arope.py:193-195`（`student(..., offset_win=ow)` fp32）、`:456-457`（`cond_t` `offset_px` fp32 → critic `:370`）、`:968/:991-996`（教师裸 bf16 不 prepare）、`distill/arope_dit.py:755,:892`（forward **内**才 `.to(float64)/PX_PER_TOKEN`）、`_KVRef` docstring `:217-231`（只提容器重建）。↔ 库/上游：torch 2.11 `fsdp/_runtime_utils.py:610-627`+`distributed/utils.py:59-63`（root cast 所有浮点 kwargs→bf16）、accelerate 1.14 `utils/dataclasses.py:2118-2124`（未设 `cast_root_forward_inputs`，默认 True `fsdp/api.py:225`）；`$U/stage3_dmd/dmd_gm.py:569-576`（教师也 FSDP 包 ⇒ 对称）；部署 `actionrope/infer.py:606`、`rollout_eval.py:477-478` 传精确 fp32；stage2 DeepSpeed 无输入 cast。 | bf16 8 位尾数：|o|<256 px 误差≤0.5 px，[256,512) ≤1 px；锚定窗跨度≤373 px ⇒ ≤1 px=1/32 token=≤0.031 rad（顶频 1 rad/token，`arope.py:68-81`），两 cell 相对≤2 px。掩码由 fp32 在调用前算好，不受影响（`dmd_arope.py:151,457`）。**时间步半边不成立**：`_per_token_timestep_pf` 对所有调用方按 latent dtype 取整（`arope_dit.py:177`）；复核者对教师侧 latent 是否 fp32 有分歧，属二阶。后果：学生/critic 与教师、与 SFT/部署间 ≤1/32 token 坐标不一致；量级远低于 DMD 信号，但未记录；step-0 "dmd 恰为 0"（`:473-474,:921,:1173`）在 FSDP 下只 ≈0；`tests/test_world_rope_window.py:176-181` 的"f32 ⇒ 训练中精确"在 FSDP 下不成立。 | 首选（局部）：偏移以 int64 定点传入——`frame_cond :150-152`、`:456-457`、probe `:487` 用 `(off*1024).round().to(torch.int64)`，`arope_dit.py:755/:892` 解码 `.to(float64)/1024/PX_PER_TOKEN`；两入口断言 `offset.dtype not in (bf16, fp16)`。备选：prepare 前 `plugin.mixed_precision_policy = dataclasses.replace(plugin.mixed_precision_policy, cast_root_forward_inputs=False)`（须重验 latent/timestep dtype 路径；一位复核者反对）。无论如何：PORT_SPEC 风险表新增一条；step-0 注释改"≈0"；`test_stage3_dmd.py:257` 的 `==0.0` 不要变成 FSDP 运行时断言。 |
+| **S3-GUARD-3** | minor | 端口 `dmd_arope.py:1237-1238`（`gng=_clip(...)` 直接 `opt_g.step()`，gng 只进日志 `:1261,:1271,:1280`）、`:1161-1172`（唯一 abort=预跑 probe，一次 backward，不覆盖 `stream_grad_replay :273-340` 的 20 次交错 backward）、`:1259-1263`（LOG_EVERY=10 ⇒ step 1 的 gnorm_g 不打印）；承诺处 `PORT_SPEC.md:540-541`(R5)、`:586`(S7a)、`configs/accelerate_fsdp_dmd_8gpu.yaml:13`、`_4gpu.yaml:11`（"从第 1 步就打印 grad_norm，恰为 0 立刻中止"）。↔ 上游 `$G/stage3_dmd/dmd_gm.py:1193-1208`（gnorm_g 是 FSDP 下唯一可信信号，仅打印）、`$U/stage3_dmd/dmd_gm.py:641-648`；`_scheme2_docs/docs/交接_蒸馏_STAGE2_NEXT_SESSION.md:283`（真实发生过：gnorm_g=0.00 每步、critic 正常、生成器没学）。 | step 0：critic==teacher（cfg=1）⇒ dmd 恰 0 合法（`test_stage3_dmd.py:256`）。step≥resume_step+1：critic 已走 critic_steps=5 步 ⇒ pred_fake≠pred_real ⇒ 代理 d/dg=grad/n≠0 ⇒ gng>0 是硬不变量；恰为 0.0 只能是 autograd/FSDP 路径断裂。后果：无学习 run 烧完 500 步而 critic loss 动、dmd 打印"健康"。 | `:1237` 与 `:1238` 之间加 `if step > resume_step and (gng == 0.0 or math.isnan(gng)): raise SystemExit(f"[gnorm] generator grad norm {gng} at step {step}: student receives no gradient (check FSDP backward setup / fsdp_use_orig_params / autocast cache / stray detach)")`（FSDP `clip_grad_norm_` 是集合通信，各 rank 同值，不会挂；inf 不 abort，clip 已归零）；同步 `config.json:109 _gate`、PORT_SPEC R5、两份 yaml 头注释；可选 CPU 测试 monkeypatch `_clip` 返回 0 断言 SystemExit。 |
+| **GEO-2** (rollout F2) | minor | 端口 `arope_dit.py:875-877`（mask_emb 加进残差流）、`:338-339,:353-357`（缓存 pre-RoPE K/V，掩码内容已在 x 中，读时只重做旋转）、`dmd_arope.py:150-151`（cell j 的掩码锚在 `max(0,j−20)`）、`:449-457`（打分掩码锚在 N−20）；过度表述 `PORT_SPEC.md:366`（"bit-for-bit"）`,368`、`dmd_arope.py:145-146`、`arope_dit.py:325-326`、`data.py:127-131`；`test_world_rope_window.py` d1 仅 a=0（`:352-359`）、a3 故意固定掩码（`:198`）；`test_stage3_dmd.py:186-209` 只查 shape。↔ 上游 `$U/gamemaster/models/gamemaster_dit.py:212-241` 无空间掩码通道，无对应物。 | 等价训练 clip 需 M(off[j]−off[a])，缓存持 M(off[j]−off[a_j])，a_j<a。CPU 测量（69/46 px/s，30×52）：窗均不一致 ~3%(k21)→~25%(k30)→~31-33%(k≥40)，最老 cell ~69%（等价 clip 全 known，缓存 31% known）；当前 cell 恒 0%，k=20→21 连续无翻转。打分：N∈[21,27] 时 20/21 个 tail cell 的 c_score≠c_gen（每 cell 最大 6-30%）。后果：学生侧 train==deploy 成立（自强迫穿过它）；是 stage2→stage3 的输入分布移位，k>27 稳态模式未训练；打分用尾锚掩码是教师 SFT 分布内的**正确**目标条件（3 位复核者中 2 位），与上游 `dmd_gm.py:674-677` 用打分器原生表示一致——不是 bug。 | 只改文档（上述 5 处："RoPE 重锚定精确（a1/a3）、查询 cell 掩码精确（d1）；缓存 cell 的掩码通道按提交时锚点固化，k>20 的历史与锚在 k−20 的 clip 不同——stage2→stage3 输入移位，自强迫训练穿过"）；加 test d3 记录数值（扩到 k=61）并断言查询 cell 0%；**不要**给教师/critic 喂混合锚掩码（会把双向教师推出分布）；可选 `--mask_zero` 消融（须同时作用 rollout 与打分，否则制造新的 c_gen≠c_score）；注明 k>27 稳态未训练或按显存抬 rollout_max。 |
+| **GEO-3** (rollout F3) | minor | 端口 `arope_dit.py:853-857`（`n_cached==min(frame_index,cap)`）、`:309-316`（`_evict` 使 n_cached==cap 成为提交的不动点）、`:833-835`（docstring 声称能拦重复 commit）、`:801,:923`（`frames_cached` 只写不读；库内无读者）、`tests/test_kv_equiv.py:276-286`（test_5 仅饱和前）。↔ 上游 `gamemaster_dit.py:631-721`（无任何 n_cached 检查，`frames_cached :721` 同样无读者）；`$G/graphgwm/models/graph_dit.py:855-857`（记录同一"饱和后形状对内容错"）。 | k≥cap 时谓词恒为 cap==cap。CPU 复现（kv_window=3）：cell 4 重复提交与跳过 cell 5 均不报错，重复提交改变后续 K/V（max|Δv|=2.5e-2）。现行调用方均每 k 恰提交一次（`dmd_arope.py:194/246,:306/318,:489/495`；`rollout_eval.py:287/300,:334`），`test_stage3_dmd.py:99-135` 的 spy 会拦。后果：无现行错算；未来编辑（重开 re-anchor、提前 `continue`��在 stage3 专门训练的驱逐区静默损坏窗口。 | forward_frame 入口加 `assert kv_cache["frames_cached"] == frame_index, f"frame_index={frame_index} 但缓存最后提交的是 cell {kv_cache['frames_cached']-1} —— 漏 commit 或重复 commit 了"`（消息须含 "frame_index" 与 "commit" 字样，保 test_4 `match="frame_index"`、test_5 `match="commit"`）；改 `:833-835` docstring；test_5 扩饱和后重复/跳过用例（kv_window=2，commit 0..3 后重提 3；跳过 4 读 5）。 |
+| **GEO-5** (rollout F5) | minor | 端口 `tests/test_world_rope_window.py:45`（`ROOT="/data/zhiyangdeng/ActionRoPE"`）、`:67-68,:91-93`（sidecar 缺失即 skip）；同硬编码 `test_kv_equiv.py:42`（仅 sys.path，实测无影响）、`test_bitparity.py:28`、`test_flowmap_model.py:35`（`:548` STAGE1B_CKPT 亦错）、`test_iface_fidelity.py:34`；正确写法 `tests/preflight.py:45`、`test_tf_port.py:32`；`RUNBOOK.md:30`（`cd /data/…`）、`:34`（容忍计数漂移）。↔ 无上游对应；`PORT_SPEC.md:45,359,535,582`（把包络门定为 sink=0 前置门）。 | 本机跑该文件 = 3 passed / 7 skipped（b1/b2/b3/c1/c2/d1/d2）；sidecar 实际在 `$R/data/latent/arope_sidecar.pt`；改路径后 10 passed（包络 372.97 px=11.655 tok=23.311 latent；b2 1.0056×/0.2%；c2 sink=1→1.50×）。后果：sink=0 包络内证据门在本机静默不跑（跑起来会过）；无 run.sh 调 pytest，不阻塞启动。 | 5 个文件改为 `ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))`；corpus fixture 加 `AROPE_REQUIRE_SIDECAR=1` 时 `pytest.fail`（RUNBOOK 启动门用该变量）；修 `RUNBOOK.md:30/34`。 |
+| **R3** | minor | 端口 `tests/test_stage2_flowmap.py:445-466`（唯一三处一致性测试；`:453` parse_args 无 `--dry_run` ⇒ `stage2_flowmap/fm_arope.py:443-444` 因 stage1b ckpt 缺失 SystemExit，本机红）；`test_stage2_cd.py`/`test_stage3_dmd.py` 无此测试；`cd_arope.py:354-362`、`dmd_arope.py:676-687` 同样需 `--dry_run`。↔ `RUNBOOK.md:15-17`、`PORT_SPEC.md:559`(R11) 只陈述规则。 | CPU 复刻：stage2_cd 23 常量 0 不一致；stage3_dmd 29 常量仅 `clamp` 假阳（`'20,980'` vs `dmd_arope.py:669-672` 派生 tuple）；flowmap 加 `--dry_run` 后 36/0。今日无漂移；未来单处改动无自动拦截。 | `test_stage2_flowmap.py:453` 加 `"--dry_run"`；克隆到 cd/dmd 测试（`parse_args(["--output",p,"--dry_run"])`，skip 派生键 betas/denoise/clamp 与 ckpt 路径，`checked>=20`），加 run.sh 别名映射（cd: DISCRETE_N→discrete_cd_N、SOLVER_STEPS→teacher_solver_steps、CFG_SCALE→real_guidance_scale、MIN_RT、WINDOW(_SINK)、KEEP_*；dmd: SINK→sink_size、SCORE_FRAMES、OFFSET_CHUNK、CLAMP、KEEP_*）并断言硬编码字面量（`--opt8bit 0 --batch_size 1 --context_extend hold --reanchor_every 0 --frankenstein 0 --scene_dropout 0.0 --use_ema 1`）。 |
+| **R4** | minor | 端口 `distill/stage3_dmd/run.sh:52`（硬编码 `accelerate_fsdp_dmd_4gpu.yaml`）、`:158`、`:147`（echo "4gpu.yaml 8 卡"自相矛盾）、`stage3_dmd/config.json:83-85,104-106`、`PORT_SPEC.md:469`（world 4/eff 4）、`dmd_arope.py:746`（`GIB_STATIC_4GPU` 在 8 卡下保守）；`configs/accelerate_fsdp_dmd_8gpu.yaml` 与 4gpu 仅 `num_processes` 不同，`stage3_onpolicy/run.sh:42` 在用。↔ 上游 `$U/../launch/stage3_cd15k.sh:10-11`（NPROC=8，部署线）、`$U/stage3_dmd/run.sh:29`（默认 3）。 | stage3 无 GA 旋钮 ⇒ eff=world=8（上游部署线亦 8，非害）；运行快照 `env.world_size` 正确（`dmd_arope.py:1035`）。仅文档漂移。 | `run.sh:52` 加 DISTILL_NGPU 分支（模板 `stage3_onpolicy/run.sh:41-45`，**不是** stage2_cd）；`:147` 打印 `等效 batch=$DISTILL_NPROC`；`config.json` 加 `_world_size` 注、`PORT_SPEC.md:469` 记 "processes=DISTILL_NGPU，eff=processes"。 |
+| **R2** | info | 端口 `stage2_cd/run.sh:32-33,:97`（`AR` 单变量喂 teacher+student）、`:80-88`（默认 `stage1_ar_tf/dits/dit_step12000` 本机不存在 ⇒ abort；`:85` 的 `AR=` 提示在 `-d $STAGE1_OUT/dits` 分支内，本机不打印）；TEACHER/STUDENT 属 stage3（`stage3_dmd/run.sh:41,43`），INIT 属 stage1_long（`:28`）。↔ 上游 `$U/stage2_cd/run.sh:24,34`（`AR_CKPT` 同一约定）。 | 失败响亮（abort），非静默；override 名已在 `run.sh:15`、`PORT_SPEC.md:301` 记录。 | 用 `AR=`；把 `:85` 提示移出分支；`RUNBOOK.md:104` 写明 `AR=` 与长教师。**不要**加 TEACHER 别名：stage3 的双向 SFT ckpt（`outputs/arope_10k_ga2/step-10000`）无 tf/windowed 元数据键，`cd_arope.py:452` 只比存在的键 ⇒ 陈旧 `export TEACHER` 会静默通过 P1 门。 |
+| **R5** | info | 端口 `stage2_cd/run.sh:76,:116-118`（RESUME=auto）、`cd_arope.py:575-597,:646-652`（不看 `finished` 标志；`:827` 写 `finished=true` 全库无读者）；过期前缀 `/data/zhiyangdeng/ActionRoPE`：`stage2_cd/run.sh:13`、`stage3_dmd/run.sh:16`、`stage1_ar/run.sh:12`、`RUNBOOK.md:30/95/128`、`stage1b_adapt/run.sh:34`、`stage3_onpolicy/run.sh:13` 及多���测试 docstring。↔ 上游 `stage2_v3_ext.sh:42`（依赖续 finished run 延长）。 | 同 OUT 冒烟后全量 = 从 step 20 续；数值无害（decay=0 硬拷到 200、warmup 正确续、beta1=0）。本机冒烟已用 `_smoke_stage2_cd`，live run start_step 0。照抄过期路径的真实后果是 `mkdir -p` 到 `/data`（系统盘）并写 ~40 GB ckpt。 | 文档：过期前缀改为 `$R` 或 `$(pwd)`；RUNBOOK 加 `_smoke_stage2_cd` 示例。可选：`:579-580` 对 `finished=true 且 step<max_steps` 打显眼日志。保持 RESUME=auto。 |
+
+### 已驳回（3/3 反驳，仅供参考）
+- **S3-AUTOCAST-2**（accelerate 自带 autocast 使 `cache_enabled=False` 失效）：torch 2.11 的 `cached_cast` 条目保留 autograd 边，实测 71/71 参数有梯度且位级等同；且所有启动路径都是 FSDP+bf16（参数为 bf16 视图，从不 cast）。`dmd_arope.py:43-47,1063-1064` 的注释在本 torch 上是保险带而非必需。
+- **F4**（resample 可超包络 11%）：需 20 个合成 cell 在一个窗内（N≥40），而 rollout_max=27 ⇒ 最多 7 个；全语料实测最大 1.033×（+0.39 token），0.05%；`config.json:35`/`dmd_arope.py:556-560` 已量化记录为 1.006×/0.2%；钳制乘子也无法做到 0%。
+- **R6**（stage3 只靠 `/dits/` 子串守卫）：这正是 `PORT_SPEC.md:562` R12 规定的缓解，且严于上游（`$U/stage3_dmd/dmd_gm.py:338` 无任何检查）；所提 `role` 门对无 role 键的 stage1 ckpt 反而默认放行。
+
+## 3. 计划运行所需的最小精确操作
+
+调用方 shell 里 `$ROOT` 为空（`run.sh` 头注释），**所有路径写绝对**。
+
+### 3a. Stage2 CD（8×H200，teacher==student==stage1_long_K3/dit_step1000）
+**当前状态**：`outputs/distill/stage2_cd` 已在跑（01:43:07 起；world 8 / GA 1 / train_batch_size 8；`--cd_tail_t 0 --teacher_solver_steps 1 --resume auto`；step 160 @ 01:56:58，5.06 s/步 ⇒ 8000 步 ≈ 11.3 h；尚无 progress.json，首存 step 500）。冒烟已完成（`_smoke_stage2_cd`，01:37-01:41，step 20 finished）。**配方正确，无需改动，不要触碰。**
+
+若需重启/续训（崩溃后），必须原样传同一组变量（GA 变了 `cd_arope.py:587-590` 会拒；`AR=` 不得更换，见 INFRA-2）：
+```
+DISTILL_NGPU=8 GRAD_ACCUM=1 \
+ACC_CFG=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/distill/configs/accelerate_zero2_8gpu_ga1.yaml \
+AR=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/outputs/distill/stage1_long_K3/dits/dit_step1000.safetensors \
+bash distill/stage2_cd/run.sh          # OUT 默认 outputs/distill/stage2_cd，RESUME=auto 续
+```
+（冒烟形式同上，另加 `MAX_STEPS=20 SAVE_EVERY=10 LOG_EVERY=1 PORT=29542 OUT=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/outputs/distill/_smoke_stage2_cd`。）
+
+运行中监控：`train_log.csv` 中 idx==2 行的 `tgt_err` 相对 step≤200 基线（CD-1 漂移表）；`cd` 降而 `cm_err` 升 = 塌成自洽糊图（`cd_arope.py:265-268`）。
+
+需用户拍板、**只能用于新 OUT 的重跑**（不改 live run）：`CD_TAIL_T=100`（补 σ=0.1011 节点，部署节点不变，4 对；改后同步 argparse/run.sh/config/test_1）；`SOLVER_STEPS=4`（默认 1 下 Euler 截断预算 Σdσ²=0.187，为上游 48 轴的 5.5×，已在 `config.json _teacher_solver_steps` 记录）。
+
+### 3b. Stage3 DMD 之前必须做的编辑（不影响 live run）
+1. **必做** `distill/eval/sweep8.sh:50`、`photodrift.sh:16`、`paired50.sh:23`：删除 `--sink_size 3 --kv_window 17 --rope_cap 20`；替换 `sweep8.sh:17`/`photodrift.sh:12`/`paired50.sh:12` 的 `/tmp/claude-2011/...` 死路径（GEO-1）。
+2. **必做** `rollout_eval.py`：读 ckpt 元数据 `kv_window/sink_size/rope_cap` 并在不一致时拒绝（GEO-1）。
+3. **强烈建议** `dmd_arope.py:1237` 后加 gnorm==0/NaN 中止（S3-GUARD-3）。
+4. **建议** `stage2_cd/run.sh:42-43`、`stage1_ar/run.sh:39-40`、`stage3_dmd/run.sh:52` 加 DISTILL_NGPU 分支 + eff-batch 断言（INFRA-1/R4）。
+5. **建议** 文档：`stage3_dmd/config.json:39,43-46`、`PORT_SPEC` R3/R10/`:666-668`（DMD-1）；`cd_arope.py:120-121`（CD-1）；GEO-2 五处措辞；5 个测试文件 ROOT（GEO-5）。
+6. **stage2 结束后** 再合入 INFRA-2 的 progress.json 身份校验（stage2+stage3）。
+
+### 3c. Stage3 DMD 冒烟（stage2 出 `ema/ema_step8000` 后）
+```
+DISTILL_NGPU=8 \
+ACC_CFG=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/distill/configs/accelerate_fsdp_dmd_8gpu.yaml \
+MAX_STEPS=20 SAVE_EVERY=10 LOG_EVERY=1 PROFILE_MEM=1 PORT=29541 \
+OUT=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/outputs/distill/_smoke_stage3_dmd \
+STUDENT=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/outputs/distill/stage2_cd/ema/ema_step8000.safetensors \
+bash distill/stage3_dmd/run.sh
+```
+（`ACC_CFG` 仅为文档一致，功能上 4gpu yaml + `--num_processes 8` 等价；`TEACHER` 默认 `outputs/arope_10k_ga2/step-10000.safetensors` 存在；`STUDENT` 绝不能是 `dits/`。）
+冒烟门：step 0 dmd≈0（FSDP 下 ~1e-3，非恰 0）；**step≥1 每步 gnorm_g 有限且 >0**（这是 streamed 重放在 FSDP 多 rank 上的���次执行，S8 仍开，必须人工看）；dmd 进 +0.15…+0.68 带；critic 下降；`_budget_gate` 通过；对 `ema_step10/20` 跑 `sweep_photodrift.sh`（已改成训练几何）确认无立即变暗。
+
+### 3d. Stage3 DMD 全量
+```
+DISTILL_NGPU=8 \
+ACC_CFG=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/distill/configs/accelerate_fsdp_dmd_8gpu.yaml \
+OUT=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/outputs/distill/stage3_dmd \
+STUDENT=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/outputs/distill/stage2_cd/ema/ema_step<N>.safetensors \
+bash distill/stage3_dmd/run.sh        # 500 步，eff batch 8（=上游 cd15k 的 8 rank），save_every 50，keep_emas 25
+```
+里程碑挑选：`rollout_eval`（位移/动作服从）+ `sweep_photodrift`（41 cell std/mean vs stage2 基线），**均按 sink 0/w20/cap 20**；sink 3/w17 只作显式 A/B。
+
+## 4. 已验证正确（摘要）
+
+**CD 数学与循环**
+- 配对 σ_t=sig[idx]>σ_tn=sig[idx+1]，idx∈[0,K−2]（`cd_arope.py:758,764-766` ≡ `cd_gm.py:105-106,435`）；部署轴 idx [0,625,833,937]，σ [1,0.75,0.5006,0.251597]，t=σ·1000 位级对齐（`flow_pf.py:263-288`，DiffSynth `flow_match.py:52`）。
+- 加噪 `(1−s)x0+sε` cell 0 钉住（`:212` ≡ `cd_gm.py:79-83,111`）；教师 Euler `lat_t+(s_tn−s_t)v`，v=ε−x0，每子步重钉（`:233-234` ≡ `:122-123`；DiffSynth `flow_match.py:391,409-411`）；cm_t 带梯度、cm_tn 在 no_grad 下由 EMA 算（`:237-245` ≡ `:126-131`）；loss=MSE(cells 1:)，目标 detach（`:248` ≡ `:133`）；oracle loss==0、上游转写位级一致（12/12 pass）。
+- cond/clean_x 在 teacher/student/EMA 三次前向中是同一对象，EMA 取 x=lat_tn 但 clean_x=x0（`:202-206,225,238,244` ≡ `:100-102,117-131`）；`first_frame_cond=True`+标量 t 复现 `_per_frame_t`（`arope_dit.py:167-194`）。
+- v 空间 CFG 只作用于冻结教师，uncond=NULL_PROMPT（stage1 scene_dropout=0.1 ⇒ 分布内），cfg==1 跳过（`:206,227-231,666-682` ≡ `:117-121`）。
+- idx 流：CPU 生成器 seed 1234+seed 全 rank 同、每 micro 一抽、按 micro 计数快进，续训拒改 GA/轴（`:741-743,758,573-593` ≡ `:417-422,435`）。
+- 仅一个 nn.Module 进 prepare（`:656`；accelerate 只保留最后一个 Module）；教师 bf16 冻结 eval（`:622`）；EMA fp32 影子非 nn.Module（`ema.py:66-86`），在 DeepSpeed 真实 step 之后每优化器步更新一次（accelerate `utils/deepspeed.py:264-284`；`:777-793`）；ema_start 语义与上游一致（步 1..200 硬拷，201 起衰减）；ZeRO-2 下参数全量非分片，zip 顺序/数量/形状断言（`ema.py:49-55`，DeepSpeed `stage_1_and_2.py:805-808`）。
+- AdamW(2e-6, wd 0, betas (0,0.999))、warmup 100 LambdaLR、clip 10 注入 DS 并生效（`:641-645,511-512`；live train.log 印证）；GA 门关闭 accelerate 静默采纳路径（`:513-519`）；`_ckpt_meta_gate` 接受 stage1_long_K3 step1000（header：teacher_forcing=True、windowed 0/0、827 BF16 张量，额外键 long_seg_dir 被忽略）。
+- 无隐藏 F=61 依赖：F 取自 latent shape，RoPE 表预算到 1024（`cd_arope.py:171,185-193`，DiffSynth `wan_video_dit.py:405-407`）。
+- 三处配方一致：stage2 22 常量、stage3 27 常量 run.sh==argparse==config.json（上游 A 的 6 处 run.sh/argparse 不一致均按 run.sh 侧解决）。
+
+**调度器**
+- 1000 格网格、权重、`training_target=noise−x0`、`add_noise_pf` 逐 cell 广播（fp32 算 1−σ 再 cast，优于上游 bf16）、`idx_of_t/t_of_idx` 往返、`cd_axis(N)` 与 `set_timesteps(4)` 陷阱规避、`ti2v_noise` 钳制 [20,980)=σ∈(0.0969,0.9959]——全部与 DiffSynth `flow_match.py` 及 `$U/gamemaster/flow_match.py` `torch.equal`（33 tests pass）。
+- bf16 时间步取整（750→752 等）是 SFT/stage1/2/3/部署共用的约定，无 train/infer 网格错配。
+
+**DMD/critic 数学、步序、FSDP**
+- `x0=noisy−σ·v`，共用一份噪声与 t（`dmd_arope.py:358-370` ≡ `dmd_gm.py:204-219`）；归一化 `|clip−pred_real|` 在 eval_idx 上均值 clamp 1e-2 + nan_to_num（`:371-377` ≡ `:221-224`）；代理 `0.5·mse(g,(g−grad).detach())` 梯度=grad/n，锚 cell 0（`:378,:391`；CPU autograd 验证）；cfg 1.0/3.0 与上游转写位级一致。
+- eval_idx=tail 1..20（`grad_start_range :260-262`，`:353`，`:454` 断言锚不打分）；critic 目标 noise−clip 无时间步权重（`:408-413` ≡ `:244-245`）；教师/critic `causal=False`、无 clean_x/kv_cache、`first_frame_cond=True` 与 SFT 一致（`:935-936,:131-132,:457`）。
+- 步序 rollout→DMD loss→student backward→clip→opt_g→slr_g→EMA→critic×5→slr_c（`:1189-1257` ≡ `dmd_gm.py:681-706`）；AdamW betas (0,0.999) wd 0.01、lr 2e-6/1e-6、warmup 100、clip 10 每模型经 FSDP 自身 `clip_grad_norm_`（`:1054-1059`）。
+- prepare 前全部 `.float()` 并断言（`:957-962`）；`fsdp_use_orig_params` 断言（`:836-842`）；教师不 prepare；EMA FSDP 包 fp32 冻结、zip 更新（`ema.py:41-63`）；`get_state_dict`+safetensors bf16 827 键（`runlib.py:297-353`）；`dit.freqs` 在��个裸模块上搬到 device（`:997-1000`）；accelerate 1.14 对 FSDP1 多模型+多优化器无限制。
+- 流式重放精确：提交在 no_grad 下且缓存 K/V detach（`arope_dit.py:356-357`），逐帧代理 `0.5·mse/n_eval` 求和=融合代理（`:325`；test_7 rel<1e-6，`config.json _measured cos 1.0`）；梯度检查点永不触及 rollout（`arope_dit.py:689-692`）。
+- N/grad_start/exit_step 从共享种子抽并由 rank 0 广播（`:1120-1148`）⇒ 各 rank 前向/反向次数一致，无集合通信错配；`_budget_gate` 在 8 rank 下保守通过；`_KVRef` 在 torch 2.11 `_apply_to_tensors` 下按身份返回（`torch/distributed/utils.py:252-253`）。
+
+**Rollout 几何与 DiT 包装**
+- RoPE 重锚定精确：q·k 只依赖位置差，f64 下窗口平移 rel 2.3e-15（`build_world_freqs`、DiffSynth `rope_apply :94-100`；test a1/a3）；`retained_cells`/`window_offsets_for_frame`/`forward_frame` 的 tpos_j=j−a、off_win_j=off[j]−off[a] 自洽（`preflight.py:77-131`，`arope_dit.py:893-908`）；`rope_cap==sink+kv_window` 四处强制。
+- 训练 rollout 与 `rollout_eval.deploy_rollout` 同一代码路径（`rollout_eval.py:79` 导入 `frame_cond`；`:308-358`）：同 renoise 链 `den=x−σ_i v; x_next=(1−σ_{i+1})den+σ_{i+1}ε`、t=0 detach 提交每 k 一次；exit_step 语义 ≡ `dmd_gm.py:140-184`（部署跑满表为已记录差异）。
+- 驱逐被训练到：`rollout_min≥sink+kv_window+1=21`，`ROLLOUT 21..27` ≡ 上游；`preflight_sliding_kv` 断言驱逐改变输出（`preflight.py:245-287`）。
+- 掩码：块因果 [S,S] 与 TF [2S,2S] 谓词 ≡ `gamemaster_dit.py:536-567`，所有 local/sink/drop_anchor 组合闭式+探针零泄漏；SDPA bool 语义 True=可见；窗口化 TF 同时约束 clean 半边（已记录，且比上游 `:549-554` 更正确）；两半共享 f_start=0 RoPE、per-cell 文本、head 只读 noisy 半边、mask_emb 两半都加（≡ `:496-586`）。
+- KV 路径缓存 pre-RoPE K、读时重旋转、最老非 sink 驱逐、commit 标志/detach、`_evict` 为 while 循环（上游单 if 的超集）、n_sink 钳制（GraphGWM 修的 bug 已含）；绝对模式 f_start 与上游 post-RoPE 缓存等价；`forward()` 的 kv 路径守卫（first_frame_cond False、无 clean_x/offset_px/mask_input、无 gc、frame_index 必填）。
+- 偏移/文本外推对齐（`extend_context` hold ≡ `dmd_gm.py:632-637`；`extend_offsets` 积分与 `parse_actions` 一致；`assert_rollout_aligned` 每次前调用）；resample 步长≤语料 per-cell p99.9；当前 cell 掩码在 k=20→21 连续（test d1）。
+- 无重复提交：`self_forcing_rollout`、`stream_grad_replay`、`deploy_rollout`、`probe_student_grad` 均每 k 恰一次（spy 测试 `test_stage3_dmd.py:99-135`、`test_kv_equiv.py` test_7/8）。
+
+**启动器/环境**
+- `common.sh:36-44` 8 卡路径正确（CUDA_VISIBLE_DEVICES 0-7、NCCL_P2P_LEVEL=NVL、OMP 16×8≤160 核）；端口无冲突（29532/29536/29541+/29545）；磁盘 26 TB 空闲，全量 stage2 ≈180 GB；数据前置全部存在（train_eybx 20983、val_eybx 1175、sidecar、text_table）；stage3 `STUDENT` 路径含 `/dits/` 双重拒绝；`_assert_run_dir_safe` 拒外阶段目录/10 min 内活 train.log。
+
+### 信息级备注（非发现）
+- `PORT_SPEC.md` 仍写 "Expect loss ≈ 0 at step 0" / S7(e)，与 `config.json _expect`（部署轴 step 0 loss 非 ≈0）矛盾，会误导冒烟判读。
+- `deploy_schedule` 吸附容差用全局 0.5·max(Δt)（≈2.47），t=3.0 会被静默吸到 idx 999 而非拒绝；固定 4 步表不受影响。
+- 窗口化 TF 的 `local` 含查询 cell 自身：匹配部署 kv_window=20 需 local=21（F=21 下=全因果）；配置 0/0，无影响。
+- F=61 TF 掩码 [2S,2S] S=23790 ⇒ 2.26 GB bool 常驻 `_mask_cache` + SDPA 转 bias ~4.5 GB 瞬时；混 21/61 cell 会同时保留两份。
+- `_KVRef` 恢复只在根 `forward()`；依赖嵌套 FSDP 单元不 cast/重建输入（accelerate 留 `cast_forward_inputs=False`）。
+- streamed 重放在 FSDP **≥2 rank** 上从未在任何参考树中执行过（GraphGWM 生产用 fused；端口 streamed==fused 检查为单卡；S8 未闭合）——3c 冒烟是首次，务必逐步看 gnorm_g。

backup_2026-09-14/code/distill/MANIFEST.md5

@@ -0,0 +1,63 @@
+# distill/ 状态快照 —— 生成于 2026-09-07 14:59:44  stage3-onpolicy 起训前
+# 本版通过：339 tests + stage2-flowmap 8000 步产物 + stage3-onpolicy 20 步 8 卡 FSDP 冒烟
+# （horizon=4、~11 s/step、峰值 75-119 GB/卡、832/827/832 存档键数正确）
+# 起训前： md5sum -c distill/MANIFEST.md5
+
+d41d8cd98f00b204e9800998ecf8427e  distill/__init__.py
+f4f75f84ac79b1d03519f2e204959546  distill/arope_dit.py
+5584bc8392d100d475c187b1d95d56d0  distill/common.sh
+e6de64ae254088f01625e19221111add  distill/configs/accelerate_fsdp_dmd_4gpu.yaml
+e2ce1fd7ca6d187cacbfa81e78535055  distill/configs/accelerate_fsdp_dmd_8gpu.yaml
+0d2bb6efe49b050aea306541b4a34d65  distill/configs/accelerate_zero2_4gpu_ga2.yaml
+85ee449b05bf9f686f14c45196455523  distill/configs/accelerate_zero2_4gpu_ga4.yaml
+6a34ea5734cf13f06e924d9754a70c12  distill/configs/accelerate_zero2_8gpu_ga1.yaml
+7548174929fd59a62a8a1d194b088575  distill/data.py
+3c4bb6e458ae3b9fe4d4ca812e581a18  distill/ema.py
+efacdb167135b5bc106b1cadc8bd997c  distill/eval/anystep_probe.py
+d34dc0469f49fc5826f8094062adf385  distill/eval/flowmap_liveness.py
+33963f63c251db48899df3c550c38e2e  distill/eval/null_branch_probe.py
+86f7ccdda50e7027e88cb0550863f451  distill/eval/rollout_eval.py
+908f04c018f2ce1d0e207fd1f88933bd  distill/eval/sweep8.sh
+dac8e1258f3cbf34db5d2b419ef61aa9  distill/eval/tf_rollout_probe.py
+a3324266dd8f3e1ff437a53fa16cd993  distill/flow_map.py
+7707dd7f57bd3aa3415c96b9baf99f88  distill/flow_pf.py
+d67ba664053d958506178a2d9fe3a470  distill/flowmap.py
+166cbbbda1d4bd5fa792cc8f75900a64  distill/onpolicy.py
+94fe49fa6b1688529dd563351197b6f8  distill/runlib.py
+0ab413868394a1467be20ec317be4ae9  distill/stage1_ar/ar_tf.py
+369dcf3e22a04fa74fb8502e2e8e99d7  distill/stage1_ar/config.json
+cff6e9ea64e3648ccef727ccebebac62  distill/stage1_ar/run.sh
+cf14a9aa4c9f52f15111ac06e6cb3936  distill/stage1b_adapt/run.sh
+45f9f10b4ab8a6299e43c291ab4f4b91  distill/stage2_cd/cd_arope.py
+2402bf288b7f78d7a30bf65e0aa3868c  distill/stage2_cd/config.json
+1f635a72fb0e938028afa356b28479cd  distill/stage2_cd/run.sh
+67f689be9600a8bd540480bb8585a0ea  distill/stage2_flowmap/config.json
+06875f5dee35b8cc28ff90e456545cf1  distill/stage2_flowmap/fm_arope.py
+22507adaec0ef9339b292a24095b3c9f  distill/stage2_flowmap/run.sh
+043f7d0fcce74f87fd185c3cca6ff9fb  distill/stage3_dmd/config.json
+20e550d8b05ac9f20412979332400797  distill/stage3_dmd/dmd_arope.py
+9de618ffe59f845e52f421665a4adf55  distill/stage3_dmd/run.sh
+dfa2de24921c087c842529242e1aeb10  distill/stage3_onpolicy/config.json
+6361310acf172703a94a804bf230483f  distill/stage3_onpolicy/onpolicy_arope.py
+6373a4762e4958eb01dab45b4a1aaa4c  distill/stage3_onpolicy/run.sh
+0a5c95ab6dfc326869d8185efb816eae  distill/tests/preflight.py
+0c42f9fed1bf89a626a313750cb97398  distill/tests/test_bitparity.py
+644e50542cba2c5accda10b496616cc0  distill/tests/test_cond_dropout.py
+8645b069fa5870292e9e80768c56085b  distill/tests/test_data.py
+b4424fe522987570e086b3c247baab74  distill/tests/test_flow_pf.py
+781f3bdeafff7cc4373f5391460ad87a  distill/tests/test_flowmap_model.py
+6f8c6eb2dda143c914e863879f838a92  distill/tests/test_flowmap_objective.py
+61d614833748b9d1fdf4a6ad369daba3  distill/tests/test_flowmap_upstream_diff.py
+56f50b9ce20f2f8a756cdbe140bafbc5  distill/tests/test_fm_arope_wiring.py
+70abc01e9dece1d8c1ac349bacf65e23  distill/tests/test_iface_fidelity.py
+2bb100531f06f83e6c2ac2df864b995a  distill/tests/test_kv_equiv.py
+f3cd288f4db613489a9cb753816ff1b7  distill/tests/test_onpolicy.py
+fdb6bfa9b7e4d365fed7e71fa8b5ce71  distill/tests/test_runlib.py
+ae80f988c038cf2a1e5558aecbe51024  distill/tests/test_stage23_upstream_diff.py
+0090187c8363f55c0e406499a211c367  distill/tests/test_stage2_cd.py
+a45eb617da973973aa8037dba48155b8  distill/tests/test_stage2_flowmap.py
+2d41d3760508df0b4608967ba0a5b366  distill/tests/test_stage3_dmd.py
+a8ac99c5a2ec7a64637cc8a46b551340  distill/tests/test_tf_port.py
+436730aa47115a5d9166f2426c659827  distill/tests/test_world_rope_window.py
+d317cfcb9b10cd1ad326ca9dca1cd41c  distill/eval/photodrift.sh
+a6d75b49e766cedc15b1ebc9e23119c2  distill/eval/sweep_photodrift.sh

backup_2026-09-14/code/distill/OPSD_OFFICIAL_NOTES.md

@@ -0,0 +1,50 @@
+# OPSD-V 官方实现（vendor/OPSD-V，commit 8ecb193，2026-08-31）与 GraphGWM 移植（hf/worldmind_ref/GraphGWM/train/code_distill/stage4_opsd）的差异笔记
+
+读的是 `vendor/OPSD-V/model/opsd_streaming_training.py::compute_opsd_generator_loss / setup_opsd_sequence / _replace_cache_chunk_with_gt`、
+`trainer/opsd_streaming.py::fwdbwd_one_step_opsd / _ema_update_lora_teacher_adapter`、`configs/train_self_forcing_lora.yaml`。
+
+## 官方机制（逐条）
+1. **部署表**：`denoising_step_list: [1000, 750, 500, 250]` + `warp_denoising_step: true`（把 t 映射到 scheduler 的 shift 网格）——与本仓库的 4 步部署表 + `deploy_schedule` 吸附到 1000 网格同构。chunk = 3 个 latent 帧（`num_frame_per_block: 3`）；本仓库 chunk = 1 cell。
+2. **两条缓存都从 GT 首 chunk 起**（`opsd_use_gt_first_chunk: true`，t=0 写入）。窗口起点 `window_start` 在 GT 长视频里**随机**（按 block 对齐；`window_start_override: 0` 时固定从 0 起），rollout 长度 `opsd_loss_frames`（180 帧）≤ `opsd_max_loss_frames`（243）。
+3. **每个 chunk、每个去噪步**：同一份 `noise_chunk` / 同一 `noisy_input` 先喂教师（no_grad，权重换成 EMA/teacher adapter），再喂学生（需要 loss 的步 enable_grad）。`opsd_teacher_trajectory_mode: student` ⇒ `teacher_noisy_input = noisy_input`（教师走学生访问的 noisy 状态）。
+4. **loss**：`flow` ⇒ `mean((student_flow − teacher_flow.detach())²)`；`opsd_loss_step_mode: all` ⇒ 每个去噪步都算；`backward_per_chunk: true` 时**逐步** `(step_loss / total_loss_terms / grad_accum).backward()`（不是逐 chunk 累加后再反传），反传后 detach 两条缓存里的张量；`total_loss_terms` = 全窗口被打分的 (chunk, step) 对数。
+5. **重新加噪**：`noisy_input = add_noise(student_chunk.detach(), transition_noise, next_t)`（学生自己的 x0 + 新噪声；stop-gradient 过渡），与本仓库的 renoise 链一致。
+6. **chunk 结束后写缓存**（`context_noise: 0` ⇒ t=0 干净写入）：学生缓存写**学生**的最终 x0；教师缓存写**教师自己**的最终 x0（`teacher_chunk`，用 EMA 权重 `_write_chunk_to_cache(use_teacher_ema=True)`）。
+7. **教师缓存的 GT 替换**（`opsd_teacher_context_mode: gt_kv`）：写完之后，若已生成 ≥2 个 chunk，把教师缓存里**倒数第二个** chunk 原位替换成 GT（`_replace_cache_chunk_with_gt`：暂存 `global_end_index/local_end_index`，在该位置用 t=0 前向重写 K/V，再恢复 end index）。于是打分第 k 个 chunk 时教师缓存 = **GT 0..k−2 + 教师自己的 k−1**。
+8. **教师权重**：`lora_ema_enabled: true, ema_weight: 0.9999, ema_start_step: 0` ⇒ 教师 = 学生（LoRA）的 EMA，每步更新；`max_iters` 3300（longlive）/ 见 self_forcing 配置。LoRA rank 256，lr 1e-5，`max_grad_norm_generator: 10.0`，gradient checkpointing 开。
+9. `opsd_loss_start_frame: 21`（前 7 个 3 帧 chunk 不打分）；`opsd_teacher_use_future_gt_context: false`（教师不看未来 GT，默认关）。
+
+## GraphGWM 移植与官方的差异（本仓库移植时要显式选择）
+| 项 | 官方 | GraphGWM | 本仓库建议 |
+|---|---|---|---|
+| 教师缓存里最近一个 chunk | **教师自己**的 x0（`teacher_chunk`） | **学生**的 k−1（`--teacher_keep_last 1`） | 加 `--teacher_last_chunk {teacher,student,none}`，默认 `teacher`（官方） |
+| 反传粒度 | 逐去噪步 `(step_loss/total_terms).backward()` | 逐 chunk（chunk 内各步求和 / 步数） | 逐步反传（显存更小），归一化按全窗口 total_terms（官方） |
+| 教师权重 | 学生 EMA 0.9999（每步更新） | 冻结（`--ema_decay 0`），理由：全参微调 + 短跑下 EMA 会漂向学生 | 默认冻结，`--ema_decay` 可开（0.9999） |
+| loss_start | 21 latent 帧 | 12 | 默认 12，可调；文档写明官方是 21 |
+| 窗口起点 | GT 长视频里随机（block 对齐） | 段首（cell 0） | 默认段首（61 cell 全用）；`--rollout_cells < 61` 时随机起点 |
+| chunk | 3 latent 帧 | 1 | 1 cell（部署也是逐 cell） |
+| 缓存写入 | 官方在**原位**替换倒数第二 chunk 为 GT（cache 是固定长度数组 + end index）。**写入顺序**：先 `_write_chunk_to_cache(teacher_chunk)`（此刻缓存里还是 GT 0..k−2 + 教师的 k−1），再把 k−1 换成 GT、**不重算最近一格** ⇒ 打分 k+1 时最近一格（教师的 k）第 1 层起的 K/V 嵌着"教师的 k−1"，不是全 GT 历史 | 追加式滑窗缓存做不到原位替换 ⇒ 真实缓存**滞后一帧** + 每 chunk 快照再把学生的 k−1 commit 进快照（K/V 在全 GT 历史上算 —— latent 层面同官方，K/V 层面不同） | 快照法 + 官方写入顺序：cell k 打分完教师把自己的 x0 commit 进**它打分用的那条快照**（上下文正是官方的），取走那格 raw K/V（`last_kv_frame`），cell k+1 用 `append_kv_frame` 接到 GT 0..k−1 的新快照末尾。latent 与逐层 K/V 都与官方相��（test_stage4_opsd::test_13 逐位核对）。student 模式保留 GraphGWM 口径 |
+| 教师前向实现 | 换权重（summon_full_params 拷贝 EMA） | 单独一份冻结 bf16 教师模块 | 单独教师模块（同 stage3 real-score 教师） |
+| RNG | 教师前向前后保存/恢复 RNG state | 无 | 无 dropout，可略；但打分状态的噪声必须显式共享 |
+
+## 补充核对（trainer/opsd_streaming.py）
+- 训练循环顺序：`zero_grad → [GA 次 fwdbwd_one_step_opsd（逐步 backward，grad 捕获到 deferred buffer）] → 恢复 grad → clip(10.0) → optimizer.step → EMA 更新（LoRA 模式：teacher adapter ← decay·teacher + (1−decay)·student，decay=0.9999）→ step+=1`。
+- 教师 adapter 起始 = 学生 adapter 的拷贝（训练前两者相同 ⇒ 教师起点 = 未 OPSD 的基座）；3300 步后 0.9999^3300≈0.72 的权重仍是初始权重 ⇒ 官方教师本质上是「几乎冻结、缓慢跟随学生」。GraphGWM 用完全冻结教师是它的 decay→1 极限，差异可忽略；本仓库默认冻结，可选 EMA。
+- 每个去噪步前保存 CPU/CUDA RNG state，教师前向后恢复，再跑学生：保证学生前向的随机数（若有 dropout 等）与「没有教师前向」时一致；本仓库 DiT 推理无 dropout，可略。
+- `context_noise=0` ⇒ `add_noise(x0, ε, t=0)` = x0（`context_noisy` 名义上加噪，实际是干净写入），与本仓库 t=0 commit 一致。
+- `generated_spans` 不含 GT 首 chunk（只记生成的 chunk）⇒ 第一个生成 chunk 打分时教师缓存 = GT 首 chunk；第 k 个生成 chunk 打分时教师缓存 = GT 首 chunk + GT 生成段 1..k−2 + 教师自己的 k−1（写缓存后再把倒数第二个替换成 GT，最近一个保留）。
+
+## 本仓库 `distill/stage4_opsd/opsd_arope.py` 的最终选择（2026-09-10）
+| 项 | 选择 | 开关 |
+|---|---|---|
+| 教师缓存最近一格 | **官方**：教师自己在最后一个去噪步、学生的 noisy 上算出的 x0（未打分 cell 也跑一次教师最后一步） | `--teacher_last_chunk teacher` / `student`（GraphGWM） |
+| loss 归一化 | **官方**：整窗所有打分 (cell, step) 的平均（逐 cell backward 时每 cell 均值 ÷ 打分 cell 数） | `--loss_norm window` / `cell`（GraphGWM） |
+| 教师权重 | 冻结 init_ckpt（官方 EMA 0.9999 ≈ 冻结；FSDP 下未实现 EMA）。开 EMA 时教师留 **fp32**：bf16 半 ulp ~2e-3 吃不进 1e-4 的增量，bf16 教师上 EMA 会静默退化成冻结（trainer 在加载处断言） | `--ema_decay`（仅单进程） |
+| 反传粒度 | 逐 cell（与官方逐步等价） | — |
+| grad_clip | **官方** 10.0（`max_grad_norm_generator`；stage3 同值）。曾是 1.0（未记录的偏离；window 归一化下会每步触发，把 AdamW 变成归一化 SGD） | `GRAD_CLIP` |
+| Adam betas | **官方** (0.0, 0.999)（`train_self_forcing_lora.yaml` beta1/beta2；stage3 同值）。曾是 (0.9, 0.95)（无出处） | `ADAM_BETAS` |
+| loss_start | 12（GraphGWM；官方 21） | `LOSS_START` |
+| 窗口起点 | 段首 cell 0（VAE 首 latent 是 TI2V 锚点，不能随机） | — |
+| 缓存写法 | 滞后一格的真值缓存 + 每 cell 浅快照 + **官方写入顺序**：最近一格是教师在上一格打分完 commit 进它打分用的快照后取走的 raw K/V（`last_kv_frame` / `append_kv_frame`，arope_dit），不是在全 GT 历史上重新 commit（那样 latent 同、第 1 层起的 K/V 不同，教师轨迹从 cell 3 起与官方分道） | — |
+| RNG | 不需要（DiT 推理无 dropout；噪声显式共享） | — |
+| 条件 dtype | offset_win / mask 按 latent dtype 量化一次后给两个模型（FSDP 根前向 cast 变空操作） | — |

backup_2026-09-14/code/distill/RUNBOOK.md

@@ -0,0 +1,248 @@
+# ActionRoPE 三阶段蒸馏 —— 串联手册
+
+因果强制（causal forcing）三阶段：**stage1 teacher forcing → stage2 CD 初始化 → stage3 DMD**。
+移植自 `teawhite/worldmind` 的 `WorldMind/distill/code/`（正典），`stream_grad_replay` 与
+`preflight_sliding_kv` 反向移植自 `GraphGWM/train/code_distill/`。
+
+设计与逐条对抗评审结论在 `PORT_SPEC.md`。本文件只讲**怎么跑**。
+
+## 铁律
+
+* **只用 GPU 0-3。** 4-7 是别人的卡。`common.sh` 会在启动前检查 0-3 是否空闲，
+  查不到占用（nvidia-smi 出错、返回行数不对）就拒绝启动 —— fail-closed，不是 fail-open。
+  阈值只能调严不能调松（`GPU_BUSY_MIB=999999` 不再是偷卡后门）。
+* **一次只跑一个阶段。** 三个阶段都要满 4 张卡。
+* 每个阶段的配方常量在**三处**重复：`run.sh` 的默认值 / `<stage>.py` 的 argparse 默认值 /
+  `config.json`。上游 WorldMind 这三处有六项不一致，每一项都对应一次已知崩溃。
+  **改配方 = 同时改这三处。**
+
+## 用户已拍板
+
+| 项 | 值 | 理由 |
+|---|---|---|
+| teacher | `outputs/arope_10k_ga2/step-10000.safetensors` | 不再续 SFT，直接开蒸 |
+| stage1 等效 batch | 8（4 卡 × GA2），lr 1e-5 | 同 lr 下微调阶段更多 optimizer step 划算 |
+| 部署采样表 | 4 步 `1000,750,500,250` | stage2 的 CD 轴与 stage3 的 rollout 都按这张表训；**训练与推理必须逐点相同** |
+
+## 跑之前
+
+```bash
+cd /opt/dlami/nvme/zhiyangdeng/ActionRoPE      # 本机；旧机是 /data/zhiyangdeng/ActionRoPE
+export PYTHONPATH=$PWD DIFFSYNTH_SKIP_DOWNLOAD=True
+
+# 全量测试（CPU 即可，不抢训练卡）—— 310 项通过 + 2 项 skip（两项 skip 是下面那两道真权重闸）
+# 计数会随 flow map 的文件继续落地而变；红了才是问题，数字对不上先看是不是刚加了测试。
+AROPE_REQUIRE_SIDECAR=1 CUDA_VISIBLE_DEVICES="" .venv/bin/python -m pytest distill/tests/ -q   # 有 AROPE_REQUIRE_SIDECAR 时语料包络门不许 skip
+
+# ★ 最吃重的两项默认都是 skip 的，都要**真实 5B 权重**，每次动过 arope_dit.py 都要跑。
+#   它们一旦失败，后面所有阶段都在量一个错的模型。
+
+# (1) AropeDiT 与 arope_forward 逐位一致（回归锚点本体）
+AROPE_BITPARITY_REAL=1 \
+AROPE_BITPARITY_CKPT=$PWD/outputs/arope_10k_ga2/step-10000.safetensors \
+CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m pytest distill/tests/test_bitparity.py -q -k real
+
+# (2) flow map 装上去的那一刻，stage1b 教师逐位没被动过（对任意 r，不只是 r == t）
+AROPE_FLOWMAP_REAL=1 \
+CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m pytest distill/tests/test_flowmap_model.py -q -s -k real_weights
+```
+
+## flow map（stage2-flowmap 的模型侧）
+
+`AropeDiT.install_flowmap(form="residual", gate=1.0)` 给模型加第二个时间步 `r`
+（`forward(..., r_timestep=)` / `forward_frame(..., r_k=)`），学均速度 `u(x_t, t, r)`。
+
+* **默认 `residual` 形式对任意 (t, r) 与装之前逐位相同**（真权重实测 max|Δ| = 0.000e+00）：
+  `delta_embedding = deepcopy(time_embedding)` 但**第二个 Linear 清零**。上游 AnyFlow 的
+  `convex` 形式只在 `r == t` 的对角线上恒等，off-diagonal 差 1.2e-01 —— 我们 50 % 的样本按构造
+  就在 off-diagonal 上，所以它只留作消融（`--flowmap_form convex`）。
+* **初始化顺序**（写反了不报错，只是隐表示变垃圾，数值闸门看不见）：
+  新 run `AropeDiT() → load_dit_ckpt(stage1b)【strict 827】→ install_flowmap()`；
+  续训 `AropeDiT() → install_flowmap() → load_dit_ckpt(dit_stepN)【strict 832】`。
+  非要反着来就得显式 `load_dit_ckpt(..., allow_missing=runlib.FLOWMAP_KEYS)`，别设成默认。
+* **checkpoint**：多 5 个平铺键 `delta_embedding.{0,2}.{weight,bias}` + `delta_gate`，
+  827 → **832**（`runlib.AROPE_DIT_KEYS_FLOWMAP`，+20.5 MB），格式其余部分一字未变。
+  存的时候 `save_dit(..., expect_keys=runlib.AROPE_DIT_KEYS_FLOWMAP)` 并写 metadata
+  `flowmap=1 / flowmap_form / deltatime_type=r / delta_gate`；读的时候用
+  `runlib.is_flowmap_ckpt(path)` 选采样器。**没有精确 827 键的导出**：
+  `u(x,t,t) = head(E(t) + g·D(P(t)))` 是两条 Sequential 的和，折不回一条 —— 于是
+  `actionrope/infer.py`、`tests/check_ckpt.py`、`eval_bg50.py` 会拒收 flow-map 权重，这是对的。
+* **没装的时候 `r_timestep` 必须是 None 或与 `timestep` 完全相同**，否则直接断言失败。
+  上游那个 stock embedder 收下 r 然后忽略它，接错线会退化成普通扩散模型且不报错。
+
+## stage1 —— teacher forcing 因果化
+
+把双向 SFT 检查点在因果掩码下继续训，得到因果 AR 扩散模型。
+
+```bash
+bash distill/stage1_ar/run.sh
+```
+
+4 卡 DeepSpeed ZeRO-2，12000 步，实测 **7.0 s/step ⇒ 约 23 h**，每卡 37 GB。
+产物 `outputs/distill/stage1_ar_tf/dits/dit_step12000.safetensors`（827 张量 / 10.00 GB）。
+
+`grad_clip` 默认 **10.0**，不是 SFT 名义上的 1.0：SFT 实测 5000 步 `grad_norm`
+p50=0.072 / max=0.481，阈值 1.0 一次都没触发；stage1 的 gnorm 约 2.5，1.0 会**步步都裁**，
+等于把优化器换成归一化 SGD、等效 lr 掉到约 1/3。10.0 同时对齐 WorldMind stage1 与 SFT 的实际行为。
+
+健康信号：`loss_ema` 与 `--val_every 500` 的曲线。**不要**把单次 `grad_norm` 尖峰当发散 ——
+它逐 run 不可复现、同样输入能摆动约 10×。
+
+20 步实机冒烟（`OUT` 必须写绝对路径，`$ROOT` 在调用方 shell 里是空的）：
+
+```bash
+MAX_STEPS=20 SAVE_EVERY=10 LOG_EVERY=1 VAL_EVERY=10 VAL_N=16 PORT=29541 \
+  OUT=/data/zhiyangdeng/ActionRoPE/outputs/distill/_smoke_stage1 \
+  bash distill/stage1_ar/run.sh
+```
+
+## stage2 —— 因果一致性蒸馏（CD）
+
+在**部署轴**上做一致性蒸馏；产物是 **EMA**（不是裸 dit），它才是 stage3 的初始化。
+
+```bash
+bash distill/stage2_cd/run.sh          # 默认吃 stage1 的 dit_step12000（4 卡 GA2）
+DISTILL_NGPU=8 AR=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/outputs/distill/stage1_long_K3/dits/dit_step1000.safetensors \
+  bash distill/stage2_cd/run.sh          # 8 卡形式：GA 自动选 1，run.sh 断言等效 batch == 8；换了 AR= 就换 OUT 或 RESUME=none
+```
+
+4 卡 ZeRO-2（或 8 卡 GA1，等效 batch 都是 8），8000 步，lr 2e-6，betas `0.0,0.999`（CF 约定：目标随 EMA 漂移，动量会追旧目标）。本机 8 卡实测 5.06 s/步 ⇒ 约 11.3 h。
+`KEEP_EMAS=25` —— stage3 的 init 要从这些里挑，别按 `dits/` 的口径只留 2 个。
+产物 `outputs/distill/stage2_cd/ema/ema_step8000.safetensors`。
+
+CD 目标：教师走**一个** Euler 步，学生与 EMA 各自在**自己的** σ 上取 x0 估计后求 MSE。
+（`cm_tn` 必须用 `σ_tn` 而不是 `σ_t` —— 这里曾经有过一个字符的 bug，现在与上游逐位一致。）
+
+**哪个节点没被约束（CD-1）。** 默认 4 节点部署轴 [1, 0.75, 0.5006, 0.2516] 上，部署终读点
+σ=0.2516 只当过一致性对的目标、从不是 source，没有任何一项 loss 约束它：学生在那里等于 stage1
+教师的一步估计**只在 step 0 成立**（student == ema == teacher），此后它是固定点族的自由分量，
+共享参数下 pair 0-2 的梯度会顺带移动它而无恢复项。上游 N=48 轴的 σ=0.096 节点结构相同，
+只是 c_out=−σ 的杠杆小 2.6×。训练中盯 `train_log.csv` 里 **idx==2 行的 `tgt_err`**
+（= ‖f_EMA(·, 0.2516) − x0‖²）相对 step≤200 的基线；`cd` 降而 `tgt_err` 升才是问题。
+`CD_TAIL_T=100` 会在轴末再挂一个 σ=0.1011 节点（4 对，部署节点不变）—— 轴是 EMA 目标的定义，
+**只能用于新 OUT 的新 run**，续训不能改轴；默认 0 是用户拍板的轴，不改。
+
+## stage3 —— self-forcing DMD
+
+```bash
+bash distill/stage3_dmd/run.sh         # 默认吃 stage2 的 ema_step8000
+```
+
+4 卡 **FSDP1**（不是 ZeRO-2），500 步，`SAVE_EVERY=50` —— 部署权重要逐个里程碑人眼挑。
+四个模型：因果学生 / 双向 critic（fake-score）/ 冻结的双向 real-score 教师 / EMA。
+
+`SCORE_FRAMES=GRAD_HORIZON=20`：连续打分尾巴取 1+20 = 21 个 cell，正好等于教师与 critic
+训练时的片长。**首次跑必须先量显存**：
+
+```bash
+MAX_STEPS=20 SAVE_EVERY=10 LOG_EVERY=1 PROFILE_MEM=1 PORT=29541 \
+  OUT=/data/zhiyangdeng/ActionRoPE/outputs/distill/_smoke_stage3 \
+  bash distill/stage3_dmd/run.sh
+```
+
+预期每卡约 66 GiB（140 GiB 的卡）。`GRAD_PASS=streamed` 是默认且不要改成 `fused` ——
+fused 在 horizon=20 时要 189 GiB/卡，必 OOM；streamed 的梯度与 fused 逐位等价。
+（`_budget_gate` 会在读 40 GB 基座权重**之前**就拦住必炸的组合。）
+
+`dmd` loss 应进入 +0.15…+0.68 区间而 `critic` 下降。**这是 min-max，单调收缩的 `dmd` 不是目标。**
+
+## 评测
+
+```bash
+# 少步学生是否还听动作 —— 背景位移 vs 指令位移
+CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m distill.eval.rollout_eval \
+  --ckpt outputs/distill/stage3_dmd/ema/ema_step400.safetensors \
+  --denoise_list 1000,750,500,250 --actions right:1.0:21
+```
+
+`rollout_eval.py` 是 `distill/` 里唯一不走 `common.sh` 的入口，它自带显存守卫；
+训练在跑时它会拒绝上卡（`--allow_busy_gpu` 会把自己硬顶在空闲量减 8 GB，撑死只炸自己）。
+纯 CPU 自检：`--self_test --device cpu`。
+
+### 崩了要先分清是路径还是曝光偏差
+
+```bash
+CUDA_VISIBLE_DEVICES=0 .venv/bin/python distill/eval/tf_rollout_probe.py \
+  --ckpt <任意阶段的 ckpt> --clip clip_Eybx_100000092_000000
+```
+
+同一条 KV-cache 路径跑两遍，唯一差别是提交进缓存的是**真值 latent** 还是**模型自己的 x0**：
+
+* **gt 平稳、self 衰减** ⇒ 部署路径没问题，崩的是误差累积（曝光偏差）。这正是 stage3
+  self-forcing 要修的东西，不是 bug。
+* **gt 也衰减** ⇒ 部署路径与训练不匹配（t=0 提交 / 窗口锚定的 RoPE 或 mask），是移植 bug，
+  别往下走。
+
+**stage1@9000 的实测**（2026-09-06，供后续对比）：gt 全程 22.7–31.8 dB 平稳；
+self 从 cell 1 的 24.39 单调掉到 cell 17 的 16.20，同时 `|x0|` 从 1.111 涨到 1.404
+（真值一直在 1.2 附近）——能量正反馈，像素上表现为亮区从边缘往里收缩、cell 20 全黑。
+路径已排除，符合预期。
+
+### sink / kv_window 的对比在 stage1 上问不出来
+
+`sink=3 + window=17` 与 `sink=0 + window=20` 在 **21 个 cell 之内保留集完全相同**
+（`[0,3) ∪ [3,20) == [0,20)`，没有驱逐），所以两者数学上等价、PSNR 逐位一样。
+它们只在 cell 21 开始驱逐后才分叉。要比这两个配置，得等模型能活过 cell 21 —— 即 stage3 之后。
+
+## 出问题时
+
+* **拒绝启动，说 GPU 有占用** —— 先 `nvidia-smi` 看是不是自己上一个 job 没退干净。
+  别去动 4-7。确认上一次真的死了再用 `FORCE_RUN_DIR=1`。
+* **续训** —— 三个阶段都默认 `RESUME=auto`（有 `progress.json` 就续）。
+  stage1 有完整的 DeepSpeed state 双槽轮转；stage2/3 只热启动权重，优化器不落盘。
+* **loss 恒 0 且不报错** —— 检查 `grad_horizon ≤ dmd_score_frames`（argparse 会拦），
+  以及 `rope_cap == sink + kv_window == F−1 == 20` 这条几何铁律。
+
+## 还没做的两件事
+
+1. **S8：2 卡 FSDP 分片一致性。** stage3 的 `arope_forward` 会直接伸手取 `dit.blocks[i].norm1`
+   这类子模块；FSDP 的 `TRANSFORMER_BASED_WRAP` 只在被包模块的 `forward()` 里 all-gather，
+   伸手取到的是分片、**不报错**。`AropeDiT` 的 monkeypatch 就是为了让每个 block 都从
+   `block(...)` 进入，但"FSDP 真的 all-gather 对了"这件事只能在 2 卡上验。
+   卡一空出来就跑，在 stage3 首launch之前。
+2. **滑动 KV 的吞吐实测。** 每层都在 float64 重算整段历史的 RoPE。
+   `PROFILE_MEM=1` 的冒烟同时会给出 s/step，量到了再信任何 wall-clock 估计。
+
+## stage4 —— OPSD（On-Policy Self-Distillation，长视野后训练）
+
+MeiGen-AI OPSD-V 经 GraphGWM `stage4_opsd/opsd_gm.py` 移植（`distill/stage4_opsd/opsd_arope.py`），并对照
+官方实现 `vendor/OPSD-V/model/opsd_streaming_training.py` 逐行核对（差异表 `distill/OPSD_OFFICIAL_NOTES.md`）。
+后训练 stage3 的 `ema/ema_step300`，**不换采样器、不改推理**。学生用部署的 4 步 renoise 表 on-policy
+rollout 61 个 cell（`data/latent_long/K3` 的长段）并提交自己的 x0；一份**冻结在初值**的教师在同一个
+(noisy, t) 上被求值，但它的缓存里是真值 0..k−2 + 最近生成的 k−1（默认官方口径：教师自己在最后一步算出的 x0，
+`TEACHER_LAST_CHUNK=student` 换成 GraphGWM 的学生 x0；`snapshot_kv_cache` 分叉滞后一格的真值缓存）；
+loss = 两者速度场的 MSE，整窗所有打分 (cell, step) 取平均（`LOSS_NORM=window`，官方；`cell` = GraphGWM），逐 cell backward。
+几何 sink 3 / window 17 / cap 20 / absolute = stage3 训练与部署的几何。**训练面 = 官方的 LoRA rank 256 / lr 1e-5**
+（`LORA_RANK`，distill/lora.py；教师 = 同一模型关补丁，存盘折回 827 键）。`LORA_RANK=0` 是全量微调，2026-09-10 实测 lr 1e-6 下
+step50 起就糊、step100 更糊（`outputs/distill/eval/stage4_opsd/ema{50,100}_60s_indist/`），别用。与官方仍不同的：教师冻结
+（官方 EMA adapter 0.9999 ≈ 冻结）、窗口起点固定在段首（VAE 首 latent 是锚点）、loss_start 12（官方 21，可 `LOSS_START=20`）。
+
+```bash
+# 三处配方（run.sh 默认 == argparse 默认 == config.json）由 tests/test_stage4_opsd.py 逐项比
+CUDA_VISIBLE_DEVICES="" .venv/bin/python distill/stage4_opsd/opsd_arope.py --smoke        # 六条 CPU 不变量，~10 s
+CUDA_VISIBLE_DEVICES="" .venv/bin/python distill/stage4_opsd/opsd_arope.py --dry_run \
+  --output /opt/dlami/nvme/zhiyangdeng/ActionRoPE/outputs/distill/_dry_opsd --max_steps 3   # tiny 三模型端到端
+CUDA_VISIBLE_DEVICES="" .venv/bin/python -m pytest -q -p no:cacheprovider \
+  distill/tests/test_stage4_opsd.py distill/tests/test_kv_equiv.py
+
+# ★ 墙钟与显存都没量过（config.json 的 _budget 全是外推：60–160 s/步、峰值 ≈ 55–60 GiB/卡）。
+#   先跑 20 步冒烟（OUT 必须写绝对路径），看 [mem step*] 的峰值与 step_time，再开 300 步：
+DISTILL_NGPU=8 MAX_STEPS=20 SAVE_EVERY=10 LOG_EVERY=1 PROFILE_MEM=1 PORT=29547 \
+  OUT=/opt/dlami/nvme/zhiyangdeng/ActionRoPE/outputs/distill/_smoke_stage4 \
+  bash distill/stage4_opsd/run.sh
+DISTILL_NGPU=8 bash distill/stage4_opsd/run.sh        # 正式：300 步，每 50 步存 dits/ + ema/，全留
+```
+
+**开跑后要读的两个启动信号**（run.sh 与 `train.log` 里都有）：
+
+1. `[first-step] opsd … gnorm …` —— opsd 必须**小但非零**（bf16 5B 上量级 1e-3…1e-1）：恰好 0 = ���师在读学生
+   自己的上下文、没有任何信号（上游第一次 launch 就这样静默跑了 6 小时）；很大 = 两者不在同一个状态上。
+   gnorm 必须有限且 > 0。trainer 自己也拦：恰好 0.0 / NaN ⇒ SystemExit（ZERO-LOSS / gnorm 守卫，从第一步起就拦，
+   OPSD 没有 stage3 那种"第 0 步恰好 0"的配方性质）。
+2. `[model] mask-channel ||W||` 学生与教师都非零（known 掩码通道 = 空间条件；0 = 那个模型对它盲）。
+
+教师冻结 ⇒ loss 缓慢下降是正常的；跌到起点的 15% 以下是 COLLAPSE（学生学会了无视自己的缓存，目标的退化极小值），
+trainer 会在日志里打 `[warn] COLLAPSE?`（不中止）。
+**质量判据不是 loss**：`distill/eval/rollout_eval.py` 在 60+ cell 上把 `dits/dit_step{N}`（上游的部署产物，学生本体）
+与 `ema/ema_step{N}`（我们额外的 fp32 EMA 影子）逐个里程碑与 pre-OPSD 的 `stage3_dmd/ema/ema_step300` 比远端漂移 ——
+OPSD 的全部主张都在 rollout 的远端，那个对比才是交付物。`--teacher_context student` 只是自检（loss 恒 0），run.sh 拒绝用它训练。

backup_2026-09-14/code/distill/arope_dit.py

@@ -0,0 +1,1091 @@
+#!/usr/bin/env python3
+"""蒸馏用的 DiT 包装：`AropeDiT` = DiffSynth `WanModel` + ARoPE（世界坐标 RoPE / known-mask 通道 /
+逐 cell 文本 cross-attn），三种前向模式共用一套权重。
+
+为什么要包一层 nn.Module
+-----------------------
+`actionrope.arope.arope_forward` 是自由函数，直接伸手拿 `dit.patch_embedding` / `dit.time_embedding`
+/ `dit.blocks[i].norm1` / `dit.head` 的参数。stage3 用 FSDP1 + TRANSFORMER_BASED_WRAP(DiTBlock)：
+block 的参数只在 **该 block 自己的 forward() 里** 才被 all-gather，在外面伸手拿到的是分片 ——
+不报错，结果是垃圾（PORT_SPEC risk 2，全项目第一号静默腐化风险）。所以：
+
+1. 整个前向搬进 `AropeDiT.forward`（根模块的 forward，根上的 patch/time/head 参数在这里被 gather）；
+2. block 一律用 `block(...)` 调用以触发 FSDP 的 pre-forward hook —— 为此在 import 时把
+   `diffsynth.models.wan_video_dit.DiTBlock.forward` 打成 `arope_block_forward`
+   （= `actionrope/arope.py:277-318` 的 `_block_forward` 逐行照抄，只把 `self_attn(input_x, freqs)`
+   换成带 attn_mask / KV 缓存的 `arope_self_attn`；`arm` 钩子去掉，ARoPE 没有臂）。
+
+三种前向模式
+-----------
+1. **并行双向**（`causal=False, clean_x=None, kv_cache=None`）—— 与 `arope_forward` **逐位相同**，
+   这是整个移植的回归锚点（`distill/tests/test_bitparity.py`）。
+2. **并行块因果**（`causal=True`）：21 个 cell 之间因果、cell 内空间双向；`clean_x` 非 None 时再加
+   teacher-forcing 双流 `[2S,2S]` 掩码（clean 流 + noisy 流；noisy cell f 只看自己 + **严格之前** 的
+   clean cell，绝不看 clean cell f —— 那就是答案，泄漏 = loss 掉到 0 而 rollout 是坏的）。
+3. **流式单 cell**（`kv_cache` 非 None）：滑窗 KV 缓存。缓存 **RoPE 之前的 K**，读的时候按
+   **窗口锚定** 的坐标重新 RoPE（时间位 j-a、世界坐标 offset[j]-offset[a]），驱逐超出 kv_window 的最老 cell。
+
+三条 BLOCKER（PORT_SPEC 对抗验证，违反任何一条都是"能跑、loss 会降、模型是错的"）
+-----------------------------------------------------------------------------
+* **B1 first_frame_cond 绝不能进单 cell 路径。** `actionrope/arope.py:227` 的 `t_tok[:, 0] = 0` 抹的是
+  "这个张量的第 0 帧"；单 cell 前向里 f==1，于是 rollout 的每一个 cell 都被告知 t=0。
+  `forward()` 在 kv_cache 非 None 时硬断言 `first_frame_cond is False`，`forward_frame` 无条件按 `t_k` 建时间步。
+* **B2 build_world_freqs 需要 f_start。** 上游签名固定查 `f_table[:f]`，单 cell 永远拿到时间位 0。
+  这里给出 distill 本地的 `build_world_freqs(..., f_start=0)` 变体（**不改** `actionrope/arope.py`），
+  f_start=0 时与上游逐位相同。
+* **B3 逐 cell 时间步。** 上游 `_per_token_timestep` 断言 `timestep.numel() == b`，stage1 要传 `[b, 21]`。
+  这里的 `_per_token_timestep_pf` 接受 标量 / `[b]` / `[b, F]`，前两种与上游逐位相同。
+
+flow map（stage2-flowmap）：第二个时间步 r
+----------------------------------------
+`install_flowmap()` 之后模型学的是 **均速度** `u(x_t, t, r)`，`forward(..., r_timestep=...)` /
+`forward_frame(..., r_k=...)` 多收一个 r。默认（`residual` + 读出层零初始化）**对任意 (t, r)**
+都与装之前逐位相同，所以它同时是 flow map 的初值和回归锚点。没装的时候 r 这条路径完全不存在，
+`r_timestep` 只准是 None 或与 `timestep` 相同 —— 上游那个不带 delta 分支的 embedder 会 **收下 r
+然后忽略它**（transformer_far_wan_model.py:241），于是接错线会退化成普通扩散模型且不报错。
+新增的 5 个键（`delta_embedding.{0,2}.{weight,bias}` + `delta_gate`）挂在 **dit 上**，827 → 832。
+
+另外两条硬规矩
+-------------
+* **绝不对 KV rollout 做梯度检查点**（risk 6）：滑窗 commit 是破坏性的（cat + 驱逐），checkpoint 重算时
+  读到的是 **驱逐之后** 的缓存 —— 缓存没满时报 `CheckpointError`，满了以后形状对、内容错、**不报错**。
+  `forward()` 在 kv_cache 非 None 时断言 `use_gradient_checkpointing` 为假；`forward_frame` 里 block 直接调用。
+* **注意力只走 SDPA**（risk 8）：import 时断言 flash-attn / sageattention 都没装 —— 装了的话
+  `flash_attention()` 会改走别的 kernel，`arope_forward` 与本文件就悄悄分叉，回归锚点失效。
+"""
+from __future__ import annotations
+
+import copy
+import functools
+import os
+import sys
+
+os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from diffsynth.core.gradient import gradient_checkpoint_forward
+from diffsynth.models import wan_video_dit as _wvd
+from diffsynth.models.wan_video_dit import WanModel, modulate, rope_apply, sinusoidal_embedding_1d
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+# _axis_freqs / _embed_context 直接复用上游实现：逐位等价是硬要求，不能各写一份
+from actionrope.arope import (  # noqa: E402
+    MASK_EMBEDDING_NAME,
+    PX_PER_TOKEN,
+    _axis_freqs,
+    _embed_context,
+    install_arope,
+)
+
+MODEL_DIR = os.path.join(ROOT, "models/Wan2.2-TI2V-5B")
+DIT_FILES = [os.path.join(MODEL_DIR, f"diffusion_pytorch_model-0000{i}-of-00003.safetensors") for i in (1, 2, 3)]
+LAT_H, LAT_W = 30, 52
+
+# ---- flow map（stage2-flowmap）：第二个时间步 r 的分支 ----------------------------------
+# 两个名字都挂在 **dit 上**（`dit.add_module` / `dit.register_buffer`），不是挂在 AropeDiT 包装上：
+# runlib.strip_dit_prefix 只认 `dit.` / `arm.` 前缀，挂在包装上会被当成"包装层自己的 buffer"静默丢掉。
+# 这是 actionrope/arope.py:196-206 给 arope_mask_embedding 用的同一条先例。
+DELTA_EMBEDDING_NAME = "delta_embedding"
+DELTA_GATE_NAME = "delta_gate"
+FLOWMAP_FORMS = ("residual", "convex")
+
+# risk 8：本文件的自注意力恒走 F.scaled_dot_product_attention；上游 flash_attention() 在这三个标志
+# 全为 False 时也落到 SDPA，两条路径才逐位相同。哪天有人 pip install flash-attn，这个断言先炸，
+# 而不是让 test_bitparity 之外的所有下游悄悄测在另一个模型上。
+assert not (_wvd.FLASH_ATTN_3_AVAILABLE or _wvd.FLASH_ATTN_2_AVAILABLE or _wvd.SAGE_ATTN_AVAILABLE), (
+    "检测到 flash_attn / sageattention：DiffSynth 的 flash_attention() 会改走它们，"
+    "而 distill/arope_dit.py 恒走 SDPA，两条路径不再逐位相同。卸载它们，或先修 arope_self_attn。"
+)
+
+
+# --------------------------------------------------------------------------
+# 世界坐标 RoPE（distill 本地变体：多一个 f_start）
+# --------------------------------------------------------------------------
+
+def rot2d(vec: torch.Tensor, ang: torch.Tensor) -> torch.Tensor:
+    """把 (dx, dy)（屏幕 u 向右、v 向下，单位任意）绕原点转 ang 弧度。vec [..., 2]、ang 可广播到 [...]。
+
+    u' = u·cos − v·sin，v' = u·sin + v·cos。v 轴朝下 ⇒ 正角在屏幕上是顺时针，纯约定。
+    """
+    c, s_ = torch.cos(ang), torch.sin(ang)
+    u, v = vec[..., 0], vec[..., 1]
+    return torch.stack([u * c - v * s_, u * s_ + v * c], dim=-1)
+
+
+def world_pos_maps(offset_tok: torch.Tensor, angle, h: int, w: int, scale=None):
+    """逐 token 的世界坐标 (pos_h, pos_w)，各 [b, f, h, w]。
+
+    cell j 的 token (i_h, j_w) 在**当前相机系**里的位置 = 绕画幅中心转 angle[b,j] 之后再加该 cell 的偏移：
+        (u, v) = R(angle) · (j_w − ctr_w, i_h − ctr_h) + (dx, dy) + (ctr_w, ctr_h)
+    angle 全 0 ⇒ 退化成可分离的 pos_h = i + dy / pos_w = j + dx（与旋转前的实现逐位相同）。
+    """
+    b, f = offset_tok.shape[:2]
+    dev, dt = offset_tok.device, torch.float64
+    gi = torch.arange(h, device=dev, dtype=dt).view(1, 1, h, 1).expand(b, f, h, w)
+    gj = torch.arange(w, device=dev, dtype=dt).view(1, 1, 1, w).expand(b, f, h, w)
+    ctr_h, ctr_w = (h - 1) / 2.0, (w - 1) / 2.0
+    local = torch.stack([gj - ctr_w, gi - ctr_h], dim=-1)              # [b, f, h, w, 2] = (u, v)
+    ang = angle.to(device=dev, dtype=dt).view(b, f, 1, 1)
+    rot = rot2d(local, ang)
+    if scale is not None:
+        # 位置整体等比缩放（判决实验）：位置铺得越开 ⇒ 这一帧名义上覆盖的世界越大。
+        # 若模型把"位置场变大"读成"视野变大"，画面内容就该按 1/scale 缩小。
+        rot = rot * torch.as_tensor(scale, device=dev, dtype=dt).view(b, f, 1, 1, 1)
+    off = offset_tok.to(dt).view(b, f, 1, 1, 2)
+    pos_w = rot[..., 0] + off[..., 0] + ctr_w
+    pos_h = rot[..., 1] + off[..., 1] + ctr_h
+    return pos_h, pos_w
+
+
+def build_world_freqs(dit: WanModel, f: int, h: int, w: int, offset_tok=None, device=None,
+                      f_start: int = 0, angle=None, scale=None) -> torch.Tensor:
+    """`actionrope.arope.build_world_freqs` + `f_start`（BLOCKER B2）。
+
+    f_start：这一段 token 的时间轴 **绝对起点**，查 `f_table[f_start:f_start+f]`。滑窗 KV 路径靠它把
+    单个 cell 放到时间位 tpos（否则 f==1 永远拿到位置 0，整条历史与当前帧的相对时间全塌成 0）。
+    f_start=0 时与上游函数逐位相同（`f_table[0:f]` 就是 `f_table[:f]`）。
+
+    offset_tok: [b, f, 2] = (dx_tok, dy_tok)，token 单位；None ⇒ 普通查表 RoPE，返回 [f*h*w, 1, 64]，
+    否则返回 [b, f*h*w, 1, 64]。dtype 跟随 dit.freqs（complex128）。
+
+    angle: [b, f] 弧度，每个 cell 的相机朝向**相对当前 cell**（当前 cell 恒为 0）。None / 全 0 ⇒ 走可分离
+    路径（逐位等于旋转支持加入之前的实现）。非零时改走逐 token 的位置图：h/w 两轴被耦合，位置不再可分离。
+    RoPE 本身对小数位置是解析的（`polar(1, pos·inv_freq)`），所以旋转出来的小数坐标天然合法；
+    模型没见过的是"位移随 token 位置变化"这种场，而不是某个没见过的数值。
+    """
+    f_table, h_table, w_table = dit.freqs
+    if device is None:
+        device = offset_tok.device if offset_tok is not None else f_table.device
+    d_h, d_w = h_table.shape[1] * 2, w_table.shape[1] * 2  # 42, 42
+    assert f_start >= 0 and f_start + f <= f_table.shape[0], (
+        f"时间轴 RoPE 表只有 {f_table.shape[0]} 行，要不到 [{f_start}, {f_start + f})"
+    )
+    ft = f_table[f_start:f_start + f]
+
+    if offset_tok is None:
+        # 走与原函数一样的查表，保证等价性；不重新算 h/w 轴
+        return torch.cat([
+            ft.view(f, 1, 1, -1).expand(f, h, w, -1),
+            h_table[:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            w_table[:w].view(1, 1, w, -1).expand(f, h, w, -1),
+        ], dim=-1).reshape(f * h * w, 1, -1).to(device)
+
+    offset_tok = offset_tok.to(device=device, dtype=torch.float64)
+    b = offset_tok.shape[0]
+    assert offset_tok.shape == (b, f, 2), f"offset_tok 形状应为 [b, {f}, 2]，得到 {tuple(offset_tok.shape)}"
+    dx = offset_tok[..., 0]  # [b, f]
+    dy = offset_tok[..., 1]
+    has_rot = angle is not None and bool(torch.any(torch.as_tensor(angle) != 0))
+    has_scale = scale is not None and bool(torch.any(torch.as_tensor(scale) != 1))
+    if has_rot or has_scale:
+        ang = (torch.zeros(b, f, device=device, dtype=torch.float64) if angle is None
+               else torch.as_tensor(angle, device=device, dtype=torch.float64))
+        assert ang.shape == (b, f), f"angle 形状应为 [b={b}, f={f}]，得到 {tuple(ang.shape)}"
+        pos_h, pos_w = world_pos_maps(offset_tok, ang, h, w, scale)      # 各 [b, f, h, w]
+        fh = _axis_freqs(pos_h, d_h)                                    # [b, f, h, w, 21]
+        fw = _axis_freqs(pos_w, d_w)
+        freqs = torch.cat([
+            ft.to(device).view(1, f, 1, 1, -1).expand(b, f, h, w, -1), fh, fw,
+        ], dim=-1)
+        return freqs.reshape(b, f * h * w, 1, -1).to(f_table.dtype)
+
+    pos_h = torch.arange(h, device=device, dtype=torch.float64).view(1, 1, h) + dy.view(b, f, 1)  # [b, f, h]
+    pos_w = torch.arange(w, device=device, dtype=torch.float64).view(1, 1, w) + dx.view(b, f, 1)  # [b, f, w]
+    fh = _axis_freqs(pos_h, d_h)  # [b, f, h, 21]
+    fw = _axis_freqs(pos_w, d_w)  # [b, f, w, 21]
+    freqs = torch.cat([
+        ft.to(device).view(1, f, 1, 1, -1).expand(b, f, h, w, -1),
+        fh.view(b, f, h, 1, -1).expand(b, f, h, w, -1),
+        fw.view(b, f, 1, w, -1).expand(b, f, h, w, -1),
+    ], dim=-1)
+    return freqs.reshape(b, f * h * w, 1, -1).to(f_table.dtype)
+
+
+def _freqs_cat_dim(freqs: torch.Tensor) -> int:
+    """freqs 的 token 轴：无 batch 的 [S,1,64] 是 0，逐样本的 [b,S,1,64] 是 1。"""
+    assert freqs.dim() in (3, 4), f"freqs 应为 [S,1,64] 或 [b,S,1,64]，得到 {tuple(freqs.shape)}"
+    return 0 if freqs.dim() == 3 else 1
+
+
+# --------------------------------------------------------------------------
+# 时间步：标量 / [b] / [b, F]
+# --------------------------------------------------------------------------
+
+def _per_token_timestep_pf(latents: torch.Tensor, timestep, n_tok_per_frame: int,
+                           first_frame_cond: bool) -> torch.Tensor:
+    """逐 token 时间步。标量 / [b] 与 `actionrope.arope._per_token_timestep` 逐位相同；
+    多出的 `[b, F]` 分支是 stage1 的逐 cell 加噪要的（BLOCKER B3）。
+
+    first_frame_cond 在展开之后才作用（把 cell 0 的 t 置 0），所以 `[b, F]` 传进来时
+    cell 0 那一列会被覆盖成 0 —— 这正是 ActionRoPE 的 TI2V 契约（cell 0 是干净锚点）。
+    返回 [b, f * n_tok_per_frame]，dtype 跟 latents。
+    """
+    b, f = latents.shape[0], latents.shape[2]
+    timestep = torch.as_tensor(timestep, device=latents.device).to(latents.dtype)
+    # 上游 _per_token_timestep 无条件先 reshape(-1)，所以 [b,1] / [1,1] 在它那里是合法的"逐样本"写法。
+    # 只在形状 **不可能** 是逐 cell（不是 [b,f]）且元素数为 1 或 b 时压平：接受的输入集合是上游的严格超集，
+    # 且 f==1 时两种解释本来就同解，不存在歧义。
+    if timestep.dim() > 1 and tuple(timestep.shape) != (b, f) and timestep.numel() in (1, b):
+        timestep = timestep.reshape(-1)
+    if timestep.dim() <= 1:
+        timestep = timestep.reshape(-1)
+        if timestep.numel() == 1:
+            timestep = timestep.expand(b)
+        assert timestep.numel() == b, f"timestep 应为标量 / [b={b}] / [b={b}, f={f}]，得到 {timestep.numel()} 个"
+        t_tok = timestep.view(b, 1, 1).expand(b, f, n_tok_per_frame).clone()
+    else:
+        assert timestep.shape == (b, f), f"逐 cell timestep 形状应为 [b={b}, f={f}]，得到 {tuple(timestep.shape)}"
+        t_tok = timestep.view(b, f, 1).expand(b, f, n_tok_per_frame).clone()
+    if first_frame_cond:
+        t_tok[:, 0] = 0
+    return t_tok.reshape(b, f * n_tok_per_frame)
+
+
+def _ts_equal(a, b) -> bool:
+    """两个时间步参数（标量 / [b] / [b,F] / python float）在数值上是否完全相同。
+
+    只在"模型没装 flow map 却收到了 r"这条断言路径上用（正常训练/推理都不走），所以允许它做
+    一次 device→CPU 的同步。同一个对象直接返回 True，不做任何拷贝。
+    """
+    if a is b:
+        return True
+    ta = torch.as_tensor(a).detach().to(device="cpu", dtype=torch.float64).reshape(-1)
+    tb = torch.as_tensor(b).detach().to(device="cpu", dtype=torch.float64).reshape(-1)
+    if ta.numel() == 1 and tb.numel() > 1:
+        ta = ta.expand_as(tb)
+    elif tb.numel() == 1 and ta.numel() > 1:
+        tb = tb.expand_as(ta)
+    return ta.shape == tb.shape and bool(torch.equal(ta, tb))
+
+
+# --------------------------------------------------------------------------
+# 注意力掩码
+# --------------------------------------------------------------------------
+
+class _KVRef:
+    """指回 KV cache 本体的自引用句柄 —— 用来绕开 FSDP 对 forward 输入容器的重建。
+
+    ★ 这是一个必须记住的 FSDP 行为：混合精度下 FSDP 会 cast forward 的输入，
+    `_apply_to_tensors` 会**遍历并重建** dict / list / tuple，于是 `forward_frame` 写进去的
+    K/V 落在一份**拷贝**上，调用方的 cache 永远是空的 —— 而且不报错（2 卡实测：
+    `外层 dict is 内层 dict → False`，commit 之后外层 cell 数恒为 0）。
+
+    但 FSDP 只 cast 张量，**非张量值按引用原样保留**（同一实测：放进去的哨兵对象 `is` 相同）。
+    所以在 cache 里塞一个指回自己的普通对象，就能在 forward 里把原始 dict 找回来。
+
+    为什么不改成"把 cache 挂成模块属性"：调用必须走包装模块的 forward（参数 all-gather 挂在
+    那里，直接调 forward_frame 拿到的是分片、同样不报错），而 34 处调用点与全部测试都按
+    `kv_cache=` 这个 kwarg 写的。用自引用句柄两边都不用动。
+    """
+
+    __slots__ = ("cache",)
+
+    def __init__(self, cache: dict):
+        self.cache = cache
+
+
+def build_causal_mask(f: int, hw: int, device, local: int = 0, sink: int = 0,
+                      drop_anchor: bool = False) -> torch.Tensor:
+    """块因果 [S, S] bool（True=可见）：cell 之间因果，cell 内空间双向。S = f*hw。
+
+    local>0 ⇒ 只看最近 local 个 cell（外加 sink 个恒可见的开头 cell），对角线恒保留。
+    stage1/2 都是 local=0 —— F=21 下"窗口 20 + sink 1"与全因果等价，而 sink=0 的 W=20 会把
+    cell 0（ActionRoPE 的干净锚点）从 cell 20 的视野里踢掉，是致命的。
+    """
+    fidx = torch.arange(f * hw, device=device) // hw  # [S] 每个 token 属于哪个 cell
+    causal = fidx.unsqueeze(1) >= fidx.unsqueeze(0)
+    if drop_anchor:      # 见 build_tf_mask 的同名参数
+        causal = causal & ((fidx.unsqueeze(0) != 0) | (fidx.unsqueeze(1) == 0))
+    if local and local > 0:
+        in_win = fidx.unsqueeze(0) > (fidx.unsqueeze(1) - local)
+        if sink and sink > 0:
+            in_win = in_win | (fidx.unsqueeze(0) < sink)
+        return (causal & in_win) | (fidx.unsqueeze(1) == fidx.unsqueeze(0))
+    return causal
+
+
+def build_tf_mask(f: int, hw: int, device, local: int = 0, sink: int = 0,
+                  drop_anchor: bool = False) -> torch.Tensor:
+    """teacher-forcing 双流掩码 [2S, 2S] bool。token [0,S) = clean 流，[S,2S) = noisy 流。
+
+    * clean q：在 clean 半流内块因果（看 clean 0..f）。
+    * noisy q（cell f）：看自己那一个 noisy cell + **严格之前** 的 clean cell（fk < fq）。
+      绝不看 clean cell f（那就是答案），也绝不看别的 noisy cell。
+    对角线（eye）恒保留。
+
+    local>0 时同时约束 **两个半流**——上游 `gamemaster_dit.py:546 vs :549-554` 只约束了 noisy_ctx，
+    clean 半流仍是全块因果，于是窗口化训练会从 clean 半流泄漏无界历史。这里顺手修掉（虽然我们不用窗口）。
+    """
+    S = f * hw
+    tok = torch.arange(2 * S, device=device)
+    cln = tok < S
+    fr = torch.where(cln, tok // hw, (tok - S) // hw)  # 两个半流各自的 cell 序号 0..f-1
+    cq, ck = cln[:, None], cln[None, :]
+    fq, fk = fr[:, None], fr[None, :]
+    clean_mask = cq & ck & (fk <= fq)
+    noise_self = (~cq) & (~ck) & (fq == fk)
+    noise_ctx = (~cq) & ck & (fk < fq)
+    # drop_anchor：把 clean cell 0（TI2V 锚点）从**每一个** noisy cell 的视野里摘掉，其余历史照旧。
+    # 这是「只 drop 场景参考、保留历史」那条 null —— 与 local/sink 拼不出来：local=F−1 sink=0 只能
+    # 把 cell 0 挤出**最后一格**的窗口，cell 19 的 20 格窗口里它还在。
+    # Δv_anchor = v(全) − v(drop_anchor) 才是"锚点相对于历史的话语权"，正是外观漂移要压的那一项。
+    if drop_anchor:
+        noise_ctx = noise_ctx & (fk != 0)
+    if local and local > 0:
+        in_win = fk > (fq - local)
+        if sink and sink > 0:
+            in_win = in_win | (fk < sink)
+        noise_ctx = noise_ctx & in_win
+        clean_mask = clean_mask & in_win
+    eye = tok[:, None] == tok[None, :]
+    return eye | clean_mask | noise_self | noise_ctx
+
+
+# --------------------------------------------------------------------------
+# 自注意力：并行（可带 mask） / 流式（滑窗 KV 缓存）
+# --------------------------------------------------------------------------
+
+def _heads(x: torch.Tensor, n: int) -> torch.Tensor:
+    return rearrange(x, "b s (n d) -> b n s d", n=n)
+
+
+def _unheads(x: torch.Tensor, n: int) -> torch.Tensor:
+    return rearrange(x, "b n s d -> b s (n d)", n=n)
+
+
+def _evict(t: torch.Tensor, kv: dict) -> torch.Tensor:
+    """滑窗驱逐：缓存超过 sink + kv_window 个 cell 就丢掉最老的 **非 sink** cell。t: [b, n*hw, dim]。"""
+    kw = kv["kv_window"]
+    if kw is None:
+        return t
+    hw, sink = kv["hw"], kv["sink_size"]
+    while t.shape[1] // hw > sink + kw:
+        t = torch.cat([t[:, :sink * hw], t[:, (sink + 1) * hw:]], dim=1)
+    return t
+
+
+def arope_self_attn(sa, x: torch.Tensor, freqs: torch.Tensor, attn_mask=None, kv=None) -> torch.Tensor:
+    """`SelfAttention.forward` 的替身。attn_mask=None 且 kv=None 时与上游 **逐位相同**
+    （上游 `flash_attention()` 在没装 flash/sage 时同样落到 SDPA，见文件头的 import 断言）。
+
+    kv 非 None ⇒ 流式：缓存 **RoPE 之前** 的 K，读的时候按窗口锚定坐标重算
+    （`kv["freqs_hist"]` 给历史、`kv["freqs_cur"]` 给当前 cell）。RoPE 是相对编码，所以重锚定对
+    注意力分数是**精确**的；但缓存的 K/V 是从带掩码通道的 x 算出来的，掩码按各 cell **提交时**的
+    锚点固化，k>20 之后的历史与一段锚在 k−20 的训练片段并不逐位相同（stage2→stage3 输入移位，
+    自强迫训练穿过，GEO-2）。
+    """
+    n = sa.num_heads
+    if kv is None:
+        q = rope_apply(sa.norm_q(sa.q(x)), freqs, n)
+        k = rope_apply(sa.norm_k(sa.k(x)), freqs, n)
+        v = sa.v(x)
+        o = F.scaled_dot_product_attention(_heads(q, n), _heads(k, n), _heads(v, n), attn_mask=attn_mask)
+        return sa.o(_unheads(o, n))
+
+    # ---- 流式单 cell ----
+    li = kv["layer"]
+    k_raw = sa.norm_k(sa.k(x))          # 进缓存的是这个（RoPE 之前）
+    v = sa.v(x)
+    q = rope_apply(sa.norm_q(sa.q(x)), kv["freqs_cur"], n)
+    k_hist, v_hist = kv["k"][li], kv["v"][li]
+    if k_hist is None:
+        k_cat, v_cat = rope_apply(k_raw, kv["freqs_cur"], n), v
+    else:
+        fh = kv["freqs_hist"]
+        assert fh is not None and fh.shape[_freqs_cat_dim(fh)] == k_hist.shape[1], (
+            f"历史 freqs 与缓存长度不符：freqs {tuple(fh.shape)} vs k_hist {tuple(k_hist.shape)}"
+        )
+        k_cat = torch.cat([rope_apply(k_hist, fh, n), rope_apply(k_raw, kv["freqs_cur"], n)], dim=1)
+        v_cat = torch.cat([v_hist, v], dim=1)
+    o = F.scaled_dot_product_attention(_heads(q, n), _heads(k_cat, n), _heads(v_cat, n), attn_mask=None)
+    if kv["commit"]:
+        nk = k_raw if k_hist is None else torch.cat([k_hist, k_raw], dim=1)
+        nv = v if v_hist is None else torch.cat([v_hist, v], dim=1)
+        # detach：缓存是跨帧的常量，梯度绝不能穿过它（stage3 的 streamed replay 正是靠这一点才是精确的）
+        kv["k"][li] = _evict(nk, kv).detach()
+        kv["v"][li] = _evict(nv, kv).detach()
+    return sa.o(_unheads(o, n))
+
+
+# --------------------------------------------------------------------------
+# DiTBlock.forward 的猴补丁（FSDP all-gather 必须经过 block(...) 这个入口）
+# --------------------------------------------------------------------------
+
+def arope_block_forward(self, x, context, t_mod, freqs, *, f: int | None = None,
+                        attn_mask=None, kv=None):
+    """= `actionrope/arope.py:277-318` 的 `_block_forward`（arm 去掉），唯一改动：
+    `self.self_attn(input_x, freqs)` → `arope_self_attn(self.self_attn, input_x, freqs, attn_mask, kv)`。
+
+    f：逐 cell 文本 cross-attn 的 cell 数（TF 双流时是 2*Fl）。context 是整条 [b,L,dim] 时用不到，
+    默认 None —— 这样 DiffSynth 自己的 `model_fn_wan_video`（`block(x, context, t_mod, freqs)`，
+    四个位置参数、整条 context）在打完补丁之后照样能跑，且逐位不变。
+    """
+    has_seq = len(t_mod.shape) == 4
+    chunk_dim = 2 if has_seq else 1
+    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+        self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=chunk_dim)
+    if has_seq:
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2),
+            shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2),
+        )
+    input_x = modulate(self.norm1(x), shift_msa, scale_msa)
+    x = self.gate(x, gate_msa, arope_self_attn(self.self_attn, input_x, freqs, attn_mask, kv))
+
+    q_in = self.norm3(x)
+    if context.shape[0] == x.shape[0]:
+        x = x + self.cross_attn(q_in, context)
+    else:
+        b = x.shape[0]
+        assert f is not None, "逐 cell context 需要显式传 f（block 的 cell 数）"
+        assert context.shape[0] == b * f, f"逐 cell context 应有 b*f={b * f} 行，得到 {context.shape[0]}"
+        q_in = rearrange(q_in, "b (f n) c -> (b f) n c", f=f)
+        y = self.cross_attn(q_in, context)
+        x = x + rearrange(y, "(b f) n c -> b (f n) c", f=f)
+
+    input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
+    x = self.gate(x, gate_mlp, self.ffn(input_x))
+    return x
+
+
+# 全局、一次性：所有 DiTBlock 实例（含 DiffSynth 自己 pipeline 里的）都换成上面这个 forward。
+_wvd.DiTBlock.forward = arope_block_forward
+
+
+# --------------------------------------------------------------------------
+# 模型加载
+# --------------------------------------------------------------------------
+
+def load_dit(model_dir: str = MODEL_DIR, device: str = "cpu") -> WanModel:
+    """只加载 DiT（不挂 T5 / VAE），bf16，先落 CPU —— 与 `actionrope/train.py:163-174` 一致。"""
+    from diffsynth.pipelines.wan_video import ModelConfig, WanVideoPipeline
+
+    files = [os.path.join(model_dir, f"diffusion_pytorch_model-0000{i}-of-00003.safetensors") for i in (1, 2, 3)]
+    pipe = WanVideoPipeline.from_pretrained(
+        torch_dtype=torch.bfloat16, device=device,
+        model_configs=[ModelConfig(path=files)],
+        tokenizer_config=None, redirect_common_files=False,
+    )
+    return pipe.dit
+
+
+def unwrap(model: nn.Module) -> "AropeDiT":
+    """拿到裸的 AropeDiT（DeepSpeed / FSDP / DDP 包装之后 `init_kv_cache`、`dit.freqs` 都要走裸模块）。"""
+    while hasattr(model, "module"):
+        model = model.module
+    return model
+
+
+# --------------------------------------------------------------------------
+# AropeDiT
+# --------------------------------------------------------------------------
+
+class AropeDiT(nn.Module):
+    """WanModel + ARoPE 的包装。三种前向模式见文件头。
+
+    causal：False = 双向教师 / critic；True = 因果学生（块因果掩码 + 允许 TF 双流 + 允许流式 KV）。
+    train_local_attn_size / train_sink_size：窗口化训练（stage1/2 都是 0，即全因果）。
+    """
+
+    def __init__(self, model_dir: str = MODEL_DIR, mask_channel: bool = True, causal: bool = False,
+                 train_local_attn_size: int = 0, train_sink_size: int = 0, dit: WanModel | None = None,
+                 flowmap_form: str | None = None):
+        super().__init__()
+        if dit is None:
+            dit = load_dit(model_dir, device="cpu")
+        # mask 通道零初始化：装上的瞬间输出与原模型完全相同
+        self.dit = install_arope(dit, mask_channel=mask_channel)
+        self.mask_channel = bool(mask_channel)
+        self.causal = bool(causal)
+        self.train_local_attn_size = int(train_local_attn_size)
+        self.train_sink_size = int(train_sink_size)
+        # 训练期把 cell 0（TI2V 锚点）从 noisy 的视野里摘掉，其余历史保留 —— Δv_anchor 的 null
+        self.train_drop_anchor = False
+        # flow map（stage2-flowmap）的第二时间步分支：None = 没装（默认，与装之前逐位相同）。
+        # 装了以后是 "residual" / "convex"。__init__ **绝不自己装**，只有 install_flowmap() 会。
+        # 传进来的 dit 已经带着 delta 分支（deepcopy / 再包一层 / 从 832 键装载过）时必须显式说明是
+        # 哪一种形式：两种形式的融合公式不同，猜错了就是"能跑、数值是另一个模型"。
+        assert flowmap_form is None or flowmap_form in FLOWMAP_FORMS, (
+            f"flowmap_form 只能是 {FLOWMAP_FORMS} 或 None，得到 {flowmap_form!r}")
+        if hasattr(self.dit, DELTA_EMBEDDING_NAME):
+            assert flowmap_form is not None, (
+                "传进来的 dit 上已经有 delta_embedding 了，请显式给 AropeDiT(..., flowmap_form=...)："
+                "residual 与 convex 的融合公式不同，猜错了不会报错，只会安静地跑成另一个模型。")
+        self.flowmap_form: str | None = flowmap_form
+        # 模块级默认；每次 forward 还可以单独传。KV rollout 上两者都必须是假（risk 6）。
+        self.use_gradient_checkpointing = False
+        self.use_gradient_checkpointing_offload = False
+        # [16380,16380] 的 bool 掩码 268 MB，按 (f, hw, tf, 窗口, device) 缓存，别每步重建
+        self._mask_cache: dict = {}
+
+    # ---- 小工具 ----
+
+    def move_rope_table(self, device) -> "AropeDiT":
+        """`dit.freqs` 是普通 tuple 属性（不是 buffer），`.to(device)` / prepare 都搬不动它。
+        DeepSpeed / FSDP prepare 之���必须在 **裸模块** 上调一次。"""
+        self.dit.freqs = tuple(t.to(device) for t in self.dit.freqs)
+        return self
+
+    def _get_mask(self, f: int, hw: int, device, tf: bool):
+        if not self.causal:
+            return None
+        key = (f, hw, bool(tf), self.train_local_attn_size, self.train_sink_size,
+               bool(self.train_drop_anchor), str(device))
+        m = self._mask_cache.get(key)
+        if m is None:
+            build = build_tf_mask if tf else build_causal_mask
+            m = build(f, hw, device, self.train_local_attn_size, self.train_sink_size,
+                      drop_anchor=bool(self.train_drop_anchor))
+            self._mask_cache[key] = m
+        return m
+
+    def _mask_embed(self, mask_input, ref: torch.Tensor):
+        """mask 通道的 embedding。ref 是 patch_embedding 的 **输出**（不是权重）——
+        autocast 下输出是 bf16 而权重可能是 fp32，拿错参照物就与 arope.py:372 不再逐位相同。"""
+        if mask_input is None:
+            return None
+        if not hasattr(self.dit, MASK_EMBEDDING_NAME):
+            raise ValueError("给了 mask_input 但 dit 没有 arope_mask_embedding，请先 install_arope(dit, mask_channel=True)")
+        return getattr(self.dit, MASK_EMBEDDING_NAME)(mask_input.to(device=ref.device, dtype=ref.dtype))
+
+    # ---- flow map：第二个时间步 r ----
+
+    @property
+    def has_flowmap(self) -> bool:
+        """装了 delta 分支没有。**以 dit 上的模块为准**，不看包装上的 flowmap_form ——
+        续训是"先 install 再 strict load"，权重里有没有这五个键才是唯一事实。"""
+        return hasattr(self.dit, DELTA_EMBEDDING_NAME)
+
+    def install_flowmap(self, form: str = "residual", gate: float = 1.0) -> "AropeDiT":
+        """装上第二时间步 r 的 delta 分支。**绝不从 __init__ 调**（见下面的初始化顺序）。
+
+        form="residual"（默认，本项目用的）
+            t = E(t) + g · D(r)，`D = deepcopy(E)` 但把 **第二个 Linear** 的 weight/bias 清零。
+            于是 D(·) ≡ 0，对 **任意** (t, r) 都与装之前逐位相同 —— 这是硬要求：我们 50 % 的样本
+            按构造就在 r ≠ t 的off-diagonal 上，而初值是 24 GPU-hours 换来的 stage1b 教师。
+            只清 **读出层**、保留第一层学好的 256→3072：分支一上来就有 r 的有意义隐表示，
+            ∂L/∂delta_embedding[2] 在第 0 步就非零，立刻开始学（arope_mask_embedding 零初始化 Conv3d 的同一套路）。
+        form="convex"（上游 AnyFlow，transformer_far_wan_model.py:330）
+            t = (1−g)·E(t) + g·D(r)，D 是 E 的原样 deepcopy。**只在 r == t 那条对角线上**与教师相同；
+            off-diagonal 是两个不同时间嵌入的 0.75/0.25 混合（CPU 实测 max|Δ| ≈ 1.1e-01）。
+            只留作"上游保真度"消融，别指望它是恒等的。
+            ⚠ 千万别在 convex 里"顺手零初始化"：那给出 t = 0.75·E(t)，把教师的时间嵌入缩水 25 %，
+            比两种做法都糟，却长得像一个安全的零初始化。
+
+        gate：`dit.delta_gate`，**persistent=True 的 buffer**，跟着权重走。上游是 persistent=False
+        （从 config 重建），于是训练与推理可以悄悄用不同的门而不报错。residual 形式下门乘的是一个
+        本来就是 0 的读出，第 0 步冗余，之后只等价于给 delta_embedding[2] 换一个有效 lr。
+
+        导出格式（runlib）：两个名字都挂在 dit 上 ⇒ state_dict 里是 `dit.delta_embedding.{0,2}.{weight,bias}`
+        和 `dit.delta_gate`，strip_dit_prefix 剥掉 `dit.` 之后正好是 5 个新键，827 → 832
+        （runlib.AROPE_DIT_KEYS_FLOWMAP / runlib.FLOWMAP_KEYS）。挂在 AropeDiT 包装上会被静默丢掉。
+
+        初始化顺序（写错了不会报错，会静默丢掉教师）：
+            新 run   : AropeDiT(...) → load_dit_ckpt(m, stage1b)【strict 827】→ install_flowmap()
+            续训     : AropeDiT(...) → install_flowmap() → load_dit_ckpt(m, dit_stepN)【strict 832】
+        反过来（先 install 再从 827 的 ckpt 装）会 deepcopy 到 **随机初始化** 的 time_embedding，
+        之后 load 又把 time_embedding 覆盖成教师的 —— delta 分支第一层从此与 time_embedding 不一致。
+        residual 形式下输出仍然是恒等的（读出是 0），但分支的隐表示是垃圾。所以装完立刻自检
+        （见下面的两个 assert），训练脚本再另外跑一次真前向的 max|Δ| == 0.0 闸门。
+        """
+        assert form in FLOWMAP_FORMS, f"flowmap_form 只能是 {FLOWMAP_FORMS}，得到 {form!r}"
+        # ★ convex 形式下门必须落在开区间 (0, 1)。上游 gate_value = 0.25。
+        #   g == 1.0（本方法的 residual 默认值）会让 t = 0·E(t) + 1·D(r) —— 时间嵌入被 r 嵌入**整个替换**，
+        #   模型对 t 再无任何��赖（实测 |u(t=900,r=100) − u(t=100,r=100)| = 0.000e+00）。
+        #   而 r == t 时 D(t) == E(t)（D 是 E 的原样 deepcopy），所以 trainer 那道
+        #   _assert_teacher_identity（convex 只查对角线）会**给出绿灯**。静默、灾难性，
+        #   只要有人写了 `--flowmap_form convex` 而忘了跟上 `--delta_gate 0.25` 就会中招。
+        #   g == 0.0 则是另一个极端：delta 分支恒被乘没，r 完全不起作用，退化成普通扩散模型。
+        if form == "convex":
+            assert 0.0 < float(gate) < 1.0, (
+                f"convex 形式的 delta_gate 必须落在开区间 (0, 1)，得到 {gate}（上游 gate_value = 0.25）。\n"
+                f"  t = (1−g)·E(t) + g·D(r)：g=1 把时间嵌入整个换成 r 嵌入（模型对 t 零依赖），"
+                f"g=0 把 r 分支整个乘没（退化成普通扩散）。两种都能骗过只查 r==t 对角线的同一性闸。\n"
+                f"  想用上游保真度消融请写 --flowmap_form convex --delta_gate 0.25。"
+            )
+        assert not self.has_flowmap, (
+            "已经装过 flow map 了。重复 install 会用当前的 time_embedding 重新 deepcopy 一份，"
+            "把已经训练过的 delta 分支整个丢掉（residual 形式下还会把它清零）—— 静默、无报错。"
+        )
+        dit = self.dit
+        emb = dit.time_embedding
+        assert isinstance(emb, nn.Sequential) and len(emb) == 3, (
+            f"time_embedding 应是 Sequential(Linear, SiLU, Linear)，得到 {emb}"
+        )
+        delta = copy.deepcopy(emb)
+        if form == "residual":
+            nn.init.zeros_(delta[2].weight)
+            nn.init.zeros_(delta[2].bias)
+        dit.add_module(DELTA_EMBEDDING_NAME, delta)
+
+        ref = emb[0].weight
+        g = torch.tensor([float(gate)], dtype=ref.dtype, device=ref.device)
+        assert float(g.item()) == float(gate), (
+            f"gate={gate} 在 {ref.dtype} 下不可精确表示（存回来是 {float(g.item())}）："
+            "门要跟着 bf16 权重一起存盘，不精确就意味着训练用的门和装载回来的门不是同一个数"
+        )
+        dit.register_buffer(DELTA_GATE_NAME, g, persistent=True)
+        self.flowmap_form = form
+
+        if form == "residual":
+            assert torch.equal(delta[0].weight, emb[0].weight) and torch.equal(delta[0].bias, emb[0].bias), (
+                "delta_embedding 第一层应当是 time_embedding 的逐位副本"
+            )
+            assert float(delta[2].weight.abs().max()) == 0.0 and float(delta[2].bias.abs().max()) == 0.0, (
+                "residual 形式的 delta_embedding 读出层必须是全 0"
+            )
+        return self
+
+    def _embed_time(self, t_tok: torch.Tensor, r_tok: torch.Tensor | None) -> torch.Tensor:
+        """[b, S] 的逐 token 时间步 → [b, S, dim] 的时间嵌入。
+
+        没装 flow map ⇒ 与改动之前 **逐字相同** 的一行（回归锚点靠这个）。
+        装了 ⇒ 按 flowmap_form 融合。三个要点：
+        * 融合发生在 **time_projection 之前**，而 time_projection[0] 就是 SiLU ——
+          `dit.time_projection(blended)` 逐字等于上游的 `time_proj(act_fn(rt_emb))`。
+        * 两个分支共用 **同一个** 正弦编码和同一个 0..1000 量纲。r 用 [0,1] 的 σ 而 t 用 [0,1000]
+          的话，正弦编码会悄悄塌成近似常数。
+        * 门与 delta 分支的输出统一 cast 到 **主分支的 dtype**，绝不让 fp32 的门把 bf16 的和提升成
+          fp32（上游 :326 把门 cast 成第三种 dtype，混合精度提升进 time_proj，bf16 下 r==t 的恒等性
+          在 1e-3 量级就破了）。
+        """
+        dit = self.dit
+        b, S = t_tok.shape
+        e_t = sinusoidal_embedding_1d(dit.freq_dim, t_tok.reshape(-1)).view(b, S, -1)
+        base = dit.time_embedding(e_t)
+        if r_tok is None:
+            return base
+        assert self.flowmap_form in FLOWMAP_FORMS, (
+            f"dit 上有 delta 分支但包装的 flowmap_form 是 {self.flowmap_form!r} —— "
+            "谁把两者拆开了（直接 add_module？绕过 install_flowmap？）")
+        e_r = sinusoidal_embedding_1d(dit.freq_dim, r_tok.reshape(-1)).view(b, S, -1)
+        delta = getattr(dit, DELTA_EMBEDDING_NAME)(e_r).to(base.dtype)
+        g = getattr(dit, DELTA_GATE_NAME).to(base.dtype)
+        if self.flowmap_form == "convex":
+            return (1.0 - g) * base + g * delta
+        return base + g * delta
+
+    def _resolve_r(self, timestep, r_timestep):
+        """r 的准入闸门，返回真正要用的 r（没装 flow map 时恒为 None）。
+
+        上游那个不带 flow map 的 embedder **收下 r_timestep 然后忽略它**
+        （transformer_far_wan_model.py:241），于是"忘了接 delta 分支"会退化成一个普通扩散模型，
+        不报错、loss 照降。这里两边都硬断言：装了就必须给 r；没装就只准不给、或者给一个与 t 相��的 r。
+        """
+        if self.has_flowmap:
+            assert r_timestep is not None, (
+                "模型装了 flow map（dit.delta_embedding 存在），r_timestep / r_k 是必填的。"
+                "扩散槽请显式传 r_timestep=timestep，别省略。"
+            )
+            return r_timestep
+        if r_timestep is None:
+            return None
+        assert _ts_equal(timestep, r_timestep), (
+            "模型没装 flow map（没有 dit.delta_embedding），r_timestep 只能是 None 或与 timestep 完全相同；"
+            "给了一个不同的 r 说明调用方以为自己在跑 flow map —— 那会静默地跑成普通扩散。"
+            "先 install_flowmap()。"
+        )
+        return None
+
+    # ---- 前向 ----
+
+    def forward(self, x: torch.Tensor, timestep, context: torch.Tensor, *,
+                context_ids=None, offset_px=None, mask_input=None, first_frame_cond: bool = True,
+                r_timestep=None, clean_x=None, aug_t=None,
+                frame_index=None, kv_cache=None, commit: bool = True,
+                offset_win=None, mask_in_k=None, angle_win=None, angle_px=None, scale_px=None,
+                use_gradient_checkpointing: bool = False,
+                use_gradient_checkpointing_offload: bool = False) -> torch.Tensor:
+        """x [b, 48, f, 30, 52] → [b, 48, f, 30, 52]（TF 模式只返回 noisy 半流，形状同样是 f 个 cell）。
+
+        并行路径
+            timestep     标量 / [b] / [b, f]（逐 cell）
+            context      [b, 512, 4096]（整条）或 [b, f, 512, 4096]（逐 cell）
+            context_ids  [b, f] 整数，同 id ⇒ 同串，用于逐 cell 去重
+            offset_px    [b, f, 2] px；None ⇒ 普通查表 RoPE
+            mask_input   [b, 1, f, 30, 52] float，known=1 / new=0
+            first_frame_cond  cell 0 是干净锚点，其 timestep 与 r 都置 0
+            r_timestep   flow map 的第二个时间步（0..1000，与 timestep 同量纲、同形状规则）。
+                         装了 flow map ⇒ **必填**（扩散槽显式传 r_timestep=timestep）；
+                         没装 ⇒ 只能是 None 或与 timestep 完全相同。first_frame_cond 对 r **同样**
+                         生效 —— 否则 TI2V 锚点那一格会掉到对角线之外，而那正是最不能变的一格。
+            clean_x      [b, 48, f, 30, 52] 干净历史流 ⇒ 打开 teacher forcing 双流（要求 causal=True）
+            aug_t        clean 半流的时间步（None ⇒ 0，即精确干净）。clean 半流的 r **硬接** r := t
+                         （= 上游把 clean_timestep 同时当两个参数传，trainer:359）：干净流恒在对角线上。
+
+        流式路径（kv_cache 非 None）：转发给 `forward_frame`。这个转发是故意的 ——
+        FSDP 的参数 all-gather 挂在 **包装模块的 forward** 上，**绝不要**在被包装的模型上直接调
+        `forward_frame`。反过来 `init_kv_cache` 必须在裸模块（`unwrap(m)`）上调。
+        """
+        gc = bool(use_gradient_checkpointing or self.use_gradient_checkpointing)
+        gc_off = bool(use_gradient_checkpointing_offload or self.use_gradient_checkpointing_offload)
+        r_eff = self._resolve_r(timestep, r_timestep)   # 没装 flow map ⇒ None，下面整条 r 路径不存在
+
+        if kv_cache is not None:
+            # ★ FSDP 会重建 forward 的输入容器（见 _KVRef 的 docstring）：拿 ref 把**调用方的
+            #   那一份** dict 找回来，否则 commit 写进拷贝、缓存永远是空的且不报错。
+            ref = kv_cache.get("_ref")
+            if ref is not None and ref.cache is not kv_cache:
+                kv_cache = ref.cache
+            # BLOCKER B1：first_frame_cond 抹的是"这个张量的第 0 帧"，单 cell 前向里就是全部
+            assert first_frame_cond is False, (
+                "流式路径必须 first_frame_cond=False：单 cell 前向 f==1，"
+                "arope.py:227 的 t_tok[:,0]=0 会把 rollout 里每一个 cell 的时间步都抹成 0"
+            )
+            assert clean_x is None, "teacher forcing 是训练路径，不能与 KV 缓存同用"
+            assert angle_px is None, "流式路径的相机朝向走 angle_win（[b, n_cached+1]，相对当前 cell）"
+            assert offset_px is None, "流式路径的世界坐标走 offset_win（[b, n_cached+1, 2]，窗口锚定）"
+            assert mask_input is None, "流式路径的 known 掩码走 mask_in_k（[b, 1, 1, 30, 52]）"
+            assert not gc, (
+                "绝不能对 KV rollout 做梯度检查点：滑窗 commit 是破坏性的，重算时读到的是驱逐后的缓存 —— "
+                "缓存满了以后形状对、内容错、不报错（PORT_SPEC risk 6）"
+            )
+            assert frame_index is not None, "流式路径必须给 frame_index（绝对 cell 序号）"
+            return self.forward_frame(x, timestep, context, frame_index, kv_cache, commit,
+                                      offset_win=offset_win, mask_in_k=mask_in_k, context_ids=context_ids,
+                                      r_k=r_eff, angle_win=angle_win)
+
+        dit = self.dit
+        b, _, f_in, lat_h, lat_w = x.shape
+        p_f, p_h, p_w = dit.patch_size
+        n_tok = (lat_h // p_h) * (lat_w // p_w)
+        tf = clean_x is not None
+        if tf:
+            assert self.causal, "teacher forcing 双流只对因果学生有意义（causal=True）"
+            assert clean_x.shape == x.shape, f"clean_x {tuple(clean_x.shape)} 应与 x {tuple(x.shape)} 同形"
+
+        # ---- 时间步：逐 token ----
+        ts_noisy_tok = _per_token_timestep_pf(x, timestep, n_tok, first_frame_cond)  # [b, f*n_tok]
+        # r 与 t 走 **同一个** 展开函数、同一个 first_frame_cond：cell 0 的 t 和 r 一起被抹成 0
+        rs_noisy_tok = (None if r_eff is None
+                        else _per_token_timestep_pf(x, r_eff, n_tok, first_frame_cond))
+        if tf:
+            # clean 半流：aug_t=None ⇒ 全 0（精确干净）。给了 aug_t 就整条按 aug_t，
+            # 不再套 first_frame_cond —— clean cell 0 本来就是锚点，与其它 clean cell 同一个噪声级别。
+            ts_clean_tok = (torch.zeros_like(ts_noisy_tok) if aug_t is None
+                            else _per_token_timestep_pf(clean_x, aug_t, n_tok, False))
+            t_tok = torch.cat([ts_clean_tok, ts_noisy_tok], dim=1)  # [b, 2*f*n_tok]
+            # clean 半流恒在 r == t 的对角线上（--tf_aug > 0 时就是 r_clean = aug_t，不是 0）
+            r_tok = None if rs_noisy_tok is None else torch.cat([ts_clean_tok, rs_noisy_tok], dim=1)
+        else:
+            t_tok = ts_noisy_tok
+            r_tok = rs_noisy_tok
+        t = self._embed_time(t_tok, r_tok)
+        t_mod = dit.time_projection(t).unflatten(2, (6, dit.dim))
+
+        # ---- 文本：TF 时在 4D 层面复制（clean / noisy 两半共用同一条 cell 文本）----
+        if tf and context.dim() == 4:
+            context = torch.cat([context, context], dim=1)                      # [b, 2f, L, 4096]
+            if context_ids is not None:
+                context_ids = torch.cat([context_ids, context_ids], dim=1)      # [b, 2f]
+        ctx = _embed_context(dit, context, context_ids)
+
+        # ---- patch embedding（+ 零初始化的 mask 通道）----
+        xt = dit.patch_embedding(x)
+        mask_emb = self._mask_embed(mask_input, xt)
+        if mask_emb is not None:
+            xt = xt + mask_emb
+        f, h, w = xt.shape[2:]
+        assert f * h * w == ts_noisy_tok.shape[1], (
+            f"patch 之后的 token 数 {f * h * w} 与时间步 {ts_noisy_tok.shape[1]} 不符（temporal patch 必须是 1）"
+        )
+        tokens = rearrange(xt, "b c f h w -> b (f h w) c").contiguous()
+        s_half = f * h * w
+        if tf:
+            # clean 半流用同一份 mask embedding（数值上与重算一遍完全一样）
+            xc = dit.patch_embedding(clean_x)
+            if mask_emb is not None:
+                xc = xc + mask_emb
+            tokens = torch.cat([rearrange(xc, "b c f h w -> b (f h w) c").contiguous(), tokens], dim=1)
+
+        # ---- RoPE：TF 时两半各自从 f_start=0 起算（cat 两份 21 帧的表，而不是造一张 42 帧的表）----
+        if offset_px is None:
+            assert angle_px is None, "给了 angle_px 就必须给 offset_px（坐标要一起旋）"
+            freqs = build_world_freqs(dit, f, h, w, None, device=x.device)
+        else:
+            # angle_px [b, f] 弧度：**绝对**朝向，锚在 cell 0（并行路径里 cell 0 是干净的 TI2V 锚点，
+            # 它的格点必须保持轴对齐）。None ⇒ 逐位走旋转支持加入之前的老路。
+            freqs = build_world_freqs(dit, f, h, w, offset_px.to(torch.float64) / PX_PER_TOKEN,
+                                      device=x.device, angle=angle_px, scale=scale_px)
+        if tf:
+            freqs = torch.cat([freqs, freqs], dim=_freqs_cat_dim(freqs))
+
+        # ---- 掩码 + blocks ----
+        attn_mask = self._get_mask(f, h * w, tokens.device, tf)
+        f_cross = 2 * f if tf else f
+        for block in dit.blocks:
+            # 用 block(...) 而不是伸手拿 block 的子模块：FSDP 的 all-gather 挂在这个入口上
+            tokens = gradient_checkpoint_forward(
+                functools.partial(block, f=f_cross, attn_mask=attn_mask, kv=None),
+                gc, gc_off,
+                tokens, ctx, t_mod, freqs,
+            )
+
+        # ---- head：TF 时只取 noisy 半流（clean 前缀连同它的时间调制一起丢掉）----
+        if tf:
+            tokens = tokens[:, s_half:]
+            t = t[:, s_half:]
+        tokens = dit.head(tokens, t)
+        return dit.unpatchify(tokens, (f, h, w))
+
+    # ---- 滑窗 KV 缓存 ----
+
+    def init_kv_cache(self, kv_window: int | None = 20, sink_size: int = 0, rope_cap: int = 20,
+                      num_layers: int | None = None) -> dict:
+        """每条 rollout 一份新缓存（换 clip 必须重建）。**在裸模块上调**（`unwrap(m).init_kv_cache(...)`）。
+
+        kv_window=None ⇒ 绝对模式：缓存无限增长，时间位就是 frame_index（等价性测试 / 短 rollout 用）。
+        kv_window 给了 ⇒ 滑窗模式：缓存 raw K，读的时候按窗口锚定坐标重算 RoPE，
+        超出 sink + kv_window 就驱逐最老的非 sink cell。
+
+        ActionRoPE 的部署配置是 sink_size=0：WorldMind 的 sink 把开头几帧永久钉在 RoPE 位 0..sink-1，
+        但这里 h/w 的 RoPE 位是 **一直在漂的世界坐标**（实测 21 个 cell 内最大位移 373 px = 11.66 token），
+        钉住的 sink 帧在 rollout 几十个 cell 之后会落在世界坐标几十个 token 之外，模型没见过。
+        sink=0 + 窗口锚定 ⇒ 每个注意力窗口的 RoPE 坐标精确等于一段训练片段的坐标（只解决坐标 OOD；
+        缓存 cell 的掩码通道按提交时锚点固化，见 arope_self_attn；latent 值离流形见 PORT_SPEC R10）。
+        """
+        n = num_layers if num_layers is not None else len(self.dit.blocks)
+        assert self.causal, "KV 缓存是因果学生的流式路径（causal=True）"
+        if kv_window is not None:
+            assert kv_window >= 1 and sink_size >= 0
+            assert rope_cap == sink_size + kv_window, (
+                f"rope_cap 必须 == sink_size + kv_window（{rope_cap} != {sink_size} + {kv_window}）："
+                "时间位 tpos 恒等于窗口内已缓存的 cell 数，超过 cap 就是训练时没见过的相对距离"
+            )
+        cache = {
+            "k": [None] * n, "v": [None] * n, "frames_cached": 0,
+            "kv_window": kv_window, "sink_size": int(sink_size), "rope_cap": int(rope_cap),
+            "hw": None, "layer": None, "commit": False, "freqs_cur": None, "freqs_hist": None,
+        }
+        # 自引用句柄：FSDP 会重建 forward 的输入容器，靠它在 forward 里找回**这一份** dict。
+        # 非张量值 FSDP 按引用保留，所以 ref 穿得过去（见 _KVRef 的 docstring 与 2 卡实测）。
+        cache["_ref"] = _KVRef(cache)
+        return cache
+
+    @staticmethod
+    def release_kv_cache(cache: dict) -> None:
+        """用完一条 rollout 必须调 —— `_ref` 与 cache 互相引用，构成**引用环**。
+
+        引用计数放不掉环，只能等循环 GC，而 GC 是惰性的：一条 21 cell 的 rollout 缓存十几 GB
+        的 K/V，几步之内显存就爬满（实测 98 → 120 → 126 GB 然后 OOM）。
+        `clear()` 把 dict 清空即断环，张量当场释放。
+        """
+        if cache is not None:
+            cache.clear()
+
+    @staticmethod
+    def snapshot_kv_cache(cache: dict) -> dict:
+        """浅快照：分叉一条历史给别人读写（stage4 OPSD 的教师上下文），**原缓存一个字节不动**。
+
+        为什么浅拷贝就是正确的分支（与上游 GraphGWM `KVCache.clone_shallow` 同一条依据）：
+        `arope_self_attn` 的提交趟写的是 `kv["k"][li] = _evict(torch.cat([k_hist, k_raw]), kv).detach()`
+        —— `torch.cat` 造**新**张量、`_evict` 要么原样返回它要么再 cat 一个新的，然后**重新绑定**
+        列表槽位；缓存里已有的张量从不被就地写。`frames_cached` 是 int、`freqs_*` / `hw` / `layer` /
+        `commit` 也都是每次前向重新赋值的瞬态字段。所以：新 dict + 拷贝过的 'k'/'v' 列表 + 指向新 dict
+        的 `_KVRef`，就是一条独立分支 —— 往快照里 commit 只改快照的槽位，原缓存的槽位仍指着旧张量。
+
+        与 `forward_frame` 入口断言 `frames_cached == frame_index` 的配合（stage4 的用法）：
+        真实历史 0..k−2 的快照 frames_cached == k−1；把学生的 cell k−1 以 frame_index=k−1 commit 进快照
+        ⇒ frames_cached == k；随后给 cell k 打分（frame_index=k）正好通过。用完必须
+        `release_kv_cache(snapshot)`：快照自己也带 `_ref` 引用环（原缓存不受影响，共享的张量仍被它持有）。
+        """
+        snap = {k: v for k, v in cache.items() if k != "_ref"}
+        snap["k"] = list(cache["k"])
+        snap["v"] = list(cache["v"])
+        snap["_ref"] = _KVRef(snap)
+        return snap
+
+    @staticmethod
+    def last_kv_frame(cache: dict) -> tuple[list, list]:
+        """取**最近一次 commit** 的那个 cell 的 raw K/V（每层 [b, hw, dim]，RoPE 之前，clone 过、detach 过）。
+
+        提交趟把新 cell cat 在末尾、`_evict` 只从 sink 之后丢最老的，所以最近一格恒是每层最后 hw 个 token。
+        clone：调用方通常紧接着 `release_kv_cache` 这份缓存，切片视图会拖着整条 [b, n*hw, dim] 不放。
+        stage4 OPSD 用它复现官方的写入顺序（见 `append_kv_frame`）。"""
+        hw = cache["hw"]
+        assert hw is not None and cache["k"][0] is not None, "缓存还没提交过任何 cell"
+        ks = [t[:, -hw:].detach().clone() for t in cache["k"]]
+        vs = [t[:, -hw:].detach().clone() for t in cache["v"]]
+        return ks, vs
+
+    @staticmethod
+    def append_kv_frame(cache: dict, k_slices: list, v_slices: list, frame_index: int) -> None:
+        """把一格**已经算好的** raw K/V 按提交趟的规则追加进缓存：逐层 cat + `_evict` + frames_cached 推进。
+
+        与 `arope_self_attn` 的 commit 路径逐字相同（cat 造新张量、重新绑定槽位、原张量不动），所以对
+        `snapshot_kv_cache` 的快照调用同样不碰原缓存。用途（stage4 OPSD，官方 gt_kv 的写入顺序）：官方教师把
+        自己的 chunk k 写进缓存时缓存里还是 [GT 0..k−2, 教师的 k−1]，**之后**才把 k−1 原位换成 GT；于是打分
+        k+1 时最近一格的 K/V（第 1 层起）是在"教师的 k−1"上下文里算出来的，不是在全 GT 历史上重算的。
+        我们的缓存 append-only，只能：cell k 打分完把教师的 x0 commit 进它打分用的那条快照（上下文正是官方的），
+        `last_kv_frame` 取走那格 raw K/V，cell k+1 时用本函数把它接到真值 0..k−1 的新快照末尾。
+        raw K/V 与位置无关（RoPE 读时按 (tpos, offset_win) 重算），搬到另一条同长度的缓存末尾是精确的。"""
+        assert cache["frames_cached"] == frame_index, (
+            f"append_kv_frame：frame_index={frame_index} 但缓存最后提交的是 cell {cache['frames_cached'] - 1}")
+        assert len(k_slices) == len(cache["k"]) == len(v_slices), "层数不符"
+        hw = k_slices[0].shape[1]
+        assert cache["hw"] in (None, hw), (cache["hw"], hw)
+        k0 = cache["k"][0]
+        n_cached = 0 if k0 is None else k0.shape[1] // hw
+        if cache["kv_window"] is not None:
+            cap = cache["sink_size"] + cache["kv_window"]
+            assert n_cached == min(frame_index, cap), (n_cached, frame_index, cap)
+        else:
+            assert n_cached == frame_index, (n_cached, frame_index)
+        cache["hw"] = hw
+        for li in range(len(cache["k"])):
+            kh, vh = cache["k"][li], cache["v"][li]
+            kn, vn = k_slices[li], v_slices[li]
+            assert kn.shape[1] == hw and vn.shape[1] == hw, (kn.shape, vn.shape)
+            nk = kn if kh is None else torch.cat([kh, kn], dim=1)
+            nv = vn if vh is None else torch.cat([vh, vn], dim=1)
+            cache["k"][li] = _evict(nk, cache).detach()
+            cache["v"][li] = _evict(nv, cache).detach()
+        cache["frames_cached"] = frame_index + 1
+
+    def forward_frame(self, x_k: torch.Tensor, t_k, ctx_k: torch.Tensor, frame_index: int,
+                      kv_cache: dict, commit: bool = True, *,
+                      offset_win=None, mask_in_k=None, context_ids=None, r_k=None,
+                      angle_win=None) -> torch.Tensor:
+        """单 cell 前向，注意力打到缓存里的历史 + 自己。**不要直接调**（走 `forward(..., kv_cache=...)`，
+        否则 FSDP 不 all-gather 根模块的参数）。
+
+        x_k:        [b, 48, 1, 30, 52]
+        t_k:        标量 / [b] / [b, 1]，**原样使用**，绝不套 first_frame_cond（BLOCKER B1）
+        r_k:        flow map 的第二时间步，规则同 t_k：**同样原样使用、同样不套 first_frame_cond**
+                    （B1 对 r 一字不差地成立）。装了 flow map ⇒ 必填；没装 ⇒ None 或与 t_k 相同。
+        ctx_k:      [b, 1, 512, 4096]（逐 cell）或 [b, 512, 4096]（整条）
+        frame_index 绝对 cell 序号 k
+        commit      True = 把这一帧的 K/V 写进缓存（干净/最终那一趟）；False = 中间去噪趟，只读不写。
+                    滑窗 commit 是 append-only 的：同一个 frame_index commit 两次会把这帧塞进去两遍，
+                    n_cached / 驱逐全乱。拦它的是入口的 `frames_cached == frame_index` 断言（每次
+                    commit 把 frames_cached 推到 frame_index+1）；下面的 n_cached 断言只在缓存饱和
+                    之前有效 —— k ≥ cap 后 `_evict` 让 n_cached 恒等于 cap，重复 / 漏 commit 形状对、
+                    内容错、不报错（GEO-3）。
+        offset_win  [b, n_cached+1, 2] px，**窗口锚定**：offset_win[:,0] 恒为 0（窗口里最老那个 cell），
+                    offset_win[:,-1] 是当前 cell。None ⇒ 普通查表 RoPE。
+        mask_in_k   [b, 1, 1, 30, 52]，用 **锚定之后** 的 offset 算的 known 掩码。
+        """
+        dit = self.dit
+        assert self.causal, "forward_frame 是因果学生的流式路径（causal=True）"
+        b, _, f1, lat_h, lat_w = x_k.shape
+        assert f1 == 1, f"forward_frame 一次只处理一个 cell，得到 f={f1}"
+        p_f, p_h, p_w = dit.patch_size
+        h, w = lat_h // p_h, lat_w // p_w
+        hw = h * w
+
+        # ★ GEO-3：每个 frame_index 恰好 commit 一次。缓存饱和后 n_cached 恒 == cap，下面���条断言
+        #   拦不住重复 / 漏 commit（形状对、内容错、不报错），只有这个计数器能。commit=False 的去噪趟
+        #   不推进它，所以同一 frame_index 上反复 commit=False 再 commit=True 一次照样通过。
+        assert kv_cache["frames_cached"] == frame_index, (
+            f"frame_index={frame_index} 但缓存最后提交的是 cell {kv_cache['frames_cached'] - 1} "
+            "—— 漏 commit 或重复 commit 了"
+        )
+
+        # 缓存长度以第 0 层为准（各层同步推进）
+        k0 = kv_cache["k"][0]
+        n_cached = 0 if k0 is None else k0.shape[1] // hw
+        sliding = kv_cache["kv_window"] is not None
+        if sliding:
+            cap = kv_cache["sink_size"] + kv_cache["kv_window"]
+            assert n_cached == min(frame_index, cap), (
+                f"缓存里有 {n_cached} 个 cell，按 frame_index={frame_index} / 窗口 {cap} 应为 "
+                f"{min(frame_index, cap)} —— 漏 commit 或重复 commit 了"
+            )
+        else:
+            assert n_cached == frame_index, (
+                f"绝对模式下缓存长度必须等于 frame_index（{n_cached} != {frame_index}）"
+            )
+        tpos = n_cached  # 当前 cell 的时间位；滑窗下恒等于 min(frame_index, rope_cap)
+
+        # ---- 时间步：原样用 t_k / r_k，**没有** first_frame_cond ----
+        r_eff = self._resolve_r(t_k, r_k)
+        t_tok = _per_token_timestep_pf(x_k, t_k, hw, first_frame_cond=False)  # [b, hw]
+        r_tok = (None if r_eff is None
+                 else _per_token_timestep_pf(x_k, r_eff, hw, first_frame_cond=False))
+        t = self._embed_time(t_tok, r_tok)
+        t_mod = dit.time_projection(t).unflatten(2, (6, dit.dim))
+
+        ctx = _embed_context(dit, ctx_k, context_ids)
+
+        xt = dit.patch_embedding(x_k)
+        mask_emb = self._mask_embed(mask_in_k, xt)
+        if mask_emb is not None:
+            xt = xt + mask_emb
+        assert xt.shape[2] == 1 and xt.shape[3] == h and xt.shape[4] == w
+        tokens = rearrange(xt, "b c f h w -> b (f h w) c").contiguous()
+
+        # ---- 窗口锚定的 RoPE ----
+        if offset_win is None:
+            off_win_tok = None
+        else:
+            assert offset_win.shape == (b, n_cached + 1, 2), (
+                f"offset_win 形状应为 [b={b}, n_cached+1={n_cached + 1}, 2]，得到 {tuple(offset_win.shape)}"
+            )
+            assert bool(torch.all(offset_win[:, 0] == 0)), (
+                "offset_win[:, 0] 必须是 0：世界坐标要锚在窗口里最老那个 cell 上"
+                "（off_win = offset_px[a:k+1] - offset_px[a]）"
+            )
+            off_win_tok = offset_win.to(torch.float64) / PX_PER_TOKEN
+        if angle_win is None:
+            ang_win = None
+        else:
+            assert off_win_tok is not None, "给了 angle_win 就必须给 offset_win（坐标要一起进当前相机系）"
+            ang_win = angle_win.to(device=tokens.device, dtype=torch.float64)
+            assert ang_win.shape == (b, n_cached + 1), (
+                f"angle_win 形状应为 [b={b}, n_cached+1={n_cached + 1}]，得到 {tuple(ang_win.shape)}"
+            )
+            assert bool(torch.all(ang_win[:, -1] == 0)), (
+                "angle_win[:, -1] 必须是 0：一切都表达在**当前 cell 的相机系**里，"
+                "当前帧自己的格点必须保持轴对齐（模型唯一见过的帧内结构）"
+            )
+        cur_off = None if off_win_tok is None else off_win_tok[:, n_cached:n_cached + 1]
+        cur_ang = None if ang_win is None else ang_win[:, n_cached:n_cached + 1]
+        freqs_cur = build_world_freqs(dit, 1, h, w, cur_off, device=tokens.device, f_start=tpos,
+                                      angle=cur_ang)
+
+        freqs_hist = None
+        if n_cached > 0:
+            n_sink = min(kv_cache["sink_size"], n_cached)
+            R = n_cached - n_sink
+            parts = []
+            if n_sink > 0:
+                # sink 段永久钉在时间位 0..n_sink-1（sink_size=0 时这一段不存在）
+                off_s = None if off_win_tok is None else off_win_tok[:, :n_sink]
+                ang_s = None if ang_win is None else ang_win[:, :n_sink]
+                parts.append(build_world_freqs(dit, n_sink, h, w, off_s, device=tokens.device, f_start=0,
+                                               angle=ang_s))
+            if R > 0:
+                off_r = None if off_win_tok is None else off_win_tok[:, n_sink:n_cached]
+                ang_r = None if ang_win is None else ang_win[:, n_sink:n_cached]
+                parts.append(build_world_freqs(dit, R, h, w, off_r, device=tokens.device, f_start=tpos - R,
+                                               angle=ang_r))
+            freqs_hist = parts[0] if len(parts) == 1 else torch.cat(parts, dim=_freqs_cat_dim(parts[0]))
+
+        kv_cache["hw"] = hw
+        kv_cache["commit"] = bool(commit)
+        kv_cache["freqs_cur"] = freqs_cur
+        kv_cache["freqs_hist"] = freqs_hist
+
+        # 直接调 block，**绝不** gradient_checkpoint_forward（risk 6：滑窗 commit 是破坏性的，
+        # 重算时读到的是驱逐之后的缓存，缓存满了以后形状对内容错、不报错）
+        for li, block in enumerate(dit.blocks):
+            kv_cache["layer"] = li
+            tokens = block(tokens, ctx, t_mod, freqs_cur, f=1, attn_mask=None, kv=kv_cache)
+
+        tokens = dit.head(tokens, t)
+        if commit:
+            kv_cache["frames_cached"] = frame_index + 1
+        return dit.unpatchify(tokens, (1, h, w))
+
+
+# 模块级别名（与 AropeDiT.release_kv_cache 一样是静态的；stage4 OPSD 直接 import 它们）
+snapshot_kv_cache = AropeDiT.snapshot_kv_cache
+last_kv_frame = AropeDiT.last_kv_frame
+append_kv_frame = AropeDiT.append_kv_frame

backup_2026-09-14/code/distill/common.sh

@@ -0,0 +1,89 @@
+#!/usr/bin/env bash
+# 三阶段蒸馏所有 launcher 的公共前置：source 它，不要执行它。
+#   source "$(dirname "${BASH_SOURCE[0]}")/../common.sh"     # 从 distill/stageN_*/run.sh 里
+# 干四件事：
+#   1) 从自身位置推 ROOT（不写死路径 —— 脚本随仓库搬走也不会指向别人的盘）
+#   2) 统一环境变量（PYTHONPATH / DIFFSYNTH_SKIP_DOWNLOAD / 分配器 / NCCL / OMP），并 unset WorldMind 的三个 env hook
+#   3) 钉死 CUDA_VISIBLE_DEVICES=0,1,2,3 —— 4-7 卡是别人的，1-3 是正式训练预留，任何情况下都不碰
+#   4) GPU 占用守卫：0-3 里任意一张 >500 MiB 就拒绝启动
+set -euo pipefail
+
+# ---- 1) ROOT：由本文件所在目录（<ROOT>/distill）上溯一级得到 ----
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+export ROOT
+PY="$ROOT/.venv/bin/python"
+ACC="$ROOT/.venv/bin/accelerate"
+export PY ACC
+[ -x "$PY" ]  || { echo "ABORT: 找不到 $PY"; exit 1; }
+[ -x "$ACC" ] || { echo "ABORT: 找不到 $ACC"; exit 1; }
+
+# ---- 2) 环境 ----
+export PYTHONPATH="$ROOT${PYTHONPATH:+:$PYTHONPATH}"
+export DIFFSYNTH_SKIP_DOWNLOAD=True
+export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
+export NCCL_P2P_LEVEL="${NCCL_P2P_LEVEL:-NVL}"      # 0-3 在同一张 NV18 NVLink mesh 上
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-16}"     # 本机 160 逻辑核；不限线程会让 4 个 rank 互相抖动
+export TOKENIZERS_PARALLELISM=false
+unset GM_XCACHE GM_COMPILE GM_SAGEATTN || true      # WorldMind 的 env hook：XCACHE 切断 autograd、COMPILE 会挂、SAGE 更慢
+
+# ---- 3) 卡的隔离：只用 0,1,2,3，且不接受外部覆盖 ----
+# CUDA 默认按 FASTEST_FIRST 枚举，理论上可以与 nvidia-smi 的 PCI 顺序不同；钉成 PCI_BUS_ID 之后
+# CUDA_VISIBLE_DEVICES=0,1,2,3 与守卫里的 `nvidia-smi -i 0,1,2,3` 一定指同样这四张物理卡。
+export CUDA_DEVICE_ORDER=PCI_BUS_ID
+# DISTILL_NGPU：用前几张卡。**默认仍是 4** —— 这是共用机器，4-7 默认留给别人。
+# 要用 8 必须显式 DISTILL_NGPU=8，而且先跟人确认：守卫只能看出「现在有没有人在用」，
+# 看不出「别人正打算用」。
+DISTILL_NGPU="${DISTILL_NGPU:-4}"
+case "$DISTILL_NGPU" in
+  1|2|4|8) : ;;
+  *) echo "ABORT: DISTILL_NGPU 只接受 1/2/4/8，得到 $DISTILL_NGPU" >&2; return 1 2>/dev/null || exit 1 ;;
+esac
+DISTILL_GPUS="$(seq -s, 0 $((DISTILL_NGPU - 1)))"
+DISTILL_NPROC="$DISTILL_NGPU"
+export CUDA_VISIBLE_DEVICES="$DISTILL_GPUS"
+export DISTILL_GPUS DISTILL_NPROC DISTILL_NGPU
+
+# ---- 4) GPU 占用守卫 ----
+# 注意：`nvidia-smi -i` 用的是物理编号，不受 CUDA_VISIBLE_DEVICES 影响，所以这里明确只查 0-3。
+# 上游 WorldMind 的写法（launch/stage3_cd15k.sh:25-28）不带 -i、查全部 8 张、阈值 20000 MiB —— 别人用 4-7 时它会误报中止，
+# 正好和我们要的隔离相反；且没有 FORCE=1 偷卡的后门。
+require_idle_gpus() {
+  local thresh="${1:-500}" out rc busy seen want
+  command -v nvidia-smi >/dev/null 2>&1 || { echo "ABORT: 没有 nvidia-smi" >&2; return 1; }
+  # ★ 必须显式收 nvidia-smi 的退出码，不能写成 busy="$(nvidia-smi ... | awk ...)"。
+  #   本函数是以 `require_idle_gpus ... || exit 1` 调用的，bash 在 `||` 左侧会**关掉 errexit**（连函数体一起），
+  #   所以命令替换失败不会中止脚本，只会让 busy 变成空串 —— 守卫就会打印"空闲"然后照样去抢卡。
+  #   驱动半死 / GPU 掉总线 / -i 编号非法 都会走到这条路上，宁可不跑也不能查不到就放行。
+  out="$(timeout 30 nvidia-smi -i "$DISTILL_GPUS" --query-gpu=index,memory.used --format=csv,noheader,nounits 2>&1)"
+  rc=$?
+  if [ "$rc" -ne 0 ]; then
+    echo "ABORT: nvidia-smi -i $DISTILL_GPUS 查询失败（退出码 $rc）：" >&2
+    printf '       %s\n' "$out" >&2
+    echo "       查不到占用就不许启动 —— 1-3 是正式训练预留、4-7 是别人的卡。" >&2
+    return 1
+  fi
+  # 必须恰好回来 0,1,2,3 这四行：驱动半死 / MIG 分区时 nvidia-smi 可能只吐一部分还 exit 0，
+  # 漏掉的那几张等于根本没查过。
+  seen="$(printf '%s\n' "$out" | awk -F', *' 'NF>=2 {printf "%s,", $1}')"
+  want="$DISTILL_GPUS,"
+  if [ "$seen" != "$want" ]; then
+    echo "ABORT: nvidia-smi 只回了 [${seen%,}]，期望 [$DISTILL_GPUS]；有卡查不到，不启动。" >&2
+    return 1
+  fi
+  # $2+0 强制数值比较；memory.used 不是纯数字（[N/A] / [Unknown Error]）也当成占用
+  busy="$(printf '%s\n' "$out" | awk -F', *' -v t="$thresh" \
+            '$2 !~ /^[0-9]+$/ || $2+0 > t+0 {printf "GPU%s=%sMiB ", $1, $2}')"
+  if [ -n "$busy" ]; then
+    echo "ABORT: GPU $DISTILL_GPUS 上有占用（>${thresh} MiB）: $busy" >&2
+    echo "       另一个蒸馏/训练任务还在跑？先停掉它。不要去动 GPU 4-7（别人的卡）。" >&2
+    return 1
+  fi
+  echo "[guard] GPU $DISTILL_GPUS 空闲（每张 <= ${thresh} MiB），可以启动。"
+  return 0
+}
+
+# 阈值只允许调**严**，不允许调松：GPU_BUSY_MIB=999999 曾经等于一个偷卡后门 ——
+# 能让第二个 launcher 直接盖在正在跑的训练上（两个 job 抢同四张卡 = 双双 OOM 或龟速）。
+# 上面第 41 行说"没有偷卡的后门"，这一行是让那句话成立的地方。
+GPU_BUSY_MIB_EFF=$(( ${GPU_BUSY_MIB:-500} < 500 ? ${GPU_BUSY_MIB:-500} : 500 ))
+require_idle_gpus "$GPU_BUSY_MIB_EFF" || exit 1

backup_2026-09-14/code/distill/configs/accelerate_fsdp_dmd_4gpu.yaml

@@ -0,0 +1,47 @@
+# 蒸馏 stage3（self-forcing DMD）专用：4 卡（GPU 0,1,2,3） FSDP1 FULL_SHARD，bf16。
+# 派生自 _wm/WorldMind/distill/code/common/accelerate_fsdp_dmd.yaml（那份是 num_processes: 3，被 run.sh 的
+# --num_processes 覆盖；这里直接写死 4）。stage1/2 用 ZeRO-2，只有 stage3 有两个可训模型（student + critic），
+# 而 accelerate 的 DeepSpeed 分支只保留传给 prepare 的最后一个 nn.Module（accelerator.py:2229-2236），故必须换 FSDP。
+#
+# fsdp_version: 1 —— 显式选 FSDP1；accelerate 1.14 的 v1 分支读 fsdp_reshard_after_forward（utils/dataclasses.py:1848），
+#   fsdp_sharding_strategy 是已废弃的旧拼写（:1836）。FULL_SHARD ⇒ 参数/梯度/优化器状态全切 4 份。
+# fsdp_use_orig_params: true —— 【风险 5，致命】两处吃紧：优化器在 prepare 之前按 model.parameters() 建好且此后不再重映射，
+#   若为 false，FSDP1 换上 FlatParameter，优化器 step 的是过期张量、永远收不到梯度 —— 模型静默地完全不学习，无报错；
+#   同时 zip-EMA 也要求 Parameter 的身份/顺序/分片完全一致。启动时断言
+#   accelerator.state.fsdp_plugin.use_orig_params 为真、开跑前 probe_student_grad 量一次梯度可达性；
+#   并且（已实现，S3-GUARD-3）dmd_arope.py 在每个生成器步的 opt_g.step() 之前检查 gnorm_g：
+#   step > resume_step 后恰为 0.0 或 NaN 直接 SystemExit（inf 不拦，clip 已归零；step 0 的 dmd 梯度
+#   恰为 0 是配方性质、由 probe 覆盖、不拦）。
+# fsdp_sync_module_states: false + fsdp_cpu_ram_efficient_loading: false —— 每个 rank 自己确定性地读同一份 safetensors，
+#   不需要 rank0 广播；开了 sync 会让 rank0 一次性实体化全部模型，prepare 阶段直接 OOM。
+# fsdp_state_dict_type: SHARDED_STATE_DICT —— 存 optimizer state 用分片；导权重走 accelerator.get_state_dict + safetensors。
+# fsdp_activation_checkpointing: false —— 梯度检查点在模型侧按路径开（student 只在并行路径上开、critic 必须开），
+#   不交给 FSDP 全局包，否则 rollout 的单帧图也会被包进去。
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_version: 1
+  fsdp_reshard_after_forward: FULL_SHARD
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: DiTBlock
+  fsdp_backward_prefetch: BACKWARD_PRE
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_offload_params: false
+  fsdp_use_orig_params: true
+  fsdp_sync_module_states: false
+  fsdp_cpu_ram_efficient_loading: false
+  fsdp_activation_checkpointing: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

backup_2026-09-14/code/distill/configs/accelerate_fsdp_dmd_8gpu.yaml

@@ -0,0 +1,49 @@
+# 8 卡版：由 accelerate_fsdp_dmd_4gpu.yaml 只改 num_processes 4 -> 8 得到。
+# ⚠️ 共用机器，4-7 默认不归我们；用它前先确认没人要用（common.sh 需要 DISTILL_NGPU=8）。
+# 蒸馏 stage3（on-policy DMD）专用：8 卡 FSDP1 FULL_SHARD，bf16。
+# 派生自 _wm/WorldMind/distill/code/common/accelerate_fsdp_dmd.yaml（那份是 num_processes: 3，被 run.sh 的
+# --num_processes 覆盖；这里直接写死 4）。stage1/2 用 ZeRO-2，只有 stage3 有两个可训模型（student + critic），
+# 而 accelerate 的 DeepSpeed 分支只保留传给 prepare 的最后一个 nn.Module（accelerator.py:2229-2236），故必须换 FSDP。
+#
+# fsdp_version: 1 —— 显式选 FSDP1；accelerate 1.14 的 v1 分支读 fsdp_reshard_after_forward（utils/dataclasses.py:1848），
+#   fsdp_sharding_strategy 是已废弃的旧拼写（:1836）。FULL_SHARD ⇒ 参数/梯度/优化器状态全切 4 份。
+# fsdp_use_orig_params: true —— 【风险 5，致命】两处吃紧：优化器在 prepare 之前按 model.parameters() 建好且此后不再重映射，
+#   若为 false，FSDP1 换上 FlatParameter，优化器 step 的是过期张量、永远收不到梯度 —— 模型静默地完全不学习，无报错；
+#   同时 zip-EMA 也要求 Parameter 的身份/顺序/分片完全一致。启动时断言
+#   accelerator.state.fsdp_plugin.use_orig_params 为真、开跑前 probe_student_grad 量一次梯度可达性；
+#   并且（已实现，S3-GUARD-3）dmd_arope.py 在每个生成器步的 opt_g.step() 之前检查 gnorm_g：
+#   step > resume_step 后恰为 0.0 或 NaN 直接 SystemExit（inf 不拦，clip 已归零；step 0 的 dmd 梯度
+#   恰为 0 是配方性质、由 probe 覆盖、不拦）。
+# fsdp_sync_module_states: false + fsdp_cpu_ram_efficient_loading: false —— 每个 rank 自己确定性地读同一份 safetensors，
+#   不需要 rank0 广播；开了 sync 会让 rank0 一次性实体化全部模型，prepare 阶段直接 OOM。
+# fsdp_state_dict_type: SHARDED_STATE_DICT —— 存 optimizer state 用分片；导权重走 accelerator.get_state_dict + safetensors。
+# fsdp_activation_checkpointing: false —— 梯度检查点在模型侧按路径开（student 只在并行路径上开、critic 必须开），
+#   不交给 FSDP 全局包，否则 rollout 的单帧图也会被包进去。
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_version: 1
+  fsdp_reshard_after_forward: FULL_SHARD
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: DiTBlock
+  fsdp_backward_prefetch: BACKWARD_PRE
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_offload_params: false
+  fsdp_use_orig_params: true
+  fsdp_sync_module_states: false
+  fsdp_cpu_ram_efficient_loading: false
+  fsdp_activation_checkpointing: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

backup_2026-09-14/code/distill/configs/accelerate_zero2_4gpu_ga2.yaml

@@ -0,0 +1,28 @@
+# 蒸馏 stage1/stage2 用：4 卡（GPU 0,1,2,3） DeepSpeed ZeRO-2，bf16，无 offload，梯度裁剪 1.0，梯度累积 2。
+# 派生自 configs/accelerate_zero2_4gpu.yaml，只改 gradient_accumulation_steps: 1 -> 2（等效 batch = 4 卡 × 1 × 2 = 8）。
+# train_micro_batch_size_per_gpu 由 accelerate 填 "auto" ⇒ 取训练脚本里 DataLoader 的 batch_size（默认 1）。
+# 训练脚本的 --grad_accum 必须等于这里的 2，否则 SystemExit（train.py:399 那条守卫要原样带进新 trainer）。
+# 不用 deepspeed_config_file：accelerate 1.14 里它与 mixed_precision 等键互斥，yaml 原生字段更省事。
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  gradient_accumulation_steps: 2
+  gradient_clipping: 1.0
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: false
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

backup_2026-09-14/code/distill/configs/accelerate_zero2_4gpu_ga4.yaml

@@ -0,0 +1,28 @@
+# 蒸馏 stage1/stage2 用：4 卡（GPU 0,1,2,3） DeepSpeed ZeRO-2，bf16，无 offload，梯度裁剪 1.0，梯度累积 4。
+# 派生自 configs/accelerate_zero2_4gpu.yaml，只改 gradient_accumulation_steps: 1 -> 4（等效 batch = 4 卡 × 1 × 4 = 16，stage1 的 eff-batch-16 备选）。
+# train_micro_batch_size_per_gpu 由 accelerate 填 "auto" ⇒ 取训练脚本里 DataLoader 的 batch_size（默认 1）。
+# 训练脚本的 --grad_accum 必须等于这里的 4，否则 SystemExit（train.py:399 那条守卫要原样带进新 trainer）。
+# 不用 deepspeed_config_file：accelerate 1.14 里它与 mixed_precision 等键互斥，yaml 原生字段更省事。
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  gradient_accumulation_steps: 4
+  gradient_clipping: 1.0
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: false
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

backup_2026-09-14/code/distill/configs/accelerate_zero2_8gpu_ga1.yaml

@@ -0,0 +1,34 @@
+# 蒸馏 stage1/stage2 用：8 卡（GPU 0-7） DeepSpeed ZeRO-2，bf16，无 offload，梯度裁剪 1.0，梯度累积 1。
+# 派生自 accelerate_zero2_4gpu_ga2.yaml，只改两处：num_processes 4 -> 8、gradient_accumulation_steps 2 -> 1。
+#
+# **等效 batch 仍然是 8**（8 卡 × micro 1 × GA 1 = 8），与 4 卡 GA2 完全相同 —— 换句话说这是同一个
+# 配方跑得快一倍，而不是另一个配方。想要等效 batch 16 得另建 ga2 的 8 卡版，别在这里改 GA。
+#
+# ⚠️ 这台是共用机器。4-7 默认不归我们；用这份配置前先确认没人要用。
+# common.sh 需要 DISTILL_NGPU=8 才会放开 CUDA_VISIBLE_DEVICES 和守卫的检查范围。
+#
+# train_micro_batch_size_per_gpu 由 accelerate 填 "auto" ⇒ 取训练脚本里 DataLoader 的 batch_size（默认 1）。
+# 训练脚本的 --grad_accum 必须等于这里的 1，否则 SystemExit。
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: false
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

backup_2026-09-14/code/distill/configs/accelerate_zero2_8gpu_ga2.yaml

@@ -0,0 +1,36 @@
+# 长视野续训用（distill/stage1_long/run.sh 在 DISTILL_NGPU=8 时选它）：8 卡（GPU 0-7） DeepSpeed ZeRO-2，
+# bf16，无 offload，梯度裁剪 1.0，梯度累积 2。
+# 派生自 accelerate_zero2_8gpu_ga1.yaml，只改 gradient_accumulation_steps 1 -> 2（longhorizon/README 建议等效 batch 16）。
+#
+# **等效 batch 是 16**（8 卡 × micro 1 × GA 2 = 16），与 4 卡 GA4（accelerate_zero2_4gpu_ga4.yaml）相同 ——
+# 这是 stage1_long 的配方，**不是** stage1/stage2 的等效 batch 8：那两条在 8 卡下用
+# accelerate_zero2_8gpu_ga1.yaml（各自的 run.sh 按 DISTILL_NGPU 自动选，并断言 卡数×GA == 8）。
+#
+# ⚠️ 这台是共用机器。4-7 默认不归我们；用这份配置前先确认没人要用。
+# common.sh 需要 DISTILL_NGPU=8 才会放开 CUDA_VISIBLE_DEVICES 和守卫的检查范围。
+#
+# train_micro_batch_size_per_gpu 由 accelerate 填 "auto" ⇒ 取训练脚本里 DataLoader 的 batch_size（默认 1）。
+# 训练脚本的 --grad_accum 必须等于这里的 2，否则 SystemExit。
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  gradient_accumulation_steps: 2
+  gradient_clipping: 1.0
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: false
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

backup_2026-09-14/code/distill/data.py

@@ -0,0 +1,399 @@
+"""蒸馏用的 batch 工具：mask / 权重图 / 偏移与文本的延长和开窗。
+
+数据集本身**一行不改**：`actionrope.dataset.AropeLatentDataset` + `arope_collate` 已经给齐
+了三个阶段需要的一切（`input_latents [b,48,21,30,52] bf16`、`context [b,21,512,4096] bf16`、
+`context_ids [b,21] i64`、`offset_px [b,21,2] f32`、`offset_tok`、`delta_tok`、`action_idx`、
+`clip_id`、`kind`）。本文件只做 batch 之上的纯函数，不碰 IO、不碰随机数据增强。
+
+四类工具
+--------
+1. `to_device` —— 与 `actionrope/train.py:344` 同一份语义（张量搬走，字符串 list 原样）。
+2. `make_mask_input` / `make_weight_map` —— stage1/2 的 `mask_input` 与 loss 权重图。
+   `mask_input` 直接用 `known_mask(offset_px / 16, 30, 52)`，**不能**写成
+   `loss_weight_map(...).eq(1)`：那样一来模型的输入通道就被 loss 配方绑架了——
+   `new_weight=1.0`（想关掉 new 区加权时的自然写法）会让权重图恒等于 1，掩码整幅变成
+   "全 known"，模型输入被悄悄改掉而不报错；`sigma`/`ramp_latent` 也各自能把这条等价关系
+   推到 float32 的舍入边缘（dist 模式下 `d == 0 ⟺ known` 只是巧合，紧贴边界的
+   `1 + 5e-10` 在 float32 里就是 1.0）。掩码是**模型输入**，权重图是**loss 配方**，
+   两件事必须各自从 offset 现算。
+3. `window_offsets` / `window_context` —— stage3 滑窗 KV 的**锚定切片**。绝对 cell k、
+   窗口最老保留 cell a 时，窗口内的世界坐标一律减去 `offset_px[a]`，于是整段窗口在 RoPE
+   看来就是一条以自身 cell 0 为原点的 21-cell 训练 clip（RoPE 是相对编码，整体平移不变）。
+   两者索引口径完全一致：`[:, a:k+1]`。
+4. `extend_offsets` / `extend_context` —— rollout 超过 21 cell 时的延长。
+   `--rollout_max 26` 意味着 cell 21..26 在数据里根本不存在：`parse_actions` 硬要求 21 个
+   cell（`infer.py:127`），`AropeLatentDataset.__getitem__` 硬要求 21 条 prompt。两条流必须
+   **同时**延长而且逐下标对齐，否则要么在 `context.shape[1] == 1+N` 的断言处崩，要么更糟：
+   几何往前走、文本停在 cell 20，rollout 照跑、loss 照降，但训的是错的东西。
+
+   两条流的延长策略**故意不同**，这不是笔误：
+   * 文本可以"保持"最后一条（`_wm/WorldMind/distill/code/stage3_dmd/dmd_gm.py:628-637`
+     的 `last = context[:, Fl-1:Fl]; pad = last.expand(...)`）。ARoPE 臂的 prompt 经
+     `prompts.strip_action` 剥掉了动作从句，只剩场景描述，而 rollout 不换场景，所以
+     "保持最后一条" 在这里是**语义正确**的，不是凑合。
+   * 偏移不能"保持"。偏移保持不动 = "玩家站住了"，那是**另一个动作**，不是中性动作
+     （V Rising 的 null action 陷阱，GraphGWM/README 记过一次）。所以偏移必须继续积分一个
+     速度：`continue`（沿用最后一段速度）或 `resample`（默认，每 5 个 cell 换一次方向，
+     方向取 `geometry.ACTION_DIRS` 9 类、速度取 sidecar 实测的 `walk_speed_px_s`）。
+
+速度 → 逐 cell 偏移的换算（不是拍脑袋的 v/4）
+------------------------------------------
+Wan2.2 因果 VAE 的分组是 cell 0 = 帧 0、cell k = 帧 4k−3…4k（`geometry.cell_frames`），
+逐 cell 偏移是 cell 内各帧偏移的**均值**（`geometry.frames_to_cells`），而逐帧偏移按
+`frame_off[i] = frame_off[i−1] + v[cell(i)] / FPS` 积分（`infer.py::parse_actions`）。
+把这两步合起来，逐 cell 增量是（v 单位 px/s，FPS=16）：
+
+    delta[1]   = 2.5 · v[1] / FPS                      （cell 0 只有帧 0，起步项不同）
+    delta[k]   = (1.5 · v[k−1] + 2.5 · v[k]) / FPS     （k ≥ 2）
+
+匀速时退化成 4·v/FPS = v/4 px per cell（一个 cell = 4 帧 = 0.25 s）。换方向的那个 cell 上
+1.5/2.5 的加权才是真值，直接写 v/4 会在每个 chunk 边界差 1.5·|Δv|/16（最大 ~13 px ≈ 0.8 个
+latent 格）。`test_data.py::test_extend_matches_frame_integration` 用帧级积分逐位核对了这条。
+"""
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+
+from actionrope.arope import known_mask, loss_weight_map
+from actionrope.geometry import ACTION_DIRS, CELLS, FPS, LAT_H, LAT_W, PX_PER_LATENT
+
+# 一个 cell 覆盖 4 帧（cell 0 除外），见 geometry.cell_frames
+CELL_FRAMES = 4
+# 帧级积分 + cell 内均值合出来的两个权重（推导见模块 docstring）
+_W_PREV, _W_CUR = 1.5, 2.5
+
+# sidecar 实测步速（data/latent/arope_sidecar.pt 的 meta["walk_speed_px_s"]，已乘 |gain|）。
+# 横纵不同是透视造成的，不是噪声；`infer.py::parse_actions` 也是按轴分别乘这个数。
+# 只作为默认值：真跑的时候优先用 walk_speed_from_sidecar() 从当前 sidecar 现读。
+WALK_SPEED_PX_S = (69.24106392992304, 46.012630371093756)
+
+# rollout 延长的默认档位，��� stage3 的 --offset_chunk_cells 5 对齐
+DEFAULT_CHUNK_CELLS = 5
+DEFAULT_SPEED_MULT = (0.8, 1.2)
+
+
+# --------------------------------------------------------------------------
+# batch 搬运
+# --------------------------------------------------------------------------
+
+def to_device(batch: dict, device, non_blocking: bool = True) -> dict:
+    """张量搬到 device，其余（clip_id / kind 这种 list[str]）原样带过去。
+
+    与 `actionrope/train.py:344-345` 同一份实现，蒸馏侧不再各写一遍。
+    """
+    return {
+        k: (v.to(device, non_blocking=non_blocking) if isinstance(v, torch.Tensor) else v)
+        for k, v in batch.items()
+    }
+
+
+# --------------------------------------------------------------------------
+# mask 通道 / loss 权重图
+# --------------------------------------------------------------------------
+
+def make_mask_input(offset_px: torch.Tensor, lat_h: int = LAT_H, lat_w: int = LAT_W,
+                    angle=None) -> torch.Tensor:
+    """known/new 掩码通道 → float32 [b, 1, f, lat_h, lat_w]，喂给 `dit.arope_mask_embedding`。
+
+    offset_px [b, f, 2] px。除以 PX_PER_LATENT=16 换成 latent 格单位，与
+    `actionrope/train.py:243-245` 逐位一致（known=1、new=0）。
+    stage3 传的是**锚定后**的窗口偏移（见 window_offsets），于是 "known" 的含义变成
+    "落在窗口最老那一 cell 的足迹内"——正是模型训练时见过的那种。
+    """
+    off = _as_bf2(offset_px, "offset_px")
+    if angle is not None and bool(torch.any(torch.as_tensor(angle) != 0)):
+        known = _known_rotated(off / PX_PER_LATENT, angle, lat_h, lat_w)
+    else:
+        known = known_mask(off / PX_PER_LATENT, lat_h, lat_w)  # bool [b, f, h, w]
+    return known.to(torch.float32).unsqueeze(1)
+
+
+def _known_rotated(off: torch.Tensor, angle, lat_h: int, lat_w: int) -> torch.Tensor:
+    """相机在锚点与当前 cell 之间转了 angle 弧度时的 known 掩码（angle = 锚点朝向 − 当前朝向）。
+
+    当前格点 p 先按原口径加上偏移，再绕画幅中心转 −angle 回到**锚点相机系**，落在锚点足迹
+    [−0.5, h−0.5] × [−0.5, w−0.5] 内 ⇒ known。angle=0 时逐位退化成 `known_mask`（旋转成恒等）。
+    旋转之后 known 区是一块**斜的**多边形 —— 训练里只出现过轴对齐的边界，这一条是已知的分布外。
+    """
+    from distill.arope_dit import rot2d
+    b, f = off.shape[:2]
+    dev, dt = off.device, torch.float64
+    ang = torch.as_tensor(angle, device=dev, dtype=dt).view(b, f, 1, 1)
+    gi = torch.arange(lat_h, device=dev, dtype=dt).view(1, 1, lat_h, 1).expand(b, f, lat_h, lat_w)
+    gj = torch.arange(lat_w, device=dev, dtype=dt).view(1, 1, 1, lat_w).expand(b, f, lat_h, lat_w)
+    ctr_h, ctr_w = (lat_h - 1) / 2.0, (lat_w - 1) / 2.0
+    o = off.to(dt).view(b, f, 1, 1, 2)
+    p = torch.stack([gj + o[..., 0] - ctr_w, gi + o[..., 1] - ctr_h], dim=-1)
+    q = rot2d(p, -ang)
+    wj, wi = q[..., 0] + ctr_w, q[..., 1] + ctr_h
+    return (wi >= -0.5) & (wi <= lat_h - 0.5) & (wj >= -0.5) & (wj <= lat_w - 0.5)
+
+
+def make_weight_map(offset_px: torch.Tensor, lat_h: int = LAT_H, lat_w: int = LAT_W,
+                    new_weight: float = 2.0, ramp_latent: float = 0.0, sigma=None) -> torch.Tensor:
+    """loss 权重图 → float32 [b, 1, f, lat_h, lat_w]（known=1，new=new_weight）。
+
+    直接转手 `actionrope.arope.loss_weight_map`，把 10k 检查点训练时用的 new 区 ×2 权重
+    原样带进 stage1。stage2 的 CD loss **故意不用**这张图（见 PORT_SPEC）。
+    """
+    off = _as_bf2(offset_px, "offset_px")
+    return loss_weight_map(off, lat_h, lat_w, new_weight=new_weight,
+                           ramp_latent=ramp_latent, sigma=sigma)
+
+
+# --------------------------------------------------------------------------
+# 滑窗锚定
+# --------------------------------------------------------------------------
+
+def window_offsets(offset_px: torch.Tensor, a: int, k: int) -> torch.Tensor:
+    """窗口 [a, k] 的锚定偏移 → [b, k−a+1, 2]，第 0 项恒为 0。
+
+    a = max(0, k − kv_window) 是窗口里最老的保留 cell。RoPE 是相对编码，整窗同时平移
+    (t, offset) 不改变任何注意力分数，所以 RoPE 重锚定是精确的（`tests/test_world_rope_window.py` a1/a3），
+    当前查询 cell 的 known 掩码也精确（d1）。但**缓存里各 cell 的 known 掩码通道是按它们各自提交时的锚点
+    a_j = max(0, j−20) < a 固化的**，所以 k>20 的历史与一条锚在 k−20 的训练 clip 并不逐位相同 ——
+    这是 stage2→stage3 的输入分布移位，靠 self-forcing 训练穿过（审计 GEO-2）。
+    """
+    off = _as_bf2(offset_px, "offset_px")
+    f = off.shape[1]
+    if not 0 <= a <= k < f:
+        raise ValueError(f"窗口下标越界: a={a}, k={k}, f={f}")
+    return off[:, a:k + 1] - off[:, a:a + 1]
+
+
+def window_context(context: torch.Tensor, context_ids: torch.Tensor, a: int, k: int):
+    """`window_offsets` 的文本孪生：同一段 [a, k] 的 context / context_ids。
+
+    单独写出来只有一个理由——把 "偏移窗口" 和 "文本窗口" 的下标口径钉在同一个函数里，
+    stage3 的 scored tail 才不会两边差一格。
+    """
+    ctx, ids = _check_context(context, context_ids)
+    f = ctx.shape[1]
+    if not 0 <= a <= k < f:
+        raise ValueError(f"窗口下标越界: a={a}, k={k}, f={f}")
+    return ctx[:, a:k + 1], ids[:, a:k + 1]
+
+
+# --------------------------------------------------------------------------
+# rollout 延长：偏移
+# --------------------------------------------------------------------------
+
+def walk_speed_from_sidecar(sidecar) -> tuple[float, float]:
+    """从 sidecar（路径或已加载的 dict）取实测步速 (px/s)。默认常量只是它的冻结副本。"""
+    from actionrope.dataset import load_sidecar
+
+    meta = load_sidecar(sidecar)["meta"]
+    ws = meta["walk_speed_px_s"]
+    if len(ws) != 2:
+        raise ValueError(f"sidecar meta 的 walk_speed_px_s 应为 2 个数，得到 {ws!r}")
+    return float(ws[0]), float(ws[1])
+
+
+def append_cell_velocity(offset_px: torch.Tensor, vel: torch.Tensor) -> torch.Tensor:
+    """把逐 cell 速度 vel [b, m, 2]（px/s）接到 offset_px [b, F0, 2] 后面 → [b, F0+m, 2]。
+
+    用模块 docstring 的 (1.5·v_prev + 2.5·v_cur)/FPS 增量积分，即帧级积分 + cell 内均值的
+    闭式解。第一个新 cell 需要"上一 cell 的速度"，而数据只给了偏移，所以从最后一段已知增量
+    反推匀速等效值 v_prev（匀速时精确；变速时那一格的误差 ≤ 1.5·|Δv|/16）。
+    """
+    off = _as_bf2(offset_px, "offset_px")
+    if vel.dim() != 3 or vel.shape[0] != off.shape[0] or vel.shape[2] != 2:
+        raise ValueError(f"vel 形状应为 [{off.shape[0]}, m, 2]，得到 {tuple(vel.shape)}")
+    b, f0 = off.shape[:2]
+    if f0 < 1:
+        raise ValueError("offset_px 至少要有 cell 0（形状 [b, ≥1, 2]）")
+    m = vel.shape[1]
+    if m == 0:
+        return off.clone()
+    vel = vel.to(device=off.device, dtype=off.dtype)
+
+    v_prev = _implied_last_velocity(off)  # [b, 2] px/s
+    out = [off]
+    cur = off[:, -1]                       # [b, 2] 最后一个已知 cell 的累计偏移
+    k = f0                                 # 即将写入的绝对 cell 下标
+    for j in range(m):
+        v = vel[:, j]
+        if k == 0:
+            raise ValueError("cell 0 的偏移恒为 0，不能由速度生成")
+        if k == 1:
+            # cell 0 只有帧 0，起步项只有 2.5 那一半
+            delta = (_W_CUR * v) / FPS
+        else:
+            delta = (_W_PREV * v_prev + _W_CUR * v) / FPS
+        cur = cur + delta
+        out.append(cur.unsqueeze(1))
+        v_prev = v
+        k += 1
+    return torch.cat(out, dim=1)
+
+
+def extend_offsets(offset_px: torch.Tensor, N: int, mode: str = "resample",
+                   chunk: int = DEFAULT_CHUNK_CELLS, rng: torch.Generator | None = None,
+                   walk_speed=None, speed_mult: tuple[float, float] = DEFAULT_SPEED_MULT,
+                   include_idle: bool = True) -> torch.Tensor:
+    """把 [b, F0, 2] 的逐 cell 偏移延长/裁剪到 rollout 需要的 [b, N+1, 2]。
+
+    N = rollout 生成的 cell 数（不含 cell 0），所以长度是 1+N，与
+    `dmd_gm.py:628-637` 的 `prompts.shape[1] == 1 + N` 口径一致。
+    1+N ≤ F0 时就是切片（原样保留真值），不会引入任何合成运动。
+
+    mode:
+      * "resample"（默认）—— 每 `chunk` 个 cell 抽一次动作：方向取 `geometry.ACTION_DIRS`
+        的 9 类之一（含 idle），速度按轴取 walk_speed × 倍率（[0.8, 1.2] 均匀），
+        与 `infer.py::parse_actions` 的 `ACTION_DIRS[i] * walk_speed * mult` 完全同构，
+        对应 ReactiveGWM 的 `dmd_action_extend_mode: preset_chunks`。逐样本独立抽。
+      * "continue" —— 沿用最后一段已知增量隐含的匀速（等价于把 delta_tok[-1] 一路重复）。
+
+    rng: CPU 上的 torch.Generator。抽样一律在 CPU 上做再搬到 offset_px.device，
+    这样 4 张卡上的种子行为与设备无关；各 rank 抽到不同方向没关系（数据本来就不同），
+    **形状**由 rank-synced 的 N 决定，FSDP 的 all-gather 不会错位。
+    """
+    off = _as_bf2(offset_px, "offset_px")
+    b, f0 = off.shape[:2]
+    if N < 0:
+        raise ValueError(f"N 应 ≥ 0，得到 {N}")
+    want = N + 1
+    if want <= f0:
+        return off[:, :want].contiguous()
+    m = want - f0
+
+    if mode == "continue":
+        v = _implied_last_velocity(off).unsqueeze(1).expand(b, m, 2)  # [b, m, 2]
+    elif mode == "resample":
+        v = _sample_chunk_velocity(b, m, chunk=chunk, rng=rng, walk_speed=walk_speed,
+                                   speed_mult=speed_mult, include_idle=include_idle)
+        v = v.to(device=off.device, dtype=off.dtype)
+    else:
+        raise ValueError(f"未知 mode {mode!r}，可选 'resample' | 'continue'")
+    ext = append_cell_velocity(off, v)
+    assert ext.shape == (b, want, 2), f"延长失败: {tuple(ext.shape)} != {(b, want, 2)}"
+    return ext
+
+
+# --------------------------------------------------------------------------
+# rollout 延长：文本
+# --------------------------------------------------------------------------
+
+def extend_context(context: torch.Tensor, context_ids: torch.Tensor, N: int, mode: str = "hold"):
+    """把 [b, F0, L, C] 的逐 cell 文本与 [b, F0] 的 id 延长/裁剪到 1+N 个 cell。
+
+    返回 (context_ext [b, 1+N, L, C], ids_ext [b, 1+N])，与 `extend_offsets` 逐下标对齐。
+
+    mode="hold"（唯一实现）：cell ≥ F0 一律复制**最后一条已写好的** cell 文本及其 id，
+    照抄 `_wm/WorldMind/distill/code/stage3_dmd/dmd_gm.py:628-637`。ARoPE 臂的 prompt 已被
+    `prompts.strip_action` 剥成"只剩场景"，rollout 又不换场景，所以保持是语义正确的；
+    id 必须跟着一起保持，`arope_forward` 靠它去重 text_embedding，id 错了就会拿错一行文本。
+    转场语料（burn / hardcut）的 cell 20 已经是烧完之后的目标场景，保持它同样正确。
+
+    注意：文本"保持"而偏移"继续走"，两者策略不同是**有意的**（见模块 docstring），
+    不是这里漏了一种模式。
+    """
+    ctx, ids = _check_context(context, context_ids)
+    b, f0 = ctx.shape[:2]
+    if N < 0:
+        raise ValueError(f"N 应 ≥ 0，得到 {N}")
+    if mode != "hold":
+        raise ValueError(f"未知 mode {mode!r}：文本只实现 'hold'（保持最后一条已写好的 cell 文本）")
+    want = N + 1
+    if want <= f0:
+        return ctx[:, :want].contiguous(), ids[:, :want].contiguous()
+    pad_n = want - f0
+    ctx_pad = ctx[:, f0 - 1:f0].expand(b, pad_n, *ctx.shape[2:])
+    ids_pad = ids[:, f0 - 1:f0].expand(b, pad_n)
+    ctx_ext = torch.cat([ctx, ctx_pad], dim=1)
+    ids_ext = torch.cat([ids, ids_pad], dim=1)
+    assert ctx_ext.shape[1] == want and ids_ext.shape[1] == want, (
+        f"文本延长失败: {ctx_ext.shape[1]}/{ids_ext.shape[1]} != {want}")
+    return ctx_ext, ids_ext
+
+
+def assert_rollout_aligned(offset_ext: torch.Tensor, context_ext: torch.Tensor,
+                           ids_ext: torch.Tensor, N: int, latent=None) -> None:
+    """rollout 前的对齐闸：三条（可选四条）流的 cell 数必须都是 1+N，且 batch 一致。
+
+    stage3 每次 rollout 之前、每次 teacher/critic forward 之前都调一下——这类错位
+    不会报错、只会让 loss 照降而训错东西，所以要用断言把它变响。
+    """
+    want = N + 1
+    b = offset_ext.shape[0]
+    if offset_ext.shape != (b, want, 2):
+        raise AssertionError(f"offset 形状 {tuple(offset_ext.shape)} != {(b, want, 2)}")
+    if context_ext.shape[:2] != (b, want):
+        raise AssertionError(f"context 形状 {tuple(context_ext.shape)} 前两维 != {(b, want)}")
+    if tuple(ids_ext.shape) != (b, want):
+        raise AssertionError(f"context_ids 形状 {tuple(ids_ext.shape)} != {(b, want)}")
+    if latent is not None and (latent.shape[0] != b or latent.shape[2] != want):
+        raise AssertionError(f"latent 形状 {tuple(latent.shape)} 与 [b={b}, ·, f={want}, ·, ·] 不符")
+
+
+# --------------------------------------------------------------------------
+# 内部
+# --------------------------------------------------------------------------
+
+def _as_bf2(off: torch.Tensor, name: str) -> torch.Tensor:
+    if not isinstance(off, torch.Tensor):
+        raise TypeError(f"{name} 应为 Tensor，得到 {type(off)}")
+    if off.dim() != 3 or off.shape[-1] != 2:
+        raise ValueError(f"{name} 形状应为 [b, f, 2]，得到 {tuple(off.shape)}")
+    return off
+
+
+def _check_context(context: torch.Tensor, context_ids: torch.Tensor):
+    if context.dim() != 4:
+        raise ValueError(f"context 形状应为 [b, f, L, C]，得到 {tuple(context.shape)}")
+    if context_ids.dim() != 2:
+        raise ValueError(f"context_ids 形状应为 [b, f]，得到 {tuple(context_ids.shape)}")
+    if context.shape[:2] != context_ids.shape:
+        raise ValueError(f"context {tuple(context.shape)[:2]} 与 context_ids "
+                         f"{tuple(context_ids.shape)} 的 (b, f) 不一致")
+    return context, context_ids
+
+
+def _implied_last_velocity(off: torch.Tensor) -> torch.Tensor:
+    """从最后一段已知增量反推匀速等效速度 [b, 2]（px/s）。
+
+    匀速时 delta[k] = 4·v/FPS（k ≥ 2），delta[1] = 2.5·v/FPS；F0 = 1（只有 cell 0）时
+    没有任何增量可用，只能取 0（此时 'continue' 退化成"站住不动"，调用方应当用 'resample'）。
+    """
+    b, f0 = off.shape[:2]
+    if f0 < 2:
+        return off.new_zeros(b, 2)
+    delta = off[:, -1] - off[:, -2]
+    w = _W_CUR if f0 == 2 else (_W_PREV + _W_CUR)  # 2.5（cell 1） / 4.0（cell ≥ 2）
+    return delta * (FPS / w)
+
+
+def _sample_chunk_velocity(b: int, m: int, chunk: int, rng: torch.Generator | None,
+                           walk_speed, speed_mult: tuple[float, float],
+                           include_idle: bool) -> torch.Tensor:
+    """逐样本、每 chunk 个 cell 抽一次动作 → 逐 cell 速度 [b, m, 2]（px/s，float32 CPU）。"""
+    if chunk < 1:
+        raise ValueError(f"chunk 应 ≥ 1，得到 {chunk}")
+    if rng is not None and rng.device.type != "cpu":
+        raise ValueError("rng 必须是 CPU 上的 torch.Generator（抽样固定在 CPU 做，保证跨设备一致）")
+    lo, hi = float(speed_mult[0]), float(speed_mult[1])
+    if not 0.0 <= lo <= hi:
+        raise ValueError(f"speed_mult 应为 0 ≤ lo ≤ hi，得到 {speed_mult!r}")
+    ws = WALK_SPEED_PX_S if walk_speed is None else walk_speed
+    ws = torch.tensor([float(ws[0]), float(ws[1])], dtype=torch.float32)
+    dirs = torch.from_numpy(np.asarray(ACTION_DIRS, dtype=np.float32))  # [9, 2] 单位方向，0 = idle
+    n_dir = dirs.shape[0]
+    lo_idx = 0 if include_idle else 1
+
+    n_chunk = math.ceil(m / chunk)
+    idx = torch.randint(lo_idx, n_dir, (b, n_chunk), generator=rng)          # [b, n_chunk]
+    mult = torch.rand((b, n_chunk, 1), generator=rng) * (hi - lo) + lo       # [b, n_chunk, 1]
+    v_chunk = dirs[idx] * ws.view(1, 1, 2) * mult                            # [b, n_chunk, 2]
+    # 每个 chunk 的速度铺满它覆盖的 cell（最后一个 chunk 可能不满，截断）
+    return v_chunk.repeat_interleave(chunk, dim=1)[:, :m].contiguous()
+
+
+__all__ = [
+    "ACTION_DIRS", "CELLS", "CELL_FRAMES", "DEFAULT_CHUNK_CELLS", "DEFAULT_SPEED_MULT",
+    "WALK_SPEED_PX_S", "append_cell_velocity", "assert_rollout_aligned", "extend_context",
+    "extend_offsets", "make_mask_input", "make_weight_map", "to_device",
+    "walk_speed_from_sidecar", "window_context", "window_offsets",
+]

backup_2026-09-14/code/distill/ema.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python3
+"""生成器 EMA：fp32 影子权重（stage2 CD 的目标网络 / stage2·stage3 的部署权重）。
+
+    from distill.ema import EmaShadow, ema_update
+    ema = EmaShadow(build_arope_dit(...), device="cuda")      # 不进 accelerator.prepare
+    ...
+    ema.update(student, decay=(0.0 if step < args.ema_start else args.ema_decay))
+    runlib.save_dit(accel, ema, runlib.ema_path(run_dir, step), {"step": step, "ema_decay": decay})
+
+★ 影子必须是 fp32。bf16 只有 8 位尾数，相对精度 ~2^-8 ≈ 0.4%：decay=0.99 时
+`pe = 0.99*pe + 0.01*p`，那 1% 的增量落在 pe 尾数之下会被舍入直接吃掉 —— 影子权重在数值上
+纹丝不动，loss 却一切正常（CD 的 loss 本来就从 0 开始涨）。这不是精度损失，是目标网络失效。
+tests/test_runlib.py::test_bf16_shadow_drops_the_update 把这条钉成了断言。
+fp32 影子只是存储；前向照样在 autocast(bf16) 下跑，算力开销不变（多 20 GB 显存，见 PORT_SPEC 预算表）。
+
+★ EmaShadow 故意**不是** nn.Module。accelerate 的 prepare 会 `for obj in result: if isinstance(obj, nn.Module): model = obj`
+（accelerator.py:2229-2236）—— 最后一个 nn.Module 静默胜出，其余一个警告都没有就没被包装。
+影子不是 nn.Module 就绝不可能被误传进 prepare、被卷进优化器的参数表、或者混进上层模块的 state_dict。
+它照样能当模型调用（__call__ 转发给内部 module）。
+
+stage3 的 EMA 是 FSDP 包装的 AropeDiT（要按分片 zip），不走 EmaShadow，直接用模块级的
+ema_update(ema, student, decay)；两条路的数学完全一样。
+"""
+from __future__ import annotations
+
+import copy
+
+import torch
+import torch.nn as nn
+
+
+def _bare(model: nn.Module) -> nn.Module:
+    """剥掉 DDP / DeepSpeedEngine / EmaShadow 的外壳。"""
+    seen = 0
+    while hasattr(model, "module") and isinstance(getattr(model, "module"), nn.Module) and seen < 8:
+        model = model.module
+        seen += 1
+    return model
+
+
+@torch.no_grad()
+def ema_update(ema: nn.Module, student, decay: float) -> None:
+    """逐参数 EMA：`pe = decay*pe + (1-decay)*p`，decay <= 0 ⇒ 硬拷贝。
+
+    zip 成立的前提：两边暴露的是**同一组、同一顺序、同一分片方式**的 Parameter。
+    ZeRO-2 不分片参数，天然成立；FSDP 要 fsdp_use_orig_params=true（PORT_SPEC R5）。
+    这里额外核对个数与逐个形状，顺序错位在绝大多数情况下会当场炸出来，而不是静默训歪。
+    student 可以是 DeepSpeedEngine / DDP —— nn.Module 外壳的 parameters() 顺序与内部一致。"""
+    pe_list = list(ema.parameters())
+    ps_list = list(student.parameters())
+    if len(pe_list) != len(ps_list):
+        raise RuntimeError(f"EMA 参数个数不匹配：影子 {len(pe_list)} vs 学生 {len(ps_list)}")
+    for i, (pe, p) in enumerate(zip(pe_list, ps_list)):
+        if pe.shape != p.shape:
+            raise RuntimeError(f"EMA 第 {i} 个参数形状不匹配：影子 {tuple(pe.shape)} vs 学生 {tuple(p.shape)}")
+        src = p.data
+        if src.device != pe.device:
+            src = src.to(pe.device, non_blocking=True)
+        if decay <= 0.0:
+            pe.data.copy_(src)                    # 硬拷贝（step < ema_start）
+        else:
+            # dtype 由 TensorIterator 就地提升：src 是 bf16 也不会额外分配一份 fp32
+            pe.data.mul_(decay).add_(src, alpha=1.0 - decay)
+
+
+class EmaShadow:
+    """fp32 影子权重。不是 nn.Module（见文件头），但能当模型直接调用。
+
+    EmaShadow(module)               —— deepcopy 一份，源模型归调用方（多占一份显存/内存）
+    EmaShadow(module, own=True)     —— 直接接管这个 module（stage2 已经单独建了第三个模型时用这个，
+                                       不再多复制一份 5B 权重）
+    device="cpu" 可以把影子放在主存里（ReactiveGWM 的 *_ema_shadow_device: cpu），
+    update 会自动把学生的分片搬过来；但前向就得先 .to("cuda")，stage2 默认放卡上。
+    """
+
+    def __init__(self, module: nn.Module, *, device=None, own: bool = False, dtype: torch.dtype = torch.float32):
+        if dtype != torch.float32:
+            raise ValueError("影子必须是 fp32：bf16 在 decay=0.99 下会把 1% 的增量整个舍掉（见文件头）")
+        src = _bare(module)
+        self.module = src if own else copy.deepcopy(src)
+        self.module.to(dtype=dtype)
+        if device is not None:
+            self.module.to(device=device)
+        self.module.requires_grad_(False).eval()
+        self.dtype = dtype
+        self._assert_fp32()
+
+    # ---- 数值 ----
+    def _assert_fp32(self) -> None:
+        bad = [n for n, p in self.module.named_parameters() if p.dtype is not torch.float32]
+        if bad:
+            raise RuntimeError(f"影子权重不是 fp32：{bad[:5]}…({len(bad)})")
+
+    @torch.no_grad()
+    def update(self, student, decay: float) -> None:
+        """一个 optimizer step 调一次（不是每个 micro-batch）。decay <= 0 ⇒ 硬拷贝。"""
+        ema_update(self.module, student, decay)
+
+    @torch.no_grad()
+    def hard_copy(self, student) -> None:
+        """影子 ← 学生。step < ema_start 期间每步都调，等价于 update(student, 0.0)。"""
+        ema_update(self.module, student, 0.0)
+
+    # ---- 当模型用 ----
+    def __call__(self, *args, **kwargs):
+        return self.module(*args, **kwargs)
+
+    def forward(self, *args, **kwargs):
+        return self.module(*args, **kwargs)
+
+    def parameters(self, recurse: bool = True):
+        return self.module.parameters(recurse=recurse)
+
+    def named_parameters(self, *args, **kwargs):
+        return self.module.named_parameters(*args, **kwargs)
+
+    def buffers(self, recurse: bool = True):
+        return self.module.buffers(recurse=recurse)
+
+    def to(self, *args, **kwargs):
+        """只搬设备。dtype 不许动：.to(torch.bfloat16) 会当场毁掉目标网络。
+
+        注意 dit.freqs 是普通的 tuple 属性（不是 buffer），Module.to 搬不动它 ——
+        影子搬完设备后调用方要照 train.py:461 那样补一句
+        `ema.module.dit.freqs = tuple(t.to(device) for t in ema.module.dit.freqs)`。"""
+        if any(isinstance(a, torch.dtype) for a in args) or kwargs.get("dtype") is not None:
+            raise RuntimeError("影子权重的 dtype 不许改（fp32），只搬设备")
+        self.module.to(*args, **kwargs)
+        return self
+
+    def eval(self):
+        self.module.eval()
+        return self
+
+    def train(self, mode: bool = True):
+        """影子永远是 eval —— 显式挡住 train()，免得被跟着学生一起 .train() 打开 dropout。"""
+        if mode:
+            raise RuntimeError("EMA 影子不训练，别对它调 train()")
+        return self.eval()
+
+    # ---- 导出 ----
+    def state_dict(self, *args, **kwargs):
+        """直接透传内部 module 的键集（不带 "module." 前缀），
+        这样 runlib.save_dit(accel, ema, ...) 剥掉 `dit.` 之后得到的就是标准的 827 键。"""
+        return self.module.state_dict(*args, **kwargs)
+
+    def load_state_dict(self, *args, **kwargs):
+        out = self.module.load_state_dict(*args, **kwargs)
+        self._assert_fp32()
+        return out
+
+    def export_state_dict(self, dtype: torch.dtype = torch.bfloat16) -> dict:
+        """给不走 runlib.save_dit 的场合用：CPU + contiguous + 指定 dtype。"""
+        return {k: v.detach().to(device="cpu", dtype=dtype).contiguous()
+                for k, v in self.module.state_dict().items()}
+
+    def __repr__(self):
+        n = sum(p.numel() for p in self.module.parameters())
+        dev = next(self.module.parameters()).device if n else "?"
+        return f"EmaShadow({type(self.module).__name__}, {n:,} params, fp32, {dev})"

backup_2026-09-14/code/distill/eval/anystep_probe.py

@@ -0,0 +1,196 @@
+"""stage2-flowmap 的**必跑前置门**：单卡量一个真实 micro batch。
+
+`distill/stage2_flowmap/config.json` 的 `_gate` 把这个探针列为 launch 前的必要条件。理由很直接：
+`~75 GB/卡` 和 `~14 s/step` 到此为止都是**推算**，从没有任何代码在真实几何（f=21、16380 token、
+5B 权重、4 次前向 + 1 次反向）上量过。用一张卡量到，好过用 4 张（或 8 张）卡去撞 OOM。
+
+三件事，都在 GPU 0 上、都不写权重：
+
+  S-FM10  真权重恒等 —— 装上 flow map 之后，对 r ∈ {t, 0, 250, 1000} 与装之前逐位相同。
+          这是整条链的回归锚点：它一旦不成立，后面训的就不是 stage1b 那个教师了。
+  S-FM11  一个真实 micro batch 的峰值显存与墙钟（走 **fm_loss 本体**，不是复刻一份）。
+  S-FM12  中心差分的 bf16 噪声底 —— 实测 (u_p − u_m) 的信噪比。fidelity 评审记录过：
+          探针在生产里被量化到 bf16，实际间隔是 8–12 而非名义的 10，分母却用字面 2·eps，
+          于是 dF/dt 有平均 16.6 % 的系统性尺度误差。这个是**继承自上游、故意不改**的，
+          但它到底有多大，得有个数。
+
+    CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/data/zhiyangdeng/ActionRoPE DIFFSYNTH_SKIP_DOWNLOAD=True \
+      .venv/bin/python distill/eval/anystep_probe.py
+
+单卡跑不出 4/8 卡的 ZeRO-2 分片行为（优化器状态是分片的），所以这里报的是**单进程上界**：
+参数 + 梯度 + 激活 + Adam 全在一张卡上。多卡只会更省。反过来说，**单卡装不下不代表多卡装不下**，
+所以 OOM 时脚本会把分片后的估算一并打出来，而不是直接判死刑。
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+
+os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import torch
+
+ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from actionrope.dataset import AropeLatentDataset, arope_collate, load_sidecar  # noqa: E402
+from distill import flowmap as fm  # noqa: E402
+from distill import runlib  # noqa: E402
+from distill.arope_dit import AropeDiT, unwrap  # noqa: E402
+from distill.data import to_device  # noqa: E402
+from distill.stage2_flowmap.fm_arope import _tf_probe, fm_loss  # noqa: E402
+
+GB = 1024 ** 3
+
+
+def _mem(dev) -> float:
+    return torch.cuda.max_memory_allocated(dev) / GB if dev.type == "cuda" else 0.0
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--init_ckpt", default=f"{ROOT}/outputs/distill/stage1b_adapt/dits/dit_step3000.safetensors")
+    ap.add_argument("--train_dir", default=f"{ROOT}/data/latent/train_eybx")
+    ap.add_argument("--sidecar", default=f"{ROOT}/data/latent/arope_sidecar.pt")
+    ap.add_argument("--text_table", default=f"{ROOT}/data/latent/text_table_actionrope.pt")
+    ap.add_argument("--flowmap_form", default="residual", choices=["residual", "convex"])
+    ap.add_argument("--delta_gate", type=float, default=1.0)
+    ap.add_argument("--flowmap_epsilon", type=float, default=5.0)
+    ap.add_argument("--cfg_scale", type=float, default=3.0, help="与 run.sh 默认一致：开着文本 fuse 才量得到真实前向数")
+    ap.add_argument("--iters", type=int, default=3, help="计时取后 iters-1 次（第 1 次含 cudnn/分配器预热）")
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--skip_backward", action="store_true", help="只量前向（调试用；默认量 fwd+bwd）")
+    args = ap.parse_args()
+
+    dev = torch.device(args.device)
+    log = print
+
+    # ---------------- S-FM10 真权重恒等 ----------------
+    log("=" * 70)
+    log("S-FM10  真权重恒等（装 flow map 前后，对任意 r 逐位相同）")
+    from safetensors.torch import load_file
+    model = AropeDiT(causal=True, mask_channel=True).to(dev).eval()
+    unwrap(model).dit.load_state_dict(load_file(args.init_ckpt), strict=True)
+    model = model.to(torch.bfloat16)
+    unwrap(model).move_rope_table(dev)
+    n_before = len(unwrap(model).dit.state_dict())
+
+    pre = _tf_probe(unwrap(model), 500.0, None)          # 装之前的参照
+    unwrap(model).install_flowmap(form=args.flowmap_form, gate=args.delta_gate)
+    unwrap(model).to(dev)
+    n_after = len(unwrap(model).dit.state_dict())
+
+    worst = 0.0
+    for r in (500.0, 0.0, 250.0, 1000.0):
+        post = _tf_probe(unwrap(model), 500.0, r)
+        d = (post.float() - pre.float()).abs().max().item()
+        worst = max(worst, d)
+        log(f"   t=500 r={r:7.1f}   max|Δ| = {d:.3e}")
+    log(f"   键数 {n_before} → {n_after}   最大偏差 {worst:.3e}")
+    if args.flowmap_form == "residual" and worst != 0.0:
+        raise SystemExit(f"S-FM10 失败：residual 形式必须对任意 r ��位相同，实测 {worst:.3e}。"
+                         f"这说明 delta 分支的读出层没有真的零初始化 —— 别往下跑。")
+    log("   ✅ 通过" if worst == 0.0 else f"   ⚠️ convex 形式 off-diagonal 有偏差（预期内）")
+
+    # ---------------- 真实 micro batch ----------------
+    log("=" * 70)
+    log("S-FM11  一个真实 micro batch 的显存与墙钟（走 fm_loss 本体）")
+    sc = load_sidecar(args.sidecar)
+    ds = AropeLatentDataset(args.train_dir, sc, args.text_table, text_mode="scene",
+                            scene_dropout=0.0, limit=8, min_rt_ratio=0.8, verbose=False)
+    batch = to_device(arope_collate([ds[0]]), dev)
+    log(f"   latent {tuple(batch['input_latents'].shape)}  context {tuple(batch['context'].shape)}")
+
+    uncond = None
+    if abs(args.cfg_scale - 1.0) >= 1e-6:
+        from actionrope.prompts import NULL_PROMPT
+        table = torch.load(args.text_table, map_location="cpu", weights_only=False)
+        row = table[NULL_PROMPT].to(torch.bfloat16)
+        uncond = row.view(1, 1, *row.shape).to(dev).expand(1, batch["context"].shape[1], -1, -1)
+    gains = {"text": args.cfg_scale - 1.0, "anchor": 0.0, "action": 0.0, "hist": 0.0, "ctx": 0.0}
+
+    model.train()
+    for p in unwrap(model).parameters():
+        p.requires_grad_(True)
+
+    # 全局 batch 必须按真实值抽：G=1 时 round(0.5*1)=0，diffusion 槽为空，sample_tr 会（正确地）断言。
+    # 真实跑法是 4 卡 GA2 或 8 卡 GA1，两者全局 batch 都是 8。这里用 world_size=8 复现，
+    # 并遍历 rank 覆盖三种角色 —— 三者代价不同：diffusion 槽（r==t）的中心差分不进 loss，
+    # 自由 (t,r) 槽最贵。报告取**最贵的那个**，因为 OOM 只需要发生一次。
+    WORLD = 8
+    gen = torch.Generator(device="cpu").manual_seed(0)
+    split_gen = fm.make_split_generator(0)
+    per_role: dict[int, list] = {}
+    ROLE_NAME = {fm.ROLE_DIFFUSION: "diffusion(r==t)", fm.ROLE_CONSISTENCY: "consistency(r=0)",
+                 fm.ROLE_RANDOM: "random(t,r)"}
+    n_iter = max(args.iters, 1)
+    for i in range(n_iter * WORLD):
+        rank = i % WORLD
+        g_i = torch.Generator(device="cpu").manual_seed(1000 + i)
+        sg_i = fm.make_split_generator(i)
+        tr = fm.sample_tr(1, world_size=WORLD, rank=rank, gen=g_i, split_gen=sg_i,
+                          diffusion_ratio=0.5, consistency_ratio=0.25, rotate=True, device=dev)
+        role = int(tr.role[0]) if hasattr(tr, "role") else fm.ROLE_RANDOM
+        torch.cuda.reset_peak_memory_stats(dev)
+        unwrap(model).zero_grad(set_to_none=True)
+        torch.cuda.synchronize(dev)
+        t0 = time.time()
+        loss_b, stats = fm_loss(model, batch, tr, gains=gains,
+                                uncond=uncond,
+                                uncond_ids=(None if uncond is None else torch.zeros_like(batch["context_ids"])),
+                                epsilon=args.flowmap_epsilon, eps_clamp=True,
+                                use_gradient_checkpointing=True)
+        loss = loss_b.mean()
+        if not args.skip_backward:
+            loss.backward()
+        torch.cuda.synchronize(dev)
+        dt, pk = time.time() - t0, _mem(dev)
+        per_role.setdefault(role, []).append((dt, pk, float(loss)))
+        if i < WORLD:                                   # 第一轮是预热，只打印不计入
+            log(f"   预热 rank{rank}  {ROLE_NAME.get(role, role):18s} {dt:5.2f}s  {pk:5.1f} GB")
+
+    log("   —— 稳态（去掉第一轮预热）——")
+    t_micro, pk = 0.0, 0.0
+    for role, rows in sorted(per_role.items()):
+        warm = rows[1:] or rows
+        tt = sum(r[0] for r in warm) / len(warm)
+        pp = max(r[1] for r in warm)
+        ll = sum(r[2] for r in warm) / len(warm)
+        log(f"   {ROLE_NAME.get(role, role):18s} n={len(warm)}  {tt:5.2f}s  峰值 {pp:5.1f} GB  loss {ll:.4f}")
+        t_micro, pk = max(t_micro, tt), max(pk, pp)
+    log(f"   最贵角色: {t_micro:.2f}s / micro   峰值显存 {pk:.1f} GB（单进程，未分片）")
+
+    # ---------------- 换算 ----------------
+    log("=" * 70)
+    log("换算（micro → step；等效 batch 恒为 8）")
+    for ngpu, ga in ((4, 2), (8, 1)):
+        s_step = t_micro * ga
+        log(f"   {ngpu} 卡 × GA{ga}:  {s_step:5.2f} s/step   8000 步 ≈ {s_step * 8000 / 3600:5.1f} h")
+    log("   注：显存是单进程上界。ZeRO-2 把 Adam 状态按卡分片，实际每卡应低于这个数；")
+    log("   反过来，这里 OOM 也不代表多卡装不下。")
+
+    # ---------------- S-FM12 bf16 噪声底 ----------------
+    log("=" * 70)
+    log("S-FM12  中心差分的 bf16 噪声底")
+    log("   已知问题（继承自上游、故意不改）：探针时间步在生产里被量化到 bf16，")
+    log("   实际间隔不是名义的 2·eps，分母却按 2·eps 除 ⇒ dF/dt 有系统性尺度误差。")
+    eps = args.flowmap_epsilon
+    for t_val in (100.0, 500.0, 900.0):
+        tt = torch.full((1,), t_val, device=dev)
+        tp = (tt + eps).to(torch.bfloat16).float()
+        tm = (tt - eps).to(torch.bfloat16).float()
+        real = float((tp - tm).item())
+        log(f"   t={t_val:6.1f}  名义间隔 {2 * eps:.1f}  bf16 实际 {real:5.2f}  "
+            f"尺度误差 {100 * (real - 2 * eps) / (2 * eps):+6.1f}%")
+    log("   （--eps_clamp 1 会按**实际**步长除，把这一项消掉；--eps_clamp 0 复现上游）")
+
+    log("=" * 70)
+    log(f"探针通过。stage2-flowmap 可以起训。峰值 {pk:.1f} GB / micro {t_micro:.2f}s")
+
+
+if __name__ == "__main__":
+    main()

backup_2026-09-14/code/distill/eval/flowmap_liveness.py

@@ -0,0 +1,183 @@
+"""flow map 到底有没有在学 —— 「一晚上空跑」的直接排查。
+
+零初始化的代价：装上 flow map 的那一刻，`delta_embedding` 的读出层全零，于是
+`u(x, t, r) ≡ u(x, t, t)` 对**任意** r 成立。这正是教师同一性门要的性质，但它也意味着
+**此刻模型对 r 完全没有依赖**。r 依赖必须靠训练长出来。
+
+如果梯度进不去那条分支，训练不会报错、loss 照样下降（因为 50 % 的样本是 r==t 的
+diffusion 槽，那本来就是普通扩散 loss），最后得到的只是一个普通扩散模型 —— 花了一晚上
+8 张卡，any-step 这件事一步没学到。
+
+这个脚本用**真权重、真 fm_loss**跑几步真优化，然后回答三个问题：
+
+  L1  delta_embedding 的读出层还是不是全零？（是 ⇒ 梯度根本没进去，死的）
+  L2  梯度范数有没有真的落在 delta 分支上？（与主干比，量级对不对）
+  L3  几步之后 u(x,t,r) 对不同 r 是否已经**分开**？（分不开 ⇒ 学不到任何 r 依赖）
+
+    CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/data/zhiyangdeng/ActionRoPE DIFFSYNTH_SKIP_DOWNLOAD=True \
+      .venv/bin/python distill/eval/flowmap_liveness.py --steps 12
+
+也可以直接查一个已经存下来的 checkpoint（训练跑起来之后用这个，最省事、也最有说服力）：
+
+    ... .venv/bin/python distill/eval/flowmap_liveness.py \
+        --ckpt outputs/distill/stage2_flowmap/ema/ema_step500.safetensors
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import torch
+
+ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from actionrope.dataset import AropeLatentDataset, arope_collate, load_sidecar  # noqa: E402
+from distill import flowmap as fm  # noqa: E402
+from distill.arope_dit import AropeDiT, unwrap  # noqa: E402
+from distill.data import to_device  # noqa: E402
+from distill.stage2_flowmap.fm_arope import _tf_probe, fm_loss  # noqa: E402
+
+DELTA_KEYS = ["delta_embedding.0.weight", "delta_embedding.0.bias",
+              "delta_embedding.2.weight", "delta_embedding.2.bias"]
+
+
+def _delta_stats(dit) -> dict:
+    sd = dit.state_dict()
+    out = {}
+    for k in DELTA_KEYS:
+        if k in sd:
+            v = sd[k].float()
+            out[k] = (float(v.abs().max()), float(v.abs().mean()))
+    if "delta_gate" in sd:
+        out["delta_gate"] = (float(sd["delta_gate"].float().abs().max()), 0.0)
+    return out
+
+
+def _r_spread(module, t_val: float = 500.0) -> float:
+    """同一个 (x, t) 下，不同 r 的输出最大分离度。0 ⇒ 模型对 r 毫无依赖。"""
+    base = _tf_probe(module, t_val, t_val)
+    worst = 0.0
+    for r in (0.0, 250.0, 750.0, 1000.0):
+        d = (_tf_probe(module, t_val, r).float() - base.float()).abs().max().item()
+        worst = max(worst, d)
+    return worst
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--ckpt", default=None, help="直接查一个已存的 flow-map ckpt（832 键），不跑优化")
+    ap.add_argument("--init_ckpt", default=f"{ROOT}/outputs/distill/stage1b_adapt/dits/dit_step3000.safetensors")
+    ap.add_argument("--train_dir", default=f"{ROOT}/data/latent/train_eybx")
+    ap.add_argument("--sidecar", default=f"{ROOT}/data/latent/arope_sidecar.pt")
+    ap.add_argument("--text_table", default=f"{ROOT}/data/latent/text_table_actionrope.pt")
+    ap.add_argument("--steps", type=int, default=12)
+    ap.add_argument("--lr", type=float, default=1e-5)
+    ap.add_argument("--device", default="cuda")
+    args = ap.parse_args()
+
+    dev = torch.device(args.device)
+    log = print
+    from safetensors.torch import load_file
+
+    # ---------------- 只查 checkpoint ----------------
+    if args.ckpt:
+        sd = load_file(args.ckpt)
+        log(f"checkpoint: {args.ckpt}")
+        log(f"  键数 {len(sd)}")
+        dead = True
+        for k in DELTA_KEYS + ["delta_gate"]:
+            if k not in sd:
+                log(f"  ✗ 缺 {k} —— 这不是 flow-map checkpoint")
+                continue
+            v = sd[k].float()
+            mx, mn = float(v.abs().max()), float(v.abs().mean())
+            log(f"  {k:32s} |max| {mx:.3e}  |mean| {mn:.3e}")
+            if k == "delta_embedding.2.weight":
+                dead = (mx == 0.0)
+        log("")
+        if dead:
+            log("  ❌ delta_embedding.2.weight 仍然全零 ⇒ r 分支是死的，训练没在学 any-step")
+            raise SystemExit(1)
+        log("  ✅ 读出层已非零 ⇒ 梯度确实进了 delta 分支，模型在学 r 依赖")
+        return
+
+    # ---------------- 跑真优化 ----------------
+    model = AropeDiT(causal=True, mask_channel=True).to(dev).eval()
+    unwrap(model).dit.load_state_dict(load_file(args.init_ckpt), strict=True)
+    model = model.to(torch.bfloat16)
+    unwrap(model).move_rope_table(dev)
+    unwrap(model).install_flowmap(form="residual", gate=1.0)
+    unwrap(model).to(dev)
+    bare = unwrap(model)
+
+    log("=" * 72)
+    log("L1  装上之后（训练前）")
+    for k, (mx, mn) in _delta_stats(bare.dit).items():
+        log(f"    {k:32s} |max| {mx:.3e}")
+    spread0 = _r_spread(bare)
+    log(f"    不同 r 的输出分离度 = {spread0:.3e}   （0 是**对的**：零初始化 ⇒ 此刻没有 r 依赖）")
+
+    sc = load_sidecar(args.sidecar)
+    ds = AropeLatentDataset(args.train_dir, sc, args.text_table, text_mode="scene",
+                            scene_dropout=0.0, limit=8, min_rt_ratio=0.8, verbose=False)
+    from actionrope.prompts import NULL_PROMPT
+    table = torch.load(args.text_table, map_location="cpu", weights_only=False)
+    row = table[NULL_PROMPT].to(torch.bfloat16)
+
+    model.train()
+    for p in bare.parameters():
+        p.requires_grad_(True)
+    opt = torch.optim.AdamW(bare.parameters(), lr=args.lr, betas=(0.9, 0.999), weight_decay=0.0)
+
+    delta_names = [n for n, _ in bare.dit.named_parameters() if n.startswith("delta_embedding")]
+    log("=" * 72)
+    log(f"L2  跑 {args.steps} 步真 fm_loss（lr={args.lr}），看梯度落在哪")
+    split_gen = fm.make_split_generator(0)
+    for i in range(args.steps):
+        batch = to_device(arope_collate([ds[i % len(ds)]]), dev)
+        uncond = row.view(1, 1, *row.shape).to(dev).expand(1, batch["context"].shape[1], -1, -1)
+        g_i = torch.Generator(device="cpu").manual_seed(500 + i)
+        tr = fm.sample_tr(1, world_size=8, rank=i % 8, gen=g_i, split_gen=fm.make_split_generator(i),
+                          diffusion_ratio=0.5, consistency_ratio=0.25, rotate=True, device=dev)
+        opt.zero_grad(set_to_none=True)
+        loss_b, stats = fm_loss(model, batch, tr, gains={"text": 2.0}, uncond=uncond,
+                                uncond_ids=torch.zeros_like(batch["context_ids"]),
+                                epsilon=5.0, eps_clamp=True, use_gradient_checkpointing=True)
+        loss_b.mean().backward()
+        gd = torch.sqrt(sum((p.grad.float() ** 2).sum() for n, p in bare.dit.named_parameters()
+                            if n.startswith("delta_embedding") and p.grad is not None))
+        ga = torch.sqrt(sum((p.grad.float() ** 2).sum() for p in bare.parameters() if p.grad is not None))
+        role = ["diffusion", "consistency", "random"][int(tr.role[0])]
+        log(f"    step {i + 1:2d}  {role:12s} t={float(tr.t[0]):6.1f} r={float(tr.r[0]):6.1f}  "
+            f"loss {float(loss_b.mean()):.4f}  |g_delta| {float(gd):.3e}  |g_all| {float(ga):.3e}")
+        opt.step()
+
+    log("=" * 72)
+    log("L1'  训练之后")
+    dead = False
+    for k, (mx, mn) in _delta_stats(bare.dit).items():
+        log(f"    {k:32s} |max| {mx:.3e}")
+        if k == "delta_embedding.2.weight" and mx == 0.0:
+            dead = True
+    log("=" * 72)
+    log("L3  r 依赖有没有长出来")
+    spread1 = _r_spread(bare)
+    log(f"    训练前 {spread0:.3e}  →  训练后 {spread1:.3e}")
+
+    log("=" * 72)
+    if dead or spread1 <= spread0:
+        log("❌ 机制没通：delta 分支没动 / 不同 r 的输出仍然一样。")
+        log("   继续跑下去只会得到一个普通扩散模型，any-step 一步没学到。停机排查。")
+        raise SystemExit(1)
+    log("✅ 机制通了：梯度进了 delta 分支，读出层已非零，不同 r 的输出已经分开。")
+    log(f"   注意 r 依赖是**从零长出来**的，前期分离度小是正常的（这几步就有 {spread1:.1e}）。")
+
+
+if __name__ == "__main__":
+    main()

backup_2026-09-14/code/distill/eval/loop_spec.txt

@@ -0,0 +1 @@
+idle:1.0:1,right:1.0:10,down:1.0:10,left:1.0:10,up:1.0:10

backup_2026-09-14/code/distill/eval/null_branch_probe.py

@@ -0,0 +1,113 @@
+"""四条 null 分支到底训出来了没有 —— stage1b 适配微调的验收探针。
+
+适配微调的**唯一目的**是把 text 之外的三条 null 分支变成分布内的合法输入。val 不掉只说明
+主目标没被破坏，说明不了 null 分支能不能用。这里直接量：
+
+    同一批 clip、同一份噪声、同一组噪声档，对每条分支各算一次 loss。
+
+  微调前（stage1 的 ckpt）  null 分支从没训过 ⇒ loss 应显著高、甚至爆掉
+  微调后（stage1b 的 ckpt）  null 分支已训 ⇒ loss 明显下降，而 full 那条基本不动
+
+`Δloss = loss(null) − loss(full)` 还有第二重意义：它就是"这条条件带来多少信息量"。
+Δ 越大说明该条件越重要 —— 于是能直接回答"该给哪条轴多大的 guidance 权重"。
+
+    CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/data/zhiyangdeng/ActionRoPE DIFFSYNTH_SKIP_DOWNLOAD=True \
+      .venv/bin/python distill/eval/null_branch_probe.py \
+        --before outputs/distill/stage1_ar_tf/dits/dit_step10000.safetensors \
+        --after  outputs/distill/stage1b_adapt/dits/dit_step3000.safetensors
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import torch
+
+os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
+ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from actionrope.dataset import AropeLatentDataset, arope_collate, load_sidecar  # noqa: E402
+from distill.arope_dit import AropeDiT, unwrap  # noqa: E402
+from distill.data import to_device  # noqa: E402
+from distill.flow_pf import sample_pf_timestep_ids  # noqa: E402
+from distill.stage1_ar.ar_tf import ar_tf_loss  # noqa: E402
+
+# 四条分支：名字 → 传给 ar_tf_loss 的 kwargs（p=1.0 ⇒ 强制命中）
+BRANCHES = {
+    "full":   {},
+    "anchor": {"anchor_dropout": 1.0},          # ★ 摘掉 cell 0、历史保留
+    "action": {"action_dropout": 1.0},          # offset ≡ 0
+    "hist":   {"hist_dropout": 1.0},            # 只剩 cell 0
+    "ctx":    {"ctx_dropout": 1.0},             # 什么都看不到
+}
+
+
+@torch.no_grad()
+def probe(ckpt: str, batches, seeds, device: str) -> dict[str, float]:
+    from safetensors.torch import load_file
+    m = AropeDiT(causal=True, mask_channel=True).to(device).eval()
+    unwrap(m).dit.load_state_dict(load_file(ckpt), strict=True)
+    m = m.to(torch.bfloat16)
+    unwrap(m).move_rope_table(device)
+
+    out = {}
+    for name, kw in BRANCHES.items():
+        tot = 0.0
+        for batch, seed in zip(batches, seeds):
+            b = to_device(batch, device)
+            x0 = b["input_latents"]
+            g = torch.Generator(device="cpu").manual_seed(seed)
+            noise = torch.randn(x0.shape, generator=g).to(dtype=x0.dtype, device=device)
+            idx = sample_pf_timestep_ids(x0.shape[0], x0.shape[2], device,
+                                         torch.Generator(device="cpu").manual_seed(seed + 7))
+            loss, _ = ar_tf_loss(m, b, noise=noise, idx_pf=idx, **kw)
+            tot += float(loss)
+        out[name] = tot / len(batches)
+    del m
+    torch.cuda.empty_cache()
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--before", required=True, help="适配微调**前**的 ckpt")
+    ap.add_argument("--after", required=True, help="适配微调**后**的 ckpt")
+    ap.add_argument("--n", type=int, default=8, help="用多少条 val clip")
+    ap.add_argument("--val_dir", default=f"{ROOT}/data/latent/val_eybx")
+    ap.add_argument("--sidecar", default=f"{ROOT}/data/latent/arope_sidecar.pt")
+    ap.add_argument("--text_table", default=f"{ROOT}/data/latent/text_table_actionrope.pt")
+    ap.add_argument("--device", default="cuda")
+    args = ap.parse_args()
+
+    sc = load_sidecar(args.sidecar)
+    ds = AropeLatentDataset(args.val_dir, sc, args.text_table, text_mode="scene",
+                            scene_dropout=0.0, limit=args.n, verbose=False)
+    batches = [arope_collate([ds[i]]) for i in range(min(args.n, len(ds)))]
+    seeds = list(range(len(batches)))
+    print(f"用 {len(batches)} 条 val clip，固定噪声/噪声档，每条分支各算一次\n")
+
+    res = {}
+    for tag, ck in (("before", args.before), ("after", args.after)):
+        res[tag] = probe(ck, batches, seeds, args.device)
+        print(f"  {tag}: {os.path.basename(ck)}  已量完")
+
+    print(f"\n{'分支':<8} {'微调前':>10} {'微调后':>10} {'变化':>10}   {'Δ vs full(后)':>14}")
+    print("-" * 62)
+    full_after = res["after"]["full"]
+    for name in BRANCHES:
+        b, a = res["before"][name], res["after"][name]
+        d = a - full_after
+        note = "" if name == "full" else f"{d:+.4f}"
+        print(f"{name:<8} {b:10.4f} {a:10.4f} {a - b:+10.4f}   {note:>14}")
+
+    print("\n读法：")
+    print("  · full 那行「变化」应接近 0 —— 补 null 分支不该破���主目标")
+    print("  · 其余三行「变化」应显著为负 —— 说明这些从没训过的输入现在训出来了")
+    print("  · 「Δ vs full」= 该条件的信息量，越大越重要，可直接用来定 guidance 权重的相对大小")
+
+
+if __name__ == "__main__":
+    main()

backup_2026-09-14/code/distill/eval/paired50.sh

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+# 50 张首帧的配对评测：同首帧、同动作脚本、同种子，只换权重。
+#
+#   bash distill/eval/paired50.sh <ckptA> <ckptB> [tagA] [tagB] [maxjobs]
+#
+# 为什么必须 50 张而不是 4 张：同一份权重在不同首帧上的净漂移散布是 9~36 px（sd≈11），
+# 而各 stage3 变体之间的差异只有 1~3 px。n=4 时均值标准误 5.5 px —— 判据分辨不出效应，
+# 实测四个 ‖Δ‖ 差几个量级的权重全部落在 22~26。n=50 时标准误降到 ~1.6 px，才谈得上比较。
+set -euo pipefail
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+A="${1:?ckptA}"; B="${2:?ckptB}"; TA="${3:-A}"; TB="${4:-B}"; MAXJOBS="${5:-6}"
+SPEC_FILE="${SPEC_FILE:-$ROOT/distill/eval/loop_spec.txt}"   # 41 cell 净零闭环（重建，见 sweep8.sh 文件头）
+SPEC="$(cat "$SPEC_FILE")"
+# stage3 学生必须按**训练几何**评（sink 0 / window 20 / cap 20 = rollout_eval.py 默认，GEO-1）；
+# sink 3 / w17 只作显式 A/B：SINK=3 KV_WINDOW=17 ROPE_CAP=20 bash distill/eval/paired50.sh ...
+# 几何旗标只在环境变量显式给出时才追加；默认一个都不传 ⇒ rollout_eval.py 的默认 sink 0 / window 20 /
+# cap 20 = stage3 训练几何。三者必须满足 rope_cap == sink + window（rollout_eval 会核）。
+GEO=()
+if [ -n "${SINK:-}" ]; then GEO+=(--sink_size "$SINK"); fi
+if [ -n "${KV_WINDOW:-}" ]; then GEO+=(--kv_window "$KV_WINDOW"); fi
+if [ -n "${ROPE_CAP:-}" ]; then GEO+=(--rope_cap "$ROPE_CAP"); fi
+OUT="$ROOT/outputs/distill/eval/paired"; mkdir -p "$OUT"
+export PYTHONPATH="$ROOT" DIFFSYNTH_SKIP_DOWNLOAD=True CUDA_DEVICE_ORDER=PCI_BUS_ID
+i=0
+for f in "$ROOT"/data/eval/eybx_bg_50/*.png; do
+  n="$(basename "$f" .png)"
+  for tc in "$TA:$A" "$TB:$B"; do
+    t="${tc%%:*}"; c="${tc#*:}"
+    [ -s "$OUT/${t}_${n}.log" ] && continue
+    CUDA_VISIBLE_DEVICES=$(( i % 8 )) nohup "$ROOT/.venv/bin/python" "$ROOT/distill/eval/rollout_eval.py" \
+      --ckpt "$c" --mode rollout --cells 41 "${GEO[@]}" \
+      --num_inference_steps 4 --actions "$SPEC" --allow_busy_gpu --seed 0 \
+      --first_frame "$f" --tiled --no-save_video \
+      --out "$OUT/${t}_${n}.mp4" > "$OUT/${t}_${n}.log" 2>&1 &
+    i=$(( i + 1 )); (( i % MAXJOBS == 0 )) && wait
+  done
+done
+wait
+echo "[paired] 完成，结果在 $OUT"

backup_2026-09-14/code/distill/eval/photodrift.sh

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+# 一个 ckpt → 41-cell rollout → 逐 cell 的像素 std/mean 表。
+#
+# stage3 的失败模式是**光度漂移**（正式 run 烧白 std 39→2.7 mean 60→171；冻 critic 过饱和
+# std 39→80），不是结构噪声，所以判据就是这两条曲线平不平。基线 stage2-7500 是 std 37~40、
+# mean 60.7→42.9。sift 能量出位移就更好，量不出（sift=None）本身就是「已经坏了」。
+# ⚠ 上面的基线数字是在旧几何（sink 3 / w17）+ 原始 loop spec 下量的；脚本现已默认训练几何（sink 0 / w20）
+#   + 重建的 loop_spec.txt，严格说不可比 —— stage2 跑完后用新默认重测 stage2 EMA 并更新这里（GEO-1）。
+#
+#   bash distill/eval/photodrift.sh <ckpt> <tag> [gpu] [steps]
+set -euo pipefail
+CK="$1"; TAG="$2"; GPU="${3:-0}"; N="${4:-4}"
+ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+SPEC_FILE="${SPEC_FILE:-$ROOT/distill/eval/loop_spec.txt}"   # 41 cell 净零闭环（重建，见 sweep8.sh 文件头）
+[ -f "$SPEC_FILE" ] || { echo "ABORT: 没有动作脚本 $SPEC_FILE"; exit 1; }
+# stage3 学生必须按**训练几何**评（sink 0 / window 20 / cap 20 = rollout_eval.py 默认，GEO-1）；
+# sink 3 / w17 只作显式 A/B：SINK=3 KV_WINDOW=17 ROPE_CAP=20 bash distill/eval/photodrift.sh ...
+# 几何旗标只在环境变量显式给出时才追加；默认一个都不传 ⇒ rollout_eval.py 的默认 sink 0 / window 20 /
+# cap 20 = stage3 训练几何。三者必须满足 rope_cap == sink + window（rollout_eval 会核）。
+GEO=()
+if [ -n "${SINK:-}" ]; then GEO+=(--sink_size "$SINK"); fi
+if [ -n "${KV_WINDOW:-}" ]; then GEO+=(--kv_window "$KV_WINDOW"); fi
+if [ -n "${ROPE_CAP:-}" ]; then GEO+=(--rope_cap "$ROPE_CAP"); fi
+OUT="$ROOT/outputs/distill/eval/deploy"; mkdir -p "$OUT"
+export PYTHONPATH="$ROOT" DIFFSYNTH_SKIP_DOWNLOAD=True CUDA_DEVICE_ORDER=PCI_BUS_ID
+CUDA_VISIBLE_DEVICES=$GPU "$ROOT/.venv/bin/python" "$ROOT/distill/eval/rollout_eval.py" \
+  --ckpt "$CK" --mode rollout --cells 41 "${GEO[@]}" \
+  --num_inference_steps "$N" --actions "$(cat "$SPEC_FILE")" --allow_raw_student \
+  --first_frame "$ROOT/data/eval/eybx_bg_50/19_medieval_castle.png" --tiled \
+  --out "$OUT/$TAG.mp4" > "$OUT/$TAG.log" 2>&1
+grep -E "实测位移" "$OUT/$TAG.log" || true
+"$ROOT/.venv/bin/python" - "$OUT/$TAG.mp4" "$TAG" <<'PY'
+import sys, imageio.v3 as iio, numpy as np
+g = iio.imread(sys.argv[1]).astype(np.float32).mean(-1)
+idx = [0,4,8,20,40,80,120,160]
+print(f"{sys.argv[2]:22s} 帧 " + " ".join(f"{i:>5d}" for i in idx))
+print(f"{'  std':22s}    " + " ".join(f"{g[i].std():5.1f}" for i in idx))
+print(f"{'  mean':22s}    " + " ".join(f"{g[i].mean():5.1f}" for i in idx))
+PY

backup_2026-09-14/code/distill/eval/rollout_eval.py

@@ -0,0 +1,1105 @@
+#!/usr/bin/env python3
+"""蒸馏产物的部署路径评测：按 4 步部署表生成一条 clip，实测背景位移 vs 指令位移。
+
+    # stage3 学生：因果 AR rollout（滑窗 KV），就是它训练时那条路
+    CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/data/zhiyangdeng/ActionRoPE DIFFSYNTH_SKIP_DOWNLOAD=True \\
+      .venv/bin/python distill/eval/rollout_eval.py \\
+        --ckpt outputs/distill/stage3_dmd/ema/ema_step400.safetensors \\
+        --replay <clip_id> \\
+        --out outputs/distill/stage3_dmd/eval/ema400_replay.mp4
+
+    # stage2 EMA 自己的部署路径：整片并行、2 步
+    ... --mode parallel --denoise_list 1000,250 \\
+        --first_frame <clip_id|png|mp4> --scene <scene key> --actions "right:1.0:21"
+
+    # 挑部署权重：一次扫完所有 ema（只加载一次 10 GB pipeline，之后只换 DiT 权重）
+    ... --sweep 'outputs/distill/stage3_dmd/ema/ema_step*.safetensors' --replay <clip_id> \\
+        --out_dir outputs/distill/stage3_dmd/eval/sweep
+
+    # 不碰真权重的自检（tiny 模型 + 合成条件，无 VAE）
+    ... --self_test
+
+为什么需要这个文件
+-----------------
+蒸馏的两个 loss 都不是质量指标：CD 的 loss 从 0 开始涨，DMD 是 min-max 目标（单调下降反而是坏消息）。
+ActionRoPE 唯一有意义的判据是**动作服从性**：命令背景往右走 373 px，画面是不是真的往左挪了 373 px。
+这里把 `actionrope/infer.py` 的 SIFT+RANSAC 实测（measure_displacement）与动作脚本解析（parse_actions）
+原样接过来，只把采样器换成蒸馏后的少步部署表 —— 于是「少步学生还听不听指令」这个问题有了数字答案。
+上游 WorldMind 判质量用的是亮度（brightness_eval.py），那是它的失败模式；这里的失败模式是动作漂移。
+
+两条部署路径（`--mode`）
+----------------------
+rollout（默认）：滑窗 KV、逐 cell 自回归，**与 stage3 训练时的 self_forcing_rollout 逐行同构**
+    —— 条件（窗口锚定的 offset_win / mask_in_k）直接 import 训练侧的 frame_cond，
+    提交趟同样是 detach 过的 x0 + t=0。差别只有一条：推理不抽退出档，每个 cell 把整张表走完。
+    它能超过 21 个 cell（这正是 stage3 要练的东西）。
+parallel：整片块因果、一次去噪全部 cell，cell 0 每步写回 —— stage2 EMA 自己的部署路径。
+
+采样规则（两条路一致，也与训练一致）
+--------------------------------
+    den   = noisy − sigma_i · v(noisy, t_i)         # x0 估计
+    noisy = (1 − sigma_{i+1})·den + sigma_{i+1}·ε   # 重新加噪到下一档（**新**噪声）
+最后一档直接取 den。这**不是** `infer.sample` 的 Euler 步（`x + (σ_next − σ)·v`）——
+那是 50 步扩散采样器；少步一致性/自强制学生训练时见到的就是上面这条 x0→重新加噪的链。
+`--ref_steps 50` 会额外跑一遍 `infer.sample`，给出同一条件下的 50 步参照（PORT_SPEC 的 S9）。
+
+一个必须守住的坑
+---------------
+部署权重是 `ema/`，不是 `dits/`。`dits/` 与 `critic/` 只为续训存在；拿 `dits/` 上线不会报错，
+只会静默降质。所以路径里带 `/dits/` 时本脚本直接拒跑，除非显式 `--allow_raw_student`（risk 12）。
+"""
+from __future__ import annotations
+
+import argparse
+import re
+import glob
+import json
+import math
+import os
+import subprocess
+import sys
+import time
+
+os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import numpy as np
+import torch
+
+ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from actionrope import geometry as G  # noqa: E402
+from actionrope import infer  # noqa: E402
+from actionrope.prompts import NULL_PROMPT  # noqa: E402
+from distill import flowmap as fm  # noqa: E402
+from distill.arope_dit import AropeDiT, unwrap  # noqa: E402
+from distill.data import extend_context, extend_offsets, make_mask_input, window_offsets  # noqa: E402
+from distill.flow_pf import DEPLOY_DENOISE_LIST, deploy_schedule  # noqa: E402
+from distill import runlib  # noqa: E402
+from distill.arope_dit import rot2d  # noqa: E402
+from distill.stage3_dmd.dmd_arope import frame_cond  # noqa: E402  ← 训练侧同一份逐 cell 条件
+from distill.tests.preflight import retained_cells, window_offsets_for_frame  # noqa: E402
+
+CELL_FRAMES = 4          # cell k(≥1) = 帧 4k−3 … 4k（Wan2.2 因果 VAE 的分组）
+
+
+# ----------------------------------------------------------------------------
+# 指令 → 逐帧 / 逐 cell 偏移
+# ----------------------------------------------------------------------------
+
+def cell_of_frames(cells: int) -> np.ndarray:
+    """每一帧属于哪个 cell（cell 0 = 帧 0；cell k = 帧 4k−3 … 4k）。= geometry.cell_frames 的逆，
+    只是把 21 的硬编码换成 cells —— rollout 要跑到 21 个 cell 以外，而 geometry 硬要求正好 21。"""
+    n_frames = CELL_FRAMES * cells - (CELL_FRAMES - 1)          # cells=21 ⇒ 81 = G.FRAMES
+    cell_of = np.zeros(n_frames, int)
+    for k in range(1, cells):
+        cell_of[CELL_FRAMES * k - 3: CELL_FRAMES * k + 1] = k
+    return cell_of
+
+
+def frames_to_cells_n(frame_off: np.ndarray, cells: int) -> np.ndarray:
+    """逐帧偏移 → 逐 cell 偏移（cell 内取均值）。cells == 21 时与 geometry.frames_to_cells 逐位相同。"""
+    cell_of = cell_of_frames(cells)
+    assert len(frame_off) == len(cell_of), f"帧数 {len(frame_off)} 与 {cells} cell 要求的 {len(cell_of)} 不符"
+    return np.stack([frame_off[cell_of == k].mean(0) for k in range(cells)], 0)
+
+
+def action_offsets(vel: np.ndarray, cells: int) -> tuple[np.ndarray, np.ndarray]:
+    """逐 cell 速度 (px/s) [cells,2] → (frame_offset_px [4·cells−3, 2], cell_offset_px [cells,2])。
+
+    与 `infer.parse_actions` + `geometry.frames_to_cells` 用的是**同一条**递推
+    （`frame_off[i] = frame_off[i−1] + v[cell(i)]/FPS`，cell 内取均值），只是把 21 的硬编码
+    换成了 cells —— cells == 21 时逐位相同（`_assert_matches_parse_actions` 会核对）。
+    """
+    assert vel.shape == (cells, 2), f"vel 形状应为 ({cells}, 2)，得到 {vel.shape}"
+    cell_of = cell_of_frames(cells)
+    frame_off = np.zeros((len(cell_of), 2), np.float64)
+    for i in range(1, len(cell_of)):
+        frame_off[i] = frame_off[i - 1] + vel[cell_of[i]] / G.FPS
+    return frame_off, frames_to_cells_n(frame_off, cells)
+
+
+def _ckpt_metadata(path: str) -> dict:
+    """读 safetensors 头里的 __metadata__（不加载权重）。save_dit 在这里写 flowmap 的形式与 gate，
+    stage3 的 do_save 还写训练几何（kv_window / sink_size / rope_cap，见 _check_ckpt_geometry）。"""
+    import json as _json
+    import struct
+    with open(path, "rb") as f:
+        n = struct.unpack("<Q", f.read(8))[0]
+        hdr = _json.loads(f.read(n))
+    return hdr.get("__metadata__", {}) or {}
+
+
+# stage3 do_save（distill/stage3_dmd/dmd_arope.py）写进 safetensors 头的训练几何键；runlib.save_dit 一律 str()。
+# stage1 / stage2 的 ckpt 没有这三个键。
+GEOMETRY_META_KEYS = ("kv_window", "sink_size", "rope_cap")
+
+
+def _check_ckpt_geometry(path: str, args, log=print) -> None:
+    """stage3 学生必须按**训练几何**评（GEO-1）：ckpt 元数据里 kv_window / sink_size / rope_cap 三键齐全
+    且任一与命令行不等 ⇒ 拒绝，除非 --allow_geometry_mismatch（只作显式 A/B，RUNBOOK 的口径）。
+
+    为什么硬拒而不是警告：sink 3 / w17 与训练的 sink 0 / w20 只在 k≤20 的保留集相同，之后的 h/w RoPE
+    与掩码通道是训练 rollout 从未产生过的输入 —— 拿它比 EMA、比 stage2 都只是在比 OOD 外推，ckpt
+    选择轴就错了。三键不全（stage1 / stage2 的 ckpt）不受限。"""
+    meta = _ckpt_metadata(path)
+    if not all(k in meta for k in GEOMETRY_META_KEYS):
+        return
+    try:
+        ck = {k: int(meta[k]) for k in GEOMETRY_META_KEYS}
+    except ValueError as e:
+        raise SystemExit(f"{path} 的几何元数据不是整数：{ {k: meta[k] for k in GEOMETRY_META_KEYS} }（{e}）")
+    cli = {k: int(getattr(args, k)) for k in GEOMETRY_META_KEYS}
+    diff = {k: (ck[k], cli[k]) for k in GEOMETRY_META_KEYS if ck[k] != cli[k]}
+    if "sink_anchor" in meta and str(meta["sink_anchor"]) != str(args.sink_anchor):
+        diff["sink_anchor"] = (str(meta["sink_anchor"]), str(args.sink_anchor))
+    if not diff:
+        log(f"[geometry] {os.path.basename(path)} 训练几何 sink {ck['sink_size']} / window {ck['kv_window']} / "
+            f"cap {ck['rope_cap']} == 命令行")
+        return
+    msg = (f"{path} 的训练几何与命令行不符：" + ", ".join(f"{k} ckpt={a} cli={b}" for k, (a, b) in diff.items())
+           + "。\n  stage3 学生只在训练几何下才是分布内的（GEO-1）；请用 --sink_size/--kv_window/--rope_cap "
+             "对齐 ckpt，或加 --allow_geometry_mismatch（只作显式 A/B，不作 ckpt 选择轴）。")
+    if args.allow_geometry_mismatch:
+        log(f"[geometry] 放行（--allow_geometry_mismatch）：{msg}")
+        return
+    raise SystemExit(msg)
+
+
+def _assert_matches_parse_actions(spec: str, walk_speed: np.ndarray) -> None:
+    """cells=21 时本文件的递推必须与 infer.parse_actions / frames_to_cells 逐位相同。"""
+    fo, vel = infer.parse_actions(spec, walk_speed)
+    fo2, co2 = action_offsets(vel, G.CELLS)
+    assert np.array_equal(fo, fo2), "action_offsets 与 infer.parse_actions 的逐帧偏移不一致"
+    assert np.array_equal(G.frames_to_cells(fo), co2), "逐 cell 偏移与 geometry.frames_to_cells 不一致"
+
+
+def parse_actions_n(spec: str, walk_speed: np.ndarray, cells: int) -> np.ndarray:
+    """与 `infer.parse_actions` 同一套语法，但总 cell 数是 `cells` 而不是硬写的 21。
+
+    上游那个函数在 `len(vel) != G.CELLS` 时直接抛错（`actionrope/infer.py:127`），而超出训练
+    视野的多段指令（"绕一圈"、"第 30 格突然转向"）正是 AR rollout 唯一值得测的东西 ——
+    直线走到底用 `extend_velocity` 重复末段就够了，不需要这个函数。
+
+    分段解析是从上游抄来的，所以两份实现可能漂移；`cells == 21` 时本函数会直接调上游核对，
+    逐位不同就抛。不动 `actionrope/`：那是训练与推理的正典，评测不该改它。
+    """
+    toks = [t.strip() for t in spec.split(",") if t.strip()]
+    segs, i = [], 0
+    while i < len(toks):
+        if ":" not in toks[i]:                       # "dx,dy:n" 被逗号切开了，粘回去
+            if i + 1 >= len(toks):
+                raise ValueError(f"动作段不完整: {toks[i]!r}")
+            segs.append(toks[i] + ";" + toks[i + 1]); i += 2
+        else:
+            segs.append(toks[i]); i += 1
+    vel = []
+    for seg in segs:
+        parts = seg.split(":")
+        if len(parts) == 3:
+            name, mult, n = parts[0].strip().lower(), float(parts[1]), int(parts[2])
+            if name not in infer.DIR_NAMES:
+                raise ValueError(f"未知方向 {name!r}，可选 {sorted(infer.DIR_NAMES)}")
+            v = G.ACTION_DIRS[infer.DIR_NAMES[name]] * walk_speed * mult
+        elif len(parts) == 2 and ";" in parts[0]:
+            dx, dy = (float(t) for t in parts[0].split(";"))
+            v, n = np.array([dx, dy], np.float64), int(parts[1])
+        else:
+            raise ValueError(f"动作段格式应为 dir:mult:cells 或 dx,dy:cells，得到 {seg!r}")
+        vel += [v] * n
+    if len(vel) != cells:
+        raise ValueError(f"动作脚本的 cell 数合计应为 {cells}（--cells），得到 {len(vel)}")
+    vel = np.stack(vel, 0)
+    if cells == G.CELLS:                             # 与上游对表，防止两份解析漂移
+        _, vel_up = infer.parse_actions(spec, walk_speed)
+        assert np.array_equal(vel, vel_up), "parse_actions_n 与 infer.parse_actions 解析结果不一致"
+    return vel
+
+
+def extend_velocity(vel: np.ndarray, cells: int) -> np.ndarray:
+    """把动作脚本的逐 cell 速度延长到 cells 个 cell：重复最后一段的速度。
+
+    评测**故意**不用训练侧默认的 `resample`：那是给训练喂多样性的，随机重抽方向之后
+    「指令」就不再是一个可核对的东西了。重复最后一段速度 = 一条明确的、可以拿来比对的指令，
+    在 cell 层面与 `data.extend_offsets(mode="continue")` 同解。
+    """
+    if cells <= len(vel):
+        return vel[:cells]
+    return np.concatenate([vel, np.repeat(vel[-1:], cells - len(vel), axis=0)], 0)
+
+
+# ----------------------------------------------------------------------------
+# 两条部署采样路径
+# ----------------------------------------------------------------------------
+
+ROT_MASK = ("rot", "plain")      # rot = 斜的 known 多边形（几何正确）；plain = 老的轴对齐掩码（消融）
+
+
+def parallel_mask(off_t: torch.Tensor, ang_t, lh: int, lw: int, rot_mask: str = "rot"):
+    """并行路径的 known 掩码。ang_t [b, f] 是**绝对**朝向（cell 0 = 0）：cell k 自己的格点先转 φ_k
+    进世界系再判定是否落在 cell 0 的足迹内。`_known_rotated` 的口径是"偏移在该 cell 自己的相机系里、
+    角度取锚点相对该 cell"，所以这里传 R(−φ_k)·off_k 与 −φ_k。ang_t=None ⇒ 老路。"""
+    if ang_t is None or rot_mask == "plain":
+        return make_mask_input(off_t, lh, lw)
+    off_local = rot2d(off_t, -ang_t)
+    return make_mask_input(off_local, lh, lw, angle=-ang_t)
+
+
+@torch.no_grad()
+def deploy_parallel(model: AropeDiT, z0: torch.Tensor, context, context_ids, off_t: torch.Tensor,
+                    deploy: tuple, *, seed: int = 0, log=print, ang_t=None,
+                    rot_mask: str = "rot", scale_t=None) -> torch.Tensor:
+    """整片并行的少步部署（stage2 EMA 自己的路径）。cell 0 每一步都写回首帧 latent。"""
+    _, sig, t_grid = deploy
+    b, c, _, lh, lw = z0.shape[0], z0.shape[1], 0, z0.shape[3], z0.shape[4]
+    cells = off_t.shape[1]
+    g = torch.Generator("cpu").manual_seed(seed)
+    latents = torch.randn(b, c, cells, lh, lw, generator=g).to(dtype=z0.dtype, device=z0.device)
+    latents[:, :, 0:1] = z0
+    mask_in = parallel_mask(off_t, ang_t, lh, lw, rot_mask)
+    t0 = time.time()
+    for i in range(len(t_grid)):
+        ts = torch.full((b,), float(t_grid[i]), device=z0.device)
+        v = model(latents, ts, context, context_ids=context_ids, offset_px=off_t,
+                  mask_input=mask_in, first_frame_cond=True, angle_px=ang_t, scale_px=scale_t)
+        den = latents - float(sig[i]) * v                       # x0 估计
+        if i + 1 == len(t_grid):
+            latents = den
+        else:
+            ns = float(sig[i + 1])
+            latents = (1.0 - ns) * den + ns * torch.randn(den.shape, generator=g).to(den)
+        latents[:, :, 0:1] = z0                                 # TI2V 锚点，每步写回
+        log(f"   parallel step {i + 1}/{len(t_grid)} t={float(t_grid[i]):.0f} σ={float(sig[i]):.4f} "
+            f"· {time.time() - t0:.1f}s")
+    return latents
+
+
+@torch.no_grad()
+def deploy_parallel_flowmap(model: AropeDiT, z0: torch.Tensor, context, context_ids,
+                            off_t: torch.Tensor, sigmas, *, seed: int = 0, log=print) -> torch.Tensor:
+    """整片并行的 **flow map** 部署（stage2-flowmap 的 EMA 走这条）。
+
+    与 renoise 版的两处**本质**差别，接错任何一处都不会报错、只会安静地产出另一个模型：
+
+      1. 模型要吃**两个**时间步。每步喂 (t_i, r_i) = (σ_i, σ_{i+1})·1000，它预测的是区间
+         [r, t] 上的**均**速度，不是 t 处的瞬时速度。
+      2. **全程不重新加噪**。每步就是 `x ← x − (σ_t − σ_r)·u`；轴的末位是 0.0，所以最后一步
+         直接落在干净估计上，不需要再单独读一次 x0。renoise 版每步都要重采一份新噪声，
+         那是一致性/扩散学生的采样方式，拿来采 flow map 会多注入 N 份噪声。
+
+    `sigmas` 长度 N+1、严格降序、末位 0.0（`flowmap.anystep_sigmas(N)`）。N 可以任意取 ——
+    这正是 any-step 的主张本身，4 步和 24 步用的是同一份权重。
+    """
+    fm.assert_anystep_axis(sigmas if torch.is_tensor(sigmas) else torch.as_tensor(sigmas), terminal=True)
+    b, c, _, lh, lw = z0.shape[0], z0.shape[1], 0, z0.shape[3], z0.shape[4]
+    cells = off_t.shape[1]
+    g = torch.Generator("cpu").manual_seed(seed)
+    x = torch.randn(b, c, cells, lh, lw, generator=g).to(dtype=z0.dtype, device=z0.device)
+    x[:, :, 0:1] = z0
+    mask_in = make_mask_input(off_t, lh, lw)
+    t0 = time.time()
+    N = len(sigmas) - 1
+    for i in range(N):
+        s_t, s_r = float(sigmas[i]), float(sigmas[i + 1])
+        tt = torch.full((b,), s_t * fm.NUM_TRAIN_TIMESTEPS, device=z0.device)
+        rr = torch.full((b,), s_r * fm.NUM_TRAIN_TIMESTEPS, device=z0.device)
+        u = model(x, tt, context, context_ids=context_ids, r_timestep=rr,
+                  offset_px=off_t, mask_input=mask_in, first_frame_cond=True)
+        x = fm.euler_step(x, u, s_t, s_r)
+        x[:, :, 0:1] = z0                                   # TI2V 锚点，每步写回
+        log(f"   flowmap parallel {i + 1}/{N}  σ {s_t:.4f}→{s_r:.4f}  · {time.time() - t0:.1f}s")
+    return x
+
+
+@torch.no_grad()
+def deploy_rollout_flowmap(model: AropeDiT, z0: torch.Tensor, ctx_ext, ids_ext, off_ext: torch.Tensor,
+                           cells: int, sigmas, *, kv_window: int = 20, sink_size: int = 0,
+                           rope_cap: int = 20, seed: int = 0, log=print, sink_anchor: str = "absolute") -> torch.Tensor:
+    """滑窗 KV 的逐 cell 自回归 **flow map** 部署。
+
+    结构与 `deploy_rollout` 逐行对应，只把每个 cell 内部的采样链换成 flow-map Euler：
+    每个 cell 沿整条 any-step 轴走完（无重新加噪），然后**照旧**用 `t=0` 提交 detach 过的
+    结果进缓存 —— 提交趟的 t=0 与 flow map 无关，它是 TF clean 流的同分布要求，不能改。
+    """
+    fm.assert_anystep_axis(sigmas if torch.is_tensor(sigmas) else torch.as_tensor(sigmas), terminal=True)
+    b, c, _, lh, lw = z0.shape
+    dev = z0.device
+    bare = unwrap(model)
+    cache = bare.init_kv_cache(kv_window=kv_window, sink_size=sink_size, rope_cap=rope_cap)
+    g = torch.Generator("cpu").manual_seed(seed)
+    N = len(sigmas) - 1
+
+    def _call(x, t, k, commit, ow, mk, r=None):
+        return model(x, t, ctx_ext[:, k:k + 1], context_ids=ids_ext[:, k:k + 1],
+                     r_timestep=(t if r is None else r),
+                     first_frame_cond=False, frame_index=k, kv_cache=cache, commit=commit,
+                     offset_win=ow, mask_in_k=mk)
+
+    ow0, mk0 = frame_cond(off_ext, 0, sink_size, kv_window, lh, lw, sink_anchor)
+    zeros = torch.zeros(b, device=dev)
+    _call(z0, zeros, 0, True, ow0, mk0)                     # 真实首帧，t=r=0 进缓存
+    frames = [z0]
+    t0 = time.time()
+    for k in range(1, cells):
+        ow, mk = frame_cond(off_ext, k, sink_size, kv_window, lh, lw, sink_anchor)
+        x = torch.randn(b, c, 1, lh, lw, generator=g).to(dtype=z0.dtype, device=dev)
+        for i in range(N):
+            s_t, s_r = float(sigmas[i]), float(sigmas[i + 1])
+            tt = torch.full((b,), s_t * fm.NUM_TRAIN_TIMESTEPS, device=dev)
+            rr = torch.full((b,), s_r * fm.NUM_TRAIN_TIMESTEPS, device=dev)
+            x = fm.euler_step(x, _call(x, tt, k, False, ow, mk, r=rr), s_t, s_r)
+        frames.append(x)
+        _call(x.detach(), zeros, k, True, ow, mk)           # 每个 cell 恰好提交一次，t=0
+        if k % 5 == 0 or k == cells - 1:
+            log(f"   flowmap rollout cell {k}/{cells - 1} · {time.time() - t0:.1f}s")
+    out = torch.cat(frames, dim=2)
+    bare.release_kv_cache(cache)      # cache 与 _ref 成环，不显���清会拖着十几 GB 不放
+    return out
+
+
+@torch.no_grad()
+def build_angles(args, cells: int, device, log=print):
+    """--rotate <每 cell 度数> / --rotate_spec "<deg>:<n>,..." ⇒ [1, cells] 的逐 cell 相机朝向（弧度）。
+
+    cell 0 恒为 0（首帧定义世界系）；朝向是**累计**的：φ_k = φ_{k−1} + 该段的每格角速度。
+    都没给 ⇒ None，整条 rollout 逐位走旋转之前的老路。
+    """
+    spec = getattr(args, "rotate_spec", None)
+    rate = float(getattr(args, "rotate", 0.0) or 0.0)
+    if not spec and rate == 0.0:
+        return None
+    per_cell = [0.0] * cells
+    if spec:
+        k = 1
+        for seg in [t.strip() for t in spec.split(",") if t.strip()]:
+            deg, n = seg.split(":")
+            for _ in range(int(n)):
+                if k < cells:
+                    per_cell[k] = float(deg)
+                k += 1
+        if k < cells:                      # 没铺满 ⇒ 末段的角速度续到底
+            last = float(spec.split(",")[-1].split(":")[0])
+            for j in range(k, cells):
+                per_cell[j] = last
+    else:
+        for j in range(1, cells):
+            per_cell[j] = rate
+    ang = torch.tensor(per_cell, dtype=torch.float32, device=device).cumsum(0).unsqueeze(0)
+    ang = ang * math.pi / 180.0
+    deg = (ang[0] * 180.0 / math.pi).tolist()
+    log(f"   --rotate：相机朝向 cell0={deg[0]:.1f}° → cell{cells - 1}={deg[-1]:.1f}°"
+        f"（每格最大 {max(per_cell):.2f}°；画幅角点离中心 ~15 token，1° 约合角点 0.26 token 位移）")
+    return ang
+
+
+def build_scales(args, cells: int, device, log=print):
+    """--zoom <每 cell 的位置缩放率>：位置整体等比缩放 s_k = (1+rate)^k（cell 0 = 1）。
+
+    判决实验用：位置铺得越开，这一帧名义上覆盖的世界就越大。若模型把"位置场变大"读成"视野变大"，
+    内容应当按 1/s 缩小 —— 那就说明旋转时看到的缩小也是同一条通道，靠推理端补不回来。
+    """
+    rate = float(getattr(args, "zoom", 0.0) or 0.0)
+    if rate == 0.0:
+        return None
+    s = torch.tensor([(1.0 + rate) ** k for k in range(cells)], dtype=torch.float32, device=device).unsqueeze(0)
+    log(f"   --zoom：位置缩放 cell0={float(s[0, 0]):.3f} → cell{cells - 1}={float(s[0, -1]):.3f}"
+        f"（若模型按视野读，内容应缩小到 1/{float(s[0, -1]):.3f} = {1 / float(s[0, -1]):.3f}）")
+    return s
+
+
+def frame_cond_rot(off_ext: torch.Tensor, ang_ext: torch.Tensor, k: int, sink_size: int,
+                   kv_window: int, lat_h: int, lat_w: int, sink_anchor: str = "absolute",
+                   rot_mask: str = "rot"):
+    """带相机旋转的逐 cell 条件：返回 (offset_win, mask_in_k, angle_win)。
+
+    一切表达在**当前 cell 的相机系**里（`angle_win[:, -1] == 0`，当前帧格点保持轴对齐）：
+      · 偏移：先把全部世界偏移转 −φ_k 进当前相机系，再走与 frame_cond 完全相同的窗口锚定；
+      · 角度：φ_j − φ_k（历史 cell 相对当前相机转了多少）；
+      · 掩码：known 的判据仍是"落在锚点 cell 的足迹内"，锚点相对当前转了 φ_a − φ_k，于是足迹是斜的。
+    ang_ext 全 0 时逐位退化成 `frame_cond`（本函数自带断言核对）。
+    """
+    idx = retained_cells(k, sink_size, kv_window) + [k]
+    sel = torch.as_tensor(idx, device=off_ext.device, dtype=torch.long)
+    phi_k = ang_ext[:, k]                                            # [b]
+    ang_win = ang_ext.index_select(1, sel) - phi_k.unsqueeze(1)      # [b, n+1]，最后一项恒 0
+    off_rot = rot2d(off_ext, -phi_k.unsqueeze(1))                    # 整条转进当前相机系
+    ow = window_offsets_for_frame(off_rot, k, sink_size, kv_window, sink_anchor=sink_anchor)
+    # 锚点 = ow 里被减成 0 的那一格（absolute ⇒ idx[0]；window 且缓存饱和 ⇒ idx[sink_size]）
+    saturated = kv_window is not None and k > sink_size + kv_window
+    a = sink_size if (sink_anchor == "window" and sink_size > 0 and saturated) else 0
+    mk = (make_mask_input(ow[:, -1:], lat_h, lat_w) if rot_mask == "plain"
+          else make_mask_input(ow[:, -1:], lat_h, lat_w, angle=ang_win[:, a:a + 1]))
+    return ow, mk, ang_win
+
+
+def deploy_rollout(model: AropeDiT, z0: torch.Tensor, ctx_ext, ids_ext, off_ext: torch.Tensor,
+                   cells: int, deploy: tuple, *, kv_window: int = 20, sink_size: int = 0,
+                   rope_cap: int = 20, seed: int = 0, log=print, euler=None, sink_anchor: str = "absolute",
+                   ang_ext: torch.Tensor | None = None, rot_mask: str = "rot") -> torch.Tensor:
+    """滑窗 KV 的逐 cell 自回归部署 —— stage3 的 self_forcing_rollout 的 no-grad 孪生。
+
+    euler 非 None ⇒ 每个 cell 不走 renoise 链，而是用这个 FlowMatchScheduler 的 Euler ODE
+    （`sched.step`，与 `infer.sample` 同一条路；longhorizon/README 的「ODE 16 步」口径），
+    从纯噪声 σ=1 一路���分到 σ=0。KV 提交与条件完全不变。
+
+    与训练时的唯一差别：推理不抽退出档，每个 cell 把整张部署表走完再提交。
+    条件（窗口锚定的 offset_win / 当前 cell 的 known 掩码）直接用训练侧的 frame_cond，
+    提交趟同样是 **detach 过的 x0 + t=0**（这个 0 让缓存历史与 TF 的 clean 流同分布）。
+    """
+    _, sig, t_grid = deploy
+    b, c, _, lh, lw = z0.shape
+    dev = z0.device
+    bare = unwrap(model)
+    cache = bare.init_kv_cache(kv_window=kv_window, sink_size=sink_size, rope_cap=rope_cap)
+    g = torch.Generator("cpu").manual_seed(seed)
+
+    def _call(x, t, k, commit, ow, mk, aw=None):
+        return model(x, t, ctx_ext[:, k:k + 1], context_ids=ids_ext[:, k:k + 1],
+                     first_frame_cond=False, frame_index=k, kv_cache=cache, commit=commit,
+                     offset_win=ow, mask_in_k=mk, angle_win=aw)
+
+    def _cond(k):
+        # ang_ext=None ⇒ 完全走旋转之前的那条路（frame_cond），逐位不变
+        if ang_ext is None:
+            ow_, mk_ = frame_cond(off_ext, k, sink_size, kv_window, lh, lw, sink_anchor)
+            return ow_, mk_, None
+        return frame_cond_rot(off_ext, ang_ext, k, sink_size, kv_window, lh, lw, sink_anchor, rot_mask)
+
+    ow0, mk0, aw0 = _cond(0)
+    _call(z0, torch.zeros(b, device=dev), 0, True, ow0, mk0, aw0)    # 真实首帧，t=0 进缓存
+    frames = [z0]
+    t0 = time.time()
+    for k in range(1, cells):
+        ow, mk, aw = _cond(k)
+        noisy = torch.randn(b, c, 1, lh, lw, generator=g).to(dtype=z0.dtype, device=dev)
+        den = None
+        if euler is not None:
+            x = noisy
+            for i, t in enumerate(euler.timesteps):
+                ts = torch.full((b,), float(t), device=dev)
+                x = euler.step(_call(x, ts, k, False, ow, mk, aw), euler.timesteps[i], x)
+            den = x
+        else:
+            for i in range(len(t_grid)):
+                ts = torch.full((b,), float(t_grid[i]), device=dev)
+                den = noisy - float(sig[i]) * _call(noisy, ts, k, False, ow, mk, aw)
+                if i + 1 < len(t_grid):
+                    ns = float(sig[i + 1])
+                    noisy = (1.0 - ns) * den + ns * torch.randn(den.shape, generator=g).to(den)
+        frames.append(den)
+        _call(den.detach(), torch.zeros(b, device=dev), k, True, ow, mk, aw)   # 每个 cell 恰好提交一次
+        if k % 5 == 0 or k == cells - 1:
+            log(f"   rollout cell {k}/{cells - 1} · {time.time() - t0:.1f}s")
+    return torch.cat(frames, dim=2)
+
+
+# ----------------------------------------------------------------------------
+# 一次评测
+# ----------------------------------------------------------------------------
+
+def measure(frames: np.ndarray, frame_off: np.ndarray, chain: int = 16) -> dict:
+    """SIFT+RANSAC 实测背景位移 vs 指令。指令背景位移 = −frame_offset_px[最后一帧]。"""
+    k = len(frames) - 1
+    meas = infer.measure_displacement(infer.to_gray(frames), k=k, chain=chain)
+    expect = (-frame_off[k]).tolist()
+    err = None if meas["sift"] is None else [meas["sift"][0] - expect[0], meas["sift"][1] - expect[1]]
+    out = {"k": k, "measured": meas, "expected_bg_shift": expect, "error_px": err}
+    if err is not None:
+        out["error_norm_px"] = float(np.hypot(*err))
+        out["expected_norm_px"] = float(np.hypot(*expect))
+        out["error_rel"] = (out["error_norm_px"] / out["expected_norm_px"]
+                            if out["expected_norm_px"] > 1e-6 else None)
+    return out
+
+
+def run_one(model, pipe, z0, cond, args, deploy, log=print, sigmas=None) -> dict:
+    """生成 + 解码 + 实测。cond 由 build_condition() 给出。
+
+    `sigmas` 非 None ⇒ 走 flow-map 的 any-step Euler（无重新加噪，模型吃 (t, r) 两个时间步）；
+    None ⇒ 走原来的 renoise 链。两者接错都不报错，所以分流在 main() 里按 checkpoint 决定。
+    """
+    t0 = time.time()
+    if args.mode == "rollout":
+        if sigmas is not None:
+            if cond.get("ang_ext") is not None:
+                raise SystemExit("--rotate 只支持 renoise / euler 的 rollout 路径，flowmap 分支没接")
+            lat = deploy_rollout_flowmap(model, z0, cond["ctx_ext"], cond["ids_ext"], cond["off_ext"],
+                                         cond["cells"], sigmas, kv_window=args.kv_window,
+                                         sink_size=args.sink_size, rope_cap=args.rope_cap,
+                                         seed=args.seed, log=log, sink_anchor=args.sink_anchor)
+        else:
+            lat = deploy_rollout(model, z0, cond["ctx_ext"], cond["ids_ext"], cond["off_ext"],
+                                 cond["cells"], deploy, kv_window=args.kv_window,
+                                 sink_size=args.sink_size, rope_cap=args.rope_cap, seed=args.seed, log=log,
+                                 euler=getattr(args, "euler_sched", None), sink_anchor=args.sink_anchor,
+                                 ang_ext=cond.get("ang_ext"), rot_mask=args.rotate_mask)
+    else:
+        if sigmas is not None:
+            lat = deploy_parallel_flowmap(model, z0, cond["context"], cond["context_ids"],
+                                          cond["off_par"], sigmas, seed=args.seed, log=log)
+        else:
+            ang_par = cond.get("ang_ext")
+            if ang_par is not None:
+                ang_par = ang_par[:, :cond["off_par"].shape[1]]
+            lat = deploy_parallel(model, z0, cond["context"], cond["context_ids"], cond["off_par"],
+                                  deploy, seed=args.seed, log=log, ang_t=ang_par,
+                                  rot_mask=args.rotate_mask, scale_t=cond.get("scale_ext"))
+    torch.cuda.synchronize() if z0.is_cuda else None
+    sample_sec = time.time() - t0
+    assert torch.isfinite(lat).all(), "采样输出里有 NaN/Inf"
+    t0 = time.time()
+    frames = infer.decode_video(pipe, lat, tiled=args.tiled)
+    dec_sec = time.time() - t0
+    res = measure(frames, cond["frame_off"], chain=args.chain)
+    res.update(sample_sec=sample_sec, decode_sec=dec_sec, frames=int(frames.shape[0]),
+               cells=cond["cells"], mode=args.mode)
+    return res, frames, lat
+
+
+def apply_late_trigger(prompts21: list[str], spec: str, cells: int) -> list[str]:
+    """把 21 条（已 strip_action 的）逐 cell 场景串延到 cells 条，并在 cell ≥ 21 处再切一次场景。
+
+    spec "burn:<C>:<start>:<n>"（或 cut:）：cell [start, start+n) 用转场从句「B burning away into C」，
+    之后的 cell 切成纯 C；start 之前保持 cell 20 的场景 B（= extend_context 的 hold）。start 必须 ≥ 21，
+    21 以内的切换走 --trigger（那是训练语料的口径）。串必须在文本表里（14 个源场景 × 13 个目的都有）。
+    """
+    kind, c, start, n = spec.split(":")
+    start, n = int(start), int(n)
+    if kind not in infer.TRIG_TMPL:
+        raise ValueError(f"--late_trigger 类型应为 burn/cut，得到 {kind!r}")
+    if not (G.CELLS <= start < cells and n >= 1 and start + n <= cells):
+        raise ValueError(f"--late_trigger 的 start/n 越界：start={start} n={n} cells={cells}（start 须 ≥ {G.CELLS}）")
+    m = re.match(r"^In a 2\.5D top down view, the player is (.+)\.$", prompts21[-1])
+    if not m:
+        raise ValueError(f"cell 20 没有场景从句（{prompts21[-1]!r}），无法构造转场；请给 --scene")
+    a, c_nl = m.group(1), infer.scene_clause(c)
+    trig = infer.TRIG_TMPL[kind].format(a=a, b=c_nl[3:] if c_nl.startswith("in ") else c_nl)
+    tmpl = "In a 2.5D top down view, the player is {scene}."
+    out = list(prompts21) + [prompts21[-1]] * (start - G.CELLS)
+    out += [tmpl.format(scene=trig)] * n + [tmpl.format(scene=c_nl)] * (cells - start - n)
+    assert len(out) == cells
+    return out
+
+
+def apply_switches(scene_a: str, spec: str, cells: int) -> list[str]:
+    """从 --scene A 起，按 spec "burn:B:10:2,burn:C:20:2,cut:D:30:2" 在任意 cell 处连续切换多次场景，
+    直接生成整条 cells 长的逐 cell 场景串（与 apply_late_trigger 同一套转场从句与模板；串都在文本表里）。
+    每段：[start, start+n) 用「上一场景 burning away into 新场景」，之后是纯新场景，直到下一次切换。"""
+    tmpl = "In a 2.5D top down view, the player is {scene}."
+    cur = infer.scene_clause(scene_a)
+    segs = []
+    for tok in [t.strip() for t in spec.split(",") if t.strip()]:
+        kind, c, start, n = tok.split(":")
+        start, n = int(start), int(n)
+        if kind not in infer.TRIG_TMPL:
+            raise ValueError(f"--switches 类型应为 burn/cut，得到 {kind!r}")
+        segs.append((start, n, kind, c))
+    segs.sort()
+    out: list[str] = []
+    for start, n, kind, c in segs:
+        if not (len(out) <= start and 1 <= start and start + n <= cells):
+            raise ValueError(f"--switches 段越界或重叠：start={start} n={n} cells={cells}，已生成 {len(out)} 条")
+        out += [tmpl.format(scene=cur)] * (start - len(out))
+        c_nl = infer.scene_clause(c)
+        trig = infer.TRIG_TMPL[kind].format(a=cur, b=c_nl[3:] if c_nl.startswith("in ") else c_nl)
+        out += [tmpl.format(scene=trig)] * n
+        cur = c_nl
+    out += [tmpl.format(scene=cur)] * (cells - len(out))
+    assert len(out) == cells
+    return out
+
+
+def build_condition(args, enc, device, log=print) -> dict:
+    """指令 → 逐帧偏移 / 逐 cell 偏移 / 逐 cell 文本 / 首帧规格。"""
+    sc = torch.load(infer.SIDECAR, map_location="cpu", weights_only=False)
+    walk_speed = np.array(sc["meta"]["walk_speed_px_s"], np.float64)
+    cells = args.cells
+    if args.replay:
+        row = infer.load_meta_row(args.replay)
+        c = sc["clips"][args.replay]
+        frame_off21 = c["frame_offset_px"].numpy().astype(np.float64)      # [81,2] 真值
+        prompts = [infer.strip_action(p) for p in row["prompts"]]          # ARoPE 臂只看场景
+        first_spec = args.replay
+        co21 = G.frames_to_cells(frame_off21)
+        assert np.array_equal(co21, frames_to_cells_n(frame_off21, G.CELLS)), "逐 cell 归约口径分叉了"
+        if cells == G.CELLS:
+            frame_off, cell_off = frame_off21, co21
+        elif cells < G.CELLS:
+            # 截短（冒烟/计时用）：cell k 覆盖帧 4k−3…4k，所以前 cells 个 cell 用前 4·cells−3 帧。
+            # 不能整段拿 81 帧去配 cells 个 cell —— frames_to_cells_n 会直接断言失败。
+            n_keep = CELL_FRAMES * cells - (CELL_FRAMES - 1)
+            frame_off = frame_off21[:n_keep]
+            cell_off = frames_to_cells_n(frame_off, cells)
+            assert np.array_equal(cell_off, co21[:cells]), "截短不该改动前 cells 个 cell 的真值"
+            log(f"   replay 截到 {cells} cell（{n_keep} 帧）")
+        else:
+            # 21 个 cell 之内**原样用真值逐帧偏移**（sidecar 的真实轨迹不是分段匀速，重建一遍会失真），
+            # 21 之外按末段隐含的匀速续走（= data.extend_offsets("continue")）
+            v_last = (co21[-1] - co21[-2]) * G.FPS / CELL_FRAMES
+            n_ext = CELL_FRAMES * (cells - G.CELLS)
+            ext = frame_off21[-1] + np.outer(np.arange(1, n_ext + 1) / G.FPS, v_last)
+            frame_off = np.concatenate([frame_off21, ext], 0)
+            cell_off = frames_to_cells_n(frame_off, cells)
+            assert np.array_equal(cell_off[:G.CELLS], co21), "延长不该改动前 21 个 cell 的真值"
+            log(f"   replay 延到 {cells} cell：末段速度 ({v_last[0]:+.1f},{v_last[1]:+.1f}) px/s 续走")
+    else:
+        if not (args.first_frame and args.actions):
+            raise SystemExit("非 --replay 模式需要 --first_frame 和 --actions")
+        n_spec = sum(int(s.split(":")[-1]) for s in args.actions.split(",") if ":" in s)
+        if n_spec == cells and cells != G.CELLS:
+            # 多段指令铺满整个 rollout（"绕一圈"、"第 30 格转向"）—— 超出训练视野的指令
+            # 才是 AR rollout 值得测的东西。直线到底仍走下面那条 21+extend 的路。
+            vel = parse_actions_n(args.actions, walk_speed, cells)
+            log(f"   动作脚本铺满 {cells} cell（不是 21+续走）")
+        else:
+            _assert_matches_parse_actions(args.actions, walk_speed)        # 与 infer.py 同一条递推
+            _, vel21 = infer.parse_actions(args.actions, walk_speed)
+            vel = extend_velocity(vel21, cells)
+        frame_off, cell_off = action_offsets(vel, cells)
+        # 文本只含场景（ARoPE 臂走 strip_action），动作走 RoPE 不走文本，所以标签取前 21 格即可
+        labels21 = [infer.velocity_to_label(v, mirror=False) for v in vel[:G.CELLS]]
+        if args.scene:
+            prompts = [infer.strip_action(p) for p in
+                       infer.compose_cell_prompts(args.scene, args.trigger, labels21)]
+        else:
+            # bg50 的 50 张背景不属于 18 个训练场景，配不出场景从句 —— eval_bg50.py:189-192
+            # 的默认口径就是空条件串，这里对齐它，别硬编一个模型没见过的场景描述。
+            prompts = [NULL_PROMPT] * G.CELLS
+            log(f"   未给 --scene，用空条件串 {NULL_PROMPT!r}（bg50 的默认口径）")
+        first_spec = args.first_frame
+    assert len(prompts) == G.CELLS, f"prompt 应为 {G.CELLS} 条，得到 {len(prompts)}"
+    context, context_ids = infer.build_context(enc, prompts, device, per_cell=True)  # [1,21,512,4096]
+    off21 = torch.tensor(cell_off[:G.CELLS], dtype=torch.float32).unsqueeze(0).to(device)
+    off_full = torch.tensor(cell_off, dtype=torch.float32).unsqueeze(0).to(device)
+    # 逐 cell 文本延到 1+N（cell ≥ 21 保持最后一条，与训练侧 extend_context 同一函数）
+    ctx_ext, ids_ext = extend_context(context, context_ids, cells - 1, mode="hold")
+    if getattr(args, "late_trigger", None):
+        prompts_n = apply_late_trigger(prompts, args.late_trigger, cells)
+        ctx_ext, ids_ext = infer.build_context(enc, prompts_n, device, per_cell=True)   # [1, cells, 512, 4096]
+        sw = [k for k in range(1, cells) if prompts_n[k] != prompts_n[k - 1]]
+        log(f"   --late_trigger {args.late_trigger}：cell {sw} 处切换，{len(set(prompts_n))} 条不同的串")
+        prompts = prompts_n
+    if getattr(args, "switches", None):
+        if not args.scene:
+            raise SystemExit("--switches 需要 --scene 作为起始场景")
+        prompts_n = apply_switches(args.scene, args.switches, cells)
+        ctx_ext, ids_ext = infer.build_context(enc, prompts_n, device, per_cell=True)   # [1, cells, 512, 4096]
+        sw = [k for k in range(1, cells) if prompts_n[k] != prompts_n[k - 1]]
+        log(f"   --switches {args.switches}：cell {sw} 处切换，{len(set(prompts_n))} 条不同的串")
+        prompts = prompts_n
+    off_ext = off_full if cells > G.CELLS else off21
+    # 整片并行路径喂的是绝对世界坐标（cell 0 恒为 0，本来就锚定）；rollout 路径在每个 cell 现算窗口锚定
+    off_par = window_offsets(off_ext[:, :min(cells, off_ext.shape[1])], 0, min(cells, off_ext.shape[1]) - 1)
+    return {"frame_off": frame_off, "cell_off": cell_off, "cells": cells, "prompts": prompts,
+            "context": ctx_ext[:, :cells], "context_ids": ids_ext[:, :cells],
+            "ctx_ext": ctx_ext, "ids_ext": ids_ext, "off_ext": off_ext, "off_par": off_par,
+            "ang_ext": build_angles(args, cells, off_ext.device, log=log),
+            "scale_ext": build_scales(args, cells, off_ext.device, log=log),
+            "first_spec": first_spec, "walk_speed": walk_speed.tolist()}
+
+
+# ----------------------------------------------------------------------------
+# 自检（tiny 模型，不碰真权重 / 不需要 VAE）
+# ----------------------------------------------------------------------------
+
+def self_test(log=print, device: str | None = None) -> None:
+    from diffsynth.models.wan_video_dit import WanModel
+    tiny = dict(dim=128, in_dim=48, ffn_dim=256, out_dim=48, text_dim=4096, freq_dim=256, eps=1e-6,
+                patch_size=(1, 2, 2), num_heads=4, num_layers=2, has_image_input=False)
+    # 尊重 --device：--self_test 不需要 GPU，别在训练用的卡上白建一个 CUDA 上下文
+    dev = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
+    if str(dev).startswith("cuda") and not torch.cuda.is_available():
+        dev = "cpu"
+    torch.manual_seed(0)
+    model = AropeDiT(dit=WanModel(**tiny).to(torch.bfloat16), mask_channel=True, causal=True)
+    model.eval().requires_grad_(False).to(dev)
+    unwrap(model).move_rope_table(torch.device(dev))
+    deploy = deploy_schedule(DEPLOY_DENOISE_LIST)
+    log(f"1. 部署表  t={[round(float(x), 1) for x in deploy[2]]} σ={[round(float(x), 4) for x in deploy[1]]}")
+
+    sc = torch.load(infer.SIDECAR, map_location="cpu", weights_only=False)
+    ws = np.array(sc["meta"]["walk_speed_px_s"], np.float64)
+    _assert_matches_parse_actions("right:1.0:21", ws)
+    fo, vel = infer.parse_actions("right:1.0:21", ws)
+    log(f"2. 指令递推 action_offsets(vel,21) 与 infer.parse_actions / frames_to_cells 逐位相同；"
+        f"frame_offset_px[80]=({fo[80, 0]:+.1f},{fo[80, 1]:+.1f}) ⇒ 期望背景位移 ({-fo[80, 0]:+.1f},{-fo[80, 1]:+.1f})")
+    N = 30
+    vel_ext = extend_velocity(vel, N)
+    fo_ext, co_ext = action_offsets(vel_ext, N)
+    assert np.allclose(fo_ext[:81], fo), "延长之后前 81 帧必须与 21 cell 的结果完全一致"
+    off_ext_t = torch.tensor(co_ext, dtype=torch.float32).unsqueeze(0).to(dev)
+    ref = extend_offsets(torch.tensor(G.frames_to_cells(fo), dtype=torch.float32).unsqueeze(0).to(dev),
+                         N - 1, mode="continue")
+    d = float((off_ext_t - ref).abs().max())
+    log(f"3. 延长口径 action_offsets(continue) vs data.extend_offsets('continue') max|Δ| {d:.3e} px "
+        f"（{N} cell，都是重复末段匀速）")
+    assert d < 1e-3, d
+
+    b, c, lh, lw, L = 1, 48, 8, 10, 4
+    z0 = torch.randn(b, c, 1, lh, lw, device=dev, dtype=torch.bfloat16)
+    ctx = torch.randn(b, G.CELLS, L, 4096, device=dev, dtype=torch.bfloat16)
+    ids = torch.zeros(b, G.CELLS, dtype=torch.int64, device=dev)
+    ctx_ext, ids_ext = extend_context(ctx, ids, N - 1, mode="hold")
+    lat = deploy_rollout(model, z0, ctx_ext, ids_ext, off_ext_t, N, deploy,
+                         kv_window=20, sink_size=0, rope_cap=20, seed=0, log=lambda *a: None)
+    assert lat.shape == (b, c, N, lh, lw) and torch.isfinite(lat).all()
+    log(f"4. rollout  {N} cell（窗口 20 ⇒ 驱逐 {max(0, N - 1 - 20)} 次）输出 {tuple(lat.shape)} 全部有限；"
+        f"|lat| 均值 {float(lat.float().abs().mean()):.4f}")
+
+    off_par = window_offsets(off_ext_t[:, :G.CELLS], 0, G.CELLS - 1)
+    lat_p = deploy_parallel(model, z0, ctx, ids, off_par, deploy, seed=0, log=lambda *a: None)
+    assert lat_p.shape == (b, c, G.CELLS, lh, lw) and torch.isfinite(lat_p).all()
+    assert torch.equal(lat_p[:, :, 0:1], z0), "parallel 路径必须每步把 cell 0 写回首帧 latent"
+    log(f"5. parallel {G.CELLS} cell 输出 {tuple(lat_p.shape)} 全部有限、cell 0 == 首帧 latent；"
+        f"|lat| 均值 {float(lat_p.float().abs().mean()):.4f}")
+
+    # 6. 用一段**已知平移**的合成画面把 measure() 的接线和符号约定钉死。
+    #    measure_displacement 的 sift 是**画面内容**在图像坐标里的位移；指令一侧比的是
+    #    −frame_offset_px（玩家往右走 ⇒ 背景往左挪），所以两者必须差一个负号。
+    rng = np.random.default_rng(0)
+    pad = 200
+    base = (rng.random((G.H + 2 * pad, G.W + 2 * pad)) * 255).astype(np.uint8)
+    d_content = np.array([6.0, -2.0])                    # 每帧内容位移（px）
+    syn = np.stack([np.stack([base[pad - int(round(d_content[1] * k)): pad - int(round(d_content[1] * k)) + G.H,
+                                   pad - int(round(d_content[0] * k)): pad - int(round(d_content[0] * k)) + G.W]] * 3,
+                             -1) for k in range(G.CELLS)], 0)
+    # 指令 = 内容位移的相反数 ⇒ 误差应当 ≈ 0
+    fo_syn = np.stack([-d_content * k for k in range(G.CELLS)], 0)
+    res = measure(syn, fo_syn, chain=8)
+    log(f"6. 实测接线 帧0→{res['k']} 合成平移真值 {(d_content * res['k']).tolist()} · "
+        f"sift={infer.fmt(res['measured']['sift'])}（{res['measured']['sift_source']}，"
+        f"{res['measured']['sift_direct_inliers']} 内点）· 指令 {infer.fmt(res['expected_bg_shift'])} · "
+        f"误差 {infer.fmt(res['error_px'])} ⇒ |误差| {res['error_norm_px']:.3f} px")
+    assert res["error_norm_px"] < 1.0, res
+    log("self_test OK")
+
+
+# ----------------------------------------------------------------------------
+# 主流程
+# ----------------------------------------------------------------------------
+
+def parse_args(argv=None):
+    ap = argparse.ArgumentParser(description="蒸馏产物的部署路径评测（少步 / AR rollout + 背景位移实测）")
+    ap.add_argument("--ckpt", default=None, help="蒸馏产物 safetensors；部署权重是 ema/，不是 dits/")
+    ap.add_argument("--sweep", default=None,
+                    help="glob（如 '.../ema/ema_step*.safetensors'）：按 step 依次评测。"
+                         "只加载一次 10 GB pipeline，之后只换 DiT 权重")
+    ap.add_argument("--allow_raw_student", action="store_true",
+                    help="放行路径里带 /dits/ 的 checkpoint。dits/ 只为续训存在，拿它上线是静默降质（risk 12）")
+    ap.add_argument("--mode", choices=["rollout", "parallel"], default="rollout",
+                    help="rollout：滑窗 KV 逐 cell 自回归（stage3 的部署路径）。"
+                         "parallel：整片块因果一次去噪（stage2 EMA 的部署路径）")
+    ap.add_argument("--causal", action=argparse.BooleanOptionalAction, default=True,
+                    help="蒸馏产物都是因果的。评双向 SFT 基线时用 --no-causal --mode parallel")
+    ap.add_argument("--denoise_list", default=",".join(str(int(x)) for x in DEPLOY_DENOISE_LIST),
+                    help="部署采样表（原始 t 值，严格降序）。默认就是用户拍板的 4 步 1000,750,500,250；"
+                         "stage2 EMA 单独部署时用 1000,250")
+    ap.add_argument("--cells", type=int, default=G.CELLS,
+                    help="生成多少个 latent cell。>21 只在 --mode rollout 下有意义（那正是 stage3 练的东西）")
+    ap.add_argument("--kv_window", type=int, default=20)
+    ap.add_argument("--sink_size", type=int, default=0)
+    ap.add_argument("--rope_cap", type=int, default=20)
+    ap.add_argument("--sink_anchor", choices=("absolute", "window"), default="absolute",
+                    help="sink_size>0 时世界坐标锚点：absolute = cell 0（当前 RoPE 位无界增长）；window = 滑动窗最老 cell，"
+                         "sink 段偏移钉 0（内容可见、RoPE 在包络内）。stage3 ckpt 元数据里有这个键时必须一致")
+    ap.add_argument("--allow_geometry_mismatch", action="store_true",
+                    help="ckpt 元数据里的训练几何（stage3 do_save 写的 kv_window/sink_size/rope_cap）与命令行"
+                         "不等时也跑。默认拒绝：stage3 学生只在训练几何下才是分布内的（GEO-1），"
+                         "sink 3 / w17 之类只作显式 A/B，不作 ckpt 选择轴。stage1/2 的 ckpt 没有这些键，不受限")
+    ap.add_argument("--ref_steps", type=int, default=0,
+                    help=">0 ⇒ 额外跑一遍 infer.sample 的 N 步扩散采样作参照（PORT_SPEC 的 S9，一般填 50）")
+    ap.add_argument("--sampler", choices=["auto", "flowmap", "renoise", "euler"], default="auto",
+                    help="auto（默认）= 按 checkpoint 自己判（832 键 + delta_embedding ⇒ flowmap）。"
+                         "两者接错都不会报错、只会安静地产出另一个模型的结果，所以显式指定时与权重不符会被拒。"
+                         "euler = 非 flow-map 权重 + --mode rollout 时，每个 cell 走 FlowMatchScheduler 的 "
+                         "Euler ODE（--euler_steps / --euler_shift，与 infer.sample 同一条 sched.step），不重新加噪")
+    ap.add_argument("--euler_steps", type=int, default=16, help="--sampler euler 的步数（longhorizon/README 用 16）")
+    ap.add_argument("--euler_shift", type=float, default=5.0, help="--sampler euler 的 shift（infer.sample 用 5）")
+    ap.add_argument("--num_inference_steps", type=int, default=4,
+                    help="flow-map 采样的步数 N（--sampler flowmap/auto 命中时才用）。"
+                         "同一份权重 4 步和 24 步都能采 —— 这正是 any-step 的主张，验收就看这两个的一致性。"
+                         "renoise 路径不看这个参数，它用 --denoise_list")
+
+    ap.add_argument("--replay", default=None, help="clip_id：用 sidecar 真值偏移 + 该 clip 首帧 + 原 prompt")
+    ap.add_argument("--first_frame", default=None, help="png / mp4 / clip_id")
+    ap.add_argument("--scene", default=None)
+    ap.add_argument("--trigger", default=None, help="burn:<B>:<start>:<n>（或 cut:），start < 21：训练语料口径的转场")
+    ap.add_argument("--rotate", type=float, default=0.0,
+                    help="相机每 cell 旋转多少**度**（累计；cell 0 恒为 0）。一切表达在当前相机系里：当前帧格点"
+                         "保持轴对齐，历史 token 按 φ_j−φ_k 转进来。0 = 不转（逐位走老路）")
+    ap.add_argument("--zoom", type=float, default=0.0,
+                    help="每 cell 的**位置缩放**率（判决实验，只接并行路径）：位置乘 (1+rate)^k。"
+                         "模型若把位置场的尺度读成视野大小，内容就会按 1/(1+rate)^k 缩小")
+    ap.add_argument("--rotate_mask", choices=sorted(ROT_MASK), default="rot",
+                    help="旋转时 known 掩码怎么算：rot = 跟着转的斜多边形（几何正确，四角变 new）；"
+                         "plain = 老的轴对齐掩码（消融：看画面缩小/摆动是不是掩码通道闹的）")
+    ap.add_argument("--rotate_spec", default=None,
+                    help='分段角速度 "2:20,0:20,-2:20"（每格度数:格数，铺不满就把末段续到底），与 --rotate 二选一')
+    ap.add_argument("--switches", default=None,
+                    help='多次场景切换脚本 "burn:B:10:2,burn:C:20:2,cut:D:30:2"（从 --scene 起，任意 cell，整条 rollout '
+                         '的逐 cell 串直接生成；与 --trigger/--late_trigger 互斥）')
+    ap.add_argument("--late_trigger", default=None,
+                    help="burn:<C>:<start>:<n>（或 cut:），start ≥ 21：rollout 视野外再切一次场景（cell 20 的场景 → C）")
+    ap.add_argument("--actions", default=None,
+                    help='动作脚本，如 "right:1.0:21"。cell 合计 == --cells 时铺满整条 rollout（如 loop_spec.txt 的 41 cell），'
+                         '否则必须正好 21 个 cell、cell ≥21 按末段速度续走')
+    ap.add_argument("--text_table", default=None, help="默认用 arope 臂的纯场景表")
+    ap.add_argument("--chain", type=int, default=16, help="SIFT 链式测量的分段长度")
+    ap.add_argument("--tiled", action="store_true", help="VAE 分块解码（长 rollout 显存吃紧时用）")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--out", default=None, help="输出 mp4（同名 _frame0.png / .json 一并写出）")
+    ap.add_argument("--out_dir", default=None, help="--sweep 时的输出目录")
+    ap.add_argument("--save_video", action=argparse.BooleanOptionalAction, default=True,
+                    help="--sweep 扫很多个 ckpt 时可以 --no-save_video 只留 json")
+    ap.add_argument("--gpu_headroom_gb", type=float, default=40.0,
+                    help="目标卡至少要有这么多空闲显存才肯跑（权重 9.3 GB + rollout + VAE 解码）")
+    ap.add_argument("--allow_busy_gpu", action="store_true",
+                    help="卡上有别的任务时也跑；会把本进程的显存**硬顶**在空闲量减 8 GB，"
+                         "撑死只炸自己，不会把同卡的训练挤 OOM")
+    ap.add_argument("--self_test", action="store_true", help="tiny 模型自检，不碰真权重、不需要 VAE")
+    args = ap.parse_args(argv)
+    if args.self_test:
+        return args
+    if not (args.ckpt or args.sweep):
+        raise SystemExit("要么 --ckpt，要么 --sweep")
+    if args.sweep and not args.out_dir:
+        raise SystemExit("--sweep 需要 --out_dir")
+    if args.ckpt and not args.out:
+        raise SystemExit("--ckpt 需要 --out")
+    if args.cells < 2:
+        raise SystemExit(f"--cells 至少 2，得到 {args.cells}")
+    if args.cells != G.CELLS and args.mode != "rollout":
+        raise SystemExit(f"--cells {args.cells} != {G.CELLS} 只在 --mode rollout 下成立"
+                         f"（parallel 是整片一次去噪，长度由文本/偏移的 cell 数决定，"
+                         f"超出 21 就没有真实训练分布可言）")
+    if args.rope_cap != args.sink_size + args.kv_window:
+        raise SystemExit(f"--rope_cap({args.rope_cap}) 必须 == sink({args.sink_size}) + window({args.kv_window})")
+    args.denoise = [float(t) for t in args.denoise_list.split(",") if t.strip()]
+    args.euler_sched = None
+    if args.sampler == "euler":
+        if args.mode != "rollout":
+            raise SystemExit("--sampler euler 只实现了 --mode rollout（parallel 请用 --ref_steps 走 infer.sample）")
+        from diffsynth.diffusion.flow_match import FlowMatchScheduler
+        args.euler_sched = FlowMatchScheduler("Wan")
+        args.euler_sched.set_timesteps(args.euler_steps, shift=args.euler_shift)
+    return args
+
+
+def _ckpt_list(args) -> list[str]:
+    paths = sorted(glob.glob(args.sweep), key=lambda p: int("".join(ch for ch in os.path.basename(p)
+                                                                   if ch.isdigit()) or 0)) if args.sweep \
+        else [args.ckpt]
+    if not paths:
+        raise SystemExit(f"--sweep 没匹配到任何文件：{args.sweep}")
+    for p in paths:
+        if "/dits/" in p.replace(os.sep, "/") and not args.allow_raw_student:
+            raise SystemExit(
+                f"拒绝评测 {p}：dits/ 是续训用的原始学生权重，部署权重是 ema/。\n"
+                "  拿 dits/ 上线不会报错，只会静默降质（PORT_SPEC risk 12）。"
+                "确实想看原始学生就加 --allow_raw_student。")
+        if not os.path.isfile(p):
+            raise SystemExit(f"checkpoint 不存在：{p}")
+    return paths
+
+
+def _require_gpu_headroom(device: str, need_gb: float, allow_busy: bool, log=print) -> None:
+    """评测默认**不许**跟正在跑的训练抢同一张卡。
+
+    为什么：本脚本要 9.3 GB 权重 + rollout（实测 N=26 峰值 12.6 GB）+ VAE 解码 81~101 帧，
+    几十 GB 的量级；蒸馏训练是 6~40 h 的长任务，被一次评测挤 OOM 的代价远大于等它跑完。
+    而文件头的示例命令写的正是 `CUDA_VISIBLE_DEVICES=0` —— 那恰好是训练在用的卡。
+
+    用 nvidia-smi 查（不是 torch.cuda.mem_get_info）：后者要先在目标卡上建 CUDA 上下文，
+    光是"看一眼"就先占几百 MB。查不到就拒绝启动，与 distill/common.sh 的 require_idle_gpus 同一套 fail-closed 口径。
+    --allow_busy_gpu 时不再拒绝，但把本进程钉死在 (空闲 − 8 GB)：共享可以，但必须是**有界**的共享。
+    """
+    if not str(device).startswith("cuda") or not torch.cuda.is_available():
+        return
+    idx = torch.device(device).index
+    if idx is None:                       # "cuda" ⇒ 当前默认卡（CUDA_VISIBLE_DEVICES 里的第 0 张）
+        idx = 0
+    phys = os.environ.get("CUDA_VISIBLE_DEVICES")
+    phys = phys.split(",")[idx].strip() if phys else str(idx)
+    try:
+        out = subprocess.run(["nvidia-smi", "-i", phys, "--query-gpu=memory.used,memory.total",
+                              "--format=csv,noheader,nounits"], capture_output=True, text=True,
+                             timeout=30, check=True).stdout.strip().splitlines()
+        used, total = (float(x) for x in out[0].split(","))
+    except Exception as e:                # 驱动半死 / 编号非法 / 超时 —— 查不到就不跑
+        raise SystemExit(f"ABORT: 查不到 GPU {phys} 的占用（{type(e).__name__}: {e}）。"
+                         "查不到就不许启动；要绕过请显式 --device cpu 或 --allow_busy_gpu。")
+    free = (total - used) / 1024.0        # MiB -> GiB
+    log(f"[gpu] 物理 GPU {phys}：已用 {used / 1024:.1f} / {total / 1024:.0f} GiB，空闲 {free:.1f} GiB")
+    if used / 1024.0 > 1.0 and not allow_busy:
+        raise SystemExit(
+            f"ABORT: GPU {phys} 上已经有 {used / 1024:.1f} GiB 被占用 —— 多半是蒸馏训练还在跑。"
+            f"评测要 ~{need_gb:.0f} GiB，别去挤它。等它跑完，或换一张空卡"
+            f"（CUDA_VISIBLE_DEVICES=<空闲卡>，**不要**动别人的卡），"
+            f"或 --allow_busy_gpu（会把本进程硬顶在 {max(free - 8, 4):.0f} GiB，只炸自己）。")
+    if free < need_gb and not allow_busy:
+        raise SystemExit(f"ABORT: GPU {phys} 只剩 {free:.1f} GiB 空闲 < --gpu_headroom_gb {need_gb:.0f}。")
+    if allow_busy:
+        cap = max(free - 8.0, 4.0)        # 留 8 GiB 给同卡的训练做碎片/峰值余量
+        torch.cuda.set_per_process_memory_fraction(min(cap / (total / 1024.0), 0.95), idx)
+        log(f"[gpu] --allow_busy_gpu：本进程显存硬顶 {cap:.0f} GiB（超了只 OOM 我自己）")
+
+
+def main(argv=None):
+    args = parse_args(argv)
+    if args.self_test:
+        self_test(device=args.device)
+        return 0
+    paths = _ckpt_list(args)
+    for p in paths:                       # 只读 safetensors 头：训练几何 vs 命令行（GEO-1），权重还没碰
+        _check_ckpt_geometry(p, args)
+    # ★ 在加载 10 GB pipeline 之前先问一句这张卡是不是别人的（fail-closed，见函数注释）
+    _require_gpu_headroom(args.device, args.gpu_headroom_gb, args.allow_busy_gpu)
+    deploy = deploy_schedule(args.denoise)
+    print(f"部署表 {args.denoise} ⇒ 网格 t={[round(float(x), 1) for x in deploy[2]]} "
+          f"σ={[round(float(x), 4) for x in deploy[1]]} · mode={args.mode} · cells={args.cells}")
+
+    # ---- 采样器分流：按 **checkpoint 本身**（832 键 + flowmap metadata）决定，不看命令行心情 ----
+    # 拿 renoise 采样器去采 flow-map 权重不会报错，只会安静地产出另一个模型的结果 ——
+    # 所以默认 auto 从权重里读，而不是让调用方记得加旗标。
+    is_fm = [runlib.is_flowmap_ckpt(p) for p in paths]
+    if args.sampler == "auto":
+        use_flowmap = is_fm[0]
+    else:
+        use_flowmap = (args.sampler == "flowmap")   # renoise / euler 都是非 flow-map 权重的采样器
+        if use_flowmap != is_fm[0]:
+            raise SystemExit(
+                f"--sampler {args.sampler} 与 checkpoint 不符：{paths[0]} "
+                f"{'是' if is_fm[0] else '不是'} flow-map 权重（832 键 + delta_embedding）。\n"
+                f"       两者接错都不会报错，只会安静地产出另一个模型的结果，所以这里直接拒绝。")
+    if len(set(is_fm)) > 1:
+        raise SystemExit("--sweep 里混了 flow-map 与非 flow-map 的 checkpoint；分开扫。")
+
+    from safetensors.torch import load_file
+    t0 = time.time()
+    if use_flowmap:
+        # 832 键喂不进 load_models 的 strict 加载（它只认 827），所以：基座 → 装壳 → 灌权重。
+        # 顺序不能反：install_flowmap 会把读出层零初始化，装完再 load 才拿得到训出来的 delta。
+        meta = _ckpt_metadata(paths[0])
+        pipe, dit, _ = infer.load_models(None, True, args.device, need_vae=True, arm="arope")
+        model = AropeDiT(dit=dit, mask_channel=True, causal=args.causal)
+        unwrap(model).install_flowmap(form=meta.get("flowmap_form", "residual"),
+                                      gate=float(meta.get("delta_gate", 1.0)))
+        unwrap(model).dit.load_state_dict(load_file(paths[0]), strict=True)
+        has_mask = True
+    else:
+        pipe, dit, _ = infer.load_models(paths[0], None, args.device, need_vae=True, arm="arope")
+        has_mask = hasattr(dit, "arope_mask_embedding")
+        model = AropeDiT(dit=dit, mask_channel=has_mask, causal=args.causal)
+    model.eval().requires_grad_(False)
+    unwrap(model).move_rope_table(torch.device(args.device))
+    sig_any = fm.anystep_sigmas(args.num_inference_steps) if use_flowmap else None
+    print(f"模型加载 {time.time() - t0:.0f}s · causal={args.causal} · mask_channel={has_mask} · "
+          f"采样器={'flowmap' if use_flowmap else ('euler' if args.euler_sched is not None else 'renoise')}"
+          + (f" · Euler {args.euler_steps} 步 shift {args.euler_shift} t={[round(float(t)) for t in args.euler_sched.timesteps]}"
+             if args.euler_sched is not None else "")
+          + (f" · N={args.num_inference_steps} 步 σ={[round(float(x), 4) for x in sig_any]}"
+             if use_flowmap else ""))
+
+    table_key = "arope"
+    enc = infer.TextEncoder(args.text_table or infer.TEXT_TABLE[table_key], args.device)
+    cond = build_condition(args, enc, args.device)
+    fo = cond["frame_off"]
+    print(f"指令：{cond['cells']} cell · frame_offset_px[末帧]=({fo[-1, 0]:+.1f},{fo[-1, 1]:+.1f}) "
+          f"⇒ 期望背景位移 ({-fo[-1, 0]:+.1f},{-fo[-1, 1]:+.1f})")
+    img = infer.load_first_frame(cond["first_spec"])
+    z0 = infer.encode_first_frame(pipe, img, tiled=args.tiled)
+
+    from safetensors.torch import load_file
+    rows = []
+    for p in paths:
+        if p != paths[0]:
+            # 只换 DiT 权重，不重建 10 GB pipeline（sweep 的全部意义）
+            unwrap(model).dit.load_state_dict(load_file(p), strict=True)
+        step = "".join(ch for ch in os.path.basename(p) if ch.isdigit())
+        stem = (os.path.splitext(args.out)[0] if args.out
+                else os.path.join(args.out_dir, os.path.splitext(os.path.basename(p))[0]))
+        os.makedirs(os.path.dirname(os.path.abspath(stem)), exist_ok=True)
+        print(f"\n=== {os.path.basename(p)} ===")
+        torch.cuda.reset_peak_memory_stats() if args.device.startswith("cuda") else None
+        res, frames, lat = run_one(model, pipe, z0, cond, args, deploy, sigmas=sig_any)
+        info = {"ckpt": p, "step": int(step) if step else None, "mode": args.mode,
+                "sampler": ("flowmap" if use_flowmap else ("euler" if args.euler_sched is not None else "renoise")),
+                "euler": ({"steps": args.euler_steps, "shift": args.euler_shift} if args.euler_sched is not None else None),
+                "num_inference_steps": (args.num_inference_steps if use_flowmap else None),
+                "anystep_sigmas": ([round(float(x), 6) for x in sig_any] if use_flowmap else None),
+                "denoise_list": (None if use_flowmap else args.denoise),
+                "cells": cond["cells"], "seed": args.seed,
+                "kv_window": args.kv_window, "sink_size": args.sink_size, "rope_cap": args.rope_cap,
+                "prompts": sorted(set(cond["prompts"])), "walk_speed_px_s": cond["walk_speed"],
+                "frame_offset_px_last": fo[-1].tolist(), **res}
+        if args.device.startswith("cuda"):
+            info["peak_mem_gb"] = torch.cuda.max_memory_allocated() / 1e9
+        m = res["measured"]
+        print(f"实测位移 帧0→{res['k']}：sift={infer.fmt(m['sift'])}（{m['sift_source']}）"
+              f" · 指令 {infer.fmt(res['expected_bg_shift'])} · 误差 {infer.fmt(res['error_px'])}"
+              + (f" · |误差| {res['error_norm_px']:.1f} px = {100 * res['error_rel']:.1f}%"
+                 if res.get("error_rel") is not None else ""))
+
+        if args.ref_steps > 0:
+            # PORT_SPEC 的 S9：同一条件下的 N 步扩散参照（infer.sample 的 Euler 步，不是少步表）
+            lat_ref = infer.sample(unwrap(model).dit, z0, cond["context"][:, :G.CELLS],
+                                   cond["context_ids"][:, :G.CELLS],
+                                   cond["off_par"][:, :G.CELLS],
+                                   make_mask_input(cond["off_par"][:, :G.CELLS], z0.shape[3], z0.shape[4]),
+                                   steps=args.ref_steps, shift=5.0, cfg=1.0, seed=args.seed)
+            fr_ref = infer.decode_video(pipe, lat_ref, tiled=args.tiled)
+            info["ref"] = {"steps": args.ref_steps, **measure(fr_ref, fo[:len(fr_ref)], chain=args.chain)}
+            rm = info["ref"]["measured"]
+            print(f"{args.ref_steps} 步参照：sift={infer.fmt(rm['sift'])} · "
+                  f"误差 {infer.fmt(info['ref']['error_px'])}")
+
+        if args.replay:
+            gt = infer.read_frames(infer.find_clip_mp4(args.replay))
+            n = min(len(gt), len(frames))
+            per = [infer.psnr(frames[i], gt[i]) for i in range(n)]
+            info["psnr_frame0"] = per[0]
+            info["psnr_mean_rest"] = float(np.mean(per[1:])) if n > 1 else None
+            info["gt_measured"] = infer.measure_displacement(infer.to_gray(gt[:n]), k=n - 1, chain=args.chain)
+            print(f"与 GT 比：帧 0 PSNR {per[0]:.1f} dB · 其余平均 {info['psnr_mean_rest']:.1f} dB · "
+                  f"GT 实测 sift={infer.fmt(info['gt_measured']['sift'])}")
+
+        if args.save_video:
+            from PIL import Image
+            from diffsynth.utils.data import save_video
+            save_video(list(frames), stem + ".mp4", fps=int(G.FPS), quality=9)
+            Image.fromarray(img).save(stem + "_frame0.png")
+            info["out"] = stem + ".mp4"
+        with open(stem + ".json", "w", encoding="utf-8") as fh:
+            json.dump(info, fh, ensure_ascii=False, indent=1, default=float)
+        rows.append(info)
+        del frames, lat
+
+    if len(rows) > 1:
+        print("\n=== sweep 汇总（按 |误差| 从小到大；部署权重就从这里挑）===")
+        ok = [r for r in rows if r.get("error_norm_px") is not None]
+        for r in sorted(ok, key=lambda r: r["error_norm_px"]):
+            print(f"  step {str(r['step']):>6}  |误差| {r['error_norm_px']:7.1f} px "
+                  f"({100 * r['error_rel']:5.1f}%)  sift={infer.fmt(r['measured']['sift'])}  "
+                  f"{os.path.basename(r['ckpt'])}")
+        for r in rows:
+            if r.get("error_norm_px") is None:
+                print(f"  step {str(r['step']):>6}  实测失败（SIFT 匹配不上；画面可能已经崩了）  "
+                      f"{os.path.basename(r['ckpt'])}")
+        csv = os.path.join(args.out_dir, "sweep.csv")
+        with open(csv, "w", encoding="utf-8") as fh:
+            fh.write("step,ckpt,err_px,err_rel,sift_dx,sift_dy,exp_dx,exp_dy,sample_sec\n")
+            for r in rows:
+                s = r["measured"]["sift"] or [float("nan")] * 2
+                e = r["expected_bg_shift"]
+                fh.write(f"{r['step']},{os.path.basename(r['ckpt'])},"
+                         f"{r.get('error_norm_px', float('nan'))},{r.get('error_rel', float('nan'))},"
+                         f"{s[0]},{s[1]},{e[0]},{e[1]},{r['sample_sec']:.1f}\n")
+        print(f"汇总写到 {csv}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

backup_2026-09-14/code/distill/eval/sgf_parity_probe.py

@@ -0,0 +1,112 @@
+"""SGF 两趟 replay 的几何等价性自检 —— 上 SGF 之前必须过的那一关。
+
+SGF 的 `self_gradient_forcing_match_context` 要求 pass 2 喂的上下文与 pass 1 **真正提交进
+KV 缓存**的东西一致。在 ARoPE 下这不是"传同一个张量"就完事的：
+  pass 1  每个 cell 锚在**自己窗口里最老那一格**（frame_cond → window_offsets_for_frame）
+  pass 2  整片并行、first_frame_cond=True、锚在**片自己的 cell 0**
+两者只在整窗平移上不同，而 RoPE 是相对编码、平移不改分数 —— 但这个论证只在
+**rollout 长度 == 打分片长**（缓存跨度与并行片完全重合）时成立。
+
+所以判据是数值的：pass 1 流式在退出档算出的 x0，与 pass 2 并行 TF 在同一档算出的 x0，
+逐 cell 比对。对不上就说明几何/掩码分叉了，这时候去训练只会得到"loss 照跑、就是不涨"。
+
+  python distill/eval/sgf_parity_probe.py --ckpt <832 键 flow map> [--cells 21] [--steps 4]
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import numpy as np
+import torch
+
+ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.insert(0, ROOT)
+os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
+
+from actionrope import geometry as G          # noqa: E402
+from actionrope import infer                  # noqa: E402
+from actionrope.prompts import NULL_PROMPT    # noqa: E402
+from distill import flowmap as fm             # noqa: E402
+from distill import onpolicy as op            # noqa: E402
+from distill.arope_dit import AropeDiT, unwrap                   # noqa: E402
+from distill.data import make_mask_input, window_context, window_offsets  # noqa: E402
+from distill.eval.rollout_eval import _ckpt_metadata             # noqa: E402
+from distill.stage3_dmd.dmd_arope import frame_cond              # noqa: E402
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(description="SGF 两趟 replay 的几何等价性自检")
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--cells", type=int, default=21, help="1+K；必须与训练时的片长一致")
+    ap.add_argument("--steps", type=int, default=4)
+    ap.add_argument("--kv_window", type=int, default=20)
+    ap.add_argument("--sink_size", type=int, default=0)
+    ap.add_argument("--rope_cap", type=int, default=20)
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--tol", type=float, default=0.05, help="逐 cell 相对误差上限")
+    args = ap.parse_args(argv)
+
+    from safetensors.torch import load_file
+    meta = _ckpt_metadata(args.ckpt)
+    pipe, dit, _ = infer.load_models(None, True, args.device, need_vae=False, arm="arope")
+    model = AropeDiT(dit=dit, mask_channel=True, causal=True)
+    unwrap(model).install_flowmap(form=meta.get("flowmap_form", "residual"),
+                                  gate=float(meta.get("delta_gate", 1.0)))
+    unwrap(model).dit.load_state_dict(load_file(args.ckpt), strict=True)
+    model.eval().requires_grad_(False)
+    dev = torch.device(args.device)
+    unwrap(model).move_rope_table(dev)
+
+    enc = infer.TextEncoder(infer.TEXT_TABLE["arope"], args.device)
+    ctx, ids = infer.build_context(enc, [NULL_PROMPT] * G.CELLS, args.device, per_cell=True)
+    N = args.cells - 1
+    ctx_ext = ctx[:, :1].repeat(1, args.cells, 1, 1)
+    ids_ext = ids[:, :1].repeat(1, args.cells)
+
+    # 一条往右匀速的轨迹（内容不重要，重要的是每个 cell 的世界坐标都不同）
+    sc = torch.load(infer.SIDECAR, map_location="cpu", weights_only=False)
+    ws = np.array(sc["meta"]["walk_speed_px_s"], np.float64)
+    cell_off = np.stack([np.array([ws[0], 0.0]) * k * 4 / G.FPS for k in range(args.cells)], 0)
+    off_ext = torch.tensor(cell_off, dtype=torch.float32, device=dev).unsqueeze(0)
+
+    g = torch.Generator("cpu").manual_seed(0)
+    z0 = torch.randn(1, 48, 1, G.LAT_H, G.LAT_W, generator=g).to(dev, torch.bfloat16)
+    sig = fm.anystep_sigmas(args.steps)
+
+    print(f"片长 {args.cells} cell · N={args.steps} 步 · kv_window={args.kv_window} "
+          f"sink={args.sink_size} · σ={[round(float(x),4) for x in sig]}")
+    ok_all = True
+    for exit_idx in range(args.steps):
+        with torch.autocast("cuda", torch.bfloat16, enabled=(dev.type == "cuda")), torch.no_grad():
+            noisy, x0_cache = op.sgf_rollout(
+                model, z0, ctx_ext, ids_ext, off_ext, N, sig, exit_idx,
+                kv_window=args.kv_window, rope_cap=args.rope_cap, sink_size=args.sink_size,
+                gen=torch.Generator("cpu").manual_seed(1), unwrap_fn=unwrap, frame_cond_fn=frame_cond)
+            off_t = window_offsets(off_ext, 0, N)
+            ctx_t, ids_t = window_context(ctx_ext, ids_ext, 0, N)
+            cond = dict(context=ctx_t, context_ids=ids_t, offset_px=off_t,
+                        mask_input=make_mask_input(off_t, G.LAT_H, G.LAT_W), first_frame_cond=True)
+            x0_p2 = op.sgf_replay(model, noisy, x0_cache, cond, sig, exit_idx, context_noise=0.0)
+
+        a = x0_cache[:, :, 1:].float()
+        b = x0_p2[:, :, 1:].float()
+        # 逐 cell 相对 L2：把 (C,H,W) 摊平再取范数（torch.norm 的 dim 只吃 2-tuple）
+        af = a[0].permute(1, 0, 2, 3).flatten(1)          # [cells, C*H*W]
+        bf = b[0].permute(1, 0, 2, 3).flatten(1)
+        rel = (af - bf).norm(dim=1) / (af.norm(dim=1) + 1e-9)
+        ok = bool((rel < args.tol).all())
+        ok_all &= ok
+        print(f"  exit_idx={exit_idx} (σ {float(sig[exit_idx]):.4f}→{float(sig[exit_idx+1]):.4f})  "
+              f"逐 cell 相对误差 中位 {rel.median():.4f} 最大 {rel.max():.4f} "
+              f"(cell {int(rel.argmax())+1})  {'✓' if ok else '✗ 分叉'}")
+    print("\n" + ("通过：pass 1 的流式几何与 pass 2 的并行几何一致，match_context 成立。"
+                  if ok_all else
+                  "★ 未通过：两趟看到的不是同一段历史。别在这个状态下训 SGF —— "
+                  "loss 会照跑，模型学到的却是另一条上下文。"))
+    return 0 if ok_all else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

backup_2026-09-14/code/distill/eval/sweep8.sh

@@ -0,0 +1,73 @@
+#!/usr/bin/env bash
+# stage2-flowmap 的里程碑扫描：把 (ckpt × 步数) 的组合按卡分片，8 张卡并行。
+#
+# rollout_eval.py 自己是单卡的（一次只加载一份 10 GB pipeline），所以并行靠分片而不是靠它。
+# 每个组合一个独立进程、钉死一张卡；重复加载 pipeline 的 11 s 相对 rollout 的 24~82 s 可以忽略。
+#
+#   bash distill/eval/sweep8.sh [OUTDIR]
+#
+# 判据两个（都不是 loss —— flow map 的 loss 是自举目标，平的）：
+#   净漂移 px   绕一圈回到起点的误差，指令是 0。衡量动作服从性。
+#   4 vs 24     同一份权重两个步数的一致性。这是 any-step 主张本身。
+#
+# 几何：默认**不传** --sink_size/--kv_window/--rope_cap ⇒ 用 rollout_eval.py 的默认 sink 0 / window 20 /
+#   cap 20 = stage3 的训练几何（distill/stage3_dmd/run.sh:57-59）。stage3 学生只在训练几何下才是
+#   分布内的（GEO-1：sink 3 / w17 与它只在 k≤20 的保留集相同，之后的 h/w RoPE 与掩码通道是训练里
+#   从未出现过的输入，拿它挑 EMA 只是在比 OOD 外推）。sink 3 / w17 只作**显式 A/B**，走环境变量：
+#     SINK=3 KV_WINDOW=17 ROPE_CAP=20 bash distill/eval/sweep8.sh
+#   rollout_eval.py 会拿 stage3 ckpt 元数据里的训练几何核对命令行，不等就拒（--allow_geometry_mismatch 放行）。
+# 动作脚本：loop_spec.txt 是丢失的原始脚本的**重建**——41 cell 的闭环，净期望背景位移恰为 0：
+#   idle 放在 cell 0（那是首帧锚点，cell 0 的速度永远不生效），之后 右/下/左/上 各 10 cell = 各 40 帧，
+#   横纵各自抵消。若把 idle 放末尾，right 段只剩 9 个生效 cell，净偏 −walk_speed_x/4 ≈ −17 px。
+set -euo pipefail
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+ROOT="$(cd "$HERE/../.." && pwd)"
+OUTDIR="${1:-$ROOT/outputs/distill/eval/sweep_stage2fm}"
+EMA_DIR="${EMA_DIR:-$ROOT/outputs/distill/stage2_flowmap/ema}"
+SPEC_FILE="${SPEC_FILE:-$ROOT/distill/eval/loop_spec.txt}"
+FIRST_FRAME="${FIRST_FRAME:-$ROOT/data/eval/eybx_bg_50/19_medieval_castle.png}"
+STEPS_LIST="${STEPS_LIST:-4 24}"
+CELLS="${CELLS:-41}"
+NGPU="${NGPU:-8}"
+
+export PYTHONPATH="$ROOT" DIFFSYNTH_SKIP_DOWNLOAD=True CUDA_DEVICE_ORDER=PCI_BUS_ID
+PY="$ROOT/.venv/bin/python"
+mkdir -p "$OUTDIR"
+[ -f "$SPEC_FILE" ] || { echo "ABORT: 没有动作脚本 $SPEC_FILE"; exit 1; }
+SPEC="$(cat "$SPEC_FILE")"
+# 几何旗标只在环境变量显式给出时才追加；默认一个都不传 ⇒ rollout_eval.py 的默认 sink 0 / window 20 /
+# cap 20 = stage3 训练几何。三者必须满足 rope_cap == sink + window（rollout_eval 会核）。
+GEO=()
+if [ -n "${SINK:-}" ]; then GEO+=(--sink_size "$SINK"); fi
+if [ -n "${KV_WINDOW:-}" ]; then GEO+=(--kv_window "$KV_WINDOW"); fi
+if [ -n "${ROPE_CAP:-}" ]; then GEO+=(--rope_cap "$ROPE_CAP"); fi
+
+# 先确认这 8 张卡是空的 —— 与 common.sh 同一套 fail-closed 口径（这里不 source 它，因为
+# rollout_eval 是单卡工具、不走 accelerate；但"别人在用就别抢"这条规矩一样适用）。
+busy="$(nvidia-smi -i "$(seq -s, 0 $((NGPU-1)))" --query-gpu=index,memory.used \
+        --format=csv,noheader,nounits 2>&1)" || { echo "ABORT: nvidia-smi 查询失败：$busy"; exit 1; }
+echo "$busy" | awk -F', *' '$2>500 {print "  GPU"$1"="$2"MiB"; f=1} END{exit f+0}' || {
+  echo "ABORT: 上面这些卡上有占用（>500 MiB）。别抢别人的卡。"; exit 1; }
+echo "[guard] GPU 0-$((NGPU-1)) 空闲"
+
+jobs_list=()
+for ck in $(ls "$EMA_DIR"/ema_step*.safetensors | sort -t_ -k3 -n); do
+  for n in $STEPS_LIST; do jobs_list+=("$ck|$n"); done
+done
+echo "[sweep] ${#jobs_list[@]} 个组合 → $NGPU 张卡，输出 $OUTDIR"
+
+i=0
+for j in "${jobs_list[@]}"; do
+  ck="${j%|*}"; n="${j#*|}"
+  gpu=$(( i % NGPU )); i=$(( i + 1 ))
+  tag="$(basename "$ck" .safetensors)_n${n}"
+  CUDA_VISIBLE_DEVICES=$gpu nohup "$PY" "$HERE/rollout_eval.py" \
+    --ckpt "$ck" --mode rollout --cells "$CELLS" "${GEO[@]}" \
+    --num_inference_steps "$n" --actions "$SPEC" \
+    --first_frame "$FIRST_FRAME" --tiled --no-save_video \
+    --out "$OUTDIR/$tag.mp4" > "$OUTDIR/$tag.log" 2>&1 &
+  # 每张卡同时最多一个（一个进程 ~17 GB，同卡叠两个也放得下，但会互相拖慢且不好读数）
+  if (( i % NGPU == 0 )); then wait; fi
+done
+wait
+echo "[sweep] 全部完成"

backup_2026-09-14/code/distill/eval/sweep_photodrift.sh

@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+# stage3 里程碑扫描：把一个 run 的所有 ema_step*.safetensors 按卡分片，每张卡跑一个
+# 41-cell rollout，汇总成一张逐 cell 的 std/mean 表。
+#
+#   bash distill/eval/sweep_photodrift.sh <run_dir> [steps] [ngpu]
+#
+# 为什么判据是 std/mean 而不是 loss：DMD 的 loss 按构造就没有数值意义（它只是把预先算好的
+# grad 塞进 autograd 的载体），flow-map 的 loss 是自举目标、是平的。stage3 的失败模式是
+# **光度漂移被自回归递推放大**，只在长 rollout 上显形 —— 所以只能这么量。
+#
+# 基线（stage2-flowmap ema_step7500，同样 41 cell / N=4）：
+#   std  39.4 37.9 37.8 38.9 36.9 40.4 39.1 28.9
+#   mean 60.7 57.2 56.5 55.4 54.5 51.3 63.5 42.9      净漂移 9.3 px
+# 及格 = 不比它差；成功 = 净漂移 < 9.3 px。
+# ⚠ 基线是旧几何（sink 3 / w17）+ 原始 loop spec 下量的；现默认训练几何（sink 0 / w20）+ 重建 loop_spec.txt，
+#   stage2 跑完后重测 stage2 EMA 并更新上面 4 行（GEO-1）。
+set -euo pipefail
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+ROOT="$(cd "$HERE/../.." && pwd)"
+RUN="${1:?用法: sweep_photodrift.sh <run_dir> [steps] [ngpu]}"
+N="${2:-4}"; NGPU="${3:-8}"
+# ★ 并发上限不是卡数：每个 rollout_eval 进程要先把 10 GB 权重读进**宿主内存**再上卡，
+#   8 个同时起会打出瞬时 RAM 尖峰，实测被 OOM killer 全部收走（4 个输出文件都是 0 字节，
+#   看起来像"跑完了没结果"）。3~4 个足够把 8 张卡喂饱，因为瓶颈是 rollout 不是加载。
+MAXJOBS="${MAXJOBS:-3}"
+EMA_DIR="$RUN/ema"
+[ -d "$EMA_DIR" ] || { echo "ABORT: 没有 $EMA_DIR"; exit 1; }
+
+busy="$(nvidia-smi -i "$(seq -s, 0 $((NGPU-1)))" --query-gpu=index,memory.used \
+        --format=csv,noheader,nounits 2>&1)" || { echo "ABORT: nvidia-smi 查询失败：$busy"; exit 1; }
+echo "$busy" | awk -F', *' '$2>500 {print "  GPU"$1"="$2"MiB"; f=1} END{exit f+0}' || {
+  echo "ABORT: 上面这些卡上有占用（>500 MiB）。别抢别人的卡。"; exit 1; }
+
+i=0
+for ck in $(ls "$EMA_DIR"/ema_step*.safetensors | sort -t_ -k3 -n); do
+  tag="$(basename "$RUN")_$(basename "$ck" .safetensors)_n$N"
+  bash "$HERE/photodrift.sh" "$ck" "$tag" $(( i % NGPU )) "$N" \
+       > "$ROOT/outputs/distill/eval/deploy/$tag.txt" 2>&1 &
+  i=$(( i + 1 ))
+  if (( i % MAXJOBS == 0 )); then wait; fi
+done
+wait
+echo "=========== 汇总（基线 std 39.4→28.9 / mean 60.7→42.9 / 漂移 9.3 px）==========="
+for f in "$ROOT/outputs/distill/eval/deploy/$(basename "$RUN")"_ema_step*_n$N.txt; do
+  [ -f "$f" ] && cat "$f"
+done

backup_2026-09-14/code/distill/eval/tf_rollout_probe.py

@@ -0,0 +1,135 @@
+"""教师强制 rollout 探针 —— 把「部署路径有 bug」和「曝光偏差」分开。
+
+同一条 KV-cache 自回归路径，只改提交进缓存的东西：
+
+  self   —— 提交模型自己的 x0 估计（= 真正的部署路径 / deploy_rollout）
+  gt     —— 提交**真值 latent**（= 训练时 TF 的 clean 流）
+
+两者的差别只有历史的来源。所以：
+  gt 干净、self 崩  ⇒ 路径没问题，崩的是误差累积（曝光偏差），正是 stage3 self-forcing 要修的
+  gt 也崩           ⇒ 部署路径与训练不匹配（t=0 提交 / 窗口锚定的 RoPE 或 mask），是移植 bug
+
+每个 cell 都拿它的 x0 估计与真值算 PSNR，于是能看出崩塌是从哪一个 cell 开始的。
+
+  CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/data/zhiyangdeng/ActionRoPE DIFFSYNTH_SKIP_DOWNLOAD=True \
+    .venv/bin/python distill/eval/tf_rollout_probe.py \
+      --ckpt outputs/distill/stage1_ar_tf/dits/dit_step9000.safetensors \
+      --clip clip_Eybx_100000092_000000
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import torch
+
+ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from diffsynth.diffusion.flow_match import FlowMatchScheduler  # noqa: E402
+
+from actionrope import geometry as G  # noqa: E402
+from actionrope import infer  # noqa: E402
+from distill.arope_dit import AropeDiT, unwrap  # noqa: E402
+from distill.stage3_dmd.dmd_arope import frame_cond  # noqa: E402
+
+
+def psnr(a: torch.Tensor, b: torch.Tensor) -> float:
+    """latent 空间的 PSNR（同一把尺子量所有 cell，只用来看相对趋势）。"""
+    mse = (a.float() - b.float()).pow(2).mean().item()
+    if mse <= 0:
+        return float("inf")
+    rng = (b.float().max() - b.float().min()).item()
+    return 20.0 * torch.log10(torch.tensor(rng / (mse ** 0.5))).item()
+
+
+@torch.no_grad()
+def probe(model, gt_lat, ctx, ids, off, t_grid, sig, *, commit_source: str,
+          kv_window: int, sink_size: int, rope_cap: int, seed: int = 0):
+    dev = gt_lat.device
+    b, c, cells, lh, lw = gt_lat.shape
+    bare = unwrap(model)
+    cache = bare.init_kv_cache(kv_window=kv_window, sink_size=sink_size, rope_cap=rope_cap)
+    g = torch.Generator("cpu").manual_seed(seed)
+
+    def call(x, t, k, commit, ow, mk):
+        return model(x, t, ctx[:, k:k + 1], context_ids=ids[:, k:k + 1],
+                     first_frame_cond=False, frame_index=k, kv_cache=cache,
+                     commit=commit, offset_win=ow, mask_in_k=mk)
+
+    ow0, mk0 = frame_cond(off, 0, sink_size, kv_window, lh, lw)
+    call(gt_lat[:, :, 0:1], torch.zeros(b, device=dev), 0, True, ow0, mk0)   # cell 0 = 真首帧
+
+    rows = []
+    for k in range(1, cells):
+        ow, mk = frame_cond(off, k, sink_size, kv_window, lh, lw)
+        x = torch.randn(b, c, 1, lh, lw, generator=g).to(dtype=gt_lat.dtype, device=dev)
+        den = None
+        for i in range(len(t_grid)):
+            ts = torch.full((b,), float(t_grid[i]), device=dev)
+            v = call(x, ts, k, False, ow, mk)
+            den = x - float(sig[i]) * v
+            if i + 1 < len(t_grid):
+                ns = float(sig[i + 1])
+                x = (1.0 - ns) * den + ns * torch.randn(den.shape, generator=g).to(den)
+        gt_k = gt_lat[:, :, k:k + 1]
+        rows.append((k, psnr(den, gt_k), den.float().abs().mean().item(), gt_k.float().abs().mean().item()))
+        # ★ 唯一的差别：提交谁的历史
+        commit_lat = gt_k if commit_source == "gt" else den.detach()
+        call(commit_lat, torch.zeros(b, device=dev), k, True, ow, mk)
+    return rows
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--clip", required=True)
+    ap.add_argument("--steps", type=int, default=50)
+    ap.add_argument("--kv_window", type=int, default=20)
+    ap.add_argument("--sink_size", type=int, default=0)
+    ap.add_argument("--rope_cap", type=int, default=None)
+    ap.add_argument("--device", default="cuda")
+    args = ap.parse_args()
+    if args.rope_cap is None:
+        args.rope_cap = args.sink_size + args.kv_window
+
+    dev = args.device
+    lat = torch.load(f"{ROOT}/data/latent/val_eybx/{args.clip}.pt", map_location="cpu",
+                     weights_only=False)["latent"][None].to(dev)          # [1,48,21,30,52]
+    sc = torch.load(infer.SIDECAR, map_location="cpu", weights_only=False)
+    off = sc["clips"][args.clip]["offset_px"][None].float().to(dev)       # [1,21,2]
+
+    row = infer.load_meta_row(args.clip)
+    prompts = [infer.strip_action(p) for p in row["prompts"]]
+    table = torch.load(f"{ROOT}/data/latent/text_table_actionrope.pt", map_location="cpu",
+                       weights_only=False)
+    keys = sorted(set(prompts))                       # context_ids 才是权威的分组依据
+    kmap = {p: i for i, p in enumerate(keys)}
+    ctx = torch.stack([table[p] for p in prompts])[None].to(dev).to(torch.bfloat16)
+    ids = torch.tensor([kmap[p] for p in prompts])[None].to(dev)
+
+    from safetensors.torch import load_file
+    model = AropeDiT(causal=True, mask_channel=True).to(dev).eval()
+    unwrap(model).dit.load_state_dict(load_file(args.ckpt), strict=True)
+    model = model.to(torch.bfloat16)
+
+    s = FlowMatchScheduler("Wan")
+    s.set_timesteps(args.steps, training=False)
+    t_grid = [float(t) for t in s.timesteps]
+    sig = [float(x) for x in s.sigmas]
+
+    print(f"clip={args.clip} cells={lat.shape[2]} steps={args.steps} "
+          f"window={args.kv_window} sink={args.sink_size}")
+    for src in ("gt", "self"):
+        rows = probe(model, lat, ctx, ids, off, t_grid, sig, commit_source=src,
+                     kv_window=args.kv_window, sink_size=args.sink_size, rope_cap=args.rope_cap)
+        print(f"\n=== 提交 {src} 历史 ===")
+        print("  cell   PSNR(x0 vs GT)   |x0|    |GT|")
+        for k, p, a, b_ in rows:
+            print(f"  {k:4d}   {p:12.2f}   {a:6.3f}  {b_:6.3f}")
+
+
+if __name__ == "__main__":
+    main()

backup_2026-09-14/code/distill/flow_map.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+"""**别名转发**。规范模块是 `distill/flowmap.py`（任务书点名的文件名）；设计文档的
+target_layout 里写的是 `flow_map.py`。两个名字都能 import，免得 trainer / sampler 按另一个
+名字写就在启动时 ImportError。
+
+新代码请一律 `from distill import flowmap`。这个文件不加任何逻辑，改动只发生在 flowmap.py。
+"""
+from distill.flowmap import *      # noqa: F401,F403
+from distill import flowmap as _fm  # noqa: E402
+
+__all__ = list(_fm.__all__)

backup_2026-09-14/code/distill/flow_pf.py

@@ -0,0 +1,345 @@
+#!/usr/bin/env python3
+"""逐 cell（per-frame）的 flow-matching 工具，全部建在 DiffSynth 的 FlowMatchScheduler("Wan") 上。
+
+ActionRoPE 的 SFT（actionrope/train.py:218-219）用的就是 `FlowMatchScheduler("Wan")` +
+`set_timesteps(1000, training=True)`，网格是：
+
+    s      = linspace(1, 0, 1001)[:-1]                    # 1000 个点
+    sigma  = 5·s / (1 + 4·s)                              # Wan shift=5
+    t      = sigma · 1000                                 # 与 sigma 一一对应
+    x_t    = (1 − sigma)·x0 + sigma·eps
+    target = eps − x0
+    weight = bsmntw 钟形（set_training_weight）
+
+蒸馏要的东西 DiffSynth 没有：它的 `add_noise` / `training_weight` 只吃**标量** timestep（内部
+argmin 反查下标），没有 `[B,F]` 的逐 cell 版本。这里补上，并且一律**按下标**索引，不再 argmin
+反查（浮点反查在 1000 点网格上是无谓的风险）。
+
+方向陷阱（写反了不会报错，只会悄悄训错噪声档）
+------------------------------------------------
+    sigmas[0]   = 1.0      纯噪声   t = 1000
+    sigmas[999] ≈ 0.00498  近干净   t ≈ 4.98
+即 **越噪 == sigma 越大 == 下标越小**。WorldMind 的 `--tf_noise_aug_max_idx` 就是把这个搞反了
+（近干净档应该是 `randint(N-k, N)`，不是 `randint(0, k)`）。
+
+模块级 scheduler 的不变量（对抗评审 [MINOR]，必须遵守）
+--------------------------------------------------------
+DiffSynth 的 `set_timesteps` 是**破坏性**的：它就地覆写 `self.sigmas` / `self.timesteps`，而且
+只在 `training=True` 时重建 `linear_timesteps_weights`（flow_match.py:365-375）。WorldMind 自己
+那份写进独立的 `infer_sigmas` / `infer_timesteps`，不碰训练网格 —— 移植时**不能**照抄那个用法。
+所以这里：
+
+  * 模块级 `_SCHED` 只在 import 时 `set_timesteps(1000, training=True)` 一次；
+  * 之后 `_SCHED.set_timesteps` 被实例级替换成一个直接抛异常的桩，物理上堵死第二次调用；
+  * `cd_axis()` / `deploy_schedule()` 这类需要别的点数的，一律**新建一个临时 scheduler**；
+  * 每个 public 函数开头都跑 `_assert_sched()`：断言 `len(timesteps)==1000`、
+    `linear_timesteps_weights is not None`，并且三个张量与 import 时的快照逐位相等。
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
+
+import torch
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+from diffsynth.diffusion.flow_match import FlowMatchScheduler  # noqa: E402
+
+__all__ = [
+    "NUM_TRAIN_TIMESTEPS", "WAN_SHIFT", "DMD_TIMESTEP_CLAMP", "DEPLOY_DENOISE_LIST",
+    "sched", "sigmas", "timesteps", "weights",
+    "sigma_of_idx", "t_of_idx", "idx_of_t", "sigma_of_t",
+    "sample_pf_timestep_ids", "add_noise_pf", "training_target", "training_weight_pf",
+    "cd_axis", "deploy_schedule", "assert_deploy_covered",
+    "ti2v_noise", "assert_ti2v",
+]
+
+NUM_TRAIN_TIMESTEPS = 1000
+WAN_SHIFT = 5.0
+# stage3 DMD 抽 t 的下标区间（WorldMind dmd_gm.py:47-50 的 int(0.02·n)/int(0.98·n)，n=1000）。
+# 是**下标**区间不是 t 值区间：idx∈[20,980) ⇒ sigma∈(0.0969, 0.9959] ⇒ t∈(96.9, 995.9]。
+DMD_TIMESTEP_CLAMP = (20, 980)
+# 用户拍板的部署采样表：4 步，原始 t 值（不是下标）。stage2 的 CD 轴和 stage3 都按这个训。
+# 用显式 t 值而不是 set_timesteps(4)：后者在 shift=5 下给出 [1000, 833, 625, 357]，最后一步在
+# sigma=0.357 处读 x0，读出来是条件均值（软/糊）。见 WorldMind dmd_gm.py:461-467 的同一坑。
+DEPLOY_DENOISE_LIST = (1000.0, 750.0, 500.0, 250.0)
+
+
+# ----------------------------------------------------------------------------
+# 模块级 scheduler —— 建一次，之后只读
+# ----------------------------------------------------------------------------
+
+def _forbidden_set_timesteps(*a, **kw):
+    raise RuntimeError(
+        "distill.flow_pf: 模块级 scheduler 是 1000 点训练网格，set_timesteps 会就地覆写它。"
+        "需要别的点数请用 cd_axis(N) / deploy_schedule()（它们新建临时实例）。"
+    )
+
+
+_SCHED = FlowMatchScheduler("Wan")
+_SCHED.set_timesteps(NUM_TRAIN_TIMESTEPS, training=True)
+_SCHED.set_timesteps = _forbidden_set_timesteps  # 实例级遮蔽类方法，堵死第二次调用
+
+# import 时的快照：_assert_sched 拿它逐位比对，任何外部改写都当场炸而不是悄悄算错
+_SNAP_SIGMAS = _SCHED.sigmas.detach().clone()
+_SNAP_TIMESTEPS = _SCHED.timesteps.detach().clone()
+_SNAP_WEIGHTS = _SCHED.linear_timesteps_weights.detach().clone()
+
+
+def _assert_sched() -> FlowMatchScheduler:
+    """每个 public 函数开头都调：确认 1000 点训练网格和钟形权重都还在、且没被人改过。"""
+    assert getattr(_SCHED, "timesteps", None) is not None, "flow_pf: scheduler 未初始化"
+    assert len(_SCHED.timesteps) == NUM_TRAIN_TIMESTEPS, \
+        f"flow_pf: 训练网格被覆写了，len(timesteps)={len(_SCHED.timesteps)} != {NUM_TRAIN_TIMESTEPS}"
+    assert len(_SCHED.sigmas) == NUM_TRAIN_TIMESTEPS, \
+        f"flow_pf: 训练网格被覆写了，len(sigmas)={len(_SCHED.sigmas)} != {NUM_TRAIN_TIMESTEPS}"
+    assert getattr(_SCHED, "linear_timesteps_weights", None) is not None, \
+        "flow_pf: linear_timesteps_weights 丢了（set_timesteps(training=False) 不重建它）"
+    assert torch.equal(_SCHED.sigmas, _SNAP_SIGMAS) and torch.equal(_SCHED.timesteps, _SNAP_TIMESTEPS) \
+        and torch.equal(_SCHED.linear_timesteps_weights, _SNAP_WEIGHTS), \
+        "flow_pf: 模块级 scheduler 的张量被就地改写过"
+    return _SCHED
+
+
+def sched() -> FlowMatchScheduler:
+    """取模块级（唯一的）训练 scheduler，取之前先验不变量。"""
+    return _assert_sched()
+
+
+def sigmas(device=None, dtype=torch.float32) -> torch.Tensor:
+    """1000 点 sigma 网格 [1000]，降序：sigmas[0]=1.0（纯噪声）→ sigmas[999]≈0.00498（近干净）。"""
+    return _assert_sched().sigmas.to(device=device, dtype=dtype)
+
+
+def timesteps(device=None, dtype=torch.float32) -> torch.Tensor:
+    """1000 点 t 网格 [1000] = sigma·1000，降序：t[0]=1000 → t[999]≈4.98。"""
+    return _assert_sched().timesteps.to(device=device, dtype=dtype)
+
+
+def weights(device=None, dtype=torch.float32) -> torch.Tensor:
+    """bsmntw 钟形训练权重 [1000]（DiffSynth set_training_weight，len==1000 时不走经验缩放分支）。"""
+    return _assert_sched().linear_timesteps_weights.to(device=device, dtype=dtype)
+
+
+# ----------------------------------------------------------------------------
+# 下标 ↔ sigma ↔ t
+# ----------------------------------------------------------------------------
+
+def _as_long(idx, device=None) -> torch.Tensor:
+    if not torch.is_tensor(idx):
+        idx = torch.as_tensor(idx)
+    if idx.dtype not in (torch.int32, torch.int64):
+        assert not idx.is_floating_point(), "flow_pf: 下标必须是整型（浮点 t 值请走 idx_of_t）"
+        idx = idx.long()
+    idx = idx.long()
+    return idx if device is None else idx.to(device)
+
+
+def _check_idx(idx: torch.Tensor) -> torch.Tensor:
+    lo, hi = int(idx.min()), int(idx.max())
+    assert 0 <= lo and hi < NUM_TRAIN_TIMESTEPS, \
+        f"flow_pf: 下标越界 [{lo},{hi}]，合法区间 [0,{NUM_TRAIN_TIMESTEPS - 1}]"
+    return idx
+
+
+def _gather(grid: torch.Tensor, idx, device, dtype) -> torch.Tensor:
+    idx = _check_idx(_as_long(idx))
+    out = grid.to(device=idx.device, dtype=torch.float32)[idx]
+    if device is not None:
+        out = out.to(device)
+    return out.to(dtype)
+
+
+def sigma_of_idx(idx, device=None, dtype=torch.float32) -> torch.Tensor:
+    """下标 → sigma，形状原样保留（标量 int 给回 0 维张量）；device=None 表示跟 idx 走。"""
+    return _gather(_assert_sched().sigmas, idx, device, dtype)
+
+
+def t_of_idx(idx, device=None, dtype=torch.float32) -> torch.Tensor:
+    """下标 → t 值（= sigma·1000），形状原样保留。stage1 的 ts_pf 就是 t_of_idx(idx_pf)。"""
+    return _gather(_assert_sched().timesteps, idx, device, dtype)
+
+
+def idx_of_t(t) -> torch.Tensor:
+    """原始 t 值 → 1000 点网格上最近的下标（WorldMind dmd_gm.py:104 的 `_sigma_of` 反查）。
+
+    形状原样保留；标量 float 给回 0 维 long 张量。部署表 [1000,750,500,250] 就是靠这个落到
+    sigma [1.0, 0.75, 0.5, 0.25]（750 在网格上是精确点，500/250 差 <5e-4）。"""
+    s = _assert_sched()
+    if not torch.is_tensor(t):
+        t = torch.as_tensor(t, dtype=torch.float32)
+    t = t.to(torch.float32)
+    grid = s.timesteps.to(device=t.device, dtype=torch.float32)
+    return (grid - t.unsqueeze(-1)).abs().argmin(dim=-1)
+
+
+def sigma_of_t(t, device=None, dtype=torch.float32) -> torch.Tensor:
+    """原始 t 值 → sigma（先 idx_of_t 再查表）。stage3 rollout 的 `_sigma_of`。"""
+    return sigma_of_idx(idx_of_t(t), device=device, dtype=dtype)
+
+
+# ----------------------------------------------------------------------------
+# 逐 cell 加噪（Diffusion Forcing / stage1 teacher forcing）
+# ----------------------------------------------------------------------------
+
+def sample_pf_timestep_ids(B: int, F: int, dev, gen: torch.Generator | None = None) -> torch.Tensor:
+    """[B,F] 独立同分布的逐 cell 噪声档下标，**严格均匀** randint(0,1000)。
+
+    不要换成 `round((1-u)·999)` 那种写法：它把下标 0 和 999 的概率各砍一半。"""
+    _assert_sched()
+    assert B > 0 and F > 0
+    g_dev = dev if gen is None else gen.device
+    idx = torch.randint(0, NUM_TRAIN_TIMESTEPS, (B, F), device=g_dev, generator=gen)
+    return idx.to(dev)
+
+
+def add_noise_pf(x0: torch.Tensor, noise: torch.Tensor, idx_pf: torch.Tensor) -> torch.Tensor:
+    """逐 cell 加噪：x_t = (1−sigma)·x0 + sigma·eps，sigma 按 [B,1,F,1,1] 广播。
+
+    x0 / noise: [B,C,F,H,W]；idx_pf: [B,F] long。返回 x0.dtype。
+
+    为什么系数在 fp32 里算完再 cast：sigma 接近 1 时 `1 − bf16(0.99593) = 0.00391`，而真值是
+    0.00407 —— bf16 直接相减是灾难性抵消，4% 的相对误差直接打在 x0 的系数上。fp32 里先算
+    `1−sigma` 再 cast 成 bf16 是精确的（bf16 表示小数本身没问题），这也正是 DiffSynth 标量
+    `add_noise` 的实际行为（0 维 fp32 标量参与类型提升，先算差再降位）。"""
+    _assert_sched()
+    assert x0.dim() == 5, f"flow_pf: x0 需要 [B,C,F,H,W]，得到 {tuple(x0.shape)}"
+    assert noise.shape == x0.shape, f"flow_pf: noise {tuple(noise.shape)} != x0 {tuple(x0.shape)}"
+    assert noise.dtype == x0.dtype, f"flow_pf: noise dtype {noise.dtype} != x0 {x0.dtype}"
+    idx_pf = _check_idx(_as_long(idx_pf, x0.device))
+    assert idx_pf.dim() == 2 and idx_pf.shape[0] == x0.shape[0] and idx_pf.shape[1] == x0.shape[2], \
+        f"flow_pf: idx_pf 需要 [B,F]=[{x0.shape[0]},{x0.shape[2]}]，得到 {tuple(idx_pf.shape)}"
+    b, f = idx_pf.shape
+    sig = _SCHED.sigmas.to(device=x0.device, dtype=torch.float32)[idx_pf].view(b, 1, f, 1, 1)
+    return (1.0 - sig).to(x0.dtype) * x0 + sig.to(x0.dtype) * noise
+
+
+def training_target(x0: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+    """flow-matching 目标：eps − x0（与噪声档无关，DiffSynth training_target 同）。"""
+    assert noise.shape == x0.shape
+    return noise - x0
+
+
+def training_weight_pf(idx_pf: torch.Tensor) -> torch.Tensor:
+    """逐 cell 的 bsmntw 钟形权重，形状与 idx_pf 相同，fp32。中段 t≈500 最重、两端最轻。"""
+    _assert_sched()
+    idx = _check_idx(_as_long(idx_pf))
+    w = _SCHED.linear_timesteps_weights.to(device=idx.device, dtype=torch.float32)
+    return w[idx]
+
+
+# ----------------------------------------------------------------------------
+# stage2 的一致性轴 / 部署采样表（都用**临时** scheduler，绝不碰模块级那个）
+# ----------------------------------------------------------------------------
+
+def cd_axis(N: int = 48, device=None, check_deploy: bool = True):
+    """stage2 CD 的离散一致性轴：`FlowMatchScheduler("Wan").set_timesteps(N)` 的 (sigmas, timesteps)。
+
+    N=48 时 sigmas = [1.0, 0.99576, …, 0.17857, 0.09615]，t = sigma·1000，都是**降序**（下标越大越干净）。
+    轴永远到不了 sigma=0（最低一对是 0.17857 → 0.09615），所以部署采样表必须结束在低噪端。
+
+    临时实例：DiffSynth 的 set_timesteps 就地覆写 sigmas/timesteps，在模块级 scheduler 上调一次
+    就把 1000 点训练网格换成 N 点的，之后每次 idx 抽样要么读错 sigma 要么直接越界（对抗评审 [MINOR]）。
+
+    check_deploy：顺带断言这条轴的 sigma 区间覆盖 DEPLOY_DENOISE_LIST 的所有读点。"""
+    _assert_sched()
+    assert N >= 2, f"flow_pf: cd_axis 需要 N>=2，得到 {N}"
+    tmp = FlowMatchScheduler("Wan")           # 一次性实例，出了这个函数就被回收
+    tmp.set_timesteps(N)                      # training=False：轴上不需要钟形权重
+    sig = tmp.sigmas.detach().clone().to(device=device, dtype=torch.float32)
+    tst = tmp.timesteps.detach().clone().to(device=device, dtype=torch.float32)
+    assert sig.shape == (N,) and tst.shape == (N,), f"flow_pf: cd_axis 点数不对 {tuple(sig.shape)}"
+    assert bool((sig[:-1] > sig[1:]).all()), "flow_pf: cd_axis 的 sigma 必须严格降序"
+    if check_deploy:
+        assert_deploy_covered(sig)
+    _assert_sched()                           # 再验一次：临时实例没污染模块级的
+    return sig, tst
+
+
+def deploy_schedule(denoise_list=None, device=None):
+    """部署采样表（原始 t 值）→ 落到 1000 点训练网格上的 (idx[K], sigma[K], t_grid[K])。
+
+    默认就是用户拍板的 4 步 `DEPLOY_DENOISE_LIST = (1000, 750, 500, 250)`。
+    返回的 t_grid 是网格上的实际值（1000 / 750 / 500.0? / 250.0? 见下），采样时应当用它而不是
+    原始输入值，这样 sigma 与 t 严格一一对应。"""
+    _assert_sched()
+    ds = tuple(DEPLOY_DENOISE_LIST if denoise_list is None else denoise_list)
+    assert len(ds) >= 1, "flow_pf: 部署表不能为空"
+    t_req = torch.as_tensor([float(x) for x in ds], dtype=torch.float32)
+    if len(ds) > 1:
+        assert bool((t_req[:-1] > t_req[1:]).all()), \
+            f"flow_pf: 部署表必须严格降序（从噪到净），得到 {ds}"
+    idx = idx_of_t(t_req)
+    sig = sigma_of_idx(idx)
+    t_grid = t_of_idx(idx)
+    # 落点误差不该超过半个格。注意网格在 t 轴上**不等距**（sigma 等距、t=sigma·1000 也等距在
+    # sigma 上，但相邻 sigma 之间的间隔是 5/(1+4s)²·0.001，从噪端 0.2 涨到净端 ~4.98），
+    # 所以阈值取全局最大间隔的一半：请求 t=250 落到 251.6 是正常的，请求 t=1.0（低于网格下界
+    # 4.98）才会被拦下。
+    gap = float((_SCHED.timesteps[:-1] - _SCHED.timesteps[1:]).max())
+    assert bool((t_grid - t_req).abs().max() <= 0.5 * gap + 1e-3), \
+        f"flow_pf: 部署 t 值 {ds} 落不到网格上（最近点 {t_grid.tolist()}，半格 {0.5 * gap:.3f}）"
+    if device is not None:
+        idx, sig, t_grid = idx.to(device), sig.to(device), t_grid.to(device)
+    return idx, sig, t_grid
+
+
+def assert_deploy_covered(sig_axis: torch.Tensor, denoise_list=None) -> None:
+    """断言 CD 一致性轴的 sigma 区间覆盖部署表的每个读点。
+
+    覆盖不到意味着：学生在部署要读 x0 的那个噪声档上，从来没被一致性目标约束过。
+    N=48 轴是 [0.09615, 1.0]，部署表最低读点 sigma=0.25 —— 覆盖住。"""
+    _, sig_dep, _ = deploy_schedule(denoise_list)
+    lo, hi = float(sig_axis.min()), float(sig_axis.max())
+    d_lo, d_hi = float(sig_dep.min()), float(sig_dep.max())
+    assert d_lo >= lo - 1e-6, \
+        f"flow_pf: 部署表最低 sigma {d_lo:.5f} 低于 CD 轴下界 {lo:.5f}，该档没被一致性约束过"
+    assert d_hi <= hi + 1e-6, \
+        f"flow_pf: 部署表最高 sigma {d_hi:.5f} 高于 CD 轴上界 {hi:.5f}"
+
+
+# ----------------------------------------------------------------------------
+# stage3 DMD 的原生 TI2V 加噪
+# ----------------------------------------------------------------------------
+
+def ti2v_noise(clip: torch.Tensor, clamp=DMD_TIMESTEP_CLAMP, gen: torch.Generator | None = None):
+    """原生 TI2V 加噪（= ActionRoPE SFT / infer 的条件形式）：cell 0 干净且 t=0，cell 1..F−1
+    共享**一个**抽出来的 t。WorldMind dmd_gm.py:53-66 的移植。
+
+    clip: [B,C,F,H,W]（可以带梯度：cell 0 用 torch.where 写回，不做原地赋值）
+    clamp: **下标**区间 [lo, hi)，默认 (20, 980) —— 即 WorldMind 的 [0.02, 0.98]·1000。
+    返回 (noisy, ts_pf[B,F] fp32, sigma[B] fp32, noise)；noise 与 clip 同 dtype/shape。"""
+    _assert_sched()
+    assert clip.dim() == 5, f"flow_pf: clip 需要 [B,C,F,H,W]，得到 {tuple(clip.shape)}"
+    b, _, f = clip.shape[:3]
+    lo, hi = int(clamp[0]), int(clamp[1])
+    assert 0 <= lo < hi <= NUM_TRAIN_TIMESTEPS, f"flow_pf: clamp 下标区间不合法 {clamp}"
+    dev = clip.device
+    g_dev = dev if gen is None else gen.device
+    idx = torch.randint(lo, hi, (b,), device=g_dev, generator=gen).to(dev)
+    noise = torch.randn(clip.shape, device=g_dev, dtype=clip.dtype, generator=gen).to(dev)
+    sig = _SCHED.sigmas.to(device=dev, dtype=torch.float32)[idx]                 # [B] fp32
+    s5 = sig.view(b, 1, 1, 1, 1)
+    noisy = (1.0 - s5).to(clip.dtype) * clip + s5.to(clip.dtype) * noise
+    ts_pf = _SCHED.timesteps.to(device=dev, dtype=torch.float32)[idx][:, None].expand(b, f).clone()
+    ts_pf[:, 0] = 0.0
+    f0 = (torch.arange(f, device=dev)[None, :] == 0).view(1, 1, f, 1, 1)
+    noisy = torch.where(f0, clip, noisy)                                        # cell 0 = 干净 GT
+    return noisy, ts_pf, sig, noise
+
+
+def assert_ti2v(noisy: torch.Tensor, clip: torch.Tensor, ts_pf: torch.Tensor) -> None:
+    """teacher / critic 每次前向前的 fail-fast 门（WorldMind dmd_gm.py:31-42 的输入部分）：
+    cell 0 必须干净且 t=0，cell 1.. 必须共享同一个 t。"""
+    assert noisy.shape == clip.shape and noisy.dim() == 5
+    assert ts_pf.shape == (clip.shape[0], clip.shape[2]), \
+        f"[gate] ts_pf 需要 [B,F]={(clip.shape[0], clip.shape[2])}，得到 {tuple(ts_pf.shape)}"
+    assert bool((ts_pf[:, 0] == 0).all()), "[gate] cell 0 必须是干净条件帧（t=0）"
+    assert torch.equal(noisy[:, :, 0], clip[:, :, 0]), "[gate] cell 0 的 latent 必须原样是干净 GT"
+    if ts_pf.shape[1] > 1:
+        assert torch.allclose(ts_pf[:, 1:], ts_pf[:, 1:2].expand_as(ts_pf[:, 1:])), \
+            "[gate] cell 1..F−1 必须共享同一个 t（原生 TI2V）"

backup_2026-09-14/code/distill/flowmap.py

@@ -0,0 +1,934 @@
+#!/usr/bin/env python3
+"""AnyFlow 式 **flow map**（any-step）蒸馏的调度器件 + 目标函数，全部建在 `distill/flow_pf.py` 的
+1000 点 Wan 训练网格上，**不引入 diffusers**。
+
+本模块只负责数学；训练循环（run 目录 / ZeRO-2 / EMA / 日志 / 保存）在
+`distill/stage2_flowmap/flowmap_arope.py`，采样器在 `distill/eval/rollout_eval.py`。
+
+一句话的算法
+------------
+把速度场 `v(x_t, t)` 换成**平均速度场** `u(x_t, t, r)`，它是 `[r, t]` 区间上的平均速度：
+
+    x_r = x_t − (σ_t − σ_r)·u(x_t, t, r)                （一步就能从 t 跳到任意 r）
+
+训练目标（AnyFlow `trainer_far_wan_anyflow_pretrain.py:232-390`）：
+
+    target = (ε − x0) − (t − r)·dF/dt                    （t, r 都是 0..1000 单位）
+    dF/dt  ≈ [u(x_t + v·ε_h/1000, t+ε_h, r) − u(x_t − v·ε_h/1000, t−ε_h, r)] / (2·ε_h·den)
+
+`r == t` 时修正项恒为 0，目标退化成普通 rectified flow 的 `ε − x0` —— 这就是 50% 的
+diffusion 槽为什么是"免费的锚"。
+
+上游三个来源
+------------
+  * `vendor/AnyFlow/far/schedulers/scheduling_flowmap_euler_discrete.py:44-112`
+        set_train_weight('beta08') / set_timesteps / step
+  * `vendor/AnyFlow/far/trainers/trainer_far_wan_anyflow_pretrain.py:232-390`
+        sample_timestep / compute_central_difference / train_causal
+  * `vendor/AnyFlow/options/train/anyflow/farwan_causal/pretrain/…4k_b16.yml:53-64`
+        shift 5.0 / weight_type beta08 / diffusion_ratio 0.5 / consistency_ratio 0.25 / epsilon 5
+
+★★★ flow-map 约定与 flow_pf 的差异（写错都**不报错**，只是悄悄训另一根轴）★★★
+================================================================================
+1. **网格公式相同，点数含义不同。**
+   flow_pf 的模块级网格是 `FlowMatchScheduler("Wan").set_timesteps(1000, training=True)`，即
+   `σ = apply_shift(linspace(1,0,1001)[:-1])`，shift=5。上游 `FlowMapDiscreteScheduler.set_timesteps(n)`
+   是**同一个公式** `apply_shift(linspace(1,0,n+1)[:-1])`，只是 n 取推理步数。所以：
+     - 训练轴（抽 t）用 flow_pf 的 1000 点网格 —— 与上游 n=1000 数学同一（fp32 vs fp64，
+       max|Δ| = 1.2e-7，见 `test_flowmap_objective.py::test_grid_matches_upstream_formula`）；
+     - 推理轴（N=4/24）用 `anystep_sigmas(N)`，它内部新建**临时** FlowMatchScheduler。
+   **绝不要**把 `linspace(1,0,n+1)[:-1]` 直接当成 σ：少了 apply_shift 就是 shift=1 的另一根轴。
+
+2. **`anystep_sigmas` 末尾补 0.0，`flow_pf.cd_axis` 不补。**
+   上游 `set_timesteps` 出来的轴**没有终点 0**（N=4 是 [1, .9375, .8333, .625]）。flow map 的
+   Euler 链最后一步必须走到 σ_r = 0 才落在干净估计上；不补 0 就停在 σ=0.625，读出来是糊的
+   条件均值（与 flow_pf:66-68 记的 DEPLOY_DENOISE_LIST 那个坑同源）。
+   反过来 CD 轴**故意**不补 0（一致性目标在低噪端读 x0），所以这两个函数返回的东西长度差 1、
+   语义不同，别互相代入。
+
+3. **权重曲线不同。** flow_pf.weights() 是 **bsmntw**（高斯钟形，峰 2.076 @ σ≈0.5005）；
+   flow map 默认 **beta08** = σ·(1−σ)^0.5 归一化到 sum=1000（峰 1.408 @ σ=0.6670）。
+   两条曲线的峰位置差 0.17 个 σ。默认**不要**用 `flow_pf.training_weight_pf`。
+
+4. **r 是原始 float，不是下标。** consistency 槽要的是**精确的 0.0**，而 1000 点网格
+   `linspace(1,0,1001)[:-1]` 按构造不含 0（最低点 σ=0.00498, t=4.98）。
+   `flow_pf.idx_of_t(0.0)` 会返回 **999**（σ=0.00498）—— 把它当 r 用就把"一致性槽"悄悄变成
+   "近一致性槽"，(t−r) 少了 4.98/1000 的量。所以本模块的 r 一律走原始值，σ_r = r/1000。
+   注意 `t/1000` 与 `flow_pf.sigmas()` 在 fp32 下**不逐位相等**（max|Δ| = 6e-8）：在网格上一律
+   用 `flow_pf.sigma_of_idx(idx)`，只有 off-grid 的 r 才用 `sigma_of_t_exact`。
+
+5. **方向。** flow_pf 的 σ 是**降序**（下标 0 = σ 1.0 = 最噪）。上游 `t = max(t1,t2)`、
+   `r = min(t1,t2)` 是在**值**上取的 ⇒ 在**下标**上是 `i_t = min`、`i_r = max`。反了的话
+   (t−r) 变号，修正项整条反向，loss 照样下降。
+
+6. **单位。** ε = 5.0 是 **0..1000 的 t 单位**：latent 扰动除以 1000（换算成 σ 单位），
+   时间扰动不除，(t−r) 也在 1000 单位，三者精确相消。把 latent 扰动写成 `v·ε` 会让 dF/dt
+   偏大 ~1000×，训练指标全程健康而流形是错的（见 `test_flowmap_objective.py` 的注入实验 I-3）。
+
+7. **逐 cell vs 逐样本 t。** `flow_pf.sample_pf_timestep_ids` 给的是 [B,F] 逐 cell 独立
+   （stage1 的 diffusion forcing）。flow map 默认是**一个标量 (t,r) 广播到 21 个 cell** +
+   cell 0 钉成干净 GT —— 这才对应部署路径（每个 cell 对着已提交的干净历史去噪）。
+
+8. **没有 `assert_deploy_covered`。** CD 必须让一致性轴覆盖部署读点（flow_pf:291-302）；
+   flow map 的训练轴与部署表**解耦**，这正是换掉 stage2 的理由。`anystep_sigmas` 不做这个断言。
+
+模块级不变量
+------------
+本模块**只读** flow_pf 的访问器，从不调用模块级 scheduler 的 set_timesteps（flow_pf 已经把它
+换成抛异常的桩）。需要 N 点轴时新建临时 `FlowMatchScheduler("Wan")`，与 `flow_pf.cd_axis` 同一个惯用法。
+"""
+from __future__ import annotations
+
+import contextlib
+import os
+import sys
+from dataclasses import dataclass, field
+
+os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
+
+import torch
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from diffsynth.diffusion.flow_match import FlowMatchScheduler  # noqa: E402
+
+from distill import flow_pf  # noqa: E402
+
+__all__ = [
+    # 常量
+    "NUM_TRAIN_TIMESTEPS", "T_MAX", "T_MIN", "FLOWMAP_EPSILON",
+    "DIFFUSION_RATIO", "CONSISTENCY_RATIO", "FUSE_GUIDANCE_SCALE",
+    "SPLIT_SEED_BASE", "NULL_AXES", "ROLE_NAMES",
+    "ROLE_DIFFUSION", "ROLE_CONSISTENCY", "ROLE_RANDOM",
+    # 网格 / 权重
+    "sigma_of_t_exact", "t_of_sigma_exact",
+    "beta08_weights", "flowmap_weights", "weight_of_idx", "weight_of_t",
+    # any-step 轴
+    "anystep_shift", "anystep_sigmas", "anystep_timesteps", "deploy_flowmap_sigmas", "assert_anystep_axis",
+    # (t, r) 抽样与跨 rank 划分
+    "TRSample", "split_counts", "make_split_generator", "advance_split_generator",
+    "split_slots", "roles_of_slots", "sample_tr",
+    # 加噪
+    "keep_frame0", "add_noise_scalar",
+    # 目标
+    "fuse_denominator", "active_gains", "fuse_predictions",
+    "probes_needed", "central_difference", "flowmap_target",
+    "flowmap_loss", "adaptive_balance",
+    # guidance 钩子
+    "null_axis", "NULL_AXIS_TABLE",
+    # 采样
+    "euler_step", "euler_chain",
+]
+
+# ----------------------------------------------------------------------------
+# 常量（= 上游 pretrain 配方；三处副本规则见 RUNBOOK，这里是 python 侧唯一一份）
+# ----------------------------------------------------------------------------
+
+NUM_TRAIN_TIMESTEPS = flow_pf.NUM_TRAIN_TIMESTEPS       # 1000
+T_MAX = 1000.0                  # σ=1，每条采样轨迹的起点；t 不许超过它（σ>1 离分布）
+T_MIN = 0.0                     # σ=0，干净；cell 0 与 TF clean 半流用的就是 t=0，在分布内
+FLOWMAP_EPSILON = 5.0           # 中心差分步长，**0..1000 的 t 单位**（yml:64 `epsilon: 5`）
+DIFFUSION_RATIO = 0.5           # yml:62
+CONSISTENCY_RATIO = 0.25        # yml:63
+FUSE_GUIDANCE_SCALE = 3.0       # yml:53；g_text = w − 1 = 2.0
+SPLIT_SEED_BASE = 4321          # split 用的 CPU 生成器种子基数（各 rank 必须一致）
+
+NULL_AXES = ("text", "anchor", "action", "hist", "ctx")
+
+ROLE_DIFFUSION, ROLE_CONSISTENCY, ROLE_RANDOM = 0, 1, 2
+ROLE_NAMES = ("diffusion", "consistency", "random")
+
+
+# ----------------------------------------------------------------------------
+# σ ↔ t（网格上一律走 flow_pf；只有 off-grid 的 r 才用下面这两个）
+# ----------------------------------------------------------------------------
+
+def sigma_of_t_exact(t):
+    """t → σ，**按定义** σ = t/1000，不做网格反查。
+
+    Wan 网格上 `timesteps ≡ sigmas · 1000` 是构造式（flow_match.py:43-52），所以这个换算在数学上
+    精确；但 fp32 下 `t/1000` 与 `flow_pf.sigmas()` **不逐位相等**（max|Δ| = 5.96e-8）。
+    所以：**在网格上取 σ 请用 `flow_pf.sigma_of_idx(idx)`**，这个函数只给 off-grid 的 r 用
+    （典型就是 consistency 槽的 r = 0.0 —— 网格里根本没有 0）。
+    """
+    if torch.is_tensor(t):
+        return t.to(torch.float32) / float(NUM_TRAIN_TIMESTEPS)
+    return float(t) / float(NUM_TRAIN_TIMESTEPS)
+
+
+def t_of_sigma_exact(sigma):
+    """σ → t = σ·1000。同上：网格上请用 `flow_pf.t_of_idx`。"""
+    if torch.is_tensor(sigma):
+        return sigma.to(torch.float32) * float(NUM_TRAIN_TIMESTEPS)
+    return float(sigma) * float(NUM_TRAIN_TIMESTEPS)
+
+
+def _assert_t_range(name: str, v, lo: float = T_MIN, hi: float = T_MAX) -> None:
+    """API 边界的硬断言：t / r 必须落在 [0, 1000]。单位搞错（喂 σ 或喂下标）在这里当场炸。"""
+    if torch.is_tensor(v):
+        mn, mx = float(v.min()), float(v.max())
+    else:
+        mn = mx = float(v)
+    assert lo - 1e-4 <= mn and mx <= hi + 1e-4, \
+        (f"flowmap: {name} 越界 [{mn:.4f}, {mx:.4f}]，合法区间 [{lo}, {hi}]（0..1000 的 t 单位）。"
+         f"喂进来的是 σ（0..1）还是下标（0..999）？")
+
+
+# ----------------------------------------------------------------------------
+# 训练权重
+# ----------------------------------------------------------------------------
+
+_W_CACHE: dict = {}
+
+
+def beta08_weights(device=None, dtype=torch.float32) -> torch.Tensor:
+    """上游 `set_train_weight('beta08')`（scheduling_flowmap_euler_discrete.py:53-56）在我们
+    1000 点网格上的复刻：
+
+        t = σ                          （上游是 timesteps/1000，Wan 网格上就等于 σ）
+        y = t^1.0 · (1 − t)^0.5
+        w = y · (1000 / Σy)
+
+    返回 [1000] 降序对齐 `flow_pf.sigmas()`：w[0] = 0.0（σ=1）、w[1] = 0.0517、
+    峰 1.4076 @ σ=0.66698（下标 714）、w[999] = 0.0182。
+
+    ⚠ w[0] 恰好是 **0.0** —— 而 σ=1 是每条采样轨迹的起点。上游连续抽 t 时这是零测度，我们在
+    1000 点网格上以 1/1000 的概率抽中它，那一步梯度全被乘没。这是**继承**的性质（bsmntw 在
+    σ=1 处同样是 0，stage1 就这么训的），要压请用 `flowmap_weights(weight_floor=...)`。
+
+    与上游的**已知且有意**的差异（实测，钉在 test_beta08_matches_upstream_formula 里）：
+    上游在 fp64 里建网格，我们从 flow_pf 的 fp32 网格算 —— 单一真相源优先于最后一位小数。
+    差异集中在 σ→1 那一端的 `(1−σ)^0.5`：σ=0.99979979 处 `1−σ ≈ 2.0e-4` 是灾难性抵消，
+    fp32 网格 1e-7 的绝对误差被放大成 **相对 1.295e-4**；但那格的权重只有 0.0517，所以
+    **绝对**误差是 6.70e-06（下标 1，全网格最大值），相对峰值 1.408 是 5e-6。其余 999 格更小。
+    归一化的 Σ 在 fp64 里做，避免 1000 项累加的误差。
+    """
+    key = ("beta08", str(device), str(dtype))
+    if key not in _W_CACHE:
+        s = flow_pf.sigmas().double()                      # [1000] 降序，只读
+        y = (s ** 1.0) * ((1.0 - s) ** 0.5)
+        w = y * (float(NUM_TRAIN_TIMESTEPS) / y.sum())
+        _W_CACHE[key] = w.to(device=device, dtype=dtype)
+    return _W_CACHE[key]
+
+
+def flowmap_weights(weight_type: str = "beta08", weight_floor: float = 0.0,
+                    device=None, dtype=torch.float32) -> torch.Tensor:
+    """按名字取 [1000] 训练权重。
+
+    beta08          上游 flow map 的默认（本 stage 的默认）
+    bsmntw/gaussian flow_pf.weights()，即 stage1/SFT 用的高斯钟形（消融用）
+    uniform         全 1（上游 'uniform'）
+
+    weight_floor：归一化**之后**做 `clamp_min`，把 σ=1 那一格从 0 抬起来。默认 0.0 = 上游逐位一致；
+    非 0 会让 Σw 略微超过 1000（这是有意的，别再归一化回去，否则中段会被压低）。
+    """
+    wt = str(weight_type).lower()
+    if wt == "beta08":
+        w = beta08_weights(device=device, dtype=dtype)
+    elif wt in ("bsmntw", "gaussian"):
+        # flow_pf.weights() == DiffSynth set_training_weight() == 上游 'gaussian' 分支
+        w = flow_pf.weights(device=device, dtype=dtype)
+    elif wt == "uniform":
+        w = torch.ones(NUM_TRAIN_TIMESTEPS, device=device, dtype=dtype)
+    else:
+        raise ValueError(f"flowmap: 未知 weight_type={weight_type!r}（beta08 / bsmntw / uniform）")
+    floor = float(weight_floor)
+    assert floor >= 0.0, f"flowmap: weight_floor 必须 >=0，得到 {floor}"
+    return w.clamp_min(floor) if floor > 0.0 else w
+
+
+def weight_of_idx(idx, weight_type: str = "beta08", weight_floor: float = 0.0) -> torch.Tensor:
+    """按**下标**取权重（形状原样保留）。上游 `get_train_weight` 是对 t 做 argmin 反查；我们的 t
+    本来就是从下标查出来的，直接索引与反查结果恒等（见 test_weight_lookup_matches_argmin）。"""
+    idx = idx if torch.is_tensor(idx) else torch.as_tensor(idx)
+    assert not idx.is_floating_point(), "flowmap: weight_of_idx 要的是整型下标（浮点 t 请用 weight_of_t）"
+    w = flowmap_weights(weight_type, weight_floor, device=idx.device)
+    return w[idx.long()]
+
+
+def weight_of_t(t, weight_type: str = "beta08", weight_floor: float = 0.0) -> torch.Tensor:
+    """按**原始 t 值**取权重（走 flow_pf.idx_of_t 的 argmin 反查）= 上游 `get_train_weight`。"""
+    _assert_t_range("t", t)
+    t = t if torch.is_tensor(t) else torch.as_tensor(t, dtype=torch.float32)
+    return weight_of_idx(flow_pf.idx_of_t(t).to(t.device), weight_type, weight_floor)
+
+
+# ----------------------------------------------------------------------------
+# any-step 采样轴（每次都新建**临时** scheduler，绝不碰 flow_pf 的模块级那个）
+# ----------------------------------------------------------------------------
+
+def anystep_shift(u):
+    """上游 `FlowMapDiscreteScheduler.apply_shift`（scheduling_flowmap_euler_discrete.py:76-80）：
+
+        σ = shift·u / (1 + (shift−1)·u)          shift = flow_pf.WAN_SHIFT = 5
+
+    u ∈ [0,1] 是**未 shift** 的均匀参数，σ 是真正的噪声水平。这两个东西差别很大：
+    u=0.5 ⇒ σ=0.833。把 u 直接当 σ 用不会报错，只会把噪声档整体压到低噪端。
+    on-policy 侧 generator 抽均匀 u、critic 抽 logit-normal u，都要经过这里。
+    """
+    u = u if torch.is_tensor(u) else torch.as_tensor(u, dtype=torch.float32)
+    sh = flow_pf.WAN_SHIFT
+    return sh * u / (1.0 + (sh - 1.0) * u)
+
+
+def anystep_sigmas(N: int, device=None, dtype=torch.float32, terminal: bool = True) -> torch.Tensor:
+    """N 步 any-step 采样轴的 σ，**长度 N+1、严格降序、末位是 0.0**。
+
+    前 N 个点 = `FlowMatchScheduler("Wan").set_timesteps(N).sigmas`
+              = 上游 `apply_shift(linspace(1,0,N+1)[:-1])`，shift=5。
+    末位补的 0.0 是 flow map 的终点：Euler 链的最后一步 `x −= (σ_{N-1} − 0)·u` 直接落在干净估计上，
+    不需要再单独读一次 `x0 = x − σ·v`。
+
+    N=4  ⇒ [1.0, 0.9375, 0.83333, 0.625, 0.0]
+    N=24 ⇒ [1.0, 0.99138, …, 0.3125, 0.17857, 0.0]
+    这一对（4 与 24）就是 any-step 结论的验收点（上游 yml:100 `num_inference_steps: [4, 24]`）。
+
+    terminal=False 用来和上游 `set_timesteps` 逐位对拍，正常采样**不要**关。
+    """
+    flow_pf.sched()                                   # 先验模块级不变量
+    assert int(N) >= 1, f"flowmap: anystep_sigmas 需要 N>=1，得到 {N}"
+    tmp = FlowMatchScheduler("Wan")                   # 一次性实例，出了函数就回收
+    tmp.set_timesteps(int(N))                         # training=False：推理轴不需要钟形权重
+    sig = tmp.sigmas.detach().clone().to(torch.float32)
+    assert sig.shape == (int(N),), f"flowmap: anystep 轴点数不对 {tuple(sig.shape)}"
+    if terminal:
+        sig = torch.cat([sig, sig.new_zeros(1)])
+    assert_anystep_axis(sig, terminal=terminal)
+    flow_pf.sched()                                   # 再验一次：临时实例没污染模块级的
+    return sig.to(device=device, dtype=dtype)
+
+
+def anystep_timesteps(N: int, device=None, dtype=torch.float32, terminal: bool = True) -> torch.Tensor:
+    """同上，返回 t = σ·1000（喂给模型的就是它）。"""
+    return anystep_sigmas(N, device=device, dtype=dtype, terminal=terminal) * float(NUM_TRAIN_TIMESTEPS)
+
+
+def deploy_flowmap_sigmas(denoise_list=None, device=None, dtype=torch.float32) -> torch.Tensor:
+    """把用户拍板的部署表（默认 `flow_pf.DEPLOY_DENOISE_LIST = (1000, 750, 500, 250)`）落到 1000 点
+    网格，再补一个终点 0.0 ⇒ [1.0, 0.75, 0.50060, 0.25161, 0.0]。
+
+    flow map 对采样表没有任何约束（训练轴与部署轴解耦），所以这里**不**调 `assert_deploy_covered`。
+    """
+    _, sig, _ = flow_pf.deploy_schedule(denoise_list)
+    sig = torch.cat([sig.to(torch.float32), torch.zeros(1)])
+    assert_anystep_axis(sig, terminal=True)
+    return sig.to(device=device, dtype=dtype)
+
+
+def assert_anystep_axis(sig: torch.Tensor, terminal: bool = True) -> None:
+    """采样轴的 fail-fast 门：严格降序、σ ∈ [0,1]、起点 ≤1，terminal 时末位恰为 0.0。"""
+    assert sig.dim() == 1 and sig.numel() >= 2, f"flowmap: 采样轴形状不对 {tuple(sig.shape)}"
+    assert bool((sig[:-1] > sig[1:]).all()), f"flowmap: 采样轴必须严格降序，得到 {sig.tolist()}"
+    assert float(sig.max()) <= 1.0 + 1e-6 and float(sig.min()) >= -1e-9, \
+        f"flowmap: 采样轴 σ 越界 [{float(sig.min())}, {float(sig.max())}]"
+    if terminal:
+        assert float(sig[-1]) == 0.0, \
+            (f"flowmap: 采样轴末位必须是精确的 0.0，得到 {float(sig[-1])} —— "
+             f"不补 0 的话最后一步停在中间噪声档，读出来是糊的条件均值")
+
+
+# ----------------------------------------------------------------------------
+# (t, r) 抽样 + 跨 rank 的 diffusion / consistency / random 划分
+# ----------------------------------------------------------------------------
+
+@dataclass
+class TRSample:
+    """一次 micro batch 里**本 rank** 的 (t, r) 与角色。所有张量都是 [b]。"""
+    idx_t: torch.Tensor          # [b] long，t 在 1000 点网格上的下标（越小越噪）
+    idx_r: torch.Tensor          # [b] long，r 的下标；consistency 槽是 -1（off-grid，见下）
+    t: torch.Tensor              # [b] fp32，0..1000
+    r: torch.Tensor              # [b] fp32，0..1000；consistency 槽恰为 0.0
+    sigma_t: torch.Tensor        # [b] fp32，= flow_pf.sigma_of_idx(idx_t)
+    sigma_r: torch.Tensor        # [b] fp32
+    is_diffusion: torch.Tensor   # [b] bool，r == t
+    role: torch.Tensor           # [b] long，ROLE_DIFFUSION / _CONSISTENCY / _RANDOM
+    slot: torch.Tensor           # [b] long，本样本在全局 batch 里拿到的槽位
+    perm: torch.Tensor = field(default_factory=lambda: torch.empty(0, dtype=torch.long))
+    n_diffusion: int = 0
+    n_consistency: int = 0
+
+    @property
+    def role_names(self) -> list[str]:
+        return [ROLE_NAMES[int(x)] for x in self.role]
+
+
+def split_counts(world_batch: int, diffusion_ratio: float = DIFFUSION_RATIO,
+                 consistency_ratio: float = CONSISTENCY_RATIO) -> tuple[int, int]:
+    """上游 trainer:243-245 的 `round(ratio · total_global_bsz)`。
+
+    G = world_size · batch_size。G=4 / 0.5 / 0.25 ⇒ (2, 1)，剩 1 个 random 槽。
+    注意 python 的 round 是 banker's rounding（round(0.5)==0），与上游同一函数，照抄。
+    """
+    G = int(world_batch)
+    assert G >= 1, f"flowmap: world_batch 必须 >=1，得到 {G}"
+    assert 0.0 <= diffusion_ratio <= 1.0 and 0.0 <= consistency_ratio <= 1.0
+    n_d = int(round(float(diffusion_ratio) * G))
+    n_c = int(round(float(consistency_ratio) * G))
+    assert n_d + n_c <= G, \
+        f"flowmap: diffusion_ratio+consistency_ratio 的槽位 {n_d}+{n_c} 超过全局 batch {G}"
+    assert n_d >= 1, \
+        (f"flowmap: 全局 batch {G} 下 diffusion 槽为 0 —— adaptive_balance 的分母会是空集，"
+         f"loss 直接 NaN。调大 --diffusion_ratio 或加大全局 batch")
+    return n_d, n_c
+
+
+def make_split_generator(seed: int = 0) -> torch.Generator:
+    """划分用的 **CPU** 生成器。所有 rank 必须用同一个种子 —— 它决定每个 micro batch 的槽位轮换，
+    rank 之间不一致就会出现"两个 rank 都以为自己是 diffusion 槽"，全局比例失衡且不报错。
+
+    与逐 rank 抽 (t,r) 的生成器（`set_seed(seed, device_specific=True)` 那条）严格分开。
+    """
+    g = torch.Generator(device="cpu")
+    g.manual_seed(int(SPLIT_SEED_BASE) + int(seed))
+    return g
+
+
+def advance_split_generator(gen: torch.Generator, n_micro: int, world_batch: int) -> torch.Generator:
+    """resume 时把划分流快进 n_micro 次抽取（cd_arope.py:745-747 的 idx_rng 惯用法）。
+
+    ⚠ 快进的次数是 **micro batch 数**（= step · grad_accum），不是优化步数。所以 resume 时
+    **不许改 --grad_accum**：改了以后同一个 step 对应的 micro 计数变了，各 rank 的流会错位。
+    """
+    G = int(world_batch)
+    for _ in range(int(n_micro)):
+        torch.randperm(G, generator=gen)
+    return gen
+
+
+def split_slots(world_batch: int, gen: torch.Generator | None = None,
+                rotate: bool = True) -> torch.Tensor:
+    """返回 [G] 的槽位分配：全局位置 i（= rank·bsz + j）拿到的槽位 `slots[i]`。
+
+    rotate=True（默认，本移植的**有意偏离**）：所有 rank 用同一个 CPU 生成器抽一个
+    `randperm(G)`，每个 micro batch 换一次。全局比例每个 micro 都精确等于上游，但角色在 rank
+    之间轮转。
+
+    rotate=False：`arange(G)`，即上游 `global_start_idx = rank·bsz + j` 的**静态**分配。
+    G=4 时那意味着 rank0/1 永远 diffusion、rank2 永远 consistency、rank3 永远 random，
+    整整 8000 步不变 —— 而且 adaptive_balance 的分母永远是同两个 rank 的 2 样本均值。
+    留着它是为了跑上游逐位差分门禁。
+    """
+    G = int(world_batch)
+    assert G >= 1
+    if not rotate:
+        return torch.arange(G, dtype=torch.long)
+    assert gen is not None, "flowmap: rotate=True 必须给一个各 rank 一致的 CPU 生成器（make_split_generator）"
+    assert gen.device.type == "cpu", "flowmap: 划分生成器必须在 CPU 上（CUDA 生成器各 rank 不同步）"
+    return torch.randperm(G, generator=gen)
+
+
+def roles_of_slots(slots: torch.Tensor, n_diffusion: int, n_consistency: int) -> torch.Tensor:
+    """槽位 → 角色（上游 trainer:248-256 的三段 if/elif/else）。"""
+    slots = slots.long()
+    role = torch.full_like(slots, ROLE_RANDOM)
+    role = torch.where(slots < int(n_diffusion) + int(n_consistency), torch.full_like(slots, ROLE_CONSISTENCY), role)
+    role = torch.where(slots < int(n_diffusion), torch.full_like(slots, ROLE_DIFFUSION), role)
+    return role
+
+
+def sample_tr(batch_size: int, world_size: int = 1, rank: int = 0, *,
+              gen: torch.Generator | None = None,
+              split_gen: torch.Generator | None = None,
+              diffusion_ratio: float = DIFFUSION_RATIO,
+              consistency_ratio: float = CONSISTENCY_RATIO,
+              rotate: bool = True, device="cpu") -> TRSample:
+    """抽本 rank 这一个 micro batch 的 (t, r)。
+
+    抽法（上游 trainer:232-241 在 1000 点网格上的离散化）：
+        i1, i2 ~ randint(0, 1000)                  逐 rank 独立（gen 是 device-specific 的那条）
+        i_t = min(i1, i2)   i_r = max(i1, i2)      σ 降序 ⇒ 下标小 = σ 大 = t 大
+        t = flow_pf.t_of_idx(i_t)                  已经过 shift 变换，(4.98, 1000]
+    角色覆写：
+        diffusion   r := t （精确相等，后面 probes_needed 靠 `==` 判定）
+        consistency r := 0.0（**off-grid**，idx_r 记 −1）
+        random      保留抽出来的一对
+
+    与上游的差异：上游 `t1,t2 ~ U(0,1)` 连续、再 apply_shift；我们在 1000 点网格上均匀抽下标，
+    正是同一个分布的离散化（网格本身就是 `apply_shift(linspace(1,0,1001)[:-1])`）。
+    """
+    b, W = int(batch_size), int(world_size)
+    assert b >= 1 and W >= 1 and 0 <= int(rank) < W
+    G = W * b
+    n_d, n_c = split_counts(G, diffusion_ratio, consistency_ratio)
+
+    perm = split_slots(G, gen=split_gen, rotate=rotate)                 # [G] CPU
+    slot = perm[int(rank) * b: int(rank) * b + b].to(device)            # [b] 本 rank 的槽位
+    role = roles_of_slots(slot, n_d, n_c)
+
+    g_dev = device if gen is None else gen.device
+    pair = torch.randint(0, NUM_TRAIN_TIMESTEPS, (b, 2), device=g_dev, generator=gen).to(device)
+    idx_t = pair.min(dim=1).values.long()                               # 下标小 ⇒ σ 大 ⇒ t 大
+    idx_r = pair.max(dim=1).values.long()
+
+    t = flow_pf.t_of_idx(idx_t, device=device)                          # [b] fp32
+    sig_t = flow_pf.sigma_of_idx(idx_t, device=device)
+    r = flow_pf.t_of_idx(idx_r, device=device)
+    sig_r = flow_pf.sigma_of_idx(idx_r, device=device)
+
+    is_role_d = role == ROLE_DIFFUSION
+    is_d = is_role_d
+    is_c = role == ROLE_CONSISTENCY
+    # diffusion：r := t（`t.clone()` 保证逐位相等，probes_needed 的 `==` 才成立）
+    r = torch.where(is_d, t, r)
+    sig_r = torch.where(is_d, sig_t, sig_r)
+    idx_r = torch.where(is_d, idx_t, idx_r)
+    # consistency：r := 精确 0.0，idx 记 −1（网格上没有 0，见文件头差异 #4）
+    zero = torch.zeros_like(r)
+    r = torch.where(is_c, zero, r)
+    sig_r = torch.where(is_c, zero, sig_r)
+    idx_r = torch.where(is_c, torch.full_like(idx_r, -1), idx_r)
+
+    # ★ is_diffusion 由**值**派生，不由角色派生：random 槽偶然抽到 i1 == i2 时也落在对角线上
+    #   （1000 点网格上 p = 1/1000；上游连续抽 t 是零测度，所以它没遇到过这个 case）。
+    #   语义要的是"修正项是否消失"，而不是"抽签抽到了哪个槽"：这两个用途都跟着值走 ——
+    #   probes_needed 靠它省两次前向，adaptive_balance 靠它决定谁进分母。
+    #   派生只会**增加**对角线样本，不会减少，所以 split_counts 的 n_diffusion>=1 保证仍然成立。
+    #   role 保留为抽样出身的标签，只用于日志（角色直方图与 diffusion mask 可以不等，这是对的）。
+    is_d = (r == t)
+    _assert_t_range("t", t)
+    _assert_t_range("r", r)
+    assert bool((t >= r - 1e-6).all()), f"flowmap: 必须 t >= r，得到 t={t.tolist()} r={r.tolist()}"
+    assert bool(is_d[is_role_d].all()), "flowmap: diffusion 槽必须逐位落在 r==t 对角线上"
+    assert not bool(is_d[is_c].any()), "flowmap: consistency 槽的 r=0 不可能等于 t（网格最低 t=4.98）"
+    return TRSample(idx_t=idx_t, idx_r=idx_r, t=t, r=r, sigma_t=sig_t, sigma_r=sig_r,
+                    is_diffusion=is_d, role=role, slot=slot, perm=perm,
+                    n_diffusion=n_d, n_consistency=n_c)
+
+
+# ----------------------------------------------------------------------------
+# 加噪（标量 t 广播到 21 个 cell + cell 0 钉成干净 GT）
+# ----------------------------------------------------------------------------
+
+def keep_frame0(x: torch.Tensor, clip: torch.Tensor) -> torch.Tensor:
+    """cell 0 换回 clip 的 cell 0（out-of-place，clip 带梯度时也安全）。
+
+    与 `stage2_cd/cd_arope.py:_keep_frame0` 逐字相同；这里再写一份是为了让 flowmap.py 不反向依赖
+    某个 stage 目录。中心差分的两个探针输入**必须**过这个：v·ε/1000 打到 cell 0 上就破坏了
+    TI2V 干净锚点，而整个条件形式都建在那个锚点上。
+    """
+    f0 = (torch.arange(x.shape[2], device=x.device) == 0).view(1, 1, -1, 1, 1)
+    return torch.where(f0, clip, x)
+
+
+def add_noise_scalar(x0: torch.Tensor, noise: torch.Tensor, idx_t: torch.Tensor,
+                     pin_first_cell: bool = True) -> torch.Tensor:
+    """x_t = (1−σ_t)·x0 + σ_t·ε，**一个标量档广播到所有 cell**，再把 cell 0 钉回干净 GT。
+
+    x0 / noise: [b,C,F,H,W]；idx_t: [b] long（每个样本一个网格下标）。
+    走 `flow_pf.add_noise_pf`（fp32 里算系数再降位，避免 σ→1 时 `1−bf16(σ)` 的灾难性抵消）。
+    """
+    assert x0.dim() == 5, f"flowmap: x0 需要 [b,C,F,H,W]，得到 {tuple(x0.shape)}"
+    idx_t = idx_t.long().to(x0.device)
+    assert idx_t.dim() == 1 and idx_t.shape[0] == x0.shape[0], \
+        f"flowmap: idx_t 需要 [b]={x0.shape[0]}，得到 {tuple(idx_t.shape)}"
+    idx_pf = idx_t[:, None].expand(x0.shape[0], x0.shape[2]).contiguous()   # [b,F] 同一个档
+    x_t = flow_pf.add_noise_pf(x0, noise, idx_pf)
+    return keep_frame0(x_t, x0) if pin_first_cell else x_t
+
+
+# ----------------------------------------------------------------------------
+# guidance 蒸馏（多轴 null 的融合）
+# ----------------------------------------------------------------------------
+
+NULL_AXIS_TABLE = {
+    # axis      来源                        干预点
+    "text":   ("stage1b --scene_dropout",  "cross-attn：context ← NULL_PROMPT 行，context_ids ← 0"),
+    "anchor": ("stage1b --anchor_dropout", "TF 掩码：bare.train_drop_anchor=True，干净 cell 0 从每个 noisy cell 视野里摘掉"),
+    "action": ("stage1b --action_dropout", "世界 RoPE + known 掩码：offset_px ← 0，mask_input 按它重建"),
+    "hist":   ("stage1b --hist_dropout",   "因果掩码：train_local_attn_size, train_sink_size = 1, 1（只看得到 cell 0）"),
+    "ctx":    ("stage1b --ctx_dropout",    "因果掩码：train_local_attn_size, train_sink_size = 1, 0（什么都看不到）"),
+}
+
+
+def active_gains(gains: dict, dropped: set | tuple | list | None = None) -> dict:
+    """把**本 micro batch 已经在主前向里被 drop 掉的轴**的 g 归零，并丢掉 g==0 的项。
+
+    这是对上游的一处**有意修正**：上游在 `pixel_loss:307-309` 里以 10% 概率把 text 换成负条件，
+    那时 `cond == uncond`，融合在算术上是恒等 —— 但 `guidance` 仍然是 3.0 并被塞进 `dF_dt` 的
+    分母（trainer:279），于是**恰好在这批样本上**修正项小了 3×。我们按实际生效的轴算 den，
+    顺带省掉那一次冗余前向。
+    """
+    dropped = set(dropped or ())
+    out = {}
+    for k, g in gains.items():
+        assert k in NULL_AXES, f"flowmap: 未知 null 轴 {k!r}，合法值 {NULL_AXES}"
+        g = float(g)
+        assert g >= 0.0, f"flowmap: g_{k} 必须 >=0（负值等于把 null 方向放大，语义未定义），得到 {g}"
+        if k not in dropped and g != 0.0:
+            out[k] = g
+    return out
+
+
+def fuse_denominator(gains: dict) -> float:
+    """den = 1 + Σ_j g_j。text-only、g_text = w−1 = 2 时 den == w == 3（上游的 `guidance`）。"""
+    return 1.0 + sum(float(g) for g in gains.values())
+
+
+def fuse_predictions(f_full: torch.Tensor, nulls: dict, gains: dict,
+                     clip_k: float = 0.0) -> tuple[torch.Tensor, float]:
+    """多轴 CFG 融合：
+
+        pred = (F_full + Σ_j g_j · F_{∅_j}) / (1 + Σ_j g_j)
+
+    text-only、g_text = 2 时，与上游 `(noise_pred − (1 − w)·noise_pred_uncond) / w`（w=3）是
+    **同一个 IEEE 值**（fp32 / fp64 下 `torch.equal` 均为 True）—— 所以多轴代码路径能过逐位差分门禁。
+
+    最优点上这会把 `F_full → v* + Σ_j g_j (v* − v_{∅_j})`，即多轴 CFG 放大后的场，永久烤进权重：
+    推理时**一次前向都不用多做**，代价是强度也再也调不回来了。
+
+    clip_k > 0：把 `|pred − F_full|` 截断到 `clip_k · std(F_full)` —— 打开额外轴时的熔断器
+    （那几条轴是**掩码差**不是 classifier-free 的分数差，放大它是外推；而且参照物是学生自己，
+    是正反馈回路）。默认 0 = 不截断 = 上游行为。
+
+    返回 (pred, den)。所有 F_{∅_j} 都应当是 `torch.no_grad()` 下算出来的。
+    """
+    assert set(nulls.keys()) >= set(gains.keys()), \
+        f"flowmap: gains 里有 {set(gains) - set(nulls)} 但 nulls 没给对应的前向"
+    den = fuse_denominator(gains)
+    assert den > 0.0, f"flowmap: den={den} 必须 >0"
+    pred = f_full
+    for k in NULL_AXES:                                    # 固定顺序 ⇒ 浮点结果可复现
+        if k in gains:
+            g = float(gains[k])
+            fn = nulls[k]
+            assert fn.shape == f_full.shape, f"flowmap: null[{k}] 形状 {tuple(fn.shape)} != {tuple(f_full.shape)}"
+            pred = pred + g * fn
+    pred = pred / den
+    if float(clip_k) > 0.0:
+        with torch.no_grad():
+            lim = float(clip_k) * f_full.float().std()
+        d = (pred - f_full).clamp(min=-lim, max=lim)
+        pred = f_full + d
+    return pred, den
+
+
+@contextlib.contextmanager
+def null_axis(axis: str, cond: dict, *, bare=None,
+              null_context=None, null_context_ids=None, mask_input_fn=None):
+    """产出某一条 null 轴的 cond（并在退出时把裸模块属性原样还回去）。
+
+    用法（trainer 侧）：
+
+        with flowmap.null_axis("anchor", cond, bare=unwrap(model)) as cond_j:
+            with torch.no_grad(), autocast:
+                nulls["anchor"] = bare(x_t, t, r_timestep=r, **cond_j)
+
+    为什么属性存取包在这里而不是抄到 trainer：三条掩码轴改的是**裸模块属性**
+    （`ar_tf.py:200-210` 的惯用法），中途异常没还原的话，之后所有前向都在错误的注意力窗口上跑
+    且不报错。放在一个 contextmanager 里，`finally` 只写一次。
+
+    text  需要 null_context（[b,L,4096] 或 [b,f,L,4096]，NULL_PROMPT 那一行）+ null_context_ids（全 0）
+    action 把 offset_px 归零并按它重建 mask_input（默认惰性 import `distill.data.make_mask_input`）
+    anchor/hist/ctx 只改 bare 的属性，cond 原样
+    """
+    assert axis in NULL_AXES, f"flowmap: 未知 null 轴 {axis!r}，合法值 {NULL_AXES}"
+    saved = None
+    try:
+        if axis == "text":
+            assert null_context is not None, "flowmap: text null 需要 null_context（NULL_PROMPT 的嵌入）"
+            ids = null_context_ids
+            if ids is None and cond.get("context_ids") is not None:
+                ids = torch.zeros_like(cond["context_ids"])
+            yield dict(cond, context=null_context, context_ids=ids)
+            return
+        if axis == "action":
+            off = cond.get("offset_px")
+            assert off is not None, "flowmap: action null 需要 cond['offset_px']"
+            off0 = torch.zeros_like(off)
+            c = dict(cond, offset_px=off0)
+            mi = cond.get("mask_input")
+            if mi is not None:
+                if mask_input_fn is None:
+                    from distill.data import make_mask_input as _mk       # 惰性 import
+                    c["mask_input"] = _mk(off0, int(mi.shape[-2]), int(mi.shape[-1])).to(
+                        device=mi.device, dtype=mi.dtype)
+                else:
+                    c["mask_input"] = mask_input_fn(off0)
+            yield c
+            return
+        # anchor / hist / ctx：掩码干预，只动裸模块属性
+        assert bare is not None, f"flowmap: {axis} null 是注意力掩码干预，必须给 bare=unwrap(model)"
+        saved = (bare.train_local_attn_size, bare.train_sink_size, bare.train_drop_anchor)
+        if axis == "ctx":
+            bare.train_local_attn_size, bare.train_sink_size = 1, 0     # 谁都看不到
+        elif axis == "hist":
+            bare.train_local_attn_size, bare.train_sink_size = 1, 1     # 只看得到 cell 0
+        else:                                                          # anchor
+            bare.train_drop_anchor = True                              # 只摘掉 cell 0
+        yield dict(cond)
+    finally:
+        if saved is not None:
+            bare.train_local_attn_size, bare.train_sink_size, bare.train_drop_anchor = saved
+
+
+# ----------------------------------------------------------------------------
+# 中心差分 dF/dt
+# ----------------------------------------------------------------------------
+
+def probes_needed(t, r) -> bool:
+    """t 与 r **逐位相等**时不用跑探针（dF/dt 恒为 0）。
+
+    不只是省两次前向：`0 × NaN = NaN` 而不是 0，所以这个跳过同时是数值熔断器 ——
+    diffusion 槽上哪怕探针炸了也进不了 loss。判定用 `==`（`sample_tr` 里 diffusion 槽是
+    `torch.where(is_d, t, r)`，逐位相等有保证），不用容差。
+    """
+    if torch.is_tensor(t) or torch.is_tensor(r):
+        tt = t if torch.is_tensor(t) else torch.as_tensor(t)
+        rr = r if torch.is_tensor(r) else torch.as_tensor(r)
+        return bool((tt != rr).any())
+    return float(t) != float(r)
+
+
+def _hi(x: torch.Tensor) -> torch.Tensor:
+    """提升到至少 fp32 做系数/残差运算；已经是 fp64 的**保持 fp64**。
+
+    不要写成 `.float()`：那会把 fp64 降成 fp32，解析门禁（fp64 闭式核对）就只剩 1e-8 的分辨率，
+    量纲 bug 的注入实验会被浮点噪声盖住。生产里 latent 是 bf16，这个函数返回 fp32，行为不变。
+    """
+    return x if x.dtype in (torch.float32, torch.float64) else x.float()
+
+
+def _hi_dtype(x: torch.Tensor) -> torch.dtype:
+    """时间标量（t / r / width）该用的 dtype：fp64 输入配 fp64，其余一律 fp32。"""
+    return torch.float64 if x.dtype == torch.float64 else torch.float32
+
+
+def _bvec(v, b: int, device, dtype=torch.float32) -> torch.Tensor:
+    if torch.is_tensor(v):
+        out = v.to(device=device, dtype=dtype).reshape(-1)
+    else:
+        out = torch.full((b,), float(v), device=device, dtype=dtype)
+    if out.numel() == 1 and b > 1:
+        out = out.expand(b).contiguous()
+    assert out.shape == (b,), f"flowmap: 期望 [b]={b} 的标量向量，得到 {tuple(out.shape)}"
+    return out
+
+
+@torch.no_grad()
+def central_difference(u_fn, x_t: torch.Tensor, x0: torch.Tensor, noise: torch.Tensor,
+                       t, r, *, epsilon: float = FLOWMAP_EPSILON, den: float = 1.0,
+                       eps_clamp: bool = True, t_min: float = T_MIN, t_max: float = T_MAX,
+                       pin_first_cell: bool = True) -> torch.Tensor:
+    """dF/dt 的中心差分（上游 `compute_central_difference`, trainer:260-277）。
+
+    单位（★ 最容易写错的一处，写错了 1000× 也不会报错）：
+        v      = ε − x0                                       速度（σ 单位，dx/dσ）
+        t_p    = min(t + ε_h, 1000)   t_m = max(t − ε_h, 0)   **时间**扰动不除 1000
+        x_p    = keep_frame0(x_t + v·(t_p − t)/1000, x0)      **latent** 扰动除 1000
+        x_m    = keep_frame0(x_t − v·(t − t_m)/1000, x0)
+        dF/dt  = (u_p − u_m) / (width · den)                  width = t_p − t_m
+
+    闭式核对（`u = A·x + b·t/1000 + c·r/1000`，轨迹 x(t) = x0 + (t/1000)·v）：
+        真值 dF/dt = A·v/1000 + b/1000，中心差分恰好给出同一个值（test 里到 8e-18）。
+
+    `den`：上游把探针**只用条件分支**跑（`u_func` 传的是 prompt_embeds，不做融合），再整体除以
+    `guidance`。所以 dF/dt 是**条件场**的导数除以 den，**不是**融合场的导数（那还要再跑 2·|axes|
+    次 null 探针）。这是上游的选择，照抄。
+
+    `eps_clamp=True`（默认，本移植的偏离）：把探针钉在 [t_min, t_max] 内���并用**实际**步长做分母。
+        t=1000（σ=1，每条轨迹的起点，我们有 1/1000 的概率抽到）时 t_p 本来会是 1005 ⇒ σ>1 离分布。
+        不触边时 width 直接取字面量 `2·ε_h`：fp32 下 `(t+5)−(t−5)` 在 1000 个网格点里有 5 个
+        不精确等于 10.0（max|Δ| = 3.05e-5），取字面量既更准也与上游逐位一致。
+    `eps_clamp=False` 复现上游（不夹、width ≡ 2·ε_h），给差分门禁用。
+
+    ⚠ bf16 噪声地板（**实测**，非推测）：latent 扰动是 `v·ε/1000 = 0.005·v`，而 bf16 的相对
+    精度是 2^-8 ≈ 0.0039。在 [4,48,21,30,52] 的随机 latent 上，实际落地的 `(x_p − x_m)` 相对
+    理想值 `2·v·ε/1000` 的**逐元素平均相对误差 29.4%**，其中 **17.4% 的元素扰动被完全舍掉**
+    （范数只差 1.3%，所以是零均值量化噪声不是偏置）。这是从上游继承的性质（它的 latent 也是
+    bf16）。模型输出对 780×48 个元素求和会把它平均掉不少，但上线前值得用真权重量一下
+    `u_p − u_m` 的信噪比（见 anystep_probe.py 的建议）。
+    """
+    assert x_t.dim() == 5 and x0.shape == x_t.shape and noise.shape == x_t.shape
+    b, dev = x_t.shape[0], x_t.device
+    tdt = _hi_dtype(x_t)
+    tv = _bvec(t, b, dev, tdt)
+    rv = _bvec(r, b, dev, tdt)
+    _assert_t_range("t", tv)
+    _assert_t_range("r", rv)
+    eh = float(epsilon)
+    assert eh > 0.0, f"flowmap: epsilon 必须 >0，得到 {eh}"
+    assert float(den) > 0.0, f"flowmap: den 必须 >0，得到 {den}"
+
+    if not probes_needed(tv, rv):
+        # 全是 diffusion 槽：一次前向都不跑
+        return torch.zeros((), device=dev, dtype=torch.float32)
+
+    if eps_clamp:
+        t_p = torch.minimum(tv + eh, torch.full_like(tv, float(t_max)))
+        t_m = torch.maximum(tv - eh, torch.full_like(tv, float(t_min)))
+        clamped = (t_p != tv + eh) | (t_m != tv - eh)
+        width = torch.where(clamped, t_p - t_m, torch.full_like(tv, 2.0 * eh))
+        assert bool((width > 0).all()), \
+            f"flowmap: 中心差分步长退化成 0（t={tv.tolist()}, ε={eh}, 区间 [{t_min},{t_max}]）"
+    else:                                              # 上游：不夹、width ≡ 2ε
+        t_p, t_m = tv + eh, tv - eh
+        width = torch.full_like(tv, 2.0 * eh)
+
+    v = (_hi(noise) - _hi(x0))                         # ε − x0，至少 fp32 里算系数
+    d_p = ((t_p - tv) / float(NUM_TRAIN_TIMESTEPS)).view(b, 1, 1, 1, 1)
+    d_m = ((tv - t_m) / float(NUM_TRAIN_TIMESTEPS)).view(b, 1, 1, 1, 1)
+    x_p = (_hi(x_t) + v * d_p).to(x_t.dtype)
+    x_m = (_hi(x_t) - v * d_m).to(x_t.dtype)
+    if pin_first_cell:
+        x_p, x_m = keep_frame0(x_p, x0), keep_frame0(x_m, x0)
+
+    u_p = u_fn(x_p, t_p, rv)
+    u_m = u_fn(x_m, t_m, rv)
+    assert u_p.shape == x_t.shape and u_m.shape == x_t.shape, \
+        f"flowmap: 探针输出形状 {tuple(u_p.shape)} != {tuple(x_t.shape)}"
+
+    dF = (_hi(u_p) - _hi(u_m)) / (width * float(den)).view(b, 1, 1, 1, 1)
+    # t == r 的样本把 dF 硬置 0：(t−r)=0 已经让它进不了 target，但 0×NaN=NaN，必须显式清掉
+    same = (tv == rv).view(b, 1, 1, 1, 1)
+    return torch.where(same, torch.zeros_like(dF), dF)
+
+
+# ----------------------------------------------------------------------------
+# 目标与 loss
+# ----------------------------------------------------------------------------
+
+def flowmap_target(x0: torch.Tensor, noise: torch.Tensor, t, r,
+                   dF_dt: torch.Tensor | float = 0.0) -> torch.Tensor:
+    """flow-map 目标（上游 trainer:381）：
+
+        target = (ε − x0) − (t − r)·dF/dt          t, r 都是 0..1000 单位
+
+    `r == t` ⇒ 修正项恒为 0 ⇒ 退化成 rectified flow 的 `ε − x0`（= `flow_pf.training_target`）。
+    fp32 里算：dF/dt 在大 (t−r) 上是重尾的，bf16 的 8 位尾数会先吃掉残差再进 MSE。
+    """
+    assert x0.dim() == 5 and noise.shape == x0.shape
+    b, dev = x0.shape[0], x0.device
+    tv = _bvec(t, b, dev, _hi_dtype(x0))
+    rv = _bvec(r, b, dev, _hi_dtype(x0))
+    _assert_t_range("t", tv)
+    _assert_t_range("r", rv)
+    v = _hi(noise) - _hi(x0)
+    if isinstance(dF_dt, (int, float)) and float(dF_dt) == 0.0:
+        return v
+    dt = (tv - rv).view(b, 1, 1, 1, 1)
+    d = dF_dt if torch.is_tensor(dF_dt) else torch.as_tensor(float(dF_dt), device=dev, dtype=v.dtype)
+    return v - dt * _hi(d)
+
+
+def flowmap_loss(pred: torch.Tensor, target: torch.Tensor, *,
+                 idx_t: torch.Tensor | None = None, t=None,
+                 weight_type: str = "beta08", weight_floor: float = 0.0,
+                 skip_first_cell: bool = True) -> tuple[torch.Tensor, dict]:
+    """逐样本的加权 MSE（上游 trainer:383-385）。返回 (loss_b [b], stats)。
+
+        se         = (pred − target)^2                    fp32
+        per_sample = se[:, :, 1:].mean(dim=(1,2,3,4))     cell 0 是干净锚点，不计
+        loss_b     = per_sample · w(σ_t)                  beta08 钟形，逐样本一个标量
+
+    `skip_first_cell` 是**我们的**约定（stage1/stage2-CD 都这么切）；上游把所有 noisy 帧都算进去，
+    因为它的干净上下文是另一段 token，本来就不在这个张量里。
+
+    权重按下标取（`idx_t`）比按值 argmin 反查（`t`）便宜且等价，两者给一个即可。
+    上游 `loss.unsqueeze(-1) * weight` 再 `.mean(-1)`：weight 是 [b,f] 但逐帧同值，等价于逐样本
+    乘一个标量 —— 这里直接写成标量乘法。
+    """
+    assert pred.shape == target.shape and pred.dim() == 5, \
+        f"flowmap: pred {tuple(pred.shape)} vs target {tuple(target.shape)}"
+    se = (_hi(pred) - _hi(target)) ** 2
+    se_l = se[:, :, 1:] if skip_first_cell else se
+    assert se_l.shape[2] >= 1, "flowmap: skip_first_cell=True 时至少要有 2 个 cell"
+    per_sample = se_l.mean(dim=(1, 2, 3, 4))                       # [b]
+    if idx_t is not None:
+        w = weight_of_idx(idx_t.to(pred.device), weight_type, weight_floor).to(se.dtype)
+    elif t is not None:
+        w = weight_of_t(t, weight_type, weight_floor).to(device=pred.device, dtype=se.dtype)
+    else:
+        w = torch.ones_like(per_sample)
+    w = w.reshape(-1)
+    assert w.shape == per_sample.shape, f"flowmap: 权重形状 {tuple(w.shape)} != [b]={tuple(per_sample.shape)}"
+    loss_b = per_sample * w
+    stats = {
+        "mse": se_l.mean().detach(),          # 完全不加权，跨配方可比（= ar_tf_loss 的 "mse"）
+        "loss_uw": per_sample.mean().detach(),
+        "tw_mean": w.mean().detach(),
+    }
+    return loss_b, stats
+
+
+def adaptive_balance(loss_b: torch.Tensor, is_diffusion: torch.Tensor,
+                     gather_fn=None, eps: float = 1e-5) -> tuple[torch.Tensor, dict]:
+    """上游 trainer:386-390 的自适应配平：把每个**非** diffusion 样本的 loss 重整到全局 diffusion 均值。
+
+        gl    = gather(loss_b.detach())          [G]  ← 集合通信，每个 rank 每个 micro 都要跑
+        gm    = gather(is_diffusion)             [G]
+        scale = gl[gm].mean() / (loss_b.detach() + eps)
+        loss_b = where(is_diffusion, loss_b, loss_b · scale)
+
+    代数性质（就是它让 (t−r)·dF/dt 这一项可训）：重整之后**每个非 diffusion 样本的 loss 值恰好
+    等于全局 diffusion 均值**（差一个 eps/(L+eps) 的相对量），梯度方向不变、量级被拉平。
+    没有它，修正项那条比普通 flow matching 重尾得多，直接主导梯度。
+
+    `torch.where` 代替上游的原地 masked 赋值：值相同，图更干净（原地赋值在 autograd 里要多留一份）。
+    `gather_fn=None` ⇒ 单进程（测试 / dry_run）；trainer 传 `accelerator.gather`。
+    ⚠ 全局 diffusion 集合为空时**抛异常**而不是产出 NaN：`split_counts` 已经断言 n_diffusion>=1，
+    这里是第二道。
+    """
+    assert loss_b.dim() == 1 and is_diffusion.shape == loss_b.shape
+    is_d = is_diffusion.bool()
+    det = loss_b.detach()
+    if gather_fn is None:
+        gl, gm = det, is_d
+    else:
+        gl = gather_fn(det)
+        # bool 的集合通信在部分后端上不受支持，统一走 float 再阈值化（值等价）
+        gm = gather_fn(is_d.to(det.dtype)) > 0.5
+    gl, gm = gl.reshape(-1), gm.reshape(-1)
+    assert gl.shape == gm.shape, f"flowmap: gather 后形状不一致 {tuple(gl.shape)} vs {tuple(gm.shape)}"
+    if not bool(gm.any()):
+        raise RuntimeError(
+            "flowmap.adaptive_balance: 全局 diffusion 集合为空 —— 分母会是 nan。"
+            "检查 --diffusion_ratio / 全局 batch（split_counts 要求 n_diffusion>=1），"
+            "以及各 rank 的 split 生成器是不是同一个种子。")
+    gm_mean = gl[gm].mean()
+    scale = gm_mean / (det + float(eps))
+    out = torch.where(is_d, loss_b, loss_b * scale)
+    info = {
+        "gm_mean": gm_mean.detach(),
+        "n_diffusion_global": int(gm.sum()),
+        "n_global": int(gm.numel()),
+        "scale_mean": scale[~is_d].mean().detach() if bool((~is_d).any()) else torch.zeros((), device=loss_b.device),
+    }
+    return out, info
+
+
+# ----------------------------------------------------------------------------
+# 采样：平均速度场上的确定性 Euler
+# ----------------------------------------------------------------------------
+
+def euler_step(x: torch.Tensor, u: torch.Tensor, sigma_t, sigma_r) -> torch.Tensor:
+    """一步 flow-map Euler（上游 scheduler.step，按 num_train_timesteps 归一化后的形式）：
+
+        x_r = x_t − (σ_t − σ_r)·u(x_t, t, r)
+
+    符号与 DiffSynth 的 `prev_sample = sample + model_output·(σ_next − σ)` 一致（σ_next < σ），
+    所以没有方向歧义。**全程无新噪声** —— 这与现有 `deploy_parallel/deploy_rollout` 的
+    "估 x0 再用新噪声重新加噪" 是**两个不同的采样器**，把 flow-map 权重喂进后者会出似是而非的垃圾。
+    """
+    s_t, s_r = float(sigma_t), float(sigma_r)
+    assert s_t > s_r - 1e-9, f"flowmap: Euler 必须从噪往净走，得到 σ_t={s_t} → σ_r={s_r}"
+    return x - (s_t - s_r) * u
+
+
+def euler_chain(u_fn, x: torch.Tensor, sigmas, pin_fn=None) -> torch.Tensor:
+    """沿一条 any-step 轴跑完整条 Euler 链（参考实现 / 解析门禁用；真正的部署采样器带 KV cache
+    与逐 cell 提交，在 `distill/eval/rollout_eval.py`）。
+
+    u_fn(x, t_val, r_val) → u；t = σ·1000。`pin_fn(x)` 每步之后调一次（TI2V 锚点重新钉 cell 0）。
+    """
+    sig = sigmas if torch.is_tensor(sigmas) else torch.as_tensor(sigmas, dtype=torch.float32)
+    assert_anystep_axis(sig, terminal=bool(float(sig[-1]) == 0.0))
+    for i in range(int(sig.numel()) - 1):
+        s_t, s_r = float(sig[i]), float(sig[i + 1])
+        u = u_fn(x, s_t * NUM_TRAIN_TIMESTEPS, s_r * NUM_TRAIN_TIMESTEPS)
+        x = euler_step(x, u, s_t, s_r)
+        if pin_fn is not None:
+            x = pin_fn(x)
+    return x

backup_2026-09-14/code/distill/longtraj.py

@@ -0,0 +1,201 @@
+"""把 21-cell 的训练片拼成连续长轨迹 —— 长 horizon 训练/评测的数据层。
+
+**为什么需要它**：stage1/2/3 全部只在 21 cell（5 秒）上训，而评测跑 41 cell、demo 想跑几分钟。
+loss 从来没有覆盖过 5 秒以上的时间跨度。数据里其实有长轨迹，只是从没被用过。
+
+**数据事实**（都是实测，不是假设）：
+  · `data/latent/train_eybx/*.pt` 每个是 21 cell = 81 帧 @16fps
+  · 同一 (session, seg_id) 下 `start_cell` 以 **21 为步长**递增 ⇒ 相邻片是同一条视频的连续窗口
+  · 实测最长连续段 **45 个 clip = 945 cell = 236 秒**；扫 4000 个样本里 239/678 组有 ≥5 个连续片
+  · ★ 但相邻片的 `video_t0` 相差 **5.25 s = 84 帧**，而每片只覆盖 81 帧
+    ⇒ **每个接缝有 3 帧（0.1875 s）的空隙**，按 69 px/s 约 13 px。源 mp4 也是 81 帧，
+      这 3 帧在切数据时就丢了，重编也拿不回来。这是这个数据集的硬限制，不要假装它不存在。
+
+**世界坐标怎么接**（分解成两段，各用最可靠的来源）：
+  片内 80 帧   用该片自己的 `frame_offset_px`（sidecar 真值，精确）
+  接缝 3 帧    用片末的逐帧速度外推（误差 ~2 px）
+  交叉校验     `data/meta/*.jsonl` 里每片有起点世界坐标 (x, z)；实测换算是
+               px_x/Δx = +29.84、px_y/Δz = −19.04，比值 1.567 与 walk_speed_px_s 的
+               69.24/46.01 = 1.505 相差 4%（深度轴透视压缩）——两个独立来源互证。
+               直接拿 x/z 算整段位移不行：逐片估计的离散有 9~19%，整段就是几十 px 的误差。
+
+  python distill/longtraj.py --build          # 建索引（读 jsonl，不碰 .pt，秒级）
+  python distill/longtraj.py --verify         # 拼接自检：接缝位移 vs x/z 交叉校验
+"""
+from __future__ import annotations
+
+import argparse
+import collections
+import glob
+import json
+import os
+
+import numpy as np
+import torch
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+LATENT_DIR = os.path.join(ROOT, "data", "latent", "train_eybx")
+META_GLOB = os.path.join(ROOT, "data", "meta", "*.jsonl")
+INDEX = os.path.join(ROOT, "data", "latent", "longtraj_index.pt")
+
+CELLS_PER_CLIP = 21
+FRAMES_PER_CLIP = 81          # 1 + 4×20
+STRIDE_CELLS = 21
+STRIDE_FRAMES = 84            # 相邻片的 video_t0 间隔
+GAP_FRAMES = STRIDE_FRAMES - FRAMES_PER_CLIP        # == 3
+# 实测标定（见文件头）；只用于交叉校验，不用于生成偏移
+PX_PER_GAME_X, PX_PER_GAME_Z = 29.840, -19.040
+
+
+def build_index(latent_dir: str = LATENT_DIR, meta_glob: str = META_GLOB,
+                out: str = INDEX, log=print) -> dict:
+    """从 jsonl 建 (session, seg) → 连续 clip 列表 的索引。不 torch.load 任何 .pt。"""
+    have = {os.path.basename(p)[:-3] for p in glob.glob(os.path.join(latent_dir, "*.pt"))}
+    log(f"[index] latent 目录里 {len(have)} 个片")
+    by = collections.defaultdict(list)
+    n_meta = n_skip = 0
+    for mp in sorted(glob.glob(meta_glob)):
+        with open(mp, encoding="utf-8") as f:
+            for line in f:
+                try:
+                    r = json.loads(line)
+                except json.JSONDecodeError:
+                    continue                      # jsonl 里有个别截断行，跳过
+                n_meta += 1
+                if r["clip_id"] not in have:
+                    continue
+                # ★ transition.jsonl 的 1900 条是**合成转场片**（带 pair/src_a/src_b/n_burn，
+                #   由两段不同素材缝出来的），没有 seg_id/x/z，也根本不是连续视频 —— 必须排除，
+                #   否则会把人工缝合处当成真实的世界运动去训。
+                if "seg_id" not in r:
+                    n_skip += 1
+                    continue
+                by[(r["session"], r["seg_id"])].append(
+                    (r["start_cell"], r["clip_id"], r.get("x"), r.get("z")))
+    runs = []
+    for (sess, seg), items in by.items():
+        items.sort()
+        cur = [items[0]]
+        for prev, nxt in zip(items, items[1:]):
+            if nxt[0] - prev[0] == STRIDE_CELLS:
+                cur.append(nxt)
+            else:
+                if len(cur) > 1:
+                    runs.append(dict(session=sess, seg_id=seg, items=cur))
+                cur = [nxt]
+        if len(cur) > 1:
+            runs.append(dict(session=sess, seg_id=seg, items=cur))
+    runs.sort(key=lambda r: -len(r["items"]))
+    idx = dict(runs=runs, cells_per_clip=CELLS_PER_CLIP, gap_frames=GAP_FRAMES,
+               stride_cells=STRIDE_CELLS)
+    torch.save(idx, out)
+    lens = np.array([len(r["items"]) for r in runs])
+    log(f"[index] 读 {n_meta} 条 meta（跳过 {n_skip} 条合成转场片）⇒ {len(runs)} 个连续段（≥2 片）")
+    log(f"[index] 段长（片）：中位 {int(np.median(lens))} · p90 {int(np.percentile(lens,90))} · 最长 {lens.max()}"
+        f"  ⇒ 最长 {lens.max()*CELLS_PER_CLIP} cell = {lens.max()*CELLS_PER_CLIP*4/16:.0f} 秒")
+    for k in (2, 3, 5, 10, 20):
+        log(f"[index]   ≥{k:2d} 片（{k*CELLS_PER_CLIP*4/16:5.1f} 秒）的段：{int((lens>=k).sum())}")
+    log(f"[index] 写入 {out}")
+    return idx
+
+
+def chain_offsets(frame_offsets: list[np.ndarray]) -> np.ndarray:
+    """把每片各自从 0 起算的逐帧偏移接成一条连续的世界轨迹 [L*81, 2]。
+
+    接缝的 3 帧空隙按**该片末尾的逐帧速度**外推（见文件头：片内用真值，只有空隙外推）。
+    """
+    out = [frame_offsets[0].copy()]
+    origin = frame_offsets[0][-1].copy()
+    for fo in frame_offsets[1:]:
+        prev = out[-1]
+        v_end = prev[-1] - prev[-2]                     # 片末逐帧速度
+        origin = origin + v_end * GAP_FRAMES            # 跨过 3 帧空隙
+        out.append(fo + origin)
+        origin = out[-1][-1].copy()
+    return np.concatenate(out, 0)
+
+
+def load_run(run: dict, n_clips: int, start: int = 0, latent_dir: str = LATENT_DIR,
+             sidecar=None):
+    """取 run 里连续的 n_clips 片，返回拼好的长轨迹。
+
+    返回 dict：latent [48, 21*n, 30, 52] · frame_off [81*n, 2] · cell_off [21*n, 2]
+             · prompts / actions / scenes（逐 cell 拼接）· seams（接缝处的绝对 cell 下标）
+    ★ latent 是**直接拼接**：第 j 片的 cell 0 在视频里只对应 1 帧（其余 cell 各 4 帧），
+      所以每个接缝那一格的"曝光"比别处短。这是数据的性质，不是 bug，但训练时要知道。
+    """
+    from actionrope import infer
+    if sidecar is None:
+        sidecar = torch.load(infer.SIDECAR, map_location="cpu", weights_only=False)
+    items = run["items"][start:start + n_clips]
+    assert len(items) == n_clips, f"该段只有 {len(run['items'])} 片，取不到 {start}+{n_clips}"
+    lats, fos, prompts, actions, scenes = [], [], [], [], []
+    for _, cid, _, _ in items:
+        d = torch.load(os.path.join(latent_dir, cid + ".pt"), map_location="cpu",
+                       weights_only=False)
+        lats.append(d["latent"])
+        fos.append(sidecar["clips"][cid]["frame_offset_px"].numpy().astype(np.float64))
+        prompts += list(d["prompts"])
+        actions += list(d["actions"])
+        scenes += list(d["scenes"])
+    frame_off = chain_offsets(fos)
+    from distill.eval.rollout_eval import frames_to_cells_n
+    n_cells = CELLS_PER_CLIP * n_clips
+    # 逐 cell 归约按每片各自的 21 格做，再拼 —— 不能整条丢给 frames_to_cells_n，
+    # 因为它假设「cell 0 = 1 帧、其余 4 帧」，而拼接后每 21 格就有一个这样的格子。
+    cell_off = np.concatenate(
+        [frames_to_cells_n(frame_off[i*FRAMES_PER_CLIP:(i+1)*FRAMES_PER_CLIP], CELLS_PER_CLIP)
+         for i in range(n_clips)], 0)
+    return dict(latent=torch.cat(lats, dim=1), frame_off=frame_off, cell_off=cell_off,
+                prompts=prompts, actions=actions, scenes=scenes,
+                seams=[i * CELLS_PER_CLIP for i in range(1, n_clips)],
+                clip_ids=[c for _, c, _, _ in items], n_cells=n_cells)
+
+
+def verify(n_runs: int = 20, n_clips: int = 5, log=print) -> bool:
+    """拼接自检：把外推出来的接缝位移和 meta 里的 (x,z) 真值对照。"""
+    from actionrope import infer
+    idx = torch.load(INDEX, map_location="cpu", weights_only=False)
+    sc = torch.load(infer.SIDECAR, map_location="cpu", weights_only=False)
+    errs, mags = [], []
+    for run in [r for r in idx["runs"] if len(r["items"]) >= n_clips][:n_runs]:
+        items = run["items"][:n_clips]
+        if any(it[2] is None for it in items):
+            continue
+        fos = [sc["clips"][c]["frame_offset_px"].numpy().astype(np.float64)
+               for _, c, _, _ in items]
+        chained = chain_offsets(fos)
+        for j in range(1, len(items)):
+            got = chained[j * FRAMES_PER_CLIP] - chained[(j - 1) * FRAMES_PER_CLIP]
+            dx = items[j][2] - items[j - 1][2]
+            dz = items[j][3] - items[j - 1][3]
+            ref = np.array([PX_PER_GAME_X * dx, PX_PER_GAME_Z * dz])
+            errs.append(np.linalg.norm(got - ref))
+            mags.append(np.linalg.norm(ref))
+    errs, mags = np.array(errs), np.array(mags)
+    rel = errs / np.maximum(mags, 1e-6)
+    log(f"[verify] {len(errs)} 个接缝：外推位移 vs (x,z) 真值")
+    log(f"[verify]   绝对误差 中位 {np.median(errs):.1f} px  p90 {np.percentile(errs,90):.1f} px")
+    log(f"[verify]   相对误差 中位 {np.median(rel):.3f}   （标定本身离散 9~19%，所以 ≲0.2 即为一致）")
+    log(f"[verify]   位移量级 中位 {np.median(mags):.1f} px / 片")
+    ok = float(np.median(rel)) < 0.25
+    log("[verify] " + ("通过：拼接的世界坐标与独立来源的 (x,z) 一致。"
+                       if ok else "★ 未通过：拼接偏移与 (x,z) 对不上，别拿它训长 horizon。"))
+    return ok
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(description="长轨迹数据层")
+    ap.add_argument("--build", action="store_true")
+    ap.add_argument("--verify", action="store_true")
+    ap.add_argument("--n_clips", type=int, default=5)
+    a = ap.parse_args(argv)
+    if a.build:
+        build_index()
+    if a.verify:
+        return 0 if verify(n_clips=a.n_clips) else 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

backup_2026-09-14/code/distill/lora.py

@@ -0,0 +1,180 @@
+"""低秩补丁层（LoRA）—— 忠实移植 AnyFlow on-policy 阶段的训练面。
+
+**为什么需要它。** 我们在 flow-map 学生上做 on-policy DMD 失败了五次，形状完全一样：权重相对
+初值只挪到 1e-4 量级，41-cell rollout 就开始退化，再往前走就烧白/烧黑。追过的原因（学习率、
+CFG 强度、critic 更新比、上下文梯度、监督覆盖面）逐一被对照实验否掉。
+
+最后核实出来的差异只有一处：**上游两个阶段都只训 LoRA，我们做的是全量微调。**
+
+    AnyFlow onpolicy yml:
+      lora_config.target_module_name = [attn1.to_q, attn1.to_k, attn1.to_v, attn1.to_out.0,
+                                        ffn.net.0.proj, ffn.net.2, delta_embedder, time_embedder]
+      param_names_to_optimize        = [lora_A, lora_B, far_patch_embedding]
+      lora_rank = lora_alpha = 256      （scale = alpha/rank = 1.0）
+    critic（discriminator）用同一份 target 但**不含** delta/time 两个嵌入。
+
+破坏的载体已经用拼接实验锁定：是 self-attn / ffn 的**权重矩阵**，bias / norm / embedding
+全部无辜 —— 正是 LoRA 会挂的那一批。
+
+★ 但「低秩之所以安全，是因为破坏藏在低秩表达不出来的方向里」这个说法**已被实测否掉**，别再用它当理由：
+  把塌掉的那份改动量做 rank-256 截断（只保留 58~73% 的能量）再加回初值，模型**照样烧白**。
+  破坏就在前 256 个主方向里，rank-256 的补丁完全有能力表达出来。
+  （脚本见 scratchpad/ranktrunc.py 的做法：对每个 target 矩阵的 Δ 做 SVD 截断，其余键保留成坏的，
+    这样变量就只剩「秩」这一条。）
+
+所以留给补丁的理由只剩两条，都比原来弱：
+  · B 零初始化 ⇒ 第 0 步与初值**逐位相同**，更新从零速率起步。补丁训练走的是另一条路径，
+    不是「全量微调那个坏终点的低秩投影」—— 上面那个实验没有触及这一条。
+  · 上游在同类学生（flow map）上就是这么跑通的，而这是我们与它之间最后一处已知差异。
+
+**为什么自己写而不用 peft。** peft 的 get_peft_model 会重写模块树并加自己的一层包装，与我们
+三种自定义前向模式（并行 / TF 双流 / 单 cell 流式 KV）和 FSDP 的交互没人验过；而我们要的东西
+很简单——包住若干个 nn.Linear。自己写还有一个实际好处：合并回去这条路是显然正确的，
+存盘拿到的仍然是 832 键的普通权重，**部署与评测路径一行都不用改**。
+
+用法：
+    install_lora(unwrap(student).dit, rank=256, alpha=256, is_student=True)
+    install_lora(unwrap(critic).dit,  rank=256, alpha=256, is_student=False)
+    opt = AdamW(lora_parameters(dit), lr=2e-6, betas=(0.0, 0.999))
+    ...
+    sd = merged_state_dict(dit)      # 832 键，B@A 已折进 base，可直接存
+"""
+from __future__ import annotations
+
+import contextlib
+import math
+import re
+
+import torch
+import torch.nn as nn
+
+# 上游 target 在我们的模块树里的对应名（见文件头的 yml 引用）
+TARGET_BLOCK = re.compile(r"^blocks\.\d+\.(self_attn\.[qkvo]|ffn\.[02])$")
+TARGET_STUDENT_EXTRA = re.compile(r"^(delta_embedding|time_embedding)\.[02]$")
+
+
+def is_target(name: str, *, is_student: bool) -> bool:
+    """cross_attn 不在上游 target 里，跟着排除。"""
+    if TARGET_BLOCK.match(name):
+        return True
+    return is_student and bool(TARGET_STUDENT_EXTRA.match(name))
+
+
+class LoRALinear(nn.Module):
+    """base(x) + (x @ A^T) @ B^T · (alpha/rank)。base 冻结，bias 从不参与。
+
+    ★ B 必须零初始化：这是「第 0 步与初值逐位相同」的唯一保证，也是低秩之外的另一层保护。
+      改成随机初始化不会报错，只会让训练从一个偏离的点开始 —— 那正是我们要避免的东西。
+    """
+
+    def __init__(self, base: nn.Linear, rank: int, alpha: int):
+        super().__init__()
+        assert isinstance(base, nn.Linear), type(base)
+        self.base = base
+        for p in self.base.parameters():
+            p.requires_grad_(False)
+        d_out, d_in = base.weight.shape
+        self.rank, self.scale = rank, alpha / rank
+        f = base.weight.dtype
+        self.lora_A = nn.Parameter(torch.empty(rank, d_in, dtype=f, device=base.weight.device))
+        self.lora_B = nn.Parameter(torch.zeros(d_out, rank, dtype=f, device=base.weight.device))
+        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+        self.enabled = True          # False ⇒ 只走 base（OPSD-V 的"教师 = 同一模型关掉 adapter"，见 lora_disabled）
+
+    def forward(self, x):
+        if not self.enabled:
+            return self.base(x)
+        return self.base(x) + torch.nn.functional.linear(
+            torch.nn.functional.linear(x, self.lora_A), self.lora_B) * self.scale
+
+    @torch.no_grad()
+    def delta_w(self) -> torch.Tensor:
+        return (self.lora_B.float() @ self.lora_A.float() * self.scale)
+
+
+def install_lora(root: nn.Module, *, rank: int = 256, alpha: int = 256,
+                 is_student: bool, log=print) -> int:
+    """把匹配到的 nn.Linear 就地换成 LoRALinear，并冻结其余所有参数。返回替换个数。"""
+    hits = [(n, m) for n, m in root.named_modules()
+            if isinstance(m, nn.Linear) and is_target(n, is_student=is_student)]
+    if not hits:
+        raise SystemExit("install_lora 没匹配到任何 Linear —— 模块名口径变了")
+    for name, mod in hits:
+        parent = root.get_submodule(name.rsplit(".", 1)[0]) if "." in name else root
+        setattr(parent, name.rsplit(".", 1)[-1], LoRALinear(mod, rank, alpha))
+    # 冻结一切非 LoRA 参数（含所有 bias / norm / embedding / patch_embedding / head）
+    n_tr = n_all = 0
+    for n, p in root.named_parameters():
+        keep = n.endswith(("lora_A", "lora_B"))
+        p.requires_grad_(keep)
+        n_all += p.numel()
+        n_tr += p.numel() if keep else 0
+    log(f"[lora] {'学生' if is_student else 'critic'}：包了 {len(hits)} 个 Linear "
+        f"(rank={rank} alpha={alpha} scale={alpha/rank:g})，"
+        f"可训 {n_tr/1e6:.1f}M / {n_all/1e6:.0f}M = {100*n_tr/n_all:.2f}%")
+    return len(hits)
+
+
+@contextlib.contextmanager
+def lora_disabled(root: nn.Module):
+    """临时关掉 root 下所有 LoRALinear 的补丁项：前向逐位等于冻结的 base（官方 OPSD-V 的教师就是这么切的，
+    _swap_generator_to_ema / teacher adapter；这里 base 冻结 ⇒ 教师 = init_ckpt）。嵌套安全，异常也会恢复。"""
+    mods = [m for m in root.modules() if isinstance(m, LoRALinear)]
+    prev = [m.enabled for m in mods]
+    for m in mods:
+        m.enabled = False
+    try:
+        yield
+    finally:
+        for m, e in zip(mods, prev):
+            m.enabled = e
+
+
+def lora_parameters(root: nn.Module):
+    ps = [p for n, p in root.named_parameters() if n.endswith(("lora_A", "lora_B"))]
+    if not ps:
+        raise SystemExit("lora_parameters 空 —— install_lora 没跑或名字对不上")
+    return ps
+
+
+@torch.no_grad()
+def merged_state_dict(root: nn.Module, src: dict | None = None) -> dict:
+    """把 B@A·scale 折进 base.weight，返回**没有 LoRA 包装**的普通 state_dict。
+
+    src 非 None 时从它取值（FSDP 下先 accelerator.get_state_dict 拿到全量，再在这里折）；
+    否则直接读 root 自己的（单进程 / CPU 路径）。
+    """
+    sd = dict(src if src is not None else root.state_dict())
+    out, n = {}, 0
+    # ★ FSDP 会把每个 DiTBlock 就地包一层，named_modules() 的名字里多出 `_fsdp_wrapped_module`（激活检查点则是
+    #   `_checkpoint_wrapped_module`），而 FULL_STATE_DICT 的键是干净的 —— 不剥掉这些壳，下面按 owner 查模块会全部落空
+    #   （2026-09-10 stage4 LoRA 冒烟第一次存盘就是这么炸的）。
+    _WRAP = {"_fsdp_wrapped_module", "_checkpoint_wrapped_module"}
+    lora = {".".join(p for p in n_.split(".") if p not in _WRAP): m
+            for n_, m in root.named_modules() if isinstance(m, LoRALinear)}
+    for k, v in sd.items():
+        if k.endswith((".lora_A", ".lora_B")):
+            continue
+        if ".base." in k:
+            owner, tail = k.split(".base.", 1)
+            # ★ 前缀无关地找模块：save_dit 的 post 钩子在 strip_dit_prefix **之前**跑，
+            #   所以这里的键还带着 "dit." 之类的外壳前缀，直接拿 owner 去查会全部落空
+            #   （表现为「一个矩阵都没合并」——下面那条断言就是为这个留的）。
+            mod = lora.get(owner)
+            if mod is None:
+                for cut in range(1, owner.count(".") + 1):
+                    cand = owner.split(".", cut)[-1]
+                    if cand in lora:
+                        mod = lora[cand]
+                        break
+            if mod is not None and tail == "weight":
+                A, B = sd.get(owner + ".lora_A"), sd.get(owner + ".lora_B")
+                if A is not None and B is not None:
+                    v = (v.float() + (B.float() @ A.float()) * mod.scale).to(v.dtype)
+                    n += 1
+            out[owner + "." + tail] = v
+        else:
+            out[k] = v
+    if n == 0 and lora:
+        raise SystemExit("merged_state_dict 一个矩阵都没合并 —— 键名口径对不上，别存这份权重")
+    return out

backup_2026-09-14/code/distill/onpolicy.py

@@ -0,0 +1,378 @@
+"""stage3-onpolicy 的算法层：AnyFlow 的 on-policy DMD，移植到 ActionRoPE 的因果 AR + 滑窗 KV。
+
+上游 = `vendor/AnyFlow/far/trainers/trainer_far_wan_anyflow_onpolicy.py`
+     + `vendor/AnyFlow/far/pipelines/pipeline_far_wan_anyflow.py::inference / training_rollout`
+
+它与我们原来那份 `stage3_dmd/dmd_arope.py`（WorldMind 口径）有三处**本质**差别，只有第一处是
+新机制，另外两处是命名/形式上的澄清：
+
+1. ★ **rollout 恒定 3 次前向，与步数 N 无关**（`pipeline_far_wan_anyflow.py:443-449`）。
+   因为学生是 flow map，「从任意 t 跳到任意 r」本来就是它的定义，所以中间步不用跑：
+
+       prev    = [σ_0,     σ_g  ]   一步跳到 g
+       current = [σ_g,     σ_{g+1}]
+       post    = [σ_{g+1}, 0    ]   一步跳到终点
+
+   三段**都带梯度**（上游 inference() 里没有任何 no_grad）。压掉的是**前向次数**（N → 3），
+   不是带梯度的图 —— 每个 in-window cell 仍有 3 次图，是旧 exit_step 版（1 次/cell）的 3 倍。
+   所以 grad_horizon 要相应调小：horizon 20 在 8×140 GB 上直接 OOM，实测 6 可跑。
+
+   `g = randrange(0, N)` 每次 rollout 抽一次（各 rank 同步）。WorldMind 的 exit_step 是
+   「跑到第 g 步就停」，要跑 g 次前向；这里是「跳过去」，永远 3 次。
+
+2. ★ **每次 rollout 的步数 N 从 `[2,4,8,16,50]` 里随机抽**（`trainer:213`）。这是 on-policy
+   阶段**保住 any-step 性质**的唯一机制：只在一张固定表上做 DMD，会把别的步数训丢。
+
+3. 上游叫 `discriminator` 的那个模型**不是 GAN 判别器** —— `discriminator_loss` 是在学生
+   自己生成的样本上做普通 flow-matching 回归（`trainer:246-251`），也就是 DMD 的 **fake score**。
+   与我们 `dmd_arope.py` 里的 `critic` 是同一个东西，这里沿用 `critic` 这个名字。
+
+三条与 ActionRoPE 契约相关、上游没有的约束（写错都不报错）：
+  * cell 0 是干净 TI2V 锚点，任何时候都不进 loss、不被加噪。
+  * 提交进 KV cache 恒用 `t=r=0` 的 detach 结果 —— 这是 TF clean 流的同分布要求，与 flow map 无关。
+  * real / fake score 都在 `r == t` 上调用（上游 `trainer:120,125` 明写 `r_timestep=timesteps`），
+    也就是当**普通扩散**用。它们不必是 flow map，我们就用 827 键的 stage1b 教师。
+"""
+from __future__ import annotations
+
+import torch
+
+from distill import flowmap as fm
+
+__all__ = [
+    "STEP_LIST", "draw_rollout_plan", "RolloutPlan",
+    "three_range_sigmas", "x0_from_u", "dmd_grad", "dmd_generator_loss",
+    "critic_flow_loss", "logit_normal_timesteps", "anystep_rollout",
+]
+
+# 上游 `rollout_cfg.num_inference_steps_list`（onpolicy yml:105）
+STEP_LIST = (2, 4, 8, 16, 50)
+
+
+class RolloutPlan:
+    """一次 rollout 的抽签结果。所有 rank 必须拿到**同一份**，否则三段的 σ 对不上、
+    DMD 的梯度会在不同 rank 上落在不同噪声档，all-reduce 之后是一锅粥（且不报错）。"""
+
+    __slots__ = ("n_steps", "grad_idx", "sigmas")
+
+    def __init__(self, n_steps: int, grad_idx: int, sigmas: torch.Tensor):
+        self.n_steps, self.grad_idx, self.sigmas = int(n_steps), int(grad_idx), sigmas
+
+    def __repr__(self):
+        return f"RolloutPlan(N={self.n_steps}, g={self.grad_idx}, σ={[round(float(x),4) for x in self.sigmas]})"
+
+
+def draw_rollout_plan(gen: torch.Generator | None = None, *, step_list=STEP_LIST,
+                      device=None, accelerator=None) -> RolloutPlan:
+    """抽 (N, g) 并算出 any-step 轴。`accelerator` 非 None 时从 rank0 广播（上游 trainer:214,219）。
+
+    `gen` 只在单进程/测试里用；多卡一律走广播 —— 各 rank 自己抽同一个种子也「通常」一致，
+    但一旦哪个 rank 多消耗一次随机数就会静默错开，广播才是硬保证。
+    """
+    if accelerator is not None and getattr(accelerator, "num_processes", 1) > 1:
+        buf = torch.zeros(2, dtype=torch.long, device=accelerator.device)
+        if accelerator.is_main_process:
+            i = int(torch.randint(len(step_list), (1,), generator=gen).item())
+            n = int(step_list[i])
+            buf[0], buf[1] = n, int(torch.randint(n, (1,), generator=gen).item())
+        torch.distributed.broadcast(buf, src=0)
+        n_steps, grad_idx = int(buf[0]), int(buf[1])
+    else:
+        n_steps = int(step_list[int(torch.randint(len(step_list), (1,), generator=gen).item())])
+        grad_idx = int(torch.randint(n_steps, (1,), generator=gen).item())
+    return RolloutPlan(n_steps, grad_idx, fm.anystep_sigmas(n_steps, device=device))
+
+
+def three_range_sigmas(plan: RolloutPlan) -> tuple[tuple[float, float] | None,
+                                                   tuple[float, float],
+                                                   tuple[float, float] | None]:
+    """把 any-step 轴切成上游那三段（`pipeline_far_wan_anyflow.py:443-445`）。
+
+    轴长 N+1（末位 0.0）。g ∈ [0, N)：
+        prev    (σ[0],   σ[g])     —— g == 0 时是空段（起点就是 g），返回 None
+        current (σ[g],   σ[g+1])   —— 恒非空，带梯度
+        post    (σ[g+1], σ[-1])    —— g == N-1 时 σ[g+1] 就是末位 0.0，空段，返回 None
+
+    `t == r` 的段上游直接 `continue`（:394），所以这里用 None 表示，别让调用方去跑一次
+    Δσ=0 的 Euler —— 那一步数值上是恒等，但会白跑一次 5B 前向。
+    """
+    s = plan.sigmas
+    g, N = plan.grad_idx, plan.n_steps
+    assert 0 <= g < N, f"grad_idx {g} 必须落在 [0, {N})"
+    assert int(s.numel()) == N + 1, f"σ 轴长应为 {N + 1}，得到 {int(s.numel())}"
+    a, b, c = float(s[0]), float(s[g]), float(s[g + 1])
+    d = float(s[-1])
+    prev = None if g == 0 else (a, b)
+    cur = (b, c)
+    post = None if (g + 1) >= N else (c, d)
+    return prev, cur, post
+
+
+def x0_from_u(x: torch.Tensor, u: torch.Tensor, sigma_t) -> torch.Tensor:
+    """从 flow-map 输出取 x0 估计：一步跳到 σ_r = 0。
+
+    = 上游 `dmd_scheduler.step(noise_pred, t, noisy, torch.zeros_like(t))`
+      （`trainer:121,126,132`）—— r 恒为 0，所以就是 `x − σ_t·u`。
+    """
+    return fm.euler_step(x, u, sigma_t, 0.0)
+
+
+def dmd_grad(pred_fake_x0: torch.Tensor, pred_real_x0: torch.Tensor,
+             generated: torch.Tensor, *, normalization: bool = True,
+             eval_mask: torch.Tensor | None = None) -> torch.Tensor:
+    """DMD 梯度（上游 `_compute_kl_grad` step 3-4，trainer:141-150）。
+
+        grad = (x0_fake − x0_real) / mean|generated − x0_real|
+
+    归一化的分母是**逐样本**的（上游 `mean(dim=[1,2,3,4], keepdim=True)`）。
+    `eval_mask`（[F] bool）非 None 时只在被打分的 cell 上归约 —— ActionRoPE 的 cell 0 是干净
+    锚点，把它算进分母会把归一化系数拉偏；上游没有这个概念，所以它没有这一项。
+
+    上游**没有** clamp(min=1e-2)（那是 WorldMind 的 dmd_gm.py 有的）。这里保持上游原样，
+    只做 nan_to_num —— 分母是 |真样本 − 真分数估计| 的均值，退化到 0 意味着 real score 已经
+    完美，那时 grad 本来就该是 0/0，nan_to_num 会把它变成 0，是对的。
+    """
+    grad = pred_fake_x0 - pred_real_x0
+    if normalization:
+        p_real = generated - pred_real_x0
+        if eval_mask is not None:
+            idx = torch.nonzero(eval_mask, as_tuple=True)[0]
+            norm = p_real.index_select(2, idx).abs().mean(dim=[1, 2, 3, 4], keepdim=True)
+        else:
+            norm = p_real.abs().mean(dim=[1, 2, 3, 4], keepdim=True)
+        grad = grad / norm
+    return torch.nan_to_num(grad)
+
+
+def dmd_generator_loss(generated: torch.Tensor, grad: torch.Tensor, weight: float = 1.0) -> torch.Tensor:
+    """把 grad 变成一个可反传的标量（上游 trainer:195）：
+
+        loss = w · MSE(generated, (generated − grad).detach())
+
+    这不是一个"距离"，是把预先算好的梯度塞进 autograd 的标准手法：
+    d(loss)/d(generated) ∝ grad。所以 **loss 的数值本身没有意义**，别拿它当收敛判据。
+    上游用 double() 算这一步，照抄 —— generated 在 bf16 里，(g − grad) 的残差会被尾数吃掉。
+    """
+    return weight * torch.nn.functional.mse_loss(
+        generated.double(), (generated.double() - grad.double()).detach(), reduction="mean")
+
+
+def logit_normal_timesteps(batch: int, *, mean: float = 0.0, std: float = 1.0,
+                           gen: torch.Generator | None = None) -> torch.Tensor:
+    """critic 训练用的噪声档抽样（上游 `discriminator_loss` trainer:236，
+    `compute_density_for_timestep_sampling(weighting_scheme='logit_normal', logit_mean=0, logit_std=1)`）。
+
+    返回 [b] 的 **shift 前** u ∈ (0,1)：u = sigmoid(N(mean, std))。调用方还要 apply_shift 再 ×1000。
+    注意 generator 侧的 DMD 噪声档用的是**均匀** rand（trainer:183），两处不是同一个分布，
+    别顺手统一 —— fake score 要在整条轨迹上准，而 DMD 的梯度只在被抽到的那一档上估。
+    """
+    return torch.sigmoid(torch.randn(batch, generator=gen) * std + mean)
+
+
+def critic_flow_loss(u_pred: torch.Tensor, x0: torch.Tensor, noise: torch.Tensor,
+                     *, eval_mask: torch.Tensor | None = None) -> torch.Tensor:
+    """fake score（上游的 `discriminator`）的训练目标（trainer:247-251）：
+
+        target = ε − x0        普通 rectified flow，**不带** flow map 的 (t−r) 修正项
+                               —— 它是在 r == t 上调用的，修正项本来就是 0。
+
+    `eval_mask` 非 None 时只在被打分的 cell 上取均值（ActionRoPE 的 cell 0 不进 loss）。
+    """
+    target = noise - x0
+    se = (u_pred.float() - target.float()) ** 2
+    if eval_mask is not None:
+        idx = torch.nonzero(eval_mask, as_tuple=True)[0]
+        se = se.index_select(2, idx)
+    return se.mean()
+
+
+# ----------------------------------------------------------------------------
+# rollout：因果 AR + 滑窗 KV，每个 cell 走三段跳
+# ----------------------------------------------------------------------------
+
+def anystep_rollout(student, initial_latent, ctx_ext, ids_ext, off_ext, num_gen_cells: int,
+                    plan: "RolloutPlan", *, frame_cond, grad_horizon: int,
+                    kv_window: int = 20, rope_cap: int = 20, sink_size: int = 0,
+                    grad_start: int | None = None, gen: torch.Generator | None = None):
+    """学生用自己提交的历史自回归生成 num_gen_cells 个 cell —— on-policy 版。
+
+    与 `stage3_dmd.self_forcing_rollout` 逐行对应，**只有一处不同**：每个 cell 内部不再沿
+    部署表逐步去噪 + 重新加噪，而是走上游的三段跳（`pipeline_far_wan_anyflow.py:443-449`）：
+
+        x ←(no_grad) jump(σ₀ → σ_g)        prev
+        x ←( grad  ) jump(σ_g → σ_{g+1})   current   ← 全 rollout 唯一带梯度的前向
+        x ←(no_grad) jump(σ_{g+1} → 0)     post
+
+    于是**每个 cell 恒定 ≤3 次前向、与 N 无关**，而且带梯度的图只有一次单帧前向 ——
+    `dmd_arope` 里那个 77.8 GB 激活的问题（以及为它写的 stream_grad_replay）在这里不存在。
+
+    返回 (generated [b,48,1+N,h,w], grad_mask [1+N] bool)，与 self_forcing_rollout 同签名，
+    所以下游的 score_tail / dmd_targets 那套可以原样接。
+    """
+    b, c, _, lat_h, lat_w = initial_latent.shape
+    dev = initial_latent.device
+    total = 1 + num_gen_cells
+    bare = unwrap_fn(student)
+    cache = bare.init_kv_cache(kv_window=kv_window, sink_size=sink_size, rope_cap=rope_cap)
+    prev_seg, cur_seg, post_seg = three_range_sigmas(plan)
+
+    def _call(x, s_t, s_r, k, commit, ow, mk):
+        """一次 flow-map 前向。**必须**走 forward(..., kv_cache=)：FSDP 的 all-gather 挂在
+        包装模块的 forward 上，直接调 forward_frame 拿到的是分片且不报错（dmd_arope risk 2）。"""
+        tt = torch.full((b,), float(s_t) * fm.NUM_TRAIN_TIMESTEPS, device=dev)
+        rr = torch.full((b,), float(s_r) * fm.NUM_TRAIN_TIMESTEPS, device=dev)
+        return student(x, tt, ctx_ext[:, k:k + 1], context_ids=ids_ext[:, k:k + 1],
+                       r_timestep=rr, first_frame_cond=False, frame_index=k,
+                       kv_cache=cache, commit=commit, offset_win=ow, mask_in_k=mk)
+
+    ow0, mk0 = frame_cond(off_ext, 0, sink_size, kv_window, lat_h, lat_w)
+    with torch.no_grad():
+        _call(initial_latent, 0.0, 0.0, 0, True, ow0, mk0)      # 真实首帧，t=r=0 进缓存
+    frame_list = [initial_latent[:, :, 0:1].detach()]
+
+    if grad_start is None:
+        grad_start = int(torch.randint(1, max(2, total - grad_horizon + 1), (1,), generator=gen).item())
+    grad_end = min(grad_start + grad_horizon, total)
+
+    for k in range(1, total):
+        ow, mk = frame_cond(off_ext, k, sink_size, kv_window, lat_h, lat_w)
+        # CPU 上抽再搬：`torch.randn(..., device=cuda, generator=cpu_gen)` 在 CUDA 上直接报
+        # "Expected a 'cuda' device type for generator"。与本文件其余噪声（generator_step /
+        # critic_step 里的 `randn(...).to(tail)`）口径一致，顺带让不同设备上可复现。
+        # 单 cell 的张量很小，多一次 H2D 拷贝可以忽略。
+        x = torch.randn(b, c, 1, lat_h, lat_w, generator=gen).to(device=dev, dtype=initial_latent.dtype)
+        in_window = grad_start <= k < grad_end
+        # ★ 三段**都**带梯度（上游 `inference()` 里没有任何 no_grad，pipeline:447-449）。
+        #   `grad_timestep` 的作用是把 N 次前向压成 **3** 次，不是「只让其中一次带梯度」——
+        #   我第一版写成了后者，那样会丢掉 prev/post 两条通路的梯度，且不会报错。
+        #   不在 grad 窗口里的 cell 仍然全程 no_grad：它们不进 loss，建图纯属浪费。
+        with (torch.enable_grad() if in_window else torch.no_grad()):
+            if prev_seg is not None:
+                x = fm.euler_step(x, _call(x, *prev_seg, k, False, ow, mk), *prev_seg)
+            x = fm.euler_step(x, _call(x, *cur_seg, k, False, ow, mk), *cur_seg)
+            if post_seg is not None:
+                x = fm.euler_step(x, _call(x, *post_seg, k, False, ow, mk), *post_seg)
+        frame_list.append(x)
+        with torch.no_grad():                                    # 提交趟：detach 的 x0 + t=r=0
+            _call(x.detach(), 0.0, 0.0, k, True, ow, mk)
+
+    generated = torch.cat(frame_list, dim=2)
+    grad_mask = torch.tensor([grad_start <= k < grad_end for k in range(total)], device=dev)
+    # ★ 必须显式释放：cache 与它内部的 `_ref` 互相引用，引用计数放不掉环，只能等惰性的循环 GC。
+    #   一条 rollout 的 K/V 是十几 GB，不放的话几步之内显存就爬满（实测 98→120→126 GB 后 OOM）。
+    bare.release_kv_cache(cache)
+    return generated, grad_mask
+
+
+def unwrap_fn(m):
+    """延迟 import，避免 distill.onpolicy ←→ distill.arope_dit 的循环。"""
+    from distill.arope_dit import unwrap
+    return unwrap(m)
+
+
+# ----------------------------------------------------------------------------
+# SGF：bounded two-pass replay（vendor/SGF，pipeline/self_gradient_forcing_training.py:117-230）
+# ----------------------------------------------------------------------------
+
+def sgf_rollout(student, initial_latent, ctx_ext, ids_ext, off_ext, num_gen_cells: int,
+                sigmas, exit_idx: int, *, context_noise: float = 0.0,
+                kv_window: int = 20, rope_cap: int = 20, sink_size: int = 0,
+                gen: torch.Generator | None = None, unwrap_fn=None, frame_cond_fn=None):
+    """SGF 的 pass 1：**全程 no_grad** 的 AR rollout，只负责「录像」。
+
+    录两样东西（上游 `_self_gradient_forcing_inference` 的 noisy_at_t / x_ctx_hat）：
+      noisy_at_t  每个 cell 在抽中的退出档 `exit_idx` 上的**带噪输入**
+      x0_cache    每个 cell 提交进 KV 缓存的那份 x0（cache_mode="exit"：退出档的 x0，不是走完的）
+
+    为什么 pass 1 可以全 no_grad —— 这正是 SGF 的论点：现在这条路上的梯度**本来就是断的**
+    （每个 cell 提交 KV 时是 `x.detach()`），所以在 pass 1 里建图只是浪费。真正的梯度由 pass 2
+    的一次并行 TF 前向补回来，那一次跨帧注意力才第一次带着梯度跑过自生成历史。
+
+    ★ 只跑到 exit_idx 就 break（上游 cache_mode="exit"，trainer:179）。每个 cell 的前向数是
+      exit_idx+1（去噪）+ 1（提交），期望约 N/2+1，与 anystep_rollout 的三段跳同量级。
+
+    返回 (noisy_at_t, x0_cache)，都是 [b,48,1+N,h,w] 且已 detach。
+    """
+    b, c, _, lat_h, lat_w = initial_latent.shape
+    dev = initial_latent.device
+    total = 1 + num_gen_cells
+    bare = unwrap_fn(student)
+    cache = bare.init_kv_cache(kv_window=kv_window, sink_size=sink_size, rope_cap=rope_cap)
+
+    def _call(x, s_t, s_r, k, commit, ow, mk):
+        tt = torch.full((b,), float(s_t) * fm.NUM_TRAIN_TIMESTEPS, device=dev)
+        rr = torch.full((b,), float(s_r) * fm.NUM_TRAIN_TIMESTEPS, device=dev)
+        return student(x, tt, ctx_ext[:, k:k + 1], context_ids=ids_ext[:, k:k + 1],
+                       r_timestep=rr, first_frame_cond=False, frame_index=k,
+                       kv_cache=cache, commit=commit, offset_win=ow, mask_in_k=mk)
+
+    ow0, mk0 = frame_cond_fn(off_ext, 0, sink_size, kv_window, lat_h, lat_w)
+    z0 = initial_latent[:, :, 0:1].detach()
+    with torch.no_grad():
+        _call(z0, 0.0, 0.0, 0, True, ow0, mk0)                  # 真实首帧，t=r=0 进缓存
+    # cell 0 是干净锚点：它的「带噪输入」就是它自己（pass 2 的 first_frame_cond 会把 t/r 抹成 0）
+    noisy_list, x0_list = [z0], [z0]
+
+    with torch.no_grad():
+        for k in range(1, total):
+            ow, mk = frame_cond_fn(off_ext, k, sink_size, kv_window, lat_h, lat_w)
+            x = torch.randn(b, c, 1, lat_h, lat_w, generator=gen).to(device=dev,
+                                                                    dtype=initial_latent.dtype)
+            noisy_at_exit, x0_exit = None, None
+            for i in range(len(sigmas) - 1):
+                s_t, s_r = float(sigmas[i]), float(sigmas[i + 1])
+                if i == exit_idx:
+                    noisy_at_exit = x                            # ★ 录的是**进入**这一档的输入
+                u = _call(x, s_t, s_r, k, False, ow, mk)
+                if i == exit_idx:
+                    x0_exit = x0_from_u(x, u, s_t)               # 退出档的 x0 估计
+                    break                                        # cache_mode="exit"
+                x = fm.euler_step(x, u, s_t, s_r)
+            assert noisy_at_exit is not None and x0_exit is not None, "没抓到退出档的状态"
+            noisy_list.append(noisy_at_exit)
+            x0_list.append(x0_exit)
+            _call(x0_exit, 0.0, 0.0, k, True, ow, mk)            # 提交：与 x0_cache 逐位相同
+    bare.release_kv_cache(cache)
+    return (torch.cat(noisy_list, dim=2).detach(),
+            torch.cat(x0_list, dim=2).detach())
+
+
+def sgf_replay(student, noisy_at_t, x0_cache, cond, sigmas, exit_idx: int, *,
+               context_noise: float = 0.0, gen=None, gradient_checkpointing: bool = True):
+    """SGF 的 pass 2：**一次带梯度的并行 TF 前向**，clean 半流喂 pass 1 自己生成并提交的历史。
+
+    这就是被补回来的那条通路 —— 跨帧注意力这次带着梯度跑过自生成上下文
+    （上游 trainer:216-230，`self_gradient_forcing_match_context=True` 分支）。
+
+    clean 半流必须与 pass 1 **真正提交进缓存的东西**一致（上游的 match_context）。我们提交的是
+    `x0_exit`，所以这里喂的就是 `x0_cache`；`context_noise>0` 时两边都要加同一档噪声。
+
+    几何：pass 2 是整片并行、first_frame_cond=True、锚在片自己的 cell 0；pass 1 的每个 cell
+    锚在自己窗口里最老那一格。二者只在**整窗平移**上不同，而 RoPE 是相对编码、平移不改分数——
+    前提是 rollout 长度 == 打分片长（都是 1+K），此时 pass 1 的缓存跨度与 pass 2 的片完全重合。
+    rollout 比片长时尾巴前段在 pass 1 里看过片外历史，那才是真的错配，所以别把 rollout 拉长。
+
+    返回 pass 2 预测的 x0 [b,48,1+N,h,w]，带梯度。
+    """
+    s_t, s_r = float(sigmas[exit_idx]), float(sigmas[exit_idx + 1])
+    b = noisy_at_t.shape[0]
+    dev = noisy_at_t.device
+    tt = torch.full((b,), s_t * fm.NUM_TRAIN_TIMESTEPS, device=dev)
+    rr = torch.full((b,), s_r * fm.NUM_TRAIN_TIMESTEPS, device=dev)
+    if context_noise > 0:
+        noise = torch.randn(x0_cache.shape, generator=gen).to(x0_cache)
+        clean_x = fm.add_noise_scalar(x0_cache, noise,
+                                      torch.full((b,), context_noise, device=dev))
+        clean_x = fm.keep_frame0(clean_x, x0_cache)              # 锚点始终精确干净
+        aug_t = torch.full((b,), context_noise * fm.NUM_TRAIN_TIMESTEPS, device=dev)
+    else:
+        clean_x, aug_t = x0_cache, None
+    # ★ 梯度检查点在这里不是可选项：pass 2 是双流 2×21×390 = 16380 token 过 30 层，
+    #   而 stage3 同时驻着四份模型（学生/critic 各带 fp32 主权重+优化器态、teacher、EMA）。
+    #   不开实测 8×140 GB 全部 OOM（137.7 GiB 已分配）。stage1 的 TF 也是开着的。
+    u = student(noisy_at_t, tt, cond["context"], context_ids=cond["context_ids"],
+                offset_px=cond["offset_px"], mask_input=cond["mask_input"],
+                first_frame_cond=True, r_timestep=rr,
+                clean_x=clean_x.detach(), aug_t=aug_t,
+                use_gradient_checkpointing=gradient_checkpointing)
+    return x0_from_u(noisy_at_t, u, s_t)

backup_2026-09-14/code/distill/runlib.py

@@ -0,0 +1,542 @@
+#!/usr/bin/env python3
+"""三阶段蒸馏共用的 run 目录 / 日志 / checkpoint 约定。
+
+    from distill import runlib
+    run_dir = runlib.resolve_run_dir(args.output)
+    log = runlib.TextTee(os.path.join(run_dir, "train.log"), enabled=accel.is_main_process)
+    runlib.save_dit(accel, student, runlib.dit_path(run_dir, step), {"step": step})
+    runlib.rotate(runlib.dit_glob(run_dir), args.keep_last_dits)
+
+目录布局（三个 stage 共用，run_dir 默认落在 <repo>/outputs/distill/<name>）：
+
+    <run_dir>/
+        config.json                            本次 run 的完整参数快照
+        progress.json                          {"step": N, ...}，续训只认它
+        train.log                              TextTee 的人读日志（与 stdout 同步）
+        train_log.csv                          CsvLog，一步一行，固定表头
+        train_log.jsonl                        JsonlLogger，一步一行 JSON（可被别的机器 tail）
+        dits/dit_step{N}.safetensors           学生原始权重，只为续训 —— 不要拿去部署
+        ema/ema_step{N}.safetensors            EMA 影子权重 = stage2 / stage3 的部署权重
+        critic/critic_step{N}.safetensors      stage3 的 fake-score critic，只为续训
+        state_A/ state_B/ + state -> 软链      accelerate 完整分片存档（只有 stage1 用）
+
+为什么 dits / ema / critic 都要 rotate：单个文件 10.0 GB。WorldMind 只 rotate 了
+state_step*（infra/run.py:rotate_states 的注释明说不动 dits/），它的 v4 stage-1 run
+光 dits/ 就吃掉 746 GB。这里 rotate() 对三种前缀一视同仁。
+
+checkpoint 格式沿用 ActionRoPE 的（不是 WorldMind 的）：单个 safetensors，827 个 bf16
+张量，DiT 键去掉 `dit.` 前缀（臂键保留 `arm.`），metadata 全是字符串，chmod 0644。
+这样 tests/check_ckpt.py、actionrope/infer.py、eval_bg50.py 不用改一行就能吃蒸馏产物。
+
+stage2-flowmap 的产物是这个格式的**唯一一处分叉**：`AropeDiT.install_flowmap()` 多挂 5 个
+张量（`delta_embedding.{0,2}.{weight,bias}` + `delta_gate`），827 → **832**（AROPE_DIT_KEYS_FLOWMAP，
++20.5 MB）。格式本身一个字没变（单文件 / 全 bf16 / 字符串 metadata / 0644 / .tmp+rename），
+只是键数变了 —— 于是上面那三个下游会**拒收**它。那是对的：flow map 要第二个时间步 r，拿旧的
+renoise 采样器去采它不会报错，只会安静地产出另一个模型的结果。用 is_flowmap_ckpt() 分流。
+
+★ save_dit 显式 cast 到 bf16 并断言导出 dtype。stage3 在 prepare 前把每个模型 .float()
+（FSDP + MixedPrecision(bf16) 下不这么做就没有 fp32 主权重），而 accelerator.get_state_dict
+在 FSDP 分支返回的是 FULL_STATE_DICT(offload_to_cpu=True, rank0_only=True) —— rank0 CPU 上的
+**fp32**。不 cast 的话每个文件 20 GB（10 次保存 × 3 个文件 = 600 GB 而不是 300 GB），
+而且 827-key / 10.0 GB 的格式承诺当场作废。
+"""
+from __future__ import annotations
+
+import csv
+import glob
+import json
+import os
+import re
+import shutil
+import struct
+import time
+
+import torch
+import torch.nn as nn
+from safetensors.torch import load_file, save_file
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+OUTPUT_ROOT = os.path.join(REPO_ROOT, "outputs", "distill")
+
+DIT_PREFIX = "dit."          # AropeDiT 里 WanModel 的属性名；导出时剥掉
+ARM_PREFIX = "arm."          # baseline 动作臂；蒸馏不用臂，保留只为格式兼容
+MASK_EMBEDDING_NAME = "arope_mask_embedding"
+MASK_KEYS = (f"{MASK_EMBEDDING_NAME}.weight", f"{MASK_EMBEDDING_NAME}.bias")
+AROPE_DIT_KEYS = 827         # 与 outputs/arope_10k_ga2/step-10000.safetensors 逐键一致
+REF_CKPT = os.path.join(REPO_ROOT, "outputs", "arope_10k_ga2", "step-10000.safetensors")
+
+# ---- flow map（stage2-flowmap）：第二时间步 r 的 delta 分支 ----------------------------------
+# AropeDiT.install_flowmap() 把 delta_embedding（= time_embedding 的 deepcopy，读出层零初始化）和
+# delta_gate（persistent buffer）挂在 **dit 上**，于是 strip_dit_prefix 剥掉 `dit.` 之后就是下面这
+# 5 个平铺键 —— 格式一个字没变（单 safetensors / 全 bf16 / 字符串 metadata / 0644 / .tmp+rename），
+# 只是键数 827 → 832（+20.5 MB，delta_embedding 是 10,229,760 个参数 = 5B 的 0.20 %）。
+#
+# ★ 没有"精确 827 键"的导出：residual 形式下 u(x,t,t) = head(E(t) + g·D(P(t)))，两条
+#   Sequential(Linear, SiLU, Linear) 的和折不回一条。所以 actionrope/infer.py、tests/check_ckpt.py、
+#   eval_bg50.py 会拒收 flow-map 权重 —— 这是 **对的**：拿旧的 renoise 采样器、不给 r 去采一个 flow map，
+#   出来的东西看着像样但是另一个模型。导出时另外写 metadata flowmap=1 / flowmap_form / delta_gate，
+#   下游按 metadata 选采样器，而不是按命令行开关。
+DELTA_EMBEDDING_NAME = "delta_embedding"
+DELTA_GATE_NAME = "delta_gate"
+FLOWMAP_KEYS = (
+    f"{DELTA_EMBEDDING_NAME}.0.weight", f"{DELTA_EMBEDDING_NAME}.0.bias",
+    f"{DELTA_EMBEDDING_NAME}.2.weight", f"{DELTA_EMBEDDING_NAME}.2.bias",
+    DELTA_GATE_NAME,
+)
+AROPE_DIT_KEYS_FLOWMAP = AROPE_DIT_KEYS + len(FLOWMAP_KEYS)   # 832
+
+
+# ----------------------------------------------------------------------------
+# 路径
+# ----------------------------------------------------------------------------
+
+def resolve_run_dir(out: str) -> str:
+    """`out` 可以是绝对路径、相对仓库根的路径、或一个裸名字（落在 outputs/distill/ 下）。"""
+    if not out or not str(out).strip():
+        raise ValueError("run dir 不能为空")
+    out = os.path.expanduser(str(out))
+    if os.path.isabs(out):
+        return os.path.normpath(out)
+    if os.sep in out:
+        return os.path.normpath(os.path.join(REPO_ROOT, out))
+    return os.path.join(OUTPUT_ROOT, out)
+
+
+def dits_dir(run_dir: str) -> str:
+    return os.path.join(run_dir, "dits")
+
+
+def ema_dir(run_dir: str) -> str:
+    return os.path.join(run_dir, "ema")
+
+
+def critic_dir(run_dir: str) -> str:
+    return os.path.join(run_dir, "critic")
+
+
+def dit_path(run_dir: str, step: int) -> str:
+    """学生原始权重：只为续训。部署一律用 ema_path()。"""
+    return os.path.join(dits_dir(run_dir), f"dit_step{int(step)}.safetensors")
+
+
+def ema_path(run_dir: str, step: int) -> str:
+    """EMA 影子权重 = stage2 / stage3 的部署权重。"""
+    return os.path.join(ema_dir(run_dir), f"ema_step{int(step)}.safetensors")
+
+
+def critic_path(run_dir: str, step: int) -> str:
+    return os.path.join(critic_dir(run_dir), f"critic_step{int(step)}.safetensors")
+
+
+def dit_glob(run_dir: str) -> str:
+    return os.path.join(dits_dir(run_dir), "dit_step*.safetensors")
+
+
+def ema_glob(run_dir: str) -> str:
+    return os.path.join(ema_dir(run_dir), "ema_step*.safetensors")
+
+
+def critic_glob(run_dir: str) -> str:
+    return os.path.join(critic_dir(run_dir), "critic_step*.safetensors")
+
+
+_STEP_RE = re.compile(r"step(\d+)")
+
+
+def _step_of(path: str) -> int | None:
+    """从 dit_step500.safetensors / ema_step500.safetensors / state_step500 里取出 500。"""
+    m = _STEP_RE.search(os.path.basename(os.path.normpath(path)))
+    return int(m.group(1)) if m else None
+
+
+def list_ckpts(pattern: str) -> list[tuple[int, str]]:
+    """[(step, path), ...] 按 step 升序。名字里解析不出 step 的条目直接忽略（也就永远不会被 rotate 删）。"""
+    out = [(s, p) for p in glob.glob(pattern) if (s := _step_of(p)) is not None]
+    return sorted(out)
+
+
+def latest_ckpt(pattern: str) -> tuple[int, str] | None:
+    c = list_ckpts(pattern)
+    return c[-1] if c else None
+
+
+def rotate(pattern: str, keep: int) -> list[str]:
+    """按 step 排序，删掉最旧的，保留最新 keep 个。dits / emas / critics / state_step* 通用。
+
+    keep <= 0 ⇒ 什么都不删（不是"全删"，别把唯一的续训点抹了）。解析不出 step 的路径不碰。
+    返回被删掉的路径。目录用 rmtree，文件用 remove。"""
+    keep = int(keep)
+    if keep <= 0:
+        return []
+    removed = []
+    for _, p in list_ckpts(pattern)[:-keep]:
+        try:
+            if os.path.isdir(p) and not os.path.islink(p):
+                shutil.rmtree(p, ignore_errors=True)
+            else:
+                os.remove(p)
+            removed.append(p)
+        except OSError:
+            pass
+    return removed
+
+
+def rotate_run(run_dir: str, keep_dits: int = 0, keep_emas: int = 0, keep_critics: int = 0) -> dict:
+    """一次把三种权重都 rotate 掉。默认 0 = 不动（显式传才删）。"""
+    return {"dits": rotate(dit_glob(run_dir), keep_dits),
+            "emas": rotate(ema_glob(run_dir), keep_emas),
+            "critics": rotate(critic_glob(run_dir), keep_critics)}
+
+
+def disk_free_gb(path: str = REPO_ROOT) -> float:
+    try:
+        return shutil.disk_usage(path).free / 1e9
+    except OSError:
+        return float("nan")
+
+
+def warn_if_low_disk(path: str = REPO_ROOT, min_gb: float = 200.0) -> str | None:
+    free = disk_free_gb(path)
+    if free == free and free < min_gb:  # 非 NaN 且不够
+        return (f"磁盘告急：{path} 只剩 {free:.0f} GB（< {min_gb:.0f} GB）。一个 dit/ema/critic 是 10.0 GB，"
+                f"一个 accelerate state 槽约 60 GB；调小 --save_every 或 --keep_last_*。")
+    return None
+
+
+# ----------------------------------------------------------------------------
+# progress.json / config.json
+# ----------------------------------------------------------------------------
+
+def _atomic_write_text(path: str, text: str) -> None:
+    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+    tmp = f"{path}.tmp"
+    with open(tmp, "w", encoding="utf-8") as fh:
+        fh.write(text)
+        fh.flush()
+        os.fsync(fh.fileno())
+    os.replace(tmp, path)
+
+
+def progress_path(run_dir: str) -> str:
+    return os.path.join(run_dir, "progress.json")
+
+
+def write_progress(run_dir: str, step: int, **extra) -> str:
+    """续训只认这个文件。原子写：保存 10 GB 权重时崩掉也不会留下半个 progress.json。"""
+    row = {"step": int(step), "wall": time.time(), "time": time.strftime("%Y-%m-%d %H:%M:%S"), **extra}
+    p = progress_path(run_dir)
+    _atomic_write_text(p, json.dumps(row, ensure_ascii=False, default=str))
+    return p
+
+
+def read_progress(run_dir: str) -> dict | None:
+    p = progress_path(run_dir)
+    if not os.path.isfile(p):
+        return None
+    try:
+        with open(p, encoding="utf-8") as fh:
+            return json.load(fh)
+    except (OSError, json.JSONDecodeError):
+        return None
+
+
+def progress_step(run_dir: str, default: int = 0) -> int:
+    pr = read_progress(run_dir)
+    return int(pr["step"]) if pr and "step" in pr else int(default)
+
+
+def save_config(run_dir: str, cfg: dict) -> str:
+    os.makedirs(run_dir, exist_ok=True)
+    p = os.path.join(run_dir, "config.json")
+    _atomic_write_text(p, json.dumps(cfg, indent=2, ensure_ascii=False, default=str))
+    return p
+
+
+# ----------------------------------------------------------------------------
+# checkpoint：导出
+# ----------------------------------------------------------------------------
+
+def _unwrap(model):
+    """剥掉 DDP / DeepSpeedEngine / EmaShadow 之类的外壳，拿到裸 module。"""
+    seen = 0
+    while hasattr(model, "module") and isinstance(getattr(model, "module"), nn.Module) and seen < 8:
+        model = model.module
+        seen += 1
+    return model
+
+
+def _gather_state_dict(accel, model) -> dict:
+    """accel=None ⇒ 单进程，直接 state_dict()。否则走 accelerator.get_state_dict（集合操作，
+    每个 rank 都必须进来；FSDP 分支是 rank0_only，非主进程拿到空 dict）。
+
+    非 nn.Module（distill.ema.EmaShadow）走本地 state_dict：它没进过 prepare，权重在每个 rank 上
+    都是整份的，get_state_dict 的 FSDP 分支反而会拿它当 FSDP 模块去遍历。"""
+    if accel is None or not isinstance(model, nn.Module):
+        return _unwrap(model).state_dict()
+    return accel.get_state_dict(model)
+
+
+def strip_dit_prefix(state_dict: dict) -> tuple[dict, list[str]]:
+    """AropeDiT 包装的 state_dict → ActionRoPE 的导出键集。
+
+    有 `dit.` 键 ⇒ 剥前缀，`arm.` 键原样保留，其余（包装层自己的 buffer 之类）丢掉并回报；
+    没有 `dit.` 键 ⇒ 传进来的本来就是裸 WanModel，原样返回。"""
+    if not any(k.startswith(DIT_PREFIX) for k in state_dict):
+        return dict(state_dict), []
+    out, dropped = {}, []
+    for k, v in state_dict.items():
+        if k.startswith(DIT_PREFIX):
+            out[k[len(DIT_PREFIX):]] = v
+        elif k.startswith(ARM_PREFIX):
+            out[k] = v
+        else:
+            dropped.append(k)
+    return out, dropped
+
+
+def save_dit(accel, model, path: str, metadata: dict | None = None, *,
+             expect_keys: int | None = AROPE_DIT_KEYS, require_mask: bool = True,
+             dtype: torch.dtype = torch.bfloat16, log=None, post=None) -> int:
+    """把（可能被 ZeRO-2 / FSDP 包着的）模型导出成 ActionRoPE 格式的 safetensors，返回字节数。
+
+    每个 rank 都要调用（get_state_dict 是集合操作），只有主进程落盘。
+
+    ★ 显式 cast 到 bf16 并断言导出 dtype。FSDP 下 accelerator.get_state_dict 返回的是 rank0 CPU 上的
+    fp32 全量权重（stage3 在 prepare 前 .float() 过每个模型），不 cast 文件就是 20 GB，827-key /
+    10.0 GB 的格式承诺当场作废，infer.py / check_ckpt.py 的 dtype 断言也会跟着变。
+
+    expect_keys / require_mask 是格式闸门：默认按 827 键 + arope_mask_embedding.{weight,bias} 校，
+    对不上直接炸（少两个 mask 键 = 模型是拿 mask_channel=False 建的，装回 infer.py 会静默走错分支）。
+    小模型自测传 expect_keys=<n>, require_mask=False。"""
+    if accel is not None:
+        accel.wait_for_everyone()
+    sd = _gather_state_dict(accel, model)
+    # post：gather 之后、格式闸门之前的一次改写。目前只有 LoRA 用它把 B@A 折回 base.weight——
+    # 不折的话 state_dict 里是 `...base.weight` + `lora_A` + `lora_B`，expect_keys 当场拦下。
+    # 每个 rank 都会走到这里，但只有主进程的 sd 非空（见 _gather_state_dict）。
+    if post is not None and sd:
+        sd = post(sd)
+    is_main = True if accel is None else bool(accel.is_main_process)
+    size = 0
+    if is_main and sd:
+        out, dropped = strip_dit_prefix(sd)
+        out = {k: v.detach().to(device="cpu", dtype=dtype).contiguous() for k, v in out.items()}
+        bad = sorted({str(v.dtype) for v in out.values()} - {str(dtype)})
+        if bad:
+            raise RuntimeError(f"导出 dtype 不是 {dtype}：{bad}")
+        if expect_keys is not None and len(out) != expect_keys:
+            raise RuntimeError(f"{path}: 导出 {len(out)} 个张量，期望 {expect_keys} 个"
+                               f"（丢弃的非 DiT 键 {dropped[:5]}…{len(dropped)}）")
+        if require_mask:
+            miss = [k for k in MASK_KEYS if k not in out]
+            if miss:
+                raise RuntimeError(f"{path}: 缺 {miss} —— 模型是拿 mask_channel=False 建的？")
+        meta = {"arm": "arope", "mask_channel": "True"}
+        meta.update({k: str(v) for k, v in (metadata or {}).items()})
+        meta.update({"n_tensors": str(len(out)), "dtype": str(dtype).replace("torch.", ""),
+                     "saved": time.strftime("%Y-%m-%d %H:%M:%S")})
+        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+        tmp = f"{path}.tmp"                       # 先写临时文件再 rename：写到一半崩掉不会留下半个 ckpt
+        save_file(out, tmp, metadata=meta)
+        os.chmod(tmp, 0o644)                      # safetensors 按 umask 写成 0600，推理可能由别的账号跑
+        os.replace(tmp, path)
+        size = os.path.getsize(path)
+        msg = f"[save] {path}  {len(out)} 个张量  {size / 1e9:.2f} GB"
+        if log is not None:
+            log(msg)
+        else:
+            print(msg, flush=True)
+        del out
+    del sd
+    if accel is not None:
+        accel.wait_for_everyone()
+    return size
+
+
+# ----------------------------------------------------------------------------
+# checkpoint：装载 / 体检
+# ----------------------------------------------------------------------------
+
+def load_dit_ckpt(model, path: str, *, assign: bool = False, allow_missing_mask: bool = True,
+                  allow_missing: tuple | list = ()) -> dict:
+    """把 ActionRoPE 格式的 safetensors 装进 AropeDiT（键带 `dit.`）或裸 WanModel（键不带）。
+
+    strict —— 默认唯一放行的缺失是 arope_mask_embedding.{weight,bias}（install_arope 已经零初始化，
+    从没装过 mask 通道的 ckpt 起步时就缺这两个键）。WorldMind 的 load_dit_ckpt 对任何 missing 都抛，
+    照抄会拒收每一个 ActionRoPE checkpoint。
+
+    allow_missing：额外放行的键（去掉 `dit.` 前缀的名字，也接受带前缀的写法）。**唯一** 的正当用法是
+    "先 install_flowmap() 再从 827 键的 stage1b 装载"这一个调用点，传 `allow_missing=FLOWMAP_KEYS`；
+    别把它设成默认，否则任何一个真的漏键都会静默通过。推荐的顺序反过来：先 strict 装 827、再 install
+    （见 AropeDiT.install_flowmap 的文档），那样一个例外都不需要。
+
+    assign=False（默认）：copy_ 进现有参数，保持模型自己的 dtype —— stage3 在 .float() 之后装 bf16
+    ckpt 就靠这个，assign=True 会把参数直接换成 bf16 张量，fp32 主权重当场消失。"""
+    sd = load_file(path)
+    tgt = _unwrap(model)
+    tgt_keys = set(tgt.state_dict().keys())
+    prefix = DIT_PREFIX if any(k.startswith(DIT_PREFIX) for k in tgt_keys) else ""
+    if prefix:
+        sd = {(k if k.startswith(prefix) else prefix + k): v for k, v in sd.items()}
+    missing, unexpected = tgt.load_state_dict(sd, strict=False, assign=assign)
+    allowed = {prefix + k for k in MASK_KEYS} if allow_missing_mask else set()
+    allowed |= {(k if prefix and k.startswith(prefix) else prefix + k) for k in allow_missing}
+    bad = [k for k in missing if k not in allowed]
+    if bad or unexpected:
+        raise RuntimeError(f"{path} 与模型键集不符：缺 {bad[:5]}…({len(bad)})，多 {list(unexpected)[:5]}…({len(unexpected)})")
+    return {"path": path, "n_tensors": len(sd), "prefix": prefix,
+            "missing": list(missing), "unexpected": list(unexpected)}
+
+
+def read_ckpt_header(path: str) -> tuple[dict, dict]:
+    """只读 safetensors 头（不碰 10 GB 数据）→ ({key: {dtype, shape, data_offsets}}, metadata)。"""
+    with open(path, "rb") as fh:
+        n = struct.unpack("<Q", fh.read(8))[0]
+        hdr = json.loads(fh.read(n))
+    meta = hdr.pop("__metadata__", {}) or {}
+    return hdr, meta
+
+
+def is_flowmap_ckpt(path: str) -> bool:
+    """这份 ckpt 是不是 flow map（832 键）。**看张量键，不看 metadata** —— metadata 是人写的字符串，
+    张量在不在是事实。下游（rollout_eval）应当据此选采样器：拿 renoise 采样器去采 flow map 不会报错，
+    只会安静地产出另一个模型的结果。"""
+    hdr, _ = read_ckpt_header(path)
+    return all(k in hdr for k in FLOWMAP_KEYS)
+
+
+def check_ckpt_format(path: str, expect_keys: int | None = AROPE_DIT_KEYS,
+                      dtype: str = "BF16", require_mask: bool = True) -> dict:
+    """导出后的秒级体检（不加载权重）：张量数 / dtype / mask 键 / metadata 全是字符串。
+
+    正品参照：outputs/arope_10k_ga2/step-10000.safetensors —— 827 个 BF16 张量，9,999,690,792 B。"""
+    hdr, meta = read_ckpt_header(path)
+    dts = sorted({v["dtype"] for v in hdr.values()})
+    res = {"path": path, "n_tensors": len(hdr), "dtypes": dts, "size": os.path.getsize(path),
+           "metadata": meta, "mask_keys": sorted(k for k in hdr if k.startswith(MASK_EMBEDDING_NAME))}
+    if expect_keys is not None and len(hdr) != expect_keys:
+        raise RuntimeError(f"{path}: {len(hdr)} 个张量，期望 {expect_keys}")
+    if dtype is not None and dts != [dtype]:
+        raise RuntimeError(f"{path}: dtype {dts}，期望 ['{dtype}']")
+    if require_mask and set(res["mask_keys"]) != set(MASK_KEYS):
+        raise RuntimeError(f"{path}: mask 键 {res['mask_keys']}，期望 {sorted(MASK_KEYS)}")
+    bad = {k: type(v).__name__ for k, v in meta.items() if not isinstance(v, str)}
+    if bad:
+        raise RuntimeError(f"{path}: metadata 必须全是字符串，{bad}")
+    return res
+
+
+# ----------------------------------------------------------------------------
+# accelerate 完整存档（只有 stage1 用）
+# ----------------------------------------------------------------------------
+
+STATE_SLOTS = ("state_A", "state_B")
+
+
+def save_state_rotating(accel, run_dir: str, trainer_state: dict) -> str:
+    """accelerator.save_state 到两个槽之一，<run_dir>/state 软链指向最新完整的那个。
+
+    原地覆盖唯一的 state/ 目录时，保存中途崩溃会把唯一的续训点也毁掉；这里总是写到「不是当前
+    state 指向」的那个槽，写完再 rename 换软链，任一时刻 state 都指向一份完整存档。不删任何目录，
+    磁盘占两份（4 卡 ZeRO-2 约 2×60 GB）。--resume <run_dir>/state 照常。"""
+    link = os.path.join(run_dir, "state")
+    cur = os.path.basename(os.path.realpath(link)) if os.path.islink(link) else None
+    slot = STATE_SLOTS[1] if cur == STATE_SLOTS[0] else STATE_SLOTS[0]
+    state_dir = os.path.join(run_dir, slot)
+    accel.save_state(state_dir)
+    if accel.is_main_process:
+        with open(os.path.join(state_dir, "trainer_state.json"), "w", encoding="utf-8") as fh:
+            json.dump(trainer_state, fh, ensure_ascii=False, default=str)
+        tmp = link + ".tmp"
+        if os.path.lexists(tmp):
+            os.remove(tmp)
+        os.symlink(slot, tmp)
+        os.replace(tmp, link)                     # 只替换软链本身，不动任何目录
+        print(f"[save] accelerator.save_state → {state_dir}（{link} → {slot}）", flush=True)
+    accel.wait_for_everyone()
+    return state_dir
+
+
+# ----------------------------------------------------------------------------
+# 日志
+# ----------------------------------------------------------------------------
+
+class TextTee:
+    """打到 stdout，同时带时间戳追加进 <run_dir>/train.log。
+
+    enabled=False（非主进程）⇒ 完全静默，也不开文件句柄。"""
+
+    def __init__(self, path: str, enabled: bool = True):
+        self.path, self.enabled, self._f = path, bool(enabled), None
+        if self.enabled:
+            os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+            self._f = open(path, "a", encoding="utf-8", buffering=1)
+
+    def __call__(self, *args):
+        if not self.enabled:
+            return
+        msg = " ".join(str(a) for a in args)
+        print(msg, flush=True)
+        self._f.write(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}\n")
+        self._f.flush()
+
+    def close(self):
+        if self._f is not None:
+            try:
+                self._f.close()
+            finally:
+                self._f = None
+
+
+class CsvLog:
+    """固定表头的追加式 CSV（actionrope/train.py:328 的写法）。已存在且非空就不再写表头。"""
+
+    def __init__(self, path: str, fields: list[str], enabled: bool = True):
+        self.enabled, self.f, self.w = bool(enabled), None, None
+        if not self.enabled:
+            return
+        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+        exists = os.path.isfile(path) and os.path.getsize(path) > 0
+        self.f = open(path, "a", encoding="utf-8", newline="")
+        self.w = csv.DictWriter(self.f, fieldnames=fields, lineterminator="\n", extrasaction="ignore")
+        if not exists:
+            self.w.writeheader()
+            self.f.flush()
+
+    def write(self, row: dict):
+        if not self.enabled:
+            return
+        self.w.writerow(row)
+        self.f.flush()
+
+    def close(self):
+        if self.f is not None:
+            try:
+                self.f.close()
+            finally:
+                self.f = None
+
+
+class JsonlLogger:
+    """一行一个 JSON 的追加式指标日志，每行都 flush，别的机器可以直接 tail。"""
+
+    def __init__(self, path: str, enabled: bool = True):
+        self.path, self.enabled, self._f = path, bool(enabled), None
+        if self.enabled:
+            os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+            self._f = open(path, "a", encoding="utf-8", buffering=1)
+
+    def log(self, row: dict):
+        if not self.enabled:
+            return
+        row = dict(row)
+        row.setdefault("wall", time.time())
+        self._f.write(json.dumps(row, ensure_ascii=False, default=str) + "\n")
+        self._f.flush()
+
+    def close(self):
+        if self._f is not None:
+            try:
+                self._f.close()
+            finally:
+                self._f = None

backup_2026-09-14/code/distill/stage1_ar/ar_tf.py

@@ -0,0 +1,695 @@
+#!/usr/bin/env python3
+"""stage1：teacher forcing 因果化 —— 把 ActionRoPE 的双向 SFT 检查点训成因果 AR 扩散模型。
+
+    # 4 卡正式训练（配方常量钉在 run.sh 里，argparse 默认值与它逐项相同）
+    bash distill/stage1_ar/run.sh
+
+    # 单卡 3 步冒烟（不走 accelerate；AdamW 直接更新 bf16 参数、没有 fp32 主权重，loss 数值不代表训练效果）
+    CUDA_VISIBLE_DEVICES=0 PYTHONPATH=/data/zhiyangdeng/ActionRoPE DIFFSYNTH_SKIP_DOWNLOAD=True \\
+      .venv/bin/python distill/stage1_ar/ar_tf.py \\
+      --init_ckpt outputs/arope_10k_ga2/step-10000.safetensors \\
+      --limit 8 --max_steps 3 --grad_accum 1 --val_every 0 --save_every 0 --log_every 1 \\
+      --output /data/zhiyangdeng/ActionRoPE/distill/_smoke_stage1
+
+    # 30 秒 dry run（tiny 模型 + 合成数据，不碰真权重；验循环 / 导出 / 日志 / 续训）
+    ... --dry_run --output <tmp>
+
+这一阶段是什么
+-------------
+outputs/arope_10k_ga2/step-10000.safetensors 是**双向**全片 SFT 模型：每个 cell 都能看见未来。
+AR rollout 要的是因果模型，所以先把它「因果化」：同一份 flow-matching 目标、同一份 new 区 ×2 权重图，
+只把自注意力换成 teacher forcing 双流掩码 —— noisy cell f 只看**干净的** 0..f-1 加自己那一格。
+产物 dits/dit_step{N}.safetensors 同时做 stage2 的 --teacher_ckpt 和 --student_ckpt。
+
+目标（= PORT_SPEC stage1，`ar_diffusion_gm.py:83-104` 融合 `actionrope/train.py:224-271`）
+-----------------------------------------------------------------------------------
+每个 cell 独立均匀抽噪声档（diffusion forcing 的加噪方式），干净历史走 clean_x 通道而不是噪声流；
+loss 是 ActionRoPE 的 Σw·se/(Σw·C) **逐 cell** 算完，再乘逐 cell 的 bsmntw 钟形权重取平均。
+
+相对 WorldMind 的三处刻意改动（每处都有代价明确的理由）
+-----------------------------------------------------
+1. **cell 0**。WorldMind 的 TF 把 cell 0 也当无条件生成来训（loss 覆盖全部 F 帧）。ActionRoPE 的契约是
+   cell 0 = 干净锚点，采样时每一步都写回（infer.py:413）。照抄会改变「这个模型是什么」，所以这里
+   cell 0 永远是干净 x0 + t=0 且**不计 loss** —— 与 stage2 的 `lat_t[:,:,0]=x0[:,:,0]`、stage3 的
+   `ts_pf[:,0]=0` 一路对齐。
+2. **loss 形状**。`se.mean(dim=(1,3,4))`（`ar_diffusion_gm.py:102`）换成 ActionRoPE 的
+   Σw·se/(Σw·C)，逐 cell 算 —— 10k 检查点就是在 new 区 ×2 的权重下训出来的，丢掉它等于换了目标。
+3. **条件**。offset_px（世界坐标 RoPE）/ mask_input（known 通道）/ context_ids（逐 cell 文本）三路条件
+   穿进同一次前向，并在 AropeDiT 内部同时喂给 TF 的**两个半流**。
+
+两个必须守住的坑
+---------------
+* `--tf_noise_aug_max_idx` 上游取的是**错的那一端**（`ar_diffusion_gm.py:93` 的 `randint(0, k)` 是
+  σ≈1 的纯噪声区，不是近干净区）。这里实现成正确的 `randint(1000-k, 1000)`，默认仍是 0（精确干净）。
+* `--grad_accum` 必须与 accelerate yaml 的 gradient_accumulation_steps 一致，否则等效 batch 不是 8
+  —— 沿用 `actionrope/train.py:399` 的 SystemExit 守卫。
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import platform
+import socket
+import sys
+import time
+import warnings
+import zlib
+
+os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
+# 分配器要在 import torch 之前设：launch 时 export 的环境变量到不了 accelerate 拉起的子进程（C 的 train.py:22-27）
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import torch
+from accelerate import Accelerator
+from accelerate.utils import DistributedType, set_seed
+
+ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from actionrope.dataset import DEFAULT_SIDECAR, TEXT_TABLES, AropeLatentDataset, arope_collate, load_sidecar  # noqa: E402
+from distill import runlib  # noqa: E402
+from distill.arope_dit import MODEL_DIR, AropeDiT, unwrap  # noqa: E402
+from distill.data import make_mask_input, make_weight_map, to_device  # noqa: E402
+from distill.flow_pf import (  # noqa: E402
+    NUM_TRAIN_TIMESTEPS,
+    add_noise_pf,
+    idx_of_t,
+    sample_pf_timestep_ids,
+    t_of_idx,
+    training_target,
+    training_weight_pf,
+)
+
+CELLS = 21          # 一条 clip 的 latent cell 数（数据集硬契约）
+STAGE = "stage1_ar_tf"
+
+
+# ----------------------------------------------------------------------------
+# 目标
+# ----------------------------------------------------------------------------
+
+def tf_clean_stream(x0: torch.Tensor, noise: torch.Tensor, max_idx: int):
+    """TF 的 clean 半流。max_idx<=0 ⇒ (x0, None) 精确干净（CF 默认）。
+
+    max_idx>0 ⇒ CF 的 history-noise 增广：给干��历史加一点点噪声，复用**同一份** noise。
+    近干净 = **高**下标（sigmas 递减，idx 999 的 σ≈0.005），所以抽 [1000-max_idx, 1000) ——
+    上游 `ar_diffusion_gm.py:93` 写的是 `randint(0, max_idx)`，那是 σ≈1 的纯噪声端，方向反了。
+    cell 0 无论如何保持精确干净 + t=0：它是 ActionRoPE 的 TI2V 锚点，两个半流都得一致。
+    """
+    if max_idx is None or int(max_idx) <= 0:
+        return x0, None
+    k = int(max_idx)
+    assert 0 < k < NUM_TRAIN_TIMESTEPS, f"--tf_noise_aug_max_idx 应在 (0, {NUM_TRAIN_TIMESTEPS}) 内，得到 {k}"
+    b, f = x0.shape[0], x0.shape[2]
+    idx_aug = torch.randint(NUM_TRAIN_TIMESTEPS - k, NUM_TRAIN_TIMESTEPS, (b, f), device=x0.device)
+    clean_x = add_noise_pf(x0, noise, idx_aug)
+    aug_t = t_of_idx(idx_aug, device=x0.device)
+    clean_x[:, :, 0] = x0[:, :, 0]
+    aug_t[:, 0] = 0.0
+    return clean_x, aug_t
+
+
+def ar_tf_loss(student, batch: dict, *, new_weight: float = 2.0, new_weight_ramp: float = 0.0,
+               teacher_forcing: bool = True, tf_noise_aug_max_idx: int = 0,
+               use_gradient_checkpointing: bool = False,
+               action_dropout: float = 0.0, hist_dropout: float = 0.0, ctx_dropout: float = 0.0,
+               anchor_dropout: float = 0.0,
+               noise: torch.Tensor | None = None, idx_pf: torch.Tensor | None = None,
+               gen: torch.Generator | None = None):
+    """一次 micro batch 的 loss。返回 (loss, stats)。
+
+    teacher_forcing=True（默认）：干净历史走 clean_x 通道，[2S,2S] 双流掩码，S=21*390=8190。
+    teacher_forcing=False：diffusion forcing 回退路径 —— 单条 8190 token 流、普通块因果掩码，
+    便宜约 2.5×（代价：stage2 也必须 --no-teacher_forcing，见 PORT_SPEC）。
+
+    条件 dropout（三个默认都是 0 ⇒ stage1 的行为一个字节都不变）
+    -------------------------------------------------------
+    只有**训练时 drop 过**的条件，才谈得上在蒸馏时对它做 classifier-free guidance：
+    否则 v(∅) 是一个没训过的输入上的外推，`v_c − v_∅` 是噪声不是信号。ActionRoPE 的 SFT
+    只有 `--scene_dropout`（文本）一条，所以 text 之外的轴都得先补出来。这三个开关就是补它们：
+
+      action_dropout  offset_px ≡ 0、掩码全 known、权重图退成均匀 ⇒ 自洽的"静止相机"输入
+                      （语料里 22782 条 clip **没有一条**全程静止，所以这个 null 必须现训）
+      anchor_dropout  只把 clean cell 0 从每个 noisy cell 的视野里摘掉，**其余历史保留**
+      hist_dropout    注意力上下文砍成 **只剩 cell 0**（local=1 sink=1）
+      ctx_dropout     注意力上下文砍成 **什么都看不到**（local=1 sink=0）
+
+    四条 guidance（★ 是实测最该生效的那条）：
+      ★ Δv_anchor = v(全) − v(drop_anchor)          锚点**相对于历史**的话语权
+        Δv_action = v(全) − v(offset≡0)             相机运动的贡献
+        Δv_hist   = v(全) − v(只剩 cell 0)           历史的贡献
+        Δv_ref    = v(只剩 cell 0) − v(无上下文)      锚点在真空里的贡献
+
+    为什么 anchor 那条才是主角：实测漂移是**漂掉的近期历史压过了首帧锚点**（AR rollout 里
+    |x0| 从 1.111 涨到 1.404、外观洗白），要压回去就得放大"锚点 vs 历史"这一项。
+    Δv_ref 量的是"真空里锚点值多少"，与它施加的场景（历史俱全）不是同一个 regime。
+    两者都留着：anchor 是主力，ref/hist 这对嵌套档给消融用。
+
+    注意 cell 0 的 TI2V 锚点（干净 latent + t=0）在**输入侧**始终保留，被 drop 的只是"别的
+    cell 能不能看到它"。这样 loss 切片（cell 0 不计）和整个 TI2V 结构都不用动。
+    """
+    x0 = batch["input_latents"]                                   # [b, 48, 21, 30, 52] bf16
+    b, c, f, lat_h, lat_w = x0.shape                              # 几何从 latent 里取，dry_run 的小分辨率也能过
+    dev = x0.device
+    off = batch["offset_px"]                                      # [b, 21, 2] f32
+
+    if noise is None:
+        noise = torch.randn_like(x0)
+    if idx_pf is None:
+        idx_pf = sample_pf_timestep_ids(b, f, dev, gen)            # [b, 21] 严格均匀，逐 cell 独立
+    noisy = add_noise_pf(x0, noise, idx_pf)                        # x_t=(1−σ)x0+σε，σ 按 [b,1,21,1,1] 广播
+    ts_pf = t_of_idx(idx_pf, device=dev)                           # [b, 21] f32，(4.98, 1000]
+    # TI2V 锚点：cell 0 用干净 x0 + t=0。模型侧 first_frame_cond=True 也会把 cell 0 的 t 抹成 0，
+    # 这里显式写一遍是为了让 stats 里的 t 与真实喂进去的一致，且 --no-teacher_forcing 回退路径同样成立。
+    noisy[:, :, 0] = x0[:, :, 0]
+    ts_pf[:, 0] = 0.0
+
+    # ---- 条件 dropout：逐 micro batch 抽签（batch_size 恒为 1，所以逐 batch == 逐样本）----
+    def _coin(p):
+        if p <= 0.0:
+            return False
+        return bool(torch.rand((), generator=gen, device="cpu").item() < p)
+
+    drop_act = _coin(action_dropout)
+    # 三个上下文档互斥，按"砍得越狠越优先"排：ctx（全空）> hist（只剩 cell0）> anchor（只摘 cell0）
+    drop_ctx = _coin(ctx_dropout)
+    drop_hist = (not drop_ctx) and _coin(hist_dropout)
+    drop_anchor = (not drop_ctx) and (not drop_hist) and _coin(anchor_dropout)
+    if drop_act:
+        off = torch.zeros_like(off)                   # 静止相机：世界 RoPE 退回普通屏幕格 RoPE
+
+    mask_in = make_mask_input(off, lat_h, lat_w)                   # [b, 1, 21, 30, 52] known=1/new=0
+    w = make_weight_map(off, lat_h, lat_w, new_weight=new_weight, ramp_latent=new_weight_ramp)
+    clean_x, aug_t = (tf_clean_stream(x0, noise, tf_noise_aug_max_idx) if teacher_forcing else (None, None))
+
+    # 上下文档只改**掩码窗口**，不动 cell 0 的锚点。掩码按 (f, hw, tf, 窗口, device) 缓存，
+    # 三档各自命中自己的 key，不会串。改的是裸模块：DeepSpeed 包装之后 self.xxx 取不到。
+    # 只在真要 drop 时才碰这两个属性 —— 测试里的桩模型没有它们，默认路径不该要求它们存在。
+    bare = unwrap(student) if (drop_ctx or drop_hist or drop_anchor) else None
+    saved = ((bare.train_local_attn_size, bare.train_sink_size, bare.train_drop_anchor)
+             if bare is not None else None)
+    if drop_ctx:
+        bare.train_local_attn_size, bare.train_sink_size = 1, 0    # 谁都看不到
+    elif drop_hist:
+        bare.train_local_attn_size, bare.train_sink_size = 1, 1    # 只看得到 cell 0
+    elif drop_anchor:
+        bare.train_drop_anchor = True                              # 只摘掉 cell 0，其余历史照旧
+    try:
+        with torch.autocast(device_type=dev.type, dtype=torch.bfloat16, enabled=(dev.type == "cuda")):
+            v = student(noisy, ts_pf, batch["context"], context_ids=batch["context_ids"],
+                        offset_px=off, mask_input=mask_in, first_frame_cond=True,
+                        clean_x=clean_x, aug_t=aug_t,
+                        use_gradient_checkpointing=use_gradient_checkpointing)
+    finally:
+        if bare is not None:
+            bare.train_local_attn_size, bare.train_sink_size, bare.train_drop_anchor = saved
+
+    target = training_target(x0, noise)                            # ε − x0（rectified flow）
+    se = (v.float() - target.float()) ** 2
+    se, ww = se[:, :, 1:], w[:, :, 1:]                             # cell 0 是条件帧，不计 loss
+    # ActionRoPE 的 Σw·se/(Σw·C)，但**逐 cell** 算：这样 new 区 ×2 的权重和逐 cell 的钟形权重能同时生效
+    per_cell = (se * ww).sum(dim=(1, 3, 4)) / (ww.sum(dim=(1, 3, 4)) * c)      # [b, 20]
+    tw = training_weight_pf(idx_pf[:, 1:]).float()                             # [b, 20] bsmntw 钟形
+    loss = (tw * per_cell).mean()
+    stats = {
+        "loss_w": per_cell.mean().detach(),          # 不乘钟形权重的加权 MSE，与 train.py 的 loss_w 同口径
+        "mse": se.mean().detach(),                   # 完全不加权的 MSE，跨配方可比
+        "new_frac": (ww > 1).float().mean().detach(),
+        "t_mean": ts_pf[:, 1:].mean().detach(),
+        "tw_mean": tw.mean().detach(),
+        # 三档 dropout 的命中率：日志里能看到实际抽签比例是不是等于设定的 p（各 rank 独立抽签，
+        # 所以单 rank 的短期比例会抖，看长期均值）
+        "d_act": torch.tensor(float(drop_act), device=dev),
+        "d_hist": torch.tensor(float(drop_hist), device=dev),
+        "d_ctx": torch.tensor(float(drop_ctx), device=dev),
+        "d_anc": torch.tensor(float(drop_anchor), device=dev),
+    }
+    return loss, stats
+
+
+# ----------------------------------------------------------------------------
+# 参数
+# ----------------------------------------------------------------------------
+
+def parse_args(argv=None):
+    ap = argparse.ArgumentParser(description="stage1：teacher forcing 因果化 ActionRoPE SFT 检查点")
+    # 默认值 = distill/stage1_ar/run.sh 的实参（PORT_SPEC risk 11：argparse 与 run.sh 不许有第二套配方）
+    ap.add_argument("--init_ckpt", default=runlib.REF_CKPT,
+                    help="要因果化的双向 SFT 检查点（827 键 bf16）")
+    ap.add_argument("--model_dir", default=MODEL_DIR)
+    ap.add_argument("--output", required=True, help="run 目录；相对路径按 runlib.resolve_run_dir 解析")
+
+    ap.add_argument("--train_dir", default=os.path.join(ROOT, "data/latent/train_eybx"))
+    ap.add_argument("--val_dir", default=os.path.join(ROOT, "data/latent/val_eybx"))
+    ap.add_argument("--sidecar", default=DEFAULT_SIDECAR)
+    ap.add_argument("--text_table", default=TEXT_TABLES["scene"],
+                    help="ARoPE 臂固定 text_mode=scene（剥掉动作从句），361 键那张表")
+    ap.add_argument("--limit", type=int, default=None, help="smoke 用：只取这么多训练 clip")
+    ap.add_argument("--long_seg_dir", default=None,
+                    help="长视野续训（longhorizon/README.md）：给 make_long_segments.py 的段目录"
+                         "（如 data/latent_long/K3）就用 AropeLongDatasetTF 替换训练集（F=61 等），"
+                         "其余配方一字不动；val 仍是 --val_dir 的 21-cell clip。None = 原 stage1 行为")
+    ap.add_argument("--scene_dropout", type=float, default=0.1)
+    ap.add_argument("--min_rt_ratio", type=float, default=0.8)
+    ap.add_argument("--val_min_rt_ratio", type=float, default=0.9)
+    ap.add_argument("--val_frac", type=float, default=0.0,
+                    help="兼容 WorldMind 的 hash 划分，只接受 0：ActionRoPE 的划分是离线的 "
+                         "train_eybx / val_eybx 清单，而且相邻 clip 只差 3 帧，按 clip 名 hash 会漏样本")
+
+    ap.add_argument("--teacher_forcing", action=argparse.BooleanOptionalAction, default=True,
+                    help="默认开：干净历史走 clean_x 通道 + [2S,2S] 双流掩码。"
+                         "--no-teacher_forcing 回退成 diffusion forcing（单流、便宜 2.5×，"
+                         "但 stage2 必须跟着一起关）")
+    ap.add_argument("--tf_noise_aug_max_idx", type=int, default=0,
+                    help="clean 半流的 history-noise 增广宽度（近干净档 = 高下标）。0 = 精确干净（CF 默认）")
+    # ---- 条件 dropout（stage1b 适配微调用；默认全 0 ⇒ stage1 行为不变）----
+    ap.add_argument("--action_dropout", type=float, default=0.0,
+                    help="以此概率把 offset_px 清零 + 掩码全 known ⇒ 造出「静止相机」这个 null 分支。"
+                         "语料里没有一条全程静止的 clip，不训就没有，蒸馏时 Δv_action 无从定义")
+    ap.add_argument("--anchor_dropout", type=float, default=0.0,
+                    help="★ 主力档：只把 clean cell 0 从每个 noisy cell 的视野里摘掉、其余历史保留 ⇒ "
+                         "Δv_anchor 的 null。这是「锚点 vs 已漂历史」那一项，实测的外观漂移要压的就是它")
+    ap.add_argument("--hist_dropout", type=float, default=0.0,
+                    help="以此概率把注意力上下文砍成**只剩 cell 0**（local=1 sink=1）⇒ Δv_hist 的 null")
+    ap.add_argument("--ctx_dropout", type=float, default=0.0,
+                    help="以此概率把注意力上下文砍成**什么都看不到**（local=1 sink=0）⇒ Δv_ref 的 null。"
+                         "与 --hist_dropout 互斥（本档优先），cell 0 的 TI2V 锚点始终保留")
+    ap.add_argument("--windowed_train", type=int, default=0,
+                    help="训练掩码的滑动窗口（cell 数）。0 = 全因果 —— F=21 下 sink=0 的 W=20 会把 "
+                         "cell 0（干净锚点）挤出 cell 20 的视野，是致命的")
+    ap.add_argument("--windowed_sink", type=int, default=0, help="[windowed_train>0 时] 恒可见的开头 cell 数")
+    ap.add_argument("--new_weight", type=float, default=2.0, help="new 区 loss 权重（10k 检查点就是这么训的）")
+    ap.add_argument("--new_weight_ramp", type=float, default=0.0,
+                    help=">0 ⇒ new 区权重按离首帧足迹的距离爬升（dist 形状）；0 = flat，与 10k 检查点一致")
+
+    ap.add_argument("--lr", type=float, default=1e-5)
+    ap.add_argument("--weight_decay", type=float, default=0.01)
+    ap.add_argument("--adam_betas", default="0.9,0.999",
+                    help="ActionRoPE SFT 的优化器口径。WorldMind 那边是 (0.0, 0.999)（CF 约定），"
+                         "但 stage1 是同一份 SFT 目标加了因果掩码，用户拍板沿用 SFT 的 betas")
+    ap.add_argument("--warmup", "--warmup_steps", dest="warmup", type=int, default=200,
+                    help="线性 warmup 步数，之后恒定 lr（每个 save_every 边界都是合法的停止点）")
+    ap.add_argument("--grad_clip", type=float, default=1.0,
+                    help="DeepSpeed 下写进 deepspeed_config.gradient_clipping（覆盖 yaml），日志里的 grad_norm 是裁剪前的")
+    ap.add_argument("--batch_size", type=int, default=1, help="每卡 micro batch；等效 batch = world × 它 × grad_accum")
+    ap.add_argument("--grad_accum", type=int, default=2, help="必须与 accelerate yaml 的 gradient_accumulation_steps 相同")
+    ap.add_argument("--max_steps", type=int, default=12000, help="优化器更新数（不是 micro batch 数）")
+    ap.add_argument("--gradient_checkpointing", action=argparse.BooleanOptionalAction, default=True)
+    ap.add_argument("--num_workers", type=int, default=2)
+    ap.add_argument("--seed", type=int, default=0)
+
+    ap.add_argument("--save_every", type=int, default=500, help="0 ⇒ 只在训练结束时导出一次")
+    ap.add_argument("--keep_last_dits", type=int, default=6, help="dits/ 只留最新几个（0 = 不删；上游 v4 光 dits/ 就 746 GB）")
+    ap.add_argument("--save_state", action="store_true", help="同时 accelerator.save_state（state_A/state_B 轮换，可续训）")
+    ap.add_argument("--val_every", type=int, default=500, help="0 ⇒ 不验证")
+    ap.add_argument("--val_n", type=int, default=16)
+    ap.add_argument("--val_timesteps", default="200,500,800")
+    ap.add_argument("--log_every", type=int, default=20)
+    ap.add_argument("--resume", default=None,
+                    help="*.safetensors ⇒ 只热启动权重（step 归 0）；目录 ⇒ accelerator.load_state 续训；"
+                         "auto ⇒ 有 <run>/state 就续，没有就从头")
+    ap.add_argument("--dry_run", action="store_true", help="tiny 模型 + 合成数据，不碰真权重 / 真数据")
+    args = ap.parse_args(argv)
+
+    if abs(args.val_frac) > 1e-12:
+        raise SystemExit("--val_frac 只接受 0：ActionRoPE 用离线的 train_eybx / val_eybx 清单划分，"
+                         "相邻 clip 只差 3 帧，按 clip 名 hash 划会把同一段画面同时放进两边。改用 --val_dir。")
+    if args.windowed_sink > 0 and args.windowed_train <= 0:
+        raise SystemExit("--windowed_sink 只在 --windowed_train > 0 时有意义")
+    try:
+        b1, b2 = (float(x) for x in args.adam_betas.split(","))
+    except ValueError:
+        raise SystemExit(f"--adam_betas 应形如 '0.9,0.999'，得到 {args.adam_betas!r}") from None
+    args.betas = (b1, b2)
+    args.val_timesteps = [float(t) for t in args.val_timesteps.split(",") if t.strip()]
+    if args.dry_run:
+        args.init_ckpt = None
+    elif not os.path.isfile(args.init_ckpt):
+        raise SystemExit(f"--init_ckpt 不存在：{args.init_ckpt}")
+    return args
+
+
+# ----------------------------------------------------------------------------
+# dry run：tiny 模型 + 合成数据（PORT_SPEC 的 S6）
+# ----------------------------------------------------------------------------
+
+TINY = dict(dim=128, in_dim=48, ffn_dim=256, out_dim=48, text_dim=4096, freq_dim=256, eps=1e-6,
+            patch_size=(1, 2, 2), num_heads=4, num_layers=2, has_image_input=False)
+TINY_SHAPE = (48, 5, 8, 10)          # [C, cells, lat_h, lat_w]
+
+
+class _SynDataset(torch.utils.data.Dataset):
+    """dry run 用：形状与 AropeLatentDataset 一致的合成样本（cell 数 / 分辨率缩小）。"""
+
+    def __len__(self):
+        return 64
+
+    def __getitem__(self, i: int) -> dict:
+        g = torch.Generator().manual_seed(i)
+        c, f, h, w = TINY_SHAPE
+        off = torch.zeros(f, 2)
+        off[1:] = torch.randn(f - 1, 2, generator=g).cumsum(0) * 20.0
+        return {
+            "input_latents": torch.randn(c, f, h, w, generator=g, dtype=torch.float32).to(torch.bfloat16),
+            "context": torch.randn(f, 8, 4096, generator=g, dtype=torch.float32).to(torch.bfloat16),
+            "context_ids": torch.zeros(f, dtype=torch.int64),
+            "offset_px": off,
+            "clip_id": f"syn_{i:04d}",
+            "kind": "syn",
+        }
+
+
+def _tiny_model() -> AropeDiT:
+    from diffsynth.models.wan_video_dit import WanModel
+    torch.manual_seed(0)
+    dit = WanModel(**TINY).to(torch.bfloat16)
+    return AropeDiT(dit=dit, mask_channel=True, causal=True)
+
+
+# ----------------------------------------------------------------------------
+# 验证
+# ----------------------------------------------------------------------------
+
+@torch.no_grad()
+def run_validation(accelerator: Accelerator, model, val_ds, tids: list[int], args) -> dict:
+    """固定 clip × 固定噪声档 × 固定噪声，各进程分摊 clip，最后 reduce（集合操作，人人都要到）。
+
+    与训练不同：一次验证把 **所有** cell 都放在同一个噪声档 tid 上（cell 0 依旧干净 + t=0），
+    这样曲线只随模型变、不随抽签抖，口径与 `actionrope/train.py:352-386` 一致。
+    """
+    model.eval()
+    device = accelerator.device
+    rank, world = accelerator.process_index, accelerator.num_processes
+    acc = torch.zeros(len(tids), 3, device=device, dtype=torch.float64)   # loss_w 和 / mse 和 / 计数
+    for i in range(rank, len(val_ds), world):
+        sample = val_ds[i]
+        batch = to_device(arope_collate([sample]), device)
+        x0 = batch["input_latents"]
+        for j, tid in enumerate(tids):
+            g = torch.Generator(device=device)
+            g.manual_seed((zlib.crc32(str(sample["clip_id"]).encode()) * 1009 + tid * 7 + args.seed) % (2 ** 31))
+            noise = torch.randn(x0.shape, generator=g, device=device, dtype=x0.dtype)
+            idx_pf = torch.full((x0.shape[0], x0.shape[2]), int(tid), device=device, dtype=torch.int64)
+            _, st = ar_tf_loss(model, batch, new_weight=args.new_weight, new_weight_ramp=args.new_weight_ramp,
+                               teacher_forcing=args.teacher_forcing, tf_noise_aug_max_idx=0,
+                               use_gradient_checkpointing=False, noise=noise, idx_pf=idx_pf)
+            acc[j, 0] += st["loss_w"].double()
+            acc[j, 1] += st["mse"].double()
+            acc[j, 2] += 1
+    acc = accelerator.reduce(acc, reduction="sum")
+    model.train()
+    out = {}
+    for j, tid in enumerate(tids):
+        n = max(acc[j, 2].item(), 1.0)
+        t = float(t_of_idx(int(tid)))
+        out[t] = {"loss_w": acc[j, 0].item() / n, "mse": acc[j, 1].item() / n, "n": int(acc[j, 2].item())}
+    return out
+
+
+# ----------------------------------------------------------------------------
+# 主流程
+# ----------------------------------------------------------------------------
+
+def main(argv=None):
+    args = parse_args(argv)
+    accelerator = Accelerator(gradient_accumulation_steps=args.grad_accum)
+    is_main = accelerator.is_main_process
+    use_ds = accelerator.distributed_type == DistributedType.DEEPSPEED
+    set_seed(args.seed, device_specific=True)   # 每个 rank 的 ε / 噪声档 / 场景 dropout 抽签都不同
+    device = accelerator.device
+
+    if use_ds:
+        # 让 --grad_clip 成为唯一入口；yaml 里的 gradient_clipping 只是默认值（train.py:394-401）
+        ds_cfg = accelerator.state.deepspeed_plugin.deepspeed_config
+        ds_cfg["gradient_clipping"] = float(args.grad_clip)
+        ga = ds_cfg.get("gradient_accumulation_steps")
+        if ga not in ("auto", args.grad_accum):
+            raise SystemExit(f"--grad_accum={args.grad_accum} 与 accelerate yaml 的 "
+                             f"gradient_accumulation_steps={ga} 不一致：要么传 --grad_accum {ga}，"
+                             f"要么换 distill/configs/accelerate_zero2_4gpu_ga{args.grad_accum}.yaml"
+                             f"（现有 ga2 / ga4 两份；等效 batch = 4 × batch_size × GA）")
+        ds_cfg["gradient_accumulation_steps"] = args.grad_accum
+
+    run_dir = runlib.resolve_run_dir(args.output)
+    if is_main:
+        os.makedirs(run_dir, exist_ok=True)
+    log = runlib.TextTee(os.path.join(run_dir, "train.log"), enabled=is_main)
+    jsonl = runlib.JsonlLogger(os.path.join(run_dir, "train_log.jsonl"), enabled=is_main)
+    eff_batch = accelerator.num_processes * args.batch_size * args.grad_accum
+    log(f"=== {STAGE} | run={run_dir} world={accelerator.num_processes} "
+        f"dist={accelerator.distributed_type} mp={accelerator.mixed_precision} ===")
+    if max(args.action_dropout, args.hist_dropout, args.ctx_dropout, args.anchor_dropout) > 0:
+        log(f"[recipe] 条件 dropout: anchor={args.anchor_dropout} action={args.action_dropout} "
+            f"hist={args.hist_dropout} ctx={args.ctx_dropout}（scene={args.scene_dropout} 由数据集做）")
+    if args.long_seg_dir:
+        log(f"[recipe] 长视野续训：训练集 = {args.long_seg_dir}（AropeLongDatasetTF），val = {args.val_dir}（21 cell）")
+    log(f"[recipe] tf={args.teacher_forcing} tf_aug={args.tf_noise_aug_max_idx} "
+        f"windowed={args.windowed_train}/{args.windowed_sink} new_weight={args.new_weight} "
+        f"lr={args.lr} betas={args.betas} wd={args.weight_decay} warmup={args.warmup} clip={args.grad_clip} "
+        f"bs={args.batch_size}×accum{args.grad_accum}×{accelerator.num_processes}卡 ⇒ 等效 batch {eff_batch} "
+        f"max_steps={args.max_steps} gc={args.gradient_checkpointing}")
+    low = runlib.warn_if_low_disk(run_dir)
+    if low:
+        log(f"[warn] {low}")
+
+    # ---- 数据 ----
+    if args.dry_run:
+        train_ds, val_ds = _SynDataset(), None
+        if args.max_steps > 30:
+            args.max_steps, args.save_every, args.warmup, args.log_every, args.val_every = 12, 6, 3, 2, 0
+        log("[dry_run] tiny 模型 + 合成数据")
+    else:
+        sidecar = load_sidecar(args.sidecar)
+        if args.long_seg_dir:
+            # 长段（F=41/61/81）：latent 是整段一次过 VAE 的，不是 21-cell clip 拼接；键集与 arope_collate 契约相同
+            from longhorizon.arope_long_dataset import AropeLongDatasetTF
+            train_ds = AropeLongDatasetTF(args.long_seg_dir, text_table=args.text_table,
+                                          scene_dropout=args.scene_dropout, limit=args.limit, verbose=is_main)
+        else:
+            train_ds = AropeLatentDataset(args.train_dir, sidecar=sidecar, text_table=args.text_table,
+                                          text_mode="scene", scene_dropout=args.scene_dropout,
+                                          limit=args.limit, seed=args.seed, verbose=is_main,
+                                          exclude_src_in=args.val_dir, min_rt_ratio=args.min_rt_ratio)
+        val_ds = None
+        if args.val_every > 0 and args.val_n > 0:
+            # 验证不做场景 dropout；文本表已经在内存里，别再读一遍 1.5 GB
+            val_ds = AropeLatentDataset(args.val_dir, sidecar=sidecar, text_table=train_ds.text_table,
+                                        text_mode="scene", scene_dropout=0.0, limit=args.val_n,
+                                        seed=args.seed, verbose=is_main, min_rt_ratio=args.val_min_rt_ratio)
+        log(f"[data] train {len(train_ds)} clip，val {len(val_ds) if val_ds is not None else 0} clip")
+    dataloader = torch.utils.data.DataLoader(
+        train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True, collate_fn=arope_collate,
+        num_workers=args.num_workers, pin_memory=True, persistent_workers=args.num_workers > 0,
+    )
+
+    # ---- 模型 ----
+    t0 = time.time()
+    if args.dry_run:
+        module = _tiny_model()
+        resume_info = None
+    else:
+        module = AropeDiT(model_dir=args.model_dir, mask_channel=True, causal=True,
+                          train_local_attn_size=args.windowed_train, train_sink_size=args.windowed_sink)
+        # 因果化的起点：双向 SFT 检查点。assign=False ⇒ copy_ 进现有参数，保持模型自己的 dtype
+        resume_info = runlib.load_dit_ckpt(module, args.init_ckpt)
+        log(f"[model] 双向 SFT ← {args.init_ckpt}（{resume_info['n_tensors']} 个张量，"
+            f"缺 {resume_info['missing']}）")
+    # 只热启动权重的续训必须在 prepare **之前**：DeepSpeed 在 initialize 时就按当时的参数建好了
+    # fp32 主权重副本，prepare 之后再 copy_ 进 bf16 参数，第一次 step 就会被主权重原样盖回去（静默）。
+    warm_start = args.resume if (args.resume and args.resume.endswith(".safetensors")) else None
+    if warm_start:
+        info = runlib.load_dit_ckpt(module, warm_start)
+        log(f"[resume] 只热启动权重 {warm_start}（step 从 0 起）：{info['n_tensors']} 个张量")
+    n_params = sum(p.numel() for p in module.parameters())
+    module.train().requires_grad_(True)
+    log(f"[model] AropeDiT causal={module.causal} mask_channel={module.mask_channel} "
+        f"{n_params:,} 参数，加载 {time.time() - t0:.0f}s")
+
+    optimizer = torch.optim.AdamW(module.parameters(), lr=args.lr, weight_decay=args.weight_decay,
+                                  betas=args.betas)
+    warm = max(int(args.warmup), 0)
+    # lr scheduler 故意不进 prepare：accelerate 包装后每步会 step num_processes 次，warmup 被悄悄压缩
+    lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda s: min(1.0, (s + 1) / warm) if warm > 0 else 1.0)
+
+    # 只把 **一个** nn.Module 交给 prepare（accelerator.py:2229-2236 只留最后一个，多传会静默丢）
+    model, optimizer, dataloader = accelerator.prepare(module, optimizer, dataloader)
+    # dit.freqs 是普通 tuple 属性（不是 buffer），prepare / .to() 都搬不动它 —— 必须在裸模块上搬一次
+    unwrap(model).move_rope_table(device)
+    if is_main and use_ds:
+        clip_eff = getattr(getattr(model, "optimizer", None), "clip_grad", None)
+        log(f"[env] grad_clip={args.grad_clip}（DeepSpeed 生效值 {clip_eff}）")
+
+    # ---- 续训（accelerate 存档：权重 + 优化器 + 数据顺序，必须在 prepare 之后）----
+    start_step = 0
+    if args.resume and not warm_start:
+        rp = args.resume
+        if rp == "auto":
+            link = os.path.join(run_dir, "state")
+            rp = link if os.path.exists(link) else None
+            log(f"[resume] auto ⇒ {rp or '没有存档，从头开始'}")
+        if rp:
+            accelerator.load_state(rp)
+            with open(os.path.join(rp, "trainer_state.json"), encoding="utf-8") as fh:
+                start_step = int(json.load(fh)["step"])
+            with warnings.catch_warnings():   # 优化器还没 step 就快进 scheduler，PyTorch 会提醒一句，这里是有意的
+                warnings.simplefilter("ignore")
+                for _ in range(start_step):
+                    lr_scheduler.step()
+            log(f"[resume] accelerator.load_state({rp})，从 step {start_step} 续训")
+
+    # ---- config.json：复现所需的一切 ----
+    if is_main:
+        cfg = {
+            "stage": STAGE, "args": vars(args), "eff_batch": eff_batch,
+            "dataset": ({"synthetic": True} if args.dry_run else {
+                "train_dir": args.long_seg_dir or args.train_dir, "long_seg_dir": args.long_seg_dir,
+                "train_used": len(train_ds),
+                "train_total": getattr(train_ds, "n_total", len(train_ds)),
+                "train_invalid": getattr(train_ds, "n_invalid", None),
+                "train_leak_excluded": getattr(train_ds, "n_leak", None),
+                "train_low_rt_excluded": getattr(train_ds, "n_low_rt", None),
+                "text_table_keys": len(train_ds.text_table),
+                "val_used": len(val_ds) if val_ds is not None else 0,
+                "val_clip_ids": val_ds.clip_ids() if val_ds is not None else []}),
+            "model": {"n_params": n_params, "causal": True, "mask_channel": True,
+                      "init_ckpt": args.init_ckpt, "resume_info": resume_info, "start_step": start_step},
+            "env": {"world_size": accelerator.num_processes, "distributed_type": str(accelerator.distributed_type),
+                    "mixed_precision": accelerator.mixed_precision, "torch": torch.__version__,
+                    "cuda": torch.version.cuda,
+                    "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
+                    "python": sys.executable, "platform": platform.platform(), "hostname": socket.gethostname(),
+                    "argv": sys.argv, "started": time.strftime("%Y-%m-%d %H:%M:%S"),
+                    "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES")},
+            "deepspeed_config": (accelerator.state.deepspeed_plugin.deepspeed_config if use_ds else None),
+        }
+        try:
+            import accelerate, diffsynth  # noqa: F401
+            cfg["env"]["accelerate"] = accelerate.__version__
+        except Exception:
+            pass
+        runlib.save_config(run_dir, cfg)
+
+    # ---- 日志 ----
+    csv_fields = ["step", "loss", "loss_w", "mse", "new_frac", "t_mean", "lr", "grad_norm",
+                  "step_time", "mem_gb", "epoch"]
+    train_log = runlib.CsvLog(os.path.join(run_dir, "train_log.csv"), csv_fields, enabled=is_main)
+    val_log = runlib.CsvLog(os.path.join(run_dir, "val_log.csv"), ["step", "timestep", "loss_w", "mse", "n"],
+                            enabled=is_main)
+    val_tids = [int(idx_of_t(t)) for t in args.val_timesteps]
+
+    def do_validation(step):
+        if val_ds is None:
+            return
+        tv = time.time()
+        res = run_validation(accelerator, model, val_ds, val_tids, args)
+        if is_main:
+            for t, r in res.items():
+                val_log.write({"step": step, "timestep": t, **r})
+            mean_w = sum(r["loss_w"] for r in res.values()) / len(res)
+            log("[val] step {}  ".format(step)
+                + "  ".join(f"t={t:.0f}: w={r['loss_w']:.4f} mse={r['mse']:.4f} (n={r['n']})" for t, r in res.items())
+                + f"  | mean w={mean_w:.4f}  {time.time() - tv:.0f}s")
+
+    def do_save(step, loss_val=None):
+        # save_dit 是集合操作（accelerator.get_state_dict），每个 rank 都必须进来
+        runlib.save_dit(accelerator, model, runlib.dit_path(run_dir, step),
+                        {"stage": STAGE, "step": step, "teacher_forcing": args.teacher_forcing,
+                         "tf_noise_aug_max_idx": args.tf_noise_aug_max_idx,
+                         "action_dropout": args.action_dropout,
+                         "hist_dropout": args.hist_dropout,
+                         "ctx_dropout": args.ctx_dropout,
+                         "anchor_dropout": args.anchor_dropout,
+                         "windowed_train": args.windowed_train, "windowed_sink": args.windowed_sink,
+                         "new_weight": args.new_weight, "text_mode": "scene", "init_ckpt": args.init_ckpt,
+                         "long_seg_dir": args.long_seg_dir or ""},
+                        expect_keys=(None if args.dry_run else runlib.AROPE_DIT_KEYS), log=log)
+        if is_main:
+            gone = runlib.rotate(runlib.dit_glob(run_dir), args.keep_last_dits)
+            if gone:
+                log(f"[rotate] 删掉 {len(gone)} 个旧 dit（保留最新 {args.keep_last_dits} 个）")
+            runlib.write_progress(run_dir, step, stage=STAGE, loss=loss_val)
+        if args.save_state:
+            runlib.save_state_rotating(accelerator, run_dir, {"step": step, "stage": STAGE})
+
+    # ---- 训练循环 ----
+    model.train()
+    step, epoch, last_saved = start_step, 0, start_step
+    if device.type == "cuda":
+        torch.cuda.reset_peak_memory_stats(device)
+    log(f"[train] 从 step {start_step} 到 {args.max_steps}")
+    micro_time, loss_ema = 0.0, None
+    while step < args.max_steps:
+        for batch in dataloader:
+            if step >= args.max_steps:
+                break
+            t_step = time.time()
+            batch = to_device(batch, device)
+            lr_now = lr_scheduler.get_last_lr()[0]     # 本步用的 lr（step 之后拿到的是下一步的）
+            with accelerator.accumulate(model):
+                loss, stats = ar_tf_loss(
+                    model, batch, new_weight=args.new_weight, new_weight_ramp=args.new_weight_ramp,
+                    teacher_forcing=args.teacher_forcing, tf_noise_aug_max_idx=args.tf_noise_aug_max_idx,
+                    action_dropout=args.action_dropout, hist_dropout=args.hist_dropout,
+                    ctx_dropout=args.ctx_dropout, anchor_dropout=args.anchor_dropout,
+                    use_gradient_checkpointing=args.gradient_checkpointing)
+                accelerator.backward(loss)   # DeepSpeed 下这一步已含 engine.step()（裁剪 + 更新）
+                grad_norm = (accelerator.clip_grad_norm_(model.parameters(), args.grad_clip)
+                             if accelerator.sync_gradients else None)
+                optimizer.step()
+                if accelerator.sync_gradients:
+                    lr_scheduler.step()      # DeepSpeed 的 scheduler 包装不认累积边界，只能在这里守门
+                optimizer.zero_grad(set_to_none=True)
+            if device.type == "cuda":
+                torch.cuda.synchronize(device)
+            micro_time += time.time() - t_step
+            if not accelerator.sync_gradients:
+                continue                     # 还在累积，不算一步
+            step += 1
+            step_time, micro_time = micro_time, 0.0
+            loss_mean = accelerator.reduce(loss.detach().float(), reduction="mean").item()
+            loss_ema = loss_mean if loss_ema is None else 0.9 * loss_ema + 0.1 * loss_mean
+            if is_main and (step % args.log_every == 0 or step == args.max_steps):
+                gn = float(grad_norm) if grad_norm is not None and math.isfinite(float(grad_norm)) else float("nan")
+                row = {"step": step, "loss": loss_mean, "loss_w": stats["loss_w"].item(),
+                       "mse": stats["mse"].item(), "new_frac": stats["new_frac"].item(),
+                       "t_mean": stats["t_mean"].item(), "lr": lr_now, "grad_norm": gn,
+                       "step_time": step_time,
+                       "mem_gb": (torch.cuda.max_memory_allocated(device) / 1024 ** 3 if device.type == "cuda" else 0.0),
+                       "epoch": epoch}
+                train_log.write(row)
+                jsonl.log({**row, "loss_ema": loss_ema})
+                log(f"step {step}/{args.max_steps}  loss {loss_mean:.4f} (ema {loss_ema:.4f})  "
+                    f"w {row['loss_w']:.4f} mse {row['mse']:.4f} t̄ {row['t_mean']:.0f} new {row['new_frac']:.2f}  "
+                    f"lr {row['lr']:.2e}  gnorm {gn:.3f}  {step_time:.2f}s  mem {row['mem_gb']:.1f}GB")
+            if args.val_every > 0 and step % args.val_every == 0:
+                do_validation(step)
+            if args.save_every > 0 and step % args.save_every == 0:
+                do_save(step, loss_mean)
+                last_saved = step
+        epoch += 1
+
+    if last_saved != step:
+        do_save(step, loss_ema)
+    if device.type == "cuda":
+        peak = torch.tensor(torch.cuda.max_memory_allocated(device) / 1024 ** 3, device=device)
+        peak_max = accelerator.reduce(peak, reduction="max").item()
+    else:
+        peak_max = 0.0
+    log(f"[done] step {step}，各进程峰值显存 {peak_max:.1f} GB。"
+        f"下一步：把 dits/dit_step{step}.safetensors 同时作为 stage2 的 --teacher_ckpt 与 --student_ckpt")
+    if is_main:
+        runlib.write_progress(run_dir, step, stage=STAGE, finished=True, peak_mem_gb=peak_max)
+    train_log.close()
+    val_log.close()
+    jsonl.close()
+    log.close()
+    accelerator.wait_for_everyone()
+    accelerator.end_training()
+    return step
+
+
+if __name__ == "__main__":
+    main()

backup_2026-09-14/code/distill/stage1_ar/config.json

@@ -0,0 +1,72 @@
+{
+  "_note": "stage1（teacher forcing 因果化）真正在用的配方常量。这三处必须逐项相同：本文件 / run.sh 的默认值 / ar_tf.py 的 argparse 默认值 —— 上游 WorldMind 的 run.sh 与 argparse 有六处不一致，每一处都对应一次已知的崩溃（PORT_SPEC risk 11）。改配方就同时改这三处。",
+  "_not_this_file": "跑起来之后 <output>/config.json 是那一次 run 的完整快照（含数据集统计 / 环境 / deepspeed_config），由 runlib.save_config 写；本文件只是签进仓库的配方记录。",
+  "_launch": "bash distill/stage1_ar/run.sh",
+  "stage": "stage1_ar_tf",
+  "script": "distill/stage1_ar/ar_tf.py",
+
+  "init_ckpt": "outputs/arope_10k_ga2/step-10000.safetensors",
+  "init_ckpt_role": "双向全片 SFT（827 键 bf16，9,999,690,792 B）；stage1 只把注意力换成因果 + TF 双流，不换目标",
+  "output": "outputs/distill/stage1_ar_tf",
+  "product": "dits/dit_step{N}.safetensors → stage2 的 --teacher_ckpt 与 --student_ckpt（同一份）",
+
+  "train_dir": "data/latent/train_eybx",
+  "val_dir": "data/latent/val_eybx",
+  "sidecar": "data/latent/arope_sidecar.pt",
+  "text_table": "data/latent/text_table_actionrope.pt",
+  "text_mode": "scene",
+  "scene_dropout": 0.1,
+  "min_rt_ratio": 0.8,
+  "val_min_rt_ratio": 0.9,
+  "val_frac": 0,
+  "_val_frac": "只接受 0：ActionRoPE 的划分是离线的 train_eybx / val_eybx 清单，相邻 clip 只差 3 帧，按 clip 名 hash 划会漏样本",
+
+  "teacher_forcing": true,
+  "tf_noise_aug_max_idx": 0,
+  "_tf_noise_aug_max_idx": "0 = 精确干净的历史（CF 默认）。>0 时本移植抽的是 [1000-k, 1000) 的近干净档；上游 ar_diffusion_gm.py:93 抽的是 [0, k)，那是 σ≈1 的纯噪声端，方向反了",
+  "windowed_train": 0,
+  "windowed_sink": 0,
+  "_windowed": "全因果。F=21 下 sink=0 的 W=20 会把 cell 0（ActionRoPE 的干净锚点）挤出 cell 20 的视野",
+  "new_weight": 2.0,
+  "new_weight_ramp": 0.0,
+  "_new_weight": "new 区 ×2 的 loss 权重图照搬 SFT —— stage1 就是那个 SFT 目标加了因果掩码",
+  "first_frame_cond": true,
+  "loss": "cell 0 干净 + t=0 且不计 loss；cells 1..20 逐 cell 独立均匀抽噪声档，Σw·se/(Σw·48) 逐 cell 算，再乘逐 cell 的 bsmntw 钟形权重取平均",
+
+  "optimizer": "AdamW",
+  "lr": 1e-05,
+  "weight_decay": 0.01,
+  "adam_betas": [0.9, 0.999],
+  "_adam_betas": "ActionRoPE SFT 的口径（用户拍板）。WorldMind 那边是 CF 约定的 (0.0, 0.999)，不照抄",
+  "warmup": 200,
+  "lr_schedule": "线性 warmup 200 步后恒定 —— 每个 save_every 边界都是合法的停止点",
+  "grad_clip": 1.0,
+
+  "batch_size": 1,
+  "grad_accum": 2,
+  "world_size": 4,
+  "eff_batch": 8,
+  "_eff_batch": "4 卡 × batch 1 × GA 2。--grad_accum 必须等于 accelerate yaml 的 gradient_accumulation_steps，否则 ar_tf.py 直接 SystemExit",
+  "max_steps": 12000,
+  "gradient_checkpointing": true,
+  "num_workers": 2,
+  "seed": 0,
+
+  "save_every": 500,
+  "keep_last_dits": 6,
+  "save_state": true,
+  "val_every": 500,
+  "val_n": 16,
+  "val_timesteps": [200, 500, 800],
+  "log_every": 20,
+  "resume": "auto",
+
+  "accelerate_config": "distill/configs/accelerate_zero2_4gpu_ga2.yaml",
+  "distributed": "DeepSpeed ZeRO-2, bf16, 4 进程（只有一个可训模型，所以不用 FSDP）",
+  "gpus": "0,1,2,3",
+  "main_process_port": 29531,
+
+  "_budget": "≈11 s/优化器步（2 × 5.5 s 的 TF micro batch）⇒ 12000 步 ≈ 38 h；显存 ~34-50 GB/卡（143 GB 卡）；磁盘 6×10 GB dits + 2 个 state 槽",
+  "_fallback": "--no-teacher_forcing 换成 diffusion forcing（单条 8190 token 流，便宜约 2.5×），代价是 stage2 必须跟着一起 --no-teacher_forcing",
+  "_gate": "开跑前先过 distill/tests/test_tf_port.py（clean f → noisy f 泄漏必须是 0.00e+00）"
+}