Repository navigation
NVFP4 native weights for DDP - #4005
Conversation
Signed-off-by: qiyuw <qiyuw@nvidia.com>
|
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. |
|
Thanks for the detailed description |
Signed-off-by: qiyuw <qiyuw@nvidia.com>
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
/ok to test fc717b2 |
Signed-off-by: qiyuw <qiyuw@nvidia.com>
|
/ok to test 51ad2bf |
|
The label "Approved" is being added by mistake due to a bug. This will be fixed in #4280. |
| # Grad buffer: [g0, g1, g2, g3, ...] numel = N | ||
| # | ||
| # We therefore maintain two index maps: | ||
| # - param_index_map: offsets into the packed param buffer (numel // 2) |
There was a problem hiding this comment.
Thinking about this more, I'm wondering if we should name the two variables param_index_map and packed_param_index_map. param_index_map has the vanilla mapping, and packed_param_index_map has the offsets in the packed param buffer.
There was a problem hiding this comment.
Do we have to make this so general? Packing two values into one byte is specific to nvfp4 case, not fp6 or fp2 even if we have it in the future. If we do the renaming, the current param_index_map will become packed_param_index_map and nvfp4_grad_index_map will become param_index_map, which requires a big refactoring of this PR.
There was a problem hiding this comment.
Well, the other way to think about this is we are keeping param_index_map un-changed from before. And adding a new packed_param_index_map or nvfp4_packed_param_index_map (or something else) to handle the specific NVFP4 case.
There was a problem hiding this comment.
Basically, I don't like the concept of a grad_index_map. Everything is just a param_index_map. And if a parameter takes less than the smallest data type we have available, then we have to do something special (and that's the special case).
There was a problem hiding this comment.
I renamed nvfp4_grad_index_map to nvfp4_unpacked_param_index_map, which serves the purpose here and avoids changing the logic.
There was a problem hiding this comment.
Sorry, I don't love this...you are changing what param_index_map means when NVFP4 is used. nvfp4_unpacked_param_index_map now does what param_index_map does for BF16 or MXFP8 training, and param_index_map stores different things based on the precision.
|
/ok to test 4cebfe8 |
|
/ok to test 2b97a34 |
| # | ||
| self.has_nvfp4_params = any(is_nvfp4tensor(p) for p in self.params) | ||
| self.nvfp4_unpacked_param_index_map = {} | ||
| grad_start_index = 0 if self.has_nvfp4_params else None |
There was a problem hiding this comment.
There is no grad_index_map now. Update these variables to mirror the name of the index map (self.nvfp4_unpacked_param_index_map).
There was a problem hiding this comment.
grad_start_index and grad_end_index are used to get the index of grad data in a grad bucket. I don't think we should rename them to unpacked_param_start_index or something.
|
This PR is actually ready to merge now (i.e. the "Approved" label is accurate). |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/24375815484 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/24381067948 |
Signed-off-by: qiyuw <qiyuw@nvidia.com> Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>
Signed-off-by: qiyuw <qiyuw@nvidia.com>
Signed-off-by: qiyuw <qiyuw@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>
Signed-off-by: qiyuw <qiyuw@nvidia.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
Signed-off-by: qiyuw <qiyuw@nvidia.com>
What does this PR do ?
Add support for NVFP4 native (packed) weights in Megatron's Distributed Data Parallel (DDP) pipeline. This allows parameters to be stored, communicated, and gathered in NVFP4 format, reducing memory footprint and communication volume during training.
Problem
Previously, NVFP4 quantization was only applied at the compute level via Transformer Engine. The DDP param buffers and all-gather operations still used full-precision (BF16) representations, missing the memory and bandwidth savings that native NVFP4 storage provides — similar to what FP8 native weights already support.
Solution
Extend the existing FP8 native weight infrastructure to support NVFP4, accounting for NVFP4's unique dual-buffer layout: the param buffer stores packed bytes (half numel) while the grad buffer uses full numel.
Important notice: This PR only support 2D quantized nvfp4 weights. With 2D quantization, rowwise and colwise tensors share the same values. Given this, we only need to all-gather the rowwise data and transpose to get colwise data later.
Key Changes
Core DDP infrastructure (
param_and_grad_buffer.py)numel // 2) in the param buffer but full numel in the grad buffernvfp4_grad_index_mapfor separate grad buffer offset trackingget_grad_index_map()so the distributed optimizer uses correct (full-numel) indices for grad/param range calculationspartition_bucketsupdated to supportreduce_scatter_with_fp32_accumulation(one bucket per group)Distributed optimizer (
distrib_optimizer.py)_get_nvfp4_params_and_shard_fp32_from_nvfp4()to collect NVFP4 params with their FP32 master weight shards_copy_main_params_to_model_paramscallsquantize_nvfp4_param_shardto cast FP32 shards back to NVFP4 after optimizer stepFP4 utilities (
fp4_utils.py)get_nvfp4_rowwise_packed_shape: compute packed byte shape for param buffermodify_nvfp4_rowwise_storage: swap NVFP4 tensor's rowwise data pointer into the DDP param bufferquantize_nvfp4_param_shard: wraps TE'squantize_master_weightsfor shard-level FP32→NVFP4 castingfp4_model_initnow passespreserve_high_precision_init_val=Trueduring training to retain BF16 init values for accurate FP32 main param initializationArguments & config
--fp4-param-gatherCLI flag (consistent with existing--fp8-param-gather)fp4_param_gathertoDistributedDataParallelConfigTransformerConfig.fp4_paramfield definition to be consistent withfp8_paramfp4_paramtoArgumentGroupFactoryexclude list to prevent duplicate argparse registrationTests (
test_fp4_param.py)Temporary aids for deterministic-mode debugging (to be removed before merge)
disable_stochastic_rounding=True)grad_normrounding inChainedOptimizer(round(grad_norm, 4))Contribution process
Pre-checks
Code review
Feel free to message or comment the @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.
For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.