Incorrect parameter norm with LayerWise optimizer and overlapped parameter gather
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
With `LayerWiseDistributedOptimizer` and `overlap_param_gather=True`, parameter norm logging can mix updated FP32 main weights on owner ranks with stale model weights on non-owner ranks.
The reported norm can be incorrect and differ across DP ranks.
Steps/Code to reproduce bug
1. Run BF16 training with LayerWise optimizer parameters distributed across at least two DP ranks.
2. Enable `--overlap-param-gather` and `--log-params-norm`.
3. Compare the logged norm after an optimizer step with the norm computed from all owners’ updated
FP32 main weights.
This follows from code inspection; a multi-GPU reproduction has not yet been run.
Expected behavior
Compute the norm from each owner’s updated main weights, summing squared norms across DP ranks
without requiring model parameter gather.
Additional context
LayerWise creates main_param only on owner ranks and does not set main_param_sharded. _get_param_data() therefore falls back to model weights on non-owner ranks, which may not have been synchronized yet.
fix: https://github.com/NVIDIA/Megatron-LM/pull/7132
@NVIDIA/mcore-oncall
Contributor guide
Assessment
This issue has not been assessed yet.