Loading…
Tricks from OpenAI gpt-oss YOU 🫵 can use with transformers
Hugging FaceAritra Roy Gosthipaty, Sergio Paniego, Vaibhav Srivastav, Pedro Cuenca, Arthur Zucker, Nathan Habib, Cyril Vallez
Summary
To support OpenAI's GPT-OSS series of models, the transformers library introduced several performance upgrades that apply across supported architectures. The release integrates zero-build custom kernels downloadable directly from the Hub, reducing external dependency bloat and compilation friction for operations like Liger RMSNorm and MegaBlocks MoE. Native support for MXFP4 quantization groups vector elements into 32-value blocks with shared scales, allowing GPT-OSS 20B to fit in roughly 16 GB of VRAM and GPT-OSS 120B in roughly 80 GB. In addition, transformers incorporates Flash Attention 3 with attention sinks, continuous batching via the generate_batch API for experimentation, and automatic memory pre-allocation to speed up model loading on GPUs.
Context
OpenAI released the GPT-OSS model family featuring novel techniques such as MXFP4 quantization, custom kernels, and attention sinks. Enabling efficient loading, execution, and fine-tuning for these large Mixture of Experts models in transformers required eliminating kernel compilation friction, managing high VRAM requirements, and improving GPU memory allocation speeds.
Approach / What changed
Transformers integrated Hub-downloadable pre-built kernels for operations like Liger RMSNorm and MegaBlocks MoE MLP, added native MXFP4 quantization with blockwise scaling across 32-element vectors, supported Flash Attention 3 with attention sinks for Hopper architectures, implemented continuous batching through the generate_batch API, and introduced GPU memory pre-allocation based on device maps during model loading.
Takeaways
- Pre-built custom kernels can be downloaded directly from the Hub by passing use_kernels=True, bypassing local compilation and dependency bloat for kernels like Liger RMSNorm and MegaBlocks MoE.
- MXFP4 uses an E2M1 4-bit layout with blockwise scaling across 32 elements, reducing VRAM footprint so that GPT-OSS 20B requires roughly 16 GB and GPT-OSS 120B requires roughly 80 GB.
- Model loading times are accelerated automatically by reading device maps and pre-allocating large GPU memory blocks before copying weights, preventing repeated small memory allocation calls.
Related reading
AI Model Co-Design: Hardware-Friendly LLM Design
Efficient large language model serving requires co-designing architectures to balance datacenter token throughput and user interactivity within hardware constraints. Roofline analysis shows that linear layers become memory-bound when hidden or intermediate projection dimensions are small, leaving arithmetic hardware underutilized even at large token counts. For prefill workloads, Chunked Pipeline Parallelism reduces first-token latency on long contexts when models use regular, repeatable layer patterns that split into balanced stages. For latency-critical decoding, decoupling attention and feed-forward network parallelization enables independent optimization, applying expert parallelism for feed-forward networks and Helix parallelism to shard the key-value cache across sequences. Following architectural design rules—such as dimension alignment, width prioritization, and low-precision NVFP4 compatibility—allows models to achieve higher utilization on modern accelerator hardware.
Elizabeth Goodman