Add support for BharatGen's Param2MoE model architecture

docs/source/en/_toctree.yml

@@ -793,6 +793,8 @@
         title: OpenAI Privacy Filter
       - local: model_doc/opt
         title: OPT
+      - local: model_doc/param2moe
+        title: Param2MoE
       - local: model_doc/pegasus
         title: Pegasus
       - local: model_doc/pegasus_x

docs/source/en/model_doc/param2moe.md

@@ -0,0 +1,114 @@
+<!--Copyright 2026 the HuggingFace 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.
+
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer.
+
+-->
+*This model was contributed to Hugging Face Transformers on 2026-02-10.*
+
+# Param2MoE
+
+## Overview
+
+Param2MoE was released by the [BharatGen AI](https://huggingface.co/bharatgenai) team as **Param-2-17B-MoE-A2.4B**, a Hybrid Mixture-of-Experts language model with 17B total parameters and only 2.4B active per token. It is pretrained from scratch on ~22 trillion tokens across two phases, with an emphasis on linguistic diversity — supporting English, Hindi, and 21 Indian languages. The model ships as an early post-training checkpoint with reasoning, tool calling, math, and code capabilities.
+
+The original model can be found [here](https://huggingface.co/bharatgenai/Param2-17B-A2.4B-Thinking).
+
+Tips:
+
+- **Expert routing**: 2 shared experts are always active alongside 6 dynamically routed experts selected from a pool of 64 per token.
+- **Memory**: Loading the full model requires ~34 GB VRAM in bfloat16.
+- **Context**: Maximum supported length is 4096 tokens.
+- **Custom code**: Pass `trust_remote_code=True` when loading with `AutoModelForCausalLM`.
+- **Decoding**: Set `skip_special_tokens=False` to preserve `<think>...</think>` reasoning tags in Thinking checkpoint outputs. Use `do_sample=False` for deterministic/evaluation runs.
+- **Safety**: The model has not undergone RLHF or safety alignment — fine-tune and evaluate before production use.
+
+## Usage examples
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from parsers import parse_model_output
+
+model_name = "bharatgenai/Param2-17B-A2.4B-Thinking"
+
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    trust_remote_code=False
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    device_map="auto"
+)
+
+conversation = [
+    {"role": "system", "content": "You are helpful assistant."},
+    {"role": "user", "content": "What is the BharatGen Mission?"}
+]
+
+inputs = tokenizer.apply_chat_template(
+    conversation=conversation,
+    return_tensors="pt",
+    add_generation_prompt=True
+).to(model.device)
+
+with torch.no_grad():
+    output = model.generate(
+        inputs,
+        max_new_tokens=300,
+        do_sample=True,
+        top_k=50,
+        top_p=0.9,
+        temperature=0.7,
+        eos_token_id=tokenizer.eos_token_id,
+        use_cache=False,
+    )
+
+generated_tokens = output[0][inputs.shape[-1]:]
+
+# IMPORTANT: skip_special_tokens=False
+generated_text = tokenizer.decode(
+    generated_tokens,
+    skip_special_tokens=False
+)
+
+parsed = parse_model_output(generated_text)
+
+print("\n========== RAW ==========\n", generated_text)
+print("\n========== REASONING ==========\n", parsed["reasoning"])
+print("\n========== TOOL CALLS ==========\n", parsed["tool_calls"])
+print("\n========== FINAL ANSWER ==========\n", parsed["final_answer"])
+```
+
+## Param2MoEConfig
+
+[[autodoc]] Param2MoEConfig
+
+## Param2MoEPreTrainedModel
+
+[[autodoc]] Param2MoEPreTrainedModel
+    - forward
+
+## Param2MoEModel
+
+[[autodoc]] Param2MoEModel
+    - forward
+
+## Param2MoEForCausalLM
+
+[[autodoc]] Param2MoEForCausalLM
+    - forward

src/transformers/models/__init__.py

@@ -319,6 +319,7 @@
     from .paddleocr_vl import *
     from .paligemma import *
     from .parakeet import *
+    from .param2moe import *
     from .patchtsmixer import *
     from .patchtst import *
     from .pe_audio import *

src/transformers/models/auto/auto_mappings.py

@@ -423,6 +423,7 @@
         ("parakeet_ctc", "ParakeetCTCConfig"),
         ("parakeet_encoder", "ParakeetEncoderConfig"),
         ("parakeet_tdt", "ParakeetTDTConfig"),
+        ("param2moe", "Param2MoEConfig"),
         ("patchtsmixer", "PatchTSMixerConfig"),
         ("patchtst", "PatchTSTConfig"),
         ("pe_audio", "PeAudioConfig"),

src/transformers/models/auto/modeling_auto.py

@@ -353,6 +353,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
         ("parakeet_ctc", "ParakeetForCTC"),
         ("parakeet_encoder", "ParakeetEncoder"),
         ("parakeet_tdt", "ParakeetForTDT"),
+        ("param2moe", "Param2MoEModel"),
         ("patchtsmixer", "PatchTSMixerModel"),
         ("patchtst", "PatchTSTModel"),
         ("pe_audio", "PeAudioModel"),
@@ -753,6 +754,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
         ("olmoe", "OlmoeForCausalLM"),
         ("openai-gpt", "OpenAIGPTLMHeadModel"),
         ("opt", "OPTForCausalLM"),
+        ("param2moe", "Param2MoEForCausalLM"),
         ("pegasus", "PegasusForCausalLM"),
         ("persimmon", "PersimmonForCausalLM"),
         ("phi", "PhiForCausalLM"),

src/transformers/models/param2moe/__init__.py

@@ -0,0 +1,28 @@
+# Copyright 2026 the HuggingFace 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_param2moe import *
+    from .modeling_param2moe import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

src/transformers/models/param2moe/configuration_param2moe.py

@@ -0,0 +1,172 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/param2moe/modular_param2moe.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_param2moe.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 the HuggingFace 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 huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="bharatgenai/Param2-17B-A2.4B-Thinking")
+@strict
+class Param2MoEConfig(PreTrainedConfig):
+    r"""
+    first_k_dense_replace (`int`, *optional*, defaults to 1):
+        Number of dense layers in the shallow layers before switching to MoE layers.
+    n_group (`int`, *optional*, defaults to 1):
+        Number of groups for routed experts.
+    use_bias (`bool`, *optional*, defaults to `False`):
+        Whether to use bias in linear layers. Also sets `mlp_bias` in `__post_init__`.
+    moe_shared_expert_intermediate_size (`int`, *optional*, defaults to 4096):
+        Intermediate size for the always-active shared expert MLP.
+    router_dtype (`str`, *optional*, defaults to `"fp32"`):
+        Data type used for router weight computation. Using float32 improves numerical
+        stability of the routing scores.
+    partial_rotary_factor (`float`, *optional*, defaults to 1.0):
+        Fraction of each attention head's dimension to apply rotary position embeddings
+        to. A value of 1.0 applies RoPE to the full head dimension.
+    num_nextn_predict_layers (`int`, *optional*, defaults to 0):
+        Number of next-n token prediction layers used for multi-token prediction (MTP).
+        Set to 0 to disable MTP.
+    mtp_loss_scaling_factor (`float`, *optional*, defaults to 0.0):
+        Scaling factor applied to the multi-token prediction auxiliary loss.
+        Set to 0.0 to disable the MTP loss contribution.
+    moe_router_enable_expert_bias (`bool`, *optional*, defaults to `True`):
+        Whether to add a per-expert learnable scalar bias to routing scores before
+        top-k selection. The bias affects routing decisions only; output weights
+        use unbiased scores to avoid distorting gradients.
+    output_dropout (`float`, *optional*, defaults to 0.0):
+        Dropout probability applied to layer outputs. Set to 0.0 to disable.
+    rope_scaling (`str` or `dict`, *optional*):
+        RoPE scaling configuration. Can be a preset name string or a dictionary
+        with scaling parameters (e.g. `{"type": "linear", "factor": 4.0}`).
+        Set to `None` to use standard (unscaled) RoPE.
+    rope_theta (`float`, *optional*, defaults to 1000000.0):
+        Base period (theta) for rotary position embeddings. Larger values extend
+        the effective context length.
+    score_function (`str`, *optional*, defaults to `"sigmoid"`):
+        Activation function used to convert router logits to routing scores.
+        `"sigmoid"` gives independent per-expert probabilities; `"softmax"` applies
+        competitive normalization across all experts.
+    torch_dtype (`str`, *optional*, defaults to `"bfloat16"`):
+        Default torch dtype for model weights when loading with `from_pretrained`.
+    use_rmsnorm (`bool`, *optional*, defaults to `True`):
+        Whether to use RMSNorm (Root Mean Square Layer Normalization) instead of
+        standard LayerNorm for all normalization layers.
+
+    Example:
+
+    ```python
+    >>> from transformers import Param2MoEModel, Param2MoEConfig
+    >>> # Initializing a Param2MoE style configuration
+    >>> configuration = Param2MoEConfig()
+    >>> # Accessing the model configuration
+    >>> model = Param2MoEModel(configuration)
+    >>> print(model.config)
+    ```
+    """
+
+    model_type = "param2moe"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.q_b_proj": "colwise",
+        "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
+        "layers.*.self_attn.kv_b_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "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"]),
+    }
+
+    vocab_size: int = 128008
+    hidden_size: int = 2048
+    intermediate_size: int = 9216
+    num_hidden_layers: int = 21
+    num_attention_heads: int = 32
+    num_key_value_heads: int | None = 8
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 4096
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int | None = 0
+    eos_token_id: int | list[int] | None = 3
+    tie_word_embeddings: bool = True
+    use_qkv_bias: bool = False
+    attention_dropout: float | None = 0.0
+    use_bias: bool = False
+    head_dim: int | None = 64
+    first_k_dense_replace: int = 1
+    n_group: int | None = 1
+    num_experts: int = 64
+    num_shared_experts: int = 2
+    routed_scaling_factor: float = 2.5
+    topk_group: int | None = 1
+    norm_topk_prob: bool | None = True
+    num_experts_per_tok: int | None = 6
+    moe_intermediate_size: int = 2048
+    moe_shared_expert_intermediate_size: int = 4096
+    rope_parameters: RopeParameters | dict | None = None
+    router_aux_loss_coef: float = 0.0
+    router_dtype: str = "fp32"
+    partial_rotary_factor: float = 1.0
+    max_window_layers: int = 20
+    num_nextn_predict_layers: int = 0
+    mtp_loss_scaling_factor: float = 0.0
+    moe_router_enable_expert_bias: bool = True
+    output_router_logits: bool = False
+    use_qk_norm: bool = True
+    output_dropout: float = 0.0
+    rope_scaling: str | dict | None = None
+    sliding_window: int | None = None
+    rope_theta: float = 1000000.0
+    score_function: str = "sigmoid"
+    torch_dtype: str = "bfloat16"
+    use_rmsnorm: bool = True
+
+    def __post_init__(self, **kwargs):
+        self.attention_bias = self.use_qkv_bias
+        self.mlp_bias = self.use_bias
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of `@strict`-powered validation. Validates the architecture of the config."""
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError(
+                f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention "
+                f"heads ({self.num_attention_heads})."
+            )
+
+
+__all__ = ["Param2MoEConfig"]