fix(clip_grads): handle empty grads_for_norm in inf-norm and p-norm paths

[Mattral]

Summary

get_grad_norm_fp32 crashes when called with an empty gradient list,
which occurs in practice when all parameters on a rank are filtered out
(frozen layers, shared params, TP duplicates). Two out of three norm-type
branches are affected; the L2 branch already has a correct guard.

Changes

megatron/core/optimizer/clip_grads.py

Branch	Bug	Fix
inf norm (line 95)	max() over empty generator → ValueError	Add default=torch.tensor(0.0) to max()
generic p-norm (line 127)	total_norm stays float 0.0; all_reduce receives a non-Tensor → TypeError	Initialise total_norm = torch.zeros(1, dtype=torch.float, device='cuda') before loop

No change to the L2 path or any non-empty-list behaviour.

Testing

• Added unit tests in tests/unit_tests/optimizer/test_clip_grads.py
  covering empty-list calls for all three norm_type paths (inf, 2.0,
  custom p). Tests run without GPU using the local fallback
  implementations.
• Existing CI passes unchanged.

Checklist

• Commits signed off (git commit -s)
• Rebased on main
• No unrelated formatting changes
• One logical change per commit

Fixes #5529

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

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[Mattral]

looking forward to have this reviewed

[janEbert]

We can simplify a bit in the norm_type != inf code path, since both norm_type == 2.0 and norm_type != 2.0 need the same zero-tensor if grads_for_norm is an empty list. norm_type == 2.0 currently solves it with the grad_norm, which we could remove if we initialized the total_norm zero-tensor earlier. The entire norm_type == 2.0 path could then be simplified around that.

E.g.:

total_norm = torch.zeros(1, dtype=torch.float, device='cuda')
if not grads_for_norm:
    pass
elif norm_type == 2.0:
    dummy_overflow_buf = torch.zeros(1, dtype=torch.int, device='cuda')
    # Use apex's multi-tensor applier for efficiency reasons.
    # Multi-tensor applier takes a function and a list of list
    # and performs the operation on that list all in one kernel.
    grad_norm, _ = multi_tensor_applier(
        l2_norm_impl,
        dummy_overflow_buf,
        [grads_for_norm],
        False,  # no per-parameter norm
    )
    # Since we will be summing across data parallel groups,
    # we need the pow(norm-type).
    total_norm = grad_norm**norm_type

else:
    for grad in grads_for_norm:
        grad_norm = torch.norm(grad, norm_type)
        total_norm += grad_norm**norm_type

[Mattral]

Thanks for the suggestion

that's a cleaner structure. I'll hoist total_norm = torch.zeros(1, dtype=torch.float, device='cuda') before the norm_type branches so both norm_type == 2.0 and the generic p-norm path share the same zero-init, and drop the redundant else: grad_norm = torch.zeros(...) guard from the L2 path. Will push an updated commit shortly.

[janEbert]

Thank you for the contribution and quick adjustment!

[Mattral]

@janEbert
Thanks for the review! Updated the commit per your suggestion. Hoisted the zero-tensor init before both branches and dropped the redundant L2 guard. Also fixed the DCO sign-off. LMK if anything else needs changing.

[janEbert]

/ok to test 0e5d2e9

[janEbert]

You'd quickly need to run bash tools/autoformat.sh.

[Mattral]

The autoformatter fix has been pushed. CI is now waiting for approvals from @NVIDIA/core-adlr and @NVIDIA/core-nemo.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29325376449

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29330988382

[janEbert]

Previous fail due to deterministic throughput being 1.29 slower than non-deterministic baseline. I doubt it's related to this PR, probably just flaky CI.