Fix gradient-norm undercounting when using EP and TP

[philipcmonk]

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

What does this PR do?

This PR fixes correctness bugs surrounding expert-parallel training when combined with tensor parallelism and/or expert-tensor parallelism.

There are two related issues, both applying only to MoE models:

• Common case: if EP > 1 and TP > 1, then the gradient norm calculations omit some experts, so gradient clipping applies less than expected. This could reduce stability during warmup and in case of gradient spikes. This usually applies only if --moe-grouped-gemm is set.

• Uncommon case: if EP = 1 and ETP != TP (which would be unusual), then the gradients themselves are not correctly reduced across the expert data-parallel group. If ETP > TP, then unrelated gradients will mix (probably with other columns of the same expert). If TP > ETP, then each expert will only receive a portion of the gradient it should, and that portion will be different on different ranks. This will cause ranks to diverge -- they will think the experts are replicated across the EDP group, but they are not.

In both cases, the logging for the parameter norms, gradient norms, and count_zeros will also be incorrect.

These come from three main causes:

• param.allreduce=False is used to indicate that a parameter is subject to the expert topology (EP/ETP/EDP) instead of the normal topology (TP/CP/DP). It was only being set if EP > 1, however it should also be set if ETP != TP, since that still implies that EDP != DP.

• param_is_not_tensor_parallel_duplicate assumed the parameter was subject to the normal topology. It should check the fixed param.allreduce flag, and if it's subject to the expert topology, use the ETP group instead of the TP group.

• TEGroupedLinear did not mark experts with param.tensor_model_parallel when it should have (to match other linear layers).

Besides fixing these, this PR also attaches expert_tp_group to the optimizer alongside the existing tp_group. We need this at all the call-sites for param_is_not_tensor_parallel_duplicate.

Finally, we add these tests which were useful in reproducing and fixing the bugs:

• Set the params and gradients to all 1s and verify the calculated L2 norms are sqrt(num_params) under a few different parallelisms.
• In the same condition, run gradient clipping and assert each gradient is now max_norm / sqrt(num_parameters).
• To verify the correct reduction of gradients, for an expert linear layer, with EP=1, set the gradients to 2**rank, then call finish_grad_sync from DDP. Assert the expected gradients come out of this. The rank-specific values are necessary to distinguish the case where gradients were being mixed with other ranks when ETP!=TP.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[philipcmonk]

/ok to test

[philipcmonk]

/ok to test

[philipcmonk]

/ok to test

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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