#1999·apex

[错误] 在 torch.no_grad() 下, FusedRMSNorm 在每次向前调用时会泄露 2 个 CUDA 向量

作者: dustnehowl创建于 2026年4月29日更新于 2026年4月29日
标签bug

Describe the Bug

apex.normalization.FusedRMSNorm leaks 2 CUDA tensors on every forward call, even under torch.no_grad(). The accumulation is linear in the number of calls and is not freed by gc.collect() or torch.cuda.empty_cache(), which causes OOM in long-running inference workloads. This is particularly impactful because Hugging Face transformers automatically uses FusedRMSNorm for T5 layer norms when apex is importable. T5-XXL has 49 LayerNorm calls per forward pass, so 49 × 2 = 98 CUDA tensors per T5EncoderModel.forward(...). In a real-world FLUX inference workload on a single B200 (180GB VRAM), this caused OOM after ~90 batches at batch_size=4 (~8 GB of GPU memory accumulated per batch). The leak disappears completely when apex is blocked from being imported (forcing transformers to fall back to native PyTorch RMSNorm), which strongly points to FusedRMSNorm as the source.