Fix gradient counting for muon+expert biases - #6099
Conversation
Signed-off-by: Philip Monk <pmonk@nvidia.com>
|
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. |
|
/ok to test |
svcnvidia-nemo-ci
added
Approved
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30483925180 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30483976702 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30484696390 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30502502574 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30502651163 |
4 participants
What does this PR do?
This is another gradient-counting bug, like #5916 and #6080. In this case, the most likely condition is: Muon/Layerwise optimizer, EP>1, TP>1, ETP=1, and experts with linear biases. These are all common except that it's rare for experts to have biases. However, it's the default if you don't pass
--disable-bias-linear. The only under or overcounted parameters are those biases, so the impact on the gradient norm is also pretty small.There are several other configurations (all requiring the layerwise optimizer and expert biases) that trigger this. The issue is LayerWiseDistributedOptimizer wrapped its child optimizers without propagating the regular and expert tensor-parallel process groups. As a result, gradient-norm duplicate filtering fell back to the global TP group, incorrectly excluding some expert parameters when ETP differed from TP. The fix attaches both process groups to each child optimizer and passes them explicitly to zero-counting, ensuring expert parameters are deduplicated using the expert topology.
Probably we should move all global-process-group fallbacks to the edge, around get_megatron_optimizer, and make all internal code require explicit process groups. For now, this is the minimal fix, plus a regression test.
Pre-checks