"""train.py 的 CPU 测试：tiny WanModel 走 GcTrainModule 前向 / 反向、TI2V 首帧、sem 项从 0 长起来、参数组、warmup、 验证三支（cond / uncond / nulltext）、导出键集与 metadata、CLI 默认值 = SPEC §7。 PYTHONPATH=code/gc25d:code/arope .venv/bin/python -m pytest -q code/gc25d/tests/test_train.py """ from __future__ import annotations import json import struct import pytest import torch import model as GM import train as TR from diffsynth.models.wan_video_dit import WanModel # tiny 几何同 code/arope/distill/tests/preflight.py::TINY（head_dim 12 ⇒ RoPE 三段列数凑得上）：latent 4×4 → 2×2 token / cell TINY = dict(dim=24, in_dim=48, ffn_dim=48, out_dim=48, text_dim=4096, freq_dim=256, eps=1e-6, patch_size=(1, 2, 2), num_heads=2, num_layers=2, has_image_input=False) F_LAT, LAT_H, LAT_W, L_TXT = 3, 4, 4, 6 N_IMG = 4 * (F_LAT - 1) + 1 def tiny_module(vocab: str = "parts", seed: int = 0, gc: bool = True) -> TR.GcTrainModule: torch.manual_seed(seed) dit = WanModel(**TINY).float() GM.install_sem_cond(dit, vocab) dit.train().requires_grad_(True) return TR.GcTrainModule(dit, first_frame_cond=True, gradient_checkpointing=gc) def tiny_batch(seed: int = 0, b: int = 1) -> dict: g = torch.Generator().manual_seed(seed) raw = torch.tensor(sorted(GM.valid_raw_ids("parts")), dtype=torch.uint8) sem = raw[torch.randint(0, len(raw), (b, N_IMG, LAT_H * 16, LAT_W * 16), generator=g)] sem[..., :32, :] = 1 # 一块大面积的静态类，免得全是噪声 return { "input_latents": torch.randn(b, 48, F_LAT, LAT_H, LAT_W, generator=g), "context": torch.randn(b, F_LAT, L_TXT, 4096, generator=g), "context_ids": torch.arange(F_LAT).repeat(b, 1), "sem_ids": sem, "clip_id": [f"c{seed}_{i}" for i in range(b)], } # ---------------------------------------------------------------------------- # GcTrainModule # ---------------------------------------------------------------------------- def test_forward_backward_finite_and_grads(): m = tiny_module() batch = tiny_batch() loss, stats = m(batch, timestep_id=300) assert torch.isfinite(loss) and loss.item() > 0 assert set(stats) == {"mse", "timestep", "sem_rms", "x_rms"} assert stats["sem_rms"].item() == 0.0 and stats["x_rms"].item() > 0 # proj_out 零初始化 ⇒ sem 项恰为 0 loss.backward() enc = getattr(m.dit, GM.SEM_COND_NAME) assert enc.proj_out.weight.grad is not None and enc.proj_out.weight.grad.abs().sum() > 0 assert m.dit.patch_embedding.weight.grad.abs().sum() > 0 assert all(p.grad is not None for n, p in m.named_parameters() if not n.startswith(TR.COND_PREFIX)) def test_loss_matches_manual_and_excludes_frame0(): m = tiny_module(gc=False) batch = tiny_batch() tid = 417 noise = torch.randn_like(batch["input_latents"]) loss, stats = m(batch, tid, noise=noise.clone()) # 手算：x_t、帧 0 换成干净 x0、t 置 0（gc_forward 里）、loss = mse(帧 1..) × w(t) x0 = batch["input_latents"] t = m.scheduler.timesteps[tid] sig = m.scheduler.sigmas[tid] x_t = (1 - sig) * x0 + sig * noise x_t[:, :, 0:1] = x0[:, :, 0:1] pred = GM.gc_forward(m.dit, x_t, t.view(1), batch["context"], batch["sem_ids"], context_ids=batch["context_ids"]) mse = ((pred[:, :, 1:] - (noise - x0)[:, :, 1:]) ** 2).mean() w = m.scheduler.training_weight(t.view(1)) assert torch.allclose(stats["mse"], mse, rtol=1e-5, atol=1e-6) assert torch.allclose(loss, mse * w, rtol=1e-5, atol=1e-6) # 帧 0 的噪声不进输入、帧 0 不计 loss ⇒ 改帧 0 的噪声 loss 不变 noise2 = noise.clone() noise2[:, :, 0] = torch.randn_like(noise2[:, :, 0]) * 7 loss2, _ = m(batch, tid, noise=noise2) assert torch.equal(loss, loss2) def test_sem_rms_grows_from_zero_and_condition_is_wired(): m = tiny_module() groups = TR.sem_param_groups(m, lr=1e-4, cond_lr=1e-2, weight_decay=0.01) opt = torch.optim.AdamW(groups) batch = tiny_batch() rms = [] for k in range(3): loss, stats = m(batch, timestep_id=500) rms.append(stats["sem_rms"].item()) opt.zero_grad() loss.backward() opt.step() assert rms[0] == 0.0 and 0 < rms[1] < rms[2], rms # 条件确实接上了：换 sem_ids（全 255）⇒ 预测变了 with torch.no_grad(): noise = torch.randn_like(batch["input_latents"]) _, s1 = m(batch, 500, noise=noise, train=False) _, s2 = m(dict(batch, sem_ids=torch.full_like(batch["sem_ids"], 255)), 500, noise=noise, train=False) assert s1["mse"].item() != s2["mse"].item() def test_vocab_ids_module_runs(): m = tiny_module(vocab="ids") batch = tiny_batch() batch["sem_ids"] = GM.collapse_parts(batch["sem_ids"]) loss, _ = m(batch, 100) assert torch.isfinite(loss) # ---------------------------------------------------------------------------- # 优化器 / 调度 # ---------------------------------------------------------------------------- def test_sem_param_groups(): m = tiny_module() groups = TR.sem_param_groups(m, lr=1e-5, cond_lr=1e-4, weight_decay=0.01) assert len(groups) == 2 assert [g["lr"] for g in groups] == [1e-5, 1e-4] and all(g["weight_decay"] == 0.01 for g in groups) names = {id(p): n for n, p in m.named_parameters()} cond = [names[id(p)] for p in groups[1]["params"]] rest = [names[id(p)] for p in groups[0]["params"]] assert len(cond) == 6 and all(n.startswith("dit.gc_sem_cond.") for n in cond) assert not any(n.startswith("dit.gc_sem_cond.") for n in rest) assert len(cond) + len(rest) == len(list(m.parameters())) opt = torch.optim.AdamW(groups, lr=1e-5, weight_decay=0.01) sch = torch.optim.lr_scheduler.LambdaLR(opt, TR.warmup_lambda(4)) lrs = [] for _ in range(6): lrs.append(sch.get_last_lr()) opt.step() sch.step() assert lrs[0] == pytest.approx([2.5e-6, 2.5e-5]) and lrs[3] == pytest.approx([1e-5, 1e-4]) and lrs[5] == lrs[3] def test_warmup_lambda(): f = TR.warmup_lambda(500) assert f(0) == pytest.approx(1 / 500) and f(499) == 1.0 and f(20000) == 1.0 assert TR.warmup_lambda(0)(0) == 1.0 def test_cli_defaults_match_spec(): a = TR.parse_args(["--output", "/opt/dlami/nvme/zhiyangdeng/maliang/outputs/_x"]) assert (a.lr, a.cond_lr, a.weight_decay, a.warmup_steps, a.grad_clip) == (1e-5, 1e-4, 0.01, 500, 1.0) assert (a.max_steps, a.batch_size, a.grad_accum, a.save_every, a.val_every, a.val_n) == (20000, 1, 2, 1000, 500, 16) assert a.val_timesteps == [200.0, 500.0, 800.0] and a.val_uncond and a.save_state and a.gradient_checkpointing assert (a.text_dropout, a.cond_dropout, a.mask_p, a.vocab, a.first_frame_cond) == (0.1, 0.05, 0.5, "parts", True) assert a.data_root.endswith("data/gc25d_v10") and a.text_table.endswith("data/gc25d_v10/text_table_gc25d.pt") # ---------------------------------------------------------------------------- # 验证三支 # ---------------------------------------------------------------------------- class _FakeAccel: device = torch.device("cpu") process_index, num_processes = 0, 1 @staticmethod def reduce(t, reduction="sum"): return t def test_val_variants_and_run_validation(): m = tiny_module(gc=False) null_ctx = torch.full((1, F_LAT, L_TXT, 4096), 0.5) null_ids = torch.full((1, F_LAT), 99, dtype=torch.int64) b = tiny_batch() vb = TR.val_variants(b, null_ctx, null_ids, uncond=True) assert set(vb) == {"cond", "uncond", "nulltext"} assert bool((vb["uncond"]["sem_ids"] == 255).all()) and torch.equal(vb["uncond"]["context"], b["context"]) assert torch.equal(vb["nulltext"]["context"], null_ctx) and vb["nulltext"]["context_ids"].tolist() == [[99] * F_LAT] assert torch.equal(vb["nulltext"]["sem_ids"], b["sem_ids"]) # val 集 = 未 collate 的样本列表（run_validation 自己 collate） val = [] for i in range(3): s = tiny_batch(seed=10 + i) val.append({k: (v[0] if isinstance(v, (torch.Tensor, list)) else v) for k, v in s.items()}) tids = [m.timestep_id_for(t) for t in (200.0, 800.0)] res = TR.run_validation(_FakeAccel(), m, m, val, tids, seed=0, null_context=null_ctx, null_ids=null_ids) assert len(res) == 2 for t, r in res.items(): assert r["n"] == 3 and all(torch.isfinite(torch.tensor(r[k])) for k in ("mse", "mse_uncond", "mse_nulltext")) assert r["gap_uncond"] == 0.0 # gc_sem_cond 零初始化 ⇒ 换成全 255 也逐位相同 assert r["gap_nulltext"] != 0.0 # 固定噪声 ⇒ 可复现 res2 = TR.run_validation(_FakeAccel(), m, m, val, tids, seed=0, null_context=null_ctx, null_ids=null_ids) assert res == res2 # ---------------------------------------------------------------------------- # 导出 # ---------------------------------------------------------------------------- def _header(path): with open(path, "rb") as fh: n = struct.unpack("