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

[fix]: sdpa-fp32 exact-math golden (flash helper force-casts bf16); isolated anchor probes; robust report over record unions

Browse Source
This commit is contained in:
SolitaryThinker
2026-07-22 14:59:23 -07:00
parent 93f80b5c4b
commit 787e17f500
2 changed files with 54 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo2/wan21/anchor.py
+40 -18
View File
@@ -208,24 +208,35 @@ def fastvideo_main_adapter(root: str, device: str) -> dict[str, Callable]:
# The generic anchor run #
# --------------------------------------------------------------------------- #
def run_anchor(adapter: dict[str, Callable], goldens_dir: str) -> list[dict]:
"""Drive one implementation over the goldens; return comparison records."""
"""Drive one implementation over the goldens; return comparison records.
Each probe is isolated: a crash becomes a fail record, not a lost run."""
records: list[dict] = []
def probe(name: str, fn: Callable[[], dict]) -> None:
try:
records.append(fn())
except Exception as e:
records.append({"name": name, "status": "fail",
"detail": f"{type(e).__name__}: {e}"})
t5 = G.load_golden(goldens_dir, "text_encoder")
for key, text in (("prompt", G.PROMPT), ("negative", G.NEGATIVE)):
ours = adapter["text_encode"](text) # [512, 4096] zero-padded
length = int(t5[f"{key}_len"])
rec = G.compare(f"text_encoder.{key}", ours[:length], t5[f"{key}_embeds"][:length],
TOLS["text_encoder"])
pad = float(np.abs(ours[length:]).max()) if length < ours.shape[0] else 0.0
rec["pad_max_abs"] = pad # padding convention check
records.append(rec)
def _t5(key=key, text=text) -> dict:
ours = adapter["text_encode"](text) # [512, 4096] zero-padded
length = int(t5[f"{key}_len"])
rec = G.compare(f"text_encoder.{key}", ours[:length], t5[f"{key}_embeds"][:length],
TOLS["text_encoder"])
rec["pad_max_abs"] = (float(np.abs(ours[length:]).max())
if length < ours.shape[0] else 0.0)
return rec
probe(f"text_encoder.{key}", _t5)
dit = G.load_golden(goldens_dir, "dit")
context = dit["context"] # [512, 4096], canonical padded form
for t in dit["timesteps"].tolist():
ours = adapter["dit_forward"](dit["x"], float(t), context)
records.append(G.compare(f"dit.t{int(t)}", ours, dit[f"out_t{int(t)}"], TOLS["dit"]))
probe(f"dit.t{int(t)}", lambda t=t: G.compare(
f"dit.t{int(t)}", adapter["dit_forward"](dit["x"], float(t), context),
dit[f"out_t{int(t)}"], TOLS["dit"]))
if "dit_forward_fp32" in adapter:
try:
@@ -233,12 +244,15 @@ def run_anchor(adapter: dict[str, Callable], goldens_dir: str) -> list[dict]:
except FileNotFoundError:
f32 = None
if f32 is not None:
ours = adapter["dit_forward_fp32"](f32["x"], float(f32["timestep"]), f32["context"])
records.append(G.compare("dit.fp32", ours, f32["out"], TOLS["dit_fp32"]))
probe("dit.fp32", lambda: G.compare(
"dit.fp32",
adapter["dit_forward_fp32"](f32["x"], float(f32["timestep"]), f32["context"]),
f32["out"], TOLS["dit_fp32"]))
vae = G.load_golden(goldens_dir, "vae")
ours = adapter["vae_decode"](vae["z"])
records.append(G.compare("vae.decode", ours, vae["video"].astype(np.float32), TOLS["vae"]))
probe("vae.decode", lambda: G.compare(
"vae.decode", adapter["vae_decode"](vae["z"]),
vae["video"].astype(np.float32), TOLS["vae"]))
return records
@@ -264,7 +278,12 @@ def report_markdown(all_records: dict[str, list[dict]], manifest: dict) -> str:
"""One table: rows = probes, columns = implementations, cells = rel L2 vs
the official goldens."""
impls = list(all_records)
names = [r["name"] for r in all_records[impls[0]]]
by_name = {impl: {r["name"]: r for r in recs} for impl, recs in all_records.items()}
names: list[str] = []
for recs in all_records.values(): # union, first-seen order
for r in recs:
if r["name"] not in names:
names.append(r["name"])
lines = ["# Wan2.1 numerics vs official implementation",
"",
f"Goldens: `{manifest.get('capture', {}).get('repo', '?')}` @ "
@@ -273,12 +292,15 @@ def report_markdown(all_records: dict[str, list[dict]], manifest: dict) -> str:
"",
"| probe | " + " | ".join(impls) + " |",
"|---|" + "---|" * len(impls)]
for i, name in enumerate(names):
for name in names:
cells = []
for impl in impls:
r = all_records[impl][i]
r = by_name[impl].get(name)
if r is None:
cells.append("—")
continue
v = r.get("rel_l2")
cell = "n/a" if v is None else f"{v:.2e}"
cell = f"{v:.2e}" if v is not None else (r.get("detail", "n/a")[:40] or "n/a")
if r["status"] == "fail":
cell = f"**FAIL** {cell}"
cells.append(cell)
fastvideo2/wan21/probe_dit.py
+14 -1
View File
@@ -114,13 +114,26 @@ def main() -> None:
print(f"[ATTN] official sdpa vs flash : {rel(e_attn, base):.3e}")
# ---- MATH: fp32, no autocast ------------------------------------------- #
# Official's flash_attention helper force-casts to bf16 even for fp32
# inputs, so a clean math golden needs the SDPA patch (which honors the
# input dtype). Record both: flash-path fp32 (bf16 attention inside) for
# reference, SDPA-path fp32 as the exact-math golden.
off_f32 = off.float()
up32 = torch.from_numpy(ctx_unpad).to(dev, torch.float32)
out_f32_flash = off_forward(off_f32, x32, [up32], autocast=False)
watt.flash_attention = sdpa_attention
if hasattr(wmodel, "flash_attention"):
wmodel.flash_attention = sdpa_attention
out_f32 = off_forward(off_f32, x32, [up32], autocast=False)
watt.flash_attention = orig_fa
if hasattr(wmodel, "flash_attention"):
wmodel.flash_attention = orig_fa
print(f"[MATH] official fp32: flash-path vs sdpa-path : {rel(out_f32_flash, out_f32):.3e}")
G.save_golden(args.goldens, "dit_fp32", {"x": x_np, "context": ctx_pad,
"context_len": np.int64(plen),
"timestep": np.float32(t750), "out": out_f32},
{"note": "official WanModel fp32 forward, no autocast — exact-math golden"})
{"note": "official WanModel fp32 forward, no autocast, SDPA attention "
"(true fp32 end to end) — the exact-math golden"})
print(f"[MATH] saved dit_fp32 golden (std {out_f32.std():.4f})")
del off, off_f32
torch.cuda.empty_cache()