NVIDIA / NVIDIA/Megatron-LM

Incorrect parameter norm with LayerWise optimizer and overlapped parameter gather

Open
#7,131 0 comments 0 reactions 1 assignee Assigned to @FDecaYed View on GitHub
bug community-request waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.