perf(training): accelerate grad clipping with fused Triton kernel - #3815
Open
piyushumate wants to merge 1 commit into
Open
perf(training): accelerate grad clipping with fused Triton kernel#3815piyushumate wants to merge 1 commit into
piyushumate wants to merge 1 commit into
Conversation
…n multi-tensor kernel (NVIDIA-NeMo#3719) In distributed and MoE LLM training (e.g. 70B, Nemotron, DeepSeek), models have 500+ gradient matrices. The existing _clip_grad_norm_impl loops through each parameter's local gradient tensor twice sequentially via Python eager ops: 1. Pass 1: g.detach().abs().max() (500 kernel launches) 2. Pass 2: g.detach().abs().div(scale).square().sum() (1,500+ kernel launches) This creates severe CPU driver overhead (~176 ms on H100 SXM 80GB) and starves GPU execution. Furthermore, using PyTorch _foreach_div allocates massive intermediate buffers (~21 GB for a 70B model). This commit implements a fused, zero-allocation 2D multi-tensor Triton kernel (_fused_grad_clipping.py) that parallelizes chunk reductions across all SMs in double-precision accumulation: - 6.49x speedup over sequential baseline (176.2 ms -> 27.1 ms on NVIDIA H100 SXM 80GB) - 1.25x speedup over PyTorch Foreach without allocating 21 GB of intermediate memory - 100% numerical parity confirmed on H100 (relative diff 0.0002%, well within BF16 1.6% precision) - Clean, automatic fallback to sequential execution on CPU/non-CUDA or when Triton is unavailable - Unit and functional tests added in tests/functional_tests/llm_pretrain_and_kd/test_fused_grad_clipping.py Resolves NVIDIA-NeMo#3719. Signed-off-by: piyushumate <piyushumate@users.noreply.github.com>
Author
|
cc - @akoumpa |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Accelerates gradient norm calculation in
_clip_grad_norm_implby up to 6.49x on NVIDIA H100 using a fused, zero-allocation 2D multi-tensor Triton reduction kernel, resolving CPU driver bottlenecks in distributed and MoE LLM training without memory overhead (#3719).Changelog
nemo_automodel/components/training/_fused_grad_clipping.pywith 2D chunk-parallel Triton kernels (_multi_tensor_max_2d_kernel,_multi_tensor_scaled_l2_2d_kernel) supportingbfloat16,float16,float32, andfloat64._clip_grad_norm_implinnemo_automodel/components/training/utils.pyto dispatch to the fused Triton kernel when running on CUDA, with clean eager fallback for CPU/non-CUDA and preservation of all DTensor sharding placements and device mesh communications.tests/functional_tests/llm_pretrain_and_kd/test_fused_grad_clipping.py.Benchmark Results (NVIDIA H100 SXM 80GB)
Workload: 500 gradient matrices totaling 10.26 Billion elements (21.00 GB) across realistic 70B/MoE model layers:
main)torch._foreach_div.Before your PR is "Ready for review"
Pre checks:
Additional Information