Optimize DistMuon packing copies

tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py

@@ -0,0 +1,79 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from torchtitan.distributed.flex_shard._optimizer_reshard_runtime import _batched_copy_
+
+
+class TestBatchedCopy(unittest.TestCase):
+    def test_rejects_mismatched_lists(self):
+        with self.assertRaisesRegex(ValueError, "equal length"):
+            _batched_copy_((torch.empty(1),), ())
+
+    def test_empty_and_single_copy_bypass_foreach(self):
+        _batched_copy_((), ())
+
+        source = torch.arange(6).reshape(2, 3)
+        destination = torch.empty_like(source)
+        with patch.object(torch, "_foreach_copy_") as foreach_copy:
+            _batched_copy_((destination,), (source,))
+        foreach_copy.assert_not_called()
+        torch.testing.assert_close(destination, source)
+
+    def test_foreach_copy_supports_mixed_sizes(self):
+        sources = (
+            torch.arange(6).reshape(2, 3),
+            torch.arange(8).reshape(2, 4),
+        )
+        destinations = tuple(torch.empty_like(source) for source in sources)
+        original_foreach_copy = torch._foreach_copy_
+        with patch.object(
+            torch,
+            "_foreach_copy_",
+            wraps=original_foreach_copy,
+        ) as foreach_copy:
+            _batched_copy_(destinations, sources)
+        foreach_copy.assert_called_once()
+        for destination, source in zip(destinations, sources, strict=True):
+            torch.testing.assert_close(destination, source)
+
+    def test_foreach_copy_supports_noncontiguous_views(self):
+        source_base = torch.arange(24).reshape(4, 6)
+        destination_base = torch.zeros_like(source_base)
+        sources = (source_base[:, ::2], source_base[:, 1::2])
+        destinations = (destination_base[:, ::2], destination_base[:, 1::2])
+
+        _batched_copy_(destinations, sources)
+
+        torch.testing.assert_close(destination_base, source_base)
+
+    def test_foreach_copy_casts_between_dtypes(self):
+        sources = (torch.arange(3), torch.arange(4))
+        destinations = (
+            torch.empty(3, dtype=torch.float32),
+            torch.empty(4, dtype=torch.float32),
+        )
+        original = torch._foreach_copy_
+        with patch.object(torch, "_foreach_copy_", wraps=original) as foreach_copy:
+            _batched_copy_(destinations, sources)
+        foreach_copy.assert_called_once()
+        for destination, source in zip(destinations, sources, strict=True):
+            torch.testing.assert_close(destination, source.to(destination.dtype))
+
+    def test_rejects_mismatched_shapes(self):
+        with self.assertRaises(RuntimeError):
+            _batched_copy_(
+                (torch.empty(3), torch.empty(4)),
+                (torch.empty(9), torch.empty(4)),
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py

@@ -684,11 +684,14 @@ def _prepare_redistributed(
         )
         prepare(item, prepared)
         spans = schedule.input_spans_by_parameter[index]
-        for span in spans:
-            packed = storage_buffer[
-                span.buffer_offset : span.buffer_offset + span.numel
-            ]
-            packed.copy_(_tensor_region_view(prepared, span.region).reshape(-1))
+        packed_views = tuple(
+            storage_buffer[span.buffer_offset : span.buffer_offset + span.numel]
+            for span in spans
+        )
+        prepared_views = tuple(
+            _tensor_region_view(prepared, span.region).reshape(-1) for span in spans
+        )
+        _batched_copy_(packed_views, prepared_views)
 
 
 def _compute_redistributed(
@@ -717,21 +720,31 @@ def _compute_redistributed(
             dtype=plan.dtype,
             device=plan.device,
         )
-        for span in received_spans:
-            received = work.compute_fragment_buffer[
+        compute_views = tuple(
+            _tensor_region_view(compute_tensor, span.region) for span in received_spans
+        )
+        received_views = tuple(
+            work.compute_fragment_buffer[
                 span.buffer_offset : span.buffer_offset + span.numel
-            ]
-            _tensor_region_view(compute_tensor, span.region).copy_(
-                received.view(span.region.shape)
-            )
+            ].view(span.region.shape)
+            for span in received_spans
+        )
+        _batched_copy_(compute_views, received_views)
 
         compute(item, compute_tensor)
 
-        for span in to_storage.input_spans_by_parameter[index]:
-            packed = work.compute_fragment_buffer[
+        output_spans = to_storage.input_spans_by_parameter[index]
+        packed_views = tuple(
+            work.compute_fragment_buffer[
                 span.buffer_offset : span.buffer_offset + span.numel
             ]
-            packed.copy_(_tensor_region_view(compute_tensor, span.region).reshape(-1))
+            for span in output_spans
+        )
+        compute_output_views = tuple(
+            _tensor_region_view(compute_tensor, span.region).reshape(-1)
+            for span in output_spans
+        )
+        _batched_copy_(packed_views, compute_output_views)
 
 
 def _finalize_redistributed(
@@ -757,16 +770,32 @@ def _finalize_redistributed(
             device=plan.device,
         )
         spans = schedule.output_spans_by_parameter[index]
-        for span in spans:
-            packed = work.storage_buffer[
+        update_views = tuple(_tensor_region_view(update, span.region) for span in spans)
+        packed_views = tuple(
+            work.storage_buffer[
                 span.buffer_offset : span.buffer_offset + span.numel
-            ]
-            _tensor_region_view(update, span.region).copy_(
-                packed.view(span.region.shape)
-            )
+            ].view(span.region.shape)
+            for span in spans
+        )
+        _batched_copy_(update_views, packed_views)
         finalize(item, update)
 
 
+def _batched_copy_(
+    destinations: tuple[Tensor, ...],
+    sources: tuple[Tensor, ...],
+) -> None:
+    """Copy aligned region views with a single foreach launch."""
+    if len(destinations) != len(sources):
+        raise ValueError("destinations and sources must have equal length")
+    if not destinations:
+        return
+    if len(destinations) == 1:
+        destinations[0].copy_(sources[0])
+        return
+    torch._foreach_copy_(destinations, sources)
+
+
 def _tensor_region_view(tensor: Tensor, region: _TensorRegion) -> Tensor:
     view = tensor[
         tuple(