Skip to content

Support configurable parameter, gradient, and optimizer sharding in experimental FSDP v2 #6137

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Autumn1998 wants to merge 11 commits into
base: main
Choose a base branch
from
Open

Support configurable parameter, gradient, and optimizer sharding in experimental FSDP v2 #6137

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
974fc0e
Add experimental FSDP sharding strategy support
Autumn1998 Jul 29, 2026
38e25d3
Fix experimental FSDP parameter group ordering
Autumn1998 Jul 30, 2026
25b2df8
Use sharding strategy names in experimental FSDP tests
Autumn1998 Jul 30, 2026
d19fe48
Fix experimental FSDP gradient placement semantics
Autumn1998 Jul 30, 2026
f22758e
Merge remote-tracking branch 'ADLR/main' into main-zero-1-2
Autumn1998 Jul 31, 2026
2d523ef
Refine experimental FSDP gradient accumulation
Autumn1998 Jul 31, 2026
bb9fa5d
Preserve gradient buffer stream lifetime
Autumn1998 Jul 31, 2026
93fcd77
Keep FSDP accumulation gradients persistent
Autumn1998 Jul 31, 2026
d4a8933
Track first microbatch per FSDP parameter group
Autumn1998 Jul 31, 2026
afa3a34
Use microbatch state for gradient accumulation
Autumn1998 Jul 31, 2026
64a447a
Update gradient accumulation microbatch test
Autumn1998 Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,13 +109,13 @@ def step_post_hook(
set_grad(parameter, original_grad)
casted_grads.clear()

fsdp_parameter_groups: set[FsdpParameterGroup] = set()
fsdp_parameter_groups: dict[FsdpParameterGroup, None] = {}
Comment thread
wujingyue marked this conversation as resolved.
for optimizer_group in hooked_optimizer.param_groups:
for parameter in optimizer_group["params"]:
parameter_group = get_containing_parameter_group(parameter)
if parameter_group is None:
continue
fsdp_parameter_groups.add(parameter_group)

@wujingyue wujingyue Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd do this instead:

if parameter_group in synced_parameter_groups:
    continue
  parameter_group.sync_model_weight_from_main_weight()
synced_parameter_groups.add(parameter_group)

@Autumn1998 Autumn1998 Jul 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I still little confusing: how does this ensure that the order in which communication is issued on each rank is the same?

@wujingyue wujingyue Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

synced_parameter_groups is only used for deduplication and is never traversed. As long as the parameter order in optimizer.param_groups is uniform, we are good.

fsdp_parameter_groups[parameter_group] = None

for parameter_group in fsdp_parameter_groups:
parameter_group.sync_model_weight_from_main_weight()

@Autumn1998 Autumn1998 Jul 30, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the optimizer post hook, we convert the main weights into model weights. In ZeRO-1, the optimizer-facing weights are stored as flat shards, while the model weights are replicated, so there is an all-gather here. We would like to move this part into the first microbatch so that we can overlap between different layers. @wujingyue Do we have a corresponding design change for this? 🤔

@wujingyue wujingyue Jul 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question. sync_main_weight_to_model_weight, invoked by the post-optimizer-step hook, does a cast followed by a redistribute today:

Megatron-LM/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py

Lines 209 to 211 in 77c8772

self.main_weight.cast(self.model_weight.dtype).redistribute(
self.model_weight.placements, out=self.model_weight
)
.

Do we prefer to

  1. keep cast in post-optimizer-step but redistribute (usually an allgather) in the first microbatch, or
  2. keep both in the first microbatch, or
  3. keep both in post-optimizer-step as is today?

My gut feeling is (1), because:

  1. compute, which is on the critical path, should be batched for shorter latency
  2. communication, which should be hidden behind compute, should be pipelined and exposed as little as possible

As the last resort, we could also introduce a user knob.

@Autumn1998 Autumn1998 Jul 30, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I also think option (1) is the better choice. In addition, we should find a clean way to determine whether it is the first microbatch. Perhaps it would make sense to first open a PR specifically to identify the “first microbatch.” 🤔

@wujingyue wujingyue Jul 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

identify the "first microbatch."

Remember #5652 tried to add is_first_microbatch but got pivoted to is_last_microbatch because FSDP2 doesn't annotate first microbatch?

However, given it's only about redistribution, maybe something like

Megatron-LM/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py

Lines 322 to 323 in ff6b92a

if self.main_grad.placements != self._accumulation_placements:
self.main_grad = self.main_grad.redistribute(self._accumulation_placements)
will work without the user having to annotate?

@Autumn1998 Autumn1998 Jul 31, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We actually don’t need the user to specify the first microbatch; we can just determine it ourselves within the param group.

Expand Down
38 changes: 31 additions & 7 deletions megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@

from ..mixed_precision import MixedPrecisionPolicy
from .dbuffer import DBuffer
from .placement import Partial, Placements, Replicate, changed_mesh_axis
from .placement import Flat, Partial, Placements, Replicate, changed_mesh_axis

_CONTAINING_PARAMETER_GROUP_ATTR = "_mfsdp_parameter_group"

Expand Down Expand Up @@ -146,6 +146,7 @@ def __init__(
dtype=grad_dtype,
device=self.main_weight.device,
)
self.main_grad.local_buffer.zero_()
assert self.main_grad.layout == self.main_weight.layout, (
"main_grad is built from main_weight tensor shapes on the same mesh, "
"and DBuffer layouts are deterministic from those shapes and mesh size."
Expand Down Expand Up @@ -222,12 +223,10 @@ def unshard_parameters(self) -> None:
gather_axis = changed_mesh_axis(
self.model_weight.placements, self._unsharded_model_weight.placements
)
if gather_axis is None:
raise RuntimeError("FSDP parameter unshard requires a changed placement axis.")
with torch.autograd._unsafe_preserve_version_counter(
self._unsharded_model_weight.local_buffer
):
if self._symm_mem_pool is not None:
if self._symm_mem_pool is not None and gather_axis is not None:
self._unsharded_model_weight.rendezvous(gather_axis)
self.model_weight.redistribute(
self._unsharded_model_weight.placements, out=self._unsharded_model_weight
Expand All @@ -250,11 +249,13 @@ def release_unsharded_storage(self) -> None:
# so keep the shared storage-release path.
self._unsharded_model_weight.release_storage()

def _install_sharded_grads(self) -> None:
def _install_sharded_grads(self, grad_buffer: DBuffer | None = None) -> None:
"""Point each sharded parameter's grad at main_grad's current DTensor view."""
assert self.main_grad is not None
if grad_buffer is None:
grad_buffer = self.main_grad
assert grad_buffer is not None
for index, sharded_parameter in enumerate(self.sharded_parameters):
sharded_parameter.grad = self.main_grad.get_dtensor(index)
sharded_parameter.grad = grad_buffer.get_dtensor(index)

def allocate_partial_grad_buffer(self) -> DBuffer:
"""Allocate the unreduced reduce-scatter input buffer."""
Expand Down Expand Up @@ -298,6 +299,29 @@ def reduce_partial_gradients(
"""
assert self.main_grad is not None

optimizer_axis = changed_mesh_axis(
self.main_grad.placements, self.main_weight.placements
)
if (
optimizer_axis is not None
and isinstance(self.main_grad.placements[optimizer_axis], Replicate)
and isinstance(self.main_weight.placements[optimizer_axis], Flat)
):
with self._symmetric_memory_context():
partial_grad = partial_grad.cast(self.main_grad.dtype)
if not is_last_microbatch:
self.main_grad.local_buffer.add_(partial_grad.local_buffer)
return
partial_grad.local_buffer.add_(self.main_grad.local_buffer)

@Autumn1998 Autumn1998 Jul 30, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For no_shard and optim, MFSDP v2 currently accumulates partial gradients into main_grad and performs the final RS on the last microbatch.
However, many kernels fuse gradient accumulation by writing directly into main_grad, bypassing FSDP’s explicit accumulation step. The current implementation is not compatible with this fused path.

@wujingyue wujingyue Jul 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I'm aware of this unimplemented optimization. I believe it can be added by allocating reduce_scatter input before calling the fused backward+copy_in kernel. In fact, for exactly this extension, I chose to launch backward and copy_in on the same stream: http://nv/mfsdp-schedule-design => alternatives considered => FSDP2.

if self._symm_mem_pool is not None:
partial_grad.rendezvous(optimizer_axis)
optimizer_grad = partial_grad.redistribute(self.main_weight.placements)
if partial_grad.placements[optimizer_axis].reduce_op == dist.ReduceOp.SUM:
optimizer_grad.local_buffer.div_(self.mesh.size(optimizer_axis))
self._install_sharded_grads(optimizer_grad)

@Autumn1998 Autumn1998 Jul 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here, we install optimizer_grad instead of main_grad because main_grad has a Partial placement, while the optimizer requires a Flat shard.
The accumulation-gradient layout differs from the optimizer-gradient layout, so we may need a separate accumulation_grad buffer.

@wujingyue wujingyue Jul 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#6041 should make this unnecessary.

@Autumn1998 Autumn1998 Jul 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let me refine this PR with #6041 🤔

@wujingyue wujingyue Jul 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#6041 has been merged FYI

@Autumn1998 Autumn1998 Jul 31, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is still fairly subtle. Take ZeRO-1 as an example. During gradient accumulation, main_grad uses a Partial placement. On the last microbatch, it is finalized with a reduce-scatter and becomes the optimizer-facing gradient with a Flat placement. Therefore, the lifecycle contains two transitions: Partial -> Flat on the last microbatch, followed by Flat -> Partial on the first microbatch of the next accumulation cycle.

These transitions should be handled differently. Partial -> Flat requires a real reduce-scatter and allocates a new output DBuffer. In contrast, Flat -> Partial should not perform communication: the Flat gradient belongs to the previous optimizer step and can be discarded after the optimizer has consumed it, so we can directly allocate a fresh Partial accumulation buffer.

The storage dependencies must still be maintained correctly. We need a record stream:

Partial buffer A
          │
          │ enqueue Reduce-Scatter (asynchronous)
          ▼
    Flat buffer B

    self.main_grad = B
          │
          └── A loses its last reference and is released by the allocator
                        │
                        ▼
                 A's storage is reused
                        │
                        ▼
           Reduce-Scatter is still reading A → corrupted gradients

Since main_grad is replaced by a new DBuffer every global batch—either returned by redistribute() or allocated directly—could its storage address become unstable? Would CUDA graph capture or fused gradient accumulation require this address to remain stable across iterations?

@Autumn1998 Autumn1998 Jul 31, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think an alternative approach is to create an accumulate grad (which uses partial/replicate placement) and separate it from the main grad (which is the flat shard), but let them share the same storage — for example, the main grad would be a view of the accumulate grad. @wujingyue What do you think?

@wujingyue wujingyue Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reduce-Scatter is still reading A → corrupted gradients

I doubt this. All these operations in the above diagram are enqueued to the same, reduce-scatter stream. So it should be safe to free after last use. https://dev-discuss.pytorch.org/t/fsdp-cudacachingallocator-an-outsider-newb-perspective/1486#p-2441-v1-single-stream-no-nonsense-cudacachingallocator-3

We need a record stream

Try our best not to. The whole blogpost is pretty much to teach people not to record_stream for FSDP.

@wujingyue wujingyue Jul 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I doubt this.

I take this back. I think you might have found the root cause of #5956 (comment).

Megatron-LM/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py

Line 345 in 8339200

self.main_grad = self.main_grad.redistribute(self.main_weight.placements)
deallocates the old main_grad on a different stream than it was allocated, which invites race conditions. I'll figure out a way forward ASAP.

cc @Achyuthan-S as well for the HFSDP CI failure he's investigating

self.main_grad.local_buffer.zero_()
return

def has_grad(parameters: tuple[nn.Parameter, ...]) -> bool:
has_any_grad = False
has_any_missing_grad = False
Expand Down
54 changes: 37 additions & 17 deletions tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,19 @@ def _flat_placements() -> Placements:
return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()])


def _strategy_placements(sharding_strategy: str) -> Placements:
placements = {
"no_shard": (Replicate(), Replicate(), Replicate()),
"optim": (Replicate(), Replicate(), Flat()),
"optim_grads": (Replicate(), Flat(), Flat()),
"optim_grads_params": (Flat(), Flat(), Flat()),
}
parameter, gradient, optimizer = placements[sharding_strategy]
return Placements(
dp_axes=[0], parameter=[parameter], gradient=[gradient], optimizer=[optimizer]
)


def _hsdp_placements() -> Placements:
"""HSDP: params/optimizer replicated across DP-outer (axis 0), sharded within
DP-inner (axis 1). main_grad rests [Partial, Flat] between microbatches and is
Expand All @@ -132,9 +145,14 @@ def _mb(num_bytes: int) -> str:
_GEMM_OP_NAME_SUBSTRING = "aten::mm"


@pytest.mark.parametrize(
"sharding_strategy", ["no_shard", "optim", "optim_grads", "optim_grads_params"]
)
@pytest.mark.parametrize("num_microbatches", [1, 3])
def test_fully_shard_sgd_losses_match_baseline(distributed_setup, num_microbatches):
"""Minimal per-module FSDP training should match single-rank SGD."""
def test_fully_shard_sgd_losses_match_baseline(
distributed_setup, num_microbatches, sharding_strategy
):
"""Every supported sharding strategy should match single-rank SGD."""
rank = distributed_setup.rank
world_size = distributed_setup.world_size
device = distributed_setup.device
Expand All @@ -147,10 +165,12 @@ def test_fully_shard_sgd_losses_match_baseline(distributed_setup, num_microbatch
model = TinyModel().to(device)
model.load_state_dict(baseline.state_dict())

fully_shard(model.fc1, mesh=mesh, placements=_flat_placements())
fully_shard(model.fc2, mesh=mesh, placements=_flat_placements())
placements = _strategy_placements(sharding_strategy)
fully_shard(model.fc1, mesh=mesh, placements=placements)
fully_shard(model.fc2, mesh=mesh, placements=placements)
baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
fully_shard_optimizer(optimizer)

micro_batch_size = 2
x = torch.randn(num_microbatches, micro_batch_size, 8, device=device)
Expand All @@ -162,19 +182,19 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]:
for step in range(5):
optimizer.zero_grad()

for microbatch, (microbatch_x, microbatch_target) in enumerate(microbatches):
loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target)
losses.append(loss.detach())
logger.debug(
"%s train parity: rank=%s, step=%s, microbatch=%s, loss=%s",
log_prefix,
rank,
step,
microbatch,
loss,
)

(loss / num_microbatches).backward()
for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches):
with microbatch(model, is_last=microbatch_index == num_microbatches - 1):
loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target)
losses.append(loss.detach())
logger.debug(
"%s train parity: rank=%s, step=%s, microbatch=%s, loss=%s",
log_prefix,
rank,
step,
microbatch_index,
loss,
)
(loss / num_microbatches).backward()

optimizer.step()
return losses
Expand Down