[DistMuon] Drop the foreach-copy compatibility fallback

tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py

@@ -9,23 +9,21 @@
 
 import torch
 
-from torchtitan.distributed.flex_shard._optimizer_reshard_runtime import (
-    _foreach_copy_or_fallback_,
-)
+from torchtitan.distributed.flex_shard._optimizer_reshard_runtime import _batched_copy_
 
 
-class TestForeachCopyOrFallback(unittest.TestCase):
+class TestBatchedCopy(unittest.TestCase):
     def test_rejects_mismatched_lists(self):
         with self.assertRaisesRegex(ValueError, "equal length"):
-            _foreach_copy_or_fallback_((torch.empty(1),), ())
+            _batched_copy_((torch.empty(1),), ())
 
     def test_empty_and_single_copy_bypass_foreach(self):
-        _foreach_copy_or_fallback_((), ())
+        _batched_copy_((), ())
 
         source = torch.arange(6).reshape(2, 3)
         destination = torch.empty_like(source)
         with patch.object(torch, "_foreach_copy_") as foreach_copy:
-            _foreach_copy_or_fallback_((destination,), (source,))
+            _batched_copy_((destination,), (source,))
         foreach_copy.assert_not_called()
         torch.testing.assert_close(destination, source)
 
@@ -41,7 +39,7 @@ def test_foreach_copy_supports_mixed_sizes(self):
             "_foreach_copy_",
             wraps=original_foreach_copy,
         ) as foreach_copy:
-            _foreach_copy_or_fallback_(destinations, sources)
+            _batched_copy_(destinations, sources)
         foreach_copy.assert_called_once()
         for destination, source in zip(destinations, sources, strict=True):
             torch.testing.assert_close(destination, source)
@@ -52,22 +50,32 @@ def test_foreach_copy_supports_noncontiguous_views(self):
         sources = (source_base[:, ::2], source_base[:, 1::2])
         destinations = (destination_base[:, ::2], destination_base[:, 1::2])
 
-        _foreach_copy_or_fallback_(destinations, sources)
+        _batched_copy_(destinations, sources)
 
         torch.testing.assert_close(destination_base, source_base)
 
-    def test_incompatible_dtype_uses_copy_fallback(self):
+    def test_foreach_copy_casts_between_dtypes(self):
+        # torch._foreach_copy_ casts like Tensor.copy_ does, so differing dtypes
+        # need no special handling; only unequal shapes would be a problem, and
+        # the op rejects those itself.
         sources = (torch.arange(3), torch.arange(4))
         destinations = (
             torch.empty(3, dtype=torch.float32),
             torch.empty(4, dtype=torch.float32),
         )
-        with patch.object(torch, "_foreach_copy_") as foreach_copy:
-            _foreach_copy_or_fallback_(destinations, sources)
-        foreach_copy.assert_not_called()
+        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

@@ -691,7 +691,7 @@ def _prepare_redistributed(
         prepared_views = tuple(
             _tensor_region_view(prepared, span.region).reshape(-1) for span in spans
         )
-        _foreach_copy_or_fallback_(packed_views, prepared_views)
+        _batched_copy_(packed_views, prepared_views)
 
 
 def _compute_redistributed(
@@ -729,7 +729,7 @@ def _compute_redistributed(
             ].view(span.region.shape)
             for span in received_spans
         )
-        _foreach_copy_or_fallback_(compute_views, received_views)
+        _batched_copy_(compute_views, received_views)
 
         compute(item, compute_tensor)
 
@@ -744,7 +744,7 @@ def _compute_redistributed(
             _tensor_region_view(compute_tensor, span.region).reshape(-1)
             for span in output_spans
         )
-        _foreach_copy_or_fallback_(packed_views, compute_output_views)
+        _batched_copy_(packed_views, compute_output_views)
 
 
 def _finalize_redistributed(
@@ -777,39 +777,32 @@ def _finalize_redistributed(
             ].view(span.region.shape)
             for span in spans
         )
-        _foreach_copy_or_fallback_(update_views, packed_views)
+        _batched_copy_(update_views, packed_views)
         finalize(item, update)
 
 
-def _foreach_copy_or_fallback_(
+def _batched_copy_(
     destinations: tuple[Tensor, ...],
     sources: tuple[Tensor, ...],
 ) -> None:
-    """Copy aligned tensor views with one foreach launch when supported."""
+    """Copy aligned region views with a single foreach launch.
+
+    Each destination is paired with a source built from the same
+    ``_TensorRegion``, so the two always have equal shape.
+    ``torch._foreach_copy_`` requires that -- unlike ``Tensor.copy_`` it does
+    not broadcast -- and raises if it is ever violated. Differing dtype or
+    device and non-contiguous views it handles correctly on its own, so none of
+    those need guarding here.
+    """
     if len(destinations) != len(sources):
         raise ValueError("destinations and sources must have equal length")
     if not destinations:
         return
     if len(destinations) == 1:
+        # Skip the foreach setup for the common single-span parameter.
         destinations[0].copy_(sources[0])
         return
-
-    reference_device = destinations[0].device
-    reference_dtype = destinations[0].dtype
-    foreach_compatible = all(
-        destination.layout is torch.strided
-        and source.layout is torch.strided
-        and destination.shape == source.shape
-        and destination.device == source.device == reference_device
-        and destination.dtype == source.dtype == reference_dtype
-        for destination, source in zip(destinations, sources, strict=True)
-    )
-    if foreach_compatible:
-        torch._foreach_copy_(destinations, sources)
-        return
-
-    for destination, source in zip(destinations, sources, strict=True):
-        destination.copy_(source)
+    torch._foreach_copy_(destinations, sources)
 
 
 def _tensor_region_view(tensor: Tensor, region: _TensorRegion) -> Tensor: