Skip the expert-parallel sentinel masking when expert parallelism is off

src/transformers/distributed/tensor_parallel.py

@@ -144,7 +144,7 @@ def context_around_forward(self, module, mesh):
     def transform_output_post_forward(self, module, output, mesh):
         return output
 
-    def install_forward(self, module, mesh):
+    def install_forward(self, module, mesh, *, is_expert_parallel=False):
         """Install pre / around / post transforms by replacing module.forward."""
         original_forward = module.forward
 
@@ -344,7 +344,7 @@ class ReplicatedWithGradAllReduce(TensorParallelLayer):
     summed across the mesh.
     """
 
-    def install_forward(self, module, mesh):
+    def install_forward(self, module, mesh, *, is_expert_parallel=False):
         # A module hook rather than `param.register_hook`: params are replaced during weight
         # loading, which happens after TP is applied, and would drop a param-level hook.
         def _all_reduce_grads(mod, grad_input, grad_output):
@@ -415,7 +415,7 @@ def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = True):
         self.sequence_dim = sequence_dim
         self.use_local_output = use_local_output
 
-    def install_forward(self, module, mesh):
+    def install_forward(self, module, mesh, *, is_expert_parallel=False):
         # Replicate the module's params (LayerNorm/RMSNorm ones-init → from_local is safe).
         for p_name, p in list(module.named_parameters(recurse=False)):
             module.register_parameter(
@@ -606,6 +606,7 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_
 
     def install_forward(self, module, mesh, *, is_expert_parallel=False):
         """Install the transforms but pass `is_expert_parallel` in the forward call."""
+        module.is_expert_parallel = is_expert_parallel
         original_forward = module.forward
         output_source = (
             Partial()
@@ -816,7 +817,9 @@ def apply_tensor_parallelism(model, tp_mesh):
                 # MLA needs to know the qk_rope_head_dim to split the projection output into KV and RoPE parts.
                 # TODO: Store qk_rope_head_dim on MLA projection modules when the models initialize them.
                 module.config = model.config.get_text_config()
-            ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh)
+            ALL_PARALLEL_STYLES[style_name].install_forward(
+                module, tp_mesh, is_expert_parallel=model.config.distributed_config.enable_expert_parallel
+            )
         module._is_hooked = True
 
     return model

src/transformers/integrations/moe.py

@@ -417,8 +417,10 @@ def grouped_mm_experts_forward(
     # In-place clamp on `expert_ids_g` keeps the per-row bias gather in-bounds (bias added at
     # sentinel positions falls in rows the kernel skips, so harmless). Safe to mutate now —
     # nothing downstream needs the sentinel info from `expert_ids_g` itself.
-    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
-    expert_ids_g.clamp_(max=self.num_experts - 1)
+    sentinel_mask = None
+    if self.is_expert_parallel:
+        sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
+        expert_ids_g.clamp_(max=self.num_experts - 1)
 
     # Select expert weights and biases
     # NOTE: We keep all experts here and rely on offsets to target the active ones.
@@ -434,7 +436,8 @@ def grouped_mm_experts_forward(
         selected_biases = self.up_proj_bias[expert_ids_g] if self.has_bias else None
 
     # Pre-mask (bwd path).
-    selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)
+    if sentinel_mask is not None:
+        selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)
 
     # --- Up projection per expert (grouped) ---
     proj_out = _grouped_linear(
@@ -462,7 +465,8 @@ def grouped_mm_experts_forward(
     weighted_out = proj_out * sample_weights_g.unsqueeze(-1)  # (S, hidden_dim)
 
     # Post-mask (fwd path).
-    weighted_out.masked_fill_(sentinel_mask, 0.0)
+    if sentinel_mask is not None:
+        weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Restore original order
     inv_perm = torch.empty_like(perm)
@@ -563,6 +567,7 @@ def __init__(self, config, *args, **kwargs):
             self.has_bias = has_bias
             self.is_transposed = is_transposed
             self.is_concatenated = is_concatenated
+            self.is_expert_parallel = False
 
         @wraps(original_forward)
         def forward(self, *args, **kwargs):