[WIP] Experimental implementation of gpt-oss (grouped GEMM MoE + FlexAttention sink/sliding)

[KhoomeiK]

This PR currently implements gpt-oss in torchtitan such that forward pass & sampling from 20B matches the HuggingFace reference implementation. You can run torchtitan/experiments/gpt_oss/scripts/compare_hf_to_tt.py to verify this.

A few notes:

• It default uses torch 2.9 grouped GEMMs for the MoE. If you can't force upgrade to torch 2.9, please set use_grouped_mm=False to use the for-loop implementation.
• It default uses my FlexAttention implementation of sinks + sliding window, which currently runs into CUDA errors when sampling more than one token. If you care about sampling, please set use_flex_attn=False, though I intend to resolve this soon.
• Parallelism is unimplemented so training is untested as of now. This is obviously my highest priority right now, but I wanted to get a baseline implementation out for the community.
• The 120B checkpoint remains untested, but architecture differences are minimal so I don't foresee any issues.
• I haven't yet spent any time adding support for the tiktokenizer, please use the HuggingFace tokenizer for now.
• I still have some open questions regarding mixed precision training and what the best numerics recipe is here.

Thanks @Chillee for pointing me to FlexAttention and @vwxyzjn for tips on the DeepSeek reference implementation!

[meta-cla]

Hi @KhoomeiK!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!