Repository navigation
gpt-oss implementationΒ #1739
Description
Description
This outlines the current status of gpt-oss features that need to be implemented in Megatron Core, leveraging Transformer Engine.
β UPDATE: All core GPT-OSS functionality is now available in Megatron Core (training) and Megatron Bridge (checkpoint conversion).
MoE Layer
Enabled Bias
- Status: β Supported
- Implementation: Available in main branch: Fixes for gpt-ossΒ #2038
Attention Mechanisms
Alternating Sliding-Window Attention Pattern
- Status: β Supported - Infrastructure exists for per-layer patterns and sliding window attention using TE
Attention Sinks
- Status: β Implemented - in Transformer Engine and cuDNN
- Reference: Streaming LLM
- Related Transformer Engine PR: [PyTorch] Add sink attention support from cuDNNΒ TransformerEngine#2148
Activation Functions
Custom SwiGLU with Clamping
- Status: β Supported
- Implementation: Megatron Core added partially fused version as "custom quick GeGLU"
FP8-aware fused kernel merged into Transformer Engine
- Related Transformer Engine PR: [Pytorch] Support for Swiglu Activation used in GPT OSSΒ TransformerEngine#2161
Positional Encodings
YaRN RoPE Scaling
- Status: β Fully Supported
- Implementation:
- YaRN scaling to 128k+ context
- Integration with existing RoPE
- YaRN for general RoPE/GPT models
- Convergence validation
- Usage:
--position-embedding-type yarnwith YaRN configuration parameters - Reference: arXiv:2309.00071
Megatron Bridge Support
Megatron Bridge provides full GPT-OSS integration:
- β Checkpoint Conversion: Hugging Face β Megatron format
- β
Pre-configured Providers:
GPTOSSProvider20BandGPTOSSProvider120B - β Quantization Support: Handles MXFP4 weight dequantization
Megatron Bridge + Megatron-LM Example
PR: #2383 provides end-to-end example scripts covering checkpoint conversion (convert_mcore_bf16_checkpoint_from_hf.py) and training/fine-tuning (training_gptoss_20b_h100_bf16_fp8.sh)
Credits: @cuichenx for core implementation, @yiakwy-xpu-ml-framework-team for example scripts
Activity
[-]GPT-OSS Implementation[/-][+]gpt-oss Implementation[/+][-]gpt-oss Implementation[/-][+]gpt-oss implementation[/+]Pikachu1412 commentedon Aug 13, 2025
So how can I train gptoss by this branch?
sbhavani commentedon Aug 13, 2025
We have a guide in NeMo Framework (using Megatron Core): https://docs.nvidia.com/nemo-framework/user-guide/latest/llms/gpt_oss.html.
Megatron-LM training scripts will be added soon
liuqianchao commentedon Oct 2, 2025
how is the progress of support gpt-oss
jgcb00 commentedon Oct 3, 2025
Hi, is there any plan to implement intra-doc masking (Attention being only compute on a document not on the full sequence length) ?
Best regards
tushtush-mc commentedon Oct 9, 2025
@cuichenx Any update on long seq optimisations?
Or for this branch to be merged in main?
JimmyAwoe commentedon Oct 23, 2025
Thanks for contribution first. So has this implementation already finished or closing? I haven't see any related argument about sliding attention window in megatron/training/argument.py.
yiakwy-xpu-ml-framework-team commentedon Nov 22, 2025
@JimmyAwoe GPTOSS model is now provided in megatron bridge.
They use megatron bridge to provide model variants now.
yiakwy-xpu-ml-framework-team commentedon Nov 24, 2025
I have given GPT OSS training support here :
#2383
@JimmyAwoe simply try it out
radna0 commentedon Jan 7, 2026
@sbhavani is gpt-oss now supported in megatron bridge? I'm getting
Model architecture 'GptOssForCausalLM' is not yet supportedthis is in megatron-bridge==0.2.0rc6 and megatron-lm==core_r0.15.0
sbhavani commentedon Jan 7, 2026
@radna0 what's your transformers version? it might work after you upgrade that.
If you're still blocked please file an issue in Megatron Bridge and we'll take a look.