From 32a783de2699882406b3b0114738d91ddb39cf3c Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 8 Oct 2026 06:22:22 -0700 Subject: [PATCH 1/2] Restore GPT-OSS's router-masked expert-parallel plan The MXFP4 experts route their tokens themselves: they take the routing data `routing_torch_dist` builds (already restricted to the rank's experts) instead of expert ids, which token dispatch cannot exchange. With the plan on `ep_dispatch_experts`, MXFP4 GPT-OSS under expert parallelism failed on its first forward (`ep_forward() got an unexpected keyword argument 'scatter_idx'`). The router-masked plan from before #49160 leaves them to their own routing. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/transformers/models/gpt_oss/configuration_gpt_oss.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/gpt_oss/configuration_gpt_oss.py b/src/transformers/models/gpt_oss/configuration_gpt_oss.py index ebe7caedbdf6..a7992ad6f8b9 100644 --- a/src/transformers/models/gpt_oss/configuration_gpt_oss.py +++ b/src/transformers/models/gpt_oss/configuration_gpt_oss.py @@ -39,12 +39,14 @@ class GptOssConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } + # MXFP4 experts route their tokens themselves, which only the router-masked plan leaves them to base_model_ep_plan = { + "layers.*.mlp.router": "ep_router", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.gate_up_proj_bias": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj_bias": "grouped_gemm", - "layers.*.mlp.experts": "ep_dispatch_experts", + "layers.*.mlp.experts": "moe_tp_experts", } num_hidden_layers: int = 36 From dd0281d9b3cdb259f61c923b66435ff2aa674221 Mon Sep 17 00:00:00 2001 From: Ilyas Moutawwakil Date: Thu, 8 Oct 2026 06:36:34 -0700 Subject: [PATCH 2/2] Keep the OpenAI privacy filter on token dispatch It inherits GPT-OSS's config through its modular file, but its experts take expert ids like every other experts class, so it keeps the `ep_dispatch_experts` plan rather than following GPT-OSS back to router masking. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../openai_privacy_filter/modular_openai_privacy_filter.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py b/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py index 2d3c6194e402..68b51b3427b6 100644 --- a/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py +++ b/src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py @@ -72,6 +72,13 @@ @strict class OpenAIPrivacyFilterConfig(GptOssConfig): model_type = "openai_privacy_filter" + base_model_ep_plan = { + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.gate_up_proj_bias": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj_bias": "grouped_gemm", + "layers.*.mlp.experts": "ep_dispatch_experts", + } vocab_size: int = 200064 hidden_size: int = 640 intermediate_size: int = 640