[Main] Numerical fix for moe single grouped weight with fp8 fp4 primary weight and grad norm spikes by zhongbozhu · Pull Request #5487 · NVIDIA/Megatron-LM

zhongbozhu · 2026-06-24T20:10:30Z

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

What does this PR do ?

Mirrors: #5464

TODOs:

Validate more combinations of toggles in E2E testing and post screen shots

Unit tests with numerical checks passed, pending E2E validation.

test_single_grouped_mxfp8_train_eval_train_matches_train_only is a newly introduced test targeting to test the reuse_grad_buff_for_mxfp8_param_ag rigorously, like adding checks for train-eval-train switches.

Unit test coverage matrix:

Precision	Primary Weight Path	Grad Accum Fusion	Comparison	Notes / Transformer Config
BF16	BF16 primary weight	Off	single grouped weight on compared with single grouped weight off	`bf16=True` `fp8=None` `fp4=None` `gradient_accumulation_fusion=False`
BF16	BF16 primary weight	On	single grouped weight on compared with single grouped weight off	`bf16=True` `fp8=None` `fp4=None` `gradient_accumulation_fusion=True`
MXFP8	BF16 primary weight, MXFP8 compute	Off	single grouped weight on compared with single grouped weight off	`bf16=True` `fp8="e4m3"` `fp8_recipe="mxfp8"` `fp8_param_gather=False` `reuse_grad_buf_for_mxfp8_param_ag=False` `gradient_accumulation_fusion=False`
MXFP8	BF16 primary weight, MXFP8 compute	On	single grouped weight on compared with single grouped weight off	`bf16=True` `fp8="e4m3"` `fp8_recipe="mxfp8"` `fp8_param_gather=False` `reuse_grad_buf_for_mxfp8_param_ag=False` `gradient_accumulation_fusion=True`
MXFP8	MXFP8 primary weight	Off	single grouped weight on compared with single grouped weight off	`bf16=True` `fp8="e4m3"` `fp8_recipe="mxfp8"` `fp8_param_gather=True` `reuse_grad_buf_for_mxfp8_param_ag=True` `gradient_accumulation_fusion=False`
MXFP8	MXFP8 primary weight	On	single grouped weight on compared with single grouped weight off	`bf16=True` `fp8="e4m3"` `fp8_recipe="mxfp8"` `fp8_param_gather=True` `reuse_grad_buf_for_mxfp8_param_ag=True` `gradient_accumulation_fusion=True`
NVFP4	BF16 primary weight, NVFP4 compute	Off	single grouped weight on compared with single grouped weight off	`bf16=True` `fp4="e2m1"` `fp4_recipe="nvfp4"` `fp4_param_gather=False` `gradient_accumulation_fusion=False`
NVFP4	BF16 primary weight, NVFP4 compute	On	single grouped weight on compared with single grouped weight off	`bf16=True` `fp4="e2m1"` `fp4_recipe="nvfp4"` `fp4_param_gather=False` `gradient_accumulation_fusion=True`
NVFP4	NVFP4 primary weight	Off	single grouped weight on compared with single grouped weight off	`bf16=True` `fp4="e2m1"` `fp4_recipe="nvfp4"` `fp4_param_gather=True` `gradient_accumulation_fusion=False`
NVFP4	NVFP4 primary weight	On	single grouped weight on compared with single grouped weight off	`bf16=True` `fp4="e2m1"` `fp4_recipe="nvfp4"` `fp4_param_gather=True` `gradient_accumulation_fusion=True`

Env: 1 x gb200 node, 4 GPUs, the unit test only uses 2 parallel ranks.

Command:

torchrun --nproc_per_node=2 --log-dir /tmp/mcore-single-weight-ut --tee 0:3 --redirects 3 -m pytest -s -q tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py

[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_ddp_param_data_remap_data_ptr[bf16]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_ddp_param_data_remap_data_ptr[nvfp4]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_mxfp8_train_eval_train_matches_train_only
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_mxfp8_single_weight_torch_dist_checkpoint_matches_discrete_baseline[save-only-single]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_mxfp8_single_weight_torch_dist_checkpoint_matches_discrete_baseline[save-single-load-discrete]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_mxfp8_single_weight_torch_dist_checkpoint_matches_discrete_baseline[save-discrete-load-single]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_with_primary_param_gather[False-bf16]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_with_primary_param_gather[False-mxfp8]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_with_primary_param_gather[False-nvfp4]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_with_primary_param_gather[True-bf16]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_with_primary_param_gather[True-mxfp8]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_with_primary_param_gather[True-nvfp4]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_without_primary_param_gather[False-bf16]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_without_primary_param_gather[False-mxfp8]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_without_primary_param_gather[False-nvfp4]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_without_primary_param_gather[True-bf16]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_without_primary_param_gather[True-mxfp8]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_without_primary_param_gather[True-nvfp4]
[default0]:PASSED tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py::TestMoESingleGroupedWeightNumerics::test_single_grouped_weight_parity_module_grouped_linear

What does this PR fix?

This PR fixes two correctness issues exposed by combining TE single grouped MoE weights with low-precision primary parameters and Megatron’s distributed optimizer.

1. Incomplete DDP Parameter Remapping

Current main primarily handles the traditional MoE layout where every expert weight is an independent tensor. With moe_single_grouped_weight, all local expert weights are represented by one contiguous TE GroupedTensor:

MXFP8 GroupedTensor
├── rowwise_data
├── scale_inv
├── columnwise_data
├── columnwise_scale_inv
├── quantizer
└── quantized_tensors: cached per-expert MXFP8 views

This layout is more efficient for grouped GEMM, but it cannot be handled as a regular torch.Tensor or List[QuantizedTensor]. The DDP remapping logic must consider both the grouped layout and the primary-weight recipe:

The complete remapping rules are:

non-grouped NVFP4       -> remap packed rowwise bytes
non-grouped quantized   -> remap TE quantized storage
regular torch.Tensor    -> replace param.data with param_data view
grouped NVFP4           -> remap packed rowwise bytes
grouped MXFP8           -> require BF16 AG through grad-buffer reuse
grouped BF16/FP16       -> remap grouped rowwise_data

Without this handling, DDP could update param_data while TE continued reading stale grouped storage, causing silent loss-curve divergence.

2. MXFP8 Reused-Buffer Gradient Pollution

With --reuse-grad-buf-for-mxfp8-param-ag, parameter AG and gradient accumulation share storage:

shared_buffer
├── param_data: temporary BF16 parameter AG view
└── grad_data:  main_grad accumulation view

After eval, checkpoint preparation, or another forced parameter sync, param_gather_dispatched could remain True. The next training step staged BF16 master weights into the shared buffer, but the forward pre-hook incorrectly treated parameter AG as already completed.

That skipped _post_param_sync(), including the required cleanup:

copy_tensor_to_quantized_param(param, param_slice)
bucket.param_data.zero_()

Backward then accumulated gradients into a buffer still containing parameter values, producing grad-norm spikes.

The PR resets parameter-sync dispatch state whenever master weights are staged. This forces the next forward to complete AG, update TE’s MXFP8 storage, and zero the shared buffer before backward.

Together, these changes make single grouped MoE weights work correctly across BF16, MXFP8, and NVFP4 primary-weight configurations while preserving the existing discrete-expert behavior.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

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

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"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
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, the Final Review label 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 Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

copy-pr-bot · 2026-06-24T20:10:34Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2026-06-24T20:10:41Z

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:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

zhongbozhu · 2026-06-24T20:27:06Z

/ok to test 9df8f4e

deepakn94 · 2026-07-16T18:25:26Z

-                            buffer_type=BufferType.PARAM,
-                        )
-                        modify_underlying_storage(param, new_param_data)
+                    if not is_grouped_tensor(param):


This is a total mess...why are there so many if conditions that all seem to be doing similar things?

explained in the updated description section

deepakn94 · 2026-07-16T18:26:15Z

-            # For MXFP8 param:
-            # we only need to map bf16 weights (layernorm, embedding, etc) to the buffer.
-            if not self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag or not is_mxfp8tensor(param):
+            # This branch remaps the parameter storage into persistent DDP param_data buffer.


I dislike how much low-precision logic is seeping into megatron/core/distributed/param_and_grad_buffer.py

deepakn94 · 2026-07-16T18:27:01Z

+                        elif is_float8tensor(param):
+                            # NVFP4 packs two FP4 values per byte, so param_data uses
+                            # packed-byte offsets instead of logical element offsets.
+                            new_param_data = self._get(


What's the difference between this block and the block starting at line 1263?

is_float8tensor checks for non-grouped parameters of moe, while the other section is for grouped tensor, more details updated in the description section

deepakn94 · 2026-07-16T18:27:43Z

Can you please summarize what the fixes are in the description for this PR?

deepakn94 · 2026-07-20T21:51:40Z

+                        # In TE2, is_float8tensor() checks QuantizedTensor, including MXFP8.
+                        # NVFP4 is handled by the branch above.
+                        elif is_float8tensor(param):
+                            # NVFP4 packs two FP4 values per byte, so param_data uses


Why are you talking about NVFP4 in the FP8 branch?

Because every code block has this following selector, which was already there before I made this PR. It's because FP4 now packs two elements in one byte, so we need some special handling. I didn't change this code block, just adding more comments about it.

( nvfp4_packed_param_start_index if self.has_nvfp4_params else param_start_index ),

deepakn94 · 2026-07-20T21:52:03Z

+                            modify_underlying_storage(param, new_param_data)
+                        # Plain torch param: replace param.data with DDP buffer view.
+                        else:
+                            # NVFP4 packs two FP4 values per byte, so param_data uses


When do we get into this else branch?

It's part of the non-grouepd tensor path: so it's discrete weight for moe.

The else here means that each weight is still a high precision buffer in bf16, so it means that we are not quantizing in the optimizer, but we only quantize the weight in forward pass.

svcnvidia-nemo-ci · 2026-07-21T17:14:52Z