DeepSpeed
961df84b - Keep float32 intermediates in the clamped SwiGLU eager fallback (#8992)

Commit
5 days ago
Keep float32 intermediates in the clamped SwiGLU eager fallback (#8992) `swiglu_clamped()` runs the Triton kernel when Triton is available and `_swiglu_clamped_eager()` otherwise , The triton kernel accumulates in float32 but the eager function not, which causes `tests/unit/v1/moe/test_autoep_expert_activation.py::TestGroupedExpertsActivation::test_fused_kernel_matches_plain_pytorch[swiglu_oai]` fail when triton is not available. The PR fixes the numeric error, has been validated in H100 and KLXPU, with the cost of larger peek memory and kernel time cost. | | fwd peak | fwd+bwd peak | fwd+bwd time | |---|---|---|---| | before | 2.50 GiB | 4.00 GiB | 13.87 ms | | after | 5.00 GiB | 10.00 GiB | 32.21 ms | Signed-off-by: guhangsong <guhangsong@outlook.com>
Author
Parents
Loading