Repository navigation
feat(refit): expose mutable MXFP8 destinations - #6826
Merged
Merged
Conversation
4 of 5 tasks
wdykas
force-pushed
the
simplify-mxfp8-refit
branch
from
August 25, 2026 00:07
d2393ca to
75c6e44
Compare
Contributor
Author
|
/ok to test 75c6e44 |
Contributor
Author
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Contributor
Author
|
/claude strict-review |
kvareddy
approved these changes
Aug 25, 2026
Signed-off-by: wdykas <wdykas@nvidia.com>
Signed-off-by: wdykas <wdykas@nvidia.com>
Signed-off-by: wdykas <wdykas@nvidia.com>
wdykas
enabled auto-merge
August 25, 2026 20:52
Signed-off-by: wdykas <wdykas@nvidia.com>
wdykas
force-pushed
the
simplify-mxfp8-refit
branch
from
August 25, 2026 20:55
bae2a0b to
f8c9186
Compare
Contributor
Author
|
/ok to test f8c9186 |
fanshiqing
reviewed
Aug 26, 2026
fanshiqing
approved these changes
Aug 26, 2026
Signed-off-by: wdykas <wdykas@nvidia.com>
Contributor
Author
|
/ok to test 096b22f |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/32923552756 |
5 of 6 tasks
ksivaman
pushed a commit
to ksivaman/Megatron-LM
that referenced
this pull request
Sep 1, 2026
Signed-off-by: wdykas <wdykas@nvidia.com>
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Enables in-place refitting of MXFP8 inference weights while preserving their underlying data and scale storage.
This is needed by external refit workflows such as Megatron Bridge and NeMo-RL, which copy updated BF16 weights into an already-quantized Megatron inference model.
Preserving storage addresses is important because those buffers may already be captured by CUDA graphs.
shape,device, anddtypemetadata toMXFP8Tensor.MXFP8Tensor.quantize_()andcopy_()for updating an existing MXFP8 destination from a logical tensor.refresh_module_caches()as a public resharding utility.The cache refresh API allows refit integrations that update weights outside
execute_reshard_plan()to refresh parameter-derived state before inference resumes.Why is this needed?
Without an in-place MXFP8 update API, refitting creates replacement quantized tensors and changes their device pointers. This can invalidate CUDA graph captures.
Additionally, some modules maintain values derived from model parameters. External refit paths need a supported way to refresh those caches after copying new weights.
Test coverage
Related work