From de803cb2505d511af3f7bb49ed65c7745ba720d7 Mon Sep 17 00:00:00 2001 From: William Lin Date: Tue, 12 May 2026 14:49:03 -0700 Subject: [PATCH] [feat]: MagiHuman DiT (transformer) port + parity tests (3/8) (#1297) --- fastvideo/configs/models/dits/__init__.py | 3 +- fastvideo/configs/models/dits/magi_human.py | 110 +++ fastvideo/models/dits/magi_human.py | 867 ++++++++++++++++++ .../helpers/magi_human_upstream.py | 334 +++++++ .../test_magi_human_distill_parity.py | 245 +++++ .../magi_human/test_magi_human_parity.py | 281 ++++++ 6 files changed, 1839 insertions(+), 1 deletion(-) create mode 100644 fastvideo/configs/models/dits/magi_human.py create mode 100644 fastvideo/models/dits/magi_human.py create mode 100644 tests/local_tests/helpers/magi_human_upstream.py create mode 100644 tests/local_tests/magi_human/test_magi_human_distill_parity.py create mode 100644 tests/local_tests/magi_human/test_magi_human_parity.py diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index b6817e254..2c3c059b2 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -5,6 +5,7 @@ from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config from fastvideo.configs.models.dits.longcat import LongCatVideoConfig from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig +from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig from fastvideo.configs.models.dits.stable_audio import StableAudioConfig from fastvideo.configs.models.dits.wanvideo import WanVideoConfig from fastvideo.configs.models.dits.hyworld import HYWorldConfig @@ -13,5 +14,5 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", - "StableAudioConfig" + "MagiHumanVideoConfig", "StableAudioConfig" ] diff --git a/fastvideo/configs/models/dits/magi_human.py b/fastvideo/configs/models/dits/magi_human.py new file mode 100644 index 000000000..2ae809886 --- /dev/null +++ b/fastvideo/configs/models/dits/magi_human.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Architecture / model config for the daVinci-MagiHuman DiT. + +The MagiHuman base DiT is a 15B-parameter single-stream transformer that +jointly denoises video, audio, and text tokens in one flat sequence. Layout +details verified against GAIR/daVinci-MagiHuman's base/ shards (2026-04-24). + +This file captures only configuration. The module implementation lives in +fastvideo/models/dits/magi_human.py and the pipeline wiring in +fastvideo/pipelines/basic/magi_human/. +""" +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def _is_block_layer(n: str, m) -> bool: + # Match "block.layers." — the FSDP shard boundary for MagiHuman. + parts = n.split(".") + return (len(parts) >= 3 and parts[0] == "block" and parts[1] == "layers" and str.isdigit(parts[2])) + + +@dataclass +class MagiHumanArchConfig(DiTArchConfig): + """MagiHuman base DiT architecture constants. + + **Scope contract:** fields here must match the `transformer/config.json` + emitted by `scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py` + 1:1, and both are sourced from the upstream Python reference + `inference/common/config.py::ModelConfig` (the HF root `config.json` + is empty so the Python source is canonical). Pipeline-level knobs + (VAE stride, fps, num_inference_steps, CFG scales, flow_shift, + t5_gemma_target_length) and data-proxy knobs (coords_style, + frame_receptive_field, ref_audio_offset, text_offset) live on + `MagiHumanBaseConfig`, NOT here. + + `param_names_mapping` is intentionally empty: the FastVideo implementation + keeps the same module tree as the reference (`adapter.*`, + `block.layers..*`, `final_linear_{video,audio}.*`, + `final_norm_{video,audio}.*`), so converted weights load directly. + """ + + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_block_layer]) + + # No renames needed — the FastVideo module mirrors the reference names. + param_names_mapping: dict = field(default_factory=dict) + reverse_param_names_mapping: dict = field(default_factory=dict) + lora_param_names_mapping: dict = field(default_factory=dict) + + # --- transformer shape --- + num_layers: int = 40 + hidden_size: int = 5120 + head_dim: int = 128 + num_query_groups: int = 8 # num_heads_kv (GQA) + + # --- modality channels --- + # video_in_channels = z_dim (48) * patch_size product (1*2*2=4), so the + # embedder receives 192 per token. text_in_channels is T5Gemma-9B's + # encoder hidden size. + video_in_channels: int = 192 + audio_in_channels: int = 64 + text_in_channels: int = 3584 + + # --- block-level architecture switches --- + # Sandwich MoE: first and last 4 layers have per-modality experts + # (video/audio/text), middle layers share a single set of weights. + mm_layers: tuple[int, ...] = (0, 1, 2, 3, 36, 37, 38, 39) + local_attn_layers: tuple[int, ...] = () + gelu7_layers: tuple[int, ...] = (0, 1, 2, 3) + post_norm_layers: tuple[int, ...] = () + enable_attn_gating: bool = True + activation_type: str = "swiglu7" + + # --- DiT patching (upstream `ModelConfig`-equivalent; NOT the VAE + # stride, which is pipeline-level). --- + patch_size: tuple[int, int, int] = (1, 2, 2) + spatial_rope_interpolation: str = "extra" + + # --- TReAD (token routing + early drop). Flattened from the upstream + # nested `tread_config` dict so it round-trips through + # `update_model_arch` cleanly. --- + tread_selection_rate: float = 0.5 + tread_start_layer_idx: int = 2 + tread_end_layer_idx: int = 25 + + # --- derived fields (populated in __post_init__) --- + num_attention_heads: int = 0 # hidden_size / head_dim + num_heads_kv: int = 0 # == num_query_groups + in_channels: int = 0 # mirror of video_in_channels (FastVideo contract) + out_channels: int = 0 # mirror of video_in_channels + + def __post_init__(self) -> None: + super().__post_init__() + self.num_attention_heads = self.hidden_size // self.head_dim + self.num_heads_kv = self.num_query_groups + self.in_channels = self.video_in_channels + self.out_channels = self.video_in_channels + # num_channels_latents is the VAE latent z_dim (48 for Wan 2.2 TI2V-5B). + # We don't declare z_dim on the arch config (it's a VAE property), + # but we still set num_channels_latents for the BaseDiT contract. + self.num_channels_latents = 48 + + +@dataclass +class MagiHumanVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=MagiHumanArchConfig) + + prefix: str = "magi_human" diff --git a/fastvideo/models/dits/magi_human.py b/fastvideo/models/dits/magi_human.py new file mode 100644 index 000000000..4682326bf --- /dev/null +++ b/fastvideo/models/dits/magi_human.py @@ -0,0 +1,867 @@ +# SPDX-License-Identifier: Apache-2.0 +"""daVinci-MagiHuman DiT (base variant). + +Ported from https://github.com/GAIR-NLP/daVinci-MagiHuman +(inference/model/dit/dit_module.py, ~950 lines in the reference). + +Architecture summary (verified against GAIR/daVinci-MagiHuman/base/ weights): + + - 40 transformer layers, hidden 5120, head_dim 128. + - GQA with 40 query heads and 8 KV heads. + - Multi-modality "sandwich": layers 0..3 and 36..39 use 3-way modality + experts (video/audio/text) packed inside each linear as + weight[..., out * 3, in]. Middle layers share a single expert. + - Per-head attention gating: the QKV projection emits an extra + num_heads_q channels that are sigmoid-gated onto the attention output. + - Activation is GELU7 on layers 0..3 (non-gated, intermediate=4*hidden) + and SwiGLU7 elsewhere (gated, intermediate=int(hidden*4*2/3)//4*4). + - Position encoding is an element-wise Fourier embedding over 9-column + coords (t,h,w + original TxHxW + reference TxHxW), not a standard + 1D/3D RoPE. + - Forward takes a flat concatenated token stream (video first, then + audio, then text) plus a modality map; the internal ModalityDispatcher + permutes by modality before each linear so per-expert chunks line up. + +Deviations from the "use fastvideo.layers primitives everywhere" guideline +in the add-model skill: + + - The packed-expert linears store weight as [out * num_experts, in]. + FastVideo's ReplicatedLinear does not model this layout; we use raw + nn.Parameter with a small wrapper below. This is deliberate and scoped + to this DiT: ReplicatedLinear still handles the adapter.* embedders + and final_linear_{video,audio} (single-expert) here. + - Self-attention is full-sequence and crosses modalities inside the flat + concat stream; DistributedAttention assumes a clean spatial-sequence + layout, so for the first port we use torch SDPA. Multi-GPU sequence + parallelism is a follow-up. + - torch.compile via magi_compiler is replaced with a plain nn.Module. + +For the full history and shape-by-shape verification notes, see +.claude/skills/add-model/SKILL.md and the scaffold PR description. +""" +from __future__ import annotations + +import math +from dataclasses import dataclass +from enum import IntEnum +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.attention import LocalAttention +from fastvideo.configs.models.dits.magi_human import ( + MagiHumanArchConfig, + MagiHumanVideoConfig, +) +from fastvideo.layers.rotary_embedding import _apply_rotary_emb +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum + + +# --------------------------------------------------------------------------- +# Enums +# --------------------------------------------------------------------------- + + +class Modality(IntEnum): + VIDEO = 0 + AUDIO = 1 + TEXT = 2 + + +# --------------------------------------------------------------------------- +# Activations +# --------------------------------------------------------------------------- + + +def swiglu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor: + """Gated swish-GLU with OpenAI-OSS-style limits and +1 linear bias.""" + in_dtype = x.dtype + x = x.to(torch.float32) + x_glu, x_linear = x[..., ::2], x[..., 1::2] + x_glu = x_glu.clamp(max=limit) + x_linear = x_linear.clamp(min=-limit, max=limit) + out_glu = x_glu * torch.sigmoid(alpha * x_glu) + return (out_glu * (x_linear + 1)).to(in_dtype) + + +def gelu7(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor: + in_dtype = x.dtype + x = x.to(torch.float32).clamp(max=limit) + return (x * torch.sigmoid(alpha * x)).to(in_dtype) + + +# --------------------------------------------------------------------------- +# Modality dispatcher +# --------------------------------------------------------------------------- + + +class ModalityDispatcher: + """Permute a flat token stream so same-modality tokens are contiguous. + + The DiT's multi-expert linears apply a different weight chunk per modality. + Instead of carrying a branch inside each Linear, we pre-permute tokens so + each chunk sees a contiguous slice, then un-permute before computing + RoPE/attention across the full sequence. + """ + + def __init__(self, modality_mapping: torch.Tensor, num_modalities: int): + self.modality_mapping = modality_mapping + self.num_modalities = num_modalities + self.permute_mapping = torch.argsort(modality_mapping) + self.inv_permute_mapping = torch.argsort(self.permute_mapping) + permuted = modality_mapping[self.permute_mapping] + self.group_size = torch.bincount(permuted, minlength=num_modalities).to(torch.int32) + self.group_size_cpu: list[int] = [int(x) for x in self.group_size.cpu().tolist()] + + def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]: + return list(torch.split(x, self.group_size_cpu, dim=0)) + + def undispatch(self, *chunks: torch.Tensor) -> torch.Tensor: + return torch.cat(chunks, dim=0) + + @staticmethod + def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor: + return x[permute_mapping] + + @staticmethod + def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor: + return x[inv_permute_mapping] + + +# --------------------------------------------------------------------------- +# Norms, rotary embed +# --------------------------------------------------------------------------- + + +class MultiModalityRMSNorm(nn.Module): + """RMSNorm with optional per-modality scale. + + When num_modality == 1, behaves identically to a standard RMSNorm with + weight initialized to zero (effective weight is 1 + weight, hence the + learnable +1 offset baked into the forward path). When num_modality > 1, + the weight tensor packs per-modality scales along its flat axis and the + dispatcher selects the right chunk per modality. + """ + + def __init__(self, dim: int, eps: float = 1e-6, num_modality: int = 1): + super().__init__() + self.dim = dim + self.eps = eps + self.num_modality = num_modality + # Always stored in fp32; matches the reference initialization. + self.weight = nn.Parameter(torch.zeros(dim * num_modality, dtype=torch.float32)) + + def _rms(self, x: torch.Tensor) -> torch.Tensor: + t = x.float() + return t * torch.rsqrt(t.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + def forward( + self, + x: torch.Tensor, + modality_dispatcher: Optional[ModalityDispatcher] = None, + ) -> torch.Tensor: + original_dtype = x.dtype + t = self._rms(x) + if self.num_modality == 1: + return (t * (self.weight + 1)).to(original_dtype) + assert modality_dispatcher is not None, ( + "MultiModalityRMSNorm with num_modality>1 requires a dispatcher" + ) + weight_chunks = self.weight.chunk(self.num_modality, dim=0) + parts = modality_dispatcher.dispatch(t) + for i in range(self.num_modality): + parts[i] = parts[i] * (weight_chunks[i] + 1) + return modality_dispatcher.undispatch(*parts).to(original_dtype) + + +def _freq_bands(num_bands: int, temperature: float = 10000.0) -> torch.Tensor: + exp = torch.arange(0, num_bands, 1, dtype=torch.int64).float() / num_bands + return 1.0 / (temperature ** exp) + + +class ElementWiseFourierEmbed(nn.Module): + """Element-wise Fourier embedding over 9-column coords (t, h, w, T, H, W, + ref_T, ref_H, ref_W). Produces a per-token positional embedding that + acts as the RoPE angle input for attention. + + Weight: `bands` of shape `[dim // 8]` (fixed at init via freq_bands). + """ + + def __init__( + self, + dim: int, + temperature: float = 10000.0, + dtype: torch.dtype = torch.float32, + ): + super().__init__() + self.dim = dim + bands = _freq_bands(dim // 8, temperature=temperature).to(dtype) + # `register_buffer` so state_dict keeps it, matching upstream naming. + self.register_buffer("bands", bands) + + def forward(self, coords: torch.Tensor) -> torch.Tensor: + # coords: [L, 9] = (t, h, w, T, H, W, ref_T, ref_H, ref_W) + coords_xyz = coords[:, :3] + sizes = coords[:, 3:6] + refs = coords[:, 6:9] + + scales = (refs - 1) / (sizes - 1) + scales[(refs == 1) & (sizes == 1)] = 1 + # Center H and W (leave time uncentered). + centers = (sizes - 1) / 2 + centers[:, 0] = 0 + coords_xyz = coords_xyz - centers + + proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * self.bands # [L, 3, B] + sin_proj = proj.sin() + cos_proj = proj.cos() + return torch.cat((sin_proj, cos_proj), dim=1).flatten(1) + + +# --------------------------------------------------------------------------- +# Packed-expert linear +# --------------------------------------------------------------------------- + + +class PackedExpertLinear(nn.Module): + """Linear where the weight is packed per-modality along the output axis. + + Shapes: + weight: [out_features * num_experts, in_features] + bias: [out_features * num_experts] (optional) + + When `num_experts == 1`, behaves exactly like `nn.Linear`. When + `num_experts > 1`, `forward` dispatches the input via the supplied + `ModalityDispatcher`, applies the per-modality weight/bias chunk, and + gathers the outputs in original order. + + Why not use `ReplicatedLinear`? Because the packed-expert layout is not + what ReplicatedLinear (or any other fastvideo.layers.linear) is wired + for. Using raw `nn.Parameter` keeps weight loading trivial (names map + 1:1 to the upstream checkpoint) and avoids quantization-path assumptions + that don't match this layout. + """ + + def __init__( + self, + in_features: int, + out_features: int, + num_experts: int = 1, + bias: bool = False, + dtype: torch.dtype = torch.bfloat16, + ): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.num_experts = num_experts + self.use_bias = bias + self.weight = nn.Parameter( + torch.empty(out_features * num_experts, in_features, dtype=dtype) + ) + if bias: + self.bias = nn.Parameter( + torch.empty(out_features * num_experts, dtype=dtype) + ) + else: + self.register_parameter("bias", None) + + def forward( + self, + x: torch.Tensor, + modality_dispatcher: Optional[ModalityDispatcher] = None, + ) -> torch.Tensor: + if self.num_experts == 1: + return F.linear(x, self.weight, self.bias) + assert modality_dispatcher is not None, ( + "PackedExpertLinear with num_experts>1 requires a dispatcher" + ) + parts = modality_dispatcher.dispatch(x) + w_chunks = self.weight.chunk(self.num_experts, dim=0) + b_chunks = ( + self.bias.chunk(self.num_experts, dim=0) + if self.bias is not None else [None] * self.num_experts + ) + for i in range(self.num_experts): + parts[i] = F.linear(parts[i], w_chunks[i], b_chunks[i]) + return modality_dispatcher.undispatch(*parts) + + +# --------------------------------------------------------------------------- +# Attention & MLP +# --------------------------------------------------------------------------- + + +@dataclass +class AttentionSubConfig: + hidden_size: int + num_heads_q: int + num_heads_kv: int + head_dim: int + num_modality: int + enable_attn_gating: bool + use_local_attn: bool = False + frame_receptive_field: int = 11 + + +class MagiAttention(nn.Module): + """Self-attention with GQA + optional per-head sigmoid gating.""" + + def __init__(self, cfg: AttentionSubConfig): + super().__init__() + self.cfg = cfg + self.gating_size = cfg.num_heads_q if cfg.enable_attn_gating else 0 + qkv_out = ( + cfg.num_heads_q * cfg.head_dim + + 2 * cfg.num_heads_kv * cfg.head_dim + + self.gating_size + ) + self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality) + self.linear_qkv = PackedExpertLinear( + cfg.hidden_size, qkv_out, num_experts=cfg.num_modality, bias=False, + ) + self.linear_proj = PackedExpertLinear( + cfg.num_heads_q * cfg.head_dim, cfg.hidden_size, + num_experts=cfg.num_modality, bias=False, + ) + self.q_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality) + self.k_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality) + + self.q_size = cfg.num_heads_q * cfg.head_dim + self.kv_size = cfg.num_heads_kv * cfg.head_dim + + self.attn = LocalAttention( + num_heads=cfg.num_heads_q, + head_size=cfg.head_dim, + num_kv_heads=cfg.num_heads_kv, + causal=False, + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ), + ) + + def configure_local_attention( + self, + *, + enabled: bool, + frame_receptive_field: int = 11, + ) -> None: + self.cfg.use_local_attn = enabled + self.cfg.frame_receptive_field = frame_receptive_field + + def _sdpa(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """Run SDPA on [L, H, D] tensors and return [L, Hq, D].""" + if q.numel() == 0: + return q.new_empty(q.shape) + out = F.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0).contiguous(), + k.transpose(0, 1).unsqueeze(0).contiguous(), + v.transpose(0, 1).unsqueeze(0).contiguous(), + enable_gqa=self.cfg.num_heads_q != self.cfg.num_heads_kv, + ) + return out.squeeze(0).transpose(0, 1).contiguous() + + def _local_window_attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + num_video_tokens: int, + num_frames: int, + ) -> torch.Tensor: + """Approximate upstream FFAHandler block accumulation with SDPA. + + SR-1080p's reference kernel sums three independently-normalized + attention contributions: + + * video frame queries -> local-window video keys; + * all video queries -> all audio+text keys; + * all audio+text queries -> full sequence keys. + + This method mirrors that accumulator semantics with ordinary SDPA + slices. It is intentionally scoped to single-process inference; layers + without ``use_local_attn`` keep the existing full LocalAttention path. + """ + if num_frames <= 0 or num_video_tokens <= 0: + return self._sdpa(q, k, v) + if num_video_tokens % num_frames != 0: + raise ValueError( + f"MagiHuman local attention expects video tokens divisible by " + f"frames, got {num_video_tokens=} and {num_frames=}." + ) + + token_per_frame = num_video_tokens // num_frames + out = torch.zeros( + q.shape[0], + self.cfg.num_heads_q, + self.cfg.head_dim, + device=q.device, + dtype=q.dtype, + ) + rf = int(self.cfg.frame_receptive_field) + + q_video = q[:num_video_tokens] + k_video = k[:num_video_tokens] + v_video = v[:num_video_tokens] + for frame_idx in range(num_frames): + q_start = frame_idx * token_per_frame + q_end = q_start + token_per_frame + k_start = max(0, (frame_idx - rf) * token_per_frame) + k_end = min(num_video_tokens, (frame_idx + rf + 1) * token_per_frame) + out[q_start:q_end] = self._sdpa( + q_video[q_start:q_end], + k_video[k_start:k_end], + v_video[k_start:k_end], + ) + + if num_video_tokens < q.shape[0]: + k_at = k[num_video_tokens:] + v_at = v[num_video_tokens:] + out[:num_video_tokens] = out[:num_video_tokens] + self._sdpa( + q[:num_video_tokens], + k_at, + v_at, + ) + out[num_video_tokens:] = self._sdpa( + q[num_video_tokens:], + k, + v, + ) + return out + + def forward( + self, + hidden_states: torch.Tensor, + rope: torch.Tensor, + permute_mapping: torch.Tensor, + inv_permute_mapping: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + num_video_tokens: int | None = None, + num_frames: int | None = None, + ) -> torch.Tensor: + orig_dtype = self.linear_qkv.weight.dtype + h = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(orig_dtype) + qkv = self.linear_qkv(h, modality_dispatcher=modality_dispatcher).float() + q, k, v, g = torch.split( + qkv, [self.q_size, self.kv_size, self.kv_size, self.gating_size], dim=-1, + ) + q = q.view(-1, self.cfg.num_heads_q, self.cfg.head_dim) + k = k.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim) + v = v.view(-1, self.cfg.num_heads_kv, self.cfg.head_dim) + g = g.view(-1, self.cfg.num_heads_q, 1) if self.gating_size else None + + q = self.q_norm(q, modality_dispatcher=modality_dispatcher) + k = self.k_norm(k, modality_dispatcher=modality_dispatcher) + + # Un-permute before RoPE + attention so positional order reflects + # the original (video, audio, text) concat — matches reference. + q = ModalityDispatcher.inv_permute(q, inv_permute_mapping) + k = ModalityDispatcher.inv_permute(k, inv_permute_mapping) + v = ModalityDispatcher.inv_permute(v, inv_permute_mapping) + if g is not None: + g = ModalityDispatcher.inv_permute(g, inv_permute_mapping) + + # Element-wise Fourier embed packs sin/cos of 3 axes into a single + # `rope` tensor. Match reference's split: + # sin_emb, cos_emb = rope.tensor_split(2, -1) + # Reference passes (cos_emb, sin_emb) but splits sin first — replicated + # exactly so weight parity holds. Partial RoPE: rope dim is + # 6 * (head_dim // 8) = 96 < head_dim (128), so the trailing 32 + # head_dim positions stay unrotated, matching the reference. + sin_emb, cos_emb = rope.tensor_split(2, -1) + rot_dim = cos_emb.shape[-1] * 2 + q_rot = _apply_rotary_emb(q[..., :rot_dim], cos_emb, sin_emb, is_neox_style=True) + k_rot = _apply_rotary_emb(k[..., :rot_dim], cos_emb, sin_emb, is_neox_style=True) + if rot_dim < q.shape[-1]: + q = torch.cat([q_rot, q[..., rot_dim:]], dim=-1) + k = torch.cat([k_rot, k[..., rot_dim:]], dim=-1) + else: + q, k = q_rot, k_rot + + # Run SDPA via FastVideo's LocalAttention so the backend selection + # (SDPA / FlashAttn / SLA / SageAttn) flows through the standard + # configurable path. GQA is handled inside the SDPA backend via + # `enable_gqa=True` when num_heads_q != num_heads_kv, so we no + # longer need the manual `repeat_interleave` here. + # Attention math runs at orig_dtype (bf16 in production and in the + # parity test, since PackedExpertLinear's default is bf16, matching + # upstream BaseLinear at dit_module.py:330). The gating multiply + # promotes back to fp32 implicitly via PyTorch's type-promotion + # rules: bf16_attn_out * sigmoid(fp32_g) -> fp32, mirroring upstream + # dit_module.py:649. + q = q.to(orig_dtype) + k = k.to(orig_dtype) + v = v.to(orig_dtype) + if self.cfg.use_local_attn: + if num_video_tokens is None or num_frames is None: + raise ValueError("MagiHuman local attention requires video token/frame metadata.") + out = self._local_window_attention( + q, + k, + v, + num_video_tokens=num_video_tokens, + num_frames=num_frames, + ) + else: + out = self.attn(q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0)).squeeze(0) + + out = ModalityDispatcher.permute(out, permute_mapping) + if g is not None: + g = ModalityDispatcher.permute(g, permute_mapping) + out = out * torch.sigmoid(g) + out = out.reshape(-1, self.cfg.num_heads_q * self.cfg.head_dim).to(orig_dtype) + return self.linear_proj(out, modality_dispatcher=modality_dispatcher) + + +@dataclass +class MLPSubConfig: + hidden_size: int + intermediate_size: int + activation: str # "swiglu7" or "gelu7" + num_modality: int + gated: bool + + +class MagiMLP(nn.Module): + def __init__(self, cfg: MLPSubConfig): + super().__init__() + self.cfg = cfg + self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality) + up_out = cfg.intermediate_size * 2 if cfg.gated else cfg.intermediate_size + self.up_gate_proj = PackedExpertLinear( + cfg.hidden_size, up_out, num_experts=cfg.num_modality, bias=False, + ) + self.down_proj = PackedExpertLinear( + cfg.intermediate_size, cfg.hidden_size, + num_experts=cfg.num_modality, bias=False, + ) + self._act = swiglu7 if cfg.activation == "swiglu7" else gelu7 + + def forward( + self, + x: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + ) -> torch.Tensor: + orig_dtype = self.up_gate_proj.weight.dtype + x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(orig_dtype) + x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).float() + x = self._act(x).to(orig_dtype) + x = self.down_proj(x, modality_dispatcher=modality_dispatcher).float() + return x + + +class MagiTransformerLayer(nn.Module): + def __init__(self, arch: MagiHumanArchConfig, layer_idx: int): + super().__init__() + num_modality = 3 if layer_idx in arch.mm_layers else 1 + self.post_norm = layer_idx in arch.post_norm_layers + self.layer_idx = layer_idx + + self.attention = MagiAttention(AttentionSubConfig( + hidden_size=arch.hidden_size, + num_heads_q=arch.num_attention_heads, + num_heads_kv=arch.num_heads_kv, + head_dim=arch.head_dim, + num_modality=num_modality, + enable_attn_gating=arch.enable_attn_gating, + use_local_attn=layer_idx in arch.local_attn_layers, + )) + + is_gelu7 = layer_idx in arch.gelu7_layers + if is_gelu7: + intermediate = arch.hidden_size * 4 + gated = False + activation = "gelu7" + else: + intermediate = (arch.hidden_size * 4 * 2 // 3) // 4 * 4 + gated = True + activation = "swiglu7" + + self.mlp = MagiMLP(MLPSubConfig( + hidden_size=arch.hidden_size, + intermediate_size=intermediate, + activation=activation, + num_modality=num_modality, + gated=gated, + )) + + if self.post_norm: + self.attn_post_norm = MultiModalityRMSNorm(arch.hidden_size, num_modality=num_modality) + self.mlp_post_norm = MultiModalityRMSNorm(arch.hidden_size, num_modality=num_modality) + + def forward( + self, + hidden_states: torch.Tensor, + rope: torch.Tensor, + permute_mapping: torch.Tensor, + inv_permute_mapping: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + num_video_tokens: int | None = None, + num_frames: int | None = None, + ) -> torch.Tensor: + attn_out = self.attention( + hidden_states, rope, permute_mapping, inv_permute_mapping, modality_dispatcher, + num_video_tokens=num_video_tokens, + num_frames=num_frames, + ) + if self.post_norm: + attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher) + hidden_states = hidden_states + attn_out + + mlp_out = self.mlp(hidden_states, modality_dispatcher=modality_dispatcher) + if self.post_norm: + mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher) + return hidden_states + mlp_out + + +# --------------------------------------------------------------------------- +# Adapter (per-modality embedders + Fourier RoPE producer) +# --------------------------------------------------------------------------- + + +class MagiAdapter(nn.Module): + def __init__(self, arch: MagiHumanArchConfig): + super().__init__() + # Embedders stay in fp32 to match the reference dtype exactly. + self.video_embedder = nn.Linear( + arch.video_in_channels, arch.hidden_size, bias=True, dtype=torch.float32, + ) + self.text_embedder = nn.Linear( + arch.text_in_channels, arch.hidden_size, bias=True, dtype=torch.float32, + ) + self.audio_embedder = nn.Linear( + arch.audio_in_channels, arch.hidden_size, bias=True, dtype=torch.float32, + ) + self.rope = ElementWiseFourierEmbed(arch.head_dim) + # RoPE cache: coords_mapping is the same tensor object across timesteps + # in the denoising loop, so data_ptr()+shape+dtype+device is a fast, + # collision-free key that avoids recomputing the Fourier embed each step. + self._cached_rope: Optional[torch.Tensor] = None + self._cached_rope_key: Optional[tuple] = None + + def _rope_cache_key(self, t: torch.Tensor) -> tuple: + return (t.data_ptr(), t.shape, t.dtype, t.device) + + def forward( + self, + x: torch.Tensor, + coords_mapping: torch.Tensor, + video_mask: torch.Tensor, + audio_mask: torch.Tensor, + text_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + key = self._rope_cache_key(coords_mapping) + if key != self._cached_rope_key: + self._cached_rope = self.rope(coords_mapping) + self._cached_rope_key = key + rope = self._cached_rope + # Embedder dtypes may differ from x's dtype when FastVideo's FSDP + # loader casts all weights to `pipeline_config.precision` (bf16). + # Match the weight dtype per modality. + v_w = self.video_embedder.weight + a_w = self.audio_embedder.weight + t_w = self.text_embedder.weight + out = torch.zeros( + x.shape[0], self.video_embedder.out_features, + device=x.device, dtype=v_w.dtype, + ) + out[text_mask] = self.text_embedder( + x[text_mask, : self.text_embedder.in_features].to(t_w.dtype) + ).to(out.dtype) + out[audio_mask] = self.audio_embedder( + x[audio_mask, : self.audio_embedder.in_features].to(a_w.dtype) + ).to(out.dtype) + out[video_mask] = self.video_embedder( + x[video_mask, : self.video_embedder.in_features].to(v_w.dtype) + ).to(out.dtype) + return out, rope + + +# --------------------------------------------------------------------------- +# Top-level DiT +# --------------------------------------------------------------------------- + + +class _TransformerBlock(nn.Module): + """Thin ModuleList wrapper to keep the 'block.layers.' state_dict + naming identical to the upstream checkpoint (which uses a magi_compile + decorator producing `block.layers..*`).""" + + def __init__(self, arch: MagiHumanArchConfig): + super().__init__() + self.layers = nn.ModuleList([ + MagiTransformerLayer(arch, i) for i in range(arch.num_layers) + ]) + + def configure_local_attention( + self, + local_attn_layers: tuple[int, ...], + frame_receptive_field: int = 11, + ) -> None: + enabled_layers = set(local_attn_layers) + for idx, layer in enumerate(self.layers): + layer.attention.configure_local_attention( + enabled=idx in enabled_layers, + frame_receptive_field=frame_receptive_field, + ) + + def forward( + self, + x: torch.Tensor, + rope: torch.Tensor, + permute_mapping: torch.Tensor, + inv_permute_mapping: torch.Tensor, + modality_dispatcher: ModalityDispatcher, + num_video_tokens: int | None = None, + num_frames: int | None = None, + ) -> torch.Tensor: + for layer in self.layers: + x = layer( + x, + rope, + permute_mapping, + inv_permute_mapping, + modality_dispatcher, + num_video_tokens=num_video_tokens, + num_frames=num_frames, + ) + return x + + +_CFG = MagiHumanVideoConfig() + + +class MagiHumanDiT(BaseDiT): + """Top-level DiT for daVinci-MagiHuman (base). + + Forward signature mirrors the reference `DiTModel.forward`: it takes a + flat token stream, its per-token coords and modality mapping, and + returns per-modality outputs packed into a max-channel-width tensor. + + This scaffold is single-GPU only; the `ulysses_scheduler().dispatch(...)` + sequence-parallel wrapping in the reference has no equivalent here yet. + """ + + # BaseDiT requires these class attrs. Source them from the config so + # they stay in sync with MagiHumanVideoConfig edits. + _fsdp_shard_conditions = _CFG._fsdp_shard_conditions + _compile_conditions = _CFG._compile_conditions + _supported_attention_backends = _CFG._supported_attention_backends + param_names_mapping = _CFG.param_names_mapping + reverse_param_names_mapping = _CFG.reverse_param_names_mapping + lora_param_names_mapping = _CFG.lora_param_names_mapping + + def __init__(self, config: MagiHumanVideoConfig, hf_config: dict | None = None, **kwargs): + super().__init__(config=config, hf_config=hf_config or {}) + arch: MagiHumanArchConfig = getattr(config, "arch_config", config) + self.arch = arch + + # BaseDiT contract instance vars. + self.hidden_size = arch.hidden_size + self.num_attention_heads = arch.num_attention_heads + self.num_channels_latents = arch.num_channels_latents + + self.adapter = MagiAdapter(arch) + self.block = _TransformerBlock(arch) + self.final_norm_video = MultiModalityRMSNorm(arch.hidden_size) + self.final_norm_audio = MultiModalityRMSNorm(arch.hidden_size) + self.final_linear_video = nn.Linear( + arch.hidden_size, arch.video_in_channels, bias=False, dtype=torch.float32, + ) + self.final_linear_audio = nn.Linear( + arch.hidden_size, arch.audio_in_channels, bias=False, dtype=torch.float32, + ) + # Dispatcher + mask cache: modality_mapping is the same tensor object + # across all timesteps in the denoising loop; data_ptr()+shape+dtype+device + # is a fast, collision-free key that avoids rebuilding ModalityDispatcher + # (which calls argsort + bincount) on every forward call. + self._cached_dispatcher: Optional[ModalityDispatcher] = None + self._cached_video_mask: Optional[torch.Tensor] = None + self._cached_audio_mask: Optional[torch.Tensor] = None + self._cached_text_mask: Optional[torch.Tensor] = None + self._cached_modality_key: Optional[tuple] = None + + def configure_local_attention( + self, + local_attn_layers: tuple[int, ...] | list[int], + frame_receptive_field: int = 11, + ) -> None: + layers = tuple(int(layer) for layer in local_attn_layers) + self.arch.local_attn_layers = layers + self.block.configure_local_attention(layers, frame_receptive_field) + + def _modality_cache_key(self, t: torch.Tensor) -> tuple: + return (t.data_ptr(), t.shape, t.dtype, t.device) + + def forward( + self, + x: torch.Tensor, + coords_mapping: torch.Tensor, + modality_mapping: torch.Tensor, + ) -> torch.Tensor: + """ + Args: + x: [L, max(V_ch, A_ch, T_ch)] + coords_mapping: [L, 9] + modality_mapping: [L] (int in {VIDEO, AUDIO, TEXT}) + Returns: + out: [L, max(V_ch, A_ch)] with video channels in video slots and + audio channels in audio slots; text slots are zero. + """ + key = self._modality_cache_key(modality_mapping) + if key != self._cached_modality_key: + self._cached_dispatcher = ModalityDispatcher(modality_mapping, num_modalities=3) + self._cached_video_mask = modality_mapping == Modality.VIDEO + self._cached_audio_mask = modality_mapping == Modality.AUDIO + self._cached_text_mask = modality_mapping == Modality.TEXT + self._cached_modality_key = key + dispatcher = self._cached_dispatcher + video_mask = self._cached_video_mask + audio_mask = self._cached_audio_mask + text_mask = self._cached_text_mask + num_video_tokens = int(video_mask.sum().item()) + if num_video_tokens: + num_frames = int(coords_mapping[:num_video_tokens, 0].max().item()) + 1 + else: + num_frames = 0 + + x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask) + # Keep the residual stream in adapter dtype (fp32) entering the block. + # Upstream daVinci-MagiHuman dit_module.py:923 casts to params_dtype, + # which is fp32 by default; each layer's pre_norm.to(bf16) handles + # the bf16 internal-compute boundary, and linear_proj outputs bf16 + # which gets promoted back to fp32 by the residual addition. Casting + # the residual to bf16 here degrades the cross-layer accumulator and + # compounds visibly over 40 layers in pipeline parity. + x = ModalityDispatcher.permute(x, dispatcher.permute_mapping) + + x = self.block( + x, rope, + permute_mapping=dispatcher.permute_mapping, + inv_permute_mapping=dispatcher.inv_permute_mapping, + modality_dispatcher=dispatcher, + num_video_tokens=num_video_tokens, + num_frames=num_frames, + ) + x = ModalityDispatcher.inv_permute(x, dispatcher.inv_permute_mapping) + + x_video = x[video_mask].to(self.final_norm_video.weight.dtype) + x_video = self.final_norm_video(x_video) + x_video = self.final_linear_video(x_video) + + x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype) + x_audio = self.final_norm_audio(x_audio) + x_audio = self.final_linear_audio(x_audio) + + max_ch = max(self.arch.video_in_channels, self.arch.audio_in_channels) + out = torch.zeros(x.shape[0], max_ch, device=x.device, dtype=x.dtype) + out[video_mask, : self.arch.video_in_channels] = x_video.to(out.dtype) + out[audio_mask, : self.arch.audio_in_channels] = x_audio.to(out.dtype) + return out + + +EntryClass = MagiHumanDiT diff --git a/tests/local_tests/helpers/magi_human_upstream.py b/tests/local_tests/helpers/magi_human_upstream.py new file mode 100644 index 000000000..ee65707f9 --- /dev/null +++ b/tests/local_tests/helpers/magi_human_upstream.py @@ -0,0 +1,334 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stubs to make upstream daVinci-MagiHuman code importable in-process. + +The upstream DiT (daVinci-MagiHuman/inference/model/dit/dit_module.py) hard- +imports SandAI's internal `magi_compiler` + a distributed-runtime init +that requires `torchrun`. Neither is available in a single-process +parity test. This module installs the minimum stubs to let the upstream +DiT load and run on a single GPU with cp_world_size == 1 (which makes +Ulysses's scatter/gather a no-op). + +Use: + from tests.local_tests.helpers.magi_human_upstream import ( + install_stubs, load_upstream_dit, + ) + install_stubs() + model = load_upstream_dit(base_shard_dir, device=torch.device("cuda")) +""" +from __future__ import annotations + +import importlib +import sys +import types +from pathlib import Path +from typing import Callable + + +# --------------------------------------------------------------------------- +# magi_compiler stubs — identity decorators. +# --------------------------------------------------------------------------- + + +def _install_magi_compiler_stub() -> None: + """Stub magi_compiler + register the `torch.ops.infra.*` ops upstream + calls via `torch.ops`. + + Upstream code decorates plain Python fns with + `@magi_register_custom_op(name="infra::flash_attn_func", ...)` and + then calls them as `torch.ops.infra.flash_attn_func(...)`. Our stub + decorator has to both (a) preserve the decorated fn for direct call + sites and (b) register the fn under the advertised torch.ops + namespace so `torch.ops.infra.*` resolves. + + For parity testing we route `infra::flash_attn_func` through + `F.scaled_dot_product_attention`, matching the FastVideo DiT's + kernel choice so drift measured in this test is architectural, + not kernel-dependent. + """ + if "magi_compiler" in sys.modules: + return + + import torch + import torch.nn.functional as F + + pkg = types.ModuleType("magi_compiler") + + def magi_compile(config_patch=None): + def decorator(cls_or_fn): + return cls_or_fn + return decorator + + # Create one Library per (namespace, schema) pair. Track by namespace + # so we don't define the same op twice on re-import. + _libs: dict[str, torch.library.Library] = {} + _defined: set[tuple[str, str]] = set() + + def _sdpa_flash_attn_func(q, k, v): + # Upstream shape: [batch=1, L, H, D]. SDPA expects [B, H, L, D] + # and no native GQA; expand K/V to match Q heads. + num_heads_q = q.shape[2] + num_heads_kv = k.shape[2] + if num_heads_q != num_heads_kv: + assert num_heads_q % num_heads_kv == 0 + repeat = num_heads_q // num_heads_kv + k = k.repeat_interleave(repeat, dim=2) + v = v.repeat_interleave(repeat, dim=2) + q2 = q.transpose(1, 2).contiguous() + k2 = k.transpose(1, 2).contiguous() + v2 = v.transpose(1, 2).contiguous() + out = F.scaled_dot_product_attention(q2, k2, v2) + return out.transpose(1, 2).contiguous() + + def _sdpa_segments(q, k, v, q_ranges, k_ranges): + # Upstream flex op shape: [L, H, D]. FFA accumulates each block's + # independently normalized attention output into the destination query + # slice. This SDPA fallback mirrors the accumulator semantics for + # SR-1080p parity tests without requiring SandAI's MagiAttention wheel. + out = torch.zeros( + q.shape[0], + q.shape[1], + q.shape[2], + dtype=q.dtype, + device=q.device, + ) + num_heads_q = q.shape[1] + num_heads_kv = k.shape[1] + for q_range, k_range in zip(q_ranges.tolist(), k_ranges.tolist()): + qs, qe = int(q_range[0]), int(q_range[1]) + ks, ke = int(k_range[0]), int(k_range[1]) + q_block = q[qs:qe] + k_block = k[ks:ke] + v_block = v[ks:ke] + if num_heads_q != num_heads_kv: + assert num_heads_q % num_heads_kv == 0 + repeat = num_heads_q // num_heads_kv + k_block = k_block.repeat_interleave(repeat, dim=1) + v_block = v_block.repeat_interleave(repeat, dim=1) + block_out = F.scaled_dot_product_attention( + q_block.transpose(0, 1).unsqueeze(0).contiguous(), + k_block.transpose(0, 1).unsqueeze(0).contiguous(), + v_block.transpose(0, 1).unsqueeze(0).contiguous(), + ) + out[qs:qe] += block_out.squeeze(0).transpose(0, 1).contiguous() + lse = torch.empty((q.shape[0], q.shape[1]), dtype=torch.float32, device=q.device) + return out, lse + + def magi_register_custom_op(name=None, mutates_args=(), infer_output_meta_fn=None, is_subgraph_boundary=False, **kwargs): + def decorator(fn): + if not name: + return fn + namespace, op_name = name.split("::", 1) + if namespace not in _libs: + _libs[namespace] = torch.library.Library(namespace, "FRAGMENT") + if (namespace, op_name) in _defined: + # Already registered in a previous test run — reuse. + return fn + # Route known ops through SDPA; leave unknown ones as direct fn. + returns = "(Tensor, Tensor)" if op_name == "flex_flash_attn_func" else "Tensor" + schema_name = f"{op_name}({_infer_schema(fn)}) -> {returns}" + try: + _libs[namespace].define(schema_name) + except Exception: + pass + if op_name == "flash_attn_func": + torch.library.impl( + _libs[namespace], op_name, "CUDA" + )(_sdpa_flash_attn_func) + torch.library.impl( + _libs[namespace], op_name, "CPU" + )(_sdpa_flash_attn_func) + elif op_name == "flex_flash_attn_func": + torch.library.impl( + _libs[namespace], op_name, "CUDA" + )(_sdpa_segments) + else: + # For ops we don't care about (compile-only wrappers), the + # Python fn path inside the module body is used directly — + # we just need `torch.ops..` to exist so module- + # load-time attribute lookups succeed. + torch.library.impl( + _libs[namespace], op_name, "CUDA" + )(fn) + _defined.add((namespace, op_name)) + return fn + return decorator + + pkg.magi_compile = magi_compile + sys.modules["magi_compiler"] = pkg + + api = types.ModuleType("magi_compiler.api") + api.magi_register_custom_op = magi_register_custom_op + sys.modules["magi_compiler.api"] = api + pkg.api = api + + config_mod = types.ModuleType("magi_compiler.config") + + class CompileConfig: + class offload_config: # pragma: no cover - pass-through + gpu_resident_weight_ratio = 1.0 + config_mod.CompileConfig = CompileConfig + sys.modules["magi_compiler.config"] = config_mod + pkg.config = config_mod + + +def _infer_schema(fn) -> str: + """Return a minimal torch.library schema string for the given fn. + + For our stub we just need *something* parseable; all real ops we + care about take `(q, k, v)` or `(q, k, v, q_ranges, k_ranges)` or + variants. Use generic `Tensor a, Tensor b, ...` arg names. + """ + import inspect + sig = inspect.signature(fn) + parts = [] + for i, name in enumerate(sig.parameters): + arg_name = name if name.isidentifier() else f"a{i}" + parts.append(f"Tensor {arg_name}") + return ", ".join(parts) + + +# --------------------------------------------------------------------------- +# distributed / CP stubs — single-GPU, cp_world_size == 1. +# --------------------------------------------------------------------------- + + +def _install_distributed_stubs() -> None: + """Monkey-patch upstream distributed + parallelism modules for cp=1.""" + # inference.infra.distributed.* + # The real module requires NCCL / parallel_state to be initialized + # from torchrun; here we short-circuit the handful of getters the DiT + # actually calls. + import inference.infra.distributed as dist_mod + + dist_mod.get_cp_world_size = lambda: 1 + dist_mod.get_cp_group = lambda: None + dist_mod.get_cp_rank = lambda: 0 + dist_mod.get_tp_rank = lambda: 0 + dist_mod.get_pp_rank = lambda: 0 + + # inference.infra.parallelism.* + # At cp_world_size=1, scatter/gather are trivially no-ops. + import inference.infra.parallelism.gather_scatter_primitive as gs + + def _scatter_noop(x, cp_split_sizes, group=None): + return x + + def _gather_noop(x, cp_split_sizes, group=None): + return x + + gs.scatter_to_context_parallel_region = _scatter_noop + gs.gather_from_context_parallel_region = _gather_noop + + # Re-import ulysses_scheduler with patched scatter/gather in place. + import inference.infra.parallelism.ulysses_scheduler as us + us.scatter_to_context_parallel_region = _scatter_noop + us.gather_from_context_parallel_region = _gather_noop + us.get_cp_world_size = lambda: 1 + us.get_cp_group = lambda: None + + # all-to-all primitives used by flash_attn_with_cp. At cp=1 they are + # entered only as a no-op path (the `if cp_world_size > 1` branch is + # skipped), so no stubs needed there. + + +def install_stubs() -> None: + """Install all stubs. Idempotent.""" + _install_magi_compiler_stub() + repo_root = Path(__file__).resolve().parents[3] + upstream = repo_root / "daVinci-MagiHuman" + path_s = str(upstream) + if path_s not in sys.path: + sys.path.insert(0, path_s) + # Reload inference.* after sys.path mutation so it picks up the real + # upstream package (not a stale one). + for name in list(sys.modules): + if name == "inference" or name.startswith("inference."): + del sys.modules[name] + import inference # noqa: F401 + _install_distributed_stubs() + + +# --------------------------------------------------------------------------- +# Upstream DiTModel loader — instantiate + load base shards. +# --------------------------------------------------------------------------- + + +def _base_arch_dict() -> dict: + """Return the upstream `ModelConfig`-equivalent dict for the base variant. + Matches `inference/common/config.py::ModelConfig` defaults for base. + """ + import torch + return dict( + num_layers=40, + hidden_size=5120, + head_dim=128, + num_query_groups=8, + video_in_channels=48 * 4, + audio_in_channels=64, + text_in_channels=3584, + checkpoint_qk_layernorm_rope=False, + params_dtype=torch.float32, + tread_config=dict( + selection_rate=0.5, start_layer_idx=2, end_layer_idx=25, + ), + mm_layers=[0, 1, 2, 3, 36, 37, 38, 39], + local_attn_layers=[], + enable_attn_gating=True, + activation_type="swiglu7", + gelu7_layers=[0, 1, 2, 3], + # derived + num_heads_q=40, + num_heads_kv=8, + post_norm_layers=[], + ) + + +def load_upstream_dit(base_shard_dir, device=None, dtype=None, local_attn_layers=None): + """Instantiate upstream `DiTModel` and load the base shards into it. + + Args: + base_shard_dir: path to `base/` (contains `model-0000*-of-00007.safetensors` + and `model.safetensors.index.json`). + device: torch device (default cuda if available). + dtype: dtype cast (default: leave checkpoint dtypes as-is). + + Returns: + An upstream `DiTModel` in `.eval()` mode with weights loaded. + """ + import glob + import json + import types as _types + + import torch + from safetensors.torch import load_file + + from inference.common.config import ModelConfig # upstream pydantic class + from inference.model.dit.dit_module import DiTModel + + arch_dict = _base_arch_dict() + if local_attn_layers is not None: + arch_dict["local_attn_layers"] = list(local_attn_layers) + # ModelConfig is a pydantic BaseModel — build via kwargs. + model_config = ModelConfig(**arch_dict) + + model = DiTModel(model_config=model_config) + + # Load all base shards into a single state dict. + base_shard_dir = Path(base_shard_dir) + shard_paths = sorted(base_shard_dir.glob("*.safetensors")) + state = {} + for p in shard_paths: + state.update(load_file(str(p))) + + missing, unexpected = model.load_state_dict(state, strict=False) + if missing: + raise RuntimeError(f"Upstream DiT missing {len(missing)} keys: {missing[:5]}") + if unexpected: + raise RuntimeError(f"Upstream DiT unexpected {len(unexpected)} keys: {unexpected[:5]}") + + device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = model.to(device=device) + if dtype is not None: + model = model.to(dtype=dtype) + model.eval() + return model diff --git a/tests/local_tests/magi_human/test_magi_human_distill_parity.py b/tests/local_tests/magi_human/test_magi_human_distill_parity.py new file mode 100644 index 000000000..14b96a442 --- /dev/null +++ b/tests/local_tests/magi_human/test_magi_human_distill_parity.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +"""DiT parity test for the daVinci-MagiHuman DMD-2 distilled checkpoint. + +The distill variant has the SAME architecture as the base model (same 40 +layers, same hidden_size, same mm_layers / gelu7_layers / local_attn_layers, +same head_dim and num_query_groups; see +`daVinci-MagiHuman/inference/common/config.py:ModelConfig`). Only the +weights differ: distill is trained for 8-step DMD-2 inference without CFG. + +This test mirrors `test_magi_human_parity.py::test_magi_human_dit_parity` +exactly, just pointing at the `distill/` subfolder of GAIR/daVinci-MagiHuman +and the matching `converted_weights/magi_human_distill/`. + +Skips cleanly when: + * `daVinci-MagiHuman/` clone is absent + * GAIR/daVinci-MagiHuman distill shards are not locally available + * Converted distill weights have not been produced yet + * CUDA is unavailable +""" +from __future__ import annotations + +import gc +import glob +import os +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.forward_context import set_forward_context + + +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + + +def _find_distill_shard_dir() -> Path | None: + """Return the local path to GAIR/daVinci-MagiHuman/distill/ shards or None.""" + override = os.getenv("MAGI_HUMAN_DISTILL_SHARD_DIR") + if override: + p = Path(override) + return p if p.is_dir() else None + try: + from huggingface_hub import snapshot_download + snap = snapshot_download( + repo_id="GAIR/daVinci-MagiHuman", + allow_patterns=[ + "distill/*.safetensors", + "distill/model.safetensors.index.json", + ], + ) + candidate = Path(snap) / "distill" + if candidate.is_dir() and any(candidate.glob("*.safetensors")): + return candidate + return None + except Exception: + return None + + +def _cleanup_gpu() -> None: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MagiHuman distill DiT parity requires CUDA.", +) +def test_magi_human_distill_dit_parity(): + repo_root = Path(__file__).resolve().parents[3] + upstream_src = repo_root / "daVinci-MagiHuman" + if not upstream_src.exists(): + pytest.skip( + "Upstream daVinci-MagiHuman/ clone missing. Run " + "`git clone --depth 1 https://github.com/GAIR-NLP/daVinci-MagiHuman.git`" + ) + + distill_shard_dir = _find_distill_shard_dir() + if distill_shard_dir is None or not distill_shard_dir.is_dir(): + pytest.skip( + "GAIR/daVinci-MagiHuman distill/ shards not available locally. " + "Set MAGI_HUMAN_DISTILL_SHARD_DIR or run the conversion once to " + "populate the HF cache." + ) + + converted_dir = Path(os.getenv( + "MAGI_HUMAN_DISTILL_DIFFUSERS_PATH", + repo_root / "converted_weights" / "magi_human_distill", + )) + transformer_dir = converted_dir / "transformer" + if not transformer_dir.is_dir(): + pytest.skip( + f"Converted distill transformer dir missing at {transformer_dir}. Run " + f"scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py " + f"--subfolder distill --cast-bf16 first." + ) + + from tests.local_tests.helpers.magi_human_upstream import ( + install_stubs, + load_upstream_dit, + ) + install_stubs() + + device = torch.device("cuda:0") + torch.manual_seed(0) + + z_dim = 48 + pT, pH, pW = 1, 2, 2 + lat_T, lat_H, lat_W = 2, 6, 6 + video_latent = torch.randn( + (1, z_dim, lat_T, lat_H, lat_W), + dtype=torch.float32, device=device, + ) + num_video_tokens = (lat_T // pT) * (lat_H // pH) * (lat_W // pW) + num_audio_tokens = 4 + num_text_tokens = 8 + audio_latent = torch.randn( + (1, num_audio_tokens, 64), + dtype=torch.float32, device=device, + ) + text_feat = torch.randn( + (1, num_text_tokens, 3584), + dtype=torch.float32, device=device, + ) + + from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import ( + build_packed_inputs, + ) + from fastvideo.models.dits.magi_human import Modality # noqa: F401 + + x, coords, mm = build_packed_inputs( + video_latent=video_latent, + audio_latent=audio_latent, + audio_feat_len=num_audio_tokens, + txt_feat=text_feat, + txt_feat_len=num_text_tokens, + patch_size=(pT, pH, pW), + coords_style="v2", + ) + assert x.shape[0] == num_video_tokens + num_audio_tokens + num_text_tokens + + total_tokens = x.shape[0] + + # Distill arch is identical to base; load_upstream_dit's _base_arch_dict + # describes both since they share num_layers / hidden_size / mm_layers + # / etc. Just point at the distill shards. + print("Loading upstream distill DiTModel from distill shards...") + upstream_model = load_upstream_dit( + distill_shard_dir, + device=device, + dtype=None, + ) + + from inference.common import VarlenHandler + cu = torch.tensor([0, total_tokens], dtype=torch.int32, device=device) + varlen = VarlenHandler( + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=total_tokens, + max_seqlen_k=total_tokens, + ) + + print("Running upstream distill forward...") + with torch.inference_mode(): + ref_out = upstream_model( + x=x.clone(), + coords_mapping=coords.clone(), + modality_mapping=mm.clone(), + varlen_handler=varlen, + local_attn_handler=None, + ).detach().float().cpu() + + del upstream_model + _cleanup_gpu() + + from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig + from fastvideo.models.dits.magi_human import MagiHumanDiT + from safetensors.torch import load_file + print("Loading FastVideo MagiHumanDiT from converted distill transformer/...") + fv_cfg = MagiHumanVideoConfig() + fv_model = MagiHumanDiT(fv_cfg) + + fv_state = {} + for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))): + fv_state.update(load_file(shard)) + missing, unexpected = fv_model.load_state_dict(fv_state, strict=False) + assert not missing, f"FastVideo distill DiT missing {len(missing)} keys: {missing[:5]}" + assert not unexpected, f"FastVideo distill DiT unexpected {len(unexpected)} keys: {unexpected[:5]}" + + fv_model = fv_model.to(device=device) + fv_model.eval() + + print("Running FastVideo distill forward...") + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + fv_out = fv_model(x.clone(), coords.clone(), mm.clone()).detach().float().cpu() + + print( + f"ref sum={ref_out.sum().item():.4f} " + f"abs_mean={ref_out.abs().mean().item():.4f} " + f"shape={tuple(ref_out.shape)}" + ) + print( + f"fv sum={fv_out.sum().item():.4f} " + f"abs_mean={fv_out.abs().mean().item():.4f} " + f"shape={tuple(fv_out.shape)}" + ) + diff = (ref_out - fv_out).abs() + print( + f"diff max={diff.max().item():.6f} " + f"mean={diff.mean().item():.6f} " + f"median={diff.median().item():.6f}" + ) + + ref_video = ref_out[:num_video_tokens] + fv_video = fv_out[:num_video_tokens] + ref_audio = ref_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64] + fv_audio = fv_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64] + ref_text = ref_out[num_video_tokens + num_audio_tokens:] + fv_text = fv_out[num_video_tokens + num_audio_tokens:] + video_diff = (ref_video - fv_video).abs() + audio_diff = (ref_audio - fv_audio).abs() + text_diff = (ref_text - fv_text).abs() + print( + f"video ref_abs={ref_video.abs().mean():.4f} " + f"diff_max={video_diff.max():.4f} diff_mean={video_diff.mean():.4f}" + ) + print( + f"audio ref_abs={ref_audio.abs().mean():.4f} " + f"diff_max={audio_diff.max():.4f} diff_mean={audio_diff.mean():.4f}" + ) + print( + f"text ref_abs={ref_text.abs().mean():.4f} " + f"diff_max={text_diff.max():.4f} diff_mean={text_diff.mean():.4f}" + ) + + assert ref_out.shape == fv_out.shape + assert_close(fv_text, ref_text, atol=1e-6, rtol=1e-6) + # Same tolerance as base DiT parity (post-Wave 14b dtype-boundary fixes, + # both DiTs are bit-exact via shared upstream BaseLinear bf16 path). + assert_close(fv_out, ref_out, atol=0.03, rtol=0.01) + ref_abs = ref_out.abs().mean().item() + fv_abs = fv_out.abs().mean().item() + rel = abs(ref_abs - fv_abs) / max(ref_abs, 1e-6) + assert rel < 0.05, f"abs_mean drift {rel:.3%} > 5%" diff --git a/tests/local_tests/magi_human/test_magi_human_parity.py b/tests/local_tests/magi_human/test_magi_human_parity.py new file mode 100644 index 000000000..a97f5cf2d --- /dev/null +++ b/tests/local_tests/magi_human/test_magi_human_parity.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical parity test: FastVideo MagiHumanDiT vs upstream DiTModel. + +Loads both models from the **same converted base checkpoint** and runs them +on identical small inputs. Asserts closeness on the joint video+audio +output tensor. + +What this catches (that the preflight test does NOT): + - Silent weight-name mismatches that `strict=False` loading would hide. + - Wrong modality-expert chunking inside `PackedExpertLinear`. + - RoPE sin/cos ordering flipped. + - Per-head gating dtype / split order. + - swiglu7 / gelu7 off-by-one on the `+1` linear bias. + +Skips cleanly when: + - `daVinci-MagiHuman/` clone is absent (no upstream source). + - GAIR/daVinci-MagiHuman base shards are not available locally. + - CUDA is unavailable. + +Tolerance: `atol=5e-3, rtol=5e-3` on bf16 forward paths. The FastVideo +attention path uses `F.scaled_dot_product_attention` while upstream uses +`flash_attn_func`; both accumulate in bf16 but via different kernels, so +small drift is expected and bounded. +""" +from __future__ import annotations + +import gc +import os +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.forward_context import set_forward_context + + +# Force TORCH_SDPA for FastVideo so the attention kernel is deterministic. +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + + +def _find_base_shard_dir() -> Path | None: + """Return the local path to GAIR/daVinci-MagiHuman/base/ with shards present, or None.""" + override = os.getenv("MAGI_HUMAN_BASE_SHARD_DIR") + if override: + p = Path(override) + return p if p.is_dir() else None + try: + from huggingface_hub import snapshot_download + snap = snapshot_download( + repo_id="GAIR/daVinci-MagiHuman", + allow_patterns=[ + "base/*.safetensors", + "base/model.safetensors.index.json", + ], + ) + candidate = Path(snap) / "base" + if candidate.is_dir() and any(candidate.glob("*.safetensors")): + return candidate + return None + except Exception: + return None + + +def _cleanup_gpu() -> None: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="MagiHuman DiT parity requires CUDA.", +) +def test_magi_human_dit_parity(): + repo_root = Path(__file__).resolve().parents[3] + upstream_src = repo_root / "daVinci-MagiHuman" + if not upstream_src.exists(): + pytest.skip( + "Upstream daVinci-MagiHuman/ clone missing. Run " + "`git clone --depth 1 https://github.com/GAIR-NLP/daVinci-MagiHuman.git`" + ) + + base_shard_dir = _find_base_shard_dir() + if base_shard_dir is None or not base_shard_dir.is_dir(): + pytest.skip( + "GAIR/daVinci-MagiHuman base/ shards not available locally. " + "Set MAGI_HUMAN_BASE_SHARD_DIR or run the conversion once to " + "populate the HF cache." + ) + + converted_dir = Path(os.getenv( + "MAGI_HUMAN_DIFFUSERS_PATH", + repo_root / "converted_weights" / "magi_human_base", + )) + transformer_dir = converted_dir / "transformer" + if not transformer_dir.is_dir(): + pytest.skip( + f"Converted transformer dir missing at {transformer_dir}. Run " + f"scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py first." + ) + + # Add upstream to sys.path and install compiler/distributed stubs. + from tests.local_tests.helpers.magi_human_upstream import ( + install_stubs, + load_upstream_dit, + ) + install_stubs() + + # --- Shared inputs (deliberately small) --- + device = torch.device("cuda:0") + torch.manual_seed(0) + + # Mirror MagiDataProxy.process_input for a tiny frame: + # video_latent: [1, z_dim, T, H, W], T=2, H=6, W=6, z_dim=48 + # -> video tokens: (T/pT)*(H/pH)*(W/pW) with patch=(1,2,2) = 2*3*3 = 18 + # audio tokens: 4 + # text tokens: 8 + # max channel width = 192 (video) + z_dim = 48 + pT, pH, pW = 1, 2, 2 + lat_T, lat_H, lat_W = 2, 6, 6 + video_latent = torch.randn( + (1, z_dim, lat_T, lat_H, lat_W), + dtype=torch.float32, device=device, + ) + num_video_tokens = (lat_T // pT) * (lat_H // pH) * (lat_W // pW) # 18 + num_audio_tokens = 4 + num_text_tokens = 8 + audio_latent = torch.randn( + (1, num_audio_tokens, 64), + dtype=torch.float32, device=device, + ) + text_feat = torch.randn( + (1, num_text_tokens, 3584), + dtype=torch.float32, device=device, + ) + + # --- Build the packed inputs the DiT consumes --- + from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import ( + build_packed_inputs, + ) + from fastvideo.models.dits.magi_human import Modality # noqa: F401 + + x, coords, mm = build_packed_inputs( + video_latent=video_latent, + audio_latent=audio_latent, + audio_feat_len=num_audio_tokens, + txt_feat=text_feat, + txt_feat_len=num_text_tokens, + patch_size=(pT, pH, pW), + coords_style="v2", + ) + assert x.shape[0] == num_video_tokens + num_audio_tokens + num_text_tokens + + total_tokens = x.shape[0] + + # --- Load upstream DiT first (so we know weights round-trip cleanly). + # Upstream reads raw base/ shards; we keep it in bf16 for speed + # and because that matches the FastVideo side after FSDP load. + print("Loading upstream DiTModel from base shards...") + upstream_model = load_upstream_dit( + base_shard_dir, + device=device, + dtype=None, # keep checkpoint dtypes (fp32 for norms, bf16 for matmuls) + ) + + # --- VarlenHandler for upstream (batch=1, total_tokens). + from inference.common import VarlenHandler + cu = torch.tensor([0, total_tokens], dtype=torch.int32, device=device) + varlen = VarlenHandler( + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=total_tokens, + max_seqlen_k=total_tokens, + ) + + # --- Forward upstream and capture output. --- + print("Running upstream forward...") + with torch.inference_mode(): + ref_out = upstream_model( + x=x.clone(), + coords_mapping=coords.clone(), + modality_mapping=mm.clone(), + varlen_handler=varlen, + local_attn_handler=None, # local_attn_layers=[] for base + ).detach().float().cpu() + + # Free upstream model before loading FastVideo (saves ~30 GB on GPU). + del upstream_model + _cleanup_gpu() + + # --- Load FastVideo MagiHumanDiT from the converted transformer/ --- + from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig + from fastvideo.models.dits.magi_human import MagiHumanDiT + from safetensors.torch import load_file + import glob + print("Loading FastVideo MagiHumanDiT from converted transformer/...") + fv_cfg = MagiHumanVideoConfig() + fv_model = MagiHumanDiT(fv_cfg) + + fv_state = {} + for shard in sorted(glob.glob(str(transformer_dir / "*.safetensors"))): + fv_state.update(load_file(shard)) + missing, unexpected = fv_model.load_state_dict(fv_state, strict=False) + assert not missing, f"FastVideo DiT missing {len(missing)} keys: {missing[:5]}" + assert not unexpected, f"FastVideo DiT unexpected {len(unexpected)} keys: {unexpected[:5]}" + + fv_model = fv_model.to(device=device) + fv_model.eval() + + print("Running FastVideo forward...") + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + fv_out = fv_model(x.clone(), coords.clone(), mm.clone()).detach().float().cpu() + + # Global stats + print( + f"ref sum={ref_out.sum().item():.4f} " + f"abs_mean={ref_out.abs().mean().item():.4f} " + f"shape={tuple(ref_out.shape)}" + ) + print( + f"fv sum={fv_out.sum().item():.4f} " + f"abs_mean={fv_out.abs().mean().item():.4f} " + f"shape={tuple(fv_out.shape)}" + ) + diff = (ref_out - fv_out).abs() + print( + f"diff max={diff.max().item():.6f} " + f"mean={diff.mean().item():.6f} " + f"median={diff.median().item():.6f}" + ) + + # Per-modality diagnostic (video, audio, text). Text rows are zero- + # padded on both sides; video and audio should carry comparable + # abs_mean. + ref_video = ref_out[:num_video_tokens] + fv_video = fv_out[:num_video_tokens] + ref_audio = ref_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64] + fv_audio = fv_out[num_video_tokens:num_video_tokens + num_audio_tokens, :64] + ref_text = ref_out[num_video_tokens + num_audio_tokens:] + fv_text = fv_out[num_video_tokens + num_audio_tokens:] + video_diff = (ref_video - fv_video).abs() + audio_diff = (ref_audio - fv_audio).abs() + text_diff = (ref_text - fv_text).abs() + print( + f"video ref_abs={ref_video.abs().mean():.4f} " + f"diff_max={video_diff.max():.4f} diff_mean={video_diff.mean():.4f}" + ) + print( + f"audio ref_abs={ref_audio.abs().mean():.4f} " + f"diff_max={audio_diff.max():.4f} diff_mean={audio_diff.mean():.4f}" + ) + print( + f"text ref_abs={ref_text.abs().mean():.4f} " + f"diff_max={text_diff.max():.4f} diff_mean={text_diff.mean():.4f}" + ) + + # --- Assertions --- + assert ref_out.shape == fv_out.shape, ( + f"shape mismatch: ref={ref_out.shape} fv={fv_out.shape}" + ) + + # Text rows are zero-padded on both sides — must match exactly. + assert_close(fv_text, ref_text, atol=1e-6, rtol=1e-6) + + # Video + audio: bf16 single-forward DiT noise floor is ~1e-3 to + # 5e-3 per element. atol=0.03 catches gross structural bugs + # (permutation flips, sign inversions, wrong modality dispatch, + # missing sub-layers) while leaving 6-10x margin over actual bf16 + # noise. Observed diff_max=0.057 will FAIL — that is the bug + # surfacing and is the intended spec for downstream root-cause + # investigation. + assert_close(fv_out, ref_out, atol=0.03, rtol=0.01) + + # Sanity: mean magnitudes should match within 5%. A gross bug + # (e.g. dropping a modality branch) would show up here. + ref_abs = ref_out.abs().mean().item() + fv_abs = fv_out.abs().mean().item() + rel = abs(ref_abs - fv_abs) / max(ref_abs, 1e-6) + assert rel < 0.05, f"abs_mean drift {rel:.3%} > 5% — possible structural bug"