[gpt_oss] Avoid materialising `(x_linear + 1)` intermediate in `swiglu` to reduce activation memory
## 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 条评论