Add support for BharatGen's Param2MoE model architecture by bhargav-patel-29 · Pull Request #46395 · huggingface/transformers

bhargav-patel-29 · 2026-06-04T04:28:33Z

What does this PR do?

This PR adds support for Param-2-17B-MoE-A2.4B, a large-scale Mixture-of-Experts (MoE) causal language model.

Param-2-17B-MoE-A2.4B uses a Hybrid Dense + MoE architecture with 17B total parameters while activating only 2.4B parameters per token, enabling high model capacity with efficient inference cost.

The model is pretrained from scratch with strong multilingual capabilities and particular emphasis on linguistic diversity and Indian language representation. It is released as a pretrained base model intended for downstream fine-tuning.

This integration enables seamless loading, inference, and downstream fine-tuning using standard Transformers APIs.

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Text models: @ArthurZucker @Cyrilvallez

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

ArthurZucker

TY looks good already let's standardize a bit more!

ArthurZucker · 2026-06-09T07:42:17Z

+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # Keep a reference for the shared expert residual addition
+        residual = hidden_states
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+
+        # Flatten to [tokens, hidden_size] for the router and experts
+        hidden_states_flat = hidden_states.view(-1, hidden_dim)
+
+        # Router: logits are captured automatically by OutputRecorder;
+        # we only consume the weights and indices here.
+        _, routing_weights, selected_experts = self.gate(hidden_states_flat)
+
+        # Routed experts: sparse computation over selected_experts
+        routed_output = self.experts(hidden_states_flat, selected_experts, routing_weights)
+        routed_output = routed_output.view(batch_size, seq_len, hidden_dim)
+
+        # Shared expert: dense computation on all tokens, added to routed output
+        hidden_states = routed_output + self.shared_experts(residual)
+        return hidden_states
+


this can be entirely inherited!

ArthurZucker · 2026-06-09T07:43:13Z

+        # ---- 1. Compute raw logits in fp32 --------------------------------
+        router_logits = F.linear(hidden_states.float(), self.weight.float())
+        # router_logits: [tokens, num_experts]
+
+        # ---- 2. Compute unbiased scores (used as output weights) ----------
+        if self.score_function == "sigmoid":
+            scores = torch.sigmoid(router_logits)
+        else:  # "softmax"
+            scores = torch.softmax(router_logits, dim=-1)
+
+        # ---- 3. Compute biased scores (used for routing decision only) ----
+        if self.expert_bias is not None:
+            scores_for_routing = scores + self.expert_bias.float()
+        else:
+            scores_for_routing = scores
+
+        # ---- 4. Group-limited top-k (no-op when n_group == 1) -------------
+        if self.n_group > 1 and self.topk_group < self.n_group:
+            num_tokens = hidden_states.shape[0]
+            experts_per_group = self.num_experts // self.n_group
+
+            # Max score per group → select topk_group best groups
+            group_scores = scores_for_routing.view(num_tokens, self.n_group, experts_per_group).max(dim=-1).values
+            group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+
+            # Build a mask that zeros out experts in non-selected groups
+            group_mask = torch.zeros_like(group_scores)
+            group_mask.scatter_(1, group_idx, 1)
+            score_mask = (
+                group_mask.unsqueeze(-1).expand(num_tokens, self.n_group, experts_per_group).reshape(num_tokens, -1)
+            )
+            scores_for_routing = scores_for_routing.masked_fill(~score_mask.bool(), 0.0)
+
+        # ---- 5. Select top-k experts using biased scores ------------------
+        _, topk_idx = torch.topk(scores_for_routing, k=self.top_k, dim=-1, sorted=False)
+
+        # ---- 6. Gather *unbiased* scores as actual output weights ---------
+        topk_weights = scores.gather(1, topk_idx)
+
+        # ---- 7. Normalize and scale ---------------------------------------
+        if self.norm_topk_prob:
+            topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-9)
+
+        topk_weights = (topk_weights * self.routed_scaling_factor).to(hidden_states.dtype)
+
+        return router_logits, topk_weights, topk_idx


we need to remove code paths and collapse to only parameters used in models that are released!
this will standardize and allow inheritance in most cases, or will isolate the difference this model vs other models 😉

Thanks for the review @ArthurZucker!

Addressed both points:

Param2MoERouter - simplified to only the paths used by the released model: sigmoid routing is hardcoded (removed softmax), expert_bias is always present, and the unused group-routing logic (n_group=1) was removed.
Param2MoESparseMoeBlock - now inherits from DeepseekV2Moe since the shared-expert + residual pattern is identical. init swaps in Param2MoERouter and Param2 config names, while forward remains a minimal override because the parent implementation depends on self.gate.weight and route_tokens_to_experts(), which don't match our router abstraction.

ArthurZucker

also this only include the LLM part right?

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

…kV2Moe Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

github-actions · 2026-06-11T08:40:31Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46395&sha=23356e

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

bhargav-patel-29 · 2026-07-03T06:21:28Z

Hey @ArthurZucker the test_tp_generation_quantized failure on Param2MoEModelTest (37% match, threshold 75%) is something I'd like your input on before merging.

What we found:
After layer-by-layer instrumentation, the issue originates in the attention block, not the MoE routing. Float8 weight-only TP quantization introduces a consistent ~12% relative error in attn_out hidden states. This cascades under greedy decoding once one token flips, all subsequent tokens diverge (autoregressive amplification), hence the low match ratio.

Key facts confirming this is not a TP correctness bug:

test_tp_forward, test_tp_backward, test_tp_generation all pass
Weight sharding is verified correct (shapes confirmed via forward hooks)
Router picks identical expert indices on TP and non-TP paths
The q_norm/k_norm TP plan entries are copied directly from Qwen3/Qwen3MoE (which pass this test), but Param2MoE's test config has fewer heads, giving worse relative quantization error per rank

I saw the NOTE(3outeille) comment in test_tensor_parallel_mixin.py acknowledging this is model-dependent. My proposed fix is to override test_tp_generation_quantized in Param2MoEModelTest with a relaxed threshold of 0.30 (above the observed 37%, well below what a real TP bug would produce).

Is this the right approach, or would you prefer a different strategy (e.g. exposing match_threshold as a parameter in _test_tp_generation_quantized_impl so model-specific overrides don't need to copy the whole function)?

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

github-actions · 2026-07-06T09:04:18Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: auto, param2moe

github-actions · 2026-07-06T09:29:24Z

CI recap

Dashboard: View test results in Grafana
Latest run: 28780172628:2
Result: failure | Jobs: 15 | Tests: 171,614 | Failures: 6 | Duration: 25h 6m

bhargav-patel-29 added 8 commits June 4, 2026 09:03

Add Param2MoE model support to Transformers

f744cbb

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

docs: update Param2MoE model documentation

2500d07

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

fix ruff linting issues

bab2fd4

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Apply Ruff formatting

1b90307

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

fix param2moe documentation

62f48f9

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

fix Param2MoEPreTrainedModel to use Param2MoERouter

7e6adc8

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

fix test file for param2moe

454c8fc

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

fix modular file for param2moe architecture

4071273

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

ArthurZucker reviewed Jun 9, 2026

View reviewed changes

bhargav-patel-29 added 7 commits June 10, 2026 10:30

Update modular and test file for param2moe

a766cd5

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

c843363

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

8b254cf

Ruff format modular and test file for param2moe

4cfde53

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Update correct date in param2moe doc file

36a764f

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

fix: collapse Param2MoERouter and inherit SparseMoeBlock from Deepsee…

f696efb

…kV2Moe Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Fix: remove unused attributes and standardize Param2MoE modular file

23356e5

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

bhargav-patel-29 added 3 commits June 22, 2026 14:28

Refactored the modular file to inherit entirely from the base class

6f49ff0

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Merge branch 'main' into main

2fa41c0

Merge branch 'main' into main

eb2d4a0

bhargav-patel-29 added 7 commits July 3, 2026 11:51

Merge branch 'main' into main

bd24b0b

Update test file for param2moe architecture

5294de7

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Ruff format test file for param2moe arch.

f723204

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Add base_model_ep_plan for param2moe

76de493

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Update auto_mappings.py

012b405

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Update attribute_map in configuration_param2moe

d37ea9d

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

update date in param2moe.md

b33f5c2

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

bhargav-patel-29 force-pushed the main branch from 687a0d9 to b33f5c2 Compare July 6, 2026 08:15

fix(param2moe): refactor router to return 3-tuple for EP compatibility

ddbd228

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add support for BharatGen's Param2MoE model architecture#46395

Add support for BharatGen's Param2MoE model architecture#46395
bhargav-patel-29 wants to merge 26 commits into
huggingface:mainfrom
Bharatgen-Tech:main

bhargav-patel-29 commented Jun 4, 2026 •

edited by github-actions Bot

Loading

Uh oh!

ArthurZucker left a comment

Uh oh!

ArthurZucker Jun 9, 2026

Uh oh!

ArthurZucker Jun 9, 2026

Uh oh!

bhargav-patel-29 Jun 10, 2026

Uh oh!

ArthurZucker left a comment

Uh oh!

github-actions Bot commented Jun 11, 2026

Uh oh!

bhargav-patel-29 commented Jul 3, 2026

Uh oh!

github-actions Bot commented Jul 6, 2026

Uh oh!

github-actions Bot commented Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

bhargav-patel-29 commented Jun 4, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

ArthurZucker Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

ArthurZucker Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

bhargav-patel-29 Jun 10, 2026

Choose a reason for hiding this comment

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jun 11, 2026

Uh oh!

bhargav-patel-29 commented Jul 3, 2026

Uh oh!

github-actions Bot commented Jul 6, 2026

Uh oh!

github-actions Bot commented Jul 6, 2026

CI recap

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

bhargav-patel-29 commented Jun 4, 2026 •

edited by github-actions Bot

Loading