[Main] Fix Numerical issue for GTP + Muon + MXFP8 param gather + use layerwise optimizer layout - #6754
Conversation
|
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:
See the contribution guide for more details. |
|
/ok to test 3f1ea27 |
| if len(buffers) == 0: | ||
| return [] | ||
|
|
||
| # LayerWiseDistributedOptimizer and DistributedOptimizer classify a whole bucket group from |
There was a problem hiding this comment.
Note: this change is necessary to pass the --over-param-gather = False unit test, otherwise the numerical error will be very high.
With overlap=False, LayerWise launches parameter AG immediately after the Muon step, before Adam fallback has stepped. Without ownership partitioning, the mixed group starts with a Muon bucket, so the entire group—including Adam parameters—is gathered with stale Adam values. DistOpt later skips that group, so updated Adam shards are never gathered.
With overlap=True, AG is deferred to the next forward pass, after both Muon and Adam have stepped. Therefore, the mixed group happens to contain current values.
partition_buckets removes this timing dependency by keeping Muon-owned and Adam-owned buckets in separate groups.
|
/claude strict-review |
|
/ok to test 1c46170 |
|
/claude strict-review |
Wohox
left a comment
Wohox
left a comment
There was a problem hiding this comment.
LGTM!
This PR is highly similar to #5479 regarding partition_buckets & high prec init value & copying back to bf16 staging buffer for AG. Discussed with @zhongbozhu #5479 need to rebase current PR and then decide the scope for change.
|
@zhongbozhu Can you help double check with this PR, the fp8 param gather also utilizes the multi-tensor copy optimization Shiqing developed? |
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
1c46170 to
814cfdd
Compare
| if self.ddp_config is not None | ||
| else config.overlap_param_gather | ||
| ) | ||
| if self.overlap_param_gather and not self.use_buffer_param_sync: |
There was a problem hiding this comment.
Note: to fix the numerical bug, this has to be
if not self.use_buffer_param_sync and (
self.overlap_param_gather or config.reuse_grad_buf_for_mxfp8_param_ag
):
which is them simplified to be layerwise_param_sync_via_bucket_group flag created at init time
| # All-gather updated params. If overlap_param_gather is True, the all-gather | ||
| # is deferred to the forward pre-hooks via DDP bucket infrastructure. | ||
| if not self.overlap_param_gather: | ||
| if self.use_buffer_param_sync: |
There was a problem hiding this comment.
Note: To fix the numerical bug, this has to be
if self.use_buffer_param_sync or self.config.reuse_grad_buf_for_mxfp8_param_ag:
which is then further simplified to be layerwise_param_sync_via_bucket_group, essentially if the bucket group should initiate the param sync (instead of letting optimizer step do the param all-gather, and then in this case we don't have all-gather & fwd pass overlap)
bucket group initiating param sync means that the bucket group chooses the buffer, launches/waits for the collective, and supports forward-hook overlap—instead of the optimizer calling allgather_params() directly, which has no param AG overlap.
The reuse_grad_buf_for_mxfp8_param_ag toggle also requires bucket-group parameter sync because the bucket group owns grad_data and performs the complete FP32 master → BF16 grad buffer → all-gather → TE MXFP8 quant workflow, so it has to be bucket group initiating the param sync, even when overlap_param_gather is False.
|
/ok to test 814cfdd |
|
/ok to test 277eceb |
svcnvidia-nemo-ci
added
Approved
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/33047704410 |
…layerwise optimizer layout (NVIDIA#6754) Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> (cherry picked from commit af7e4f6)
…layerwise optimizer layout (NVIDIA#6754) Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> (cherry picked from commit af7e4f6)
Rebased onto main after NVIDIA#6754 ([Main] Fix Numerical issue for GTP + Muon + MXFP8 param gather + use layerwise optimizer layout). Conflict resolution notes: * ``use_layer_wise_param_layout`` (bool) becomes ``layer_wise_param_layout`` (``'padded' | 'decoupled' | 'legacy'``, default ``'decoupled'``). NVIDIA#6754 gave the old ``False`` value -- the legacy non-DistOpt single-buffer path -- real MXFP8 param-gather support and a padded-vs-legacy parity test, so this PR no longer reuses that value for the new compact decoupled layout. * Dropped the "padded layout does not support FP8/FP4 parameter gather" validation guard: NVIDIA#6754 added that support. * Dropped this PR's copy of the ``get_high_precision_init_val`` FP32-master seeding in ``Float16OptimizerWithFloat16Params``; NVIDIA#6754 landed the same fix. Kept the shared ``pop_high_precision_init_val()`` helper so optimizer.py, distrib_optimizer.py and Megatron-FSDP all use one implementation. * Dropped this PR's ``partition_buckets`` owner-splitting in the force_single_bucket_group branch and the ``_bucket_distopt`` / ``_merged_use_distributed_optimizer`` helpers: NVIDIA#6754's partition-by-optimizer- owner already separates LayerWise from DistOpt buffers before any merging strategy runs. Kept the per-buffer ``ddp_config`` threading in the fp8-merge branch and reduced the rest to one invariant assert. * ``layerwise_param_sync_via_bucket_group`` (NVIDIA#6754) is forced to ``overlap_param_gather`` on the decoupled layout: a compact-layout param's ping-pong owner is unrelated to its DDP-buffer offset, so the DistOpt-style fixed-shard bucket-group gather does not apply there. * Golden values: kept this PR's regenerated ``dist_dist_muon`` values (the default is still the decoupled layout) and applied main's NVIDIA#6539 convention of dropping the leading NaN iteration-time entry. * Reverted the unrelated autoformatter churn in ``models/{gpt,hybrid,base}.py`` and ``models/dist_utils.py`` that main has not applied.
7 participants
What does this PR do?
This PR fixed a numerical issue under the specific combination as shown in the title: GTP + Muon + MXFP8 param gather + use layerwise optimizer layout
The fix is simple, let
use_layer_wise_param_layoutfor Muon layerwise optimizer respectreuse_grad_buf_for_mxfp8_param_agand choose not to remap parameters into the flat DDP buffer and reuse the gradient buffer for AG.This bug also fixed the corner case where Muon + layerwise optimizer + mxfp8 primary weight used to only work with
--overlap-param-gather, so we need to make sure overlap False also just works.How
_stage_layerwise_mxfp8_paramsworks for distributed layerwise optimizer for muonUnit test:
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.