Avoid extra MFSDP v2 model-weight sync memcpy

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()
 

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
         )

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,))