feat(models): add NVIDIA Cosmos 3 Edge text reasoner backbone - #448
Open
justinchuby wants to merge 1 commit into
Open
feat(models): add NVIDIA Cosmos 3 Edge text reasoner backbone#448justinchuby wants to merge 1 commit into
justinchuby wants to merge 1 commit into
Conversation
Add support for the text reasoner (language tower) of the cosmos3_edge
vision-language checkpoint (nvidia/Cosmos3-Edge,
Cosmos3EdgeForConditionalGeneration).
The language tower is a standard grouped-query-attention decoder with two
Cosmos-specific traits handled here:
- Non-gated squared-ReLU FFN (hidden_act="relu2",
down_proj(relu2(up_proj(x)))), mapped onto the existing FCMLP component
(Nemotron precedent) instead of the GLU-style gated MLP.
- 3D multimodal RoPE (mrope_section=[24, 20, 20]); for text-only inference
the three sections are identical, reducing to standard 1D RoPE.
preprocess_weights renames the self_attn.to_{q,k,v,out} projections to the
q/k/v/o_proj component names, nests the top-level text tower (layers.*,
embed_tokens, norm) under model., keeps lm_head at the top level, and drops
the vision encoder (model.visual.*), the multimodal projector
(model.projector.*), and the per-layer k_norm_und_for_gen key-norm — the
latter being a two-tower (Mixture-of-Transformers) artifact that normalizes
the understanding tower's keys for the generator (diffusion) tower and is
not applied in the reasoner's own causal self-attention.
Registered as cosmos3_edge / cosmos3_edge_text and exported from
models/__init__.py. L1 graph-build verified via the parametrized
CAUSAL_LM_CONFIGS matrix; end-to-end build from the real config.json
produces a 28-layer GQA decoder with the expected non-gated relu2 FFN.
L4/L5 numerical parity is deferred (NVIDIA's custom edge modeling code is
not in transformers), recorded in _COVERAGE_SKIP. The cosmos3_omni
diffusion world-model variants (Cosmos3-Nano/-Super) are out of scope for
this decoder-only path.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds support in mobius for NVIDIA Cosmos3-Edge’s decoder-only text reasoner backbone by introducing a dedicated model class that reuses the existing CausalLMModel stack while swapping in a non-gated squared-ReLU FFN and applying Cosmos-specific weight-key remapping.
Changes:
- Added
Cosmos3EdgeTextModelwithFCMLP(non-gatedrelu2) and apreprocess_weights()mapping to rename attention projections, nest the text tower undermodel., and drop vision/projector +k_norm_und_for_genweights. - Registered
cosmos3_edge/cosmos3_edge_textin the model registry and exported the model frommobius.models. - Added L1 graph-build coverage via
tests/_test_configs.py, and recorded L4/L5 waiver rationale intests/model_coverage_test.py; documented inCHANGELOG.md.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/model_coverage_test.py | Records Cosmos3-Edge as L1-only coverage with explicit L4/L5 parity rationale. |
| tests/_test_configs.py | Adds cosmos3_edge to the causal LM config matrix (relu2 + mRoPE section) for L1 graph-build. |
| src/mobius/models/cosmos.py | Implements the Cosmos3-Edge text reasoner backbone model + weight preprocessing rules. |
| src/mobius/models/init.py | Exports Cosmos3EdgeTextModel. |
| src/mobius/_registry.py | Registers cosmos3_edge / cosmos3_edge_text to the new model and adds default IDs. |
| CHANGELOG.md | Documents the new model support and scope limitations. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds mobius support for the text reasoner backbone of NVIDIA's
cosmos3_edgevision-language checkpoint (nvidia/Cosmos3-Edge,Cosmos3EdgeForConditionalGeneration).The
Cosmos3series splits into two very different families:cosmos3_edgenvidia/Cosmos3-Edgecosmos3_omniCosmos3-Nano/-SuperThe text tower
A grouped-query-attention decoder (28 layers, 16 q / 8 kv heads, head_dim 128) with two Cosmos-specific traits:
down_proj(relu2(up_proj(x)))(hidden_act="relu2"), mapped onto the existingFCMLPcomponent (Nemotron precedent) instead of the GLU-style gated MLP.mrope_section=[24, 20, 20]; for text-only inference the three sections are identical, reducing to standard 1D RoPE (same simplification as the Qwen-VL text decoders).preprocess_weights:self_attn.to_{q,k,v,out}→q/k/v/o_proj,layers.*,embed_tokens,norm) undermodel., keepslm_headat the top level,model.visual.*), the projector (model.projector.*), and the per-layerk_norm_und_for_genkey-norm.On
k_norm_und_for_genThis per-layer key-norm is a two-tower (MoT) artifact: it normalizes the understanding tower's keys for consumption by the generator (diffusion) tower (hence
..._for_gen). It is not applied in the reasoner's own causal self-attention, so it is dropped for the standalone text decoder. Flagging for reviewer confirmation, since NVIDIA's custom edge modeling code is not yet intransformers.Testing
cosmos3_edgeto the parametrizedCAUSAL_LM_CONFIGSmatrix (relu2 + mrope).pytest tests/build_graph_test.py -k cosmos3_edge→ 8 passed.build("nvidia/Cosmos3-Edge", load_weights=False)produces a 28-layer GQA decoder (236 initializers) with the non-gated relu2 FFN (mlp.up_projpresent,mlp.gate_projabsent). A weight-mapping check confirms every checkpoint tensor lands on a real ONNX initializer with no orphans; vision/projector/k_norm_und_for_gencorrectly dropped.tests/model_coverage_test.py→ passes (both types recorded in_COVERAGE_SKIPwith a reason).Scope / follow-ups
cosmos3_edgemodeling code lands intransformers(onlycosmos3_omniships today) — needed to verify thek_norm_und_for_geninterpretation and generate golden data.SigLIPVisionModelto build on.mobius build --model nvidia/Cosmos3-Edgeroutes to the diffusers autodetect (the repo ships amodel_index.json); the text reasoner is reachable via the Pythonbuild()API / registry.🤖 Generated with GitHub Copilot CLI