ITADN

[gpt_oss] Avoid materialising `(x_linear + 1)` intermediate in `swiglu` to reduce activation memory

#3437Openjanbernloehr 创建于 2026-05-27
J
janbernloehrcommented
## Summary `torchtitan/models/gpt_oss/moe.py::swiglu` returns `out_glu * (x_linear + 1)`, which materialises an intermediate tensor the same shape as `x_linear`. Under FSDP2 + EP and full activation checkpointing, swiglu is recomputed in backward; that intermediate contributes a multi-GiB peak per rank that is otherwise unnecessary, and OOMs on 80 GiB H100 when the per-rank routed-token tensor is ~5 GiB. ## Reproduction - Repo: https://github.com/pytorch/torchtitan - Revision: `963c20c` (and current main) - Hardware: 8 nodes × 8 H100 80 GB. - Config: `gpt_oss_120b` with EP=64, seqlen=8192, full activation checkpointing, batch_size=1. - Command: ``` torchrun --nnodes 8 --nproc-per-node 8 -m torchtitan.train \ --module gpt_oss --config gpt_oss_120b \ --parallelism.expert_parallel_degree 64 \ --training.seq_len 8192 --training.steps 50 \ --training.local_batch_size 1 ``` - Expected: training proceeds past step 2. - Actual: `torch.OutOfMemoryError: Tried to allocate 4.20 GiB. GPU 0 has a total capacity of 79.11 GiB of which 2.45 GiB is free.` Stack ends in `torch.utils.checkpoint.unpack_hook → recompute_fn → moe.py:362 self.experts(...) → moe.py:151 _experts_forward → swiglu` at the `out_glu * (x_linear + 1)` site. Step 1 fits at 36.5% memory; step 2 backward recompute is the failure point. ## Proposed change ```python - return out_glu * (x_linear + 1) + return torch.addcmul(out_glu, out_glu, x_linear) ``` `torch.addcmul(out_glu, out_glu, x_linear)` computes `out_glu + out_glu * x_linear == out_glu * (1 + x_linear)` in a single fused kernel without materialising the intermediate. No numerical change, no autograd-contract change; benefits any gpt_oss training under activation checkpointing where the swiglu is on the recompute path. ## Validation Patched in our container build (sed against the upstream source before `pip install`) and ran the previously-failing 8×8 H100 EP=64 / seqlen=8192 / full-AC training for 50 steps. Validation pipeline status to be appended once it finishes.
1 条评论