From c73ccc75ed2e1db351b7e98ea7c30dcc9b0010bd Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Thu, 23 Jul 2026 04:50:49 -0700 Subject: [PATCH] =?UTF-8?q?[feat]:=20VSA=20training=20gate=20=E2=80=94=20c?= =?UTF-8?q?apture/anchor=20vsa=20mode,=20per-step=20ramped=20metadata=20th?= =?UTF-8?q?rough=20FinetuneStep?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- fastvideo2/train/finetune.py | 10 +++-- fastvideo2/train/gates/capture_train_main.py | 24 ++++++++--- fastvideo2/train/gates/train_anchor.py | 45 ++++++++++++-------- 3 files changed, 52 insertions(+), 27 deletions(-) diff --git a/fastvideo2/train/finetune.py b/fastvideo2/train/finetune.py index 1adb263cd..792381779 100644 --- a/fastvideo2/train/finetune.py +++ b/fastvideo2/train/finetune.py @@ -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() diff --git a/fastvideo2/train/gates/capture_train_main.py b/fastvideo2/train/gates/capture_train_main.py index f198c83c3..80a502c1c 100644 --- a/fastvideo2/train/gates/capture_train_main.py +++ b/fastvideo2/train/gates/capture_train_main.py @@ -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 " diff --git a/fastvideo2/train/gates/train_anchor.py b/fastvideo2/train/gates/train_anchor.py index 738d96fac..58537a0ff 100644 --- a/fastvideo2/train/gates/train_anchor.py +++ b/fastvideo2/train/gates/train_anchor.py @@ -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"))