Cosmos3 super reasoner

src/maxtext/checkpoint_conversion/utils/hf_model_configs.py

@@ -1933,6 +1933,56 @@ def __init__(self, **kwargs):
 }
 cosmos3_nano_reasoner_config = PTConfig(**cosmos3_nano_reasoner_dict)
 
+cosmos3_super_reasoner_dict = {
+    "architectures": ["Cosmos3ForConditionalGeneration"],
+    "model_type": "cosmos3_omni",
+    "text_config": {
+        "attention_bias": False,
+        "attention_dropout": 0.0,
+        "bos_token_id": 151643,
+        "dtype": "bfloat16",
+        "eos_token_id": 151645,
+        "head_dim": 128,
+        "hidden_act": "silu",
+        "hidden_size": 5120,
+        "initializer_range": 0.02,
+        "intermediate_size": 25600,
+        "max_position_embeddings": 262144,
+        "model_type": "qwen3_vl_text",
+        "num_attention_heads": 64,
+        "num_hidden_layers": 64,
+        "num_key_value_heads": 8,
+        "pad_token_id": None,
+        "rms_norm_eps": 1e-06,
+        "rope_parameters": {
+            "mrope_interleaved": True,
+            "mrope_section": [24, 20, 20],
+            "rope_theta": 5000000,
+            "rope_type": "default",
+        },
+        "tie_word_embeddings": False,
+        "use_cache": True,
+        "vocab_size": 151936,
+    },
+    "vision_config": {
+        "deepstack_visual_indexes": [8, 16, 24],
+        "depth": 27,
+        "hidden_act": "gelu_pytorch_tanh",
+        "hidden_size": 1152,
+        "in_channels": 3,
+        "initializer_range": 0.02,
+        "intermediate_size": 4304,
+        "model_type": "qwen3_vl_vision",
+        "num_heads": 16,
+        "num_position_embeddings": 2304,
+        "out_hidden_size": 5120,
+        "patch_size": 16,
+        "spatial_merge_size": 2,
+        "temporal_patch_size": 2,
+    },
+}
+cosmos3_super_reasoner_config = PTConfig(**cosmos3_super_reasoner_dict)
+
 
 # {maxtext model name: hf model config}
 HF_MODEL_CONFIGS = {
@@ -1964,6 +2014,7 @@ def __init__(self, **kwargs):
     "qwen3-vl-4b": qwen3_vl_4b_config,
     "qwen3-vl-30b-a3b": qwen3_vl_30b_a3b_config,
     "cosmos3-nano-reasoner": cosmos3_nano_reasoner_config,
+    "cosmos3-super-reasoner": cosmos3_super_reasoner_config,
     "llama3.1-8b": llama31_8b_config,
     "llama3.1-8b-Instruct": llama31_8b_config,
     "llama3.1-70b": llama31_70b_config,

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -4280,6 +4280,7 @@ def mhc_concat_scale(input_tensors, target_shape=None):
     "qwen3-vl-4b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING,
     "qwen3-vl-30b-a3b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING,
     "cosmos3-nano-reasoner": COSMOS3_MAXTEXT_TO_HF_PARAM_MAPPING,
+    "cosmos3-super-reasoner": COSMOS3_MAXTEXT_TO_HF_PARAM_MAPPING,
     "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING,
     "llama3.1-8b-Instruct": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING,
     "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING,
@@ -4335,6 +4336,7 @@ def mhc_concat_scale(input_tensors, target_shape=None):
     "qwen3-vl-4b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "qwen3-vl-30b-a3b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "cosmos3-nano-reasoner": COSMOS3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
+    "cosmos3-super-reasoner": COSMOS3_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "llama3.1-8b-Instruct": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN,
     "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN,

src/maxtext/configs/models/cosmos3-super-reasoner.yml

@@ -0,0 +1,41 @@
+# Model config for NVIDIA Cosmos3-Super Reasoner
+
+# Core Architectural Parameters
+decoder_block: "qwen3"
+base_emb_dim: 5120
+base_mlp_dim: 25600
+base_num_query_heads: 64
+base_num_kv_heads: 8
+base_num_decoder_layers: 64
+head_dim: 128
+mlp_activations: ["silu", "linear"]
+vocab_size: 151936
+normalization_layer_epsilon: 1.0e-6
+use_qk_norm: true
+logits_via_embedding: false
+
+# RoPE Settings
+rope_max_timescale: 5000000
+
+# General Model Settings
+enable_dropout: false
+
+# Multimodal Settings
+use_multimodal: true
+use_mrope: true
+mrope_section: [24, 20, 20]
+
+# Vision Encoder Configuration
+image_size_for_vit: 768
+hidden_size_for_vit: 1152
+intermediate_size_for_vit: 4304
+num_attention_heads_for_vit: 16
+num_hidden_layers_for_vit: 27
+num_channels_for_vit: 3
+patch_size_for_vit: 16
+temporal_patch_size_for_vit: 2
+spatial_merge_size_for_vit: 2
+out_hidden_size_for_vit: 5120
+num_position_embeddings_for_vit: 2304
+deepstack_visual_indexes_for_vit: [8, 16, 24]
+vision_encoder_block: "qwen3_vl"

src/maxtext/configs/types.py

@@ -278,6 +278,7 @@ class ProfilerType(str, Enum):
     "qwen3-vl-4b",
     "qwen3-vl-30b-a3b",
     "cosmos3-nano-reasoner",
+    "cosmos3-super-reasoner",
     "qwen3-next-80b-a3b",
     "qwen3-omni-30b-a3b",
     "qwen3-custom-30b-a3b",
@@ -4308,6 +4309,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "qwen3.5-397b-a17b",
           "maxtext-omni-gemma3-qwen3",
           "cosmos3-nano-reasoner",
+          "cosmos3-super-reasoner",
       )
       if self.model_name not in valid_mm_models and self.model_name != "default":
         raise ValueError(f"Multimodal is only supported for {valid_mm_models}, not {self.model_name}")

src/maxtext/layers/nnx_decoders.py

@@ -1480,6 +1480,7 @@ def _apply_embedding(
             "qwen3.5-397b-a17b",
             "maxtext-omni-gemma3-qwen3",
             "cosmos3-nano-reasoner",
+            "cosmos3-super-reasoner",
         }:
           y = mm_utils.merge_mm_embeddings(
               text_embeddings=y,
@@ -1499,6 +1500,7 @@ def _apply_embedding(
             "qwen3.5-35b-a3b",
             "qwen3.5-397b-a17b",
             "cosmos3-nano-reasoner",
+            "cosmos3-super-reasoner",
         }:
           y = mm_utils.merge_mm_embeddings(
               text_embeddings=y,

src/maxtext/multimodal/processor.py

@@ -41,6 +41,7 @@
     "maxtext-omni-gemma3-qwen3": ("gemma3", "qwen3"),
     # Cosmos
     "cosmos3-nano-reasoner": ("qwen3_vl", "qwen3"),
+    "cosmos3-super-reasoner": ("qwen3_vl", "qwen3"),
 }
 
 

src/maxtext/utils/globals.py

@@ -92,6 +92,7 @@
     "olmo3-7b-pt": "allenai/Olmo-3-1025-7B",
     "olmo3-32b": "allenai/Olmo-3-32B-Think",
     "cosmos3-nano-reasoner": "nvidia/Cosmos3-Nano",
+    "cosmos3-super-reasoner": "nvidia/Cosmos3-Super",
     # "default" is not HF model, but adding to to avoid confusing warning about tokenizer_path
     "default": os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers/tokenizer.llama2"),
 }