Skip to content

Restore GPT-OSS's router-masked expert-parallel plan #49431

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
IlyasMoutawwakil wants to merge 2 commits into from
Closed

Restore GPT-OSS's router-masked expert-parallel plan #49431

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
32a783d
Restore GPT-OSS's router-masked expert-parallel plan
IlyasMoutawwakil Oct 8, 2026
dd0281d
Keep the OpenAI privacy filter on token dispatch
IlyasMoutawwakil Oct 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion src/transformers/models/gpt_oss/configuration_gpt_oss.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@ArthurZucker ArthurZucker Oct 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

mmm this should not be related. MXFP4 should be acting on the plan AFAIK

@IlyasMoutawwakil IlyasMoutawwakil Oct 12, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in an ideal world yes, but the current mxfp4 integration is based on kernels that aren't modular enough by default to do that. i can adapt this mxfp4 (which we will get rid of once we get to finegrained kernels) but it has to import some non-surfaced ops in the gpt oss mxfp4 kernels which turns iut to be much uglier.

@IlyasMoutawwakil IlyasMoutawwakil Oct 12, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that's why i said this is temporary, until fg kernels is merged which does not have the limitations and rigidity of "triton-kernels", lmk what u prefer

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
Expand Down
7 changes: 7 additions & 0 deletions src/transformers/models/openai_privacy_filter/modular_openai_privacy_filter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading