Add unpadded shards for MFSDP v2 RowAtomic/BlockAtomic layouts by XiaodaNV · Pull Request #8116 · NVIDIA/Megatron-LM

XiaodaNV · 2026-10-11T11:36:03Z

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

Add an opt-in unpadded layout for MFSDP v2 RowAtomic/BlockAtomic parameter groups.
Tensors are packed back to back and rank boundaries snap to the nearest row
(32-row block for MXFP8) boundary, so shards are generally unequal but carry zero
padding. The default padded layout is unchanged.

Motivation: the padded layout rounds every group up to dp_size * LCM(row sizes)
(* 32 for MXFP8). For groups mixing many row widths this is large; the
DeepSeek-V3-shaped dense group measured for #7585 pads 13% at DP=4 and over 80% at
DP=16.

Changes:

layout_builder.build_for_row_atomic_unpadding: padding-free builder; shards
differ by at most one row/block and may be empty. Shape validation and per-tensor
alignment are shared with the padded builder.
GlobalLayout.validate_for_row_atomic only checks that no rank boundary splits a
row/block, instead of requiring equal segments; padded layouts still pass.
DBuffer: the list-based all_gather / reduce_scatter paths slice per-peer chunks
via get_rank_range, so uneven layouts work on multi-axis (HSDP/HFSDP) meshes.
QuantizedDBuffer: scale planes derive from the unpadded data layout.
Knob: ShardLayout.PADDED | UNPADDED enum, fully_shard(shard_layout=...);
DistributedDataParallelConfig.megatron_fsdp_padded_shards (default True) and
--no-megatron-fsdp-padded-shards for MCore. Incompatible with nccl_ub
(symmetric memory needs equal shards).

Uneven NCCL collectives are somewhat slower than the even *_into_tensor paths; the
NCCL team is optimizing them. Only users opting into UNPADDED are affected.

Tests: existing DBuffer, QuantizedDBuffer, SGD/HSDP/HFSDP parity, and MXFP8 training
tests are parametrized over ShardLayout; layout builder tests added; adapter flag
covered. tests/unit_tests/distributed/mfsdp_v2 passes on 4xGB200.

Issue tracking

Linked issue: Related to #7585, #6587

copy-pr-bot · 2026-10-11T11:36:07Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2026-10-11T11:36:12Z

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.

wujingyue · 2026-10-13T22:16:40Z

+        Unlike ``build_for_row_atomic``, this never pads. Tensors are packed back to
+        back in tensor-id order, so ``size`` is the sum of all tensor numels.


back to back

Could the alignment be an performance issue? I forgot what alignment nccl likes.

Yes, there is a cost. In our measurements the uneven all-gather / reduce-scatter (list-based dist.all_gather / dist.reduce_scatter) are somewhat slower than the even *_into_tensor / *_tensor paths, and in end-to-end runs the exposed AG/RS show some impact.

The NCCL team is already working on optimizing the uneven collectives. The padded layout stays the default, so this only affects users who opt into ShardLayout.UNPADDED.

Non-blocking: I was talking about alignment, not balance. Also, even with equal-size inputs, dist.all_gather is known to be slower than all_gather_into_tensor/all_gather_single. dist.all_gather does extra data copy.

wujingyue · 2026-10-13T22:18:39Z

        return layout

+    @classmethod
+    def build_for_row_atomic_unpadding(


Consider for future PRs: there may be some code sharable with the even version.

Done: shape validation and per-tensor alignment are shared via _row_atomic_alignments; the two packing algorithms stay separate.

wujingyue · 2026-10-14T00:22:08Z

+        """Slice ``buffer`` into each peer's range under ``placements`` along ``mesh_axis``."""
+        return [
+            buffer.local_buffer.narrow(0, peer_range.start - buffer.offset, peer_range.numel)
+            for peer_range in self.layout.get_peer_ranges(self.mesh, placements, mesh_axis)


get_peer_ranges appears to be very specific for this use case. Before we have more use cases to justify the API, consider getting the peers from mesh.get_group(mesh_axis) and then calling get_rank_range for each.

Done. _peer_chunks now uses dist.get_process_group_ranks(mesh.get_group(axis)) + get_rank_range.

wujingyue · 2026-10-14T00:33:06Z

                             "model gradients during FSDP. If 'auto', then the main gradient data-type will "
                             "be used for the gradient communication / reduction data-type. When using NCCL "
                             "v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels.")
+    group.add_argument('--megatron-fsdp-uneven-shards', action='store_true', default=False,


"" instead of ''. See the style guide.

wujingyue · 2026-10-14T00:35:06Z

+            assert not (args.megatron_fsdp_uneven_shards and args.nccl_ub), \
+                '--megatron-fsdp-uneven-shards is not compatible with --use-nccl-ub'


Nit: this is already checked in ParameterGroup. Should we expose all invalid config combinations in validate_args? That could be a lot of redundancy. Really a question, not a request for change.

Removed both asserts. The v2 adapter already rejects the nccl_ub conflict, and v1 ignores the flag like the other v2-only fields.

wujingyue · 2026-10-14T00:35:59Z

+
+
+@pytest.mark.parametrize("use_out", [False, True])
+def test_quantized_dbuffer_uneven_allgathers_every_plane(distributed_setup, use_out):


Would you like to parameterize the existing test instead of having a new one? They look almost identical.

Done, parametrized over ShardLayout.

wujingyue · 2026-10-14T00:37:43Z

        device: torch.device | str,
        *,
        block_size: int = 1,
+        uneven_shards: bool = False,


Nit: https://www.teamten.com/lawrence/programming/prefer-enums-over-booleans.html

Done: shard_layout: ShardLayout = ShardLayout.PADDED (PADDED / UNPADDED) from fully_shard down to DBuffer.empty. The config layer keeps a bool, megatron_fsdp_padded_shards (default True).

wujingyue · 2026-10-14T00:49:44Z

@@ -503,14 +520,7 @@ def allgather(
                    raise NotImplementedError(
                        "Symmetric-memory allgather() requires equal shard sizes."
                    )
-                # Unequal segments only arise with TensorAtomic on a 1-D mesh, so ``view``
-                # spans the whole global buffer and the layout's rank segments index it.
-                chunks = [
-                    view.local_buffer.narrow(
-                        0, self.layout.rank_to_offset[rank], self.layout.rank_size(rank)
-                    )
-                    for rank in range(self.layout.dp_size)
-                ]
+                chunks = self._peer_chunks(view, input_placements, axis)
                dist.all_gather(chunks, local_buffer, group=group)
            local_buffer = view.local_buffer


This loop body has got longer over time. Can we extract this into a separate method? Maybe all_gather_one_axis and let all_gather call it?

(Can be separate PR)

Agreed, will do in a follow-up PR.

copy-pr-bot · 2026-10-19T13:44:20Z

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.

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

…en the padded and unpadded builders; drop GlobalLayout.get_peer_ranges; replace the uneven_shards bool with a ShardLayout enum (PADDED / UNPADDED) Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

nemo-automation-bot · 2026-10-28T00:54:58Z

🔄 Merge queue validation started!

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

XiaodaNV requested review from a team as code owners October 11, 2026 11:36

XiaodaNV marked this pull request as draft October 11, 2026 11:36

lhb8125 mentioned this pull request Oct 12, 2026

Megatron-LM Monthly Activity Report — October 2026 (Final, through October 31, 2026) #7882

Open

XiaodaNV changed the title ~~Add unpadded uneven shards for MFSDP v2 BlockAtomic layouts~~ Add unpadded uneven shards for MFSDP v2 RowAtomic/BlockAtomic layouts Oct 12, 2026

XiaodaNV force-pushed the mfsdp-v2-blockatomic-uneven branch from 83fe378 to 7b80509 Compare October 13, 2026 06:48

XiaodaNV marked this pull request as ready for review October 13, 2026 06:49

github-actions Bot added the complexity: medium label Oct 13, 2026

wujingyue approved these changes Oct 14, 2026

View reviewed changes

XiaodaNV force-pushed the mfsdp-v2-blockatomic-uneven branch from 7b80509 to 30c6ac7 Compare October 19, 2026 13:42

XiaodaNV marked this pull request as draft October 19, 2026 13:44

copy-pr-bot Bot temporarily deployed to test October 19, 2026 13:44 Inactive

copy-pr-bot Bot temporarily deployed to public October 19, 2026 13:47 Inactive

XiaodaNV force-pushed the mfsdp-v2-blockatomic-uneven branch 2 times, most recently from 53880c0 to c68b1cd Compare October 22, 2026 12:23

XiaodaNV marked this pull request as ready for review October 22, 2026 12:25

copy-pr-bot Bot temporarily deployed to test October 22, 2026 12:26 Inactive

github-actions Bot added complexity: medium and removed complexity: medium labels Oct 22, 2026

copy-pr-bot Bot temporarily deployed to public October 22, 2026 12:29 Inactive

wujingyue approved these changes Oct 23, 2026

View reviewed changes

XiaodaNV requested a review from svcnvidia-nemo-ci October 23, 2026 03:21

XiaodaNV force-pushed the mfsdp-v2-blockatomic-uneven branch from ee35bae to 0ca597d Compare October 23, 2026 04:01

copy-pr-bot Bot temporarily deployed to test October 23, 2026 04:02 Inactive

copy-pr-bot Bot temporarily deployed to public October 23, 2026 04:05 Inactive

Add unpadded uneven shards for MFSDP v2 RowAtomic/BlockAtomic layouts

c014a1a

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

XiaodaNV force-pushed the mfsdp-v2-blockatomic-uneven branch from 0ca597d to 9da9339 Compare October 23, 2026 11:21

copy-pr-bot Bot temporarily deployed to test October 23, 2026 11:22 Inactive

copy-pr-bot Bot temporarily deployed to public October 23, 2026 11:25 Inactive

Share the shape validation and per-tensor alignment computation betwe…

fab7790

…en the padded and unpadded builders; drop GlobalLayout.get_peer_ranges; replace the uneven_shards bool with a ShardLayout enum (PADDED / UNPADDED) Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>

XiaodaNV force-pushed the mfsdp-v2-blockatomic-uneven branch from 9da9339 to fab7790 Compare October 26, 2026 08:15

copy-pr-bot Bot temporarily deployed to test October 26, 2026 08:16 Inactive

copy-pr-bot Bot temporarily deployed to public October 26, 2026 08:19 Inactive

Phlip79 approved these changes Oct 26, 2026

View reviewed changes

guihong-nv approved these changes Oct 26, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Oct 26, 2026

jaredcasper approved these changes Oct 27, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Oct 27, 2026

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

Merged via the queue into NVIDIA:main with commit 47715bb Oct 28, 2026
163 of 169 checks passed

XiaodaNV deleted the mfsdp-v2-blockatomic-uneven branch October 28, 2026 03:40

		Unlike ``build_for_row_atomic``, this never pads. Tensors are packed back to
		back in tensor-id order, so ``size`` is the sum of all tensor numels.

		assert not (args.megatron_fsdp_uneven_shards and args.nccl_ub), \
		'--megatron-fsdp-uneven-shards is not compatible with --use-nccl-ub'



		@pytest.mark.parametrize("use_out", [False, True])
		def test_quantized_dbuffer_uneven_allgathers_every_plane(distributed_setup, use_out):

Conversation

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

Uh oh!

Issue tracking

Uh oh!

copy-pr-bot Bot commented Oct 11, 2026

Uh oh!

github-actions Bot commented Oct 11, 2026

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!

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!

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!

copy-pr-bot Bot commented Oct 19, 2026

Uh oh!

nemo-automation-bot Bot commented Oct 28, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

XiaodaNV commented Oct 11, 2026 •

edited

Loading