feat(models): add NVIDIA Cosmos3-Edge vision-language model - #448
Open
justinchuby wants to merge 2 commits into
Open
feat(models): add NVIDIA Cosmos3-Edge vision-language model#448justinchuby wants to merge 2 commits into
justinchuby wants to merge 2 commits 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
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
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. |
Extend the cosmos3_edge support from a text-only reasoner to the full
vision-language model (nvidia/Cosmos3-Edge,
Cosmos3EdgeForConditionalGeneration) as a LLaVA-style 3-model
onnxruntime-genai split (decoder + vision_encoder + embedding):
- decoder: squared-ReLU GQA text reasoner taking inputs_embeds, built
with 3D multimodal RoPE (mrope_section=[24,20,20]).
- vision_encoder: SigLIP vision tower + new Cosmos3EdgeMultiModalProjector
(pre-shuffle LayerNorm -> 2x2 pixel-shuffle -> linear_fc1 -> GELU ->
linear_fc2).
- embedding: token embedding + image-feature fusion at image_token_id=19.
preprocess_weights routes the single HF checkpoint to the three
sub-models: model.visual.* / model.projector.* -> vision (SigLIP
mlp.fc1/fc2 -> up_proj/down_proj), embed_tokens -> embedding, the
top-level text tower -> decoder (self_attn.to_{q,k,v,out} ->
{q,k,v,o}_proj), and drops the generator-tower k_norm_und_for_gen
key-norm. Built via a new Cosmos3EdgeVLTask ("cosmos3-edge-vl").
A vision config hook reconstructs image_size from num_patches and pulls
the projector's merger_intermediate_size from projector_config. The
decoder-only text reasoner remains available as cosmos3_edge_text.
L1 graph-build tested only: NVIDIA publishes no modeling code for
cosmos3_edge (not in transformers, no remote-code module), so exact
pixel-shuffle ordering and numerical parity are unverifiable.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
justinchuby
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.
Suppressed comments (3)
src/mobius/_configs/per_model/_cosmos3_edge_vision.py:40
- Using round(sqrt(num_patches)) can silently produce the wrong grid for non-square num_patches values, yielding an inconsistent image_size and downstream shape errors. Using math.isqrt() and validating grid*grid == num_patches makes this fail fast with a clear error.
patch_size = getattr(hf_vision, "patch_size", None) or fields.get("patch_size")
if num_patches is not None and patch_size is not None:
grid = round(math.sqrt(num_patches))
fields["image_size"] = grid * patch_size
src/mobius/components/_multimodal.py:159
- Cosmos3EdgeMultiModalProjector implicitly assumes spatial_merge_size > 0 and that grid_size is divisible by spatial_merge_size. If either is violated, the Reshape in forward will fail at runtime with a hard-to-debug shape error. Add explicit validation in init with a clear ValueError.
self._grid = grid_size
self._ms = spatial_merge_size
self._vision_hidden = vision_hidden_size
merged_dim = vision_hidden_size * spatial_merge_size * spatial_merge_size
# Pre-shuffle LayerNorm over the raw vision hidden size.
src/mobius/components/init.py:177
- There are multiple adjacent import blocks from mobius.components._multimodal. Consolidating them into a single grouped import reduces duplication and avoids isort/ruff churn when adding/removing exports.
from mobius.components._multimodal import (
Cosmos3EdgeMultiModalProjector as Cosmos3EdgeMultiModalProjector,
)
from mobius.components._multimodal import (
Gemma3MultiModalProjector as Gemma3MultiModalProjector,
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 the full NVIDIA Cosmos3-Edge vision-language model
(
nvidia/Cosmos3-Edge,Cosmos3EdgeForConditionalGeneration) to mobius as aLLaVA-style 3-model onnxruntime-genai split (
decoder+vision_encoder+embedding), building on the text reasoner backbone from the first commit.Architecture
squared-ReLU FFN (
hidden_act="relu2",up_proj → relu2 → down_proj) and3D multimodal RoPE (
mrope_section=[24, 20, 20]); takesinputs_embeds.Cosmos3EdgeMultiModalProjector(pre-shuffleLayerNorm→ 2x2 pixel-shuffle→
linear_fc1→ GELU →linear_fc2).image_token_id=19.Built through a new
Cosmos3EdgeVLTask(cosmos3-edge-vl). The decoder-onlytext reasoner remains available as
cosmos3_edge_text.Weight routing (
preprocess_weights)Single HF checkpoint → three sub-models:
model.visual.*/model.projector.*→ vision (SigLIPmlp.fc1/fc2→up_proj/down_proj)embed_tokens→ embeddinglayers.*/norm/lm_head) → decoder (self_attn.to_{q,k,v,out}→{q,k,v,o}_proj)k_norm_und_for_gen(generator-tower key-norm) → droppedVerified against the real
nvidia/Cosmos3-Edgesafetensors index: every producedweight key lands on a graph initializer, and every weight-bearing initializer is
covered (only computed RoPE/const tensors are unmatched, as expected).
Confidence
L1 (graph-build) only. NVIDIA does not publish modeling code for
cosmos3_edge(not in
transformers, no remote-code module), so the exact pixel-shuffle orderingand numerical parity are unverifiable; L4/L5 parity is deferred. The
cosmos3_omnivariants (Cosmos3-Nano/-Super) are two-tower diffusion worldmodels tracked separately.
Testing
tests/build_graph_test.py -k cosmos3— all pass (text + VLM 3-model split).build_graph_test.py+model_coverage_test.py+cli_test.pysuites pass.ruff format+ruff checkclean.🤖 Do not merge — awaiting review.