sarvam mla support

Author: mohit-sarvam

src/megatron/bridge/models/__init__.py

@@ -196,6 +196,7 @@
 from megatron.bridge.models.t5_provider import T5ModelProvider
 from megatron.bridge.models.sarvam import (
     SarvamMoEBridge,
+    SarvamMLABridge,
 )
 
 
@@ -349,4 +350,5 @@
     "NemotronNano12Bv2Provider",
     "NemotronNano12Bv2VLModelProvider",
     "SarvamMoEBridge",
+    "SarvamMLABridge",
 ]

src/megatron/bridge/models/sarvam/__init__.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 
 from megatron.bridge.models.sarvam.sarvam_moe_bridge import SarvamMoEBridge
+from megatron.bridge.models.sarvam.sarvam_mla_bridge import SarvamMLABridge
 
 
 __all__ = [
     "SarvamMoEBridge",
+    "SarvamMLABridge",
 ]

src/megatron/bridge/models/sarvam/common.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
+
+
+def get_common_config(hf_pretrained: PreTrainedCausalLM) -> dict:
+    """
+    Returns a dictionary of common configurations for the Sarvam family of models.
+    """
+    hf_config = hf_pretrained.config
+
+    config = {}
+
+    config["num_layers"] = hf_config.num_hidden_layers
+    config["hidden_size"] = hf_config.hidden_size
+    config["ffn_hidden_size"] = hf_config.intermediate_size
+    config["moe_ffn_hidden_size"] = hf_config.moe_intermediate_size
+    config["num_attention_heads"] = hf_config.num_attention_heads
+    config["num_moe_experts"] = hf_config.num_experts
+    config["moe_router_topk"] = hf_config.num_experts_per_tok
+    config["moe_shared_expert_intermediate_size"] = (
+        hf_config.num_shared_experts * hf_config.moe_intermediate_size
+    )
+    config["moe_layer_freq"] = [0] * hf_config.first_k_dense_replace + [1] * (
+        hf_config.num_hidden_layers - hf_config.first_k_dense_replace
+    )
+    config["vocab_size"] = hf_config.vocab_size
+    config["seq_length"] = hf_config.max_position_embeddings
+    config["generation_config"] = getattr(hf_pretrained, "generation_config", None)
+    config["rotary_base"] = hf_config.rope_theta
+
+    return config

src/megatron/bridge/models/sarvam/sarvam_mla_bridge.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
+from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
+from megatron.bridge.models.conversion.param_mapping import (
+    AutoMapping,
+    GatedMLPMapping,
+)
+from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
+from megatron.bridge.models.sarvam.common import get_common_config
+from megatron.bridge.models.sarvam.sarvam_provider import SarvamMLAModelProvider
+from megatron.core.models.gpt.gpt_model import GPTModel
+
+
+@MegatronModelBridge.register_bridge(source="SarvamMLAForCausalLM", target=GPTModel)
+class SarvamMLABridge(MegatronModelBridge):
+    """
+    Megatron Hub Bridge for Sarvam MLA Causal LM.
+
+    This bridge handles the conversion between HuggingFace SarvamMLAForCausalLM
+    and Megatron-Core GPTModel formats. Sarvam MLA models use multi-latent attention
+    architecture.
+    """
+
+    def provider_bridge(
+        self, hf_pretrained: PreTrainedCausalLM
+    ) -> SarvamMLAModelProvider:
+        hf_config = hf_pretrained.config
+        config = get_common_config(hf_pretrained)
+
+        config["fp16"] = (
+            self.dtype_from_hf(hf_config, default=torch.float32) == torch.float16
+        )
+        config["bf16"] = (
+            self.dtype_from_hf(hf_config, default=torch.float32) == torch.bfloat16
+        )
+        config["params_dtype"] = self.dtype_from_hf(hf_config, default=torch.float32)
+        config["kv_channels"] = hf_config.hidden_size // hf_config.num_attention_heads
+
+        # MLA
+        config["kv_lora_rank"] = hf_config.kv_lora_rank
+        config["qk_head_dim"] = hf_config.qk_nope_head_dim
+        config["qk_pos_emb_head_dim"] = hf_config.qk_rope_head_dim
+        config["v_head_dim"] = hf_config.v_head_dim
+        
+        if hasattr(hf_config, "rope_scaling") and hf_config.rope_scaling is not None:
+            config["rotary_scaling_factor"] = hf_config.rope_scaling["factor"]
+            config["mscale"] = hf_config.rope_scaling["mscale"]
+            config["mscale_all_dim"] = hf_config.rope_scaling["mscale_all_dim"]
+        else:
+            config["rotary_scaling_factor"] = 1.0
+            config["mscale"] = 1.0
+            config["mscale_all_dim"] = 1.0
+
+        return SarvamMLAModelProvider(**config)
+
+    def mapping_registry(self) -> MegatronMappingRegistry:
+
+        param_mappings = {
+            # Embed
+            "embedding.word_embeddings.weight": "model.embed_tokens.weight",
+
+            # Attention
+            "decoder.layers.*.input_layernorm.weight": "model.layers.*.input_layernorm.weight",
+            "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight",
+            #  In sarvam, HF weight `model.layers.*.post_attention_layernorm.weight` is mapped to the following mcore weights depending on the layer type:
+            #  (a) `decoder.layers.*.pre_mlp_layernorm.weight`, if the layer is MoE
+            #  (b) `decoder.layers.*.mlp.linear_fc1.layer_norm_weight`, if the layer is dense
+            "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.post_attention_layernorm.weight",
+            "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.layers.*.post_attention_layernorm.weight",
+            "decoder.layers.*.self_attention.linear_q_proj.weight": "model.layers.*.self_attn.q_proj.weight",
+            "decoder.layers.*.self_attention.linear_kv_down_proj.weight": "model.layers.*.self_attn.kv_a_proj_with_mqa.weight",
+            "decoder.layers.*.self_attention.linear_kv_up_proj.weight": "model.layers.*.self_attn.kv_b_proj.weight",
+            "decoder.layers.*.self_attention.linear_kv_up_proj.layer_norm_weight": "model.layers.*.self_attn.kv_a_layernorm.weight",
+            # Mcore local spec
+            "decoder.layers.*.self_attention.kv_layernorm.weight": "model.layers.*.self_attn.kv_a_layernorm.weight",
+
+            # Dense MLP
+            "decoder.layers.*.mlp.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight",
+
+            # Moe
+            "decoder.layers.*.mlp.experts.linear_fc2.weight*": "model.layers.*.mlp.experts.*.down_proj.weight",
+            "decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.layers.*.mlp.shared_experts.down_proj.weight",
+            "decoder.layers.*.mlp.router.expert_bias": "model.layers.*.mlp.gate.e_score_correction_bias",
+            "decoder.layers.*.mlp.router.weight": "model.layers.*.mlp.gate.weight",
+
+            # LM Head
+            "decoder.final_layernorm.weight": "model.norm.weight",
+            "output_layer.weight": "lm_head.weight",
+        }
+
+        mapping_list = []
+        for megatron_param, hf_param in param_mappings.items():
+            mapping_list.append(
+                AutoMapping(hf_param=hf_param, megatron_param=megatron_param)
+            )
+
+        mapping_list.extend(
+            [
+                GatedMLPMapping(
+                    megatron_param="decoder.layers.*.mlp.linear_fc1.weight",
+                    gate="model.layers.*.mlp.gate_proj.weight",
+                    up="model.layers.*.mlp.up_proj.weight",
+                ),
+                GatedMLPMapping(
+                    megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*",
+                    gate="model.layers.*.mlp.experts.*.gate_proj.weight",
+                    up="model.layers.*.mlp.experts.*.up_proj.weight",
+                ),
+                GatedMLPMapping(
+                    megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight",
+                    gate="model.layers.*.mlp.shared_experts.gate_proj.weight",
+                    up="model.layers.*.mlp.shared_experts.up_proj.weight",
+                ),
+            ]
+        )
+
+        return MegatronMappingRegistry(*mapping_list)

src/megatron/bridge/models/sarvam/sarvam_moe_bridge.py

@@ -13,17 +13,17 @@
 # limitations under the License.
 
 import torch
-from megatron.core.models.gpt.gpt_model import GPTModel
-
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
 from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
 from megatron.bridge.models.conversion.param_mapping import (
     AutoMapping,
-    GatedMLPMapping,
     ConcatenatedQKVMapping,
+    GatedMLPMapping,
 )
 from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
+from megatron.bridge.models.sarvam.common import get_common_config
 from megatron.bridge.models.sarvam.sarvam_provider import SarvamMoEModelProvider
+from megatron.core.models.gpt.gpt_model import GPTModel
 
 
 @MegatronModelBridge.register_bridge(source="SarvamMoEForCausalLM", target=GPTModel)
@@ -36,59 +36,51 @@ class SarvamMoEBridge(MegatronModelBridge):
     architecture with QKV layernorm.
     """
 
-    def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> SarvamMoEModelProvider:
+    def provider_bridge(
+        self, hf_pretrained: PreTrainedCausalLM
+    ) -> SarvamMoEModelProvider:
         hf_config = hf_pretrained.config
+        config = get_common_config(hf_pretrained)
 
-        provider = SarvamMoEModelProvider(
-            num_layers=hf_config.num_hidden_layers,
-            hidden_size=hf_config.hidden_size,
-            ffn_hidden_size=hf_config.intermediate_size,
-            moe_ffn_hidden_size=hf_config.moe_intermediate_size,  # Maps to moe_intermediate_size in HF
-            num_attention_heads=hf_config.num_attention_heads,
-            kv_channels=hf_config.head_dim,
-            num_query_groups=hf_config.num_key_value_heads,
-            num_moe_experts=hf_config.num_experts,
-            moe_router_topk=hf_config.num_experts_per_tok,  # Maps to num_experts_per_tok in HF
-            moe_shared_expert_intermediate_size=hf_config.num_shared_experts * hf_config.moe_intermediate_size,
-            moe_router_enable_expert_bias=hf_config.moe_router_enable_expert_bias,
-            moe_layer_freq=[0] * hf_config.first_k_dense_replace + [1] * (hf_config.num_hidden_layers - hf_config.first_k_dense_replace),
-            vocab_size=hf_config.vocab_size,
-            seq_length=hf_config.max_position_embeddings,
-            generation_config=hf_pretrained.generation_config,
-            rotary_base=hf_config.rope_theta,
-            fp16=(self.dtype_from_hf(hf_config, default=torch.float32) == torch.float16),
-            bf16=(self.dtype_from_hf(hf_config, default=torch.float32) == torch.bfloat16),
-            params_dtype=self.dtype_from_hf(hf_config, default=torch.float32),
+        config["fp16"] = (
+            self.dtype_from_hf(hf_config, default=torch.float32) == torch.float16
+        )
+        config["bf16"] = (
+            self.dtype_from_hf(hf_config, default=torch.float32) == torch.bfloat16
         )
+        config["params_dtype"] = self.dtype_from_hf(hf_config, default=torch.float32)
 
-        return provider
+        # GQA
+        config["num_query_groups"] = hf_config.num_key_value_heads
+        config["kv_channels"] = hf_config.head_dim
+
+        return SarvamMoEModelProvider(**config)
 
     def mapping_registry(self) -> MegatronMappingRegistry:
-       # Return MegatronMappingRegistry containing parameter mappings from Megatron to HF format
-        # First create simple 1:1 parameter mappings using a dictionary for readability
 
-        # Dictionary maps Megatron parameter names -> HF parameter names
         param_mappings = {
             # Embed
             "embedding.word_embeddings.weight": "model.word_embeddings.weight",
 
             # Attention
             "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight",
-            "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.layers.*.post_attention_layernorm.weight",
+            #  In sarvam, HF weight `model.layers.*.post_attention_layernorm.weight` is mapped to the following mcore weights depending on the layer type:
+            #  (a) `decoder.layers.*.pre_mlp_layernorm.weight`, if the layer is MoE
+            #  (b) `decoder.layers.*.mlp.linear_fc1.layer_norm_weight`, if the layer is dense
             "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.post_attention_layernorm.weight",
-            "decoder.layers.*.mlp.router.expert_bias": "model.layers.*.mlp.gate.expert_bias",
-            "decoder.layers.*.mlp.router.weight": "model.layers.*.mlp.gate.weight",
-
+            "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.layers.*.post_attention_layernorm.weight",
             "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.attention.query_layernorm.weight",
             "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.attention.key_layernorm.weight",
             "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.attention.dense.weight",
-            
+
             # Dense MLP
             "decoder.layers.*.mlp.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight",
+
+            # MoE
+            "decoder.layers.*.mlp.router.expert_bias": "model.layers.*.mlp.gate.expert_bias",
+            "decoder.layers.*.mlp.router.weight": "model.layers.*.mlp.gate.weight",
             "decoder.layers.*.mlp.experts.linear_fc2.weight*": "model.layers.*.mlp.experts.*.down_proj.weight",
             "decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.layers.*.mlp.shared_experts.down_proj.weight",
-            
-            "final_layernorm.weight": "final_layernorm.weight",
 
             # LM Head
             "decoder.final_layernorm.weight": "model.norm.weight",
@@ -97,7 +89,9 @@ def mapping_registry(self) -> MegatronMappingRegistry:
 
         mapping_list = []
         for megatron_param, hf_param in param_mappings.items():
-            mapping_list.append(AutoMapping(hf_param=hf_param, megatron_param=megatron_param))
+            mapping_list.append(
+                AutoMapping(hf_param=hf_param, megatron_param=megatron_param)
+            )
 
         mapping_list.extend(
             [
@@ -120,7 +114,6 @@ def mapping_registry(self) -> MegatronMappingRegistry:
                     gate="model.layers.*.mlp.shared_experts.gate_proj.weight",
                     up="model.layers.*.mlp.shared_experts.up_proj.weight",
                 ),
-                
             ]
         )