code/tests/test_arope.py

"""ARoPE 模型补丁的自测（一张 H200 足够，DiT bf16 ~10 GB）。

    CUDA_VISIBLE_DEVICES=0 /opt/dlami/nvme/zhiyangdeng/ActionRoPE/.venv/bin/python -m pytest \
        /opt/dlami/nvme/zhiyangdeng/ActionRoPE/tests/test_arope.py -s -v

每个用例都能单独跑（-k），模型只在 module 级 fixture 里加载一次。
实测数字追加写到 $AROPE_TEST_NUMBERS（默认 $AROPE_TEST_DIR=outputs/test_artifacts 下的 json），供报告汇总。
"""
from __future__ import annotations

import json
import os
import time

import pytest
import torch

os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")

ROOT = "/opt/dlami/nvme/zhiyangdeng/ActionRoPE"
MODEL_DIR = f"{ROOT}/models/Wan2.2-TI2V-5B"
DIT_FILES = [f"{MODEL_DIR}/diffusion_pytorch_model-0000{i}-of-00003.safetensors" for i in (1, 2, 3)]
# 测试产物只落在项目内（outputs/test_artifacts），不写任何机器相关的临时目录
TEST_DIR = os.environ.get("AROPE_TEST_DIR", os.path.join(ROOT, "outputs", "test_artifacts"))
NUMBERS_PATH = os.environ.get("AROPE_TEST_NUMBERS", os.path.join(TEST_DIR, "arope_numbers.json"))

# SPEC 的几何常量
C, F, LAT_H, LAT_W = 48, 21, 30, 52
TOK_H, TOK_W = 15, 26
L_TXT, D_TXT = 512, 4096


def record(**kv):
    """把实测数字追加进 json（多个用例各写各的键）。"""
    data = {}
    try:
        with open(NUMBERS_PATH) as fp:
            data = json.load(fp)
    except (FileNotFoundError, json.JSONDecodeError):
        pass
    data.update({k: (float(v) if isinstance(v, (int, float)) and not isinstance(v, bool) else v) for k, v in kv.items()})
    os.makedirs(os.path.dirname(NUMBERS_PATH), exist_ok=True)
    with open(NUMBERS_PATH, "w") as fp:
        json.dump(data, fp, indent=2, ensure_ascii=False)
    print("\n[numbers]", json.dumps(kv, ensure_ascii=False))


@pytest.fixture(scope="module")
def dit():
    from diffsynth.pipelines.wan_video import ModelConfig, WanVideoPipeline

    pipe = WanVideoPipeline.from_pretrained(
        torch_dtype=torch.bfloat16, device="cuda",
        model_configs=[ModelConfig(path=DIT_FILES)],
        tokenizer_config=None, redirect_common_files=False,
    )
    model = pipe.dit
    model.eval().requires_grad_(False)
    # 原始键集：test_f 用它做对照，且必须在任何 install 之前抓
    model._arope_test_baseline_keys = set(model.state_dict().keys())
    return model


def make_inputs(seed=0, batch=1, device="cuda"):
    g = torch.Generator(device="cpu").manual_seed(seed)
    latents = torch.randn(batch, C, F, LAT_H, LAT_W, generator=g).to(device=device, dtype=torch.bfloat16)
    context = torch.randn(batch, L_TXT, D_TXT, generator=g).to(device=device, dtype=torch.bfloat16)
    # 训练链路里 timestep 是 scheduler.timesteps[id].to(bf16)，形状 [1]
    timestep = torch.tensor([500.0], device=device, dtype=torch.bfloat16)
    return latents, context, timestep


def err_stats(out, ref):
    out, ref = out.float(), ref.float()
    diff = (out - ref).abs()
    return {
        "max_abs": diff.max().item(),
        "mean_abs": diff.mean().item(),
        "rel_l2": (diff.norm() / ref.norm().clamp_min(1e-12)).item(),
        "ref_absmax": ref.abs().max().item(),
        "ref_std": ref.std().item(),
    }


# ---------------------------------------------------------------- (a)

@torch.no_grad()
def test_a_equivalence_with_model_fn(dit):
    from diffsynth.pipelines.wan_video import model_fn_wan_video
    from actionrope.arope import arope_forward

    latents, context, timestep = make_inputs()
    ref = model_fn_wan_video(dit=dit, latents=latents, timestep=timestep, context=context,
                             fuse_vae_embedding_in_latents=True)
    out = arope_forward(dit, latents, timestep, context, offset_px=None, mask_input=None)
    assert out.shape == (1, C, F, LAT_H, LAT_W)
    s = err_stats(out, ref)
    record(a_equiv_max_abs=s["max_abs"], a_equiv_rel_l2=s["rel_l2"], a_ref_absmax=s["ref_absmax"],
           a_ref_std=s["ref_std"], a_bitwise_equal=bool(torch.equal(out, ref)))
    assert torch.isfinite(out).all()
    assert s["rel_l2"] <= 1e-2


# ---------------------------------------------------------------- (b)

@torch.no_grad()
def test_b_per_cell_context_equivalence(dit):
    from actionrope.arope import arope_forward

    latents, context, timestep = make_inputs()
    ref = arope_forward(dit, latents, timestep, context)
    ctx_cells = context.unsqueeze(1).expand(1, F, L_TXT, D_TXT).contiguous()

    # 按内容去重
    out = arope_forward(dit, latents, timestep, ctx_cells)
    s = err_stats(out, ref)
    # 按 context_ids 去重（dataset 会给）
    ids = torch.zeros(1, F, dtype=torch.long)
    out2 = arope_forward(dit, latents, timestep, ctx_cells, context_ids=ids)
    s2 = err_stats(out2, ref)
    record(b_percell_max_abs=s["max_abs"], b_percell_rel_l2=s["rel_l2"],
           b_percell_ids_max_abs=s2["max_abs"], b_percell_ids_rel_l2=s2["rel_l2"],
           b_bitwise_equal=bool(torch.equal(out, ref)))
    assert s["rel_l2"] <= 1e-2 and s2["rel_l2"] <= 1e-2

    # 两个不同的串真的会分到不同 cell：把后半段 cell 换成另一串，输出必须变
    other = torch.randn_like(context)
    ctx_mix = ctx_cells.clone()
    ctx_mix[:, 11:] = other
    out3 = arope_forward(dit, latents, timestep, ctx_mix, context_ids=torch.tensor([[0] * 11 + [1] * 10]))
    s3 = err_stats(out3, ref)
    record(b_mixed_context_rel_l2=s3["rel_l2"])
    assert s3["rel_l2"] > 1e-3


# ---------------------------------------------------------------- (c)

def test_c_world_freqs_shift(dit):
    from actionrope.arope import build_world_freqs

    f, h, w = F, TOK_H, TOK_W
    table = build_world_freqs(dit, f, h, w, None, device="cuda")  # 查表，[S,1,64]
    assert table.shape == (f * h * w, 1, 64)
    ref_table = torch.cat([
        dit.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
        dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
        dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
    ], dim=-1).reshape(f * h * w, 1, -1).to("cuda")
    assert torch.equal(table, ref_table)

    # 零偏移走现算路径，必须与查表逐位相同
    zero = build_world_freqs(dit, f, h, w, torch.zeros(1, f, 2), device="cuda")
    assert zero.shape == (1, f * h * w, 1, 64)
    zero_max_diff = (zero[0] - table).abs().max().item()
    record(c_zero_offset_bitwise_equal=bool(torch.equal(zero[0], table)), c_zero_offset_max_diff=zero_max_diff)
    assert zero_max_diff <= 1e-6

    # 每帧右移 1 token：w 轴分量 = 查表的 w 索引 +1（列 j 拿到 w_table[j+1]）
    shift = build_world_freqs(dit, f, h, w, torch.tensor([[[1.0, 0.0]]]).expand(1, f, 2), device="cuda")
    shift = shift.view(f, h, w, 64)
    manual = torch.cat([
        dit.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
        dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
        dit.freqs[2][1:w + 1].view(1, 1, w, -1).expand(f, h, w, -1),
    ], dim=-1).to("cuda")
    d = (shift - manual).abs().max().item()
    record(c_shift_w_plus1_max_diff=d, c_shift_w_plus1_bitwise_equal=bool(torch.equal(shift, manual)))
    assert d <= 1e-6
    # f 分量直接查表，逐位相同；h 分量是现算的，与 CPU 上算的表只有 cos/sin 库函数的 ULP 级差异
    assert torch.equal(shift[..., :22], table.view(f, h, w, 64)[..., :22])
    d_h = (shift[..., 22:43] - table.view(f, h, w, 64)[..., 22:43]).abs().max().item()
    record(c_recomputed_h_axis_max_diff=d_h)
    assert d_h <= 1e-6

    # 每帧下移 1 token（dy=+1）：h 轴分量 = 查表的 h 索引 +1
    shift_h = build_world_freqs(dit, f, h, w, torch.tensor([[[0.0, 1.0]]]).expand(1, f, 2), device="cuda").view(f, h, w, 64)
    manual_h = dit.freqs[1][1:h + 1].view(1, h, 1, -1).expand(f, h, w, -1).to("cuda")
    assert (shift_h[..., 22:43] - manual_h).abs().max().item() <= 1e-6

    # 逐帧不同的偏移：第 k 帧右移 k token，只有对应帧被移
    off = torch.zeros(1, f, 2)
    off[0, :, 0] = torch.arange(f)
    var = build_world_freqs(dit, f, h, w, off, device="cuda").view(f, h, w, 64)
    for k in (0, 3, 20):
        manual_k = dit.freqs[2][k:k + w].view(1, w, -1).expand(h, w, -1).to("cuda")
        assert (var[k, ..., 43:] - manual_k).abs().max().item() <= 1e-6

    # 小数偏移：相位是连续的（0.5 token 落在 0 与 1 之间）
    half = build_world_freqs(dit, 1, 1, 1, torch.tensor([[[0.5, 0.0]]]), device="cuda").view(64)
    ang = torch.angle(half[43:])
    inv = 1.0 / (10000.0 ** (torch.arange(0, 42, 2).double() / 42))
    assert torch.allclose(ang.cpu(), 0.5 * inv, atol=1e-9)


# ---------------------------------------------------------------- (d)

def test_d_known_mask_geometry():
    from actionrope.arope import known_mask, loss_weight_map

    # (+64, 0) px = +2 token：每帧最右 2 列 new
    off = torch.tensor([[[64.0, 0.0]]]).expand(1, F, 2) / 32
    m = known_mask(off, TOK_H, TOK_W)
    assert m.shape == (1, F, TOK_H, TOK_W) and m.dtype == torch.bool
    assert m[..., :, :TOK_W - 2].all() and (~m[..., :, TOK_W - 2:]).all()

    # (0, −32) px = −1 token：最上 1 行 new
    off = torch.tensor([[[0.0, -32.0]]]).expand(1, F, 2) / 32
    m = known_mask(off, TOK_H, TOK_W)
    assert (~m[..., 0, :]).all() and m[..., 1:, :].all()

    # 零偏移：全 known
    assert known_mask(torch.zeros(1, F, 2), TOK_H, TOK_W).all()

    # 逐帧不同（帧 k 右移 k token）：帧 k 最右 k 列 new
    off = torch.zeros(1, F, 2)
    off[0, :, 0] = torch.arange(F)
    m = known_mask(off, TOK_H, TOK_W)
    for k in range(F):
        n_new = (~m[0, k]).any(dim=0).sum().item()
        assert n_new == min(k, TOK_W), (k, n_new)

    # loss_weight_map：latent 分辨率，偏移除以 16
    w = loss_weight_map(torch.tensor([[[64.0, 0.0]]]).expand(1, F, 2), new_weight=2.0)
    assert w.shape == (1, 1, F, LAT_H, LAT_W) and w.dtype == torch.float32
    assert (w[..., :, :LAT_W - 4] == 1).all() and (w[..., :, LAT_W - 4:] == 2).all()
    w = loss_weight_map(torch.tensor([[[0.0, -32.0]]]).expand(1, F, 2), new_weight=3.0)
    assert (w[..., :2, :] == 3).all() and (w[..., 2:, :] == 1).all()
    assert (loss_weight_map(torch.zeros(2, F, 2)) == 1).all()

    # 半格边界取闭区间：+16 px = +0.5 token ⇒ 最右列世界坐标 25.5，仍 known
    assert known_mask(torch.tensor([[[16.0, 0.0]]]) / 32, TOK_H, TOK_W).all()
    assert not known_mask(torch.tensor([[[17.0, 0.0]]]) / 32, TOK_H, TOK_W)[..., -1].any()


# ---------------------------------------------------------------- (e)

def test_e_training_step(dit):
    from actionrope.arope import arope_forward, install_arope, loss_weight_map

    n_params_before = sum(p.numel() for p in dit.parameters())
    install_arope(dit, mask_channel=True)
    n_params_after = sum(p.numel() for p in dit.parameters())
    assert hasattr(dit, "arope_mask_embedding")
    assert dit.arope_mask_embedding.weight.dtype == torch.bfloat16
    assert dit.arope_mask_embedding.weight.device.type == "cuda"

    latents, context, timestep = make_inputs(seed=1)
    off = torch.zeros(1, F, 2, device="cuda")
    off[0, :, 0] = torch.linspace(0, 96, F)  # 全程向右走 3 token
    off[0, :, 1] = torch.linspace(0, -20, F)
    mask = loss_weight_map(off.cpu()).to("cuda").eq(1.0).float()  # known=1 / new=0
    ctx_cells = context.unsqueeze(1).expand(1, F, L_TXT, D_TXT).contiguous()

    # 零初始化 ⇒ mask 通道不改变输出（eval、无梯度）
    with torch.no_grad():
        o_none = arope_forward(dit, latents, timestep, context, offset_px=None)
        o_ref = arope_forward(dit, latents, timestep, context, offset_px=off)
        o_mask = arope_forward(dit, latents, timestep, context, offset_px=off, mask_input=mask)
    assert torch.equal(o_ref, o_mask)
    # 世界 RoPE 真的在起作用：非零偏移必须改变输出，但量级不变
    s_off = err_stats(o_ref, o_none)
    assert s_off["rel_l2"] > 1e-3
    assert abs(o_ref.float().std().item() / o_none.float().std().item() - 1) < 0.2
    record(e_zero_init_mask_no_change=True, e_params_added=n_params_after - n_params_before,
           e_params_total=n_params_after, e_offset_vs_plain_rel_l2=s_off["rel_l2"],
           e_out_std_plain=o_none.float().std().item(), e_out_std_offset=o_ref.float().std().item())

    # 全参数训练一步：前向 + 反向，梯度检查点开
    try:
        dit.train().requires_grad_(True)
        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()
        t0 = time.time()
        out = arope_forward(dit, latents, timestep, ctx_cells, offset_px=off, mask_input=mask,
                            use_gradient_checkpointing=True,
                            context_ids=torch.zeros(1, F, dtype=torch.long))
        target = torch.randn_like(out)
        wmap = loss_weight_map(off.cpu()).to(out.device)
        # 与 SPEC 的 loss 形式一致：Σ w·(pred−target)² / Σ w，w 在通道维广播，所以分母也要乘通道数
        loss = ((out.float() - target.float()) ** 2 * wmap).sum() / (wmap.sum() * out.shape[1])
        loss.backward()
        torch.cuda.synchronize()
        dt = time.time() - t0
        peak_gb = torch.cuda.max_memory_allocated() / 1024 ** 3
        assert torch.isfinite(loss)

        g_mask = dit.arope_mask_embedding.weight.grad
        assert g_mask is not None and torch.isfinite(g_mask).all()
        mask_grad_norm = g_mask.float().norm().item()
        assert mask_grad_norm > 0
        # 抽查其它参数的梯度非零
        checks = {
            "patch_embedding.weight": dit.patch_embedding.weight,
            "blocks.0.self_attn.q.weight": dit.blocks[0].self_attn.q.weight,
            "blocks.15.cross_attn.k.weight": dit.blocks[15].cross_attn.k.weight,
            "blocks.29.ffn.2.weight": dit.blocks[29].ffn[2].weight,
            "head.head.weight": dit.head.head.weight,
            "text_embedding.0.weight": dit.text_embedding[0].weight,
            "time_embedding.0.weight": dit.time_embedding[0].weight,
        }
        grad_norms = {}
        for name, p in checks.items():
            assert p.grad is not None, name
            grad_norms[name] = p.grad.float().norm().item()
            assert grad_norms[name] > 0, name
        n_with_grad = sum(1 for p in dit.parameters() if p.grad is not None and p.grad.abs().sum() > 0)
        n_total = sum(1 for p in dit.parameters())
        record(e_loss=loss.item(), e_peak_mem_gb=peak_gb, e_fwd_bwd_sec=dt,
               e_mask_conv_grad_norm=mask_grad_norm, e_grad_norms=grad_norms,
               e_params_with_nonzero_grad=f"{n_with_grad}/{n_total}")
        assert n_with_grad == n_total
    finally:
        # 还原：其它用例要在干净的 eval 模型上跑
        for p in dit.parameters():
            p.grad = None
        dit.eval().requires_grad_(False)
        install_arope(dit, mask_channel=False)
        torch.cuda.empty_cache()


# ---------------------------------------------------------------- (f)

def test_f_state_dict_keys_unchanged(dit):
    from actionrope.arope import install_arope

    baseline = dit._arope_test_baseline_keys
    install_arope(dit, mask_channel=False)
    assert set(dit.state_dict().keys()) == baseline
    assert not hasattr(dit, "arope_mask_embedding")

    install_arope(dit, mask_channel=True)
    keys = set(dit.state_dict().keys())
    extra = keys - baseline
    assert extra == {"arope_mask_embedding.weight", "arope_mask_embedding.bias"}
    assert baseline <= keys

    install_arope(dit, mask_channel=False)
    assert set(dit.state_dict().keys()) == baseline
    record(f_keys_unchanged=True, f_extra_keys_with_mask=sorted(extra))