ActionRoPE-internal / code /tests /test_arm_adaln.py

add docs+code

880dff9 verified about 2 months ago

10.9 kB

"""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)

	"""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)