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

[feat]: VSA training gate — capture/anchor vsa mode, per-step ramped metadata through FinetuneStep

Browse Source
This commit is contained in:
SolitaryThinker
2026-07-23 04:50:49 -07:00
parent d96fd59525
commit c73ccc75ed
3 changed files with 52 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo2/train/finetune.py
+6 -4
View File
@@ -49,12 +49,14 @@ class FinetuneStep:
p.requires_grad_(True)
def step(self, latents: Any, embeds: Any, noise: Any, timesteps: Any,
sigmas: Any) -> tuple[float, float]:
"""All tensor args bf16 on device (sigmas broadcastable). Returns
(loss, grad_norm) exactly as main logs them."""
sigmas: Any, vsa: Any = None) -> tuple[float, float]:
"""All tensor args bf16 on device (sigmas broadcastable). ``vsa`` is
the per-step VSAMeta for VSA training (main ramps sparsity per step,
so metadata is rebuilt each step). Returns (loss, grad_norm) exactly
as main logs them."""
import torch
noisy = (1.0 - sigmas) * latents + sigmas * noise
pred = self.model(noisy, embeds, timesteps.to(torch.bfloat16))
pred = self.model(noisy, embeds, timesteps.to(torch.bfloat16), vsa=vsa)
target = noise - latents
loss = torch.mean((pred.float() - target.float()) ** 2) / self.grad_accum
loss.backward()
fastvideo2/train/gates/capture_train_main.py
+18 -6
View File
@@ -31,7 +31,9 @@ import json
import os
import sys
os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
MODE = sys.argv[1] if len(sys.argv) > 1 and not sys.argv[1].startswith("-") else "finetune"
VSA = MODE == "vsa"
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")
@@ -39,6 +41,11 @@ SEED = 42
STEPS = 5
NUM_LATENT_T = 8
DATASET = "wlsaidhi/crush-smol_processed_t2v"
# vsa gate trains the FastWan checkpoint (has to_gate_compress weights) so
# all parameters are deterministic; base-checkpoint VSA finetune random-
# 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")
def _hash(t) -> str:
@@ -54,7 +61,7 @@ def main() -> None:
here = os.path.dirname(os.path.abspath(__file__))
out = os.path.normpath(os.path.join(here, "..", "..", "evidence", "goldens",
"train-finetune-main"))
f"train-{MODE}-main"))
os.makedirs(out, exist_ok=True)
data_root = snapshot_download(DATASET, repo_type="dataset", token=False)
@@ -70,8 +77,8 @@ def main() -> None:
argv = [
"capture",
"--model_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
"--pretrained_model_name_or_path", "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
"--model_path", MODEL,
"--pretrained_model_name_or_path", MODEL,
"--data_path", data_path,
"--dataloader_num_workers", "1",
"--num_gpus", "1", "--sp_size", "1", "--tp_size", "1",
@@ -92,6 +99,9 @@ def main() -> None:
"--weight_only_checkpointing_steps", "100000",
"--training_state_checkpointing_steps", "100000",
]
if VSA:
argv += ["--VSA_sparsity", "0.8", "--VSA_decay_rate", "0.2",
"--VSA_decay_interval_steps", "1"]
sys.argv = argv
parser = FlexibleArgumentParser()
parser = TrainingArgs.add_cli_args(parser)
@@ -140,7 +150,8 @@ def main() -> None:
tb = orig_step(self, tb)
records[-1].update(loss=float(tb.total_loss.item() if torch.is_tensor(tb.total_loss)
else tb.total_loss),
grad_norm=float(tb.grad_norm))
grad_norm=float(tb.grad_norm),
vsa_sparsity=float(getattr(tb, "current_vsa_sparsity", 0.0) or 0.0))
return tb
TP.TrainingPipeline._get_next_batch = get_batch
@@ -182,7 +193,8 @@ def main() -> None:
commit = subprocess.run(["git", "-C", src, "rev-parse", "HEAD"],
capture_output=True, text=True).stdout.strip()
with open(os.path.join(out, "manifest.json"), "w") as f:
json.dump({"fastvideo_commit": commit, "dataset": DATASET,
json.dump({"fastvideo_commit": commit, "dataset": DATASET, "mode": MODE,
"model": MODEL,
"seed": SEED, "steps": STEPS, "num_latent_t": NUM_LATENT_T,
"config": "1gpu bs1 accum1 lr5e-5 wd1e-4 betas(0.9,0.999) "
"clip1.0 uniform-t cfg_rate0 dit_fp32 mixed_bf16 "
fastvideo2/train/gates/train_anchor.py
+28 -17
View File
@@ -20,10 +20,10 @@ import sys
import numpy as np
def _gold_dir() -> str:
def _gold_dir(mode: str) -> str:
here = os.path.dirname(os.path.abspath(__file__))
return os.path.normpath(os.path.join(here, "..", "..", "evidence", "goldens",
"train-finetune-main"))
f"train-{mode}-main"))
def rel(a: np.ndarray, b: np.ndarray) -> float:
@@ -33,26 +33,29 @@ def rel(a: np.ndarray, b: np.ndarray) -> float:
return float(np.linalg.norm(a - b)) / (d if d else 1.0)
def run() -> int:
def run(mode: str = "finetune") -> int:
import torch
from fastvideo2.loading import resolve_weights
from fastvideo2.train.finetune import FinetuneStep
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()
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"
# main trains its own WanTransformer3DModel port from the Diffusers repo;
# our bitwise-equal vendor of it loads the same weights fp32
root = resolve_weights(WAN21_T2V_1_3B) # Wan-AI/Wan2.1-T2V-1.3B-Diffusers
model = WanModelFV.from_pretrained(root, torch_dtype=torch.float32,
subfolder="transformer").to(device)
# main trains its own WanTransformer3DModel port; our bitwise-equal
# vendor loads the same weights fp32 (VSA gate: the FastWan checkpoint,
# which carries deterministic to_gate_compress weights)
card = FASTWAN_T2V_1_3B if mode == "vsa" else WAN21_T2V_1_3B
cls = WanModelFVVSA if mode == "vsa" else WanModelFV
root = resolve_weights(card)
model = cls.from_pretrained(root, torch_dtype=torch.float32,
subfolder="transformer").to(device)
trainer = FinetuneStep(model)
w0 = trainer.master["proj_out.weight"][:8, :8].cpu().numpy().copy()
@@ -71,6 +74,12 @@ def run() -> int:
timesteps = torch.tensor(rec["timesteps"], device=device)
sigmas = torch.tensor(rec["sigmas"], device=device,
dtype=torch.bfloat16).view(-1, 1, 1, 1)
vsa_meta = None
if mode == "vsa":
from fastvideo2.layers.vsa import build_vsa_meta
_, _, tl, hl, wl = z["latents"].shape
vsa_meta = build_vsa_meta((tl, hl // 2, wl // 2),
rec["vsa_sparsity"], device)
if i == 0: # input-construction triage: does OUR noisy match theirs?
noisy = (1.0 - sigmas) * latents + sigmas * noise
rows.append(("noisy.step0", 0.0 if _hash(noisy) == rec["noisy_hash"] else 1.0))
@@ -80,10 +89,11 @@ def run() -> int:
if "pred" in z.files:
with torch.no_grad():
p = trainer.model(noisy, embeds,
timesteps.to(torch.bfloat16))
timesteps.to(torch.bfloat16), vsa=vsa_meta)
rows.append(("pred.step0", rel(p.detach().to(torch.float32).cpu().numpy(),
z["pred"])))
loss, gnorm = trainer.step(latents, embeds, noise, timesteps, sigmas)
loss, gnorm = trainer.step(latents, embeds, noise, timesteps, sigmas,
vsa=vsa_meta)
rows.append((f"loss.step{i}", abs(loss - rec["loss"])))
rows.append((f"gnorm.step{i}", abs(gnorm - rec["grad_norm"])))
print(f" step{i}: ours loss={loss:.8f} main={rec['loss']:.8f} "
@@ -115,10 +125,10 @@ def run() -> int:
print(f" {n:14s} {v:.6e} {'OK' if ok else 'FAIL'} ({why})")
from fastvideo2.verify import GateResult, append_ledger, env_fingerprint
append_ledger([GateResult(gate="anchor.train-finetune-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={"step0": 0.0,
"later": "main self-noise band (see manifest)"},
@@ -128,4 +138,5 @@ def run() -> int:
if __name__ == "__main__":
sys.exit(run())
args = [a for a in sys.argv[1:] if not a.startswith("-")]
sys.exit(run(mode=args[0] if args else "finetune"))