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

[fix] correct GPU adapters against real fastvideo API (cross-check findings)

Browse Source 
Adversarial cross-check of the written-not-run torch adapters against the real
fastvideo source confirmed the interface contracts (DiT returns bare velocity
tensor; timestep=sigma*1000; encode().mode() + bare decode; .last_hidden_state;
no fused solver kernel) but caught a wrong construction layer. Fixed in code:

- Construction: WanTransformer3DModel / AutoencoderKLWan have NO from_pretrained.
  Replace it with the real FastVideo loaders (TransformerLoader / VAELoader /
  TextEncoderLoader + TokenizerLoader, each load(model_path, fastvideo_args)).
  The loader resolves the class from the checkpoint config — so UMT5-vs-T5 is
  chosen correctly instead of hardcoded (was BLOCKER #1/#4/#5).
- Text encoder: wrap the forward in set_forward_context(...) — the (U)MT5
  attention reads global state via get_forward_context(); a bare call mis-encodes
  (was BLOCKER #2). Drop the wrong padding="max_length".
- VAE: apply latent normalization the DiT expects — (z-mean)*inv_std on encode,
  inverse on decode, with latents_std stored as its reciprocal; shift_factor
  before decode (was BLOCKER #3). Skipping it yields washed-out video, not error.

The remaining unknowns are genuinely box-dependent (FastVideoArgs fields,
shift_factor placement, exact tokenizer kwargs, FSDP) — GPU_BRINGUP.md reconciled
to mark what's now fixed-in-code vs what still needs the box. 204 CPU tests pass.
This commit is contained in:
SolitaryThinker
2026-06-18 09:41:15 +00:00
parent 15b6278e9f
commit 7fd58f26ea
2 changed files with 84 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
v2/platform/backends/GPU_BRINGUP.md
+25 -16
View File
@@ -33,20 +33,22 @@ resolved implementations differ:
## Ordered bring-up checklist (on the GPU box)
1. **Set weights.** Fill `ComponentSpec.checkpoint` for each component (HF id or local path) — e.g. a
small helper that stamps a model root onto the wan21 card. *(Risk A — without this, builders raise.)*
1. **Set weights + args.** Fill `ComponentSpec.checkpoint` for each component (HF id or local path),
and provide the `FastVideoArgs` the real loaders need (`_fastvideo_args` builds a minimal one from
the path — confirm its required fields/precision). *(Risk A — without these, builders raise.)*
2. **Detect.** `Platform.detect()` returns `cuda(smXY)`; confirm the arch string. `component_matrix()`
shows the three cuda components `available=True`.
3. **Build in isolation.** Build each component; assert type + a single forward's output shape (no
full denoise yet). Confirm the loader call (`from_pretrained` vs `TransformerLoader().load`). *(A)*
4. **One DiT step.** `dit(x, pe, sigma) -> velocity`; check finite + shape `[C,T,H,W]`. Confirm the
`timestep = sigma*1000` convention *(B)* and that the output is **velocity (noise−clean)**, not x0
or epsilon — a sign flip denoises backward. *(C)*
3. **Build in isolation.** Build each component via the FastVideo loaders (`TransformerLoader` /
`VAELoader` / `TextEncoderLoader` + `TokenizerLoader`); assert type + a single forward's output
shape (no full denoise yet). *(A)*
4. **One DiT step.** `dit(x, pe, sigma) -> velocity`; check finite + shape `[C,T,H,W]`. The
`timestep = sigma*1000` convention and the velocity (noise−clean) semantics were cross-checked as
MATCHING the real `forward`/scheduler — confirm numerically. *(B, C)*
5. **One solver step.** `flow_match_step` finite + right shape (math mirrors `loop/sampler.py`).
6. **VAE.** `encode` → handle the `DiagonalGaussianDistribution` (`.mode()`) + `latents_mean/std`
normalization to match the DiT's training latent scale; `decode` → video in `[-1,1]`. *(D)*
7. **Text.** Confirm `UMT5EncoderModel` vs `T5EncoderModel`, the tokenizer (`max_length=512`), and
that the forward runs inside `set_forward_context(...)`. *(E)*
6. **VAE.** Normalization is now applied in-adapter (`(z-mean)*inv_std` on encode, inverse on decode,
`latents_std` as reciprocal). Confirm the `shift_factor` placement/sign and dtype on the box. *(D)*
7. **Text.** Class resolution (UMT5 vs T5, from config) and `set_forward_context(...)` are now wired.
Confirm the exact tokenizer kwargs (`text_len`/`max_length`, special tokens) from the model config. *(E)*
8. **End-to-end.** Full t2v denoise → VAE decode → compare to a known-good fastvideo generation
(SSIM / the ssim regression harness).
9. **RL / SDE path.** `flow_sde_step` returns a finite `(prev, log_prob, mean, eff_std)`; run a
@@ -57,14 +59,21 @@ resolved implementations differ:
## Risks / open unknowns (the `# BRINGUP` points)
Cross-checked against the real source (`crosscheck-gpu-adapters`): the **interface contracts matched**
(DiT returns a bare velocity tensor; `timestep=sigma*1000`; `encode().mode()` + bare `decode`;
`.last_hidden_state`; no fused solver kernel). The **construction layer was wrong and is now fixed**
in code (real loaders instead of the nonexistent `from_pretrained`; UMT5-vs-T5 resolved from config;
`set_forward_context` wired; latent normalization applied). What remains is genuinely box-dependent:
| | Risk | Failure mode if wrong |
|---|---|---|
| **A** | checkpoint/config + exact loader call (`from_pretrained` vs `TransformerLoader().load`) | builder raises (blocking) |
| **B** | sigma→timestep scaling (`sigma*1000`, continuous vs discrete index) | silent garbage, not a crash |
| **C** | DiT output is velocity (noise−clean), correct sign/orientation | denoises backward |
| **D** | VAE returns a distribution + `latents_mean/std/shift` normalization | washed-out / saturated video |
| **E** | `UMT5EncoderModel` vs `T5EncoderModel`; tokenizer; `set_forward_context` | wrong/empty conditioning |
| **A** | `FastVideoArgs` construction the real loaders need (exact fields/precision); checkpoint path | builder raises (blocking) |
| **B** | sigma→timestep scaling — cross-checked as `sigma*1000`; confirm numerically | silent garbage, not a crash |
| **C** | DiT output velocity (noise−clean) — cross-checked as matching; confirm sign on box | denoises backward |
| **D** | `shift_factor` placement/sign + dtype of the (now-applied) latent normalization | washed-out / saturated video |
| **E** | exact tokenizer kwargs (`text_len`/`max_length`, special tokens) from config — class + `set_forward_context` now fixed in code | wrong/empty conditioning |
| **F** | torch training surface (`mse_grad_step`) not implemented | RL/distill on cuda is a separate workstream |
| **G** | numpy↔torch marshalling per DiT call (H2D/D2H + dtype round-trip) | correct but slow; torch-native surface is the perf follow-up |
Items A–E are correctness; F–G are scope/perf. None block the **CPU mini** (all gated `available=False`).
Multi-GPU FSDP sharding via the loaders also needs on-box verification (single-GPU bring-up first).
v2/platform/backends/torch_adapters.py
+59 -35
View File
@@ -16,8 +16,6 @@ numpy<->torch at its boundary. A torch-native end-to-end surface is the perf fol
"""
from __future__ import annotations
import importlib
import numpy as np
import torch
@@ -27,10 +25,22 @@ import torch
NUM_TRAIN_TIMESTEPS = 1000
def _resolve(load_id: str):
"""``"pkg.mod:Class"`` → the class object (the card's ``load_id``)."""
module_path, _, cls_name = load_id.partition(":")
return getattr(importlib.import_module(module_path), cls_name)
def _fastvideo_args(spec):
# BRINGUP risk A: the real loaders take a FastVideoArgs (model config / precision / parallelism),
# not just a path. Confirm the exact constructor + required fields on the box.
from fastvideo.fastvideo_args import FastVideoArgs
return FastVideoArgs(model_path=spec.checkpoint)
def _load_via_fastvideo(loader_attr: str, spec):
"""Construct a real component via the FastVideo loaders (models/loader/component_loader.py) — the
REAL entry point. ``WanTransformer3DModel`` / ``AutoencoderKLWan`` have NO ``from_pretrained``; the
loader reads the checkpoint config, resolves the right class (so e.g. UMT5 vs T5 is chosen from the
config, not hardcoded), instantiates with that config, and loads weights (incl. FSDP). The card's
``load_id`` is therefore a hint, not the instantiation path."""
from fastvideo.models.loader import component_loader as _cl
loader = getattr(_cl, loader_attr)()
return loader.load(spec.checkpoint, _fastvideo_args(spec))
def _torch_dtype(name: str):
@@ -117,29 +127,45 @@ class TorchWanDiT:
# VAE — AutoencoderKLWan (fastvideo/models/vaes/wanvae.py) #
# --------------------------------------------------------------------------- #
class TorchWanVAE:
"""vae.decode(latent[C,T,H,W]) -> video; vae.encode(video) -> latent (numpy in/out)."""
"""vae.decode(latent[C,T,H,W]) -> video; vae.encode(video) -> latent (numpy in/out).
The DiT operates in NORMALIZED latent space, so encode applies ``(z - mean) * inv_std`` and decode
inverts it. FastVideo stores ``latents_std`` as its RECIPROCAL — the encode transform is
``(z - latents_mean) * latents_std`` with ``latents_std = 1/std`` (training_utils.py:801-806).
Skipping this yields washed-out / saturated video, not an error (BRINGUP risk D).
"""
def __init__(self, module, *, device, dtype):
self.module = module.to(device=device, dtype=dtype).eval()
self.device, self.dtype = device, dtype
@torch.no_grad()
def decode(self, latent):
z = _to_torch(latent, device=self.device, dtype=self.dtype).unsqueeze(0)
# BRINGUP risk D: if the DiT works in normalized latent space, undo latents_mean/std/shift
# (wanvae.py:1111) before decode. decode -> video [B,3,T,H,W] in [-1,1] (wanvae.py:1262).
video = self.module.decode(z)
return _to_numpy(video.squeeze(0))
def _mean_invstd(self, like):
mean = torch.tensor(self.module.latents_mean, device=like.device,
dtype=torch.float32).view(1, -1, 1, 1, 1)
inv_std = (1.0 / torch.tensor(self.module.latents_std, device=like.device,
dtype=torch.float32)).view(1, -1, 1, 1, 1)
return mean, inv_std
@torch.no_grad()
def encode(self, video):
x = _to_torch(video, device=self.device, dtype=self.dtype).unsqueeze(0)
dist = self.module.encode(x)
# BRINGUP risk D: encode returns a DiagonalGaussianDistribution (wanvae.py:1196), NOT a tensor.
dist = self.module.encode(x) # DiagonalGaussianDistribution (wanvae.py:1178)
z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist)
# ... and apply latents_mean/std normalization to match the DiT's training latent scale.
mean, inv_std = self._mean_invstd(z)
z = (z.float() - mean) * inv_std # → the normalized latent the DiT expects
return _to_numpy(z.squeeze(0))
@torch.no_grad()
def decode(self, latent):
z = _to_torch(latent, device=self.device, dtype=self.dtype).unsqueeze(0).float()
mean, inv_std = self._mean_invstd(z)
z = z / inv_std + mean # invert the encode-side normalization
shift = getattr(self.module, "shift_factor", None) # BRINGUP risk D: confirm placement/sign
if shift is not None:
z = z + (shift.to(z) if torch.is_tensor(shift) else shift)
video = self.module.decode(z.to(self.dtype)) # → video [B,3,T,H,W] in [-1,1] (wanvae.py:1240)
return _to_numpy(video.squeeze(0))
# --------------------------------------------------------------------------- #
# Text encoder — (U)MT5 (fastvideo/models/encoders/t5.py) #
@@ -154,13 +180,17 @@ class TorchT5Encoder:
@torch.no_grad()
def encode(self, text):
toks = self.tokenizer(text or "", return_tensors="pt", padding="max_length",
max_length=self.max_length, truncation=True)
# BRINGUP risk E: the exact tokenizer kwargs come from the model config (text_len/max_length,
# special tokens). Don't force padding="max_length" — that is not in the real kwargs.
toks = self.tokenizer(text or "", return_tensors="pt", max_length=self.max_length, truncation=True)
ids = toks.input_ids.to(self.device)
mask = toks.attention_mask.to(self.device)
# BRINGUP risk E: Wan2.1 t2v uses UMT5EncoderModel (card load_id says T5EncoderModel), and the
# real forward must run inside FastVideo's set_forward_context(...) (text_encoding.py:273).
out = self.module(input_ids=ids, attention_mask=mask)
# REQUIRED: the FastVideo T5/UMT5 attention reads global state via get_forward_context(), so the
# forward must run inside set_forward_context(...) (forward_context.py:54; text_encoding.py:273).
# A bare call reads stale/None context and fails or mis-encodes.
from fastvideo.forward_context import set_forward_context
with set_forward_context(current_timestep=0, attn_metadata=None):
out = self.module(input_ids=ids, attention_mask=mask)
hidden = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
return _to_numpy(hidden.squeeze(0))
@@ -175,28 +205,22 @@ def _device_dtype(spec, instance, platform):
def build_torch_dit(spec, instance, platform):
cls = _resolve(spec.load_id)
ckpt = _require_checkpoint(spec)
_require_checkpoint(spec)
device, dtype = _device_dtype(spec, instance, platform)
# BRINGUP risk A: the exact construction path. FastVideo models load via from_pretrained or a
# TransformerLoader().load(ckpt, fastvideo_args); confirm which on the box.
module = cls.from_pretrained(ckpt)
module = _load_via_fastvideo("TransformerLoader", spec) # NOT from_pretrained (it has none)
return TorchWanDiT(module, device=device, dtype=dtype)
def build_torch_vae(spec, instance, platform):
cls = _resolve(spec.load_id)
ckpt = _require_checkpoint(spec)
_require_checkpoint(spec)
device, dtype = _device_dtype(spec, instance, platform)
module = cls.from_pretrained(ckpt) # BRINGUP risk A
module = _load_via_fastvideo("VAELoader", spec)
return TorchWanVAE(module, device=device, dtype=dtype)
def build_torch_text_encoder(spec, instance, platform):
cls = _resolve(spec.load_id)
ckpt = _require_checkpoint(spec)
_require_checkpoint(spec)
device, dtype = _device_dtype(spec, instance, platform)
from transformers import AutoTokenizer # BRINGUP risk E: confirm tokenizer class
tokenizer = AutoTokenizer.from_pretrained(ckpt)
module = cls.from_pretrained(ckpt) # BRINGUP risk A/E
module = _load_via_fastvideo("TextEncoderLoader", spec) # resolves UMT5 vs T5 from the config
tokenizer = _load_via_fastvideo("TokenizerLoader", spec)
return TorchT5Encoder(module, tokenizer, device=device, dtype=dtype)