Add SchedulePolicy.unshard_all_planes_in_backward for MXFP8 recomputation in MFSDP v2 by XiaodaNV · Pull Request #8432 · NVIDIA/Megatron-LM

XiaodaNV · 2026-10-22T14:04:25Z

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

What does this PR do?

Lets an FsdpModule opt in to gathering all MXFP8 planes in backward, for recomputation inside the module that does not re-enter its forward hooks. Follow-up to #7907, which gathers only the MXFP8 planes each compute phase needs.

Problem. Backward unshards gather only columnwise planes. A recomputation that does not re-enter the owning FsdpModule's pre_forward therefore finds no rowwise planes and fails in Transformer Engine (Requested row-wise usage, but MXFP8Tensor is missing row-scaled FP8 data), e.g. MTP eh_proj under the root with --recompute-granularity full.

Change. SchedulePolicy(unshard_all_planes_in_backward=True) makes the module's backward unshards, including its prefetch by the preceding module, gather rowwise planes together with columnwise ones. The default is unchanged, and weights of other dtypes are unaffected.

Tests (test_quantization.py, zero3, reentrant and non-reentrant):

test_mxfp8_backward_unshard_gathers_rowwise_for_recompute_inside_fsdp_module: recomputing a root-owned proj fails without the opt-in and has both planes ready with it.
test_mxfp8_recompute_inside_fsdp_module_matches_no_recompute: bit-exact against no recompute with the opt-in.

TODO (follow-ups).

The Megatron adapter does not enable the opt-in yet. It should enable it for FsdpModules whose MXFP8 weights are used by recomputation that does not re-enter them, e.g. the root for MTP full recompute (eh_proj).
Recomputation that re-enters a module still gathers rowwise planes on demand, without overlap. On top of Record MFSDP execution order for prefetch #8478 and Unify preset FSDP execution order traversal #8485, prefetching them needs:
1. Record recomputation in the execution order, distinguishable from the original forward (e.g. Phase.RECOMPUTE).
2. Prefetch each event for its own phase: recompute events gather rowwise planes and don't end the backward prefetch window.
3. Advance the order and prefetch from the recomputation's pre_forward.
4. Use PrefetchOrder.RECORDED in the adapter, since registration order has no recompute events.

github-actions · 2026-10-22T14:04:36Z

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.

copy-pr-bot · 2026-10-22T14:04:37Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

svcnvidia-nemo-ci

LGTM. The owner hooks preserve phase-specific unsharding, and the regression test covers both checkpoint modes, bit-exact training parity, and post-backward plane release. Static review only; distributed GPU tests were not run.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

wujingyue

A checkpoint can recompute a module whose weights another
FsdpModule owns without re-entering the owner's pre_forward: MTP eh_proj is owned by the root because the MTP layer is not an FSDP unit.

Can you say more about this? Does the recomputation happen outside the root similar to #7777?

I ask this because it has been a contract that no parameters can escape:

As a contract, no parameters can escape its lowest FsdpModule ancestor to avoid issues like #4899.

I wonder if this can/should be fixed in the model code.

XiaodaNV · 2026-10-23T02:10:49Z

A checkpoint can recompute a module whose weights another
FsdpModule owns without re-entering the owner's pre_forward: MTP eh_proj is owned by the root because the MTP layer is not an FSDP unit.

Can you say more about this? Does the recomputation happen outside the root similar to #7777?

I ask this because it has been a contract that no parameters can escape:

As a contract, no parameters can escape its lowest FsdpModule ancestor to avoid issues like #4899.

I wonder if this can/should be fixed in the model code.

Thanks for the question. This isn't the #7777 case, and no parameter escapes its owner; my description was misleading there.

In #7777, the replay is attached outside the owner, so it runs before the owner's pre-backward and reads sharded params. Here the checkpoint is created inside the owner's forward:

MTP: inside the root's forward, via MultiTokenPredictionLayer._checkpointed_forward;
--recompute-modules mlp: inside TransformerLayer.forward.

So the recompute runs inside the owner's backward. The weights are gathered, but only the columnwise MXFP8 planes, since BACKWARD skips the rowwise ones after #7907. Recompute is forward compute and needs rowwise, hence "missing row-scaled FP8 data". We only detect recompute when it re-enters the owner's __call__, which doesn't happen when the checkpoint sits inside it.

Fixing this in the model code would need one of two things:

Make eh_proj and the MLP their own FSDP units. This adds smaller collectives every step.
Require every checkpoint to wrap a whole FsdpModule. Selective recompute can't satisfy that, and nothing would catch a violation.

This PR hooks the owned leaves, but the hook gathers the owner's whole group, not per-leaf params. Ownership, sharding, and the per-step collectives without recompute are unchanged. Only a recompute adds one rowwise gather per owner, which recompute needs anyway.

wujingyue

This PR hooks the owned leaves, but the hook gathers the owner's whole group, not per-leaf params. Ownership, sharding, and the per-step collectives without recompute are unchanged. Only a recompute adds one rowwise gather per owner, which recompute needs anyway.

This is similar to MFSDP v1's enable_fine_grained_hooks that always make me sweat.

How bad is it to tell (via fully_shard for example) the parent module's (in this case, MTP) pre_backward to unshard rowwise in addition to columnwise?

wujingyue · 2026-10-23T22:56:48Z

+            torch.manual_seed(2026)
+            model = _RecomputedBlockWithOwnedProj(device, use_reentrant)
+        with fully_shard_context(device=device):
+            fully_shard(model.mlp, mesh=mesh, placements=_ZERO3_PLACEMENTS)


I don't think this is needed to reproduce the error. Just a root FsdpModule with part of it checkpointed should be sufficient. The root cause is that FsdpModule doesn't rowwise-allgather its owned parameters and doesn't know it's needed for activation checkpointing.

Changed. A single root FsdpModule that checkpoints only proj reproduces it.

nemo-automation-bot · 2026-10-31T06:17:25Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/41175974043

Backward unshards MXFP8 weights with columnwise planes only and relies on a recomputed FsdpModule's pre_forward to gather rowwise planes. A checkpoint can recompute a module whose parameters another FsdpModule owns without re-entering that owner's pre_forward, e.g. MTP eh_proj under the root, or TransformerLayer MLP linears under selective MLP recompute, so TE fails with missing row-scaled FP8 data. Register a forward pre-hook on each module directly holding an FsdpModule's parameters that unshards the owner for forward compute during recomputation. Groups serving both phases are already ready; sharded MXFP8 groups gather their rowwise planes. Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

…e-enter the module Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

…cks in tests Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

nemo-automation-bot · 2026-11-01T03:36:49Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/41269872069

nemo-automation-bot · 2026-11-01T03:51:22Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/41270797905

nemo-automation-bot · 2026-11-01T04:57:39Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/41275047493

XiaodaNV requested a review from a team as a code owner October 22, 2026 14:04

svcnvidia-nemo-ci marked this pull request as draft October 22, 2026 14:04

copy-pr-bot Bot temporarily deployed to test October 22, 2026 14:05 Inactive

XiaodaNV changed the base branch from main to pull-request/7907 October 22, 2026 14:06

copy-pr-bot Bot temporarily deployed to public October 22, 2026 14:08 Inactive

XiaodaNV mentioned this pull request Oct 22, 2026

Gather only the MXFP8 weight planes each compute phase needs in MFSDPv2 #7907

Merged

1 task

XiaodaNV marked this pull request as ready for review October 22, 2026 14:18

github-actions Bot added the complexity: low label Oct 22, 2026

svcnvidia-nemo-ci reviewed Oct 22, 2026

View reviewed changes

XiaodaNV changed the base branch from pull-request/7907 to main October 23, 2026 00:47

wujingyue reviewed Oct 23, 2026

View reviewed changes

XiaodaNV force-pushed the mfsdp-v2-mxfp8-recompute-owned-weights branch from 77105d7 to 89dcc45 Compare October 23, 2026 10:38

copy-pr-bot Bot temporarily deployed to test October 23, 2026 10:40 Inactive

copy-pr-bot Bot temporarily deployed to public October 23, 2026 10:42 Inactive

XiaodaNV requested review from svcnvidia-nemo-ci and wujingyue October 23, 2026 10:43

wujingyue reviewed Oct 23, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to public October 24, 2026 04:46 Inactive

XiaodaNV force-pushed the mfsdp-v2-mxfp8-recompute-owned-weights branch from c8520f6 to 5eeed8c Compare October 24, 2026 04:55

copy-pr-bot Bot temporarily deployed to public October 24, 2026 05:00 Inactive

XiaodaNV mentioned this pull request Oct 30, 2026

Record MFSDP execution order for prefetch #8478

Merged

1 task done

copy-pr-bot Bot had a problem deploying to test October 30, 2026 02:47 Error

copy-pr-bot Bot temporarily deployed to public October 30, 2026 02:49 Inactive

XiaodaNV force-pushed the mfsdp-v2-mxfp8-recompute-owned-weights branch from 89f28f1 to ac42c30 Compare October 30, 2026 02:56

copy-pr-bot Bot deployed to test October 30, 2026 02:57 Active

copy-pr-bot Bot temporarily deployed to public October 30, 2026 03:00 Inactive

XiaodaNV enabled auto-merge October 30, 2026 03:08

XiaodaNV added this pull request to the merge queue Oct 31, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Oct 31, 2026

XiaodaNV force-pushed the mfsdp-v2-mxfp8-recompute-owned-weights branch from ac42c30 to 7133fc6 Compare October 31, 2026 11:35

copy-pr-bot Bot deployed to test October 31, 2026 11:36 Active

copy-pr-bot Bot temporarily deployed to public October 31, 2026 11:39 Inactive

shjwudp mentioned this pull request Oct 31, 2026

Enable MFSDP v2 recompute weight planes in the MCore adapter #8803

Closed

XiaodaNV enabled auto-merge October 31, 2026 12:11

XiaodaNV added 6 commits October 31, 2026 18:26

Refine testacase

4df89bc

Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

Add SchedulePolicy.unshard_rowwise_in_backward for MXFP8 recompute

bc51c1f

Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

Rename unshard_rowwise_in_backward to unshard_all_planes_in_backward

d803a81

Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

Scope unshard_all_planes_in_backward to recomputation that does not r…

4730de1

…e-enter the module Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

Reuse is_in_backward from module_utils and split plane allocation che…

efe6b4a

…cks in tests Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

Conversation

XiaodaNV commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Uh oh!

github-actions Bot commented Oct 22, 2026

Uh oh!

copy-pr-bot Bot commented Oct 22, 2026

Uh oh!

svcnvidia-nemo-ci left a comment

Choose a reason for hiding this comment

Uh oh!

wujingyue left a comment • 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!

XiaodaNV commented Oct 23, 2026

Uh oh!

wujingyue left a comment • 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!

wujingyue Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

XiaodaNV Oct 24, 2026

Choose a reason for hiding this comment

Uh oh!

nemo-automation-bot Bot commented Oct 31, 2026

Uh oh!

Uh oh!

nemo-automation-bot Bot commented Nov 1, 2026

Uh oh!

nemo-automation-bot Bot commented Nov 1, 2026

Uh oh!

nemo-automation-bot Bot commented Nov 1, 2026

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

XiaodaNV commented Oct 22, 2026 •

edited

Loading

wujingyue left a comment •

edited

Loading

wujingyue left a comment •

edited

Loading