Cosmos 3 代码讲解：Mixture-of-Transformers 双塔统一五模态世界模型

# model/generator/mot/unified_mot.py (simplified)
import torch
import torch.nn as nn
import torch.nn.functional as F


class PackedAttentionMoT(nn.Module):
    def __init__(self, config):
        super().__init__()
        hidden_size = config.hidden_size
        num_heads = config.num_attention_heads
        num_kv_heads = config.num_key_value_heads

        # --- Understanding path projections (causal attention) ---
        self.q_proj = nn.Linear(hidden_size, num_heads * head_dim)
        self.k_proj = nn.Linear(hidden_size, num_kv_heads * head_dim)
        self.v_proj = nn.Linear(hidden_size, num_kv_heads * head_dim)
        self.o_proj = nn.Linear(num_heads * head_dim, hidden_size)

        # --- Generation path projections (full attention) ---
        self.q_proj_moe_gen = nn.Linear(hidden_size, num_heads * head_dim)
        self.k_proj_moe_gen = nn.Linear(hidden_size, num_kv_heads * head_dim)
        self.v_proj_moe_gen = nn.Linear(hidden_size, num_kv_heads * head_dim)
        self.o_proj_moe_gen = nn.Linear(num_heads * head_dim, hidden_size)

        # Optional QK normalization
        self.q_norm = RMSNorm(head_dim)
        self.k_norm = RMSNorm(head_dim)
        self.q_norm_moe_gen = RMSNorm(head_dim)
        self.k_norm_moe_gen = RMSNorm(head_dim)
        # Cross-attention: gen queries attending to und KVs
        self.k_norm_und_for_gen = RMSNorm(head_dim)

    def forward(self, hidden_states, attention_mask, position_ids, moe_gen_mask):
        # Split hidden states into und and gen paths
        und_mask = ~moe_gen_mask
        und_h = hidden_states[und_mask]
        gen_h = hidden_states[moe_gen_mask]

        # --- Understanding path: causal QKV ---
        und_q = self.q_norm(self.q_proj(und_h))
        und_k = self.k_norm(self.k_proj(und_h))
        und_v = self.v_proj(und_h)
        und_out = scaled_dot_product_attention(und_q, und_k, und_v, is_causal=True)

        # --- Generation path: full attention + cross-attention to und ---
        gen_q = self.q_norm_moe_gen(self.q_proj_moe_gen(gen_h))
        gen_k = self.k_norm_moe_gen(self.k_proj_moe_gen(gen_h))
        gen_v = self.v_proj_moe_gen(gen_h)
        # Self-attention among gen tokens (full, not causal)
        gen_out = scaled_dot_product_attention(gen_q, gen_k, gen_v, is_causal=False)
        # Cross-attention: gen queries attend to und key/values
        und_k_for_gen = self.k_norm_und_for_gen(self.k_proj(und_h.detach()))
        gen_cross = scaled_dot_product_attention(gen_q, und_k_for_gen, und_v.detach(), is_causal=False)

        gen_out = gen_out + gen_cross  # combine self + cross
        gen_out = self.o_proj_moe_gen(gen_out)
        und_out = self.o_proj(und_out)

        # Reassemble into original order
        output = torch.zeros_like(hidden_states)
        output[und_mask] = und_out
        output[moe_gen_mask] = gen_out
        return output

class MoTDecoderLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.input_layernorm = nn.LayerNorm(config.hidden_size)
        self.input_layernorm_moe_gen = nn.LayerNorm(config.hidden_size)
        self.self_attn = PackedAttentionMoT(config)
        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size)
        self.post_attention_layernorm_moe_gen = nn.LayerNorm(config.hidden_size)
        self.mlp = MLP(config)  # und path
        self.mlp_moe_gen = MLP(config)  # gen path

    def forward(self, hidden_states, attention_mask, position_ids, moe_gen_mask):
        residual = hidden_states
        # Dual layernorm: und and gen get different norms
        und_normed = self.input_layernorm(hidden_states)
        gen_normed = self.input_layernorm_moe_gen(hidden_states)
        # Merge back: pick norm based on path
        normed = torch.where(moe_gen_mask.unsqueeze(-1), gen_normed, und_normed)
        attn_out = self.self_attn(normed, attention_mask, position_ids, moe_gen_mask)
        hidden_states = residual + attn_out

        residual = hidden_states
        # Dual post-attention norm
        und_normed = self.post_attention_layernorm(hidden_states)
        gen_normed = self.post_attention_layernorm_moe_gen(hidden_states)
        normed = torch.where(moe_gen_mask.unsqueeze(-1), gen_normed, und_normed)
        # Dual FFN
        ff_out = torch.where(
            moe_gen_mask.unsqueeze(-1),
            self.mlp_moe_gen(normed),
            self.mlp(normed),
        )
        hidden_states = residual + ff_out
        return hidden_states

# 形状追踪（Nano 规格：H=4096, n_heads=32, head_dim=128）
input_hidden = torch.randn(1, 500, 4096)   # (B, L, H) = (1, 100+400, 4096)
moe_gen_mask = torch.cat([                  # 前100个False, 后400个True
    torch.zeros(100, dtype=torch.bool),
    torch.ones(400, dtype=torch.bool),
])

# 拆分成两路
und_h = input_hidden[0, :100]               # (100, 4096) — und 路径
gen_h = input_hidden[0, 100:]               # (400, 4096) — gen 路径

# Und 注意力: 32 heads × 128 head_dim
und_q = q_proj(und_h).reshape(100, 32, 128)  # (100, 32, 128)
und_k = k_proj(und_h).reshape(100, 4, 128)   # (100, 4, 128) — GQA: 4 KV heads
und_v = v_proj(und_h).reshape(100, 4, 128)
# causal attention → und_out: (100, 4096)

# Gen 注意力: 同样 32 heads, 不同权重
gen_q = q_proj_moe_gen(gen_h).reshape(400, 32, 128)
gen_k = k_proj_moe_gen(gen_h).reshape(400, 4, 128)
gen_v = v_proj_moe_gen(gen_h).reshape(400, 4, 128)
# full attention → gen_self_out: (400, 4096)

# Cross-attention: gen query × und key/value
gen_cross = cross_attention(gen_q,                           # (400, 32, 128)
    k_norm_und_for_gen(k_proj(und_h)),                       # und 的 key
    v_proj(und_h))                                           # und 的 value
# → gen_cross_out: (400, 4096)

gen_out = gen_self_out + gen_cross  # 组合
# 最后 FFN: 每路各走各的 MLP
# und: MLP(4096 → 12288 → 4096)  (Nano 规格)
# gen: MLP_moe_gen(4096 → 12288 → 4096)
# 输出: (1, 500, 4096) — und 和 gen 重新拼回

# model/generator/mot/omni_mot_model.py (simplified)
class OmniMoTModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            [MoTDecoderLayer(config) for _ in range(config.num_hidden_layers)]
        )
        self.norm = nn.LayerNorm(config.hidden_size)

    def forward(self, input_ids, moe_gen_mask, position_ids, attention_mask=None):
        h = self.embed_tokens(input_ids)
        for layer in self.layers:
            h = layer(h, attention_mask, position_ids, moe_gen_mask)
        h = self.norm(h)
        return h  # (B, L, H) final hidden states

# 更完整的 OmniMoTModel 调用示意
class OmniMoTForCausalLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.model = OmniMoTModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.gen_head = DiffusionHead(config)  # projects gen hidden → vector field

    def forward(self, input_ids, moe_gen_mask, labels=None, gen_targets=None):
        hidden = self.model(input_ids, moe_gen_mask, ...)
        # Und tokens go to lm_head (next-token prediction)
        und_logits = self.lm_head(hidden[~moe_gen_mask])
        # Gen tokens go to gen_head (vector field prediction)
        gen_pred = self.gen_head(hidden[moe_gen_mask])
        loss = 0
        if labels is not None:
            loss += cross_entropy(und_logits, labels)  # reasoner loss
        if gen_targets is not None:
            loss += mse_loss(gen_pred, gen_targets)    # generator loss
        return loss

# model/generator/tokenizers/dc_ae/ (simplified)
class DCAEEncoder(nn.Module):
    def __init__(self, latent_channels=64, downsample_factor=32, temporal_downsample=4):
        super().__init__()
        # Stack of 3D convolutions with increasing channels
        self.stem = nn.Conv3d(3, 128, kernel_size=3, padding=1)
        self.blocks = nn.ModuleList([
            DownsampleBlock3D(128, 256,  stride=(1, 2, 2)),  # spatial /2
            DownsampleBlock3D(256, 512,  stride=(1, 2, 2)),  # spatial /4
            DownsampleBlock3D(512, 512,  stride=(1, 2, 2)),  # spatial /8
            DownsampleBlock3D(512, 512,  stride=(2, 2, 2)),  # temporal/2, spatial/16
            DownsampleBlock3D(512, latent_channels, stride=(2, 2, 2)),  # temporal/4, spatial/32
        ])
        self.causal = True  # causal padding throughout

    def forward(self, video):
        # video: (B, T, H, W, 3) normalized to [-1, 1]
        # permute to (B, C, T, H, W) for Conv3d
        x = video.permute(0, 4, 1, 2, 3)
        x = self.stem(x)
        for block in self.blocks:
            x = block(x)
        return x  # (B, latent_channels, T//4, H//32, W//32)


class DCAEDecoder(nn.Module):
    def __init__(self, latent_channels=64):
        super().__init__()
        # Symmetric to encoder: upsample blocks
        self.blocks = nn.ModuleList([
            UpsampleBlock3D(latent_channels, 512, stride=(2, 2, 2)),
            UpsampleBlock3D(512, 512, stride=(2, 2, 2)),
            UpsampleBlock3D(512, 512, stride=(1, 2, 2)),
            UpsampleBlock3D(512, 256, stride=(1, 2, 2)),
            UpsampleBlock3D(256, 128, stride=(1, 2, 2)),
        ])
        self.head = nn.Conv3d(128, 3, kernel_size=3, padding=1)

    def forward(self, latent):
        for block in self.blocks:
            latent = block(latent)
        out = self.head(latent)
        return out.permute(0, 2, 3, 4, 1)  # (B, T, H, W, 3)

# dc_ae/modules/causal_conv3d.py (simplified concept)
class CausalConv3d(nn.Conv3d):
    def __init__(self, in_channels, out_channels, kernel_size,
                 stride=1, dilation=1):
        # Standard Conv3d with causal padding
        # kernel_size assumed (k_t, k_h, k_w)
        pad_t = (kernel_size[0] - 1) * dilation[0]  # only pad past, not future
        pad_h = (kernel_size[1] - 1) * dilation[1] // 2
        pad_w = (kernel_size[2] - 1) * dilation[2] // 2
        # padding: (W_left, W_right, H_top, H_bottom, T_front, T_back)
        # T_back=0 ensures causal: current frame only sees past frames
        padding = (pad_w, pad_w, pad_h, pad_h, pad_t, 0)
        super().__init__(in_channels, out_channels, kernel_size,
                         stride=stride, padding=padding, dilation=dilation)

    def forward(self, x):
        # Standard Conv3d with causal padding already set in __init__
        return super().forward(x)

# dc_ae/modules/quant.py (simplified concept)
class DCAEQuantizer(nn.Module):
    def __init__(self, latent_channels=64, quant_type="lpips"):
        super().__init__()
        self.conv_in = nn.Conv3d(latent_channels, latent_channels, 1)
        self.conv_out = nn.Conv3d(latent_channels, latent_channels, 1)
        self.quant_type = quant_type  # "lpips" or "vq"

    def forward(self, x):
        # Apply channel-wise scaling based on perceptual importance
        x = self.conv_in(x)
        if self.quant_type == "vq":
            # Vector quantization (for discrete latent space)
            x, _ = self.vector_quantization(x)
        # "lpips" mode: no actual quantization, just learnable scaling
        x = self.conv_out(x)
        return x

# DCAE training loss components
L_dcae = L_recon + λ_perc · L_perceptual + λ_gan · L_gan + λ_t · L_temporal + β · L_kl

# configs/base/defaults/model_config.py (simplified)
@dataclass
class RectifiedFlowTrainingConfig:
    shift: int = 5                     # distribution shift for noise schedule
    train_time_image_distribution: str = "logitnormal"
    train_time_video_distribution: str = "logitnormal"
    train_time_action_distribution: str = "logitnormal"
    loss_scale: float = 1.0
    action_loss_weight: float = 10.0   # action loss weighted higher
    
@dataclass
class RectifiedFlowInferenceConfig:
    scheduler_type: str = "unipc"      # UniPC multi-step
    num_train_timesteps: int = 1000
    shift: int = 1

def rectified_flow_loss(model, video_latent, text_input_ids, moe_gen_mask):
    # video_latent: (B, C, T_lat, H_lat, W_lat) from DCAE encoder
    # Flatten spatial-temporal dims to sequence
    B, C, T, H, W = video_latent.shape
    x_1 = video_latent.reshape(B, C, -1).transpose(1, 2)  # (B, T*H*W, C)
    
    # Sample noise and time
    noise = torch.randn_like(x_1)
    t = sample_logit_normal((B, 1, 1), shift=5)  # logit-normal distribution
    
    # Linear interpolation: x_t = (1-t) * noise + t * data
    x_t = (1 - t) * noise + t * x_1
    
    # Forward through MoT (packed with text tokens)
    packed_input = pack_sequence(text_input_ids, x_t, moe_gen_mask)
    hidden = model(packed_input, moe_gen_mask)
    
    # Get gen token predictions (vector field)
    gen_hidden = hidden[moe_gen_mask]
    v_pred = gen_head(gen_hidden)  # (B*T*H*W, C)
    
    # Target: v = x_1 - noise (direction from noise to data)
    v_target = x_1 - noise
    
    loss = F.mse_loss(v_pred, v_target.reshape(-1, C))
    return loss

# data/generator/sequence_packing/runtime.py (simplified)
def pack_sequence(text_tokens, gen_latents, config):
    # text_tokens: (B, L_text) token ids for und path
    # gen_latents: (B, N_gen, C) flattened video/audio latent tokens for gen path
    
    # 1. Create moe_gen_mask: True for gen tokens, False for und tokens
    und_len = text_tokens.shape[1]
    gen_len = gen_latents.shape[1]
    total_len = und_len + gen_len
    
    moe_gen_mask = torch.cat([
        torch.zeros(und_len, dtype=torch.bool),
        torch.ones(gen_len, dtype=torch.bool),
    ])
    
    # 2. Concatenate text + latent along sequence dimension
    packed_input_ids = torch.cat([text_tokens], dim=1)  # text only (latents go through different path)
    
    # 3. Create combined hidden states
    packed_hidden = torch.cat([
        text_embeddings,      # (B, und_len, H)
        gen_proj(gen_latents),  # (B, gen_len, H)
    ], dim=1)
    
    return packed_hidden, moe_gen_mask

# 注意力 mask 示意
# und tokens (causal):  ✓ 每行只能看当前列及以左
# gen tokens (full):    ✓ 任意 gen 看所有 gen
# gen→und cross:       ✓ gen 看所有 und
# und→gen:             ✗ 不看（因果上 und 不能看到未来的生成结果）

# mask 形状: (total_len, total_len)
# und1 und2 gen1 gen2
# und1  ✓   ✗   ✗   ✗
# und2  ✓   ✓   ✗   ✗
# gen1  ✓   ✓   ✓   ✓
# gen2  ✓   ✓   ✓   ✓

# scripts/inference.py (simplified concept)
@torch.no_grad()
def text_to_video(prompt, model, dcae, tokenizer, num_steps=35, guidance_scale=6.0):
    # 1. Tokenize text prompt
    text_ids = tokenizer.encode(prompt)  # (B, L_text)
    
    # 2. Initialize noise in latent space
    # For 5s @24fps video: T=120 → T_lat=30 (DCAE /4), H=320→10 (/32), W=512→16 (/32)
    noise = torch.randn(1, 64, 30, 10, 16)  # (B, C, T_lat, H_lat, W_lat)
    latent = noise.clone()
    
    # 3. UniPC sampling loop
    scheduler = UniPCScheduler(num_train_timesteps=1000, num_inference_steps=35)
    timesteps = scheduler.get_timesteps()
    
    for t in timesteps:
        # Pack und (text) + gen (noisy latent) tokens
        gen_tokens = flatten_latent(latent)  # (B, N_gen, C)
        packed_hidden, moe_gen_mask = pack_sequence(text_ids, gen_tokens, ...)
        
        # MoT forward: get gen path vector field
        hidden = model(packed_hidden, moe_gen_mask)
        gen_hidden = hidden[moe_gen_mask]
        v_pred = model.gen_head(gen_hidden)
        
        # Reshape back to latent space
        v_pred = unflatten_to_latent(v_pred, latent.shape)
        
        # Classifier-free guidance
        v_uncond = ...  # forward with empty prompt
        v_guided = v_uncond + guidance_scale * (v_pred - v_uncond)
        
        # Scheduler step
        latent = scheduler.step(v_guided, t, latent)
    
    # 4. Decode latent to video
    video = dcae.decoder(latent)  # (B, T, H, W, 3)
    return video

# configs/base/defaults/model_config.py (simplified structure)
@dataclass
class OmniMoTModelConfig:
    # Transformer backbone
    hidden_size: int = 4096
    intermediate_size: int = 12288
    num_hidden_layers: int = 36
    num_attention_heads: int = 32
    num_key_value_heads: int = 8
    rms_norm_eps: float = 1e-6
    
    # MoT specific
    moe_gen_enabled: bool = True  # toggle MoT on/off
    # Use separate layernorm for gen path
    moe_gen_layernorm: bool = True
    # Cross-attention config
    moe_gen_cross_attention: bool = True
    moe_gen_k_norm_und: bool = True
    
@dataclass
class DiffusionExpertConfig:
    # Rectified Flow training
    shift: int = 5
    loss_scale: float = 1.0
    action_loss_weight: float = 10.0
    
    # Inference
    scheduler_type: str = "unipc"
    num_inference_steps: int = 35
    guidance_scale: float = 6.0
    
@dataclass
class TokenizerConfig:
    video_tokenizer: str = "dcae"  # or "cosmos_videovae"
    audio_tokenizer: str = "avae"
    tokenizer_type: str = "CausalVAE3D"
    latent_channels: int = 64
    base_channels: int = 128
    channel_multipliers: tuple = (1, 2, 4, 4)
    down_samples: tuple = (1, 2, 2, 2)  # temporal down factors
    spatial_compression: int = 32
    temporal_compression: int = 4

@dataclass
class ReasonerConfig:
    backbone: str = "qwen3_vl"  # or "nemotron_3_dense_vl"
    freeze_backbone: bool = False
    enable_gradient_checkpointing: bool = True

# configs/experiment/nano_text_to_video.yaml
defaults:
  - base_text_to_video
  - _self_

model:
  hidden_size: 4096
  num_hidden_layers: 36
  num_attention_heads: 32
  moe_gen_enabled: true

training:
  per_device_train_batch_size: 1
  gradient_accumulation_steps: 8
  num_train_epochs: 1

# trainer/train.py (simplified concept)
class CosmosTrainer:
    def __init__(self, config):
        self.model = OmniMoTForCausalLM(config)
        self.dcae = load_frozen_dcae()  # frozen tokenizer
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=3e-4)
        
    def train_step(self, batch):
        # batch: video, text, audio, action (depending on task)
        video_latent = self.dcae.encoder(batch.video)  # (B, C, T/4, H/32, W/32)
        
        # Pack sequence: text tokens (und) + video latents (gen)
        packed, moe_gen_mask = pack_sequence(batch.text, video_latent)
        
        # Rectified Flow loss on gen path
        loss = self.compute_rf_loss(packed, moe_gen_mask, video_latent)
        
        loss.backward()
        self.optimizer.step()
        return loss.item()

# 冻结 und 路径，只训练 gen 路径
for name, param in model.named_parameters():
    if 'moe_gen' not in name:
        param.requires_grad = False

# scripts/inference.py (simplified routing)
def run_inference(args):
    if args.mode == "text_to_video":
        return text_to_video(args.prompt, ...)
    elif args.mode == "image_to_video":
        return image_to_video(args.image, args.prompt, ...)
    elif args.mode == "video_to_video":
        return video_to_video(args.input_video, args.edit_prompt, ...)
    elif args.mode == "text_to_action":
        return text_to_action(args.prompt, ...)
    elif args.mode == "text_to_audio":
        return text_to_audio(args.prompt, ...)
    elif args.mode == "text_and_image_to_video":
        return text_and_image_to_video(args.prompt, args.image, ...)

# image_to_video 的核心逻辑
def image_to_video(image, prompt, model, dcae, ...):
    # 1. 把输入图过一遍 DCAE encoder，拿到 latent
    img_latent = dcae.encoder(image.unsqueeze(1))  # (B, C, 1, H_lat, W_lat)
    
    # 2. 噪声 latent 只在 t>0 的时间步上初始化
    noise = torch.randn(B, C, T_lat - 1, H_lat, W_lat)
    # 首帧直接用 img_latent，后续帧从噪声开始去噪
    latent = torch.cat([img_latent, noise], dim=2)
    
    # 3. 序列打包时，首帧 token 走 und 路径（提供条件），
    #    后续帧走 gen 路径（生成）
    # 这样 und 路径看到了"开头长什么样"，gen 路径负责"延续下去"

# auxiliary/guardrail/ (simplified concept)
class SafetyGuardrail:
    def __init__(self):
        self.prompt_classifier = LlamaGuard3ForSequenceClassification(...)
        self.output_classifier = LlamaGuard3ForSequenceClassification(...)
        self.face_blur = FaceBlurModule()
        self.content_filter = ContentFilter()
        
    def check_prompt(self, text):
        # 检查输入 prompt 是否包含不安全内容
        result = self.prompt_classifier(text)
        if result.unsafe:
            return False, result.violated_categories
        return True, None
    
    def filter_output(self, video):
        # 对生成的视频做后处理
        video = self.face_blur(video)       # 人脸模糊
        video = self.content_filter(video)  # 内容过滤（暴力/色情检测）
        return video