From 97a02f742be5961f220719bb10a9b95d0ee01c5e Mon Sep 17 00:00:00 2001 From: Mehta Date: Thu, 20 Aug 2026 14:51:03 -0400 Subject: [PATCH 1/3] Optimize DistMuon packing copies --- .../test_optimizer_reshard_runtime.py | 73 ++++++++++++++++ .../flex_shard/_optimizer_reshard_runtime.py | 85 ++++++++++++++----- 2 files changed, 138 insertions(+), 20 deletions(-) create mode 100644 tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py diff --git a/tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py b/tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py new file mode 100644 index 0000000000..d87382e13d --- /dev/null +++ b/tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py @@ -0,0 +1,73 @@ +# 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 ( + _foreach_copy_or_fallback_, +) + + +class TestForeachCopyOrFallback(unittest.TestCase): + def test_rejects_mismatched_lists(self): + with self.assertRaisesRegex(ValueError, "equal length"): + _foreach_copy_or_fallback_((torch.empty(1),), ()) + + def test_empty_and_single_copy_bypass_foreach(self): + _foreach_copy_or_fallback_((), ()) + + 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,)) + 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: + _foreach_copy_or_fallback_(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]) + + _foreach_copy_or_fallback_(destinations, sources) + + torch.testing.assert_close(destination_base, source_base) + + def test_incompatible_dtype_uses_copy_fallback(self): + 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() + for destination, source in zip(destinations, sources, strict=True): + torch.testing.assert_close(destination, source.to(destination.dtype)) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py b/torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py index d5245b2df2..a67c52b0e5 100644 --- a/torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py +++ b/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 + ) + _foreach_copy_or_fallback_(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 + ) + _foreach_copy_or_fallback_(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 + ) + _foreach_copy_or_fallback_(packed_views, compute_output_views) def _finalize_redistributed( @@ -757,16 +770,48 @@ 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 + ) + _foreach_copy_or_fallback_(update_views, packed_views) finalize(item, update) +def _foreach_copy_or_fallback_( + destinations: tuple[Tensor, ...], + sources: tuple[Tensor, ...], +) -> None: + """Copy aligned tensor views with one foreach launch when supported.""" + 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 + + 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) + + def _tensor_region_view(tensor: Tensor, region: _TensorRegion) -> Tensor: view = tensor[ tuple( From 64c8fd35c862efc06866dfa5c929d36174d4bacf Mon Sep 17 00:00:00 2001 From: Mehta Date: Sun, 23 Aug 2026 13:49:36 -0400 Subject: [PATCH 2/3] Simplify DistMuon batched copies --- .../test_optimizer_reshard_runtime.py | 35 ++++++++++------- .../flex_shard/_optimizer_reshard_runtime.py | 39 ++++++++----------- 2 files changed, 38 insertions(+), 36 deletions(-) diff --git a/tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py b/tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py index d87382e13d..d42fd13290 100644 --- a/tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py +++ b/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,33 @@ 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() diff --git a/torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py b/torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py index a67c52b0e5..129e5c5f4a 100644 --- a/torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py +++ b/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: From e024f340cf79e9886ccc4bbf9d98e5ddb95ac935 Mon Sep 17 00:00:00 2001 From: Mehta Date: Sun, 23 Aug 2026 14:05:40 -0400 Subject: [PATCH 3/3] Simplify DistMuon batched copies --- .../flex_shard/test_optimizer_reshard_runtime.py | 3 --- .../flex_shard/_optimizer_reshard_runtime.py | 11 +---------- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py b/tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py index d42fd13290..c93c257c80 100644 --- a/tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py +++ b/tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py @@ -55,9 +55,6 @@ def test_foreach_copy_supports_noncontiguous_views(self): torch.testing.assert_close(destination_base, source_base) 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), diff --git a/torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py b/torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py index 129e5c5f4a..876d5e2817 100644 --- a/torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py +++ b/torchtitan/distributed/flex_shard/_optimizer_reshard_runtime.py @@ -785,21 +785,12 @@ def _batched_copy_( destinations: tuple[Tensor, ...], sources: tuple[Tensor, ...], ) -> None: - """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. - """ + """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: - # Skip the foreach setup for the common single-span parameter. destinations[0].copy_(sources[0]) return torch._foreach_copy_(destinations, sources)