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

[qgallouedec]

grouped_mm_experts_forward runs the expert-parallel sentinel dance unconditionally:

transformers/src/transformers/integrations/moe.py

Lines 420 to 421 in d56c55b

	sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
	expert_ids_g.clamp_(max=self.num_experts - 1)

transformers/src/transformers/integrations/moe.py

Line 437 in d56c55b

selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)

transformers/src/transformers/integrations/moe.py

Line 465 in d56c55b

weighted_out.masked_fill_(sentinel_mask, 0.0)

Sentinel ids are produced only by RouterParallel._prepare_output_fn, i.e. only under expert parallelism. Without it the router's topk cannot return an index >= num_experts, so sentinel_mask is all-False and the two masked_fill_ are pure memory traffic, on tensors the size of the expert activations, twice in the forward and twice again in the backward.

This skips the whole thing when expert parallelism is off. The flag defaults to False where the experts module is built, and is set to True in MoEParamShard.shard_param when it shards the expert dimension, i.e. on the grouped_gemm plan entry that every expert-parallel plan carries. That is also where num_experts is rewritten to the per-rank count the sentinel mask compares against, so the two cannot disagree. A module-level plan entry would not do: llama4's expert-parallel plan has ep_router and grouped_gemm but no moe_tp_experts. The expert-parallel path is unchanged.

Numbers

One Qwen3MoeSparseMoeBlock from Qwen/Qwen3-30B-A3B (128 experts, top-8, hidden=2048), seq 4096, bf16, fwd+bwd, single H100, grouped_mm:

	fwd+bwd	peak memory
main	6.18 ms	3.46 GB
this PR	5.22 ms (1.18×)	3.46 GB

bench_sentinel.py — the script that produced the table

# bench_sentinel.py — cost of the expert-parallel sentinel masking when EP is off.
import statistics
import time

import torch
from transformers import AutoConfig
from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock

MODEL, SEQ, ITERS = "Qwen/Qwen3-30B-A3B", 4096, 30

cfg = AutoConfig.from_pretrained(MODEL)
cfg._experts_implementation = "grouped_mm"
block = Qwen3MoeSparseMoeBlock(cfg).cuda().to(torch.bfloat16)
x = torch.randn(1, SEQ, cfg.hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True)


def run():
    torch.cuda.synchronize()
    t = time.perf_counter()
    out = block(x)
    (out[0] if isinstance(out, tuple) else out).sum().backward()
    torch.cuda.synchronize()
    return time.perf_counter() - t


for _ in range(5):
    run()
times = [run() for _ in range(ITERS)]
torch.cuda.reset_peak_memory_stats()
run()
print(f"is_expert_parallel={getattr(block.experts, 'is_expert_parallel', 'n/a')}  "
      f"{statistics.median(times) * 1e3:6.2f} ms/iter   {torch.cuda.max_memory_allocated() / 2**30:5.2f} GB")

End to end on 2×H100 with FSDP2 (Qwen/Qwen3-30B-A3B, seq 4096, per-device batch 2, LoRA r=16 on attention, chunked_nll):

	step	tokens/s/GPU	peak
main	1.927 s	4252	41.3 GB
this PR	1.798 s	4555 (1.07×)	41.3 GB

The same 1.07× was measured before #47579 refactored this code, on an 8×H100 run of the same model, and 1.05× on Qwen3.6-35B-A3B.
In an 8×H100 profile of that model, aten::masked_fill_ was 2.0% of all GPU kernel time at per-device batch 1 and 3.8% at batch 4, entirely from these two calls.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

CI was right, this was broken. The gate is never True.

install_forward has a single call site and it doesn't pass is_expert_parallel, so it always takes the False default. The masking was skipped unconditionally, including under EP, which is exactly when it's needed.

Repro, 2 GPUs, torchrun --nproc_per_node=2 repro.py:

# torchrun --nproc_per_node=2 repro.py
import os
import tempfile

import torch
import torch.distributed as dist

from transformers import AutoModelForCausalLM, Qwen3MoeConfig
from transformers.distributed import DistributedConfig

dist.init_process_group(backend="nccl")
rank, world = dist.get_rank(), dist.get_world_size()
torch.cuda.set_device(rank)

path = os.path.join(tempfile.gettempdir(), "tiny_qwen3_moe")
if rank == 0:
    config = Qwen3MoeConfig(
        vocab_size=256, hidden_size=128, intermediate_size=256, moe_intermediate_size=64,
        num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2,
        num_experts=8, num_experts_per_tok=2, decoder_sparse_step=1, tie_word_embeddings=False,
    )
    torch.manual_seed(0)
    AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16).save_pretrained(
        path, save_original_format=True
    )
dist.barrier()

model_ep = AutoModelForCausalLM.from_pretrained(
    path, dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=world, enable_expert_parallel=True),
)
dist.barrier()
model_ref = AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16).to(model_ep.device)

torch.manual_seed(0)
input_ids = torch.randint(0, 256, (2, 64), device=model_ep.device)
with torch.no_grad():
    logits_ref = model_ref(input_ids).logits
    logits_ep = model_ep(input_ids).logits

if rank == 0:
    experts = model_ep.model.layers[0].mlp.experts
    print(f"experts implementation: {model_ep.config._experts_implementation}")
    print(f"is_expert_parallel:     {getattr(experts, 'is_expert_parallel', '<absent>')}")
    print(f"NaN in EP logits:       {bool(torch.isnan(logits_ep).any())}")
    print(f"max abs diff vs non-EP: {(logits_ref.float() - logits_ep.float()).abs().max().item()}")

dist.destroy_process_group()

On this branch:

experts implementation: grouped_mm
is_expert_parallel:     False
NaN in EP logits:       True
max abs diff vs non-EP: nan

With the flag actually plumbed through:

experts implementation: grouped_mm
is_expert_parallel:     True
NaN in EP logits:       False
max abs diff vs non-EP: 0.0

On main (no gate, masking unconditional) the same script gives no NaN and 0.0.

Fix is to pass the flag from apply_tensor_parallelism, where model.config.distributed_config.enable_expert_parallel is already set, plus an is_expert_parallel=False kwarg on the base TensorParallelLayer.install_forward and the two overrides that share that call site. 6 lines.

Worth flagging separately: the EP mixin tests pass on CPU even without the fix, so the CPU suite doesn't cover this path. Only the GPU run above surfaces it.

[IlyasMoutawwakil]

thanks for the optimization ! i remember measuring it for inference only and seeing only noise 🥲

[IlyasMoutawwakil]

not sure but i think we can probably know if EP is enabled by checking for DTensor + sharding dim (0) of gate_up/down proj

[IlyasMoutawwakil]

btw there are other implementations that do the same, in finegrained-fp8 and deepgemm

[qgallouedec]

@IlyasMoutawwakil I tried the DTensor+Shard(0) derivation and it can't work: _use_local_dtensor_params wraps the experts forward under EP, so inside the forward the weights are already plain local tensors and the check reads False when EP is on.

(to verify: an in-forward probe shows type(gate_up_proj) == Tensor under EP, and the skipped clamp makes batched_mm crash with an OOB gather.)

Kept the explicit attribute instead, applied your other two points: the batched_mm clamp is now gated too, and the same skip is applied in finegrained-fp8 and deepgemm.

[qgallouedec]

CI was red on nine kernel tests, all the same cause: the tests build their experts stand-in as a SimpleNamespace, so it never goes through use_experts_implementation.__init__
Fixed in eb9dafc, full tests/kernels/ run should be green now.

[VI-Arthur]

LGTM, much needed indeed cc @IlyasMoutawwakil

[qgallouedec]

@VI-Arthur needs another look, the version you approved had a bug.

Passing the flag through install_forward also reached

transformers/src/transformers/distributed/tensor_parallel.py

Lines 601 to 602 in 2ec1d59

	if not is_expert_parallel:
	top_k_weights = _AllReduceBackward.apply(top_k_weights, tp_group)

where a true value skips the _AllReduceBackward on top_k_weights: router grads would stop being summed under EP. That path is back to main.

Flag is now set in MoEParamShard.shard_param under shards_expert_dim, next to the num_experts rewrite the mask compares against. A module-level entry doesn't work, llama4's ep_plan has no moe_tp_experts. That's the bit I'd like your eyes on.

Also dropped the parameter from the three overrides that ignored it, per your comment. tensor_parallel.py is now main + 2 lines.

[IlyasMoutawwakil]

lgtm ! just make sure the new flag will stay relevant in the a2a PR where there are no sentinels iiuc

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 34508019559:1
Result: failure | Jobs: 16 | Tests: 187,513 | Failures: 1 | Duration: 16h 24m