perf(training): fused multi-tensor gradient-norm reduction

[akoumpa]

Draft. The kernel is correct and fast in isolation, but the end-to-end gain is small — see Is this worth it? below. Opening for review of the approach rather than as a merge-ready change.

What

Gradient clipping reduces every parameter's gradient separately, ~7 kernel launches each. On DiffusionGemma-26B-A4B (8×H100, ep=8) that is 5,434 launches and 74 ms/step in _clip_grad_norm_impl alone — more launches than the entire MoE GEMM path. This replaces it with a Triton multi-tensor reduction: one launch per dtype, for the 2-norm and inf-norm.

Measured

In isolation, 740 tensors:

per-tensor loop	16.77 ms
fused multi-tensor	1.13 ms (14.8×)

End-to-end, 100 steps, DiffusionGemma-26B-A4B SFT with dispatcher=deepep:

	step time
before (NEMO_AUTOMODEL_FUSED_GRAD_NORM=0)	10.406 s
after	10.370 s
gain	0.35%

wandb: before · after

Is this worth it?

Honestly: the end-to-end win is marginal. I originally sized this change off a profile that attributed 87 ms/step to gradient clipping, but that profile was captured under tracing, which inflates launch-heavy CPU-side work. The isolated benchmark puts the true cost at ~16.8 ms, so removing it is ~0.15% of a 10.4 s step; the measured 0.35% is consistent with that.

What the change does buy, beyond the 0.35%:

• ~5,400 fewer kernel launches per step, which matters more on a CPU-bound step than this GPU-bound one;
• a strictly more accurate norm (fp64 throughout, versus the current fp32-scaled formulation);
• it scales with parameter count rather than with model FLOPs.

Reviewers may reasonably decide 300 lines and a Triton kernel is not worth 0.35%. That call is easier to make with the number stated plainly than buried.

Numerics

Steps 0–11 are bit-identical to the unpatched run (loss, dllm_loss, grad_norm). They diverge from step 12 (loss 1.0903 vs 1.0890, grad_norm 4.9544 vs 4.9792) and drift after that — the expected signature of a more accurate reduction in a system with discrete MoE routing: a last-bit change in the norm changes the clip scale, which flips routing, which amplifies. The 12 identical steps are what establish the implementation is right.

Against an fp64 reference: exact for BF16, 2.5e-16 (fp32), 6.9e-15 (2000-tensor accumulation). The first version of this kernel reduced tiles in fp32 and was quietly ~1e-7 — the tolerance in the test is set tight enough to catch that.

Correctness properties

• Precision — values widen in registers on load, so BF16 is never squared in BF16 (|g| > 256 overflows once squared). The in-tile tl.sum runs in fp64 too. Widening in registers matters as much as widening at all: upcasting whole gradients would allocate 2–8× gradient memory. Measured peak overhead: 0.12% of gradient bytes.
• Determinism — chunk decomposition is host-side from sizes only, no atomics, so the norm is bit-identical run to run. An atomic accumulate would be simpler but would make the gradient norm — and every optimizer step after it — non-reproducible.
• Non-finite — NaN/Inf propagate through the fp64 sum-of-squares naturally, but tl.maximum follows IEEE maxNum and ignores NaN, so an inf-norm over NaN gradients would have returned finite and silently defeated error_if_nonfinite. The kernel counts NaNs and poisons the partial. Both paths tested for NaN/±Inf.
• dtype safety — tensors are bucketed by dtype; a mixed batch previously would have reinterpreted BF16 storage as FP32.
• Unsupported inputs are routed to a slice-bounded reference path rather than converted: non-CUDA, non-contiguous (where .contiguous() would allocate a full-size copy), and dtypes outside the supported set — notably FP8, whose reduction belongs in a scaled path rather than a raw square.
• DTensor — the caller unwraps to local and owns the cross-rank scalar reduction; the kernel never sees a DTensor.

NEMO_AUTOMODEL_FUSED_GRAD_NORM=0 forces the per-tensor path.

Not included

Unit tests. The validation above ran as a standalone script on 8×H100; it should be folded into tests/unit_tests/ before this leaves draft.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.