From 6771f4ec2ee5be7f1b5d80b01caefdf64b9a2ff0 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sat, 6 Jun 2026 20:52:36 -0700 Subject: [PATCH] [feat] cosmos3: native DiT + framework parity (Cosmos3VFMTransformer) --- fastvideo/models/dits/cosmos3.py | 983 +++++++++++++----- .../cosmos3/test_cosmos3_dit_parity.py | 256 +++++ .../cosmos3/test_cosmos3_state_dict_keys.py | 235 +++-- 3 files changed, 1099 insertions(+), 375 deletions(-) create mode 100644 tests/local_tests/cosmos3/test_cosmos3_dit_parity.py diff --git a/fastvideo/models/dits/cosmos3.py b/fastvideo/models/dits/cosmos3.py index 9c6c4fae9..c4b5c9ac3 100644 --- a/fastvideo/models/dits/cosmos3.py +++ b/fastvideo/models/dits/cosmos3.py @@ -1,36 +1,38 @@ # SPDX-License-Identifier: Apache-2.0 -"""Cosmos3 VFM Transformer (Phase 2b.1 skeleton). +"""FastVideo-native Cosmos3 omni DiT (``Cosmos3VFMTransformer``). -This module lands the parameter-named module tree of -``Cosmos3VFMTransformer`` together with two standalone math utilities used -by both the FastVideo port and the Phase 2a parity tests: +Numerical-parity port of the official ``cosmos_framework`` ``Cosmos3VFMNetwork`` +(the ``Qwen3-VL-text`` MoT backbone + the VFM vision / action / sound heads). +The module tree mirrors the published *diffusers* checkpoint layout so a +converter can strict-load with a near-identity ``param_names_mapping``: -* ``compute_mrope_position_ids_text`` / ``compute_mrope_position_ids_vision`` — - the unified-3D mRoPE position ID generators ported verbatim from - ``vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py`` lines - 113-177 (HEAD ``8536f5b1421f``). +* top level: ``embed_tokens`` / ``norm`` / ``norm_moe_gen`` / ``lm_head`` / + ``proj_in`` / ``proj_out`` / ``time_embedder.linear_{1,2}`` plus the dormant + ``action_*`` / ``audio_*`` heads (constructed for strict-load parity); +* per layer ``layers.{i}``: dual-pathway ``self_attn`` with understanding + (``to_{q,k,v}`` / ``to_out``) and generation (``add_{q,k,v}_proj`` / + ``to_add_out``) projections + per-head QK-norms (``norm_{q,k}`` for und, + ``norm_added_{q,k}`` for gen), ``mlp`` (und) and ``mlp_moe_gen`` (gen) SwiGLU + blocks, and four RMSNorms. -* ``Cosmos3VFMTransformer.patchify`` / - ``Cosmos3VFMTransformer.unpatchify`` — the - ``[B,C,T,H,W] <-> [B, T*hp*wp, p*p*C]`` patch tokenizer, ported from - the same reference lines 1009-1036. +The forward replicates the framework contract for the video path: +``patchify -> proj_in -> (+ 3d-rope additive latent pos emb) -> scatter-add +timestep embeds onto noisy patches -> dual-pathway decoder (causal text + +full vision two-way attention, GQA, per-head QK-norm before RoPE, unified +3D-MRoPE) -> und/gen final norms -> proj_out on noisy vision patches -> +unpatchify``. -Every per-layer ``forward()`` raises ``NotImplementedError("Phase 2b.2")`` -because the layer math (RoPE application, QK-norm, GQA SDPA, UND/GEN -attention plumbing) is intentionally deferred to Phase 2b.2. Only the -module composition + tensor shape contract is exercised here, enough to -make 4 Tier-A parity tests pass: +RoPE is applied with the framework's exact ``rotate_half`` math (contiguous +split-half: ``q * cos + rotate_half(q) * sin``) rather than the interleaved +``apply_rotary_emb`` helper, and attention uses plain SDPA so the whole module +runs on CPU / float32 for parity testing. No diffusers / transformers +model-class imports happen at runtime — the DiT is fully native. -* ``test_cosmos3_mrope_parity.py::test_compute_mrope_position_ids_text_and_vision`` -* ``test_cosmos3_patchify_unpatchify_parity.py::test_patchify_unpatchify_roundtrip`` -* ``test_cosmos3_patchify_unpatchify_parity.py::test_patchify_default_patch_size`` -* ``test_cosmos3_state_dict_keys.py::test_fastvideo_cosmos3_dit_module_tree_param_names`` - -The constructor signature mirrors the upstream -``Cosmos3VFMTransformer(od_config, *, temporal_compression_factor=None)`` -contract so the Phase 2a tests (which were authored against the upstream -shape) can drive it directly. Phase 2b.2+ will adapt this to FastVideo's -``BaseDiT(config, hf_config)`` loader path via the ``TransformerLoader``. +This module also keeps two pure, standalone math utilities used by the Tier-A +scaffold parity tests: the unified-3D mRoPE position-ID generators +(``compute_mrope_position_ids_text`` / ``compute_mrope_position_ids_vision``) +and the batched ``patchify`` / ``unpatchify`` ``[B,C,T,H,W] <-> [B,N,p*p*C]`` +roundtrip helpers. """ from __future__ import annotations @@ -41,25 +43,17 @@ import torch import torch.nn as nn import torch.nn.functional as F -from fastvideo.configs.models.dits.cosmos3 import Cosmos3ArchConfig, Cosmos3VideoConfig +from fastvideo.configs.models.dits.cosmos3 import Cosmos3VideoConfig +from fastvideo.layers.layernorm import RMSNorm +from fastvideo.layers.linear import ReplicatedLinear from fastvideo.models.dits.base import BaseDiT - EntryClass = ["Cosmos3VFMTransformer"] -def _tf_config_get(config: Any, key: str, default: Any) -> Any: - """Read a value from a dict, dataclass, or simple namespace. - - Mirrors ``transformer_cosmos3._tf_config_get`` at reference line 77. - """ - if config is None: - return default - if hasattr(config, "get"): - return config.get(key, default) - return getattr(config, key, default) - - +# =========================================================================== +# Standalone mRoPE position-ID generators (unified 3D mRoPE) +# =========================================================================== def compute_mrope_position_ids_text( num_tokens: int, temporal_offset: int, @@ -67,9 +61,7 @@ def compute_mrope_position_ids_text( """Generate 3D mRoPE position IDs for text tokens. Text tokens broadcast a single monotonically-increasing position-ID - sequence across all three (t, h, w) axes. Verbatim port of - ``transformer_cosmos3.compute_mrope_position_ids_text`` (reference - lines 113-124). + sequence across all three (t, h, w) axes. """ ids = torch.arange(num_tokens, dtype=torch.long) + temporal_offset mrope_ids = ids.unsqueeze(0).expand(3, -1).contiguous() @@ -90,13 +82,10 @@ def compute_mrope_position_ids_vision( ) -> tuple[torch.Tensor, int]: """Generate 3D mRoPE position IDs for vision tokens. - Builds a ``(t, h, w)`` position grid (Qwen3-VL style, spatial indices - reset per temporal segment) flattened in t-major order. Optionally - modulates the temporal axis by ``base_fps / tcf * (1 / (fps / tcf))`` - so two clips at different FPS retain wall-clock-aligned temporal - positions. Verbatim port of - ``transformer_cosmos3.compute_mrope_position_ids_vision`` (reference - lines 127-177). + Builds a ``(t, h, w)`` position grid (Qwen3-VL style, spatial indices reset + per temporal segment) flattened in t-major order. Optionally modulates the + temporal axis by ``base_fps / tcf * (1 / (fps / tcf))`` so two clips at + different FPS retain wall-clock-aligned temporal positions. """ fps_modulation = enable_fps_modulation and fps is not None @@ -104,15 +93,14 @@ def compute_mrope_position_ids_vision( assert fps is not None tps = fps / temporal_compression_factor effective_base_tcf = (base_temporal_compression_factor - if base_temporal_compression_factor is not None - else temporal_compression_factor) + if base_temporal_compression_factor is not None else temporal_compression_factor) base_tps = base_fps / effective_base_tcf frame_indices = torch.arange(grid_t, dtype=torch.float32) - t_index = (((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset) - .view(-1, 1).expand(-1, grid_h * grid_w).flatten()) + t_index = (((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset).view(-1, 1).expand( + -1, grid_h * grid_w).flatten()) else: - t_index = (torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() - + int(temporal_offset) + start_frame_offset) + t_index = (torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() + + int(temporal_offset) + start_frame_offset) h_index = (torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten()) w_index = (torch.arange(grid_w, dtype=torch.long).view(1, 1, -1).expand(grid_t, grid_h, -1).flatten()) @@ -126,83 +114,266 @@ def compute_mrope_position_ids_vision( return mrope_ids, next_offset -class _Qwen3VLTextRMSNorm(nn.Module): - """Qwen3-VL / T5-style RMSNorm with an ``eps`` and a learnable ``weight``. +# =========================================================================== +# RoPE helpers (match cosmos_framework qwen3_vl rotate_half + apply_rotary_pos_emb) +# =========================================================================== +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + """Contiguous split-half rotation, identical to qwen3_vl.rotate_half.""" + x1 = x[..., :x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2:] + return torch.cat((-x2, x1), dim=-1) - Mirrors ``Qwen3VLTextRMSNorm`` at reference line 89-107 so checkpoint - keys match the canonical ``*.norm*.weight`` / ``*.{q,k}_norm.weight`` - parameter naming. + +def _apply_rotary_pos_emb( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + unsqueeze_dim: int = 1, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply RoPE exactly as cosmos_framework qwen3_vl.apply_rotary_pos_emb. + + ``q`` / ``k`` are ``[N, heads, head_dim]`` and ``cos`` / ``sin`` are + ``[N, head_dim]`` (per-token); ``unsqueeze_dim=1`` broadcasts cos/sin over + the head axis. """ - - def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.bfloat16) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype)) - self.variance_epsilon = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - raise NotImplementedError("Phase 2b.2") + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (_rotate_half(q) * sin) + k_embed = (k * cos) + (_rotate_half(k) * sin) + return q_embed, k_embed -class _TimestepEmbedder(nn.Module): - """Sinusoidal-to-vector timestep embedder used by the GEN pathway. +class Cosmos3TextRotaryEmbedding(nn.Module): + """Unified 3D-MRoPE rotary embedding (qwen3_vl Cosmos3 flavor). - Skeleton of ``TimestepEmbedder`` at reference line 257-283. The - ``linear_1``/``linear_2`` parameter names match the upstream - checkpoint so the Phase 5 remap can copy weights without renaming. + Reproduces ``Qwen3VLTextRotaryEmbedding``: ``inv_freq`` from the default + RoPE init (``1 / theta**(arange(0, head_dim, 2) / head_dim)``), interleaved + MRoPE mixing of the T/H/W frequency bands per ``mrope_section``, and a final + ``cat([freqs, freqs])`` so cos/sin are ``[N, head_dim]``. """ def __init__( self, - hidden_size: int, - frequency_embedding_size: int = 256, - max_period: int = 10000, - target_dtype: torch.dtype = torch.bfloat16, + head_dim: int, + rope_theta: float, + mrope_section: list[int], ) -> None: super().__init__() - self.linear_1 = nn.Linear(frequency_embedding_size, hidden_size, bias=True) - self.act = nn.SiLU() - self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True) + self.head_dim = head_dim + self.mrope_section = list(mrope_section) + inv_freq = 1.0 / (rope_theta**(torch.arange(0, head_dim, 2, dtype=torch.int64).float() / head_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.attention_scaling = 1.0 # default RoPE has unit attention scaling + + def _apply_interleaved_mrope(self, freqs: torch.Tensor) -> torch.Tensor: + """freqs: [3, N, head_dim//2] -> [N, head_dim//2] (interleaved T/H/W).""" + freqs_t = freqs[0].clone() + for dim, offset in enumerate((1, 2), start=1): # H, W + length = self.mrope_section[dim] * 3 + idx = slice(offset, length, 3) + freqs_t[..., idx] = freqs[dim, ..., idx] + return freqs_t + + @torch.no_grad() + def forward(self, position_ids: torch.Tensor, device: torch.device, + dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + """position_ids: [N] (1D) or [3, N] (mrope) -> (cos, sin) each [N, head_dim].""" + if position_ids.ndim == 1: + position_ids = position_ids[None, :].expand(3, -1) # [3, N] + position_ids = position_ids.to(device) + inv_freq = self.inv_freq.to(device).float() # [head_dim//2] + inv_freq_expanded = inv_freq[None, :, None].expand(3, -1, 1) # [3, head_dim//2, 1] + position_ids_expanded = position_ids[:, None, :].float() # [3, 1, N] + freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) # [3, N, head_dim//2] + freqs = self._apply_interleaved_mrope(freqs) # [N, head_dim//2] + emb = torch.cat((freqs, freqs), dim=-1) # [N, head_dim] + cos = (emb.cos() * self.attention_scaling).to(dtype) + sin = (emb.sin() * self.attention_scaling).to(dtype) + return cos, sin + + +# =========================================================================== +# 3D-RoPE additive latent position embedding (VideoRopePosition3DEmb) +# =========================================================================== +class Cosmos3VideoRopePosition3DEmb(nn.Module): + """Additive 3D-RoPE-style latent position embedding. + + Mirrors ``cosmos_framework`` ``VideoRopePosition3DEmb`` (used when + ``position_embedding_type == "3d_rope"``). Produces an additive + ``[N_vision, head_dim]`` tensor concatenated over per-latent (t, h, w) + grids. ``enable_fps_modulation`` is supported for completeness but the + video path keeps it disabled by default. + """ + + def __init__( + self, + head_dim: int, + len_h: int, + len_w: int, + len_t: int, + base_fps: int = 24, + base_temporal_compression_factor: int = 4, + temporal_compression_factor: int = 4, + h_extrapolation_ratio: float = 1.0, + w_extrapolation_ratio: float = 1.0, + t_extrapolation_ratio: float = 1.0, + enable_fps_modulation: bool = False, + ) -> None: + super().__init__() + self.base_tps = base_fps / base_temporal_compression_factor + self.temporal_compression_factor = temporal_compression_factor + self.max_h = len_h + self.max_w = len_w + self.max_t = len_t + self.enable_fps_modulation = enable_fps_modulation + + dim = head_dim + dim_h = dim // 6 * 2 + dim_w = dim_h + dim_t = dim - 2 * dim_h + assert dim == dim_h + dim_w + dim_t, f"bad dim: {dim} != {dim_h} + {dim_w} + {dim_t}" + self._dim_h = dim_h + self._dim_t = dim_t + + self.register_buffer( + "dim_spatial_range", + torch.arange(0, dim_h, 2)[:(dim_h // 2)].float() / dim_h, + persistent=True, + ) + self.register_buffer( + "dim_temporal_range", + torch.arange(0, dim_t, 2)[:(dim_t // 2)].float() / dim_t, + persistent=True, + ) + self.h_ntk_factor = h_extrapolation_ratio**(dim_h / (dim_h - 2)) + self.w_ntk_factor = w_extrapolation_ratio**(dim_w / (dim_w - 2)) + self.t_ntk_factor = t_extrapolation_ratio**(dim_t / (dim_t - 2)) + + def _generate(self, t: int, h: int, w: int, device: torch.device, fps: torch.Tensor | None, + start_frame_offset: int) -> torch.Tensor: + tps = (fps / self.temporal_compression_factor) if fps is not None else None + + h_theta = 10000.0 * self.h_ntk_factor + w_theta = 10000.0 * self.w_ntk_factor + t_theta = 10000.0 * self.t_ntk_factor + + spatial = self.dim_spatial_range.to(device).float() + temporal = self.dim_temporal_range.to(device).float() + h_spatial_freqs = 1.0 / (h_theta**spatial) + w_spatial_freqs = 1.0 / (w_theta**spatial) + temporal_freqs = 1.0 / (t_theta**temporal) + + max_needed = max(t, h, w) + seq = torch.arange(max_needed, device=device, dtype=torch.float32) + half_emb_h = torch.outer(seq[:h], h_spatial_freqs) # [h, dim_h/2] + half_emb_w = torch.outer(seq[:w], w_spatial_freqs) # [w, dim_w/2] + frame_indices = seq[:t] # [t] + + if self.enable_fps_modulation and tps is not None: + scaled_time = (frame_indices + start_frame_offset) / tps[:1] * self.base_tps + half_emb_t = torch.outer(scaled_time, temporal_freqs) + else: + half_emb_t = torch.outer(frame_indices, temporal_freqs) # [t, dim_t/2] + + emb_t = half_emb_t[:, None, None, :].expand(t, h, w, -1) + emb_h = half_emb_h[None, :, None, :].expand(t, h, w, -1) + emb_w = half_emb_w[None, None, :, :].expand(t, h, w, -1) + rope = torch.cat([emb_t, emb_h, emb_w] * 2, dim=-1) # [t, h, w, head_dim] + return rope.reshape(t * h * w, -1).float() + + def forward( + self, + token_shapes: list[tuple[int, int, int]], + fps: torch.Tensor | None = None, + start_frame_offset: int = 0, + ) -> torch.Tensor: + device = self.dim_spatial_range.device + out = [] + for i, (t, h, w) in enumerate(token_shapes): + video_fps = fps[i:i + 1] if fps is not None else None + out.append(self._generate(t, h, w, device, video_fps, start_frame_offset)) + return torch.cat(out, dim=0) # [N_vision, head_dim] + + +# =========================================================================== +# Timestep embedder (DiT-style; matches modeling_utils.TimestepEmbedder) +# =========================================================================== +class Cosmos3TimestepEmbedder(nn.Module): + """Sinusoidal timestep -> MLP embedder. Checkpoint keys: ``linear_{1,2}``.""" + + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256) -> None: + super().__init__() self.frequency_embedding_size = frequency_embedding_size self.hidden_size = hidden_size + self.linear_1 = ReplicatedLinear(frequency_embedding_size, hidden_size, bias=True) + self.act = nn.SiLU() + self.linear_2 = ReplicatedLinear(hidden_size, hidden_size, bias=True) - half = frequency_embedding_size // 2 - freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half) - self.register_buffer("freqs", freqs, persistent=False) + @staticmethod + def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding def forward(self, t: torch.Tensor) -> torch.Tensor: - raise NotImplementedError("Phase 2b.2") + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + h, _ = self.linear_1(t_freq) + h = self.act(h) + out, _ = self.linear_2(h) + return out -class _Cosmos3GatedMLP(nn.Module): - """Gated-MLP block with ``gate_proj`` / ``up_proj`` / ``down_proj`` names. - - Skeleton of ``Cosmos3GatedMLP`` at reference line 288-326. The three - linear-layer parameter names are load-bearing for the checkpoint - remap; Phase 2b.2 will swap them for FastVideo-native - ``ReplicatedLinear`` shards. - """ +# =========================================================================== +# SwiGLU MLP (Qwen3-VL-text dense MLP) +# =========================================================================== +class Cosmos3MLP(nn.Module): + """SwiGLU MLP: ``down(act(gate(x)) * up(x))``. Keys: gate/up/down_proj.""" def __init__(self, hidden_size: int, intermediate_size: int) -> None: super().__init__() - self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) - self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) - self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + self.gate_proj = ReplicatedLinear(hidden_size, intermediate_size, bias=False) + self.up_proj = ReplicatedLinear(hidden_size, intermediate_size, bias=False) + self.down_proj = ReplicatedLinear(intermediate_size, hidden_size, bias=False) + self.act_fn = nn.SiLU() def forward(self, x: torch.Tensor) -> torch.Tensor: - raise NotImplementedError("Phase 2b.2") + gate, _ = self.gate_proj(x) + up, _ = self.up_proj(x) + out, _ = self.down_proj(self.act_fn(gate) * up) + return out -class _Cosmos3SkeletonAttention(nn.Module): - """Shared skeleton for UND causal- and GEN cross-attention. +# =========================================================================== +# Dual-pathway packed two-way attention +# =========================================================================== +def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, is_causal: bool, scale: float) -> torch.Tensor: + """SDPA on ``[S, heads, head_dim]`` with GQA broadcast. Returns same layout.""" + n_heads = q.shape[1] + n_kv = k.shape[1] + q_ = q.permute(1, 0, 2).unsqueeze(0) # [1, heads, S, head_dim] + k_ = k.permute(1, 0, 2).unsqueeze(0) + v_ = v.permute(1, 0, 2).unsqueeze(0) + if n_kv != n_heads: + k_ = k_.repeat_interleave(n_heads // n_kv, dim=1) + v_ = v_.repeat_interleave(n_heads // n_kv, dim=1) + out = F.scaled_dot_product_attention(q_, k_, v_, is_causal=is_causal, scale=scale) + return out.squeeze(0).permute(1, 0, 2) # [S, heads, head_dim] - Exposes the q/k/v/o projections at the upstream parameter-name - locations (``*.{q,k,v,o}_proj`` and ``*.{q,k}_norm``) with the - GQA-correct in/out dimensions - (``q_proj`` writes ``num_attention_heads * head_dim`` channels; - ``k_proj``/``v_proj`` write ``num_key_value_heads * head_dim`` - channels). Phase 2b.2 will replace the bare ``nn.Linear`` shells with - ``ColumnParallelLinear``-equivalent FastVideo layers and implement - the RoPE + SDPA forward. + +class Cosmos3DualAttention(nn.Module): + """Dual-pathway packed attention (understanding + generation). + + Understanding (text) tokens use causal self-attention; generation (vision) + tokens use full attention where the gen query attends to ALL tokens + (und ++ gen). GQA with per-head QK-norm applied *before* RoPE. Checkpoint + keys: und ``to_{q,k,v}`` / ``to_out`` / ``norm_{q,k}``; gen + ``add_{q,k,v}_proj`` / ``to_add_out`` / ``norm_added_{q,k}``. """ def __init__( @@ -211,244 +382,488 @@ class _Cosmos3SkeletonAttention(nn.Module): num_attention_heads: int, num_key_value_heads: int, head_dim: int, - rms_norm_eps: float, - dtype: torch.dtype, + eps: float, + attention_bias: bool, ) -> None: super().__init__() self.hidden_size = hidden_size - self.num_heads = num_attention_heads - self.num_kv_heads = num_key_value_heads + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads self.head_dim = head_dim + self.scaling = head_dim**-0.5 - self.q_proj = nn.Linear(hidden_size, num_attention_heads * head_dim, bias=False) - self.k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) - self.v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) - self.o_proj = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=False) + q_dim = num_attention_heads * head_dim + kv_dim = num_key_value_heads * head_dim - self.q_norm = _Qwen3VLTextRMSNorm(head_dim, eps=rms_norm_eps, dtype=dtype) - self.k_norm = _Qwen3VLTextRMSNorm(head_dim, eps=rms_norm_eps, dtype=dtype) + # Understanding pathway + self.to_q = ReplicatedLinear(hidden_size, q_dim, bias=attention_bias) + self.to_k = ReplicatedLinear(hidden_size, kv_dim, bias=attention_bias) + self.to_v = ReplicatedLinear(hidden_size, kv_dim, bias=attention_bias) + self.to_out = ReplicatedLinear(q_dim, hidden_size, bias=attention_bias) + self.norm_q = RMSNorm(head_dim, eps=eps) + self.norm_k = RMSNorm(head_dim, eps=eps) - def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: - raise NotImplementedError("Phase 2b.2") + # Generation pathway + self.add_q_proj = ReplicatedLinear(hidden_size, q_dim, bias=attention_bias) + self.add_k_proj = ReplicatedLinear(hidden_size, kv_dim, bias=attention_bias) + self.add_v_proj = ReplicatedLinear(hidden_size, kv_dim, bias=attention_bias) + self.to_add_out = ReplicatedLinear(q_dim, hidden_size, bias=attention_bias) + self.norm_added_q = RMSNorm(head_dim, eps=eps) + self.norm_added_k = RMSNorm(head_dim, eps=eps) + + def forward( + self, + und_seq: torch.Tensor, # [N_und, hidden] + gen_seq: torch.Tensor, # [N_gen, hidden] + cos_und: torch.Tensor, + sin_und: torch.Tensor, + cos_gen: torch.Tensor, + sin_gen: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + n_heads = self.num_attention_heads + n_kv = self.num_key_value_heads + d = self.head_dim + + q_und, _ = self.to_q(und_seq) + k_und, _ = self.to_k(und_seq) + v_und, _ = self.to_v(und_seq) + q_gen, _ = self.add_q_proj(gen_seq) + k_gen, _ = self.add_k_proj(gen_seq) + v_gen, _ = self.add_v_proj(gen_seq) + + q_und = q_und.view(-1, n_heads, d) + k_und = k_und.view(-1, n_kv, d) + v_und = v_und.view(-1, n_kv, d) + q_gen = q_gen.view(-1, n_heads, d) + k_gen = k_gen.view(-1, n_kv, d) + v_gen = v_gen.view(-1, n_kv, d) + + # Per-head QK-norm BEFORE RoPE + q_und = self.norm_q(q_und) + k_und = self.norm_k(k_und) + q_gen = self.norm_added_q(q_gen) + k_gen = self.norm_added_k(k_gen) + + # RoPE (heads-second layout; unsqueeze cos/sin over the head axis at dim=1) + q_und, k_und = _apply_rotary_pos_emb(q_und, k_und, cos_und, sin_und, unsqueeze_dim=1) + q_gen, k_gen = _apply_rotary_pos_emb(q_gen, k_gen, cos_gen, sin_gen, unsqueeze_dim=1) + + # Causal self-attention over understanding (text) tokens + und_out = _sdpa(q_und, k_und, v_und, is_causal=True, scale=self.scaling) + und_out = und_out.reshape(und_out.shape[0], n_heads * d) + + # Full attention: gen query attends to ALL tokens (und ++ gen) + all_k = torch.cat([k_und, k_gen], dim=0) + all_v = torch.cat([v_und, v_gen], dim=0) + gen_out = _sdpa(q_gen, all_k, all_v, is_causal=False, scale=self.scaling) + gen_out = gen_out.reshape(gen_out.shape[0], n_heads * d) + + und_out, _ = self.to_out(und_out) + gen_out, _ = self.to_add_out(gen_out) + return und_out, gen_out -class _Cosmos3UndDecoderLayer(nn.Module): - """UND (understanding) decoder layer skeleton. +# =========================================================================== +# Dual-pathway decoder layer +# =========================================================================== +class Cosmos3DecoderLayer(nn.Module): + """MoT decoder layer: dual-pathway attention + dual SwiGLU MLP, 4 RMSNorms.""" - The submodule names ``self_attn``, ``input_layernorm``, - ``post_attention_layernorm``, ``mlp`` mirror the upstream layer - composition at reference line 624-678. - """ - - def __init__(self, arch: Cosmos3ArchConfig, dtype: torch.dtype) -> None: + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + intermediate_size: int, + eps: float, + attention_bias: bool, + ) -> None: super().__init__() - self.self_attn = _Cosmos3SkeletonAttention( - hidden_size=arch.hidden_size, - num_attention_heads=arch.num_attention_heads, - num_key_value_heads=arch.num_key_value_heads, - head_dim=arch.head_dim, - rms_norm_eps=arch.rms_norm_eps, - dtype=dtype, + self.self_attn = Cosmos3DualAttention( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + eps=eps, + attention_bias=attention_bias, ) - self.input_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype) - self.post_attention_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype) - self.mlp = _Cosmos3GatedMLP(arch.hidden_size, arch.intermediate_size) + self.mlp = Cosmos3MLP(hidden_size, intermediate_size) + self.mlp_moe_gen = Cosmos3MLP(hidden_size, intermediate_size) + self.input_layernorm = RMSNorm(hidden_size, eps=eps) + self.input_layernorm_moe_gen = RMSNorm(hidden_size, eps=eps) + self.post_attention_layernorm = RMSNorm(hidden_size, eps=eps) + self.post_attention_layernorm_moe_gen = RMSNorm(hidden_size, eps=eps) - def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: - raise NotImplementedError("Phase 2b.2") + def forward( + self, + und_seq: torch.Tensor, + gen_seq: torch.Tensor, + cos_und: torch.Tensor, + sin_und: torch.Tensor, + cos_gen: torch.Tensor, + sin_gen: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Pre-attention norm + und_norm = self.input_layernorm(und_seq) + gen_norm = self.input_layernorm_moe_gen(gen_seq) + + und_attn, gen_attn = self.self_attn(und_norm, gen_norm, cos_und, sin_und, cos_gen, sin_gen) + und_res = und_seq + und_attn + gen_res = gen_seq + gen_attn + + # Pre-MLP norm + dual SwiGLU + und_ln = self.post_attention_layernorm(und_res) + gen_ln = self.post_attention_layernorm_moe_gen(gen_res) + und_out = und_res + self.mlp(und_ln) + gen_out = gen_res + self.mlp_moe_gen(gen_ln) + return und_out, gen_out -class _Cosmos3GenDecoderLayer(nn.Module): - """GEN (generation) decoder layer skeleton. +# =========================================================================== +# Domain-aware linear (dormant action head; per-domain weight/bias embeddings) +# =========================================================================== +class _DomainAwareLinear(nn.Module): + """One weight/bias pair per embodiment domain (matches framework keys).""" - Substitutes ``cross_attention`` for ``self_attn`` (mirrors upstream - reference line 681-751); otherwise identical to the UND layer in - naming. The ``layer_idx`` is captured for the eventual cached-KV - path so Phase 2b.2 can route per-layer UND keys/values correctly. - """ - - def __init__(self, layer_idx: int, arch: Cosmos3ArchConfig, dtype: torch.dtype) -> None: + def __init__(self, input_size: int, output_size: int, num_domains: int) -> None: super().__init__() - self.layer_idx = layer_idx - self.cross_attention = _Cosmos3SkeletonAttention( - hidden_size=arch.hidden_size, - num_attention_heads=arch.num_attention_heads, - num_key_value_heads=arch.num_key_value_heads, - head_dim=arch.head_dim, - rms_norm_eps=arch.rms_norm_eps, - dtype=dtype, - ) - self.input_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype) - self.post_attention_layernorm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype) - self.mlp = _Cosmos3GatedMLP(arch.hidden_size, arch.intermediate_size) + self.input_size = int(input_size) + self.output_size = int(output_size) + self.num_domains = int(num_domains) + self.fc = nn.Embedding(self.num_domains, self.output_size * self.input_size) + self.bias = nn.Embedding(self.num_domains, self.output_size) - def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: - raise NotImplementedError("Phase 2b.2") - - -class _Cosmos3LanguageModel(nn.Module): - """Understanding-pathway container. - - Holds ``embed_tokens``, ``layers``, and a final ``norm``. The - submodule names mirror the upstream reference at line 757-831 so the - checkpoint remap target keys ``language_model.embed_tokens.*``, - ``language_model.layers.{i}.*``, and ``language_model.norm.*`` are - populated as-is by the FastVideo module tree. - """ - - def __init__(self, arch: Cosmos3ArchConfig, dtype: torch.dtype) -> None: - super().__init__() - self.embed_tokens = nn.Embedding(arch.vocab_size, arch.hidden_size) - self.layers = nn.ModuleList( - [_Cosmos3UndDecoderLayer(arch, dtype) for _ in range(arch.num_hidden_layers)]) - self.norm = _Qwen3VLTextRMSNorm(arch.hidden_size, eps=arch.rms_norm_eps, dtype=dtype) - - def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: - raise NotImplementedError("Phase 2b.2") + def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: + if domain_id.ndim == 0: + domain_id = domain_id.unsqueeze(0) + domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1) + weight = self.fc(domain_id).view(domain_id.shape[0], self.input_size, self.output_size) + bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size) + if x.ndim == 2: + return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias + return torch.bmm(x, weight) + bias.unsqueeze(1) +# =========================================================================== +# Full DiT +# =========================================================================== class Cosmos3VFMTransformer(BaseDiT): - """Cosmos3 VFM Transformer — UND language model + GEN denoising layers. + """FastVideo-native Cosmos3 omni DiT. - Phase 2b.1 ships only the module-tree skeleton and the two pure math - utilities (``patchify``/``unpatchify``). All per-layer ``forward`` - methods raise ``NotImplementedError("Phase 2b.2")``; the DiT-level - ``forward`` likewise. The constructor signature - ``(od_config, *, temporal_compression_factor=None)`` matches the - upstream reference at ``transformer_cosmos3.py`` line 899 so the - Phase 2a parity tests can drive it directly. - - Phase 2b.2 will: - * Replace the per-layer ``nn.Linear``s with FastVideo-native - ``ReplicatedLinear`` / ``DistributedAttention`` plumbing. - * Implement ``forward`` (UND once, GEN per-step with cached K/V). - * Adapt the constructor to the FastVideo - ``BaseDiT(config, hf_config)`` loader contract via - ``TransformerLoader``. + Mirrors the published checkpoint's transformer key surface and the + ``cosmos_framework`` ``Cosmos3VFMNetwork`` forward math for the video path. + Runs on CPU / float32 (plain SDPA + native-math RoPE). """ _fsdp_shard_conditions = Cosmos3VideoConfig().arch_config._fsdp_shard_conditions _compile_conditions = Cosmos3VideoConfig().arch_config._compile_conditions param_names_mapping = Cosmos3VideoConfig().arch_config.param_names_mapping - reverse_param_names_mapping: dict[str, str] = {} + reverse_param_names_mapping: dict = {} - def __init__( - self, - od_config: object | None = None, - *, - temporal_compression_factor: int | None = None, - ) -> None: - nn.Module.__init__(self) + def __init__(self, config: Cosmos3VideoConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + arch = config.arch_config - model_config = getattr(od_config, "tf_model_config", None) if od_config is not None else None - rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {} + self.hidden_size = arch.hidden_size + self.num_attention_heads = arch.num_attention_heads + self.num_key_value_heads = arch.num_key_value_heads + self.head_dim = arch.head_dim + self.num_hidden_layers = arch.num_hidden_layers + self.intermediate_size = arch.intermediate_size + self.vocab_size = arch.vocab_size + self.rms_norm_eps = arch.rms_norm_eps + self.attention_bias = arch.attention_bias - self.hidden_size = int(_tf_config_get(model_config, "hidden_size", 4096)) - self.num_hidden_layers = int(_tf_config_get(model_config, "num_hidden_layers", 36)) - self.num_attention_heads = int(_tf_config_get(model_config, "num_attention_heads", 32)) - self.num_key_value_heads = int(_tf_config_get(model_config, "num_key_value_heads", 8)) - self.head_dim = int(_tf_config_get(model_config, "head_dim", 128)) - self.intermediate_size = int(_tf_config_get(model_config, "intermediate_size", 12288)) - self.vocab_size = int(_tf_config_get(model_config, "vocab_size", 151936)) - self.rms_norm_eps = float(_tf_config_get(model_config, "rms_norm_eps", 1e-6)) - self.rope_theta = float(_tf_config_get(model_config, "rope_theta", 5_000_000)) - self.mrope_section = list(rope_scaling.get("mrope_section", [24, 20, 20])) - self.latent_patch_size = int(_tf_config_get(model_config, "latent_patch_size", 2)) - self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48)) - self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001)) - self.base_fps = float(_tf_config_get(model_config, "base_fps", 24.0)) - if temporal_compression_factor is None: - resolved_tcf = int(_tf_config_get(model_config, "temporal_compression_factor", 4)) - else: - resolved_tcf = int(temporal_compression_factor) - self.temporal_compression_factor = resolved_tcf - self.enable_fps_modulation = bool(_tf_config_get(model_config, "enable_fps_modulation", True)) - self.temporal_modality_margin = int( - _tf_config_get(model_config, "unified_3d_mrope_temporal_modality_margin", 15000)) + # VAE / patch geometry + self.latent_patch_size = arch.latent_patch_size + self.latent_channel = arch.latent_channel + # Alias kept for the standalone patchify/unpatchify helpers + scaffold tests. + self.latent_channel_size = arch.latent_channel + self.patch_latent_dim = arch.patch_latent_dim + self.num_channels_latents = arch.latent_channel + self.timestep_scale = arch.timestep_scale + self.position_embedding_type = arch.position_embedding_type - self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size - self.num_channels_latents = self.latent_channel_size + # ---- Backbone ---- + self.embed_tokens = nn.Embedding(self.vocab_size, self.hidden_size) + self.layers = nn.ModuleList([ + Cosmos3DecoderLayer( + hidden_size=self.hidden_size, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + head_dim=self.head_dim, + intermediate_size=self.intermediate_size, + eps=self.rms_norm_eps, + attention_bias=self.attention_bias, + ) for _ in range(self.num_hidden_layers) + ]) + self.norm = RMSNorm(self.hidden_size, eps=self.rms_norm_eps) + self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps) + self.lm_head = nn.Linear(self.hidden_size, self.vocab_size, bias=False) - dtype = getattr(od_config, "dtype", torch.bfloat16) if od_config is not None else torch.bfloat16 - - arch = Cosmos3ArchConfig( - hidden_size=self.hidden_size, - num_hidden_layers=self.num_hidden_layers, - num_attention_heads=self.num_attention_heads, - num_key_value_heads=self.num_key_value_heads, + # Backbone RoPE (unified 3D-MRoPE) + self.rotary_emb = Cosmos3TextRotaryEmbedding( head_dim=self.head_dim, - intermediate_size=self.intermediate_size, - vocab_size=self.vocab_size, - rms_norm_eps=self.rms_norm_eps, - rope_theta=self.rope_theta, - mrope_section=self.mrope_section, - latent_patch_size=self.latent_patch_size, - latent_channel=self.latent_channel_size, - timestep_scale=self.timestep_scale, - base_fps=self.base_fps, - temporal_compression_factor=self.temporal_compression_factor, - enable_fps_modulation=self.enable_fps_modulation, - temporal_modality_margin=self.temporal_modality_margin, - in_channels=self.latent_channel_size, - out_channels=self.latent_channel_size, + rope_theta=arch.rope_theta, + mrope_section=arch.mrope_section, ) - self.language_model = _Cosmos3LanguageModel(arch, dtype) - self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size) - self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim) - self.time_embedder = _TimestepEmbedder(self.hidden_size, target_dtype=dtype) - self.gen_layers = nn.ModuleList( - [_Cosmos3GenDecoderLayer(i, arch, dtype) for i in range(self.num_hidden_layers)]) - self.norm_moe_gen = _Qwen3VLTextRMSNorm(self.hidden_size, eps=self.rms_norm_eps, dtype=dtype) + # ---- Vision head ---- + self.proj_in = ReplicatedLinear(self.patch_latent_dim, self.hidden_size, bias=True) + self.proj_out = ReplicatedLinear(self.hidden_size, self.patch_latent_dim, bias=True) + self.time_embedder = Cosmos3TimestepEmbedder(self.hidden_size) + # Additive latent position embedding only for the legacy 3d_rope variant. + self.latent_pos_embed: Cosmos3VideoRopePosition3DEmb | None = None + if self.position_embedding_type == "3d_rope": + self.latent_pos_embed = Cosmos3VideoRopePosition3DEmb( + head_dim=self.hidden_size, + len_h=getattr(arch, "max_latent_h", 32), + len_w=getattr(arch, "max_latent_w", 32), + len_t=getattr(arch, "max_latent_t", 32), + base_fps=int(arch.base_fps), + base_temporal_compression_factor=arch.temporal_compression_factor, + temporal_compression_factor=arch.temporal_compression_factor, + enable_fps_modulation=arch.enable_fps_modulation, + ) + + # ---- Dormant action head (constructed for strict-load parity) ---- + if getattr(arch, "action_gen", False): + self.action_dim = arch.action_dim + self.num_embodiment_domains = arch.num_embodiment_domains + self.action_proj_in = _DomainAwareLinear(self.action_dim, self.hidden_size, self.num_embodiment_domains) + self.action_proj_out = _DomainAwareLinear(self.hidden_size, self.action_dim, self.num_embodiment_domains) + self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size)) + + # ---- Dormant audio / sound head (constructed for strict-load parity) ---- + if getattr(arch, "sound_gen", False): + self.sound_dim = arch.sound_dim + self.audio_proj_in = ReplicatedLinear(self.sound_dim, self.hidden_size, bias=True) + self.audio_proj_out = ReplicatedLinear(self.hidden_size, self.sound_dim, bias=True) + self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size)) + + self.__post_init__() + + # ------------------------------------------------------------------ + # Standalone batched patchify / unpatchify (scaffold-test contract) + # ------------------------------------------------------------------ def _pad_to_patch_size(self, h: int, w: int) -> tuple[int, int, int, int]: - """Return ``(hp, wp, H_padded, W_padded)`` for ``latent_patch_size`` padding. - - Mirrors ``Cosmos3VFMTransformer._pad_to_patch_size`` at reference - lines 1002-1007. - """ + """Return ``(hp, wp, H_padded, W_padded)`` for ``latent_patch_size`` padding.""" p = self.latent_patch_size h_padded = ((h + p - 1) // p) * p w_padded = ((w + p - 1) // p) * p return h_padded // p, w_padded // p, h_padded, w_padded def patchify(self, latents: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor: - """``[B, C, t, h, w] -> [B, t*hp*wp, p*p*C]``. - - Pads ``h``/``w`` up to a multiple of ``latent_patch_size`` before - reshaping. Verbatim port of - ``Cosmos3VFMTransformer.patchify`` (reference lines 1009-1021). - """ + """``[B, C, t, h, w] -> [B, t*hp*wp, p*p*C]`` (zero-pad h/w to patch multiples).""" batch_size = latents.shape[0] p = self.latent_patch_size c = self.latent_channel_size hp, wp, h_padded, w_padded = self._pad_to_patch_size(h, w) - if h_padded != h or w_padded != w: latents = F.pad(latents, (0, w_padded - w, 0, h_padded - h)) - x = latents.reshape(batch_size, c, t, hp, p, wp, p) x = x.permute(0, 2, 3, 5, 4, 6, 1) return x.reshape(batch_size, t * hp * wp, p * p * c) def unpatchify(self, tokens: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor: - """``[B, t*hp*wp, p*p*C] -> [B, C, t, h, w]``, cropping ``h``/``w`` padding. - - Verbatim port of ``Cosmos3VFMTransformer.unpatchify`` (reference - lines 1023-1036). - """ + """``[B, t*hp*wp, p*p*C] -> [B, C, t, h, w]`` (crop h/w padding).""" batch_size = tokens.shape[0] p = self.latent_patch_size c = self.latent_channel_size hp, wp, h_padded, w_padded = self._pad_to_patch_size(h, w) - x = tokens.reshape(batch_size, t, hp, wp, p, p, c) x = x.permute(0, 6, 1, 2, 4, 3, 5) x = x.reshape(batch_size, c, t, h_padded, w_padded) - if h_padded != h or w_padded != w: x = x[:, :, :, :h, :w] return x - def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: - raise NotImplementedError("Phase 2b.2") + # ------------------------------------------------------------------ + # Framework-faithful packed patchify / unpatchify (einsum ordering) + # ------------------------------------------------------------------ + def _patchify_and_pack( + self, + tokens_vision: list[torch.Tensor], + token_shapes: list[tuple[int, int, int]], + ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]: + p = self.latent_patch_size + packed = [] + original_shapes: list[tuple[int, int, int]] = [] + for latent, (_t, _h, _w) in zip(tokens_vision, token_shapes): + latent = latent.squeeze(0) if latent.dim() == 5 else latent # [C, T, H, W] + _, t_actual, h_actual, w_actual = latent.shape + original_shapes.append((t_actual, h_actual, w_actual)) + h_padded = ((h_actual + p - 1) // p) * p + w_padded = ((w_actual + p - 1) // p) * p + if h_padded != h_actual or w_padded != w_actual: + padded = torch.zeros((self.latent_channel, t_actual, h_padded, w_padded), + device=latent.device, + dtype=latent.dtype) + padded[:, :, :h_actual, :w_actual] = latent + latent = padded + h_patches = h_padded // p + w_patches = w_padded // p + latent = latent.reshape(self.latent_channel, t_actual, h_patches, p, w_patches, p) + latent = torch.einsum("cthpwq->thwpqc", latent).reshape(-1, p * p * self.latent_channel) + packed.append(latent) + return torch.cat(packed, dim=0), original_shapes + + def _unpatchify_and_unpack( + self, + packed_preds: torch.Tensor, + token_shapes: list[tuple[int, int, int]], + noisy_frame_indexes: list[torch.Tensor], + original_shapes: list[tuple[int, int, int]] | None, + ) -> list[torch.Tensor]: + p = self.latent_patch_size + outputs = [] + start_idx = 0 + for i, (t_c, _h_c, _w_c) in enumerate(token_shapes): + if original_shapes is not None: + _t_orig, h_orig, w_orig = original_shapes[i] + h_padded = ((h_orig + p - 1) // p) * p + w_padded = ((w_orig + p - 1) // p) * p + h_patches = h_padded // p + w_patches = w_padded // p + else: + h_orig, w_orig = _h_c * p, _w_c * p + h_patches, w_patches = _h_c, _w_c + + nfi = noisy_frame_indexes[i] + t_n = len(nfi) + out = torch.zeros((self.latent_channel, t_c, h_orig, w_orig), + device=packed_preds.device, + dtype=packed_preds.dtype) + num_patches = t_n * h_patches * w_patches + if num_patches > 0: + end_idx = start_idx + num_patches + patches = packed_preds[start_idx:end_idx] + patches = patches.reshape(t_n, h_patches, w_patches, p, p, self.latent_channel) + latent = torch.einsum("thwpqc->cthpwq", patches) + latent = latent.reshape(self.latent_channel, t_n, h_patches * p, w_patches * p) + latent = latent[:, :, :h_orig, :w_orig] + out[:, nfi] = latent + start_idx = end_idx + outputs.append(out.unsqueeze(0)) # [1, C, T, H, W] + return outputs + + def _scatter_timestep_embeds( + self, + packed_tokens: torch.Tensor, + packed_timestep_embeds: torch.Tensor, + noisy_frame_indexes: list[torch.Tensor], + token_shapes: list[tuple[int, ...]], + ) -> torch.Tensor: + """Add timestep embeds onto noisy patches (matches framework scatter_add).""" + start_noisy_index = 0 + flat_idx = [] + for noisy_i, shape_i in zip(noisy_frame_indexes, token_shapes): + spatial = math.prod(shape_i[1:]) + spatial_idx = torch.arange(spatial, device=packed_tokens.device) + ni = (noisy_i * spatial).unsqueeze(-1).expand(-1, spatial) + ni = ni.clone() + spatial_idx + start_noisy_index + flat_idx.append(ni.flatten()) + start_noisy_index += math.prod(shape_i) + flat = torch.cat(flat_idx, dim=0) + flat = flat.unsqueeze(-1).expand(-1, packed_tokens.shape[1]) + return packed_tokens.scatter_add(0, flat, packed_timestep_embeds) + + # ------------------------------------------------------------------ + # forward + # ------------------------------------------------------------------ + def forward( # type: ignore[override] + self, + text_ids: torch.Tensor, + text_indexes: torch.Tensor, + position_ids: torch.Tensor, + sequence_length: int, + split_lens: list[int], + attn_modes: list[str], + vision_tokens: list[torch.Tensor], + vision_token_shapes: list[tuple[int, int, int]], + vision_sequence_indexes: torch.Tensor, + vision_timesteps: torch.Tensor, + vision_mse_loss_indexes: torch.Tensor, + vision_noisy_frame_indexes: list[torch.Tensor], + fps_vision: torch.Tensor | None = None, + **kwargs: Any, + ) -> dict[str, Any]: + """Video-path forward returning ``{"last_hidden_state", "preds_vision"}``. + + Inputs mirror the framework ``PackedSequence`` fields for a single + text(causal)+vision(full) sample. This is the surface the parity test + and the converter exercise; a pipeline-facing wrapper composes these + from a ``PackedSequence`` builder. + """ + device = self.embed_tokens.weight.device + + # 1. Text embedding scattered into the packed sequence buffer + text_embed = self.embed_tokens(text_ids) # [N_text, hidden] + target_dtype = text_embed.dtype + packed_sequence = text_embed.new_zeros((sequence_length, self.hidden_size)) + packed_sequence[text_indexes] = text_embed + + # 2. Vision: patchify -> proj_in -> (+ additive 3d-rope pos emb) -> timestep scatter + original_shapes: list[tuple[int, int, int]] | None = None + if vision_tokens: + packed_vision, original_shapes = self._patchify_and_pack(vision_tokens, vision_token_shapes) + packed_vision, _ = self.proj_in(packed_vision) + + if self.latent_pos_embed is not None: + pos_emb = self.latent_pos_embed(vision_token_shapes, fps=fps_vision).to(target_dtype) + packed_vision = packed_vision + pos_emb + + if vision_mse_loss_indexes.numel() > 0: + timesteps = vision_timesteps * self.timestep_scale + ts_embeds = self.time_embedder(timesteps).to(target_dtype) + packed_vision = self._scatter_timestep_embeds(packed_vision, ts_embeds, vision_noisy_frame_indexes, + vision_token_shapes) + + packed_sequence[vision_sequence_indexes] = packed_vision + + # 3. RoPE for the full packed sequence (cos/sin per token); split und/gen + cos, sin = self.rotary_emb(position_ids, device=device, dtype=target_dtype) # [N, head_dim] + und_idx, gen_idx = self._mode_indices(split_lens, attn_modes, device) + cos_und, sin_und = cos[und_idx], sin[und_idx] + cos_gen, sin_gen = cos[gen_idx], sin[gen_idx] + + # 4. Dual-pathway decoder over (und = text causal, gen = vision full) + und_seq = packed_sequence[und_idx] + gen_seq = packed_sequence[gen_idx] + for layer in self.layers: + und_seq, gen_seq = layer(und_seq, gen_seq, cos_und, sin_und, cos_gen, sin_gen) + + # 5. Final norms (und vs gen) and re-scatter into a joint buffer + und_seq = self.norm(und_seq) + gen_seq = self.norm_moe_gen(gen_seq) + last_hidden_state = packed_sequence.new_zeros((sequence_length, self.hidden_size)) + last_hidden_state[und_idx] = und_seq + last_hidden_state[gen_idx] = gen_seq + + # 6. Vision prediction: proj_out on noisy patches -> unpatchify + output: dict[str, Any] = {} + if vision_tokens and vision_mse_loss_indexes.numel() > 0: + preds, _ = self.proj_out(last_hidden_state[vision_mse_loss_indexes]) + output["preds_vision"] = self._unpatchify_and_unpack(preds, vision_token_shapes, + vision_noisy_frame_indexes, original_shapes) + + output["last_hidden_state"] = last_hidden_state + return output + + @staticmethod + def _mode_indices(split_lens: list[int], attn_modes: list[str], + device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + """Token indexes for causal (und) and full (gen) splits, in pack order.""" + und, gen = [], [] + start = 0 + for split_len, mode in zip(split_lens, attn_modes): + rng = range(start, start + split_len) + if mode == "causal": + und.extend(rng) + elif mode == "full": + gen.extend(rng) + start += split_len + return (torch.tensor(und, dtype=torch.long, device=device), + torch.tensor(gen, dtype=torch.long, device=device)) diff --git a/tests/local_tests/cosmos3/test_cosmos3_dit_parity.py b/tests/local_tests/cosmos3/test_cosmos3_dit_parity.py new file mode 100644 index 000000000..835212493 --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_dit_parity.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical-parity test: FastVideo Cosmos3 DiT vs official ``Cosmos3VFMNetwork``. + +Builds a tiny official-framework ``Cosmos3VFMNetwork`` AND a tiny FastVideo +``Cosmos3VFMTransformer`` from the SAME tiny config, copies the framework +weights into the FastVideo DiT via an explicit framework->fastvideo name map, +runs BOTH forwards on identical deterministic inputs (CPU / float32), and +asserts ``torch.allclose`` on the vision prediction output (``preds_vision``) +and the per-token ``last_hidden_state``. + +The official model is the parity ORACLE. It runs on CPU / float32 via the SDPA +monkey-patch in ``test_cosmos3_reference_forward`` (flash2/flash3/natten are +CUDA-only). The FastVideo DiT runs natively on CPU with plain SDPA. + +Run: + cd && -m pytest \ + tests/local_tests/cosmos3/test_cosmos3_dit_parity.py -q +""" +from __future__ import annotations + +import pytest +import torch + +# The official framework provides the parity oracle. +cosmos_framework = pytest.importorskip( + "cosmos_framework", + reason="cosmos_framework not installed; run in fv-cosmos3 env.", +) + +# Reuse the reference harness's tiny-model builder + SDPA monkey-patch. +from .test_cosmos3_reference_forward import ( # noqa: E402 + _apply_sdpa_patches, + _build_tiny_cosmos3, + _build_tiny_packed_seq, +) + +pytestmark = [pytest.mark.local] + +# Ensure the CPU/float32 SDPA patches are installed (idempotent). +_apply_sdpa_patches() + + +# --------------------------------------------------------------------------- +# Tiny config shared by both models (must match _build_tiny_cosmos3). +# --------------------------------------------------------------------------- +def _build_tiny_fastvideo_dit() -> "Cosmos3VFMTransformer": # noqa: F821 + from fastvideo.configs.models.dits.cosmos3 import ( + Cosmos3ArchConfig, + Cosmos3VideoConfig, + ) + from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer + + arch = Cosmos3ArchConfig( + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=8, + intermediate_size=32, + vocab_size=64, + rms_norm_eps=1e-6, + attention_bias=False, + latent_patch_size=2, + latent_channel=16, + rope_theta=5_000_000.0, + mrope_section=[24, 20, 20], + position_embedding_type="3d_rope", + base_fps=24.0, + temporal_compression_factor=4, + enable_fps_modulation=False, + # Dormant heads present in the checkpoint surface (constructed for + # strict-load parity; not exercised by this video-path forward). + action_gen=True, + action_dim=64, + max_action_dim=64, + num_embodiment_domains=32, + sound_gen=True, + sound_dim=64, + ) + cfg = Cosmos3VideoConfig(arch_config=arch) + model = Cosmos3VFMTransformer(cfg, hf_config={}) + return model.to(torch.float32).eval() + + +# --------------------------------------------------------------------------- +# Framework -> FastVideo weight name map. +# --------------------------------------------------------------------------- +def _framework_to_fastvideo_state_dict(vfm, num_layers: int) -> dict[str, torch.Tensor]: + """Translate framework param names into the FastVideo DiT param names. + + Framework (Cosmos3VFMNetwork): + language_model.model.{embed_tokens,norm,norm_moe_gen} + language_model.lm_head + language_model.model.layers.{i}.self_attn.{q,k,v,o}_proj(+ _moe_gen) + language_model.model.layers.{i}.self_attn.{q,k}_norm(+ _moe_gen) + language_model.model.layers.{i}.{mlp,mlp_moe_gen}.{gate,up,down}_proj + language_model.model.layers.{i}.{input,post_attention}_layernorm(+ _moe_gen) + vae2llm / llm2vae / time_embedder.mlp.{0,2} + + FastVideo (Cosmos3VFMTransformer): + embed_tokens / norm / norm_moe_gen / lm_head + layers.{i}.self_attn.{to_q,to_k,to_v,to_out} (und) + layers.{i}.self_attn.{add_q,add_k,add_v}_proj / to_add_out (gen) + layers.{i}.self_attn.{norm_q,norm_k,norm_added_q,norm_added_k} + layers.{i}.{mlp,mlp_moe_gen}.{gate,up,down}_proj + layers.{i}.{input,post_attention}_layernorm(+ _moe_gen) + proj_in / proj_out / time_embedder.linear_{1,2} + """ + src = dict(vfm.named_parameters()) + out: dict[str, torch.Tensor] = {} + + def take(name: str) -> torch.Tensor: + return src[name].detach().clone() + + # ---- Top-level backbone ---- + out["embed_tokens.weight"] = take("language_model.model.embed_tokens.weight") + out["norm.weight"] = take("language_model.model.norm.weight") + out["norm_moe_gen.weight"] = take("language_model.model.norm_moe_gen.weight") + out["lm_head.weight"] = take("language_model.lm_head.weight") + + # ---- Vision adapters ---- + out["proj_in.weight"] = take("vae2llm.weight") + out["proj_in.bias"] = take("vae2llm.bias") + out["proj_out.weight"] = take("llm2vae.weight") + out["proj_out.bias"] = take("llm2vae.bias") + + # ---- Timestep embedder (mlp.0/mlp.2 -> linear_1/linear_2) ---- + out["time_embedder.linear_1.weight"] = take("time_embedder.mlp.0.weight") + out["time_embedder.linear_1.bias"] = take("time_embedder.mlp.0.bias") + out["time_embedder.linear_2.weight"] = take("time_embedder.mlp.2.weight") + out["time_embedder.linear_2.bias"] = take("time_embedder.mlp.2.bias") + + # ---- Per layer ---- + und_attn = {"q_proj": "to_q", "k_proj": "to_k", "v_proj": "to_v", "o_proj": "to_out"} + gen_attn = { + "q_proj_moe_gen": "add_q_proj", + "k_proj_moe_gen": "add_k_proj", + "v_proj_moe_gen": "add_v_proj", + "o_proj_moe_gen": "to_add_out", + } + und_norm = {"q_norm": "norm_q", "k_norm": "norm_k"} + gen_norm = {"q_norm_moe_gen": "norm_added_q", "k_norm_moe_gen": "norm_added_k"} + + for i in range(num_layers): + fw = f"language_model.model.layers.{i}" + fv = f"layers.{i}" + for s, d in und_attn.items(): + out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight") + for s, d in gen_attn.items(): + out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight") + for s, d in und_norm.items(): + out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight") + for s, d in gen_norm.items(): + out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight") + for mlp in ("mlp", "mlp_moe_gen"): + for proj in ("gate_proj", "up_proj", "down_proj"): + out[f"{fv}.{mlp}.{proj}.weight"] = take(f"{fw}.{mlp}.{proj}.weight") + for ln in ("input_layernorm", "input_layernorm_moe_gen", "post_attention_layernorm", + "post_attention_layernorm_moe_gen"): + out[f"{fv}.{ln}.weight"] = take(f"{fw}.{ln}.weight") + + return out + + +def _copy_weights(vfm, dit) -> None: + """Copy framework weights into the FastVideo DiT (shape-checked).""" + mapped = _framework_to_fastvideo_state_dict(vfm, num_layers=dit.num_hidden_layers) + dst = dict(dit.named_parameters()) + # Every mapped tensor must land on an existing FastVideo param with a matching shape. + for name, tensor in mapped.items(): + assert name in dst, f"FastVideo DiT missing param for mapped key {name!r}" + assert dst[name].shape == tensor.shape, (f"shape mismatch for {name}: " + f"dit={tuple(dst[name].shape)} fw={tuple(tensor.shape)}") + with torch.no_grad(): + for name, tensor in mapped.items(): + dst[name].copy_(tensor.to(dst[name].dtype)) + + +def _fastvideo_inputs_from_packed_seq(ps) -> dict: + """Build the FastVideo DiT forward kwargs from a framework PackedSequence.""" + v = ps.vision + return dict( + text_ids=ps.text_ids, + text_indexes=ps.text_indexes, + position_ids=ps.position_ids, + sequence_length=int(ps.sequence_length), + split_lens=list(ps.split_lens), + attn_modes=list(ps.attn_modes), + vision_tokens=list(v.tokens), + vision_token_shapes=list(v.token_shapes), + vision_sequence_indexes=v.sequence_indexes, + vision_timesteps=v.timesteps, + vision_mse_loss_indexes=v.mse_loss_indexes, + vision_noisy_frame_indexes=list(v.noisy_frame_indexes), + fps_vision=None, + ) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- +class TestCosmos3DiTParity: + + def _run_both(self, seed_model: int = 42, seed_data: int = 7): + vfm = _build_tiny_cosmos3(seed=seed_model) + dit = _build_tiny_fastvideo_dit() + _copy_weights(vfm, dit) + ps = _build_tiny_packed_seq(n_text=4, seed=seed_data) + + with torch.no_grad(): + fw_out = vfm(packed_seq=ps) + fv_out = dit(**_fastvideo_inputs_from_packed_seq(ps)) + return fw_out, fv_out + + def test_weight_map_is_complete(self): + """The framework->fastvideo map must cover EVERY FastVideo DiT parameter + that is exercised by the video path (i.e. all non-dormant params). + + Dormant action/audio heads have no framework counterpart in this tiny + vision-only setup, so they are excluded from the copy; everything else + must be covered. + """ + vfm = _build_tiny_cosmos3(seed=42) + dit = _build_tiny_fastvideo_dit() + mapped = set(_framework_to_fastvideo_state_dict(vfm, num_layers=dit.num_hidden_layers)) + dit_params = set(n for n, _ in dit.named_parameters()) + dormant = { + n + for n in dit_params + if n.startswith(("action_", "audio_")) + } + uncovered = dit_params - mapped - dormant + assert not uncovered, f"FastVideo DiT params not covered by weight map: {sorted(uncovered)}" + + def test_preds_vision_parity(self): + fw_out, fv_out = self._run_both() + fw_pv = fw_out["preds_vision"][0] # [1, C, T, H, W] + fv_pv = fv_out["preds_vision"][0] + assert fw_pv.shape == fv_pv.shape, f"shape mismatch: fw={fw_pv.shape} fv={fv_pv.shape}" + max_abs = (fw_pv - fv_pv).abs().max().item() + print(f"\n[preds_vision] max abs diff = {max_abs:.3e}") + torch.testing.assert_close(fv_pv, fw_pv, atol=1e-4, rtol=1e-3) + + def test_last_hidden_state_parity(self): + fw_out, fv_out = self._run_both() + fw_lhs = fw_out["last_hidden_state"] # [N, hidden] + fv_lhs = fv_out["last_hidden_state"] + assert fw_lhs.shape == fv_lhs.shape, f"shape mismatch: fw={fw_lhs.shape} fv={fv_lhs.shape}" + max_abs = (fw_lhs - fv_lhs).abs().max().item() + print(f"\n[last_hidden_state] max abs diff = {max_abs:.3e}") + torch.testing.assert_close(fv_lhs, fw_lhs, atol=1e-4, rtol=1e-3) + + def test_parity_holds_across_seeds(self): + """Re-running with a different random init still matches (not a fluke).""" + fw_out, fv_out = self._run_both(seed_model=99, seed_data=13) + torch.testing.assert_close(fv_out["preds_vision"][0], fw_out["preds_vision"][0], atol=1e-4, rtol=1e-3) diff --git a/tests/local_tests/cosmos3/test_cosmos3_state_dict_keys.py b/tests/local_tests/cosmos3/test_cosmos3_state_dict_keys.py index 45b31e9da..24f279ea3 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_state_dict_keys.py +++ b/tests/local_tests/cosmos3/test_cosmos3_state_dict_keys.py @@ -1,28 +1,38 @@ # SPDX-License-Identifier: Apache-2.0 -"""Cosmos3 state-dict-key contract for checkpoint conversion (Tier A scaffold). +"""Cosmos3 DiT param-name / checkpoint-key-surface contract. -Reference: ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:319-409`` -(``Cosmos3OmniDiffusersPipeline._remap_ckpt_key``). The vllm-omni loader's -remap is the canonical source of truth: it defines which Diffusers -checkpoint keys become which target module-tree parameter names. +The FastVideo native Cosmos3 DiT (``fastvideo.models.dits.cosmos3. +Cosmos3VFMTransformer``) deliberately mirrors the *published diffusers +checkpoint* transformer key surface so the converter at +``scripts/checkpoint_conversion/cosmos3_convert.py`` can strict-load with a +near-identity ``param_names_mapping``. -The FastVideo conversion script at -``scripts/checkpoint_conversion/cosmos3_convert.py`` (to be authored in -Phase 5) must produce a state-dict that loads cleanly into the FastVideo -Cosmos3 DiT module tree. This test pins the target side of that contract: +This pins the FastVideo side of that contract: a tiny 1-layer DiT must expose +exactly the published checkpoint's parameter names. The checkpoint key surface +(validated against ``nvidia/Cosmos3-Nano``; ``{i}`` ranges over the layers): - * ``embed_tokens``/``norm`` live under ``language_model.`` - * ``norm_moe_gen`` is a top-level (not per-layer) module - * Per-layer keys split into UND (``language_model.layers.{i}``) and - GEN (``gen_layers.{i}``) sub-trees: - - UND attention: ``self_attn.{q,k,v,o}_proj``, ``self_attn.{q,k}_norm`` - - GEN cross-attention: ``cross_attention.{q,k,v,o}_proj``, - ``cross_attention.{q,k}_norm`` (from the ``*_moe_gen`` source keys) - - UND/GEN norms: ``input_layernorm``, ``post_attention_layernorm`` - - UND/GEN MLPs: ``mlp.{gate,up,down}_proj`` - * Top-level adapters ``vae2llm``, ``llm2vae``, ``time_embedder`` live - under ``transformer.`` (no remapping inside the transformer namespace). - * ``lm_head.weight`` is skipped (mapped to ``None``). + Top level: + embed_tokens.weight, norm.weight, norm_moe_gen.weight, lm_head.weight, + proj_in.{weight,bias}, proj_out.{weight,bias}, + time_embedder.linear_1.{weight,bias}, time_embedder.linear_2.{weight,bias} + Dormant heads (present for strict-load): + action_proj_in.fc.weight, action_proj_in.bias.weight, + action_proj_out.fc.weight, action_proj_out.bias.weight, + action_modality_embed, + audio_proj_in.{weight,bias}, audio_proj_out.{weight,bias}, + audio_modality_embed + Per layer ``layers.{i}``: + self_attn.{to_q,to_k,to_v,to_out,add_q_proj,add_k_proj,add_v_proj, + to_add_out,norm_q,norm_k,norm_added_q,norm_added_k}.weight, + mlp.{gate_proj,up_proj,down_proj}.weight, + mlp_moe_gen.{gate_proj,up_proj,down_proj}.weight, + {input_layernorm,input_layernorm_moe_gen,post_attention_layernorm, + post_attention_layernorm_moe_gen}.weight + +The earlier scaffold pinned the dead vllm-omni layout +(``language_model.layers`` / ``gen_layers`` / ``cross_attention`` / +``vae2llm`` / ``llm2vae``); that structure is gone — the native DiT is a single +dual-pathway ``layers`` ModuleList matching the diffusers checkpoint. """ from __future__ import annotations @@ -31,86 +41,129 @@ import pytest pytestmark = [pytest.mark.local] -EXPECTED_REMAPS: dict[str, str | None] = { - "model.embed_tokens.weight": "transformer.language_model.embed_tokens.weight", - "model.norm.weight": "transformer.language_model.norm.weight", - "model.norm_moe_gen.weight": "transformer.norm_moe_gen.weight", - "model.layers.3.self_attn.q_proj.weight": "transformer.language_model.layers.3.self_attn.q_proj.weight", - "model.layers.3.self_attn.q_proj_moe_gen.weight": "transformer.gen_layers.3.cross_attention.q_proj.weight", - "model.layers.3.self_attn.k_norm_moe_gen.weight": "transformer.gen_layers.3.cross_attention.k_norm.weight", - "model.layers.3.input_layernorm.weight": "transformer.language_model.layers.3.input_layernorm.weight", - "model.layers.3.input_layernorm_moe_gen.weight": "transformer.gen_layers.3.input_layernorm.weight", - "model.layers.3.mlp.gate_proj.weight": "transformer.language_model.layers.3.mlp.gate_proj.weight", - "model.layers.3.mlp_moe_gen.up_proj.weight": "transformer.gen_layers.3.mlp.up_proj.weight", - "vae2llm.weight": "transformer.vae2llm.weight", - "llm2vae.weight": "transformer.llm2vae.weight", - "time_embedder.linear_1.weight": "transformer.time_embedder.linear_1.weight", - "lm_head.weight": None, -} - - -def test_fastvideo_cosmos3_remap_matches_reference() -> None: - """Asserts the FastVideo Cosmos3 checkpoint remap table matches the - canonical vllm-omni reference at pipeline_cosmos3.py:319-409. - - Once the FastVideo Cosmos3 pipeline (or a standalone remap utility - in ``scripts/checkpoint_conversion/cosmos3_convert.py``) exposes a - ``_remap_ckpt_key`` callable, this test verifies every key in - ``EXPECTED_REMAPS`` maps to the expected target name (or ``None`` - for skipped keys like ``lm_head.weight``). - """ +def _build_tiny_dit(): + """Construct a tiny 1-layer FastVideo Cosmos3 DiT, or skip if unavailable.""" try: - from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore - Cosmos3OmniDiffusersPipeline, + from fastvideo.configs.models.dits.cosmos3 import ( + Cosmos3ArchConfig, + Cosmos3VideoConfig, ) - except ImportError: - pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)") + from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer + except ImportError: # pragma: no cover - import guard + pytest.skip("FastVideo Cosmos3 DiT not importable in this environment") - remap = getattr(Cosmos3OmniDiffusersPipeline, "_remap_ckpt_key", None) - if remap is None: - pytest.skip("Cosmos3OmniDiffusersPipeline._remap_ckpt_key not yet defined") + import torch - actual = {key: remap(key) for key in EXPECTED_REMAPS} - assert actual == EXPECTED_REMAPS + arch = Cosmos3ArchConfig( + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=8, + intermediate_size=32, + vocab_size=64, + latent_patch_size=2, + latent_channel=16, + position_embedding_type="3d_rope", + enable_fps_modulation=False, + action_gen=True, + action_dim=64, + max_action_dim=64, + num_embodiment_domains=32, + sound_gen=True, + sound_dim=64, + ) + cfg = Cosmos3VideoConfig(arch_config=arch) + return Cosmos3VFMTransformer(cfg, hf_config={}).to(torch.float32) def test_fastvideo_cosmos3_dit_module_tree_param_names() -> None: - """Asserts the FastVideo Cosmos3 DiT module tree produces the param - names that the checkpoint converter is expected to write. + """The native DiT param-name set must equal the published checkpoint surface.""" + model = _build_tiny_dit() + names = {name for name, _ in model.named_parameters()} - Once ``fastvideo.models.dits.cosmos3.Cosmos3VFMTransformer`` is - constructible with a tiny config (see ``conftest._tiny_cosmos3_config``), - this test instantiates a 1-layer model and verifies that the - parameter-name set contains the expected UND/GEN/adapter prefixes: - * ``language_model.layers.0.self_attn.q_proj.weight`` - * ``gen_layers.0.cross_attention.q_proj.weight`` - * ``norm_moe_gen.weight`` - * ``vae2llm.*`` and ``llm2vae.*`` - """ - try: - from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer # type: ignore - except ImportError: - pytest.skip("FastVideo Cosmos3 not yet implemented (Phase 2b)") + expected_top = { + "embed_tokens.weight", + "norm.weight", + "norm_moe_gen.weight", + "lm_head.weight", + "proj_in.weight", + "proj_in.bias", + "proj_out.weight", + "proj_out.bias", + "time_embedder.linear_1.weight", + "time_embedder.linear_1.bias", + "time_embedder.linear_2.weight", + "time_embedder.linear_2.bias", + } + expected_dormant = { + "action_proj_in.fc.weight", + "action_proj_in.bias.weight", + "action_proj_out.fc.weight", + "action_proj_out.bias.weight", + "action_modality_embed", + "audio_proj_in.weight", + "audio_proj_in.bias", + "audio_proj_out.weight", + "audio_proj_out.bias", + "audio_modality_embed", + } + layer_suffixes = { + "self_attn.to_q.weight", + "self_attn.to_k.weight", + "self_attn.to_v.weight", + "self_attn.to_out.weight", + "self_attn.add_q_proj.weight", + "self_attn.add_k_proj.weight", + "self_attn.add_v_proj.weight", + "self_attn.to_add_out.weight", + "self_attn.norm_q.weight", + "self_attn.norm_k.weight", + "self_attn.norm_added_q.weight", + "self_attn.norm_added_k.weight", + "mlp.gate_proj.weight", + "mlp.up_proj.weight", + "mlp.down_proj.weight", + "mlp_moe_gen.gate_proj.weight", + "mlp_moe_gen.up_proj.weight", + "mlp_moe_gen.down_proj.weight", + "input_layernorm.weight", + "input_layernorm_moe_gen.weight", + "post_attention_layernorm.weight", + "post_attention_layernorm_moe_gen.weight", + } + expected_layer = {f"layers.0.{s}" for s in layer_suffixes} + expected = expected_top | expected_dormant | expected_layer - from types import SimpleNamespace - import torch + assert names == expected, (f"Cosmos3 DiT param surface mismatch.\n" + f" missing: {sorted(expected - names)}\n" + f" unexpected: {sorted(names - expected)}") - from .conftest import _tiny_cosmos3_config - config = _tiny_cosmos3_config(num_hidden_layers=1) - model = Cosmos3VFMTransformer(SimpleNamespace(tf_model_config=config, dtype=torch.float32)) - names = set(name for name, _ in model.named_parameters()) - - required_prefixes = [ - "language_model.layers.0.self_attn.q_proj", - "gen_layers.0.cross_attention.q_proj", - "norm_moe_gen", +def test_fastvideo_cosmos3_dit_no_dead_vllm_omni_layout() -> None: + """The dead vllm-omni layout (split language_model/gen_layers/cross_attention, + vae2llm/llm2vae) must NOT appear in the native DiT param tree.""" + model = _build_tiny_dit() + names = {name for name, _ in model.named_parameters()} + dead_fragments = ( + "language_model.", + "gen_layers.", + "cross_attention.", "vae2llm", "llm2vae", - ] - for prefix in required_prefixes: - assert any(name.startswith(prefix) for name in names), ( - f"Cosmos3 DiT module tree missing expected prefix {prefix!r}; " - f"converter at scripts/checkpoint_conversion/cosmos3_convert.py " - f"will fail to load weights for this subtree." - ) + ) + offenders = sorted(n for n in names if any(frag in n for frag in dead_fragments)) + assert not offenders, f"native DiT still exposes dead vllm-omni param names: {offenders}" + + +def test_fastvideo_cosmos3_dit_per_layer_counts() -> None: + """Per-layer block count: 22 weights/layer; full 1-layer model = 44 params. + + (12 attention + 6 dual-MLP + 4 layernorm per layer; + 12 top-level + + 10 dormant-head params.) + """ + model = _build_tiny_dit() + names = [name for name, _ in model.named_parameters()] + per_layer = [n for n in names if n.startswith("layers.0.")] + assert len(per_layer) == 22, f"expected 22 per-layer params, got {len(per_layer)}" + assert len(names) == 44, f"expected 44 total params for tiny 1-layer DiT, got {len(names)}"