SUPPORT: Pruning support for gpt-oss-20b model

[sandeshrajbhandari]

Was there any planned support for openai/gpt-oss-20b model? I've tried a few things, but it seems the codebase isn't configured to handle the model yet.

Tried to enable pruning for gpt-oss-20b model using the following in MODEL_ATTRS in model_utils.py.

"GptOssForCausalLM": {
        "moe_block": "mlp",
        "gate_proj": "gate_up_proj",
        "up_proj": "gate_up_proj",
        "down_proj": "down_proj",
        "experts": "experts",
        "fused": True,
        "router": "router",
        "num_experts": "num_local_experts",
        "num_experts_per_tok": "num_experts_per_tok",
    },

and for observer.py

@dataclass
class GptOssMoEObserverHookConfig(MoETransformerObserverConfig):
    module_class_name_to_hook_regex: Optional[str] = r"GptOssMLP"
    num_experts_attr_name: str = "config.num_local_experts"
    top_k_attr_name: str = "config.num_experts_per_tok" 
    fused_experts: bool = True

[ZongfangLiu]

Have you resolved yet? I tried but it seems like the activation is extremely large, I'm debugging on it.

[Tsukatskki]

I have developed one version with support of gpt-oss models. I tested it with https://huggingface.co/unsloth/gpt-oss-20b-BF16. Both pruning and router finetuning works fine for me on ada6000.
The entry of whole is reap/prune.sh, feel free to check it out.

https://github.com/Tsukatskki/reap.git

[mklasby]

@Tsukatskki curious if you saw a benefit from the router fine-tuning? Nice extension

[ZongfangLiu]

@Tsukatskki curious if you saw a benefit from the router fine-tuning? Nice extension

Hi Mike, thanks for your awesome work. For your infomation: https://arxiv.org/abs/2603.02217 the paper tried router fine-tuning. There are miner gains from my point of view.

[Tsukatskki]

@mklasby Thanks! From what I observed, router fine-tuning was actually necessary rather than just beneficial. If I applied REAP without router fine-tuning, the router and experts became misaligned, and the model often produced fragmented, incoherent, discontinuous sentences. So at least for GPT-style models in my setup, REAP-only was not really acceptable — router fine-tuning seemed essential to restore stable routing and usable generation quality.