Fix mixed-dtype gradient norms and clipping

[aroshanghias-nvd]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Group gradients by dtype before fused L2-norm and clipping-scale calls, so mixed BF16/FP32 lists produce the intended norm and in-place scaling.

The fused backend tested here interprets a list using its first tensor's dtype. For example, [BF16[128], FP32[128]] filled with 1/64 returns 0.2165063471 instead of 0.25. Reversing the order can read beyond the logical BF16 view. Scaling has an independent defect: even with coefficient 0.5 supplied directly, mixed lists leave some elements unscaled or modify padding beyond the gradient view. The reproductions use padded backing allocations.

The change:

• Buckets L2-norm inputs locally and sums their squared norms before the existing distributed reductions. This preserves collective order for empty ranks and ranks with different dtype sets, and retains the existing homogeneous squared-norm assignment.
• Buckets both scalar- and tensor-coefficient scaling calls, and avoids calling fused kernels with empty lists.
• Adds numerical regressions for both mixed orders, homogeneous controls, ordinary and decoupled gradients, clipping active/inactive, non-vector-aligned lengths, zero gradients with nonzero padding, and asymmetric rank ownership.
• Adds mixed-dtype norm/clipping bit-exact replay under side-stream contention to the existing registered optimizer determinism tests.

The existing single-device gradient-list contract is retained. No kernel implementation or optimizer-state layout changes are required.

Validation

• Native two-rank GB200 investigation with PyTorch 2.11.0+cu130 and the Transformer Engine backend: all 48 norm boundary cases per rank matched first-dtype byte interpretation, with 15 incorrect logical norms. Both scaling APIs failed on both mixed orders. Bucketed controls preserved padding and matched the logical reference.
• Matched caller regression: the original had 24 numerical failures in 96 cases per rank; the corrected caller passed those cases plus four empty-list cases per rank. Mixed-buffer and FP32-buffer DDP/DistributedOptimizer/FusedAdam fixtures each completed eight arms per rank.
• The production file on the PR base is byte-identical to the original file used for that investigation. The PR adds typing and repository formatting to the validated implementation.
• Repository tools/autoformat.sh: Black, isort, pylint, and Ruff passed. Mypy could not resolve Torch/TE/Apex/pytest imports in the isolated local lint environment.
• python3 tools/check_kernel_determinism_coverage.py --base-ref origin/main: passed.
• Exact PR tests on two GB200 ranks with the native TE backend: the original caller produced 29 expected failures out of 105 unit cases per rank; the fixed caller passed all 105 with no skips or errors. Both mixed-dtype determinism cases passed on both ranks, including four bit-exact replays with side-stream contention. Test dependencies were staged separately from the pinned runtime; unrelated top-level suite fixtures were excluded for this focused run.

This establishes the component defect and correction; it does not claim a measured full-training quality impact or that the fused API promises heterogeneous-list support.

Issue tracking

Linked issue: none; standalone bug fix.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have updated the bit-exact determinism tests and their existing registry entry
• I have added proper typing to my code
• I have added relevant documentation
• I have run tools/autoformat.sh on my PR

No training recipe or public API changes; no functional-test goldens or user documentation updates are included.

[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.