Add FSDP v2 MXFP8 mixed precision support by Autumn1998 · Pull Request #6 · shjwudp/Megatron-LM

Autumn1998 · 2026-05-12T09:06:47Z

What does this PR do ?

⚠️ 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 the @mcore-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
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

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

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!

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

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.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

shjwudp · 2026-05-12T10:04:58Z

@@ -56,6 +56,14 @@ def __init__(
        self.device = params[0].device
        self.dtype = params[0].dtype
        self.requires_grad = params[0].requires_grad
+        self.mp_policy = mp_policy
+        self.is_fp8_group = self.mp_policy.is_fp8_param(params[0])


Can dtype be used instead of is_fp8_group?

The dtype approach probably won't work — the dtype of an FP8 tensor is actually uint8 since it's a fake tensor. We probably shouldn't use the dtype of a fake tensor as the basis for this check.

I see, I think we can consider using an enum to represent the dtypes supported by M-FSDP. The flag is_fp8_group can be misleading, as it conflates different FP8 recipes without distinction.

shjwudp · 2026-05-12T10:09:13Z

        if s != "no_shard":
-            wbuf = self._create_buffer(self.dtype, shard_weights)
+            model_weight_dtype = self.mp_policy.model_weight_buffer_dtype(self.params[0])


Could model_weight_buffer_dtype be implemented as a standalone function? We may need to determine the parameter data type directly from the parameters or their parent module, especially in cases where mixed formats like NVFP4 and MXFP8 are used, rather than assuming a single uniform recipe.

This function should be functionally equivalent to the previous version. Since it accepts a param argument, we actually first determine the type based on the type of that param However, we need to ensure that the recipe is consistent within the same param group

Could the model_weight_buffer_dtype be static method?

shjwudp

Looks great overall: it implements MXFP8 parameter initialization, pre‑ and post‑unshard/reshard handling, and batched distributed quantization from the main weight to the FP8 weight. It may be worth double‑checking compatibility with fine‑grained activation recompute.

In particular, this also addresses the EDP grad NaN issue, which is great work!

Let’s merge first so we don’t block progress, but there are still some details that need fixing, and we should also add a design doc to go with it.

shjwudp · 2026-05-14T13:39:39Z

@@ -56,6 +56,14 @@ def __init__(
        self.device = params[0].device
        self.dtype = params[0].dtype
        self.requires_grad = params[0].requires_grad
+        self.mp_policy = mp_policy
+        self.is_fp8_group = self.mp_policy.is_fp8_param(params[0])


I see, I think we can consider using an enum to represent the dtypes supported by M-FSDP. The flag is_fp8_group can be misleading, as it conflates different FP8 recipes without distinction.

shjwudp · 2026-05-14T13:51:07Z

+            if work is None:
+                work = weight_work
+
+            if self.is_fp8_group and full_weight_buffer is not None:


When should full_weight_buffer be None?

For the self.is_fp8_group condition, a more intuitive way to determine this would be based on whether a backward-specific weight is required. For example:

if self.requires_backward_specific_weight:

shjwudp · 2026-05-14T13:52:52Z

+            start_offsets,
+            self.dp_group,
+            model_param_shards,
+        )


Can we abstract L265-298 into a function and put it in mixed_precision.py?

shjwudp · 2026-05-14T14:00:48Z

+            return self.main_grads_dtype
+        return torch.bfloat16 if self.is_fp8_param(tensor) else tensor.dtype
+
+    def initial_main_weight(self, tensor: torch.Tensor) -> torch.Tensor:


Would dequantized be a better name here?

get_dequantized_value or get_high_precision_value

shjwudp · 2026-05-15T02:43:39Z

+                or self.model_weight_buffer is None
+                and self.main_weight_buffer is None
+            ):
+                continue


Looks have memory issue.

shjwudp · 2026-05-15T02:50:34Z

+            full_weight_buffer, weight_work = weight_buffer.unshard(
+                async_op=async_op,
+                bind_params=not self.is_fp8_group,
+            )


Handle activation recompute case:

Per-layer recompute

module [A, B] forward, unshard and bind [A, B] rowwise data

module B backward, unshard and bind B colwise data

Fine-grained recompute

module C forward, unshard and bind C rowwise data

module B backward, unshard and bind B colwise data

*C is sub module of B

shjwudp · 2026-05-15T03:02:11Z

+            return self.main_grads_dtype
+        return torch.bfloat16 if self.is_fp8_param(tensor) else tensor.dtype
+
+    def initial_main_weight(self, tensor: torch.Tensor) -> torch.Tensor:


get_dequantized_value or get_high_precision_value

shjwudp · 2026-05-15T03:04:05Z

        if s != "no_shard":
-            wbuf = self._create_buffer(self.dtype, shard_weights)
+            model_weight_dtype = self.mp_policy.model_weight_buffer_dtype(self.params[0])


Could the model_weight_buffer_dtype be static method?

Add FSDP v2 MXFP8 mixed precision support

8a2515b

shjwudp reviewed May 12, 2026

View reviewed changes

Autumn1998 added 3 commits May 14, 2026 00:28

Align FSDP v2 MXFP8 weight handling

3a80853

Merge mfsdp refactor main into FSDP v2 MXFP8

876a87b

Clean up FSDP v2 bucket allocator ownership

4e219fd

shjwudp approved these changes May 14, 2026

View reviewed changes

shjwudp merged commit 0c911b8 into shjwudp:mfsdp_refactor_main May 14, 2026
1 check passed

shjwudp reviewed May 15, 2026

View reviewed changes

Conversation

Autumn1998 commented May 12, 2026

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Autumn1998 May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shjwudp May 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

shjwudp left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shjwudp May 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

(Step 1): Add PR label `Expert Review`

Autumn1998 May 13, 2026 •

edited

Loading

shjwudp May 15, 2026 •

edited

Loading

shjwudp May 15, 2026 •

edited

Loading