MFSDP v2 should support HSDP and HFSDP placements

[wujingyue]

MFSDP v2 currently rejects the placement differences needed for HSDP/HFSDP.

megatron_fsdp/experimental/parameter_group.py requires main_grad and main_weight to have identical placements and raises until HSDP/HFSDP support is implemented. That blocks hybrid sharding modes where weights, gradients, and optimizer state intentionally use different inner-DP / outer-DP layouts.

Implementation detail:

• First microbatch: all-gather across DP-outer.
• Last microbatch: reduce-scatter or all-reduce across DP-outer.
• DP-outer communication should probably run on a separate stream so it can overlap with DP-inner communication.

Expected outcome:

• Define MFSDP v2 layout semantics for HSDP and HFSDP.
• Support HSDP outer-DP replication / gradient synchronization.
• Support HFSDP optimizer-state sharding across inner and outer DP ranks.
• Add focused tests for the required placement and collective transitions.

Development PRs:

• Add microbatch context helper #5652

[Achyuthan-S]

Hi @wujingyue — I dug into the experimental MFSDP v2 path and I'd be glad to help with this if it's open for outside contribution.

The blocker looks fairly well-scoped from the code:

• FsdpParameterGroup guards on main_grad.placements == main_weight.placements (experimental/parameter_group.py), which is exactly the invariant HSDP/HFSDP need to relax.
• DBuffer.redistribute() currently composes only a single-axis transition, so a 2-D [Partial, Partial] → [Replicate, Flat] reduction isn't yet expressible.
• reduce_gradients() reduces across all mesh axes at once, with no inner-reduce-scatter / outer-reduce split.
• Remove DBuffer mesh axis validation #5441 already dropped the DBuffer mesh-axis validation, so the collectives can run on the outer/inner axes — that groundwork looks ready.

If a PR would be welcome, I'd suggest landing this incrementally rather than as one change:

1. HSDP first — outer-DP-replicated params/optimizer + outer all-reduce of gradients on a single stream. Smallest slice that replaces the guard with real support, and straightforward to test.
2. HFSDP (optimizer sharded across both inner and outer DP) as a follow-up.
3. The separate outer-DP overlap stream as its own PR, since it's orthogonal to correctness and overlaps your comm-overlap work (Overlap experimental FSDP communication #5416 / Add forward all-gather overlap #5513).

I'd add the first 2-D mesh unit tests along the way — the current fully_shard tests are all 1-D.

Two questions before I start:

• Are you already working on this, or would a PR be welcome?
• Anything in flight (e.g. the optimizer work in Add fully_shard_optimizer for mixed-precision FSDP #5411) I should build on or wait for to avoid conflicts?

Happy to open a draft PR for step 1 if that direction sounds right.

[wujingyue]

Thanks for the interest, @Achyuthan-S!

I'd wait for #5513 and the microbatch-context work from #5378 to land. #5513 adds FsdpContext, and the #5378 work adds state for deciding which microbatch we are processing.

I don't think you need #5411 to move forward. You can work around it by using a precision-aware optimizer, e.g. SGD with foreach=False.

[Achyuthan-S]

Thanks @wujingyue, that's really helpful — I'll wait on #5513 and the microbatch-context work before building the overlap/optimized path.

One thing: #5378 is closed and was based on pull-request/4976, so I don't have a live PR to track for the microbatch-context piece. Is that work being re-homed onto main in a new PR I should follow?

In the meantime, would you be open to a correctness-only HSDP slice that's independent of the microbatch context? I'd reduce gradients across DP-outer every backward on the default stream — numerically identical to a last-microbatch reduction since AVG is linear — which only needs the main_grad/main_weight guard relaxed plus a multi-axis reduce_gradients/redistribute, no FsdpContext or is_first_microbatch. Once #5513 and the microbatch work land, I'd follow up to move the outer-DP reduction to the last microbatch on a separate stream.

Either way I'll start on the HSDP/HFSDP layout semantics and the first 2-D-mesh unit tests, since those are unblocked. Thanks again for the pointers on the optimizer workaround.

[wujingyue]

@Achyuthan-S

Is that work being re-homed onto main in a new PR I should follow?

I'll make the new PR, and when it's ready I'll link it to this issue.

I'd reduce gradients across DP-outer every backward on the default stream

I worry about the microbatch-specific logic more than streams. So even just for functionality, I wouldn't recommend proceeding without the microbatch PR.

If you follow https://google.github.io/eng-practices/review/developer/small-cls.html, there might be a chance I reuse some code from you.

[Achyuthan-S]

Thanks @wujingyue , that's clear — I'll hold the HSDP functionality until your microbatch PR lands and watch the issue for the link. Agreed the microbatch-boundary logic (first-microbatch outer all-gather, last-microbatch outer reduce) is the substantive part; reduce-every-backward sidesteps exactly what's worth getting right, so I'll drop that framing.

I'll keep to small, self-contained CLs. Two unblocked pieces I'll prep in the meantime:

1. The first 2-D-mesh parity test for fully_shard (current tests are all 1-D) — it asserts final numerics, so it holds regardless of reduction cadence.
2. A short write-up of the HSDP/HFSDP layout semantics (parameter / gradient / optimizer placement tables across the outer/inner axes) for the issue's first deliverable.

One scoping question so I don't collide with your microbatch PR: would a small standalone CL for the multi-axis gradient reduction — composing inner reduce-scatter + outer reduce, [Partial, Partial] -> [Replicate, Flat] / [Flat, Flat] — be useful, or will that reduction live inside the microbatch PR? Happy to shape it whichever way slots in best.

[wujingyue]

Two unblocked pieces I'll prep in the meantime

SGTM

[Achyuthan-S]

Hey @wujingyue , now that #5652 has landed I'll pick up the DP-outer gradient reduction on top of it — reduce-scatter across DP-inner every backward, reduce DP-outer only at is_last_microbatch, with the main_grad/main_weight guard relaxed for the outer-Partial accumulation — and put it up as a small draft PR. Shout if that's already in progress on your end.