From 29686707abaea22d7ff0366e2b02f235b4b41f84 Mon Sep 17 00:00:00 2001 From: Aashay Sarvam Date: Tue, 10 Mar 2026 11:54:30 +0000 Subject: [PATCH 1/8] Add SarvamMLA model (sarvamai/sarvam-105b) Add native support for the sarvam_mla model type using the modular pattern, inheriting from DeepSeek V3. The model uses Multi-head Latent Attention (MLA) with Mixture of Experts (MoE), supporting 105B parameters with 128 routed experts and 8 active per token. New files: - configuration_sarvam_mla.py: Config with attribute mapping, rope normalization, and head_dim handling for Hub compatibility - modular_sarvam_mla.py: 48-line modular file inheriting DeepSeek V3 - modeling_sarvam_mla.py: Auto-generated from modular (736 lines) - test_modeling_sarvam_mla.py: 140 passing unit tests - sarvam_mla.md: Documentation with usage examples Modified files: - Auto-registration in configuration_auto.py, modeling_auto.py - Model import in models/__init__.py - Weight conversion mapping (qwen2_moe pattern) in conversion_mapping.py - Documentation index in _toctree.yml Made-with: Cursor --- docs/source/en/_toctree.yml | 2 + docs/source/en/model_doc/sarvam_mla.md | 79 ++ src/transformers/conversion_mapping.py | 1 + src/transformers/models/__init__.py | 1 + .../models/auto/configuration_auto.py | 2 + src/transformers/models/auto/modeling_auto.py | 4 + .../models/sarvam_mla/__init__.py | 27 + .../sarvam_mla/configuration_sarvam_mla.py | 147 ++++ .../models/sarvam_mla/modeling_sarvam_mla.py | 736 ++++++++++++++++++ .../models/sarvam_mla/modular_sarvam_mla.py | 90 +++ tests/models/sarvam_mla/__init__.py | 0 .../sarvam_mla/test_modeling_sarvam_mla.py | 298 +++++++ 12 files changed, 1387 insertions(+) create mode 100644 docs/source/en/model_doc/sarvam_mla.md create mode 100644 src/transformers/models/sarvam_mla/__init__.py create mode 100644 src/transformers/models/sarvam_mla/configuration_sarvam_mla.py create mode 100644 src/transformers/models/sarvam_mla/modeling_sarvam_mla.py create mode 100644 src/transformers/models/sarvam_mla/modular_sarvam_mla.py create mode 100644 tests/models/sarvam_mla/__init__.py create mode 100644 tests/models/sarvam_mla/test_modeling_sarvam_mla.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 07aad5be5b57..d4327e2e0db0 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -946,6 +946,8 @@ title: Segment Anything - local: model_doc/sam_hq title: Segment Anything High Quality + - local: model_doc/sarvam_mla + title: SarvamMLA - local: model_doc/superglue title: SuperGlue - local: model_doc/superpoint diff --git a/docs/source/en/model_doc/sarvam_mla.md b/docs/source/en/model_doc/sarvam_mla.md new file mode 100644 index 000000000000..d5076c307b1f --- /dev/null +++ b/docs/source/en/model_doc/sarvam_mla.md @@ -0,0 +1,79 @@ + + +# SarvamMLA + +## Overview + +SarvamMLA is a 105B parameter Mixture of Experts (MoE) language model developed by [Sarvam AI](https://www.sarvam.ai/). It uses Multi-head Latent Attention (MLA) combined with sparse MoE routing, architecturally similar to DeepSeek-V3. + +Key architectural features: + +- **Multi-head Latent Attention (MLA)**: Low-rank KV compression with decoupled RoPE, reducing KV cache memory while maintaining performance. +- **Sparse Mixture of Experts**: 128 routed experts with 8 active per token, plus 1 shared expert. The first layer uses a dense MLP. +- **DeepSeek YaRN RoPE**: Extended context support up to 131K tokens via YaRN rotary position embeddings. +- **Sigmoid routing with group-based top-k**: Token-choice routing using sigmoid scores with expert bias correction and group-aware selection. + +## Usage + +```python +from transformers import AutoModelForCausalLM, AutoTokenizer + +model = AutoModelForCausalLM.from_pretrained( + "sarvamai/sarvam-105b", + device_map="auto", + torch_dtype="auto", +) +tokenizer = AutoTokenizer.from_pretrained("sarvamai/sarvam-105b") + +inputs = tokenizer("Hello, how are you?", return_tensors="pt").to(model.device) +outputs = model.generate(**inputs, max_new_tokens=50) +print(tokenizer.decode(outputs[0], skip_special_tokens=True)) +``` + +For running on limited GPU memory, use quantization: + +```python +from transformers import AutoModelForCausalLM, BitsAndBytesConfig + +quantization_config = BitsAndBytesConfig(load_in_4bit=True) +model = AutoModelForCausalLM.from_pretrained( + "sarvamai/sarvam-105b", + quantization_config=quantization_config, + device_map="auto", +) +``` + +## SarvamMLAConfig + +[[autodoc]] SarvamMLAConfig + +## SarvamMLAModel + +[[autodoc]] SarvamMLAModel + - forward + +## SarvamMLAForCausalLM + +[[autodoc]] SarvamMLAForCausalLM + - forward + +## SarvamMLAForSequenceClassification + +[[autodoc]] SarvamMLAForSequenceClassification + - forward + +## SarvamMLAForTokenClassification + +[[autodoc]] SarvamMLAForTokenClassification + - forward diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 983472525964..64fdb37721f9 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -59,6 +59,7 @@ "hunyuan_v1_moe": "qwen2_moe", "flex_olmo": "qwen2_moe", "olmoe": "qwen2_moe", + "sarvam_mla": "qwen2_moe", "exaone_moe": "qwen2_moe", "rt_detr_v2": "rt_detr", "pp_doclayout_v2": "rt_detr", diff --git a/src/transformers/models/__init__.py b/src/transformers/models/__init__.py index 5f45081ac4a0..fbf5c229da9b 100644 --- a/src/transformers/models/__init__.py +++ b/src/transformers/models/__init__.py @@ -365,6 +365,7 @@ from .sam3_tracker_video import * from .sam3_video import * from .sam_hq import * + from .sarvam_mla import * from .seamless_m4t import * from .seamless_m4t_v2 import * from .seed_oss import * diff --git a/src/transformers/models/auto/configuration_auto.py b/src/transformers/models/auto/configuration_auto.py index 476b5362343f..f891aa50937f 100644 --- a/src/transformers/models/auto/configuration_auto.py +++ b/src/transformers/models/auto/configuration_auto.py @@ -417,6 +417,7 @@ ("sam_hq", "SamHQConfig"), ("sam_hq_vision_model", "SamHQVisionConfig"), ("sam_vision_model", "SamVisionConfig"), + ("sarvam_mla", "SarvamMLAConfig"), ("seamless_m4t", "SeamlessM4TConfig"), ("seamless_m4t_v2", "SeamlessM4Tv2Config"), ("seed_oss", "SeedOssConfig"), @@ -934,6 +935,7 @@ ("sam_hq", "SAM-HQ"), ("sam_hq_vision_model", "SamHQVisionModel"), ("sam_vision_model", "SamVisionModel"), + ("sarvam_mla", "SarvamMLA"), ("seamless_m4t", "SeamlessM4T"), ("seamless_m4t_v2", "SeamlessM4Tv2"), ("seed_oss", "SeedOss"), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index 764d3b770e86..ac3d13366af6 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -397,6 +397,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("sam_hq", "SamHQModel"), ("sam_hq_vision_model", "SamHQVisionModel"), ("sam_vision_model", "SamVisionModel"), + ("sarvam_mla", "SarvamMLAModel"), ("seamless_m4t", "SeamlessM4TModel"), ("seamless_m4t_v2", "SeamlessM4Tv2Model"), ("seed_oss", "SeedOssModel"), @@ -722,6 +723,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("roc_bert", "RoCBertForCausalLM"), ("roformer", "RoFormerForCausalLM"), ("rwkv", "RwkvForCausalLM"), + ("sarvam_mla", "SarvamMLAForCausalLM"), ("seed_oss", "SeedOssForCausalLM"), ("smollm3", "SmolLM3ForCausalLM"), ("solar_open", "SolarOpenForCausalLM"), @@ -1280,6 +1282,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("roberta-prelayernorm", "RobertaPreLayerNormForSequenceClassification"), ("roc_bert", "RoCBertForSequenceClassification"), ("roformer", "RoFormerForSequenceClassification"), + ("sarvam_mla", "SarvamMLAForSequenceClassification"), ("seed_oss", "SeedOssForSequenceClassification"), ("smollm3", "SmolLM3ForSequenceClassification"), ("squeezebert", "SqueezeBertForSequenceClassification"), @@ -1484,6 +1487,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("roberta-prelayernorm", "RobertaPreLayerNormForTokenClassification"), ("roc_bert", "RoCBertForTokenClassification"), ("roformer", "RoFormerForTokenClassification"), + ("sarvam_mla", "SarvamMLAForTokenClassification"), ("seed_oss", "SeedOssForTokenClassification"), ("smollm3", "SmolLM3ForTokenClassification"), ("squeezebert", "SqueezeBertForTokenClassification"), diff --git a/src/transformers/models/sarvam_mla/__init__.py b/src/transformers/models/sarvam_mla/__init__.py new file mode 100644 index 000000000000..931ca79e5fa9 --- /dev/null +++ b/src/transformers/models/sarvam_mla/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2026 Sarvam AI and the HuggingFace Inc. team. 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 typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_sarvam_mla import * + from .modeling_sarvam_mla import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py b/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py new file mode 100644 index 000000000000..cce687afddab --- /dev/null +++ b/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py @@ -0,0 +1,147 @@ +# Copyright 2026 Sarvam AI and the HuggingFace Inc. team. 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. +"""SarvamMLA model configuration""" + +from ...utils import auto_docstring +from ..deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config + + +@auto_docstring(checkpoint="sarvamai/sarvam-105b") +class SarvamMLAConfig(DeepseekV3Config): + r""" + n_group (`int`, *optional*, defaults to 16): + Number of groups for routed experts. + rope_interleave (`bool`, *optional*, defaults to `True`): + Whether to interleave the rotary position embeddings. + first_k_dense_replace (`int`, *optional*, defaults to 1): + Number of dense layers in shallow layers(embed->dense->moe->moe...->lm_head). + \--k dense layers--/ + + Example: + + ```python + >>> from transformers import SarvamMLAModel, SarvamMLAConfig + + >>> # Initializing a SarvamMLA style configuration + >>> configuration = SarvamMLAConfig() + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "sarvam_mla" + attribute_map = { + "n_routed_experts": "num_experts", + "n_shared_experts": "num_shared_experts", + "num_local_experts": "num_experts", + } + + def __init__( + self, + vocab_size=262144, + hidden_size=4096, + intermediate_size=16384, + moe_intermediate_size=2048, + num_hidden_layers=32, + num_attention_heads=64, + num_key_value_heads=None, + num_shared_experts=1, + num_experts=128, + routed_scaling_factor=2.5, + kv_lora_rank=512, + q_lora_rank=None, + qk_rope_head_dim=64, + v_head_dim=128, + qk_nope_head_dim=128, + n_group=16, + topk_group=2, + num_experts_per_tok=8, + first_k_dense_replace=1, + norm_topk_prob=True, + hidden_act="silu", + max_position_embeddings=4096, + initializer_range=0.006, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + bos_token_id=None, + eos_token_id=1, + tie_word_embeddings=False, + rope_parameters=None, + rope_interleave=True, + attention_bias=False, + attention_dropout=0.0, + **kwargs, + ): + self.num_experts = num_experts + self.num_shared_experts = num_shared_experts + + # head_dim in the Hub config.json is set to kv_lora_rank + qk_rope_head_dim + # for vLLM MLA compatibility, but internally the model uses qk_rope_head_dim + # for rotary embeddings. Remove it from kwargs to prevent overriding. + kwargs.pop("head_dim", None) + kwargs.pop("q_head_dim", None) + + # The Hub config uses "deepseek_yarn" as rope type; normalize to "yarn" + # which is the standard type in ROPE_INIT_FUNCTIONS. + if rope_parameters is not None and rope_parameters.get("type") == "deepseek_yarn": + rope_parameters = dict(rope_parameters) + rope_parameters["type"] = "yarn" + rope_scaling = kwargs.pop("rope_scaling", None) + if rope_scaling is not None: + if rope_scaling.get("type") == "deepseek_yarn": + rope_scaling = dict(rope_scaling) + rope_scaling["type"] = "yarn" + if rope_parameters is None: + rope_parameters = rope_scaling + + super().__init__( + vocab_size=vocab_size, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + moe_intermediate_size=moe_intermediate_size, + num_hidden_layers=num_hidden_layers, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + n_shared_experts=num_shared_experts, + n_routed_experts=num_experts, + routed_scaling_factor=routed_scaling_factor, + kv_lora_rank=kv_lora_rank, + q_lora_rank=q_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + qk_nope_head_dim=qk_nope_head_dim, + n_group=n_group, + topk_group=topk_group, + num_experts_per_tok=num_experts_per_tok, + first_k_dense_replace=first_k_dense_replace, + norm_topk_prob=norm_topk_prob, + hidden_act=hidden_act, + max_position_embeddings=max_position_embeddings, + initializer_range=initializer_range, + rms_norm_eps=rms_norm_eps, + use_cache=use_cache, + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + rope_parameters=rope_parameters, + rope_interleave=rope_interleave, + attention_bias=attention_bias, + attention_dropout=attention_dropout, + **kwargs, + ) + + +__all__ = ["SarvamMLAConfig"] diff --git a/src/transformers/models/sarvam_mla/modeling_sarvam_mla.py b/src/transformers/models/sarvam_mla/modeling_sarvam_mla.py new file mode 100644 index 000000000000..5cca52de5272 --- /dev/null +++ b/src/transformers/models/sarvam_mla/modeling_sarvam_mla.py @@ -0,0 +1,736 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/sarvam_mla/modular_sarvam_mla.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_sarvam_mla.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 Sarvam AI and the HuggingFace Inc. team. 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 math +from collections.abc import Callable +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationMixin +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...masking_utils import create_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import ( + GenericForSequenceClassification, + GenericForTokenClassification, + GradientCheckpointingLayer, +) +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_sarvam_mla import SarvamMLAConfig + + +@use_kernel_forward_from_hub("RMSNorm") +class SarvamMLARMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + SarvamMLARMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class SarvamMLARotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: SarvamMLAConfig, device=None): + super().__init__() + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + + self.rope_type = self.config.rope_parameters["rope_type"] + rope_init_fn: Callable = self.compute_default_rope_parameters + if self.rope_type != "default": + rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + inv_freq, self.attention_scaling = rope_init_fn(self.config, device) + + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + + @staticmethod + def compute_default_rope_parameters( + config: SarvamMLAConfig | None = None, + device: Optional["torch.device"] = None, + seq_len: int | None = None, + ) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PreTrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = config.rope_parameters["rope_theta"] + dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with maybe_autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class SarvamMLAMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class SarvamMLATopkRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.n_routed_experts = config.n_routed_experts + + self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) + self.register_buffer("e_score_correction_bias", torch.zeros(self.n_routed_experts)) + + def forward(self, hidden_states): + hidden_states = hidden_states.view(-1, self.config.hidden_size) + router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) + return router_logits + + +@use_experts_implementation +class SarvamMLANaiveMoe(nn.Module): + """Collection of expert weights stored as 3D tensors.""" + + def __init__(self, config): + super().__init__() + self.num_experts = config.num_local_experts + self.hidden_dim = config.hidden_size + self.intermediate_dim = config.moe_intermediate_size + self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) + self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) + self.act_fn = ACT2FN[config.hidden_act] + + def forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ) -> torch.Tensor: + final_hidden_states = torch.zeros_like(hidden_states) + with torch.no_grad(): + expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + + for expert_idx in expert_hit: + expert_idx = expert_idx[0] + if expert_idx == self.num_experts: + continue + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + current_state = hidden_states[token_idx] + gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1) + current_hidden_states = self.act_fn(gate) * up + current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx]) + current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) + + return final_hidden_states + + +class SarvamMLAMoE(nn.Module): + """ + A mixed expert module containing shared experts. + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.experts = SarvamMLANaiveMoe(config) + self.gate = SarvamMLATopkRouter(config) + self.shared_experts = SarvamMLAMLP( + config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts + ) + self.n_routed_experts = config.n_routed_experts + self.n_group = config.n_group + self.topk_group = config.topk_group + self.norm_topk_prob = config.norm_topk_prob + self.routed_scaling_factor = config.routed_scaling_factor + self.top_k = config.num_experts_per_tok + + def route_tokens_to_experts(self, router_logits): + router_logits = router_logits.sigmoid() + router_logits_for_choice = router_logits + self.gate.e_score_correction_bias + group_scores = ( + router_logits_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) + group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(-1, self.n_group, self.n_routed_experts // self.n_group) + .reshape(-1, self.n_routed_experts) + ) + scores_for_choice = router_logits_for_choice.masked_fill(~score_mask.bool(), 0.0) + topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] + topk_weights = router_logits.gather(1, topk_indices) + if self.norm_topk_prob: + denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 + topk_weights /= denominator + topk_weights = topk_weights * self.routed_scaling_factor + return topk_indices, topk_weights + + def forward(self, hidden_states): + residuals = hidden_states + orig_shape = hidden_states.shape + router_logits = self.gate(hidden_states) + topk_indices, topk_weights = self.route_tokens_to_experts(router_logits) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) + hidden_states = hidden_states + self.shared_experts(residuals) + return hidden_states + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +@use_kernel_func_from_hub("rotary_pos_emb") +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + r""" + TODO let's just use the original freqcis computation to not have the view + transpose + reshape! This is not optimized! + Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + b, h, s, d = q.shape + q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + b, h, s, d = k.shape + k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def yarn_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class SarvamMLAAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: SarvamMLAConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.attention_dropout = config.attention_dropout + self.num_heads = config.num_attention_heads + + self.q_lora_rank = config.q_lora_rank + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_head_dim = config.qk_head_dim + + self.is_causal = True + if self.q_lora_rank is None: + self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False) + else: + self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=config.attention_bias) + self.q_a_layernorm = SarvamMLARMSNorm(config.q_lora_rank) + self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) + + self.kv_a_proj_with_mqa = nn.Linear( + config.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = SarvamMLARMSNorm(self.kv_lora_rank) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + ) + + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + + self.scaling = self.qk_head_dim ** (-0.5) + if self.config.rope_parameters.get("rope_type", "default") != "default": + mscale_all_dim = self.config.rope_parameters.get("mscale_all_dim", 0) + scaling_factor = self.config.rope_parameters["factor"] + if mscale_all_dim: + mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) + self.scaling = self.scaling * mscale * mscale + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + batch_size, seq_length = hidden_states.shape[:-1] + query_shape = (batch_size, seq_length, -1, self.qk_head_dim) + key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim) + + if self.q_lora_rank is None: + q_states = self.q_proj(hidden_states) + else: + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q_states = q_states.view(query_shape).transpose(1, 2) + q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + + k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2) + k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) + + cos, sin = position_embeddings + if self.config.rope_interleave: # support using interleaved weights for efficiency + q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) + else: + q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) + k_rot = k_rot.expand(*k_pass.shape[:-1], -1) + + query_states = torch.cat((q_pass, q_rot), dim=-1) + key_states = torch.cat((k_pass, k_rot), dim=-1) + + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class SarvamMLADecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: SarvamMLAConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = SarvamMLAAttention(config=config, layer_idx=layer_idx) + + if layer_idx >= config.first_k_dense_replace: + self.mlp = SarvamMLAMoE(config) + else: + self.mlp = SarvamMLAMLP(config) + + self.input_layernorm = SarvamMLARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = SarvamMLARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + # Self Attention + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +@auto_docstring +class SarvamMLAPreTrainedModel(PreTrainedModel): + config: SarvamMLAConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["SarvamMLADecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + + _can_compile_fullgraph = True + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": SarvamMLADecoderLayer, + "attentions": SarvamMLAAttention, + } + _keep_in_fp32_modules_strict = ["e_score_correction_bias"] + _keys_to_ignore_on_load_unexpected = [r"model\.layers\.61.*"] + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, SarvamMLATopkRouter): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + init.zeros_(module.e_score_correction_bias) + elif isinstance(module, SarvamMLANaiveMoe): + init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range) + init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) + + +@auto_docstring +class SarvamMLAModel(SarvamMLAPreTrainedModel): + def __init__(self, config: SarvamMLAConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [SarvamMLADecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = SarvamMLARMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = SarvamMLARotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + if position_ids is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens + position_ids = position_ids.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class SarvamMLAForCausalLM(SarvamMLAPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + _tp_plan = {"lm_head": "colwise_gather_output"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config): + super().__init__(config) + self.model = SarvamMLAModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutputWithPast: + r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, SarvamMLAForCausalLM + + >>> model = SarvamMLAForCausalLM.from_pretrained("meta-sarvam_mla/SarvamMLA-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-sarvam_mla/SarvamMLA-2-7b-hf") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class SarvamMLAForSequenceClassification(GenericForSequenceClassification, SarvamMLAPreTrainedModel): + pass + + +class SarvamMLAForTokenClassification(GenericForTokenClassification, SarvamMLAPreTrainedModel): + pass + + +__all__ = [ + "SarvamMLAPreTrainedModel", + "SarvamMLAModel", + "SarvamMLAForCausalLM", + "SarvamMLAForSequenceClassification", + "SarvamMLAForTokenClassification", +] diff --git a/src/transformers/models/sarvam_mla/modular_sarvam_mla.py b/src/transformers/models/sarvam_mla/modular_sarvam_mla.py new file mode 100644 index 000000000000..6963a86cd7b2 --- /dev/null +++ b/src/transformers/models/sarvam_mla/modular_sarvam_mla.py @@ -0,0 +1,90 @@ +# Copyright 2026 Sarvam AI and the HuggingFace Inc. team. 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 ...modeling_layers import GenericForSequenceClassification, GenericForTokenClassification +from ..deepseek_v3.modeling_deepseek_v3 import ( + DeepseekV3Attention, + DeepseekV3DecoderLayer, + DeepseekV3ForCausalLM, + DeepseekV3MLP, + DeepseekV3MoE, + DeepseekV3Model, + DeepseekV3NaiveMoe, + DeepseekV3PreTrainedModel, + DeepseekV3RMSNorm, + DeepseekV3RotaryEmbedding, + DeepseekV3TopkRouter, +) +from .configuration_sarvam_mla import SarvamMLAConfig + + +class SarvamMLARMSNorm(DeepseekV3RMSNorm): + pass + + +class SarvamMLARotaryEmbedding(DeepseekV3RotaryEmbedding): + pass + + +class SarvamMLAMLP(DeepseekV3MLP): + pass + + +class SarvamMLATopkRouter(DeepseekV3TopkRouter): + pass + + +class SarvamMLANaiveMoe(DeepseekV3NaiveMoe): + pass + + +class SarvamMLAMoE(DeepseekV3MoE): + pass + + +class SarvamMLAAttention(DeepseekV3Attention): + pass + + +class SarvamMLADecoderLayer(DeepseekV3DecoderLayer): + pass + + +class SarvamMLAPreTrainedModel(DeepseekV3PreTrainedModel): + pass + + +class SarvamMLAModel(DeepseekV3Model): + pass + + +class SarvamMLAForCausalLM(DeepseekV3ForCausalLM): + pass + + +class SarvamMLAForSequenceClassification(GenericForSequenceClassification, SarvamMLAPreTrainedModel): + pass + + +class SarvamMLAForTokenClassification(GenericForTokenClassification, SarvamMLAPreTrainedModel): + pass + + +__all__ = [ + "SarvamMLAPreTrainedModel", + "SarvamMLAModel", + "SarvamMLAForCausalLM", + "SarvamMLAForSequenceClassification", + "SarvamMLAForTokenClassification", +] diff --git a/tests/models/sarvam_mla/__init__.py b/tests/models/sarvam_mla/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/models/sarvam_mla/test_modeling_sarvam_mla.py b/tests/models/sarvam_mla/test_modeling_sarvam_mla.py new file mode 100644 index 000000000000..fc4b7a56d673 --- /dev/null +++ b/tests/models/sarvam_mla/test_modeling_sarvam_mla.py @@ -0,0 +1,298 @@ +# Copyright 2026 Sarvam AI and the HuggingFace Inc. team. 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. +"""Testing suite for the PyTorch SarvamMLA model.""" + +import unittest + +from parameterized import parameterized + +from transformers import SarvamMLAConfig, is_torch_available +from transformers.testing_utils import ( + require_torch, + torch_device, +) + +from ...generation.test_utils import GenerationTesterMixin +from ...test_configuration_common import ConfigTester +from ...test_modeling_common import ModelTesterMixin, ids_tensor +from ...test_pipeline_mixin import PipelineTesterMixin + + +if is_torch_available(): + import torch + + from transformers import ( + Cache, + SarvamMLAForCausalLM, + SarvamMLAForSequenceClassification, + SarvamMLAForTokenClassification, + SarvamMLAModel, + ) + + +class SarvamMLAModelTester: + if is_torch_available(): + causal_lm_class = SarvamMLAForCausalLM + + def __init__( + self, + parent, + batch_size=13, + seq_length=7, + is_training=True, + use_input_mask=True, + use_token_type_ids=False, + use_labels=True, + vocab_size=99, + hidden_size=32, + intermediate_size=32, + moe_intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + num_shared_experts=1, + num_experts=8, + routed_scaling_factor=2.5, + kv_lora_rank=16, + q_lora_rank=None, + qk_rope_head_dim=16, + v_head_dim=32, + qk_nope_head_dim=32, + n_group=2, + topk_group=1, + num_experts_per_tok=8, + first_k_dense_replace=1, + norm_topk_prob=True, + hidden_act="silu", + max_position_embeddings=512, + initializer_range=0.02, + attention_probs_dropout_prob=0.0, + type_vocab_size=16, + type_sequence_label_size=2, + num_labels=3, + num_choices=4, + pad_token_id=0, + scope=None, + ): + self.parent = parent + self.batch_size = batch_size + self.seq_length = seq_length + self.is_training = is_training + self.use_input_mask = use_input_mask + self.use_token_type_ids = use_token_type_ids + self.use_labels = use_labels + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.num_shared_experts = num_shared_experts + self.num_experts = num_experts + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.attention_probs_dropout_prob = attention_probs_dropout_prob + self.type_vocab_size = type_vocab_size + self.type_sequence_label_size = type_sequence_label_size + self.num_labels = num_labels + self.num_choices = num_choices + self.pad_token_id = pad_token_id + self.scope = scope + + def prepare_config_and_inputs(self): + input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) + + input_mask = None + if self.use_input_mask: + input_mask = torch.tril(torch.ones_like(input_ids).to(torch_device)) + + token_type_ids = None + if self.use_token_type_ids: + token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size) + + sequence_labels = None + token_labels = None + choice_labels = None + if self.use_labels: + sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size) + token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels) + choice_labels = ids_tensor([self.batch_size], self.num_choices) + + config = self.get_config() + + return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels + + def get_config(self): + return SarvamMLAConfig( + vocab_size=self.vocab_size, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + moe_intermediate_size=self.moe_intermediate_size, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + num_shared_experts=self.num_shared_experts, + num_experts=self.num_experts, + routed_scaling_factor=self.routed_scaling_factor, + kv_lora_rank=self.kv_lora_rank, + q_lora_rank=self.q_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + v_head_dim=self.v_head_dim, + qk_nope_head_dim=self.qk_nope_head_dim, + n_group=self.n_group, + topk_group=self.topk_group, + num_experts_per_tok=self.num_experts_per_tok, + first_k_dense_replace=self.first_k_dense_replace, + norm_topk_prob=self.norm_topk_prob, + hidden_act=self.hidden_act, + max_position_embeddings=self.max_position_embeddings, + initializer_range=self.initializer_range, + use_cache=True, + pad_token_id=self.pad_token_id, + attention_dropout=self.attention_probs_dropout_prob, + ) + + def create_and_check_model( + self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels + ): + model = SarvamMLAModel(config=config) + model.to(torch_device) + model.eval() + result = model(input_ids, attention_mask=input_mask) + result = model(input_ids) + self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size)) + + def prepare_config_and_inputs_for_common(self): + config_and_inputs = self.prepare_config_and_inputs() + ( + config, + input_ids, + token_type_ids, + input_mask, + sequence_labels, + token_labels, + choice_labels, + ) = config_and_inputs + inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask} + return config, inputs_dict + + +@require_torch +class SarvamMLAModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase): + all_model_classes = ( + ( + SarvamMLAModel, + SarvamMLAForCausalLM, + SarvamMLAForSequenceClassification, + SarvamMLAForTokenClassification, + ) + if is_torch_available() + else () + ) + all_generative_model_classes = (SarvamMLAForCausalLM,) if is_torch_available() else () + pipeline_model_mapping = ( + { + "feature-extraction": SarvamMLAModel, + "text-classification": SarvamMLAForSequenceClassification, + "token-classification": SarvamMLAForTokenClassification, + "text-generation": SarvamMLAForCausalLM, + "zero-shot": SarvamMLAForSequenceClassification, + } + if is_torch_available() + else {} + ) + + model_split_percents = [0.5, 0.7, 0.8] + + _torch_compile_train_cls = SarvamMLAForCausalLM if is_torch_available() else None + + def setUp(self): + self.model_tester = SarvamMLAModelTester(self) + self.config_tester = ConfigTester(self, config_class=SarvamMLAConfig, hidden_size=37) + + def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_length, config): + self.assertIsInstance(past_key_values, Cache) + + expected_common_shape = ( + batch_size, + getattr(config, "num_key_value_heads", config.num_attention_heads), + seq_length, + ) + expected_key_shape = expected_common_shape + (config.qk_nope_head_dim + config.qk_rope_head_dim,) + expected_value_shape = expected_common_shape + (config.v_head_dim,) + + for layer in past_key_values.layers: + self.assertEqual(layer.keys.shape, expected_key_shape) + self.assertEqual(layer.values.shape, expected_value_shape) + + @parameterized.expand([("random",), ("same",)]) + @unittest.skip("SarvamMLA uses MLA, not compatible with assisted decoding") + def test_assisted_decoding_matches_greedy_search(self, assistant_type): + pass + + @unittest.skip("SarvamMLA uses MLA, not compatible with assisted decoding") + def test_prompt_lookup_decoding_matches_greedy_search(self, assistant_type): + pass + + @unittest.skip("SarvamMLA uses MLA, not compatible with assisted decoding") + def test_assisted_decoding_sample(self): + pass + + @unittest.skip("SarvamMLA uses MLA, not compatible with standard cache format") + def test_beam_search_generate_dict_outputs_use_cache(self): + pass + + @unittest.skip("SarvamMLA uses MLA, not compatible with standard cache format") + def test_greedy_generate_dict_outputs_use_cache(self): + pass + + @unittest.skip(reason="SDPA can't dispatch on flash due to unsupported head dims") + def test_sdpa_can_dispatch_on_flash(self): + pass + + @unittest.skip("SarvamMLA has MoE, output can be different") + def test_model_outputs_equivalence(self, **kwargs): + pass + + def test_config(self): + self.config_tester.run_common_tests() + + def test_model(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_model(*config_and_inputs) + + def test_sarvam_mla_sequence_classification_model(self): + config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() + config.num_labels = 3 + input_ids = input_dict["input_ids"] + attention_mask = input_ids.ne(1).to(torch_device) + sequence_labels = ids_tensor([self.model_tester.batch_size], self.model_tester.num_labels) + model = SarvamMLAForSequenceClassification(config) + model.to(torch_device) + model.eval() + result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) + self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) From 8cac1eb72e574942384ef7d6e7e6bd1dd0cad7ae Mon Sep 17 00:00:00 2001 From: Aashay Sarvam Date: Tue, 10 Mar 2026 12:19:05 +0000 Subject: [PATCH 2/8] Fix sarvam_mla docs placement: move from Vision to Text models section Made-with: Cursor --- docs/source/en/_toctree.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index d4327e2e0db0..a47ef9ec4ce5 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -778,6 +778,8 @@ title: RoFormer - local: model_doc/rwkv title: RWKV + - local: model_doc/sarvam_mla + title: SarvamMLA - local: model_doc/seed_oss title: Seed-Oss - local: model_doc/solar_open @@ -946,8 +948,6 @@ title: Segment Anything - local: model_doc/sam_hq title: Segment Anything High Quality - - local: model_doc/sarvam_mla - title: SarvamMLA - local: model_doc/superglue title: SuperGlue - local: model_doc/superpoint From b941609be9a6be26f3ad196bbdd2e61987d7e687 Mon Sep 17 00:00:00 2001 From: Aashay Sarvam Date: Fri, 13 Mar 2026 08:15:33 +0000 Subject: [PATCH 3/8] Address reviewer feedback: use DeepseekV3 classes directly Per vasqu's review: - Remove modular_sarvam_mla.py and modeling_sarvam_mla.py (no need to re-implement identical DeepSeek V3 architecture) - Point auto mappings directly to DeepseekV3 model classes - Move rope type normalization (deepseek_yarn -> yarn) to convert_rope_params_to_dict override - Remove test file (DeepseekV3 tests cover the architecture) - Slim down docs to config-only autodoc Made-with: Cursor --- docs/source/en/model_doc/sarvam_mla.md | 37 +- src/transformers/models/auto/modeling_auto.py | 8 +- .../models/sarvam_mla/__init__.py | 1 - .../sarvam_mla/configuration_sarvam_mla.py | 28 +- .../models/sarvam_mla/modeling_sarvam_mla.py | 736 ------------------ .../models/sarvam_mla/modular_sarvam_mla.py | 90 --- tests/models/sarvam_mla/__init__.py | 0 .../sarvam_mla/test_modeling_sarvam_mla.py | 298 ------- 8 files changed, 20 insertions(+), 1178 deletions(-) delete mode 100644 src/transformers/models/sarvam_mla/modeling_sarvam_mla.py delete mode 100644 src/transformers/models/sarvam_mla/modular_sarvam_mla.py delete mode 100644 tests/models/sarvam_mla/__init__.py delete mode 100644 tests/models/sarvam_mla/test_modeling_sarvam_mla.py diff --git a/docs/source/en/model_doc/sarvam_mla.md b/docs/source/en/model_doc/sarvam_mla.md index d5076c307b1f..1a48ffd6917a 100644 --- a/docs/source/en/model_doc/sarvam_mla.md +++ b/docs/source/en/model_doc/sarvam_mla.md @@ -15,7 +15,7 @@ See the License for the specific language governing permissions and limitations ## Overview -SarvamMLA is a 105B parameter Mixture of Experts (MoE) language model developed by [Sarvam AI](https://www.sarvam.ai/). It uses Multi-head Latent Attention (MLA) combined with sparse MoE routing, architecturally similar to DeepSeek-V3. +SarvamMLA is a 105B parameter Mixture of Experts (MoE) language model developed by [Sarvam AI](https://www.sarvam.ai/). It uses Multi-head Latent Attention (MLA) combined with sparse MoE routing, architecturally identical to DeepSeek-V3. Key architectural features: @@ -24,6 +24,8 @@ Key architectural features: - **DeepSeek YaRN RoPE**: Extended context support up to 131K tokens via YaRN rotary position embeddings. - **Sigmoid routing with group-based top-k**: Token-choice routing using sigmoid scores with expert bias correction and group-aware selection. +This model uses the DeepSeek-V3 architecture with a custom configuration. See the [DeepSeek-V3 documentation](deepseek_v3) for model and forward reference. + ## Usage ```python @@ -41,39 +43,6 @@ outputs = model.generate(**inputs, max_new_tokens=50) print(tokenizer.decode(outputs[0], skip_special_tokens=True)) ``` -For running on limited GPU memory, use quantization: - -```python -from transformers import AutoModelForCausalLM, BitsAndBytesConfig - -quantization_config = BitsAndBytesConfig(load_in_4bit=True) -model = AutoModelForCausalLM.from_pretrained( - "sarvamai/sarvam-105b", - quantization_config=quantization_config, - device_map="auto", -) -``` - ## SarvamMLAConfig [[autodoc]] SarvamMLAConfig - -## SarvamMLAModel - -[[autodoc]] SarvamMLAModel - - forward - -## SarvamMLAForCausalLM - -[[autodoc]] SarvamMLAForCausalLM - - forward - -## SarvamMLAForSequenceClassification - -[[autodoc]] SarvamMLAForSequenceClassification - - forward - -## SarvamMLAForTokenClassification - -[[autodoc]] SarvamMLAForTokenClassification - - forward diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index ac3d13366af6..5538f7fd0d9c 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -397,7 +397,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("sam_hq", "SamHQModel"), ("sam_hq_vision_model", "SamHQVisionModel"), ("sam_vision_model", "SamVisionModel"), - ("sarvam_mla", "SarvamMLAModel"), + ("sarvam_mla", "DeepseekV3Model"), ("seamless_m4t", "SeamlessM4TModel"), ("seamless_m4t_v2", "SeamlessM4Tv2Model"), ("seed_oss", "SeedOssModel"), @@ -723,7 +723,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("roc_bert", "RoCBertForCausalLM"), ("roformer", "RoFormerForCausalLM"), ("rwkv", "RwkvForCausalLM"), - ("sarvam_mla", "SarvamMLAForCausalLM"), + ("sarvam_mla", "DeepseekV3ForCausalLM"), ("seed_oss", "SeedOssForCausalLM"), ("smollm3", "SmolLM3ForCausalLM"), ("solar_open", "SolarOpenForCausalLM"), @@ -1282,7 +1282,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("roberta-prelayernorm", "RobertaPreLayerNormForSequenceClassification"), ("roc_bert", "RoCBertForSequenceClassification"), ("roformer", "RoFormerForSequenceClassification"), - ("sarvam_mla", "SarvamMLAForSequenceClassification"), + ("sarvam_mla", "DeepseekV3ForSequenceClassification"), ("seed_oss", "SeedOssForSequenceClassification"), ("smollm3", "SmolLM3ForSequenceClassification"), ("squeezebert", "SqueezeBertForSequenceClassification"), @@ -1487,7 +1487,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("roberta-prelayernorm", "RobertaPreLayerNormForTokenClassification"), ("roc_bert", "RoCBertForTokenClassification"), ("roformer", "RoFormerForTokenClassification"), - ("sarvam_mla", "SarvamMLAForTokenClassification"), + ("sarvam_mla", "DeepseekV3ForTokenClassification"), ("seed_oss", "SeedOssForTokenClassification"), ("smollm3", "SmolLM3ForTokenClassification"), ("squeezebert", "SqueezeBertForTokenClassification"), diff --git a/src/transformers/models/sarvam_mla/__init__.py b/src/transformers/models/sarvam_mla/__init__.py index 931ca79e5fa9..f9447754575c 100644 --- a/src/transformers/models/sarvam_mla/__init__.py +++ b/src/transformers/models/sarvam_mla/__init__.py @@ -19,7 +19,6 @@ if TYPE_CHECKING: from .configuration_sarvam_mla import * - from .modeling_sarvam_mla import * else: import sys diff --git a/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py b/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py index cce687afddab..2a54b739c784 100644 --- a/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py +++ b/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py @@ -31,7 +31,7 @@ class SarvamMLAConfig(DeepseekV3Config): Example: ```python - >>> from transformers import SarvamMLAModel, SarvamMLAConfig + >>> from transformers import SarvamMLAConfig >>> # Initializing a SarvamMLA style configuration >>> configuration = SarvamMLAConfig() @@ -89,23 +89,10 @@ def __init__( # head_dim in the Hub config.json is set to kv_lora_rank + qk_rope_head_dim # for vLLM MLA compatibility, but internally the model uses qk_rope_head_dim - # for rotary embeddings. Remove it from kwargs to prevent overriding. + # for rotary embeddings. kwargs.pop("head_dim", None) kwargs.pop("q_head_dim", None) - # The Hub config uses "deepseek_yarn" as rope type; normalize to "yarn" - # which is the standard type in ROPE_INIT_FUNCTIONS. - if rope_parameters is not None and rope_parameters.get("type") == "deepseek_yarn": - rope_parameters = dict(rope_parameters) - rope_parameters["type"] = "yarn" - rope_scaling = kwargs.pop("rope_scaling", None) - if rope_scaling is not None: - if rope_scaling.get("type") == "deepseek_yarn": - rope_scaling = dict(rope_scaling) - rope_scaling["type"] = "yarn" - if rope_parameters is None: - rope_parameters = rope_scaling - super().__init__( vocab_size=vocab_size, hidden_size=hidden_size, @@ -143,5 +130,16 @@ def __init__( **kwargs, ) + def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation: set | None = None, **kwargs): + # The Hub config uses "deepseek_yarn" as rope type; normalize to "yarn" + # which is the standard type in ROPE_INIT_FUNCTIONS. + rope_scaling = kwargs.get("rope_scaling", None) + if rope_scaling is not None and rope_scaling.get("type") == "deepseek_yarn": + kwargs["rope_scaling"] = dict(rope_scaling) + kwargs["rope_scaling"]["type"] = "yarn" + if self.rope_parameters and self.rope_parameters.get("type") == "deepseek_yarn": + self.rope_parameters["type"] = "yarn" + return super().convert_rope_params_to_dict(ignore_keys_at_rope_validation=ignore_keys_at_rope_validation, **kwargs) + __all__ = ["SarvamMLAConfig"] diff --git a/src/transformers/models/sarvam_mla/modeling_sarvam_mla.py b/src/transformers/models/sarvam_mla/modeling_sarvam_mla.py deleted file mode 100644 index 5cca52de5272..000000000000 --- a/src/transformers/models/sarvam_mla/modeling_sarvam_mla.py +++ /dev/null @@ -1,736 +0,0 @@ -# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 -# This file was automatically generated from src/transformers/models/sarvam_mla/modular_sarvam_mla.py. -# Do NOT edit this file manually as any edits will be overwritten by the generation of -# the file from the modular. If any change should be done, please apply the change to the -# modular_sarvam_mla.py file directly. One of our CI enforces this. -# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 -# Copyright 2026 Sarvam AI and the HuggingFace Inc. team. 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 math -from collections.abc import Callable -from typing import Optional - -import torch -import torch.nn.functional as F -from torch import nn - -from ... import initialization as init -from ...activations import ACT2FN -from ...cache_utils import Cache, DynamicCache -from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub -from ...masking_utils import create_causal_mask -from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_layers import ( - GenericForSequenceClassification, - GenericForTokenClassification, - GradientCheckpointingLayer, -) -from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast -from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update -from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel -from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple -from ...utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults -from ...utils.output_capturing import capture_outputs -from .configuration_sarvam_mla import SarvamMLAConfig - - -@use_kernel_forward_from_hub("RMSNorm") -class SarvamMLARMSNorm(nn.Module): - def __init__(self, hidden_size, eps: float = 1e-6) -> None: - """ - SarvamMLARMSNorm is equivalent to T5LayerNorm - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) - - def extra_repr(self): - return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" - - -class SarvamMLARotaryEmbedding(nn.Module): - inv_freq: torch.Tensor # fix linting for `register_buffer` - - def __init__(self, config: SarvamMLAConfig, device=None): - super().__init__() - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings - - self.config = config - - self.rope_type = self.config.rope_parameters["rope_type"] - rope_init_fn: Callable = self.compute_default_rope_parameters - if self.rope_type != "default": - rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - inv_freq, self.attention_scaling = rope_init_fn(self.config, device) - - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) - - @staticmethod - def compute_default_rope_parameters( - config: SarvamMLAConfig | None = None, - device: Optional["torch.device"] = None, - seq_len: int | None = None, - ) -> tuple["torch.Tensor", float]: - """ - Computes the inverse frequencies according to the original RoPE implementation - Args: - config ([`~transformers.PreTrainedConfig`]): - The model configuration. - device (`torch.device`): - The device to use for initialization of the inverse frequencies. - seq_len (`int`, *optional*): - The current sequence length. Unused for this type of RoPE. - Returns: - Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the - post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). - """ - base = config.rope_parameters["rope_theta"] - dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads - - attention_factor = 1.0 # Unused in this type of RoPE - - # Compute the inverse frequencies - inv_freq = 1.0 / ( - base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) - ) - return inv_freq, attention_factor - - @torch.no_grad() - @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) - def forward(self, x, position_ids): - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) - position_ids_expanded = position_ids[:, None, :].float() - - device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" - with maybe_autocast(device_type=device_type, enabled=False): # Force float32 - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1) - cos = emb.cos() * self.attention_scaling - sin = emb.sin() * self.attention_scaling - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) - - -class SarvamMLAMLP(nn.Module): - def __init__(self, config, intermediate_size=None): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj - - -class SarvamMLATopkRouter(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.n_routed_experts = config.n_routed_experts - - self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size))) - self.register_buffer("e_score_correction_bias", torch.zeros(self.n_routed_experts)) - - def forward(self, hidden_states): - hidden_states = hidden_states.view(-1, self.config.hidden_size) - router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) - return router_logits - - -@use_experts_implementation -class SarvamMLANaiveMoe(nn.Module): - """Collection of expert weights stored as 3D tensors.""" - - def __init__(self, config): - super().__init__() - self.num_experts = config.num_local_experts - self.hidden_dim = config.hidden_size - self.intermediate_dim = config.moe_intermediate_size - self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) - self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) - self.act_fn = ACT2FN[config.hidden_act] - - def forward( - self, - hidden_states: torch.Tensor, - top_k_index: torch.Tensor, - top_k_weights: torch.Tensor, - ) -> torch.Tensor: - final_hidden_states = torch.zeros_like(hidden_states) - with torch.no_grad(): - expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts) - expert_mask = expert_mask.permute(2, 1, 0) - expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() - - for expert_idx in expert_hit: - expert_idx = expert_idx[0] - if expert_idx == self.num_experts: - continue - top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) - current_state = hidden_states[token_idx] - gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1) - current_hidden_states = self.act_fn(gate) * up - current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx]) - current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None] - final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) - - return final_hidden_states - - -class SarvamMLAMoE(nn.Module): - """ - A mixed expert module containing shared experts. - """ - - def __init__(self, config): - super().__init__() - self.config = config - self.experts = SarvamMLANaiveMoe(config) - self.gate = SarvamMLATopkRouter(config) - self.shared_experts = SarvamMLAMLP( - config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts - ) - self.n_routed_experts = config.n_routed_experts - self.n_group = config.n_group - self.topk_group = config.topk_group - self.norm_topk_prob = config.norm_topk_prob - self.routed_scaling_factor = config.routed_scaling_factor - self.top_k = config.num_experts_per_tok - - def route_tokens_to_experts(self, router_logits): - router_logits = router_logits.sigmoid() - router_logits_for_choice = router_logits + self.gate.e_score_correction_bias - group_scores = ( - router_logits_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group) - .topk(2, dim=-1)[0] - .sum(dim=-1) - ) - group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] - group_mask = torch.zeros_like(group_scores) - group_mask.scatter_(1, group_idx, 1) - score_mask = ( - group_mask.unsqueeze(-1) - .expand(-1, self.n_group, self.n_routed_experts // self.n_group) - .reshape(-1, self.n_routed_experts) - ) - scores_for_choice = router_logits_for_choice.masked_fill(~score_mask.bool(), 0.0) - topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] - topk_weights = router_logits.gather(1, topk_indices) - if self.norm_topk_prob: - denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20 - topk_weights /= denominator - topk_weights = topk_weights * self.routed_scaling_factor - return topk_indices, topk_weights - - def forward(self, hidden_states): - residuals = hidden_states - orig_shape = hidden_states.shape - router_logits = self.gate(hidden_states) - topk_indices, topk_weights = self.route_tokens_to_experts(router_logits) - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) - hidden_states = hidden_states + self.shared_experts(residuals) - return hidden_states - - -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) - - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask - - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() - - return attn_output, attn_weights - - -def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): - r""" - TODO let's just use the original freqcis computation to not have the view - transpose + reshape! This is not optimized! - Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - position_ids (`torch.Tensor`): - The position indices of the tokens corresponding to the query and key tensors. For example, this can be - used to pass offsetted position ids when working with a KV-cache. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - - b, h, s, d = q.shape - q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) - - b, h, s, d = k.shape - k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) - - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - -def yarn_get_mscale(scale=1, mscale=1): - if scale <= 1: - return 1.0 - return 0.1 * mscale * math.log(scale) + 1.0 - - -class SarvamMLAAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config: SarvamMLAConfig, layer_idx: int): - super().__init__() - self.config = config - self.layer_idx = layer_idx - self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads - self.attention_dropout = config.attention_dropout - self.num_heads = config.num_attention_heads - - self.q_lora_rank = config.q_lora_rank - self.qk_rope_head_dim = config.qk_rope_head_dim - self.kv_lora_rank = config.kv_lora_rank - self.v_head_dim = config.v_head_dim - self.qk_nope_head_dim = config.qk_nope_head_dim - self.qk_head_dim = config.qk_head_dim - - self.is_causal = True - if self.q_lora_rank is None: - self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False) - else: - self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=config.attention_bias) - self.q_a_layernorm = SarvamMLARMSNorm(config.q_lora_rank) - self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) - - self.kv_a_proj_with_mqa = nn.Linear( - config.hidden_size, - self.kv_lora_rank + self.qk_rope_head_dim, - bias=config.attention_bias, - ) - self.kv_a_layernorm = SarvamMLARMSNorm(self.kv_lora_rank) - self.kv_b_proj = nn.Linear( - self.kv_lora_rank, - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), - bias=False, - ) - - self.o_proj = nn.Linear( - self.num_heads * self.v_head_dim, - config.hidden_size, - bias=config.attention_bias, - ) - - self.scaling = self.qk_head_dim ** (-0.5) - if self.config.rope_parameters.get("rope_type", "default") != "default": - mscale_all_dim = self.config.rope_parameters.get("mscale_all_dim", 0) - scaling_factor = self.config.rope_parameters["factor"] - if mscale_all_dim: - mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) - self.scaling = self.scaling * mscale * mscale - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - attention_mask: torch.Tensor | None, - past_key_values: Cache | None = None, - **kwargs: Unpack[FlashAttentionKwargs], - ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: - batch_size, seq_length = hidden_states.shape[:-1] - query_shape = (batch_size, seq_length, -1, self.qk_head_dim) - key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim) - - if self.q_lora_rank is None: - q_states = self.q_proj(hidden_states) - else: - q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) - q_states = q_states.view(query_shape).transpose(1, 2) - q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - - compressed_kv = self.kv_a_proj_with_mqa(hidden_states) - k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - - k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2) - k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - - k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) - - cos, sin = position_embeddings - if self.config.rope_interleave: # support using interleaved weights for efficiency - q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) - else: - q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) - k_rot = k_rot.expand(*k_pass.shape[:-1], -1) - - query_states = torch.cat((q_pass, q_rot), dim=-1) - key_states = torch.cat((k_pass, k_rot), dim=-1) - - if past_key_values is not None: - key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) - - if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: - value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) - - attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( - self.config._attn_implementation, eager_attention_forward - ) - - attn_output, attn_weights = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - **kwargs, - ) - - if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: - attn_output = attn_output[:, :, :, : self.v_head_dim] - - attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() - attn_output = self.o_proj(attn_output) - return attn_output, attn_weights - - -class SarvamMLADecoderLayer(GradientCheckpointingLayer): - def __init__(self, config: SarvamMLAConfig, layer_idx: int): - super().__init__() - self.hidden_size = config.hidden_size - - self.self_attn = SarvamMLAAttention(config=config, layer_idx=layer_idx) - - if layer_idx >= config.first_k_dense_replace: - self.mlp = SarvamMLAMoE(config) - else: - self.mlp = SarvamMLAMLP(config) - - self.input_layernorm = SarvamMLARMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = SarvamMLARMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - use_cache: bool | None = False, - position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> torch.Tensor: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - # Self Attention - hidden_states, _ = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - position_embeddings=position_embeddings, - **kwargs, - ) - hidden_states = residual + hidden_states - - # Fully Connected - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - return hidden_states - - -@auto_docstring -class SarvamMLAPreTrainedModel(PreTrainedModel): - config: SarvamMLAConfig - base_model_prefix = "model" - supports_gradient_checkpointing = True - _no_split_modules = ["SarvamMLADecoderLayer"] - _skip_keys_device_placement = ["past_key_values"] - _supports_flash_attn = True - _supports_sdpa = True - _supports_flex_attn = True - - _can_compile_fullgraph = True - _supports_attention_backend = True - _can_record_outputs = { - "hidden_states": SarvamMLADecoderLayer, - "attentions": SarvamMLAAttention, - } - _keep_in_fp32_modules_strict = ["e_score_correction_bias"] - _keys_to_ignore_on_load_unexpected = [r"model\.layers\.61.*"] - - @torch.no_grad() - def _init_weights(self, module): - super()._init_weights(module) - if isinstance(module, SarvamMLATopkRouter): - init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) - init.zeros_(module.e_score_correction_bias) - elif isinstance(module, SarvamMLANaiveMoe): - init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range) - init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) - - -@auto_docstring -class SarvamMLAModel(SarvamMLAPreTrainedModel): - def __init__(self, config: SarvamMLAConfig): - super().__init__(config) - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) - self.layers = nn.ModuleList( - [SarvamMLADecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] - ) - self.norm = SarvamMLARMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = SarvamMLARotaryEmbedding(config=config) - self.gradient_checkpointing = False - - # Initialize weights and apply final processing - self.post_init() - - @merge_with_config_defaults - @capture_outputs - @auto_docstring - def forward( - self, - input_ids: torch.LongTensor | None = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - inputs_embeds: torch.FloatTensor | None = None, - use_cache: bool | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPast: - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - - if inputs_embeds is None: - inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) - - if use_cache and past_key_values is None: - past_key_values = DynamicCache(config=self.config) - - if position_ids is None: - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens - position_ids = position_ids.unsqueeze(0) - - causal_mask = create_causal_mask( - config=self.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=past_key_values, - position_ids=position_ids, - ) - - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) - - for decoder_layer in self.layers[: self.config.num_hidden_layers]: - hidden_states = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_embeddings=position_embeddings, - position_ids=position_ids, - past_key_values=past_key_values, - use_cache=use_cache, - **kwargs, - ) - - hidden_states = self.norm(hidden_states) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=past_key_values, - ) - - -@auto_docstring -class SarvamMLAForCausalLM(SarvamMLAPreTrainedModel, GenerationMixin): - _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} - _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} - - def __init__(self, config): - super().__init__(config) - self.model = SarvamMLAModel(config) - self.vocab_size = config.vocab_size - self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - - # Initialize weights and apply final processing - self.post_init() - - @can_return_tuple - @auto_docstring - def forward( - self, - input_ids: torch.LongTensor | None = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - inputs_embeds: torch.FloatTensor | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - logits_to_keep: int | torch.Tensor = 0, - **kwargs: Unpack[TransformersKwargs], - ) -> CausalLMOutputWithPast: - r""" - Example: - - ```python - >>> from transformers import AutoTokenizer, SarvamMLAForCausalLM - - >>> model = SarvamMLAForCausalLM.from_pretrained("meta-sarvam_mla/SarvamMLA-2-7b-hf") - >>> tokenizer = AutoTokenizer.from_pretrained("meta-sarvam_mla/SarvamMLA-2-7b-hf") - - >>> prompt = "Hey, are you conscious? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." - ```""" - outputs: BaseModelOutputWithPast = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - **kwargs, - ) - - hidden_states = outputs.last_hidden_state - # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep - logits = self.lm_head(hidden_states[:, slice_indices, :]) - - loss = None - if labels is not None: - loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - -class SarvamMLAForSequenceClassification(GenericForSequenceClassification, SarvamMLAPreTrainedModel): - pass - - -class SarvamMLAForTokenClassification(GenericForTokenClassification, SarvamMLAPreTrainedModel): - pass - - -__all__ = [ - "SarvamMLAPreTrainedModel", - "SarvamMLAModel", - "SarvamMLAForCausalLM", - "SarvamMLAForSequenceClassification", - "SarvamMLAForTokenClassification", -] diff --git a/src/transformers/models/sarvam_mla/modular_sarvam_mla.py b/src/transformers/models/sarvam_mla/modular_sarvam_mla.py deleted file mode 100644 index 6963a86cd7b2..000000000000 --- a/src/transformers/models/sarvam_mla/modular_sarvam_mla.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright 2026 Sarvam AI and the HuggingFace Inc. team. 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 ...modeling_layers import GenericForSequenceClassification, GenericForTokenClassification -from ..deepseek_v3.modeling_deepseek_v3 import ( - DeepseekV3Attention, - DeepseekV3DecoderLayer, - DeepseekV3ForCausalLM, - DeepseekV3MLP, - DeepseekV3MoE, - DeepseekV3Model, - DeepseekV3NaiveMoe, - DeepseekV3PreTrainedModel, - DeepseekV3RMSNorm, - DeepseekV3RotaryEmbedding, - DeepseekV3TopkRouter, -) -from .configuration_sarvam_mla import SarvamMLAConfig - - -class SarvamMLARMSNorm(DeepseekV3RMSNorm): - pass - - -class SarvamMLARotaryEmbedding(DeepseekV3RotaryEmbedding): - pass - - -class SarvamMLAMLP(DeepseekV3MLP): - pass - - -class SarvamMLATopkRouter(DeepseekV3TopkRouter): - pass - - -class SarvamMLANaiveMoe(DeepseekV3NaiveMoe): - pass - - -class SarvamMLAMoE(DeepseekV3MoE): - pass - - -class SarvamMLAAttention(DeepseekV3Attention): - pass - - -class SarvamMLADecoderLayer(DeepseekV3DecoderLayer): - pass - - -class SarvamMLAPreTrainedModel(DeepseekV3PreTrainedModel): - pass - - -class SarvamMLAModel(DeepseekV3Model): - pass - - -class SarvamMLAForCausalLM(DeepseekV3ForCausalLM): - pass - - -class SarvamMLAForSequenceClassification(GenericForSequenceClassification, SarvamMLAPreTrainedModel): - pass - - -class SarvamMLAForTokenClassification(GenericForTokenClassification, SarvamMLAPreTrainedModel): - pass - - -__all__ = [ - "SarvamMLAPreTrainedModel", - "SarvamMLAModel", - "SarvamMLAForCausalLM", - "SarvamMLAForSequenceClassification", - "SarvamMLAForTokenClassification", -] diff --git a/tests/models/sarvam_mla/__init__.py b/tests/models/sarvam_mla/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/tests/models/sarvam_mla/test_modeling_sarvam_mla.py b/tests/models/sarvam_mla/test_modeling_sarvam_mla.py deleted file mode 100644 index fc4b7a56d673..000000000000 --- a/tests/models/sarvam_mla/test_modeling_sarvam_mla.py +++ /dev/null @@ -1,298 +0,0 @@ -# Copyright 2026 Sarvam AI and the HuggingFace Inc. team. 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. -"""Testing suite for the PyTorch SarvamMLA model.""" - -import unittest - -from parameterized import parameterized - -from transformers import SarvamMLAConfig, is_torch_available -from transformers.testing_utils import ( - require_torch, - torch_device, -) - -from ...generation.test_utils import GenerationTesterMixin -from ...test_configuration_common import ConfigTester -from ...test_modeling_common import ModelTesterMixin, ids_tensor -from ...test_pipeline_mixin import PipelineTesterMixin - - -if is_torch_available(): - import torch - - from transformers import ( - Cache, - SarvamMLAForCausalLM, - SarvamMLAForSequenceClassification, - SarvamMLAForTokenClassification, - SarvamMLAModel, - ) - - -class SarvamMLAModelTester: - if is_torch_available(): - causal_lm_class = SarvamMLAForCausalLM - - def __init__( - self, - parent, - batch_size=13, - seq_length=7, - is_training=True, - use_input_mask=True, - use_token_type_ids=False, - use_labels=True, - vocab_size=99, - hidden_size=32, - intermediate_size=32, - moe_intermediate_size=16, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=4, - num_shared_experts=1, - num_experts=8, - routed_scaling_factor=2.5, - kv_lora_rank=16, - q_lora_rank=None, - qk_rope_head_dim=16, - v_head_dim=32, - qk_nope_head_dim=32, - n_group=2, - topk_group=1, - num_experts_per_tok=8, - first_k_dense_replace=1, - norm_topk_prob=True, - hidden_act="silu", - max_position_embeddings=512, - initializer_range=0.02, - attention_probs_dropout_prob=0.0, - type_vocab_size=16, - type_sequence_label_size=2, - num_labels=3, - num_choices=4, - pad_token_id=0, - scope=None, - ): - self.parent = parent - self.batch_size = batch_size - self.seq_length = seq_length - self.is_training = is_training - self.use_input_mask = use_input_mask - self.use_token_type_ids = use_token_type_ids - self.use_labels = use_labels - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.moe_intermediate_size = moe_intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.num_shared_experts = num_shared_experts - self.num_experts = num_experts - self.routed_scaling_factor = routed_scaling_factor - self.kv_lora_rank = kv_lora_rank - self.q_lora_rank = q_lora_rank - self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim - self.qk_nope_head_dim = qk_nope_head_dim - self.n_group = n_group - self.topk_group = topk_group - self.num_experts_per_tok = num_experts_per_tok - self.first_k_dense_replace = first_k_dense_replace - self.norm_topk_prob = norm_topk_prob - self.hidden_act = hidden_act - self.max_position_embeddings = max_position_embeddings - self.initializer_range = initializer_range - self.attention_probs_dropout_prob = attention_probs_dropout_prob - self.type_vocab_size = type_vocab_size - self.type_sequence_label_size = type_sequence_label_size - self.num_labels = num_labels - self.num_choices = num_choices - self.pad_token_id = pad_token_id - self.scope = scope - - def prepare_config_and_inputs(self): - input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size) - - input_mask = None - if self.use_input_mask: - input_mask = torch.tril(torch.ones_like(input_ids).to(torch_device)) - - token_type_ids = None - if self.use_token_type_ids: - token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size) - - sequence_labels = None - token_labels = None - choice_labels = None - if self.use_labels: - sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size) - token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels) - choice_labels = ids_tensor([self.batch_size], self.num_choices) - - config = self.get_config() - - return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels - - def get_config(self): - return SarvamMLAConfig( - vocab_size=self.vocab_size, - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, - moe_intermediate_size=self.moe_intermediate_size, - num_hidden_layers=self.num_hidden_layers, - num_attention_heads=self.num_attention_heads, - num_key_value_heads=self.num_key_value_heads, - num_shared_experts=self.num_shared_experts, - num_experts=self.num_experts, - routed_scaling_factor=self.routed_scaling_factor, - kv_lora_rank=self.kv_lora_rank, - q_lora_rank=self.q_lora_rank, - qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim, - qk_nope_head_dim=self.qk_nope_head_dim, - n_group=self.n_group, - topk_group=self.topk_group, - num_experts_per_tok=self.num_experts_per_tok, - first_k_dense_replace=self.first_k_dense_replace, - norm_topk_prob=self.norm_topk_prob, - hidden_act=self.hidden_act, - max_position_embeddings=self.max_position_embeddings, - initializer_range=self.initializer_range, - use_cache=True, - pad_token_id=self.pad_token_id, - attention_dropout=self.attention_probs_dropout_prob, - ) - - def create_and_check_model( - self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels - ): - model = SarvamMLAModel(config=config) - model.to(torch_device) - model.eval() - result = model(input_ids, attention_mask=input_mask) - result = model(input_ids) - self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size)) - - def prepare_config_and_inputs_for_common(self): - config_and_inputs = self.prepare_config_and_inputs() - ( - config, - input_ids, - token_type_ids, - input_mask, - sequence_labels, - token_labels, - choice_labels, - ) = config_and_inputs - inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask} - return config, inputs_dict - - -@require_torch -class SarvamMLAModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase): - all_model_classes = ( - ( - SarvamMLAModel, - SarvamMLAForCausalLM, - SarvamMLAForSequenceClassification, - SarvamMLAForTokenClassification, - ) - if is_torch_available() - else () - ) - all_generative_model_classes = (SarvamMLAForCausalLM,) if is_torch_available() else () - pipeline_model_mapping = ( - { - "feature-extraction": SarvamMLAModel, - "text-classification": SarvamMLAForSequenceClassification, - "token-classification": SarvamMLAForTokenClassification, - "text-generation": SarvamMLAForCausalLM, - "zero-shot": SarvamMLAForSequenceClassification, - } - if is_torch_available() - else {} - ) - - model_split_percents = [0.5, 0.7, 0.8] - - _torch_compile_train_cls = SarvamMLAForCausalLM if is_torch_available() else None - - def setUp(self): - self.model_tester = SarvamMLAModelTester(self) - self.config_tester = ConfigTester(self, config_class=SarvamMLAConfig, hidden_size=37) - - def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_length, config): - self.assertIsInstance(past_key_values, Cache) - - expected_common_shape = ( - batch_size, - getattr(config, "num_key_value_heads", config.num_attention_heads), - seq_length, - ) - expected_key_shape = expected_common_shape + (config.qk_nope_head_dim + config.qk_rope_head_dim,) - expected_value_shape = expected_common_shape + (config.v_head_dim,) - - for layer in past_key_values.layers: - self.assertEqual(layer.keys.shape, expected_key_shape) - self.assertEqual(layer.values.shape, expected_value_shape) - - @parameterized.expand([("random",), ("same",)]) - @unittest.skip("SarvamMLA uses MLA, not compatible with assisted decoding") - def test_assisted_decoding_matches_greedy_search(self, assistant_type): - pass - - @unittest.skip("SarvamMLA uses MLA, not compatible with assisted decoding") - def test_prompt_lookup_decoding_matches_greedy_search(self, assistant_type): - pass - - @unittest.skip("SarvamMLA uses MLA, not compatible with assisted decoding") - def test_assisted_decoding_sample(self): - pass - - @unittest.skip("SarvamMLA uses MLA, not compatible with standard cache format") - def test_beam_search_generate_dict_outputs_use_cache(self): - pass - - @unittest.skip("SarvamMLA uses MLA, not compatible with standard cache format") - def test_greedy_generate_dict_outputs_use_cache(self): - pass - - @unittest.skip(reason="SDPA can't dispatch on flash due to unsupported head dims") - def test_sdpa_can_dispatch_on_flash(self): - pass - - @unittest.skip("SarvamMLA has MoE, output can be different") - def test_model_outputs_equivalence(self, **kwargs): - pass - - def test_config(self): - self.config_tester.run_common_tests() - - def test_model(self): - config_and_inputs = self.model_tester.prepare_config_and_inputs() - self.model_tester.create_and_check_model(*config_and_inputs) - - def test_sarvam_mla_sequence_classification_model(self): - config, input_dict = self.model_tester.prepare_config_and_inputs_for_common() - config.num_labels = 3 - input_ids = input_dict["input_ids"] - attention_mask = input_ids.ne(1).to(torch_device) - sequence_labels = ids_tensor([self.model_tester.batch_size], self.model_tester.num_labels) - model = SarvamMLAForSequenceClassification(config) - model.to(torch_device) - model.eval() - result = model(input_ids, attention_mask=attention_mask, labels=sequence_labels) - self.assertEqual(result.logits.shape, (self.model_tester.batch_size, self.model_tester.num_labels)) From 07067094c19f5df9b48149ea0c34c6f8ed0cd41e Mon Sep 17 00:00:00 2001 From: Aashay Sarvam Date: Mon, 16 Mar 2026 07:49:27 +0000 Subject: [PATCH 4/8] Refactor SarvamMLA config to use modular pattern Move SarvamMLAConfig definition into modular_sarvam_mla.py and auto-generate configuration_sarvam_mla.py from it, following the canonical transformers modular pattern. Made-with: Cursor --- .../sarvam_mla/configuration_sarvam_mla.py | 142 ++++++++++------- .../models/sarvam_mla/modular_sarvam_mla.py | 144 ++++++++++++++++++ 2 files changed, 235 insertions(+), 51 deletions(-) create mode 100644 src/transformers/models/sarvam_mla/modular_sarvam_mla.py diff --git a/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py b/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py index 2a54b739c784..e64ec6f41519 100644 --- a/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py +++ b/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py @@ -1,3 +1,9 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/sarvam_mla/modular_sarvam_mla.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_sarvam_mla.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2026 Sarvam AI and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -11,14 +17,13 @@ # 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. -"""SarvamMLA model configuration""" +from ...configuration_utils import PreTrainedConfig from ...utils import auto_docstring -from ..deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config @auto_docstring(checkpoint="sarvamai/sarvam-105b") -class SarvamMLAConfig(DeepseekV3Config): +class SarvamMLAConfig(PreTrainedConfig): r""" n_group (`int`, *optional*, defaults to 16): Number of groups for routed experts. @@ -41,10 +46,25 @@ class SarvamMLAConfig(DeepseekV3Config): ```""" model_type = "sarvam_mla" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } attribute_map = { - "n_routed_experts": "num_experts", - "n_shared_experts": "num_shared_experts", - "num_local_experts": "num_experts", + "num_local_experts": "n_routed_experts", } def __init__( @@ -56,8 +76,8 @@ def __init__( num_hidden_layers=32, num_attention_heads=64, num_key_value_heads=None, - num_shared_experts=1, - num_experts=128, + n_shared_experts=1, + n_routed_experts=128, routed_scaling_factor=2.5, kv_lora_rank=512, q_lora_rank=None, @@ -77,6 +97,7 @@ def __init__( pad_token_id=0, bos_token_id=None, eos_token_id=1, + pretraining_tp=1, tie_word_embeddings=False, rope_parameters=None, rope_interleave=True, @@ -84,62 +105,81 @@ def __init__( attention_dropout=0.0, **kwargs, ): - self.num_experts = num_experts - self.num_shared_experts = num_shared_experts + # Hub config.json uses num_experts/num_shared_experts; map to parent names + n_routed_experts = kwargs.pop("num_experts", n_routed_experts) + n_shared_experts = kwargs.pop("num_shared_experts", n_shared_experts) - # head_dim in the Hub config.json is set to kv_lora_rank + qk_rope_head_dim - # for vLLM MLA compatibility, but internally the model uses qk_rope_head_dim - # for rotary embeddings. + # head_dim in Hub config.json is kv_lora_rank + qk_rope_head_dim (for vLLM + # MLA compat), but DeepseekV3Config computes it as qk_rope_head_dim. kwargs.pop("head_dim", None) kwargs.pop("q_head_dim", None) + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.head_dim = qk_rope_head_dim + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.rope_interleave = rope_interleave + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads - super().__init__( - vocab_size=vocab_size, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - moe_intermediate_size=moe_intermediate_size, - num_hidden_layers=num_hidden_layers, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - n_shared_experts=num_shared_experts, - n_routed_experts=num_experts, - routed_scaling_factor=routed_scaling_factor, - kv_lora_rank=kv_lora_rank, - q_lora_rank=q_lora_rank, - qk_rope_head_dim=qk_rope_head_dim, - v_head_dim=v_head_dim, - qk_nope_head_dim=qk_nope_head_dim, - n_group=n_group, - topk_group=topk_group, - num_experts_per_tok=num_experts_per_tok, - first_k_dense_replace=first_k_dense_replace, - norm_topk_prob=norm_topk_prob, - hidden_act=hidden_act, - max_position_embeddings=max_position_embeddings, - initializer_range=initializer_range, - rms_norm_eps=rms_norm_eps, - use_cache=use_cache, - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - tie_word_embeddings=tie_word_embeddings, - rope_parameters=rope_parameters, - rope_interleave=rope_interleave, - attention_bias=attention_bias, - attention_dropout=attention_dropout, - **kwargs, - ) + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.rope_parameters = rope_parameters + + self.tie_word_embeddings = tie_word_embeddings + self.pad_token_id = pad_token_id + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + super().__init__(**kwargs) def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation: set | None = None, **kwargs): # The Hub config uses "deepseek_yarn" as rope type; normalize to "yarn" # which is the standard type in ROPE_INIT_FUNCTIONS. - rope_scaling = kwargs.get("rope_scaling", None) + rope_scaling = kwargs.get("rope_scaling") if rope_scaling is not None and rope_scaling.get("type") == "deepseek_yarn": kwargs["rope_scaling"] = dict(rope_scaling) kwargs["rope_scaling"]["type"] = "yarn" if self.rope_parameters and self.rope_parameters.get("type") == "deepseek_yarn": self.rope_parameters["type"] = "yarn" - return super().convert_rope_params_to_dict(ignore_keys_at_rope_validation=ignore_keys_at_rope_validation, **kwargs) + rope_scaling = kwargs.pop("rope_scaling", None) + self.rope_parameters = rope_scaling or self.rope_parameters + self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} + + # Standardize and validate the correctness of rotary position embeddings parameters + self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta)) + self.standardize_rope_params() + self.validate_rope(ignore_keys=ignore_keys_at_rope_validation) + + # Convert to float because RoPE fn expect a float. Models on the hub were saved as int + for key in ["beta_fast", "beta_slow", "factor"]: + if key in self.rope_parameters: + self.rope_parameters[key] = float(self.rope_parameters[key]) + return kwargs __all__ = ["SarvamMLAConfig"] diff --git a/src/transformers/models/sarvam_mla/modular_sarvam_mla.py b/src/transformers/models/sarvam_mla/modular_sarvam_mla.py new file mode 100644 index 000000000000..081a7d6bc8f7 --- /dev/null +++ b/src/transformers/models/sarvam_mla/modular_sarvam_mla.py @@ -0,0 +1,144 @@ +# Copyright 2026 Sarvam AI and the HuggingFace Inc. team. 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. +"""SarvamMLA model configuration""" + +from ...utils import auto_docstring +from ..deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config + + +@auto_docstring(checkpoint="sarvamai/sarvam-105b") +class SarvamMLAConfig(DeepseekV3Config): + r""" + n_group (`int`, *optional*, defaults to 16): + Number of groups for routed experts. + rope_interleave (`bool`, *optional*, defaults to `True`): + Whether to interleave the rotary position embeddings. + first_k_dense_replace (`int`, *optional*, defaults to 1): + Number of dense layers in shallow layers(embed->dense->moe->moe...->lm_head). + \--k dense layers--/ + + Example: + + ```python + >>> from transformers import SarvamMLAConfig + + >>> # Initializing a SarvamMLA style configuration + >>> configuration = SarvamMLAConfig() + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "sarvam_mla" + + def __init__( + self, + vocab_size=262144, + hidden_size=4096, + intermediate_size=16384, + moe_intermediate_size=2048, + num_hidden_layers=32, + num_attention_heads=64, + num_key_value_heads=None, + n_shared_experts=1, + n_routed_experts=128, + routed_scaling_factor=2.5, + kv_lora_rank=512, + q_lora_rank=None, + qk_rope_head_dim=64, + v_head_dim=128, + qk_nope_head_dim=128, + n_group=16, + topk_group=2, + num_experts_per_tok=8, + first_k_dense_replace=1, + norm_topk_prob=True, + hidden_act="silu", + max_position_embeddings=4096, + initializer_range=0.006, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + bos_token_id=None, + eos_token_id=1, + pretraining_tp=1, + tie_word_embeddings=False, + rope_parameters=None, + rope_interleave=True, + attention_bias=False, + attention_dropout=0.0, + **kwargs, + ): + # Hub config.json uses num_experts/num_shared_experts; map to parent names + n_routed_experts = kwargs.pop("num_experts", n_routed_experts) + n_shared_experts = kwargs.pop("num_shared_experts", n_shared_experts) + + # head_dim in Hub config.json is kv_lora_rank + qk_rope_head_dim (for vLLM + # MLA compat), but DeepseekV3Config computes it as qk_rope_head_dim. + kwargs.pop("head_dim", None) + kwargs.pop("q_head_dim", None) + + super().__init__( + vocab_size=vocab_size, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + moe_intermediate_size=moe_intermediate_size, + num_hidden_layers=num_hidden_layers, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + n_shared_experts=n_shared_experts, + n_routed_experts=n_routed_experts, + routed_scaling_factor=routed_scaling_factor, + kv_lora_rank=kv_lora_rank, + q_lora_rank=q_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + qk_nope_head_dim=qk_nope_head_dim, + n_group=n_group, + topk_group=topk_group, + num_experts_per_tok=num_experts_per_tok, + first_k_dense_replace=first_k_dense_replace, + norm_topk_prob=norm_topk_prob, + hidden_act=hidden_act, + max_position_embeddings=max_position_embeddings, + initializer_range=initializer_range, + rms_norm_eps=rms_norm_eps, + use_cache=use_cache, + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pretraining_tp=pretraining_tp, + tie_word_embeddings=tie_word_embeddings, + rope_parameters=rope_parameters, + rope_interleave=rope_interleave, + attention_bias=attention_bias, + attention_dropout=attention_dropout, + **kwargs, + ) + + def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation: set | None = None, **kwargs): + # The Hub config uses "deepseek_yarn" as rope type; normalize to "yarn" + # which is the standard type in ROPE_INIT_FUNCTIONS. + rope_scaling = kwargs.get("rope_scaling", None) + if rope_scaling is not None and rope_scaling.get("type") == "deepseek_yarn": + kwargs["rope_scaling"] = dict(rope_scaling) + kwargs["rope_scaling"]["type"] = "yarn" + if self.rope_parameters and self.rope_parameters.get("type") == "deepseek_yarn": + self.rope_parameters["type"] = "yarn" + return super().convert_rope_params_to_dict( + ignore_keys_at_rope_validation=ignore_keys_at_rope_validation, **kwargs + ) + + +__all__ = ["SarvamMLAConfig"] From 0f5d73dfbd571d62d5c430368d4a68a2ba08189c Mon Sep 17 00:00:00 2001 From: Aashay Sarvam Date: Tue, 17 Mar 2026 07:22:18 +0000 Subject: [PATCH 5/8] Address reviewer feedback: simplify modular config, rebase on main - Remove torch_dtype="auto" from docs (now default) - Simplify modular_sarvam_mla.py to only override defaults that differ from DeepseekV3Config (no __init__, no workarounds) - Add @strict(accept_kwargs=True) for config validation (#41250) - Regenerate configuration_sarvam_mla.py with dataclass fields and __post_init__ pattern - Hub config.json changes needed: remove head_dim/q_head_dim, change rope_scaling.type to "yarn", update architectures Made-with: Cursor --- docs/source/en/model_doc/sarvam_mla.md | 1 - .../sarvam_mla/configuration_sarvam_mla.py | 148 ++++++------------ .../models/sarvam_mla/modular_sarvam_mla.py | 113 ++----------- 3 files changed, 63 insertions(+), 199 deletions(-) diff --git a/docs/source/en/model_doc/sarvam_mla.md b/docs/source/en/model_doc/sarvam_mla.md index 1a48ffd6917a..1190f17725a1 100644 --- a/docs/source/en/model_doc/sarvam_mla.md +++ b/docs/source/en/model_doc/sarvam_mla.md @@ -34,7 +34,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer model = AutoModelForCausalLM.from_pretrained( "sarvamai/sarvam-105b", device_map="auto", - torch_dtype="auto", ) tokenizer = AutoTokenizer.from_pretrained("sarvamai/sarvam-105b") diff --git a/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py b/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py index e64ec6f41519..f5701edc656d 100644 --- a/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py +++ b/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py @@ -18,11 +18,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +from huggingface_hub.dataclasses import strict + from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @auto_docstring(checkpoint="sarvamai/sarvam-105b") +@strict(accept_kwargs=True) class SarvamMLAConfig(PreTrainedConfig): r""" n_group (`int`, *optional*, defaults to 16): @@ -67,105 +71,50 @@ class SarvamMLAConfig(PreTrainedConfig): "num_local_experts": "n_routed_experts", } - def __init__( - self, - vocab_size=262144, - hidden_size=4096, - intermediate_size=16384, - moe_intermediate_size=2048, - num_hidden_layers=32, - num_attention_heads=64, - num_key_value_heads=None, - n_shared_experts=1, - n_routed_experts=128, - routed_scaling_factor=2.5, - kv_lora_rank=512, - q_lora_rank=None, - qk_rope_head_dim=64, - v_head_dim=128, - qk_nope_head_dim=128, - n_group=16, - topk_group=2, - num_experts_per_tok=8, - first_k_dense_replace=1, - norm_topk_prob=True, - hidden_act="silu", - max_position_embeddings=4096, - initializer_range=0.006, - rms_norm_eps=1e-6, - use_cache=True, - pad_token_id=0, - bos_token_id=None, - eos_token_id=1, - pretraining_tp=1, - tie_word_embeddings=False, - rope_parameters=None, - rope_interleave=True, - attention_bias=False, - attention_dropout=0.0, - **kwargs, - ): - # Hub config.json uses num_experts/num_shared_experts; map to parent names - n_routed_experts = kwargs.pop("num_experts", n_routed_experts) - n_shared_experts = kwargs.pop("num_shared_experts", n_shared_experts) - - # head_dim in Hub config.json is kv_lora_rank + qk_rope_head_dim (for vLLM - # MLA compat), but DeepseekV3Config computes it as qk_rope_head_dim. - kwargs.pop("head_dim", None) - kwargs.pop("q_head_dim", None) - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.moe_intermediate_size = moe_intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.n_shared_experts = n_shared_experts - self.n_routed_experts = n_routed_experts - self.routed_scaling_factor = routed_scaling_factor - self.kv_lora_rank = kv_lora_rank - self.q_lora_rank = q_lora_rank - self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim - self.qk_nope_head_dim = qk_nope_head_dim - self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim - self.head_dim = qk_rope_head_dim - self.n_group = n_group - self.topk_group = topk_group - self.num_experts_per_tok = num_experts_per_tok - self.first_k_dense_replace = first_k_dense_replace - self.norm_topk_prob = norm_topk_prob - self.rope_interleave = rope_interleave - - # for backward compatibility - if num_key_value_heads is None: - num_key_value_heads = num_attention_heads - - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.pretraining_tp = pretraining_tp - self.use_cache = use_cache - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.rope_parameters = rope_parameters - - self.tie_word_embeddings = tie_word_embeddings - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - super().__init__(**kwargs) - - def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation: set | None = None, **kwargs): - # The Hub config uses "deepseek_yarn" as rope type; normalize to "yarn" - # which is the standard type in ROPE_INIT_FUNCTIONS. - rope_scaling = kwargs.get("rope_scaling") - if rope_scaling is not None and rope_scaling.get("type") == "deepseek_yarn": - kwargs["rope_scaling"] = dict(rope_scaling) - kwargs["rope_scaling"]["type"] = "yarn" - if self.rope_parameters and self.rope_parameters.get("type") == "deepseek_yarn": - self.rope_parameters["type"] = "yarn" + vocab_size: int = 262144 + hidden_size: int = 4096 + intermediate_size: int = 16384 + moe_intermediate_size: int = 2048 + num_hidden_layers: int = 32 + num_attention_heads: int = 64 + num_key_value_heads: int | None = None + n_shared_experts: int = 1 + n_routed_experts: int = 128 + routed_scaling_factor: float = 2.5 + kv_lora_rank: int = 512 + q_lora_rank: int | None = None + qk_rope_head_dim: int = 64 + v_head_dim: int | None = 128 + qk_nope_head_dim: int = 128 + n_group: int | None = 16 + topk_group: int | None = 2 + num_experts_per_tok: int | None = 8 + first_k_dense_replace: int | None = 1 + norm_topk_prob: bool | None = True + hidden_act: str = "silu" + max_position_embeddings: int = 4096 + initializer_range: float = 0.006 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 + pretraining_tp: int | None = 1 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + rope_interleave: bool | None = True + attention_bias: bool = False + attention_dropout: float | int | None = 0.0 + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.head_dim = self.qk_rope_head_dim + super().__post_init__(**kwargs) + + def convert_rope_params_to_dict(self, **kwargs): rope_scaling = kwargs.pop("rope_scaling", None) self.rope_parameters = rope_scaling or self.rope_parameters self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} @@ -173,7 +122,6 @@ def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation: set | None # Standardize and validate the correctness of rotary position embeddings parameters self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta)) self.standardize_rope_params() - self.validate_rope(ignore_keys=ignore_keys_at_rope_validation) # Convert to float because RoPE fn expect a float. Models on the hub were saved as int for key in ["beta_fast", "beta_slow", "factor"]: diff --git a/src/transformers/models/sarvam_mla/modular_sarvam_mla.py b/src/transformers/models/sarvam_mla/modular_sarvam_mla.py index 081a7d6bc8f7..e092dbb2764b 100644 --- a/src/transformers/models/sarvam_mla/modular_sarvam_mla.py +++ b/src/transformers/models/sarvam_mla/modular_sarvam_mla.py @@ -11,13 +11,15 @@ # 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. -"""SarvamMLA model configuration""" + +from huggingface_hub.dataclasses import strict from ...utils import auto_docstring from ..deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config @auto_docstring(checkpoint="sarvamai/sarvam-105b") +@strict(accept_kwargs=True) class SarvamMLAConfig(DeepseekV3Config): r""" n_group (`int`, *optional*, defaults to 16): @@ -42,103 +44,18 @@ class SarvamMLAConfig(DeepseekV3Config): model_type = "sarvam_mla" - def __init__( - self, - vocab_size=262144, - hidden_size=4096, - intermediate_size=16384, - moe_intermediate_size=2048, - num_hidden_layers=32, - num_attention_heads=64, - num_key_value_heads=None, - n_shared_experts=1, - n_routed_experts=128, - routed_scaling_factor=2.5, - kv_lora_rank=512, - q_lora_rank=None, - qk_rope_head_dim=64, - v_head_dim=128, - qk_nope_head_dim=128, - n_group=16, - topk_group=2, - num_experts_per_tok=8, - first_k_dense_replace=1, - norm_topk_prob=True, - hidden_act="silu", - max_position_embeddings=4096, - initializer_range=0.006, - rms_norm_eps=1e-6, - use_cache=True, - pad_token_id=0, - bos_token_id=None, - eos_token_id=1, - pretraining_tp=1, - tie_word_embeddings=False, - rope_parameters=None, - rope_interleave=True, - attention_bias=False, - attention_dropout=0.0, - **kwargs, - ): - # Hub config.json uses num_experts/num_shared_experts; map to parent names - n_routed_experts = kwargs.pop("num_experts", n_routed_experts) - n_shared_experts = kwargs.pop("num_shared_experts", n_shared_experts) - - # head_dim in Hub config.json is kv_lora_rank + qk_rope_head_dim (for vLLM - # MLA compat), but DeepseekV3Config computes it as qk_rope_head_dim. - kwargs.pop("head_dim", None) - kwargs.pop("q_head_dim", None) - - super().__init__( - vocab_size=vocab_size, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - moe_intermediate_size=moe_intermediate_size, - num_hidden_layers=num_hidden_layers, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - n_shared_experts=n_shared_experts, - n_routed_experts=n_routed_experts, - routed_scaling_factor=routed_scaling_factor, - kv_lora_rank=kv_lora_rank, - q_lora_rank=q_lora_rank, - qk_rope_head_dim=qk_rope_head_dim, - v_head_dim=v_head_dim, - qk_nope_head_dim=qk_nope_head_dim, - n_group=n_group, - topk_group=topk_group, - num_experts_per_tok=num_experts_per_tok, - first_k_dense_replace=first_k_dense_replace, - norm_topk_prob=norm_topk_prob, - hidden_act=hidden_act, - max_position_embeddings=max_position_embeddings, - initializer_range=initializer_range, - rms_norm_eps=rms_norm_eps, - use_cache=use_cache, - pad_token_id=pad_token_id, - bos_token_id=bos_token_id, - eos_token_id=eos_token_id, - pretraining_tp=pretraining_tp, - tie_word_embeddings=tie_word_embeddings, - rope_parameters=rope_parameters, - rope_interleave=rope_interleave, - attention_bias=attention_bias, - attention_dropout=attention_dropout, - **kwargs, - ) - - def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation: set | None = None, **kwargs): - # The Hub config uses "deepseek_yarn" as rope type; normalize to "yarn" - # which is the standard type in ROPE_INIT_FUNCTIONS. - rope_scaling = kwargs.get("rope_scaling", None) - if rope_scaling is not None and rope_scaling.get("type") == "deepseek_yarn": - kwargs["rope_scaling"] = dict(rope_scaling) - kwargs["rope_scaling"]["type"] = "yarn" - if self.rope_parameters and self.rope_parameters.get("type") == "deepseek_yarn": - self.rope_parameters["type"] = "yarn" - return super().convert_rope_params_to_dict( - ignore_keys_at_rope_validation=ignore_keys_at_rope_validation, **kwargs - ) + vocab_size: int = 262144 + hidden_size: int = 4096 + intermediate_size: int = 16384 + num_hidden_layers: int = 32 + num_attention_heads: int = 64 + num_key_value_heads: int | None = None + n_routed_experts: int = 128 + q_lora_rank: int | None = None + n_group: int | None = 16 + topk_group: int | None = 2 + first_k_dense_replace: int | None = 1 + initializer_range: float = 0.006 __all__ = ["SarvamMLAConfig"] From 926004bbb5a9e6bdf972c1177592068daa1acfc9 Mon Sep 17 00:00:00 2001 From: vasqu Date: Tue, 17 Mar 2026 18:59:49 +0100 Subject: [PATCH 6/8] add date --- docs/source/en/model_doc/sarvam_mla.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/en/model_doc/sarvam_mla.md b/docs/source/en/model_doc/sarvam_mla.md index 1190f17725a1..6ebc0bfe3c3d 100644 --- a/docs/source/en/model_doc/sarvam_mla.md +++ b/docs/source/en/model_doc/sarvam_mla.md @@ -10,6 +10,7 @@ distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, e See the License for the specific language governing permissions and limitations under the License. --> +*This model was released on 2026-03-06 and added to Hugging Face Transformers on 2026-03-17.* # SarvamMLA From 3d99694399be7e99f164bab08923806170baad4c Mon Sep 17 00:00:00 2001 From: vasqu Date: Tue, 17 Mar 2026 19:16:13 +0100 Subject: [PATCH 7/8] add exception there --- utils/check_config_attributes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/utils/check_config_attributes.py b/utils/check_config_attributes.py index 64ef814f0b48..be07b36ae1bb 100644 --- a/utils/check_config_attributes.py +++ b/utils/check_config_attributes.py @@ -138,6 +138,7 @@ "GptOssConfig": True, "LwDetrConfig": True, "NemotronHConfig": True, + "SarvamMLAConfig": True, # Uses DeepseekV3 under the hood } # Common and important attributes, even if they do not always appear in the modeling files (can be a regex pattern) From 9ce210d9f1f761f8858df1c091237c35d4807299 Mon Sep 17 00:00:00 2001 From: vasqu Date: Wed, 18 Mar 2026 14:46:03 +0100 Subject: [PATCH 8/8] fix --- .../models/sarvam_mla/configuration_sarvam_mla.py | 15 --------------- .../models/sarvam_mla/modular_sarvam_mla.py | 3 +++ 2 files changed, 3 insertions(+), 15 deletions(-) diff --git a/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py b/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py index f5701edc656d..e4ee4061b338 100644 --- a/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py +++ b/src/transformers/models/sarvam_mla/configuration_sarvam_mla.py @@ -114,20 +114,5 @@ def __post_init__(self, **kwargs): self.head_dim = self.qk_rope_head_dim super().__post_init__(**kwargs) - def convert_rope_params_to_dict(self, **kwargs): - rope_scaling = kwargs.pop("rope_scaling", None) - self.rope_parameters = rope_scaling or self.rope_parameters - self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {} - - # Standardize and validate the correctness of rotary position embeddings parameters - self.rope_parameters.setdefault("rope_theta", kwargs.pop("rope_theta", self.default_theta)) - self.standardize_rope_params() - - # Convert to float because RoPE fn expect a float. Models on the hub were saved as int - for key in ["beta_fast", "beta_slow", "factor"]: - if key in self.rope_parameters: - self.rope_parameters[key] = float(self.rope_parameters[key]) - return kwargs - __all__ = ["SarvamMLAConfig"] diff --git a/src/transformers/models/sarvam_mla/modular_sarvam_mla.py b/src/transformers/models/sarvam_mla/modular_sarvam_mla.py index e092dbb2764b..a2a18e8f85fe 100644 --- a/src/transformers/models/sarvam_mla/modular_sarvam_mla.py +++ b/src/transformers/models/sarvam_mla/modular_sarvam_mla.py @@ -57,5 +57,8 @@ class SarvamMLAConfig(DeepseekV3Config): first_k_dense_replace: int | None = 1 initializer_range: float = 0.006 + def convert_rope_params_to_dict(self, **kwargs): + raise AttributeError("No BC behavior needed!") + __all__ = ["SarvamMLAConfig"]