diff --git a/v2/platform/backends/GPU_BRINGUP.md b/v2/platform/backends/GPU_BRINGUP.md index 13dc16921..657b207ff 100644 --- a/v2/platform/backends/GPU_BRINGUP.md +++ b/v2/platform/backends/GPU_BRINGUP.md @@ -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). diff --git a/v2/platform/backends/torch_adapters.py b/v2/platform/backends/torch_adapters.py index 75de0d199..ff9c3aa28 100644 --- a/v2/platform/backends/torch_adapters.py +++ b/v2/platform/backends/torch_adapters.py @@ -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)