Add GPT-OSS-20B model support

[CarlG0123]

Summary

• Adds support for OpenAI's GPT-OSS-20B, a 20B-parameter MoE model with 32 experts
• Implements custom MoE architecture (GPT-OSS-specific activation, router with bias, softmax-after-top-k routing)
• Includes a direct safetensors loader that dequantizes MXFP4 weights on CPU, enabling the model to run on consumer hardware (tested on MacBook Pro with 38.7GB RAM)

Key architecture differences from Mixtral

• Custom GLU activation: gate.clamp(max=7) * sigmoid(gate * 1.702) * (up.clamp(-7, 7) + 1) (not standard SiLU)
• Router has bias and applies softmax after top-k selection (Mixtral applies before)
• Expert projections have biases (Mixtral does not)
• Merged expert weights with interleaved gate/up columns

Files changed

File	Description
transformer_lens/components/mlps/gpt_oss_moe.py	GptOssExpert and GptOssMoE classes
transformer_lens/pretrained/weight_conversions/openai.py	HF model → TransformerLens weight converter
transformer_lens/loading_from_pretrained.py	Model registration, config extraction, weight dispatch
transformer_lens/factories/mlp_factory.py	Factory routing for GptOssForCausalLM
transformer_lens/pretrained/weight_conversions/__init__.py	Export convert_gpt_oss_weights
run_gpt_oss.py	Standalone loader using direct safetensors + MXFP4 dequantization

Usage

# Via TransformerLens from_pretrained (requires ~80GB RAM for HF model path)
model = HookedTransformer.from_pretrained("openai/gpt-oss-20b")

# Via direct safetensors loader (works on <40GB RAM machines)
python run_gpt_oss.py              # all 24 layers
python run_gpt_oss.py --layers 12  # fewer layers for less memory

Test plan

• 1-layer model: forward pass produces valid logits (no NaN)
• 3-layer model: loads and runs inference correctly
• Full 24-layer model: loads all layers (~10 min), runs inference on MacBook Pro (38.7GB RAM)
• Imports don't break existing models (tested GPT-2 config loading)
• CI tests pass (no existing tests modified)

🤖 Generated with Claude Code