nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feat]: VSA+DMD2 gate — sparse student rollout, sparsity-0 VSA scoring (main's dense-copy semantics)

Browse Source
This commit is contained in:
SolitaryThinker
2026-07-23 05:45:52 -07:00
parent 26fe3b4226
commit 77afe2d747
3 changed files with 55 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo2/train/dmd2.py
+19 -14
View File
@@ -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)
fastvideo2/train/gates/capture_train_main.py
+10 -3
View File
@@ -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"),
fastvideo2/train/gates/train_anchor.py
+26 -15
View File
@@ -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))