"""adaln 臂（baseline/adaln.py）的自测：baseline/SPEC.md 的测试 1–4，单卡，模型只加载一次。 CUDA_VISIBLE_DEVICES=7 PYTHONPATH=/opt/dlami/nvme/zhiyangdeng/ActionRoPE DIFFSYNTH_SKIP_DOWNLOAD=True \ /opt/dlami/nvme/zhiyangdeng/ActionRoPE/.venv/bin/python -m pytest \ /opt/dlami/nvme/zhiyangdeng/ActionRoPE/tests/test_arm_adaln.py -s -v 1. 零初始化 ⇒ arope_forward(arm=arm) 与 arope_forward() 逐位相同 2. 扰动新增参数 + 非零动作 ⇒ 输出变化；梯度回到新增参数与 DiT（use_gradient_checkpointing=True） 3. 新增参数量与结构公式一致，并与旧报告的 linear / xattn 对照 4. state_dict 键集：加 `arm.` 前缀后与 DiT 键不冲突、可 strict 加载 实测数字写到 outputs/test_artifacts/arm_adaln_numbers.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)] TEST_DIR = os.environ.get("AROPE_TEST_DIR", os.path.join(ROOT, "outputs", "test_artifacts")) NUMBERS_PATH = os.environ.get("ARM_ADALN_TEST_NUMBERS", os.path.join(TEST_DIR, "arm_adaln_numbers.json")) 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): 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) model._baseline_keys = set(model.state_dict().keys()) return model @pytest.fixture(scope="module") def arm(dit): from baseline.adaln import AdaLNArm a = AdaLNArm() a.install(dit) a.eval().requires_grad_(False) return a 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 = torch.tensor([500.0], device=device, dtype=torch.bfloat16) return latents, context, timestep def make_action_inputs(batch=1, device="cuda", zero=False): """SPEC 的 action_inputs：全程向右走 ~11 token、向上 ~3 token（cell 0 = 0）。""" off = torch.zeros(batch, F, 2) if not zero: off[:, :, 0] = torch.linspace(0, 356, F) / 32 # 356 px ≈ 1.0× 步速 21 cell off[:, :, 1] = torch.linspace(0, -96, F) / 32 delta = torch.zeros_like(off) delta[:, 1:] = off[:, 1:] - off[:, :-1] return { "offset_px": (off * 32).to(device), "offset_tok": off.to(device), "delta_tok": delta.to(device), "action_idx": torch.full((batch, F), 7, dtype=torch.long, device=device), } def err_stats(out, ref): out, ref = out.float(), ref.float() diff = (out - ref).abs() return { "max_abs": diff.max().item(), "rel_l2": (diff.norm() / ref.norm().clamp_min(1e-12)).item(), } # ---------------------------------------------------------------- 1 @torch.no_grad() def test_1_zero_init_equivalence(dit, arm): from actionrope.arope import arope_forward assert arm.zero_init_check() latents, context, timestep = make_inputs() ai = make_action_inputs() ref = arope_forward(dit, latents, timestep, context) out = arope_forward(dit, latents, timestep, context, arm=arm, action_inputs=ai) assert out.shape == ref.shape and torch.isfinite(out).all() s = err_stats(out, ref) bitwise = bool(torch.equal(out, ref)) # delta_tok 缺省 ⇒ 自己差分，同样逐位不变 ai2 = {k: v for k, v in ai.items() if k != "delta_tok"} out2 = arope_forward(dit, latents, timestep, context, arm=arm, action_inputs=ai2) record(equiv_max_abs=s["max_abs"], equiv_rel_l2=s["rel_l2"], equiv_bitwise=bitwise, equiv_no_delta_bitwise=bool(torch.equal(out2, ref))) assert bitwise or s["rel_l2"] <= 1e-6 assert torch.equal(out2, ref) # ---------------------------------------------------------------- 2 def test_2_perturbed_output_changes_and_grads(dit, arm): from actionrope.arope import arope_forward latents, context, timestep = make_inputs(seed=1) ai = make_action_inputs() ai_zero = make_action_inputs(zero=True) ctx_cells = context.unsqueeze(1).expand(1, F, L_TXT, D_TXT).contiguous() # 随机扰动全部新增参数（proj 末层用较小的 std，免得把 t_mod 冲到失真） saved = {k: v.clone() for k, v in arm.state_dict().items()} g = torch.Generator(device="cpu").manual_seed(123) with torch.no_grad(): for name, p in arm.named_parameters(): std = 0.01 if name.startswith("proj") else 0.05 p.copy_(torch.randn(p.shape, generator=g).to(p.device, p.dtype) * std) assert not arm.zero_init_check() try: with torch.no_grad(): ref = arope_forward(dit, latents, timestep, context) out = arope_forward(dit, latents, timestep, context, arm=arm, action_inputs=ai) out_zero = arope_forward(dit, latents, timestep, context, arm=arm, action_inputs=ai_zero) s = err_stats(out, ref) s_act = err_stats(out, out_zero) # 动作本身（而不只是偏置）在起作用 std_ratio = out.float().std().item() / ref.float().std().item() record(perturbed_rel_l2=s["rel_l2"], perturbed_max_abs=s["max_abs"], perturbed_action_vs_zero_action_rel_l2=s_act["rel_l2"], perturbed_out_std_ratio=std_ratio) assert torch.isfinite(out).all() assert s["rel_l2"] > 1e-3 assert s_act["rel_l2"] > 1e-3 # 前向 + 反向，梯度检查点开，DiT 与 arm 一起训练 dit.train().requires_grad_(True) arm.train().requires_grad_(True) torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() t0 = time.time() out = arope_forward(dit, latents, timestep, ctx_cells, arm=arm, action_inputs=ai, use_gradient_checkpointing=True, context_ids=torch.zeros(1, F, dtype=torch.long)) target = torch.randn_like(out) loss = ((out.float() - target.float()) ** 2)[:, :, 1:].mean() # 帧 0 不计 loss.backward() torch.cuda.synchronize() dt = time.time() - t0 peak_gb = torch.cuda.max_memory_allocated() / 1024 ** 3 assert torch.isfinite(loss) arm_grads = {} for name, p in arm.named_parameters(): assert p.grad is not None and torch.isfinite(p.grad).all(), name arm_grads[name] = p.grad.float().norm().item() assert arm_grads[name] > 0, name checks = { "patch_embedding.weight": dit.patch_embedding.weight, "blocks.0.self_attn.q.weight": dit.blocks[0].self_attn.q.weight, "blocks.0.modulation": dit.blocks[0].modulation, "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, "time_projection.1.weight": dit.time_projection[1].weight, } dit_grads = {} for name, p in checks.items(): assert p.grad is not None, name dit_grads[name] = p.grad.float().norm().item() assert dit_grads[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()) grad_ok = n_with_grad == n_total and all(v > 0 for v in arm_grads.values()) record(grad_ok=grad_ok, loss=loss.item(), peak_mem_gb=peak_gb, fwd_bwd_sec=dt, arm_grad_norms=arm_grads, dit_grad_norms=dit_grads, dit_params_with_nonzero_grad=f"{n_with_grad}/{n_total}") assert grad_ok finally: for p in list(dit.parameters()) + list(arm.parameters()): p.grad = None dit.eval().requires_grad_(False) arm.load_state_dict(saved) arm.eval().requires_grad_(False) torch.cuda.empty_cache() assert arm.zero_init_check() # ---------------------------------------------------------------- 3 def test_3_param_count(dit, arm): from baseline.adaln import AdaLNArm n = arm.n_new_params() expected = AdaLNArm.expected_n_params(dit.dim, n_axes=4) n_dit = sum(p.numel() for p in dit.parameters()) print(f"\n[adaln] 新增参数 {n:,} ({n / 1e6:.2f}M)，DiT {n_dit / 1e6:.1f}M，占 {100 * n / n_dit:.2f}%；" f"对照：linear ≈ 0.18M，xattn ≈ 98.6M，上游 LTX-2.3(dim 4096, 6 轴) ≈ 118M") record(n_new_params=n, n_new_params_expected=expected, n_dit_params=n_dit) assert n == expected assert 60e6 < n < 70e6 # ---------------------------------------------------------------- 4 def test_4_state_dict_keys(dit, arm): from baseline.adaln import AdaLNArm keys = list(arm.state_dict().keys()) assert len(keys) == len(set(keys)) # 只有权重，没有 buffer（action_scale 是 persistent=False） assert all(k.startswith(("embedder.", "proj.")) for k in keys), keys prefixed = {f"arm.{k}" for k in keys} assert not (prefixed & dit._baseline_keys) assert set(dit.state_dict().keys()) == dit._baseline_keys # 装 arm 不动 DiT 的键 fresh = AdaLNArm() fresh.install(dit) missing, unexpected = fresh.load_state_dict(arm.state_dict(), strict=True) assert not missing and not unexpected for k, v in fresh.state_dict().items(): assert torch.equal(v, arm.state_dict()[k]) # 带前缀的 ckpt 视图能剥回来 strict 加载 ckpt = {f"arm.{k}": v for k, v in arm.state_dict().items()} fresh.load_state_dict({k[len("arm."):]: v for k, v in ckpt.items() if k.startswith("arm.")}, strict=True) record(state_dict_keys=sorted(keys), state_dict_strict_load_ok=True)