diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 6a9455637a33..3151f6264a12 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -541,6 +541,8 @@ def shard_param(self, module, param, mesh): f"{expert_parallel_size} expert-parallel ranks." ) module.num_experts = global_num_experts // expert_parallel_size + # The experts forward masks sentinel rows only when its experts are actually split. + module._is_expert_parallel = True module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, mesh, [self.placement], src_data_rank=None), requires_grad=meta.requires_grad, diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index ab266873f821..880526719c7e 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -504,6 +504,7 @@ def _dispatch_routed_input( num_experts: int, m_alignment: int, use_psum_layout: bool, + is_expert_parallel: bool, ) -> tuple: """Sort tokens by expert id and build the M-grouped padded layout. @@ -533,8 +534,10 @@ def _dispatch_routed_input( # keeps any per-row gather (e.g. bias) in-bounds — bias added at sentinel positions falls # in rows the kernel skips, so harmless. Safe to mutate now: the layout was built from the # unclamped tensor and nothing downstream needs the sentinel info from `expert_ids_g` itself. - sentinel_mask = (expert_ids_g >= num_experts).unsqueeze(-1) - expert_ids_g.clamp_(max=num_experts - 1) + sentinel_mask = None + if is_expert_parallel: + sentinel_mask = (expert_ids_g >= num_experts).unsqueeze(-1) + expert_ids_g.clamp_(max=num_experts - 1) return ( sorted_hidden_states_g, sample_weights_g, @@ -550,7 +553,7 @@ def _dispatch_routed_input( def _combine_routed_output( out_padded: torch.Tensor, sorted_weights: torch.Tensor, - sentinel_mask: torch.Tensor, + sentinel_mask: torch.Tensor | None, perm: torch.Tensor, sorted_to_padded: torch.Tensor, num_tokens: int, @@ -563,7 +566,8 @@ def _combine_routed_output( weighted = out * sorted_weights.to(out.dtype).unsqueeze(-1) # Sentinel rows past the valid expert blocks may carry NaN from allocator # reuse (`0 * NaN = NaN`); zero them so the top-k reduction stays finite. - weighted.masked_fill_(sentinel_mask, 0.0) + if sentinel_mask is not None: + weighted.masked_fill_(sentinel_mask, 0.0) inv_perm = torch.empty_like(perm) inv_perm[perm] = torch.arange(perm.size(0), device=out.device) # Deterministic reshape+sum (index_add_ with duplicates is non-deterministic on CUDA). @@ -646,7 +650,13 @@ def deepgemm_bf16_experts_forward( grouped_layout, total_padded_rows, ) = _dispatch_routed_input( - hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100() + hidden_states, + top_k_index, + top_k_weights, + self.num_experts, + deepgemm.m_alignment, + is_sm100(), + is_expert_parallel=self._is_expert_parallel, ) weight_up = self.gate_up_proj if self.has_gate else self.up_proj @@ -732,7 +742,13 @@ def deepgemm_fp8_fp4_experts_forward( grouped_layout, total_padded_rows, ) = _dispatch_routed_input( - hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100() + hidden_states, + top_k_index, + top_k_weights, + self.num_experts, + deepgemm.m_alignment, + is_sm100(), + is_expert_parallel=self._is_expert_parallel, ) sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 412fa42630e0..084b17733b16 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -463,7 +463,7 @@ def fp8_batched_mm_experts_forward( # EP sentinel handling: leave `expert_ids` unclamped — the batched kernel early-returns on # `expert_id >= NUM_EXPERTS`, leaving sentinel output rows uninitialized. The post-mask below # zeroes them before the per-token reduction so `uninit * 0 = NaN` can't poison the sum. - sentinel_mask = (expert_ids >= self.num_experts).unsqueeze(-1) + sentinel_mask = (expert_ids >= self.num_experts).unsqueeze(-1) if self._is_expert_parallel else None weight_up = self.gate_up_proj if self.has_gate else self.up_proj weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv @@ -501,7 +501,8 @@ def fp8_batched_mm_experts_forward( # Post-mask sentinel rows: kernel left them uninitialized, so zero them out # before the reduction below (uninit may be NaN; NaN * 0 = NaN). - weighted_out.masked_fill_(sentinel_mask, 0.0) + if sentinel_mask is not None: + weighted_out.masked_fill_(sentinel_mask, 0.0) # Accumulate results using deterministic reshape+sum instead of index_add_ # (index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd) @@ -551,7 +552,7 @@ def fp8_grouped_mm_experts_forward( # valid rows, so sentinel-tail `proj_out` rows are uninit; without the post-mask below, # `proj_out[sentinel] * 0 = NaN * 0 = NaN` would poison the per-token reduction. FP8 # quantized weights are inference-only, so no bwd pre-mask is needed. - sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1) + sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1) if self._is_expert_parallel else None weight_up = self.gate_up_proj if self.has_gate else self.up_proj weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv @@ -590,7 +591,8 @@ def fp8_grouped_mm_experts_forward( weighted_out = proj_out * sample_weights_g.to(proj_out.dtype).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) diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py index 99447bef58cd..ba61ef85bf7f 100644 --- a/src/transformers/integrations/moe.py +++ b/src/transformers/integrations/moe.py @@ -125,11 +125,12 @@ def batched_mm_experts_forward( sample_weights = top_k_weights.reshape(-1) # (S,) expert_ids = top_k_index.reshape(-1) # (S,) - # Clamp EP sentinels so `gate_up_proj[expert_ids]` stays in-bounds. Routing weights are already - # zero at sentinel slots (RouterParallel masks them at dispatch), so the weighted mul drops - # those contributions — we pay the wasted GEMM compute because batched_mm has no offset to skip. - # Out-of-place to avoid mutating the caller's routing tensor (a contiguous `reshape(-1)` aliases it). - expert_ids = expert_ids.clamp(0, self.num_experts - 1) + if self._is_expert_parallel: + # Clamp EP sentinels so `gate_up_proj[expert_ids]` stays in-bounds. Routing weights are already + # zero at sentinel slots (RouterParallel masks them at dispatch), so the weighted mul drops + # those contributions — we pay the wasted GEMM compute because batched_mm has no offset to skip. + # Out-of-place to avoid mutating the caller's routing tensor (a contiguous `reshape(-1)` aliases it). + expert_ids = expert_ids.clamp(0, self.num_experts - 1) # Select gate_up or just up projection weights and biases if self.has_gate: @@ -417,8 +418,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 +437,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( @@ -442,7 +446,8 @@ def grouped_mm_experts_forward( ) # (S, 2 * intermediate_dim) or (S, intermediate_dim) depending on whether we have gating # Zero the sentinel-tail rows the kernel left uninitialized (fwd output and bwd `d_input`). - proj_out = proj_out.masked_fill(sentinel_mask, 0.0) + if sentinel_mask is not None: + proj_out = proj_out.masked_fill(sentinel_mask, 0.0) # Apply gating or activation if self.has_gate: @@ -462,7 +467,8 @@ def grouped_mm_experts_forward( ) # (S, hidden_dim) # Same: zero the uninitialized sentinel-tail rows. - proj_out = proj_out.masked_fill(sentinel_mask, 0.0) + if sentinel_mask is not None: + proj_out = proj_out.masked_fill(sentinel_mask, 0.0) # Apply routing weights weighted_out = proj_out * sample_weights_g.unsqueeze(-1) # (S, hidden_dim) @@ -566,6 +572,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): diff --git a/tests/kernels/test_finegrained_fp8.py b/tests/kernels/test_finegrained_fp8.py index c2e59a29731e..2e16b16bcead 100644 --- a/tests/kernels/test_finegrained_fp8.py +++ b/tests/kernels/test_finegrained_fp8.py @@ -264,7 +264,7 @@ def test_batched_mm_non_gated_uses_up_proj_and_activation(self): def test_batched_mm_passes_sentinel_expert_ids_unclamped(self): # EP sentinels (expert_ids >= num_experts) reach the kernel unclamped; the post-mask zeroes the # matching output rows before the per-token reduction (the kernel leaves them uninitialized). - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16) + experts = make_fp8_experts(num_experts=4, hidden=8, inter=16, is_expert_parallel=True) hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) top_k_index = torch.tensor([[0, 4], [1, 4], [2, 4]], device=torch_device) # 4 == num_experts -> sentinel top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) @@ -326,7 +326,7 @@ def test_grouped_mm_kernel_inputs_and_output(self): def test_grouped_mm_sentinels_dropped_from_histogram(self): # Sentinels are left unclamped so the sort pushes them to the tail and histc(max=num_experts-1) # drops them from tokens_per_expert -> no wasted GEMM rows; the post-mask zeroes their output. - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16) + experts = make_fp8_experts(num_experts=4, hidden=8, inter=16, is_expert_parallel=True) hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) top_k_index = torch.tensor([[0, 4], [1, 4], [2, 4]], device=torch_device) # three sentinels (== num_experts) top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) diff --git a/tests/kernels/test_utils.py b/tests/kernels/test_utils.py index 1a3d110d9162..dc467c451389 100644 --- a/tests/kernels/test_utils.py +++ b/tests/kernels/test_utils.py @@ -32,7 +32,18 @@ def _build_experts( - *, num_experts, hidden, inter, has_gate, has_bias, is_transposed, weight_dtype, scale_dtype, hidden_act, **extra + *, + num_experts, + hidden, + inter, + has_gate, + has_bias, + is_transposed, + weight_dtype, + scale_dtype, + hidden_act, + is_expert_parallel, + **extra, ): def weight(out_dim, in_dim): # Non-transposed weights are (E, out, in); transposed are (E, in, out). @@ -56,6 +67,7 @@ def apply_gate(gate_up): has_gate=has_gate, has_bias=has_bias, is_transposed=is_transposed, + _is_expert_parallel=is_expert_parallel, act_fn=act_fn, _apply_gate=apply_gate, down_proj=weight(hidden, inter), @@ -81,6 +93,7 @@ def make_experts( hidden_act="silu", is_concatenated=True, weight_dtype=torch.bfloat16, + is_expert_parallel=False, ): """BF16 experts stand-in (no scales) for the sonic-moe and DeepGEMM BF16 forwards. Carries `config.hidden_act` / `is_concatenated` (read by sonic-moe).""" @@ -94,6 +107,7 @@ def make_experts( weight_dtype=weight_dtype, scale_dtype=None, hidden_act=hidden_act, + is_expert_parallel=is_expert_parallel, config=types.SimpleNamespace(hidden_act=hidden_act), is_concatenated=is_concatenated, ) @@ -111,6 +125,7 @@ def make_fp8_experts( scale_dtype=torch.float32, activation_scheme="dynamic", block_size=(128, 128), + is_expert_parallel=False, ): """FP8/FP4 experts stand-in (per-projection `_scale_inv`) for the DeepGEMM FP8 and finegrained-fp8 forwards, plus the `_deepgemm_disabled` multi-device flag.""" @@ -124,6 +139,7 @@ def make_fp8_experts( weight_dtype=weight_dtype, scale_dtype=scale_dtype, hidden_act=hidden_act, + is_expert_parallel=is_expert_parallel, activation_scheme=activation_scheme, block_size=block_size, _deepgemm_disabled=False,