Repository: huggingface/transformers
Pull request: Fix non-deterministic output in MoE experts with batched_mm/grouped_mm

# What does this PR do?

Fixes non-deterministic output when using `do_sample=False` with MoE models and `experts_implementation='batched_mm'` or `'grouped_mm'`.

The issue is `index_add_` with duplicate indices on CUDA. When `top_k > 1`, `token_idx` contains duplicates (`[0,0,1,1,2,2,...]`) and CUDA's `atomicAdd` doesn't guarantee operation order. Since FP addition isn't associative, this causes different rounding errors on each run.

The fix replaces `index_add_` with `view().sum()` - this works because each token appears exactly `top_k` times in sequence, so we can reshape and sum deterministically.

Fixes #43360

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] 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. #43360
- [ ] Did you make sure to update the documentation with your changes?
- [ ] Did you write any new necessary tests?

## Who can review?

@IlyasMoutawwakil

Changed files:
filename: src/transformers/integrations/moe.py
patch: @@ -71,8 +71,8 @@ def batched_mm_experts_forward(
     device = hidden_states.device
     num_top_k = top_k_index.size(-1)
     num_tokens = hidden_states.size(0)
+    hidden_dim = hidden_states.size(-1)
     num_experts = self.gate_up_proj.size(0)
-    final_hidden_states = torch.zeros_like(hidden_states)
 
     # Flatten top_k_index to get expert_ids per selected sample
     expert_ids = top_k_index.reshape(-1)
@@ -117,10 +117,11 @@ def batched_mm_experts_forward(
     # Apply routing weights
     out_per_sample = out_per_sample * sample_weights.unsqueeze(-1)  # (S, hidden_dim)
 
-    # Accumulate results back to the final_hidden_states using original token indices
-    final_hidden_states.index_add_(0, token_idx, out_per_sample.to(final_hidden_states.dtype))
+    # 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)
 
-    return final_hidden_states
+    return final_hidden_states.to(hidden_states.dtype)
 
 
 def grouped_mm_experts_forward(
@@ -137,8 +138,8 @@ def grouped_mm_experts_forward(
     device = hidden_states.device
     num_top_k = top_k_index.size(-1)
     num_tokens = hidden_states.size(0)
+    hidden_dim = hidden_states.size(-1)
     num_experts = self.gate_up_proj.size(0)
-    final_hidden_states = torch.zeros_like(hidden_states)
 
     # Flatten top_k_index to get expert_ids per selected sample
     expert_ids = top_k_index.reshape(-1)
@@ -199,10 +200,11 @@ def grouped_mm_experts_forward(
     # Restore original order
     out_per_sample = out_per_sample_g[inv_perm]
 
-    # Accumulate results back to the final_hidden_states using original token indices
-    final_hidden_states.index_add_(0, token_idx, out_per_sample.to(final_hidden_states.dtype))
+    # 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)
 
-    return final_hidden_states
+    return final_hidden_states.to(hidden_states.dtype)
 
 
 class ExpertsInterface(GeneralInterface):