Gpt-oss training on A100 - OOM error

John6666
Wed, 26 Nov 2025 23:18:24 +0000

When considering the difference compared to the H100 GPU environment shown in OpenAI’s Notebook, the overhead of executing torch.compile might be the culprit. Next in line are differences in PyTorch, CUDA, or Transformers versions.

Either way, I think VRAM is quite tight and failures are likely to occur from even minor differences, so personally, I recommend using QLoRA.

Main causes of OOM on a single H100

1. You’re really doing BF16 LoRA on a 20B MoE, not FP4 training

The cookbook loads with Mxfp4Config(dequantize=True) and torch_dtype=torch.bfloat16, which dequantizes MXFP4 weights to BF16 on GPU. (OpenAI Cookbook)

Unsloth’s analysis: any non-Unsloth training stack must upcast to BF16, and then “all other training methods will require a minimum of 65GB VRAM to train the 20B model.” (Unsloth)

2. Attention uses full L×L “eager” kernels

The fine-tuning notebook sets attn_implementation="eager", which in current Transformers uses full sequence × sequence attention instead of the more memory-efficient sliding-window behavior GPT-OSS was designed for. (OpenAI Cookbook)

At typical settings (e.g. batch 4, seq 1024–2048), this creates attention tensors on the order of hundreds of MiB per layer, which is exactly what fails in your OOM trace.

3. 80GB H100 is “just enough”, so any extra overhead breaks it

Officially, the notebook is only guaranteed for one H100-80GB:

“This notebook is designed to be run on a single H100 GPU with 80GB of memory. If you have access to a smaller GPU, you can reduce the batch size and sequence length…” (OpenAI Cookbook)

In reality, things like:

torch.compile/Dynamo buffers,

mixed-precision upcasts to FP32,

logging overhead,

and allocator fragmentation
can easily push peak usage from “~75–78GB” to “>80GB”, causing OOM even on H100.

4. This is a known pattern, not just your setup

Independent Japanese notes on GPT-OSS fine-tuning say for gpt-oss-20B LoRA in BF16 you realistically need ≥60GB VRAM (e.g. A100×2, H100), and that this is why they strongly recommend QLoRA. (Zenn)

Practical solutions / workarounds on a single H100

A. Best fix: switch to an optimized QLoRA stack (Unsloth)

Unsloth’s GPT-OSS guide: gpt-oss-20B fine-tuning fits in ~14GB VRAM; they recommend ≥16GB for stable runs. (Unsloth)

Their key point: Unsloth keeps GPT-OSS in a quantized / optimized form and uses Flex Attention, cutting VRAM by ≈70–80% compared to BF16 LoRA.

If you can change stack, this is by far the cleanest fix: your H100 then has massive headroom instead of being at the cliff edge.

B. If you must stay on HF/TRL BF16 LoRA notebook

1. Shrink “tokens per step” hard

Set e.g.:

per_device_train_batch_size = 1

max_length = 512 (start small)

gradient_accumulation_steps = 8–16 to recover effective batch size.

Then gradually increase max_length and/or accumulate steps while watching nvidia-smi.

2. Remove avoidable overhead

Do not use torch.compile on this model until a plain run is stable.

Turn off heavy logging (report_to="none").

Set allocator flags to reduce fragmentation, e.g.:

export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True,max_split_size_mb:512"

3. Consider better attention kernels on H100

Where supported, replacing attn_implementation="eager" with a Hopper-optimized backend (e.g. kernels-community/vllm-flash-attn3) reduces how much L×L attention is materialized in HBM. (OpenAI Cookbook)

Even with all of this, BF16 LoRA on gpt-oss-20B will always be near the 80GB limit; QLoRA / Unsloth is the only way to make it “comfortable”.

Key links

OpenAI Cookbook – Fine-tuning gpt-oss with Transformers (original notebook, H100-80GB target) (OpenAI Cookbook)

Unsloth guide – “gpt-oss: How to run and fine-tune” (explains 65GB vs 14GB VRAM for 20B) (Unsloth Docs)

Zenn hardware outline for GPT-OSS fine-tuning (20B LoRA BF16 ≈ 60GB+ VRAM, QLoRA ≈ 14GB+) (Zenn)