From 77afe2d7470c149aacd4700fa10b9e83103dddbf Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Thu, 23 Jul 2026 05:45:52 -0700 Subject: [PATCH] =?UTF-8?q?[feat]:=20VSA+DMD2=20gate=20=E2=80=94=20sparse?= =?UTF-8?q?=20student=20rollout,=20sparsity-0=20VSA=20scoring=20(main's=20?= =?UTF-8?q?dense-copy=20semantics)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- fastvideo2/train/dmd2.py | 33 +++++++++------- fastvideo2/train/gates/capture_train_main.py | 13 +++++-- fastvideo2/train/gates/train_anchor.py | 41 +++++++++++++------- 3 files changed, 55 insertions(+), 32 deletions(-) diff --git a/fastvideo2/train/dmd2.py b/fastvideo2/train/dmd2.py index 7d6fd65fb..881ea624a 100644 --- a/fastvideo2/train/dmd2.py +++ b/fastvideo2/train/dmd2.py @@ -102,11 +102,12 @@ class DMD2Step: ).unflatten(0, pred_noise.shape[:2]) # --- forwards (BTCHW state, BCTHW model calls — main's convention) ----- # - def _fwd(self, model: Any, noisy_btchw: Any, t: Any, embeds: Any) -> Any: - out = model(noisy_btchw.permute(0, 2, 1, 3, 4), embeds, t) + def _fwd(self, model: Any, noisy_btchw: Any, t: Any, embeds: Any, + vsa: Any = None) -> Any: + out = model(noisy_btchw.permute(0, 2, 1, 3, 4), embeds, t, vsa=vsa) return out.permute(0, 2, 1, 3, 4) - def student_rollout(self, draws: dict, embeds: Any) -> Any: + def student_rollout(self, draws: dict, embeds: Any, vsa: Any = None) -> Any: """main's `_generator_multi_step_simulation_forward`, replaying the recorded draws: target_idx (int), init_noise, step_noises (list).""" import torch @@ -118,7 +119,7 @@ class DMD2Step: for k in range(len(self.denoising_steps) - 1): t_k = torch.tensor([self.denoising_steps[k]], device=x.device, dtype=torch.long) - pred = self._fwd(self.student.model, x, t_k, embeds) + pred = self._fwd(self.student.model, x, t_k, embeds, vsa=vsa) clean = self.pred_video(pred, x, t_k) t_next = torch.tensor([self.denoising_steps[k + 1]], device=x.device, dtype=torch.long) @@ -127,32 +128,36 @@ class DMD2Step: noisy_input = noise_latents[target_idx - 1] if target_idx > 0 else x_copy t_tgt = torch.tensor([self.denoising_steps[target_idx]], device=x.device, dtype=torch.long) - pred = self._fwd(self.student.model, noisy_input, t_tgt, embeds) + pred = self._fwd(self.student.model, noisy_input, t_tgt, embeds, vsa=vsa) return self.pred_video(pred, noisy_input, t_tgt) - def dmd_loss(self, x0_student: Any, draws: dict, cond: Any, uncond: Any) -> Any: + def dmd_loss(self, x0_student: Any, draws: dict, cond: Any, uncond: Any, + vsa_dense: Any = None) -> Any: + """``vsa_dense``: for VSA+DMD2, main scores with sparsity-0.0 VSA + metadata (same kernel, all blocks, gate mixing active) — NOT flash.""" import torch import torch.nn.functional as F with torch.no_grad(): t = draws["dmd_timestep"].to(x0_student.device) # warped+clamped, long noisy = self.add_noise(x0_student, draws["dmd_noise"], t) - fake = self.pred_video(self._fwd(self.critic.model, noisy, t, cond), - noisy, t) - real_c = self.pred_video(self._fwd(self.teacher, noisy, t, cond), - noisy, t) - real_u = self.pred_video(self._fwd(self.teacher, noisy, t, uncond), - noisy, t) + fake = self.pred_video(self._fwd(self.critic.model, noisy, t, cond, + vsa=vsa_dense), noisy, t) + real_c = self.pred_video(self._fwd(self.teacher, noisy, t, cond, + vsa=vsa_dense), noisy, t) + real_u = self.pred_video(self._fwd(self.teacher, noisy, t, uncond, + vsa=vsa_dense), noisy, t) real = real_c + (real_c - real_u) * self.guidance # DMD2 CFG grad = (fake - real) / torch.abs(x0_student - real).mean() grad = torch.nan_to_num(grad) return 0.5 * F.mse_loss(x0_student.float(), (x0_student.float() - grad.float()).detach()) - def critic_loss(self, x0_student: Any, draws: dict, cond: Any) -> Any: + def critic_loss(self, x0_student: Any, draws: dict, cond: Any, + vsa_dense: Any = None) -> Any: import torch t = draws["critic_timestep"].to(x0_student.device) noise = draws["critic_noise"] noisy = self.add_noise(x0_student, noise, t) - pred = self._fwd(self.critic.model, noisy, t, cond) + pred = self._fwd(self.critic.model, noisy, t, cond, vsa=vsa_dense) target = noise - x0_student return torch.mean((pred - target) ** 2) # bf16, no .float() (legacy) diff --git a/fastvideo2/train/gates/capture_train_main.py b/fastvideo2/train/gates/capture_train_main.py index 919c9d7f2..32d75c787 100644 --- a/fastvideo2/train/gates/capture_train_main.py +++ b/fastvideo2/train/gates/capture_train_main.py @@ -32,8 +32,8 @@ import os import sys MODE = sys.argv[1] if len(sys.argv) > 1 and not sys.argv[1].startswith("-") else "finetune" -VSA = MODE == "vsa" -DMD2 = MODE == "dmd2" +VSA = MODE in ("vsa", "vsa_dmd2") +DMD2 = MODE in ("dmd2", "vsa_dmd2") os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" if VSA else "FLASH_ATTN" os.environ.setdefault("WANDB_MODE", "offline") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") @@ -47,6 +47,8 @@ DATASET = "wlsaidhi/crush-smol_processed_t2v" # initializes the gate projections — an init-RNG parity gate for later MODEL = ("FastVideo/FastWan2.1-T2V-1.3B-Diffusers" if VSA else "Wan-AI/Wan2.1-T2V-1.3B-Diffusers") +# vsa_dmd2: teacher/critic ALSO FastWan (deterministic gate weights); the +# shipped recipe uses the base model + random-init gates — init-RNG gate later def _hash(t) -> str: @@ -104,7 +106,7 @@ def main() -> None: argv += ["--VSA_sparsity", "0.8", "--VSA_decay_rate", "0.2", "--VSA_decay_interval_steps", "1"] if DMD2: - base = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" + base = MODEL if MODE == "vsa_dmd2" else "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" argv[argv.index(str(NUM_LATENT_T))] = "4" argv += ["--real_score_model_path", base, "--fake_score_model_path", base, "--dmd_denoising_steps", "1000,757,522", @@ -114,6 +116,8 @@ def main() -> None: "--min_timestep_ratio", "0.02", "--max_timestep_ratio", "0.98", "--flow_shift", "8.0"] argv[argv.index("5e-5")] = "2e-6" # learning_rate per the shipped recipe + if MODE == "vsa_dmd2": + argv += ["--VSA_sparsity", "0.8"] sys.argv = argv parser = FlexibleArgumentParser() parser = TrainingArgs.add_cli_args(parser) @@ -281,6 +285,8 @@ def main() -> None: tb = orig_dstep(self, tb) rec = _cur() rec["gen_loss"] = float(tb.generator_loss) + md = getattr(tb, "attn_metadata_vsa", None) + rec["vsa_sparsity"] = float(getattr(md, "VSA_sparsity", 0.0) or 0.0) if md is not None else 0.0 rec["embeds"] = _np(tb.encoder_hidden_states) if tb.encoder_hidden_states is not None else None rec["closed"] = True return tb @@ -362,6 +368,7 @@ def main() -> None: meta_steps.append({ "targets": [ro.get("target_idx") for ro in rec["rollouts"]], "dmd_t": rec.get("dmd", {}).get("t_final"), + "vsa_sparsity": rec.get("vsa_sparsity", 0.0), "critic_t": rec["critic"]["t_final"], "gen_loss": rec["gen_loss"], "fake_loss": rec["fake_loss"], "x0_student_hash": rec.get("x0_student_hash"), diff --git a/fastvideo2/train/gates/train_anchor.py b/fastvideo2/train/gates/train_anchor.py index 524b9a693..788f91664 100644 --- a/fastvideo2/train/gates/train_anchor.py +++ b/fastvideo2/train/gates/train_anchor.py @@ -139,7 +139,7 @@ def run(mode: str = "finetune") -> int: -def run_dmd2() -> int: +def run_dmd2(mode: str = "dmd2") -> int: """DMD2 anchor: replay main's recorded per-phase draws through DMD2Step. Order per step mirrors train_one_step: student phase (rollout -> dmd loss -> student AdamW) only on interval steps, then critic phase (its OWN @@ -151,20 +151,23 @@ def run_dmd2() -> int: from fastvideo2.loading import resolve_weights from fastvideo2.train.dmd2 import DMD2Step - from fastvideo2.wan21.card import WAN21_T2V_1_3B - from fastvideo2.wan21.model_fv import WanModelFV + from fastvideo2.wan21.card import FASTWAN_T2V_1_3B, WAN21_T2V_1_3B + from fastvideo2.wan21.model_fv import WanModelFV, WanModelFVVSA - gold = _gold_dir("dmd2") + gold = _gold_dir(mode) with open(os.path.join(gold, "manifest.json")) as f: manifest = json.load(f) with open(os.path.join(gold, "steps.json")) as f: steps = json.load(f) device = "cuda" - root = resolve_weights(WAN21_T2V_1_3B) + is_vsa = mode == "vsa_dmd2" + card = FASTWAN_T2V_1_3B if is_vsa else WAN21_T2V_1_3B + cls = WanModelFVVSA if is_vsa else WanModelFV + root = resolve_weights(card) def load(): - return WanModelFV.from_pretrained(root, torch_dtype=torch.float32, - subfolder="transformer").to(device) + return cls.from_pretrained(root, torch_dtype=torch.float32, + subfolder="transformer").to(device) step_ctx = DMD2Step(load(), load(), load()) pz = np.load(os.path.join(gold, "params.npz")) @@ -182,6 +185,13 @@ def run_dmd2() -> int: z = np.load(os.path.join(gold, f"step{i}.npz")) cond = tt(z["embeds"]) n_roll = len(rec["targets"]) + vsa_student = vsa_dense = None + if is_vsa: + from fastvideo2.layers.vsa import build_vsa_meta + _, tl, _, hl, wl = z["ro0_init"].shape # BTCHW draws + grid = (tl, hl // 2, wl // 2) + vsa_student = build_vsa_meta(grid, rec.get("vsa_sparsity", 0.0), device) + vsa_dense = build_vsa_meta(grid, 0.0, device) def roll_draws(j): return {"target_idx": rec["targets"][j], @@ -191,7 +201,7 @@ def run_dmd2() -> int: student_step = rec["gen_loss"] != 0.0 if student_step: - x0 = step_ctx.student_rollout(roll_draws(0), cond) + x0 = step_ctx.student_rollout(roll_draws(0), cond, vsa=vsa_student) if i == 0 or rows == []: pass if rec.get("x0_student_hash"): @@ -199,7 +209,8 @@ def run_dmd2() -> int: 0.0 if _hash(x0) == rec["x0_student_hash"] else 1.0)) dmd = step_ctx.dmd_loss( x0, {"dmd_timestep": torch.tensor([rec["dmd_t"]], device=device), - "dmd_noise": tt(z["dmd_noise"])}, cond, uncond) + "dmd_noise": tt(z["dmd_noise"])}, cond, uncond, + vsa_dense=vsa_dense) dmd.backward() step_ctx.student.apply_grads_and_step() rows.append((f"gen.step{i}", abs(float(dmd.detach().item()) - rec["gen_loss"]))) @@ -209,13 +220,13 @@ def run_dmd2() -> int: j = n_roll - 1 # critic-phase rollout is the last one recorded with torch.no_grad(): - x0c = step_ctx.student_rollout(roll_draws(j), cond) + x0c = step_ctx.student_rollout(roll_draws(j), cond, vsa=vsa_student) if not student_step and rec.get("x0_student_hash"): rows.append((f"x0.step{i}", 0.0 if _hash(x0c) == rec["x0_student_hash"] else 1.0)) closs = step_ctx.critic_loss( x0c, {"critic_timestep": torch.tensor([rec["critic_t"]], device=device), - "critic_noise": tt(z["critic_noise"])}, cond) + "critic_noise": tt(z["critic_noise"])}, cond, vsa_dense=vsa_dense) closs.backward() step_ctx.critic.apply_grads_and_step() rows.append((f"fake.step{i}", abs(float(closs.detach().item()) - rec["fake_loss"]))) @@ -249,10 +260,10 @@ def run_dmd2() -> int: print(f" {n:18s} {v:.6e} {'OK' if ok else 'FAIL'} ({why})") from fastvideo2.verify import GateResult, append_ledger, env_fingerprint - append_ledger([GateResult(gate="anchor.train-dmd2-main", + append_ledger([GateResult(gate=f"anchor.train-{mode}-main", status="pass" if not failed else "fail", - model_id=WAN21_T2V_1_3B.model_id, - card_digest=WAN21_T2V_1_3B.digest(), + model_id=card.model_id, + card_digest=card.digest(), metrics={n: v for n, v in rows}, tolerances={"x0.step0": 0.0, "later": "band"}, env=env_fingerprint(), @@ -263,4 +274,4 @@ def run_dmd2() -> int: if __name__ == "__main__": args = [a for a in sys.argv[1:] if not a.startswith("-")] mode = args[0] if args else "finetune" - sys.exit(run_dmd2() if mode == "dmd2" else run(mode=mode)) + sys.exit(run_dmd2(mode) if mode in ("dmd2", "vsa_dmd2") else run(mode=mode))