Repository: huggingface/transformers
Pull request: Support non-gated experts

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (issue)


## 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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? 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](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Changed files:
filename: src/transformers/integrations/moe.py
patch: @@ -118,40 +118,50 @@ def batched_mm_experts_forward(
 
     # Handle invalid expert IDs from Expert Parallelism (EP)
     # When EP is enabled, tokens assigned to experts on other devices are marked with sentinel value >= num_experts
-    valid_mask = expert_ids < self.num_experts
-    expert_ids_clamped = expert_ids.clamp(0, self.num_experts - 1)
+    invalid_mask = expert_ids >= self.num_experts
+    expert_ids = expert_ids.clamp(0, self.num_experts - 1)
 
     # Get current hidden states for selected samples
     selected_hidden_states = hidden_states[token_idx]
 
-    # Select expert weights and biases for selected samples (using clamped IDs for safe indexing)
-    selected_gate_up = self.gate_up_proj[expert_ids_clamped]
-    selected_down = self.down_proj[expert_ids_clamped]
-    selected_gate_up_bias = self.gate_up_proj_bias[expert_ids_clamped] if self.has_bias else None
-    selected_down_bias = self.down_proj_bias[expert_ids_clamped] if self.has_bias else None
+    # Select gate_up or just up projection weights and biases
+    if self.has_gate:
+        selected_weights = self.gate_up_proj[expert_ids]
+        selected_biases = self.gate_up_proj_bias[expert_ids] if self.has_bias else None
+    else:
+        selected_weights = self.up_proj[expert_ids]
+        selected_biases = self.up_proj_bias[expert_ids] if self.has_bias else None
 
     # --- Up projection per expert (batched) ---
-    gate_up_out = _batched_linear(
-        selected_hidden_states, selected_gate_up, selected_gate_up_bias, is_transposed=self.is_transposed
-    )  # (S, 2 * intermediate_dim)
+    up_proj_out = _batched_linear(
+        selected_hidden_states, selected_weights, bias=selected_biases, is_transposed=self.is_transposed
+    )  # (S, 2 * intermediate_dim) or  (S, intermediate_dim) depending on whether we have gating
+
+    # Apply gating or just activation
+    if self.has_gate:
+        up_proj_out = self._apply_gate(up_proj_out)  # (S, intermediate_dim)
+    else:
+        # for non-gated experts we just apply the activation function
+        up_proj_out = self.act_fn(up_proj_out)  # (S, intermediate_dim)
 
-    # Apply gating
-    gated_out = self._apply_gate(gate_up_out)  # (S, intermediate_dim)
+    # Select down projection weights and biases for selected samples
+    selected_weights = self.down_proj[expert_ids]
+    selected_biases = self.down_proj_bias[expert_ids] if self.has_bias else None
 
     # --- Down projection per expert (batched) ---
-    out_per_sample = _batched_linear(
-        gated_out, selected_down, selected_down_bias, is_transposed=self.is_transposed
+    down_proj_out = _batched_linear(
+        up_proj_out, selected_weights, bias=selected_biases, is_transposed=self.is_transposed
     )  # (S, hidden_dim)
 
     # Apply routing weights and zero out invalid expert contributions
-    if sample_weights.shape != expert_ids_clamped.shape:
-        sample_weights = sample_weights.gather(0, expert_ids_clamped)
-    out_per_sample = out_per_sample * sample_weights.unsqueeze(-1)  # (S, hidden_dim)
-    out_per_sample = out_per_sample * valid_mask.unsqueeze(-1).to(out_per_sample.dtype)
+    weighted_out = down_proj_out * sample_weights.unsqueeze(-1)  # (S, hidden_dim)
+    weighted_out.masked_fill_(invalid_mask.unsqueeze(-1), 0.0)  # Zero out invalid expert contributions
 
     # Accumulate results using deterministic reshape+sum instead of index_add_
-    # (index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd)
-    final_hidden_states = out_per_sample.view(num_tokens, num_top_k, hidden_dim).sum(dim=1)
+    # index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd
+    # index_add_ accumulates in-place using the dtype of the output tensor (fp16/bf16)
+    # reshape+sum accumulates in fp32 which is more stable for low precision training/inference.
+    final_hidden_states = weighted_out.view(num_tokens, num_top_k, hidden_dim).sum(dim=1)
 
     return final_hidden_states.to(hidden_states.dtype)
 
@@ -352,53 +362,57 @@ def grouped_mm_experts_forward(
     sample_weights_g = sample_weights[perm]
     selected_hidden_states_g = selected_hidden_states[perm]
 
-    # Select expert weights and biases for selected samples
-    # NOTE: We keep all experts here and rely on offsets to target the active ones.
-    # I have already implemented a version that only passes the active experts, but
-    # to do so I had to use torch.unique which breaks the graph capture (data-dependent).
-    # Also there were no speedup gains from it in my experiments, even in eager mode.
-    selected_gate_up = self.gate_up_proj
-    selected_down = self.down_proj
-    selected_gate_up_bias = self.gate_up_proj_bias[expert_ids_g] if self.has_bias else None
-    selected_down_bias = self.down_proj_bias[expert_ids_g] if self.has_bias else None
-
     # Compute offsets for grouped_mm
     # using histc instead of bincount to avoid cuda graph issues
     # With deterministic algorithms, CPU only supports float input, CUDA only supports int input.
     histc_input = expert_ids_g.float() if device.type == "cpu" else expert_ids_g.int()
     num_tokens_per_expert = torch.histc(histc_input, bins=self.num_experts, min=0, max=self.num_experts - 1)
     offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32)
 
+    # Select expert weights and biases
+    # NOTE: We keep all experts here and rely on offsets to target the active ones.
+    # I have already implemented a version that only passes the active experts, but
+    # to do so I had to use torch.unique which breaks the graph capture (data-dependent).
+    # Also there were no speedup gains from it in my experiments, even in eager mode.
+    if self.has_gate:
+        selected_weights = self.gate_up_proj
+        selected_biases = self.gate_up_proj_bias[expert_ids_g] if self.has_bias else None
+    else:
+        selected_weights = self.up_proj
+        selected_biases = self.up_proj_bias[expert_ids_g] if self.has_bias else None
+
     # --- Up projection per expert (grouped) ---
-    gate_up_out = _grouped_linear(
-        selected_hidden_states_g,
-        selected_gate_up,
-        offs=offsets,
-        bias=selected_gate_up_bias,
-        is_transposed=self.is_transposed,
-    )  # (S, 2 * intermediate_dim)
+    up_proj_out = _grouped_linear(
+        selected_hidden_states_g, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed
+    )  # (S, 2 * intermediate_dim) or  (S, intermediate_dim) depending on whether we have gating
 
-    # Apply gating
-    gated_out = self._apply_gate(gate_up_out)  # (S, intermediate_dim)
+    # Apply gating or just activation
+    if self.has_gate:
+        up_proj_out = self._apply_gate(up_proj_out)  # (S, intermediate_dim)
+    else:
+        # for non-gated experts we just apply the activation function
+        up_proj_out = self.act_fn(up_proj_out)  # (S, intermediate_dim)
+
+    # Select down projection weights and biases
+    selected_weights = self.down_proj
+    selected_biases = self.down_proj_bias[expert_ids_g] if self.has_bias else None
 
     # --- Down projection per expert (grouped) ---
-    out_per_sample_g = _grouped_linear(
-        gated_out,
-        selected_down,
-        offs=offsets,
-        bias=selected_down_bias,
-        is_transposed=self.is_transposed,
+    down_proj_out = _grouped_linear(
+        up_proj_out, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed
     )  # (S, hidden_dim)
 
     # Apply routing weights
-    out_per_sample_g = out_per_sample_g * sample_weights_g.unsqueeze(-1)  # (S, hidden_dim)
+    weighted_out = down_proj_out * sample_weights_g.unsqueeze(-1)  # (S, hidden_dim)
 
     # Restore original order
-    out_per_sample = out_per_sample_g[inv_perm]
+    weighted_out = weighted_out[inv_perm]  # (S, hidden_dim)
 
     # Accumulate results using deterministic reshape+sum instead of index_add_
-    # (index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd)
-    final_hidden_states = out_per_sample.view(num_tokens, num_top_k, hidden_dim).sum(dim=1)
+    # index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd
+    # index_add_ accumulates in-place using the dtype of the output tensor (fp16/bf16)
+    # reshape+sum accumulates in fp32 which is more stable for low precision training/inference.
+    final_hidden_states = weighted_out.view(num_tokens, num_top_k, hidden_dim).sum(dim=1)
 
     return final_hidden_states.to(hidden_states.dtype)
 
@@ -444,7 +458,11 @@ def _default_apply_gate(self, gate_up_out: torch.Tensor) -> torch.Tensor:
 
 
 def use_experts_implementation(
-    experts_class: type[torch.nn.Module] | None = None, *, is_transposed: bool = False, has_bias: bool = False
+    experts_class: type[torch.nn.Module] | None = None,
+    *,
+    is_transposed: bool = False,
+    has_bias: bool = False,
+    has_gate: bool = True,
 ) -> type[torch.nn.Module]:
     """Decorator to modify experts class to support different experts implementations.
 
@@ -468,6 +486,7 @@ def wrapper(experts_class: type[torch.nn.Module]) -> type[torch.nn.Module]:
         def __init__(self, config, *args, **kwargs):
             original_init(self, config, *args, **kwargs)
             self.config = config
+            self.has_gate = has_gate
             self.has_bias = has_bias
             self.is_transposed = is_transposed
 
@@ -480,6 +499,7 @@ def forward(self, *args, **kwargs):
 
         if not hasattr(experts_class, "_apply_gate"):
             experts_class._apply_gate = _default_apply_gate
+
         experts_class.__init__ = __init__
         experts_class.forward = forward
         return experts_class

filename: src/transformers/integrations/tensor_parallel.py
patch: @@ -1143,7 +1143,9 @@ class ParallelInterface(GeneralInterface):
 # =============================================================================
 
 
-def gather_full_tensor(local_tensor: torch.Tensor, shard_dim: int, device_mesh) -> torch.Tensor:
+def gather_full_tensor(
+    local_tensor: torch.Tensor, shard_dim: int, device_mesh: dist.device_mesh.DeviceMesh
+) -> torch.Tensor:
     """
     All-gather a sharded tensor along the specified dimension to reconstruct the full tensor.
 
@@ -1157,7 +1159,7 @@ def gather_full_tensor(local_tensor: torch.Tensor, shard_dim: int, device_mesh)
     """
     world_size = device_mesh.size()
     # In case of TP+DP configuration, the TP group should be used for gathering, not the full DP group
-    process_group = device_mesh.get_group("tp") if "tp" in device_mesh.mesh_dim_names else None
+    process_group = device_mesh.get_group("tp") if "tp" in (device_mesh.mesh_dim_names or {}) else None
 
     # Normalize negative dimension
     if shard_dim < 0:

filename: tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py
patch: @@ -136,10 +136,6 @@ def test_assisted_decoding_matches_greedy_search(self, assistant_type):
     def test_assisted_decoding_sample(self):
         pass
 
-    @unittest.skip("Requires torch>=2.9.0 for grouped MM")
-    def test_eager_matches_batched_and_grouped_inference(self):
-        pass
-
     @unittest.skip("DSA indexer mask shape mismatch with static cache")
     def test_generate_from_inputs_embeds_with_static_cache(self):
         pass