|
| 1 | +# Copyright (c) Microsoft Corporation. |
| 2 | +# Licensed under the MIT License. |
| 3 | + |
| 4 | +"""NVIDIA Cosmos 3 model support. |
| 5 | +
|
| 6 | +Currently supports the **text reasoner backbone** of the ``cosmos3_edge`` |
| 7 | +checkpoint (``nvidia/Cosmos3-Edge``, ``Cosmos3EdgeForConditionalGeneration``). |
| 8 | +
|
| 9 | +Cosmos3-Edge is a vision-language model whose language tower is a standard |
| 10 | +grouped-query-attention decoder with two Cosmos-specific traits: |
| 11 | +
|
| 12 | +- **Non-gated feed-forward network** — ``down_proj(relu2(up_proj(x)))`` using |
| 13 | + a squared-ReLU activation (``hidden_act="relu2"``), rather than the |
| 14 | + GLU-style gated MLP used by Llama/Qwen. This maps onto :class:`FCMLP`. |
| 15 | +- **3D multimodal RoPE** (``mrope_section=[24, 20, 20]``). For text-only |
| 16 | + inference the three sections are identical, so it reduces to standard 1D |
| 17 | + RoPE (same simplification used by the Qwen-VL text decoders). |
| 18 | +
|
| 19 | +The HuggingFace weights use ``self_attn.to_{q,k,v,out}`` projection names and |
| 20 | +place the text tower at the top level (``layers.*``, ``embed_tokens``, |
| 21 | +``norm``, ``lm_head``) with the vision encoder and projector under |
| 22 | +``model.visual.*`` / ``model.projector.*``. :meth:`preprocess_weights` renames |
| 23 | +the projections, prefixes the text tower with ``model.`` and drops the |
| 24 | +vision/projector weights. |
| 25 | +
|
| 26 | +The ``k_norm_und_for_gen`` per-layer key-norm weight is an artifact of the |
| 27 | +two-tower (Mixture-of-Transformers) design: it normalizes the *understanding* |
| 28 | +tower's keys for consumption by the *generator* (diffusion) tower, and is not |
| 29 | +applied in the reasoner's own causal self-attention. It is therefore dropped |
| 30 | +for the standalone text decoder. |
| 31 | +
|
| 32 | +The ``cosmos3_omni`` variants (``nvidia/Cosmos3-Nano`` / ``-Super``) are |
| 33 | +two-tower diffusion world models exported as diffusers pipelines and are out |
| 34 | +of scope for this decoder-only path. |
| 35 | +""" |
| 36 | + |
| 37 | +from __future__ import annotations |
| 38 | + |
| 39 | +from typing import TYPE_CHECKING |
| 40 | + |
| 41 | +from mobius._configs import ArchitectureConfig |
| 42 | +from mobius.components import FCMLP |
| 43 | +from mobius.models.base import CausalLMModel |
| 44 | + |
| 45 | +if TYPE_CHECKING: |
| 46 | + import torch |
| 47 | + |
| 48 | + |
| 49 | +class Cosmos3EdgeTextModel(CausalLMModel): |
| 50 | + """Cosmos3-Edge text reasoner backbone (decoder-only). |
| 51 | +
|
| 52 | + Extracts the language tower from the ``cosmos3_edge`` vision-language |
| 53 | + checkpoint. Replaces the gated MLP with a non-gated squared-ReLU |
| 54 | + :class:`FCMLP` and renames/strips weights so the standard |
| 55 | + :class:`CausalLMModel` backbone can consume them. |
| 56 | + """ |
| 57 | + |
| 58 | + def __init__(self, config: ArchitectureConfig): |
| 59 | + super().__init__(config) |
| 60 | + # Cosmos3-Edge uses a non-gated FFN (up_proj -> relu2 -> down_proj), |
| 61 | + # unlike the GLU-style gated MLP of the base CausalLMModel. |
| 62 | + for layer in self.model.layers: |
| 63 | + layer.mlp = FCMLP( |
| 64 | + config.hidden_size, |
| 65 | + config.intermediate_size, |
| 66 | + activation=config.hidden_act or "relu2", |
| 67 | + bias=config.mlp_bias, |
| 68 | + ) |
| 69 | + |
| 70 | + def preprocess_weights( |
| 71 | + self, state_dict: dict[str, torch.Tensor] |
| 72 | + ) -> dict[str, torch.Tensor]: |
| 73 | + renamed: dict[str, torch.Tensor] = {} |
| 74 | + for key, value in state_dict.items(): |
| 75 | + # Drop the vision encoder and multimodal projector — this is the |
| 76 | + # standalone text decoder. |
| 77 | + if key.startswith(("model.visual.", "model.projector.")): |
| 78 | + continue |
| 79 | + # Drop the generator-tower key-norm (see module docstring). |
| 80 | + if "k_norm_und_for_gen" in key: |
| 81 | + continue |
| 82 | + |
| 83 | + new_key = ( |
| 84 | + key.replace("self_attn.to_q.", "self_attn.q_proj.") |
| 85 | + .replace("self_attn.to_k.", "self_attn.k_proj.") |
| 86 | + .replace("self_attn.to_v.", "self_attn.v_proj.") |
| 87 | + .replace("self_attn.to_out.", "self_attn.o_proj.") |
| 88 | + ) |
| 89 | + # The text tower is stored at the top level; the mobius backbone |
| 90 | + # nests it under ``model.``. ``lm_head`` stays at the top level. |
| 91 | + if new_key == "lm_head.weight": |
| 92 | + pass |
| 93 | + elif new_key.startswith(("layers.", "embed_tokens.")) or new_key == "norm.weight": |
| 94 | + new_key = f"model.{new_key}" |
| 95 | + |
| 96 | + renamed[new_key] = value |
| 97 | + return super().preprocess_weights(renamed) |
0 commit comments