From 85cbc0ceee3fe8cf883434684ada6045244943b9 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Thu, 16 Jul 2026 00:00:20 +0000 Subject: [PATCH] Avoid extra MFSDP v2 model-weight sync memcpy Signed-off-by: Jingyue Wu --- .../src/megatron_fsdp/experimental/dbuffer.py | 41 ++++++++++--------- .../experimental/parameter_group.py | 7 ++++ .../distributed/mfsdp_v2/test_dbuffer.py | 23 +++++++++++ 3 files changed, 52 insertions(+), 19 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py index 9b6e6dc44c3..a240b148d62 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py @@ -244,15 +244,24 @@ def distribute_tensors( return buffer def _create_or_validate_out( - self, placements: Iterable[Placement], out: "DBuffer | None" + self, + out: "DBuffer | None", + *, + placements: Iterable[Placement] | None = None, + dtype: torch.dtype | None = None, ) -> "DBuffer": - placements = tuple(placements) + if placements is None: + placements = self.placements + else: + placements = tuple(placements) + if dtype is None: + dtype = self.dtype if out is None: return DBuffer( mesh=self.mesh, placements=placements, tensor_shapes=self.layout.tensor_shapes, - dtype=self.dtype, + dtype=dtype, device=self.device, ) @@ -262,24 +271,18 @@ def _create_or_validate_out( raise ValueError(f"Expected out placements {placements!r}, got {out.placements!r}.") if out.layout != self.layout: raise ValueError(f"Expected out layout {self.layout!r}, got {out.layout!r}.") - if out.dtype != self.dtype: - raise ValueError(f"Expected out dtype {self.dtype}, got {out.dtype}.") + if out.dtype != dtype: + raise ValueError(f"Expected out dtype {dtype}, got {out.dtype}.") if out.device != self.device: raise ValueError(f"Expected out device {self.device}, got {out.device}.") return out - def cast(self, dtype: torch.dtype) -> "DBuffer": + def cast(self, dtype: torch.dtype, *, out: "DBuffer | None" = None) -> "DBuffer": """Return this buffer with the same layout and placements in ``dtype``.""" - if self.dtype == dtype: + if self.dtype == dtype and out is None: return self - destination = DBuffer( - mesh=self.mesh, - placements=self.placements, - tensor_shapes=self.layout.tensor_shapes, - dtype=dtype, - device=self.device, - ) + destination = self._create_or_validate_out(out, dtype=dtype) destination.local_buffer.copy_(self.local_buffer) return destination @@ -304,7 +307,7 @@ def redistribute( if changed_axis is None: if out is None: return self - out = self._create_or_validate_out(new_placements, out) + out = self._create_or_validate_out(out, placements=new_placements) out.local_buffer.copy_(self.local_buffer) return out @@ -334,7 +337,7 @@ def allgather(self, mesh_axis: int, *, out: "DBuffer | None" = None) -> "DBuffer placements = list(self.placements) placements[mesh_axis] = Replicate() _validate_placements(placements) - out = self._create_or_validate_out(placements, out) + out = self._create_or_validate_out(out, placements=placements) dist.all_gather_into_tensor( output_tensor=out.local_buffer, input_tensor=self.local_buffer, @@ -351,7 +354,7 @@ def allreduce(self, mesh_axis: int, *, out: "DBuffer | None" = None) -> "DBuffer placements = list(self.placements) placements[axis] = Replicate() - out = self._create_or_validate_out(placements, out) + out = self._create_or_validate_out(out, placements=placements) out.local_buffer.copy_(self.local_buffer) dist.all_reduce( out.local_buffer, op=partial_placement.reduce_op, group=self.mesh.get_group(axis) @@ -372,7 +375,7 @@ def reduce_scatter( placements = list(self.placements) placements[axis] = new_placement _validate_placements(placements) - out = self._create_or_validate_out(placements, out) + out = self._create_or_validate_out(out, placements=placements) dist.reduce_scatter_tensor( output=out.local_buffer, input=self.local_buffer, @@ -400,7 +403,7 @@ def scatter( self.mesh, placements ) else: - out = self._create_or_validate_out(placements, out) + out = self._create_or_validate_out(out, placements=placements) destination_offset = out.offset destination_numel = out.local_buffer.numel() diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index eeec848416b..a44e799b04b 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py @@ -201,6 +201,13 @@ def sync_model_weight_from_main_weight(self) -> None: if self.main_weight is self.model_weight: return + if self.main_weight.placements == self.model_weight.placements: + self.main_weight.cast(self.model_weight.dtype, out=self.model_weight) + return + + # main_weight is typically the higher-precision optimizer dtype, while + # model_weight is the lower-precision compute dtype. Cast before redistributing + # so cross-rank communication moves the smaller compute-dtype payload. self.main_weight.cast(self.model_weight.dtype).redistribute( self.model_weight.placements, out=self.model_weight ) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_dbuffer.py b/tests/unit_tests/distributed/mfsdp_v2/test_dbuffer.py index 8631113d480..1180cd55c80 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_dbuffer.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_dbuffer.py @@ -178,6 +178,29 @@ def test_cast_preserves_layout_and_casts_values(distributed_setup): ) +def test_cast_with_out_reuses_destination_and_casts_values(distributed_setup): + """DBuffer.cast writes casted values into an existing destination buffer.""" + mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,)) + tensors = _same_tensors_on_all_ranks(distributed_setup.device) + buffer = DBuffer.distribute_tensors(tensors, mesh, [Replicate()]) + destination = DBuffer( + mesh=mesh, + placements=[Replicate()], + tensor_shapes=buffer.layout.tensor_shapes, + dtype=torch.bfloat16, + device=distributed_setup.device, + ) + destination_data_ptr = destination.local_buffer.data_ptr() + + result = buffer.cast(torch.bfloat16, out=destination) + + assert result is destination + assert destination.local_buffer.data_ptr() == destination_data_ptr + _assert_dbuffer_local_tensors_close( + destination, [tensor.to(dtype=torch.bfloat16) for tensor in tensors] + ) + + def test_release_and_reallocate_storage_preserves_buffer_views(distributed_setup): """DBuffer storage can be released and reallocated without replacing existing views.""" mesh = init_device_mesh(distributed_setup.device.type, (distributed_setup.world_size,))