From 00aec75499aa573114d7e9a83c4deb7cd39a97fb Mon Sep 17 00:00:00 2001 From: Yin Date: Mon, 31 Aug 2026 18:25:48 +0000 Subject: [PATCH] Cosmos3-super-reasoner bring up. --- .../utils/hf_model_configs.py | 51 +++++++++++++++++++ .../utils/param_mapping.py | 2 + .../configs/models/cosmos3-super-reasoner.yml | 41 +++++++++++++++ src/maxtext/configs/types.py | 2 + src/maxtext/layers/nnx_decoders.py | 2 + src/maxtext/multimodal/processor.py | 1 + src/maxtext/utils/globals.py | 1 + 7 files changed, 100 insertions(+) create mode 100644 src/maxtext/configs/models/cosmos3-super-reasoner.yml diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index 5dc3315d77..a731f89bb3 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/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, diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 40e2508e58..8598f823ed 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/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, diff --git a/src/maxtext/configs/models/cosmos3-super-reasoner.yml b/src/maxtext/configs/models/cosmos3-super-reasoner.yml new file mode 100644 index 0000000000..59b003c51d --- /dev/null +++ b/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" diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 353656c665..3855f7a25f 100644 --- a/src/maxtext/configs/types.py +++ b/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}") diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 1e06020e72..c1d567eea3 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/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, diff --git a/src/maxtext/multimodal/processor.py b/src/maxtext/multimodal/processor.py index affb99ff0d..470de9ab5d 100644 --- a/src/maxtext/multimodal/processor.py +++ b/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"), } diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index 0c73e6f020..8782f4c3e0 100644 --- a/src/maxtext/utils/globals.py +++ b/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"), }