From 93f15a54edee5529731d498c7f1b1035d47e8abc Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Fri, 31 Jul 2026 04:15:27 -0700 Subject: [PATCH] Add standalone distributed Muon optimizer --- .../test_bucketed_optimizer_redistribution.py | 259 ++++ ...est_deepseek_v3_distributed_muon_config.py | 194 +++ tests/unit_tests/test_distributed_muon.py | 1027 ++++++++++++++ .../unit_tests/test_distributed_muon_math.py | 109 ++ tests/unit_tests/test_muon_parameter_prep.py | 189 +++ .../distributed_optimizers/__init__.py | 1 + .../bucketed_redistribution.py | 1226 +++++++++++++++++ .../components/distributed_optimizers/muon.py | 650 +++++++++ .../muon_parameter_prep.py | 208 +++ torchtitan/components/optimizer.py | 41 +- .../models/deepseek_v3/config_registry.py | 167 ++- 11 files changed, 4062 insertions(+), 9 deletions(-) create mode 100644 tests/unit_tests/test_bucketed_optimizer_redistribution.py create mode 100644 tests/unit_tests/test_deepseek_v3_distributed_muon_config.py create mode 100644 tests/unit_tests/test_distributed_muon.py create mode 100644 tests/unit_tests/test_distributed_muon_math.py create mode 100644 tests/unit_tests/test_muon_parameter_prep.py create mode 100644 torchtitan/components/distributed_optimizers/__init__.py create mode 100644 torchtitan/components/distributed_optimizers/bucketed_redistribution.py create mode 100644 torchtitan/components/distributed_optimizers/muon.py create mode 100644 torchtitan/components/distributed_optimizers/muon_parameter_prep.py diff --git a/tests/unit_tests/test_bucketed_optimizer_redistribution.py b/tests/unit_tests/test_bucketed_optimizer_redistribution.py new file mode 100644 index 0000000000..4ebc0a4ed2 --- /dev/null +++ b/tests/unit_tests/test_bucketed_optimizer_redistribution.py @@ -0,0 +1,259 @@ +# 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 dataclasses import dataclass +from unittest.mock import Mock, patch + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor +from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( + _build_bucket_plans, + _build_owned_redistribution_plan, + _lower_packed_all_to_all, + _MatrixBlock, + _MatrixBlockRoute, + _PackedAllToAllSchedule, + _RedistributionGroup, + _RedistributionPlan, + BucketSpec, +) + + +class TestBucketedOptimizerRedistribution(unittest.TestCase): + def test_bucket_planner_preserves_empty_local_storage_block(self): + @dataclass(frozen=True) + class Item: + fqn: str + tensor: DTensor + + tensor = Mock(spec=DTensor) + tensor.shape = torch.Size((2, 3)) + tensor.to_local.return_value = torch.empty(0, 3) + item = Item("layers.0.weight", tensor) + blocks = ( + ((3,), _MatrixBlock(offsets=(0, 0), shape=(2, 3))), + ((7,), _MatrixBlock(offsets=(2, 0), shape=(0, 3))), + ) + group = _RedistributionGroup( + process_group=object(), + participants=(3, 7), + local_participant=7, + ) + mesh = Mock(spec=DeviceMesh) + mesh.ndim = 1 + + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "get_process_group_ranks", + return_value=[3, 7], + ), patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution." + "_redistribution_group", + return_value=group, + ), patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution." + "_dtensor_storage_blocks", + return_value=blocks, + ): + result = _build_bucket_plans( + (item,), + ( + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={item.fqn: 0}, + mesh=mesh, + ), + ), + fqn=lambda value: value.fqn, + compute_locally=lambda _value: False, + storage_dtensor=lambda value: value.tensor, + ) + + plan = result.plans[0] + self.assertEqual(result.ordered_items, (item,)) + self.assertEqual( + plan.storage_to_compute_schedule.input_spans_by_parameter[0][0].numel, + 0, + ) + self.assertEqual( + plan.compute_to_storage_schedule.output_spans_by_parameter[0][0].numel, + 0, + ) + + def test_transport_neutral_routes_lower_to_packed_all_to_all(self): + first = _MatrixBlock(offsets=(0, 0), shape=(2, 3)) + second = _MatrixBlock(offsets=(2, 0), shape=(2, 3)) + plan = _RedistributionPlan( + participants=(3, 7), + logical_shape=(4, 3), + storage_to_compute_routes=( + _MatrixBlockRoute(first, (3,), (7,)), + _MatrixBlockRoute(second, (7,), (7,)), + ), + compute_to_storage_routes=( + _MatrixBlockRoute(first, (7,), (3,)), + _MatrixBlockRoute(second, (7,), (7,)), + ), + ) + + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "get_process_group_ranks", + return_value=[3, 7], + ): + forward = _lower_packed_all_to_all( + (plan,), + direction="storage_to_compute", + process_group=object(), + local_participant=7, + ) + reverse = _lower_packed_all_to_all( + (plan,), + direction="compute_to_storage", + process_group=object(), + local_participant=7, + ) + + self.assertIsInstance(forward, _PackedAllToAllSchedule) + self.assertEqual(forward.input_split_sizes, (0, 6)) + self.assertEqual(forward.output_split_sizes, (6, 6)) + self.assertEqual( + tuple(span.block for span in forward.output_spans_by_parameter[0]), + (first, second), + ) + self.assertEqual(reverse.input_split_sizes, (6, 6)) + self.assertEqual(reverse.output_split_sizes, (0, 6)) + self.assertEqual( + tuple(span.block for span in reverse.output_spans_by_parameter[0]), + (second,), + ) + + def test_equivalent_replicas_prefer_local_copy_source(self): + block = _MatrixBlock(offsets=(0, 0), shape=(2, 3)) + plan = _build_owned_redistribution_plan( + (((3, 7), block),), + participants=(3, 7), + owner=7, + logical_shape=(2, 3), + ) + self.assertEqual(plan.storage_to_compute_routes[0].source_participants, (3, 7)) + self.assertEqual(plan.compute_to_storage_routes[0].destination_participants, (3, 7)) + + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "get_process_group_ranks", + return_value=[3, 7], + ): + schedule = _lower_packed_all_to_all( + (plan,), + direction="storage_to_compute", + process_group=object(), + local_participant=7, + ) + + self.assertEqual(schedule.input_split_sizes, (0, 6)) + self.assertEqual(schedule.output_split_sizes, (0, 6)) + + def test_copy_fanout_and_reduction_routes_are_explicit(self): + block = _MatrixBlock(offsets=(0, 0), shape=(2, 3)) + fanout = _RedistributionPlan( + participants=(3, 7), + logical_shape=(2, 3), + storage_to_compute_routes=( + _MatrixBlockRoute(block, (3,), (3, 7)), + ), + compute_to_storage_routes=( + _MatrixBlockRoute(block, (3, 7), (3,)), + ), + ) + reduction = _RedistributionPlan( + participants=(3, 7), + logical_shape=(2, 3), + storage_to_compute_routes=( + _MatrixBlockRoute( + block, + (3, 7), + (3,), + reduce_op=dist.ReduceOp.SUM, + ), + ), + compute_to_storage_routes=( + _MatrixBlockRoute(block, (3,), (3, 7)), + ), + ) + + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "get_process_group_ranks", + return_value=[3, 7], + ): + schedule = _lower_packed_all_to_all( + (fanout,), + direction="storage_to_compute", + process_group=object(), + local_participant=3, + ) + with self.assertRaisesRegex(ValueError, "cannot lower reduction"): + _lower_packed_all_to_all( + (reduction,), + direction="storage_to_compute", + process_group=object(), + local_participant=3, + ) + + self.assertEqual(schedule.input_split_sizes, (6, 6)) + self.assertEqual(schedule.output_split_sizes, (6, 0)) + + def test_routes_require_an_exact_nonoverlapping_partition(self): + def plan(blocks): + routes = tuple( + _MatrixBlockRoute(block, (3,), (7,)) for block in blocks + ) + return _RedistributionPlan( + participants=(3, 7), + logical_shape=(2, 3), + storage_to_compute_routes=routes, + compute_to_storage_routes=routes, + ) + + invalid_partitions = ( + ( + (_MatrixBlock((0, 0), (3, 3)),), + ValueError, + "outside", + ), + ( + ( + _MatrixBlock((0, 0), (2, 3)), + _MatrixBlock((0, 0), (2, 3)), + ), + NotImplementedError, + "overlapping", + ), + ( + (_MatrixBlock((0, 0), (1, 3)),), + ValueError, + "do not cover", + ), + ) + for blocks, error, message in invalid_partitions: + with self.subTest(message=message), self.assertRaisesRegex(error, message): + plan(blocks) + + split_routes = ( + _MatrixBlockRoute(_MatrixBlock((0, 0), (1, 3)), (3,), (3,)), + _MatrixBlockRoute(_MatrixBlock((1, 0), (1, 3)), (7,), (7,)), + ) + with self.assertRaisesRegex(ValueError, "compute destination"): + _RedistributionPlan( + participants=(3, 7), + logical_shape=(2, 3), + storage_to_compute_routes=split_routes, + compute_to_storage_routes=split_routes, + ) diff --git a/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py b/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py new file mode 100644 index 0000000000..b1f93a80bf --- /dev/null +++ b/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py @@ -0,0 +1,194 @@ +# 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 + +import torch +from torch.distributed.tensor import Shard +from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( + BucketConfig, + assign_balanced_owners, +) +from torchtitan.components.distributed_optimizers.muon import Owned +from torchtitan.components.distributed_optimizers.muon_parameter_prep import ( + BatchedMatrixComputeView, + MuonComputeSharding, +) +from torchtitan.components.optimizer import ( + OptimizersContainer, + register_moe_load_balancing_hook, +) +from torchtitan.models.deepseek_v3.config_registry import ( + deepseek_v3_16b_distributed_muon, +) + + +class TestDeepSeekV3DistributedMuonConfig(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.config = deepseek_v3_16b_distributed_muon() + assert cls.config.model_spec is not None + with torch.device("meta"): + cls.model = cls.config.model_spec.model.build() + + @classmethod + def tearDownClass(cls): + del cls.model + + def test_balanced_owner_assignment(self): + self.assertEqual( + assign_balanced_owners( + [("a", "b"), ("c",)], + {"a": 8, "b": 4, "c": 4}, + num_ranks=2, + initial_memory_by_rank=(0, 4), + ), + ({"a": 0, "b": 1}, {"c": 0}), + ) + + owners = {"a": 0} + spec = BucketConfig( + patterns=("a",), + owner_rank_by_fqn=owners, + mesh_axis="dp_shard", + ) + owners["a"] = 1 + self.assertEqual(spec.owner_rank_by_fqn, {"a": 0}) + + def test_parameter_routing(self): + optimizer_config = self.config.optimizer + impl_kwargs = OptimizersContainer._build_impl_kwargs(optimizer_config) + groups_by_optimizer, _ = OptimizersContainer._build_param_groups( + self.model, + optimizer_config.param_groups, + impl_kwargs, + ) + + model_names = set(dict(self.model.named_parameters())) + self.assertEqual(len(model_names), 377) + + expected_muon_names = set() + for suffix, count in ( + (".attention.wq.weight", 27), + (".attention.wkv_a.weight", 27), + (".attention.wkv_b.weight", 27), + (".moe.routed_experts.inner_experts.w1_EFD", 26), + (".moe.routed_experts.inner_experts.w2_EDF", 26), + (".moe.routed_experts.inner_experts.w3_EFD", 26), + ): + names = {name for name in model_names if name.endswith(suffix)} + self.assertEqual(len(names), count, suffix) + expected_muon_names.update(names) + + muon_groups = groups_by_optimizer["DistributedMuon"] + muon_names = { + name for group in muon_groups for name in group["param_names"] + } + self.assertEqual(len(muon_names), 159) + self.assertEqual(muon_names, expected_muon_names) + + adamw_names = { + name + for group in groups_by_optimizer["AdamW"] + for name in group["param_names"] + } + self.assertEqual(len(adamw_names), 218) + self.assertEqual(adamw_names, model_names - expected_muon_names) + self.assertEqual(len(muon_names | adamw_names), 377) + self.assertFalse(muon_names & adamw_names) + wo_names = { + name for name in model_names if name.endswith(".attention.wo.weight") + } + self.assertEqual(len(wo_names), 27) + self.assertTrue(wo_names <= adamw_names) + + groups_by_suffix = { + suffix: next( + group + for group in muon_groups + if group["param_names"][0].endswith(suffix) + ) + for suffix in ( + ".attention.wq.weight", + ".attention.wkv_a.weight", + ".attention.wkv_b.weight", + ".moe.routed_experts.inner_experts.w1_EFD", + ) + } + expected_compute_sharding = { + ".attention.wq.weight": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=16, + matrices_flattened_into_dim=0, + ), + placement=Shard(0), + ), + ".attention.wkv_a.weight": MuonComputeSharding(placement=Owned()), + ".attention.wkv_b.weight": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=16, + matrices_flattened_into_dim=0, + ), + placement=Shard(0), + ), + ".moe.routed_experts.inner_experts.w1_EFD": MuonComputeSharding( + placement=Shard(0) + ), + } + for suffix, group in groups_by_suffix.items(): + self.assertEqual( + group["compute_sharding"], + expected_compute_sharding[suffix], + ) + + def test_bucket_and_parallelism_config(self): + optimizer_config = self.config.optimizer + bucket_configs = optimizer_config.optimizer_init_kwargs["DistributedMuon"][ + "bucket_configs" + ] + self.assertEqual( + set(optimizer_config.optimizer_init_kwargs["DistributedMuon"]), + {"bucket_configs"}, + ) + self.assertEqual( + [config.name for config in bucket_configs], + [f"layers.{layer_id}" for layer_id in range(27)], + ) + for layer_id, config in enumerate(bucket_configs): + prefix = f"layers.{layer_id}" + expected = tuple( + f"{prefix}.attention.{projection}.weight" + for projection in ("wq", "wkv_a", "wkv_b") + ) + if layer_id: + expected += tuple( + f"{prefix}.moe.routed_experts.inner_experts.{projection}" + for projection in ("w1_EFD", "w2_EDF", "w3_EFD") + ) + self.assertEqual(config.patterns, expected) + self.assertEqual(config.mesh_axis, "dp_shard") + self.assertEqual( + config.owner_rank_by_fqn, + {f"{prefix}.attention.wkv_a.weight": layer_id % 8}, + ) + + parallelism = self.config.parallelism + self.assertEqual(parallelism.data_parallel_replicate_degree, 1) + self.assertEqual(parallelism.data_parallel_shard_degree, 8) + self.assertEqual(parallelism.expert_parallel_degree, 4) + self.assertEqual(parallelism.tensor_parallel_degree, 1) + self.assertEqual(parallelism.context_parallel_degree, 1) + self.assertEqual(parallelism.pipeline_parallel_degree, 1) + self.assertFalse(parallelism.enable_sequence_parallel) + self.assertEqual(parallelism.spmd_backend, "spmd_types") + self.assertIs( + self.config.model_spec.post_optimizer_build_fn, + register_moe_load_balancing_hook, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/test_distributed_muon.py b/tests/unit_tests/test_distributed_muon.py new file mode 100644 index 0000000000..5162e9d53f --- /dev/null +++ b/tests/unit_tests/test_distributed_muon.py @@ -0,0 +1,1027 @@ +# 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 +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import distribute_tensor, DTensor, Replicate, Shard +from torch.distributed.tensor.placement_types import _StridedShard +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) +from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( + BucketConfig, + BucketSpec, +) +from torchtitan.components.distributed_optimizers.muon import ( + _has_dim0_sharded_storage, + _has_replicated_storage, + DistributedMuon, + Owned, +) +from torchtitan.components.checkpoint_utils import ( + get_flat_optim_state_dict, + init_optim_state, +) +from torchtitan.components.distributed_optimizers.muon_parameter_prep import ( + build_distributed_muon, + BatchedMatrixComputeView, + MuonComputeSharding, +) + + +class TestDistributedMuonStoragePolicy(unittest.TestCase): + def test_rejects_placement_subclasses(self): + class UnsupportedShard(Shard): + pass + + class UnsupportedReplicate(Replicate): + pass + + class FakeParameter: + ndim = 3 + + parameter = FakeParameter() + parameter.placements = (UnsupportedShard(0),) + self.assertFalse(_has_dim0_sharded_storage(parameter)) + parameter.placements = (UnsupportedReplicate(),) + self.assertFalse(_has_replicated_storage(parameter)) + +class _DistributedMuonTestBase(DTensorTestBase): + @property + def world_size(self): + return 2 + + @property + def device_type(self): + return "cuda" + + @property + def mesh(self): + if not hasattr(self, "_mesh"): + self._mesh = init_device_mesh( + self.device_type, + (self.world_size,), + mesh_dim_names=("dp_shard",), + ) + return self._mesh + + @property + def device(self): + return torch.device("cuda", self.rank) + + def _parameter(self, value: torch.Tensor) -> torch.nn.Parameter: + return torch.nn.Parameter( + distribute_tensor(value.clone(), self.mesh, (Shard(0),)) + ) + + def _optimizer( + self, + redistributed: torch.nn.Parameter, + local_blocks: torch.nn.Parameter, + ) -> DistributedMuon: + return build_distributed_muon( + [ + { + "params": [redistributed], + "param_names": ["layers.0.redistributed"], + "compute_sharding": MuonComputeSharding(placement=Owned()), + }, + { + "params": [local_blocks], + "param_names": ["layers.0.local_blocks"], + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=2, matrices_flattened_into_dim=0 + ), + placement=Shard(0), + ), + }, + ], + bucket_configs=[ + BucketConfig( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.redistributed": 1}, + mesh_axis="dp_shard", + name="layers.0", + ) + ], + lr=0.03, + weight_decay=0.2, + momentum=0.8, + nesterov=True, + ns_steps=2, + ) + + def _set_grads( + self, + redistributed: torch.nn.Parameter, + local_blocks: torch.nn.Parameter, + redistributed_grad: torch.Tensor, + local_blocks_grad: torch.Tensor, + ) -> None: + redistributed.grad = distribute_tensor( + redistributed_grad.clone(), self.mesh, (Shard(0),) + ) + local_blocks.grad = distribute_tensor( + local_blocks_grad.clone(), self.mesh, (Shard(0),) + ) + + def _assert_matches_reference( + self, + optimizer: DistributedMuon, + redistributed: torch.nn.Parameter, + local_blocks: torch.nn.Parameter, + reference_optimizer: torch.optim.Muon, + reference_redistributed: torch.nn.Parameter, + reference_local_blocks: tuple[ + torch.nn.Parameter, torch.nn.Parameter + ], + ) -> None: + rank = self.mesh.get_local_rank() + expected_redistributed = reference_redistributed.detach().chunk( + self.world_size, dim=0 + )[rank] + expected_local_blocks = reference_local_blocks[rank].detach() + torch.testing.assert_close(redistributed.to_local(), expected_redistributed) + torch.testing.assert_close(local_blocks.to_local(), expected_local_blocks) + + for param in (redistributed, local_blocks): + self.assertIsInstance(param, DTensor) + self.assertEqual(param.placements, (Shard(0),)) + + redistributed_momentum = optimizer.state[redistributed]["momentum_buffer"] + self.assertIsInstance(redistributed_momentum, DTensor) + self.assertEqual(redistributed_momentum.placements, (Shard(0),)) + expected_redistributed_momentum = reference_optimizer.state[ + reference_redistributed + ]["momentum_buffer"].detach().chunk(self.world_size, dim=0)[rank] + torch.testing.assert_close( + redistributed_momentum.to_local(), expected_redistributed_momentum + ) + + local_blocks_momentum = optimizer.state[local_blocks]["momentum_buffer"] + self.assertIsInstance(local_blocks_momentum, DTensor) + self.assertEqual(local_blocks_momentum.placements, (Shard(0),)) + expected_local_blocks_momentum = reference_optimizer.state[ + reference_local_blocks[rank] + ]["momentum_buffer"].detach() + torch.testing.assert_close( + local_blocks_momentum.to_local(), expected_local_blocks_momentum + ) + + +@unittest.skipUnless(torch.cuda.device_count() >= 2, "requires two CUDA devices") +class TestDistributedMuon(_DistributedMuonTestBase): + @with_comms + def test_constructor_strictly_validates_strided_storage_shards(self): + mesh = init_device_mesh(self.device_type, (self.world_size, 1)) + + def make_parameter(value, dim): + placements = ( + _StridedShard(dim, split_factor=self.world_size), + Shard(dim), + ) + parameter = torch.nn.Parameter( + distribute_tensor(value, mesh, placements) + ) + self.assertEqual(parameter.placements, placements) + return parameter + + def build(parameter, name, compute_placement, owner_rank=None): + fqn = f"layers.0.{name}" + owners = {} if owner_rank is None else {fqn: owner_rank} + return build_distributed_muon( + [ + { + "params": [parameter], + "param_names": [fqn], + "compute_sharding": MuonComputeSharding( + placement=compute_placement + ), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn=owners, + mesh=self.mesh, + ) + ], + ) + + local_blocks = make_parameter( + torch.arange(24, device=self.device).reshape(4, 2, 3).float(), 0 + ) + optimizer = build(local_blocks, "local_blocks", Shard(0)) + self.assertIs( + optimizer._plans[0].local_items[0].param, local_blocks + ) + local_blocks.grad = distribute_tensor( + torch.ones(4, 2, 3, device=self.device), + mesh, + local_blocks.placements, + ) + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "all_to_all_single" + ) as collective: + optimizer.step() + collective.assert_not_called() + self.assertEqual( + optimizer.state[local_blocks]["momentum_buffer"].placements, + local_blocks.placements, + ) + local_blocks._local_tensor = local_blocks.to_local().clone() + with self.assertRaisesRegex(RuntimeError, "local storage changed"): + optimizer.step() + + dim1_sharded = make_parameter( + torch.arange(24, device=self.device).reshape(2, 4, 3).float(), 1 + ) + with self.assertRaisesRegex( + ValueError, "must already match storage sharding" + ): + build(dim1_sharded, "dim1_sharded", Shard(0)) + + owned = make_parameter( + torch.arange(12, device=self.device).reshape(4, 3).float(), 0 + ) + with self.assertRaisesRegex( + ValueError, "requires replicated or 1D Shard" + ): + build(owned, "owned", Owned(), owner_rank=0) + + @with_comms + def test_constructor_requires_exact_bucket_coverage_without_creating_state(self): + redistributed = self._parameter( + torch.arange(12, device=self.device).reshape(4, 3).float() + ) + local_blocks = self._parameter( + torch.arange(12, 24, device=self.device).reshape(4, 3).float() + ) + + optimizer = self._optimizer(redistributed, local_blocks) + self.assertEqual(len(optimizer.state), 0) + redistribution = optimizer._plans[0].redistribution_plans[0] + self.assertTrue( + all( + route.destination_participants == (1,) + for route in redistribution.storage_to_compute_routes + ) + ) + redistributed_before = redistributed.to_local().clone() + redistributed.grad = distribute_tensor( + torch.ones(4, 3, device=self.device), self.mesh, (Shard(0),) + ) + with self.assertRaisesRegex(RuntimeError, "every configured gradient"): + optimizer.step() + self.assertEqual(len(optimizer.state), 0) + torch.testing.assert_close(redistributed.to_local(), redistributed_before) + redistributed.grad = None + + with self.assertRaisesRegex(ValueError, "must match one bucket"): + build_distributed_muon( + [ + { + "params": [redistributed, local_blocks], + "param_names": [ + "layers.0.redistributed", + "layers.0.local_blocks", + ], + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=2, matrices_flattened_into_dim=0 + ), + placement=Shard(0), + ), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("*.redistributed",), + owner_rank_by_fqn={}, + mesh=self.mesh, + ) + ], + ) + + with self.assertRaisesRegex(ValueError, "must match one bucket"): + build_distributed_muon( + [ + { + "params": [redistributed, local_blocks], + "param_names": [ + "layers.0.redistributed", + "layers.0.local_blocks", + ], + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=2, matrices_flattened_into_dim=0 + ), + placement=Shard(0), + ), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={}, + mesh=self.mesh, + ), + BucketSpec( + patterns=("*.local_blocks",), + owner_rank_by_fqn={}, + mesh=self.mesh, + ), + ], + ) + + local_blocks_before = local_blocks.to_local().clone() + init_optim_state(optimizer) + self.assertEqual(len(optimizer.state), 2) + torch.testing.assert_close(redistributed.to_local(), redistributed_before) + torch.testing.assert_close(local_blocks.to_local(), local_blocks_before) + flat_state = get_flat_optim_state_dict(optimizer) + self.assertIn( + "state.layers.0.redistributed.momentum_buffer", flat_state + ) + self.assertIn("state.layers.0.local_blocks.momentum_buffer", flat_state) + + @with_comms + def test_constructor_rejects_storage_shards_that_split_matrices(self): + parameter = self._parameter( + torch.arange(36, device=self.device).reshape(12, 3).float() + ) + with self.assertRaisesRegex(ValueError, "not aligned"): + build_distributed_muon( + [ + { + "params": [parameter], + "param_names": ["layers.0.wq.weight"], + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=3, matrices_flattened_into_dim=0 + ), + placement=Shard(0), + ), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={}, + mesh=self.mesh, + ) + ], + ) + + @with_comms + def test_constructor_requires_valid_owner_assignments(self): + with self.assertRaises(TypeError): + Owned(0) + + first = self._parameter( + torch.arange(12, device=self.device).reshape(4, 3).float() + ) + second = self._parameter( + torch.arange(12, 24, device=self.device).reshape(4, 3).float() + ) + params = [ + { + "params": [first, second], + "param_names": ["layers.0.first", "layers.0.second"], + "compute_sharding": MuonComputeSharding(placement=Owned()), + } + ] + + with self.assertRaisesRegex(TypeError, "compute_sharding"): + build_distributed_muon( + [{"params": [first], "param_names": ["layers.0.first"]}], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={}, + mesh=self.mesh, + ) + ], + ) + + with self.assertRaisesRegex(ValueError, "batch of complete Muon matrices"): + build_distributed_muon( + [ + { + "params": [first], + "param_names": ["layers.0.first"], + "compute_sharding": MuonComputeSharding( + placement=Shard(0) + ), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={}, + mesh=self.mesh, + ) + ], + ) + + with self.assertRaisesRegex(ValueError, "owned Muon parameter"): + build_distributed_muon( + [ + { + "params": [first], + "param_names": ["layers.0.first"], + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=2 + ), + placement=Owned(), + ), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.first": 0}, + mesh=self.mesh, + ) + ], + ) + + with self.assertRaisesRegex(ValueError, "exactly cover"): + build_distributed_muon( + params, + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.first": 0}, + mesh=self.mesh, + ) + ], + ) + self.assertIn("compute_sharding", params[0]) + + with self.assertRaisesRegex(ValueError, "outside its process group"): + build_distributed_muon( + params, + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={ + "layers.0.first": 0, + "layers.0.second": self.world_size, + }, + mesh=self.mesh, + ) + ], + ) + + @with_comms + def test_constructor_rejects_cross_rank_plan_mismatch(self): + redistributed = self._parameter( + torch.arange(12, device=self.device).reshape(4, 3).float() + ) + with self.assertRaisesRegex(RuntimeError, "plans differ across ranks"): + build_distributed_muon( + [ + { + "params": [redistributed], + "param_names": ["layers.0.redistributed"], + "compute_sharding": MuonComputeSharding( + placement=Owned() + ), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.redistributed": self.rank}, + mesh=self.mesh, + ) + ], + ) + + with self.assertRaisesRegex(RuntimeError, "plans differ across ranks"): + build_distributed_muon( + [ + { + "params": [redistributed], + "param_names": ["layers.0.redistributed"], + "compute_sharding": MuonComputeSharding( + placement=Owned() + ), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.redistributed": 0}, + mesh=self.mesh, + ) + ], + lr=0.01 if self.rank == 0 else 0.02, + ) + + @with_comms + def test_constructor_accepts_uneven_storage_shards(self): + redistributed = self._parameter( + torch.arange(15, device=self.device).reshape(5, 3).float() + ) + optimizer = build_distributed_muon( + [ + { + "params": [redistributed], + "param_names": ["layers.0.redistributed"], + "compute_sharding": MuonComputeSharding(placement=Owned()), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.redistributed": 0}, + mesh=self.mesh, + ) + ], + ) + + schedule = optimizer._plans[0].storage_to_compute_schedule + self.assertEqual( + schedule.input_buffer_numel, 9 if self.rank == 0 else 6 + ) + + @with_comms + def test_shard1_owned_matches_plain_muon(self): + value = torch.arange(15, device=self.device).reshape(3, 5).float().div_(10) + placement = (Shard(1),) + parameter = torch.nn.Parameter( + distribute_tensor(value.clone(), self.mesh, placement) + ) + optimizer = build_distributed_muon( + [ + { + "params": [parameter], + "param_names": ["layers.0.weight"], + "compute_sharding": MuonComputeSharding(placement=Owned()), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.weight": 1}, + mesh=self.mesh, + ) + ], + lr=0.03, + momentum=0.8, + ns_steps=2, + ) + grad = value.flip((0, 1)).contiguous() + parameter.grad = distribute_tensor(grad, self.mesh, placement) + + reference = torch.nn.Parameter(value.clone()) + reference.grad = grad.clone() + reference_optimizer = torch.optim.Muon( + [reference], + lr=0.03, + momentum=0.8, + ns_steps=2, + ) + + all_to_all_single = dist.all_to_all_single + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution." + "dist.all_to_all_single", + wraps=all_to_all_single, + ) as collective: + optimizer.step() + reference_optimizer.step() + + self.assertEqual(collective.call_count, 2) + expected_parameter = distribute_tensor( + reference.detach(), self.mesh, placement + ) + expected_momentum = distribute_tensor( + reference_optimizer.state[reference]["momentum_buffer"], + self.mesh, + placement, + ) + momentum = optimizer.state[parameter]["momentum_buffer"] + self.assertEqual(parameter.placements, placement) + self.assertEqual(momentum.placements, placement) + torch.testing.assert_close( + parameter.to_local(), expected_parameter.to_local() + ) + torch.testing.assert_close( + momentum.to_local(), + expected_momentum.to_local(), + ) + + @with_comms + def test_replicated_storage_matches_plain_muon_without_redistribution(self): + values = [ + torch.arange(offset, offset + 12, device=self.device) + .reshape(4, 3) + .float() + .div_(10) + for offset in (1, 13) + ] + owned, batched = ( + torch.nn.Parameter( + distribute_tensor(value.clone(), self.mesh, (Replicate(),)) + ) + for value in values + ) + optimizer = build_distributed_muon( + [ + { + "params": [owned], + "param_names": ["layers.0.owned"], + "compute_sharding": MuonComputeSharding(placement=Owned()), + }, + { + "params": [batched], + "param_names": ["layers.0.batched"], + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=2, matrices_flattened_into_dim=0 + ), + placement=Shard(0), + ), + }, + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={}, + mesh=self.mesh, + ) + ], + lr=0.03, + weight_decay=0.2, + momentum=0.8, + nesterov=True, + ns_steps=2, + ) + + grads = [value.flip(0).contiguous() for value in values] + for param, grad in zip((owned, batched), grads, strict=True): + param.grad = distribute_tensor(grad, self.mesh, (Replicate(),)) + + references = [ + torch.nn.Parameter(values[0].clone()), + torch.nn.Parameter(values[1].view(2, 2, 3).clone()), + ] + reference_optimizer = torch.optim.Muon( + references, + lr=0.03, + weight_decay=0.2, + momentum=0.8, + nesterov=True, + ns_steps=2, + ) + references[0].grad = grads[0] + references[1].grad = grads[1].view(2, 2, 3) + + all_to_all_single = dist.all_to_all_single + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "all_to_all_single", + wraps=all_to_all_single, + ) as collective: + optimizer.step() + reference_optimizer.step() + + collective.assert_not_called() + for param, reference in zip( + (owned, batched), references, strict=True + ): + self.assertEqual(param.placements, (Replicate(),)) + self.assertEqual(param.grad.placements, (Replicate(),)) + momentum = optimizer.state[param]["momentum_buffer"] + self.assertEqual(momentum.placements, (Replicate(),)) + torch.testing.assert_close( + param.to_local(), reference.view(param.shape) + ) + torch.testing.assert_close( + momentum.to_local(), + reference_optimizer.state[reference]["momentum_buffer"].view( + param.shape + ), + ) + + @with_comms + def test_step_matches_plain_muon_and_continues_from_state_dict(self): + redistributed_value = ( + torch.arange(12, device=self.device) + .reshape(4, 3) + .float() + .div_(10) + .add_(1) + ) + local_blocks_value = ( + torch.arange(12, 24, device=self.device) + .reshape(4, 3) + .float() + .div_(10) + ) + redistributed = self._parameter(redistributed_value) + local_blocks = self._parameter(local_blocks_value) + optimizer = self._optimizer(redistributed, local_blocks) + self.assertEqual(len(optimizer.state), 0) + + reference_redistributed = torch.nn.Parameter(redistributed_value.clone()) + reference_local_blocks = tuple( + torch.nn.Parameter(block.clone()) + for block in local_blocks_value.chunk(self.world_size, dim=0) + ) + reference_optimizer = torch.optim.Muon( + [reference_redistributed, *reference_local_blocks], + lr=0.03, + weight_decay=0.2, + momentum=0.8, + nesterov=True, + ns_steps=2, + ) + + first_redistributed_grad = ( + torch.arange(1, 13, device=self.device) + .reshape(4, 3) + .float() + .div_(17) + ) + first_local_blocks_grad = ( + torch.arange(13, 25, device=self.device) + .reshape(4, 3) + .float() + .div_(19) + ) + self._set_grads( + redistributed, + local_blocks, + first_redistributed_grad, + first_local_blocks_grad, + ) + reference_redistributed.grad = first_redistributed_grad.clone() + for parameter, grad in zip( + reference_local_blocks, + first_local_blocks_grad.chunk(self.world_size, dim=0), + ): + parameter.grad = grad.clone() + + all_to_all_single = dist.all_to_all_single + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "all_to_all_single", + wraps=all_to_all_single, + ) as collective: + optimizer.step() + self.assertEqual(collective.call_count, 2) + reference_optimizer.step() + self._assert_matches_reference( + optimizer, + redistributed, + local_blocks, + reference_optimizer, + reference_redistributed, + reference_local_blocks, + ) + + state_dict = optimizer.state_dict() + self.assertTrue( + all( + "compute_sharding" not in group + and "_compute_placement" not in group + for group in state_dict["param_groups"] + ) + ) + resumed_redistributed = self._parameter(reference_redistributed.detach()) + resumed_local_blocks = self._parameter( + torch.cat([parameter.detach() for parameter in reference_local_blocks]) + ) + resumed_optimizer = self._optimizer( + resumed_redistributed, resumed_local_blocks + ) + resumed_optimizer.load_state_dict(state_dict) + + second_redistributed_grad = first_redistributed_grad.flip(0).contiguous() + second_local_blocks_grad = first_local_blocks_grad.flip(0).contiguous() + self._set_grads( + resumed_redistributed, + resumed_local_blocks, + second_redistributed_grad, + second_local_blocks_grad, + ) + reference_redistributed.grad = second_redistributed_grad.clone() + for parameter, grad in zip( + reference_local_blocks, + second_local_blocks_grad.chunk(self.world_size, dim=0), + ): + parameter.grad = grad.clone() + + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "all_to_all_single", + wraps=all_to_all_single, + ) as collective: + resumed_optimizer.step() + self.assertEqual(collective.call_count, 2) + reference_optimizer.step() + self._assert_matches_reference( + resumed_optimizer, + resumed_redistributed, + resumed_local_blocks, + reference_optimizer, + reference_redistributed, + reference_local_blocks, + ) + + +@unittest.skipUnless(torch.cuda.device_count() >= 4, "requires four CUDA devices") +class TestDistributedMuonBucketMeshes(_DistributedMuonTestBase): + @property + def world_size(self): + return 4 + + @property + def mesh(self): + if not hasattr(self, "_mesh"): + self._mesh = init_device_mesh( + self.device_type, + (2, 2), + mesh_dim_names=("fsdp", "tp"), + ) + return self._mesh + + @with_comms + def test_distinct_bucket_meshes_use_mesh_local_owners(self): + fsdp_mesh = self.mesh["fsdp"] + tp_mesh = self.mesh["tp"] + meshes = (fsdp_mesh, tp_mesh) + values = ( + torch.arange(15, device=self.device).reshape(5, 3).float().div_(10), + torch.arange(20, device=self.device).reshape(4, 5).float().div_(10), + ) + params = [ + torch.nn.Parameter( + distribute_tensor(value.clone(), mesh, (Shard(0),)) + ) + for value, mesh in zip(values, meshes, strict=True) + ] + names = ("layers.0.fsdp", "layers.1.tp") + optimizer = build_distributed_muon( + [ + { + "params": [param], + "param_names": [name], + "compute_sharding": MuonComputeSharding(placement=Owned()), + } + for param, name in zip(params, names, strict=True) + ], + bucket_spec=[ + BucketSpec( + patterns=(name,), + owner_rank_by_fqn={name: 1}, + mesh=mesh, + ) + for name, mesh in zip(names, meshes, strict=True) + ], + ns_steps=1, + ) + + for param, value, mesh in zip(params, values, meshes, strict=True): + param.grad = distribute_tensor(torch.ones_like(value), mesh, (Shard(0),)) + + all_to_all_single = dist.all_to_all_single + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "all_to_all_single", + wraps=all_to_all_single, + ) as collective: + optimizer.step() + + for plan, mesh in zip(optimizer._plans, meshes, strict=True): + participants = tuple(dist.get_process_group_ranks(mesh.get_group())) + route = plan.redistribution_plans[0].storage_to_compute_routes[0] + self.assertEqual(route.destination_participants, (participants[1],)) + self.assertEqual( + sum( + call.kwargs["group"] is mesh.get_group() + for call in collective.call_args_list + ), + 2, + ) + + +@unittest.skipUnless(torch.cuda.device_count() >= 2, "requires two CUDA devices") +class TestDistributedMuonPipeline(_DistributedMuonTestBase): + @with_comms + def test_local_only_bucket_does_not_reuse_inflight_slot(self): + values = [ + torch.arange(offset, offset + 12, device=self.device) + .reshape(4, 3) + .float() + .div_(10) + for offset in (0, 12, 24) + ] + distributed_0, local_blocks, distributed_2 = map(self._parameter, values) + optimizer = build_distributed_muon( + [ + { + "params": [distributed_0, distributed_2], + "param_names": [ + "layers.0.redistributed", + "layers.2.redistributed", + ], + "compute_sharding": MuonComputeSharding(placement=Owned()), + }, + { + "params": [local_blocks], + "param_names": ["layers.1.local_blocks"], + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=2, matrices_flattened_into_dim=0 + ), + placement=Shard(0), + ), + }, + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.redistributed": 0}, + mesh=self.mesh, + ), + BucketSpec( + patterns=("layers.1.*",), + owner_rank_by_fqn={}, + mesh=self.mesh, + ), + BucketSpec( + patterns=("layers.2.*",), + owner_rank_by_fqn={"layers.2.redistributed": 0}, + mesh=self.mesh, + ), + ], + lr=0.03, + momentum=0.8, + ns_steps=1, + ) + grads = [torch.full_like(value, index + 1) for index, value in enumerate(values)] + for param, grad in zip( + (distributed_0, local_blocks, distributed_2), grads, strict=True + ): + param.grad = distribute_tensor(grad, self.mesh, (Shard(0),)) + + rank = self.mesh.get_local_rank() + references = [ + torch.nn.Parameter(values[0].clone()), + torch.nn.Parameter(values[1].chunk(self.world_size, dim=0)[rank].clone()), + torch.nn.Parameter(values[2].clone()), + ] + reference = torch.optim.Muon( + references, lr=0.03, momentum=0.8, ns_steps=1 + ) + references[0].grad = grads[0].clone() + references[1].grad = grads[1].chunk(self.world_size, dim=0)[rank].clone() + references[2].grad = grads[2].clone() + + all_to_all_single = dist.all_to_all_single + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "all_to_all_single", + wraps=all_to_all_single, + ) as collective: + optimizer.step() + reference.step() + + self.assertEqual(collective.call_count, 4) + splits = [ + ( + tuple(call.kwargs["input_split_sizes"]), + tuple(call.kwargs["output_split_sizes"]), + ) + for call in collective.call_args_list + ] + self.assertEqual(splits[0], splits[1]) + self.assertEqual(splits[2], splits[3]) + self.assertNotEqual(splits[0], splits[2]) + torch.testing.assert_close( + distributed_0.to_local(), references[0].chunk(self.world_size, dim=0)[rank] + ) + torch.testing.assert_close(local_blocks.to_local(), references[1]) + torch.testing.assert_close( + distributed_2.to_local(), references[2].chunk(self.world_size, dim=0)[rank] + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/test_distributed_muon_math.py b/tests/unit_tests/test_distributed_muon_math.py new file mode 100644 index 0000000000..63d8df5106 --- /dev/null +++ b/tests/unit_tests/test_distributed_muon_math.py @@ -0,0 +1,109 @@ +# 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 + +import torch +from torchtitan.components.distributed_optimizers.muon import _compute_muon_update + + +class TestDistributedMuonMath(unittest.TestCase): + def test_two_steps_match_torch_muon(self): + optimizer_kwargs = { + "lr": 0.03, + "weight_decay": 0.2, + "momentum": 0.8, + "nesterov": True, + "ns_coefficients": (3.4445, -4.7750, 2.0315), + "eps": 1e-7, + "ns_steps": 3, + "adjust_lr_fn": "original", + } + + for shape in ((3, 5), (5, 3)): + with self.subTest(shape=shape): + generator = torch.Generator().manual_seed(4) + initial = torch.randn( + shape, generator=generator, dtype=torch.bfloat16 + ) + gradients = [ + torch.randn(shape, generator=generator, dtype=torch.bfloat16) + for _ in range(2) + ] + + reference_param = torch.nn.Parameter(initial.clone()) + reference = torch.optim.Muon( + [reference_param], **optimizer_kwargs + ) + actual_param = initial.clone() + actual_momentum = torch.zeros_like(actual_param) + + for gradient in gradients: + reference_param.grad = gradient.clone() + reference.step() + + actual_momentum.lerp_( + gradient, 1 - optimizer_kwargs["momentum"] + ) + prepared = torch.lerp( + gradient, + actual_momentum, + optimizer_kwargs["momentum"], + ) + update = _compute_muon_update( + prepared, + out=torch.empty_like(prepared), + lr=optimizer_kwargs["lr"], + ns_coefficients=optimizer_kwargs["ns_coefficients"], + ns_steps=optimizer_kwargs["ns_steps"], + eps=optimizer_kwargs["eps"], + adjust_lr_fn=optimizer_kwargs["adjust_lr_fn"], + ) + actual_param.mul_( + 1 + - optimizer_kwargs["lr"] + * optimizer_kwargs["weight_decay"] + ) + actual_param.add_(update) + + torch.testing.assert_close( + actual_param, reference_param, rtol=0, atol=0 + ) + torch.testing.assert_close( + actual_momentum, + reference.state[reference_param]["momentum_buffer"], + rtol=0, + atol=0, + ) + + def test_batched_update_matches_independent_matrices(self): + kwargs = { + "lr": 0.03, + "ns_coefficients": (3.4445, -4.7750, 2.0315), + "ns_steps": 3, + "eps": 1e-7, + "adjust_lr_fn": "match_rms_adamw", + } + + for shape in ((4, 3, 5), (4, 5, 3)): + with self.subTest(shape=shape): + generator = torch.Generator().manual_seed(5) + prepared = torch.randn(shape, generator=generator) + batched = _compute_muon_update( + prepared, out=torch.empty_like(prepared), **kwargs + ) + independent = torch.stack( + [ + _compute_muon_update( + matrix, out=torch.empty_like(matrix), **kwargs + ) + for matrix in prepared + ] + ) + + torch.testing.assert_close( + batched, independent, rtol=0, atol=0 + ) diff --git a/tests/unit_tests/test_muon_parameter_prep.py b/tests/unit_tests/test_muon_parameter_prep.py new file mode 100644 index 0000000000..cea29a7a6d --- /dev/null +++ b/tests/unit_tests/test_muon_parameter_prep.py @@ -0,0 +1,189 @@ +# 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 import mock + +import torch +from torch.distributed.tensor import DTensor, Shard +from torch.distributed.tensor.placement_types import _StridedShard +from torchtitan.components.distributed_optimizers.muon import DistributedMuon, Owned +from torchtitan.components.distributed_optimizers.muon_parameter_prep import ( + BatchedMatrixComputeView, + build_distributed_muon, + MuonComputeSharding, +) + + +class TestMuonParameterPrep(unittest.TestCase): + def test_batched_matrix_view_validation(self): + for num_matrices in (0, -1, True, 1.5): + with self.subTest(num_matrices=num_matrices): + with self.assertRaisesRegex(ValueError, "positive integer"): + BatchedMatrixComputeView(num_matrices, 0) + for matrices_flattened_into_dim in (True, "0"): + with self.subTest( + matrices_flattened_into_dim=matrices_flattened_into_dim + ): + with self.assertRaisesRegex(ValueError, "must be an integer"): + BatchedMatrixComputeView(3, matrices_flattened_into_dim) + with self.assertRaisesRegex( + ValueError, "only matrices_flattened_into_dim=0" + ): + BatchedMatrixComputeView(3, 1) + + def test_builder_compiles_layout_without_mutating_caller_group(self): + view = BatchedMatrixComputeView( + num_matrices=3, matrices_flattened_into_dim=0 + ) + compute_sharding = MuonComputeSharding( + view_before_placement=view, + placement=Shard(0), + ) + storage = torch.arange(24).reshape(6, 4) + other_storage = torch.empty(9, 5) + group = { + "params": [storage, other_storage], + "param_names": [ + "layers.0.wq.weight", + "layers.0.wkv_b.weight", + ], + "compute_sharding": compute_sharding, + } + identity_storage = torch.empty(4, 3) + identity_group = { + "params": [identity_storage], + "param_names": ["layers.0.wkv_a.weight"], + "compute_sharding": MuonComputeSharding(placement=Owned()), + } + bucket_spec = () + + with mock.patch.object(DistributedMuon, "__init__", return_value=None) as init: + optimizer = build_distributed_muon( + [group, identity_group], + bucket_spec=bucket_spec, + lr=0.1, + ) + + self.assertIsInstance(optimizer, DistributedMuon) + core_groups = init.call_args.args[0] + prepared = init.call_args.kwargs["_prepared_compute_views"] + init.assert_called_once_with( + core_groups, + bucket_spec=bucket_spec, + _prepared_compute_views=prepared, + lr=0.1, + ) + self.assertIsNot(core_groups[0], group) + self.assertIsNot(core_groups[1], identity_group) + self.assertIs(group["compute_sharding"], compute_sharding) + self.assertNotIn("compute_sharding", core_groups[0]) + self.assertEqual(core_groups[0]["_compute_placement"], Shard(0)) + self.assertEqual(core_groups[1]["_compute_placement"], Owned()) + self.assertFalse(any(value is view for value in core_groups[0].values())) + self.assertEqual( + prepared["layers.0.wq.weight"].global_compute_shape, + torch.Size((3, 2, 4)), + ) + self.assertEqual( + prepared["layers.0.wq.weight"].local_compute_tensor.shape, + torch.Size((3, 2, 4)), + ) + self.assertEqual( + prepared["layers.0.wkv_b.weight"].global_compute_shape, + torch.Size((3, 3, 5)), + ) + self.assertEqual( + prepared["layers.0.wkv_b.weight"].local_compute_tensor.shape, + torch.Size((3, 3, 5)), + ) + self.assertEqual( + prepared["layers.0.wq.weight"].local_compute_tensor.data_ptr(), + storage.data_ptr(), + ) + self.assertEqual( + prepared["layers.0.wkv_a.weight"].global_compute_shape, + identity_storage.shape, + ) + self.assertEqual( + prepared["layers.0.wkv_a.weight"].local_compute_tensor.shape, + identity_storage.shape, + ) + self.assertIs( + prepared["layers.0.wkv_a.weight"].local_compute_tensor, + identity_storage, + ) + + def test_builder_validates_global_shape_and_aligned_names(self): + for shape, message in ( + ((2, 3, 4), "requires rank-2 storage"), + ((5, 4), "is not divisible"), + ): + with self.subTest(shape=shape): + with self.assertRaisesRegex(ValueError, message): + build_distributed_muon( + [ + { + "params": [torch.empty(shape)], + "param_names": ["layers.0.wq.weight"], + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + 2, 0 + ), + placement=Shard(0), + ), + } + ], + bucket_spec=(), + ) + + with self.assertRaisesRegex(ValueError, "must be aligned"): + build_distributed_muon( + [ + { + "params": [torch.empty(6, 4)], + "param_names": [], + "compute_sharding": MuonComputeSharding( + placement=Shard(0) + ), + } + ], + bucket_spec=(), + ) + + def test_builder_requires_compute_sharding(self): + with self.assertRaisesRegex(TypeError, "must be a MuonComputeSharding"): + build_distributed_muon( + [{"params": [], "param_names": [], "compute_sharding": object()}], + bucket_spec=(), + ) + + def test_builder_rejects_strided_storage_shard_for_batched_matrices(self): + param = mock.Mock(spec=DTensor) + param.shape = torch.Size((6, 4)) + param.placements = ( + _StridedShard(0, split_factor=2), + Shard(0), + ) + + with self.assertRaisesRegex(ValueError, "Shard or Replicate"): + build_distributed_muon( + [ + { + "params": [param], + "param_names": ["layers.0.wq.weight"], + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView(3), + placement=Shard(0), + ), + } + ], + bucket_spec=(), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/components/distributed_optimizers/__init__.py b/torchtitan/components/distributed_optimizers/__init__.py new file mode 100644 index 0000000000..c175afdbad --- /dev/null +++ b/torchtitan/components/distributed_optimizers/__init__.py @@ -0,0 +1 @@ +"""Distributed optimizer implementations and redistribution runtimes.""" diff --git a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py new file mode 100644 index 0000000000..3357800f51 --- /dev/null +++ b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py @@ -0,0 +1,1226 @@ +# 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. + +"""Bucketed storage-to-compute redistribution for optimizer steps.""" + +from __future__ import annotations + +import fnmatch +import hashlib +import heapq +import math +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field +from types import ModuleType +from typing import Any, Generic, TypeVar + +import torch +import torch.distributed as dist +from torch import Tensor +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor, Shard + + +__all__ = ["BucketConfig", "BucketSpec", "assign_balanced_owners"] + + +@dataclass(frozen=True, slots=True) +class BucketConfig: + """Static bucket configuration resolved after runtime meshes exist.""" + + patterns: tuple[str, ...] + owner_rank_by_fqn: Mapping[str, int] + mesh_axis: str + name: str = "" + + def __post_init__(self) -> None: + object.__setattr__(self, "patterns", tuple(self.patterns)) + object.__setattr__(self, "owner_rank_by_fqn", dict(self.owner_rank_by_fqn)) + + def bind(self, mesh: DeviceMesh) -> BucketSpec: + return BucketSpec( + patterns=self.patterns, + owner_rank_by_fqn=self.owner_rank_by_fqn, + mesh=mesh, + name=self.name, + ) + + +@dataclass(frozen=True, slots=True) +class BucketSpec: + """One ordered optimizer-work bucket selected by canonical FQN. + + Patterns use case-sensitive ``fnmatch`` syntax. Every optimizer FQN must + match exactly one bucket, and sequence order controls execution order. + ``mesh`` is the bucket's exact one-dimensional communication mesh. + ``owner_rank_by_fqn`` must exactly cover parameters requiring whole-tensor + redistribution and uses mesh-local ranks. Compute-ready parameters have no + owner entry. ``name`` is diagnostic metadata only. + """ + + patterns: tuple[str, ...] + owner_rank_by_fqn: Mapping[str, int] + mesh: DeviceMesh + name: str = "" + + def __post_init__(self) -> None: + if self.mesh.ndim != 1: + raise ValueError("BucketSpec mesh must be one-dimensional") + object.__setattr__(self, "patterns", tuple(self.patterns)) + object.__setattr__(self, "owner_rank_by_fqn", dict(self.owner_rank_by_fqn)) + + +def _bind_bucket_configs( + configs: Sequence[BucketConfig], + storage_by_fqn: Mapping[str, DTensor], +) -> tuple[BucketSpec, ...]: + specs = [] + for config in configs: + candidates = tuple(config.owner_rank_by_fqn) or tuple( + fqn + for fqn in storage_by_fqn + if any(fnmatch.fnmatchcase(fqn, pattern) for pattern in config.patterns) + ) + if not candidates: + raise ValueError(f"bucket {config.name!r} matched no storage tensor") + + meshes = [] + for fqn in candidates: + if fqn not in storage_by_fqn: + raise ValueError(f"bucket {config.name!r} references unknown {fqn!r}") + storage_mesh = storage_by_fqn[fqn].device_mesh + if storage_mesh.mesh_dim_names is None or ( + config.mesh_axis not in storage_mesh.mesh_dim_names + ): + raise ValueError( + f"bucket {config.name!r} mesh axis {config.mesh_axis!r} " + f"is not present on storage for {fqn!r}" + ) + meshes.append(storage_mesh[config.mesh_axis]) + + mesh = meshes[0] + if any(not torch.equal(candidate.mesh, mesh.mesh) for candidate in meshes[1:]): + raise ValueError( + f"bucket {config.name!r} resolves to inconsistent communication meshes" + ) + specs.append(config.bind(mesh)) + return tuple(specs) + + +def assign_balanced_owners( + bucket_fqns: Sequence[Sequence[str]], + memory_estimate_by_fqn: Mapping[str, int], + *, + num_ranks: int, + initial_memory_by_rank: Sequence[int] | None = None, +) -> tuple[dict[str, int], ...]: + """Greedily balance selected parameters across group-local ranks.""" + initial_memory_by_rank = initial_memory_by_rank or (0,) * num_ranks + rank_loads = list(zip(initial_memory_by_rank, range(num_ranks), strict=True)) + heapq.heapify(rank_loads) + owners_by_bucket = [] + for bucket in bucket_fqns: + bucket_owners = {} + candidates = (fqn for fqn in bucket if fqn in memory_estimate_by_fqn) + for fqn in sorted( + candidates, key=lambda name: (-memory_estimate_by_fqn[name], name) + ): + load, rank = heapq.heappop(rank_loads) + bucket_owners[fqn] = rank + heapq.heappush( + rank_loads, (load + memory_estimate_by_fqn[fqn], rank) + ) + owners_by_bucket.append(bucket_owners) + return tuple(owners_by_bucket) + + +_ItemT = TypeVar("_ItemT") + + +def _resolve_buckets( + items: Sequence[_ItemT], + specs: Sequence[BucketSpec], + *, + fqn: Callable[[_ItemT], str], +) -> tuple[tuple[_ItemT, ...], ...]: + resolved: list[list[_ItemT]] = [[] for _ in specs] + for item in items: + name = fqn(item) + matches = [ + index + for index, spec in enumerate(specs) + if any(fnmatch.fnmatchcase(name, pattern) for pattern in spec.patterns) + ] + if len(matches) != 1: + raise ValueError( + f"optimizer parameter {name!r} must match one bucket" + ) + resolved[matches[0]].append(item) + return tuple(tuple(bucket) for bucket in resolved) + + +@dataclass(frozen=True, slots=True) +class _MatrixBlock: + """A rectangular logical compute unit, independent of placement.""" + + offsets: tuple[int, ...] + shape: tuple[int, ...] + + @property + def numel(self) -> int: + return math.prod(self.shape) + + +@dataclass(frozen=True, slots=True) +class _MatrixBlockRoute: + """Map one logical block from storage holders to compute holders. + + ``None`` means the sources hold equivalent copies and one may be selected. + A reduction requires contributions from every source participant. + """ + + block: _MatrixBlock + source_participants: tuple[int, ...] + destination_participants: tuple[int, ...] + reduce_op: dist.ReduceOp | None = None + + +@dataclass(frozen=True, slots=True) +class _RedistributionPlan: + """Transport-neutral exact block partitions in both directions.""" + + participants: tuple[int, ...] + logical_shape: tuple[int, ...] + storage_to_compute_routes: tuple[_MatrixBlockRoute, ...] + compute_to_storage_routes: tuple[_MatrixBlockRoute, ...] + + def __post_init__(self) -> None: + all_routes = self.storage_to_compute_routes + self.compute_to_storage_routes + if any( + not route.source_participants or not route.destination_participants + for route in all_routes + ): + raise ValueError("redistribution routes require sources and destinations") + for direction, routes in ( + ("storage-to-compute", self.storage_to_compute_routes), + ("compute-to-storage", self.compute_to_storage_routes), + ): + _validate_matrix_block_partition( + tuple(route.block for route in routes), + self.logical_shape, + direction=direction, + ) + + compute_destinations = { + destination + for route in self.storage_to_compute_routes + for destination in route.destination_participants + } + for destination in compute_destinations: + _validate_matrix_block_partition( + tuple( + route.block + for route in self.storage_to_compute_routes + if destination in route.destination_participants + ), + self.logical_shape, + direction=f"compute destination {destination}", + ) + if any( + source not in compute_destinations + for route in self.compute_to_storage_routes + for source in route.source_participants + ): + raise ValueError("compute-to-storage source has no complete compute tensor") + + +def _validate_matrix_block_partition( + blocks: tuple[_MatrixBlock, ...], + logical_shape: tuple[int, ...], + *, + direction: str, +) -> None: + if any(size < 0 for size in logical_shape): + raise ValueError("logical tensor shape must be nonnegative") + for block in blocks: + if len(block.offsets) != len(logical_shape) or len(block.shape) != len( + logical_shape + ): + raise ValueError(f"{direction} block rank does not match logical tensor") + if any( + offset < 0 or size < 0 or offset + size > logical_size + for offset, size, logical_size in zip( + block.offsets, block.shape, logical_shape, strict=True + ) + ): + raise ValueError(f"{direction} block is outside the logical tensor") + + positive_blocks = tuple(block for block in blocks if block.numel) + for index, first in enumerate(positive_blocks): + for second in positive_blocks[index + 1 :]: + if all( + max(first_offset, second_offset) + < min( + first_offset + first_size, + second_offset + second_size, + ) + for first_offset, first_size, second_offset, second_size in zip( + first.offsets, + first.shape, + second.offsets, + second.shape, + strict=True, + ) + ): + raise NotImplementedError( + "overlapping logical matrix blocks are not supported" + ) + + if sum(block.numel for block in blocks) != math.prod(logical_shape): + raise ValueError(f"{direction} blocks do not cover the logical tensor") + + +def _build_owned_redistribution_plan( + storage_blocks: Sequence[tuple[tuple[int, ...], _MatrixBlock]], + *, + participants: tuple[int, ...], + owner: int, + logical_shape: tuple[int, ...], +) -> _RedistributionPlan: + """Build mirrored routes from one canonical block-to-holders mapping.""" + return _RedistributionPlan( + participants=participants, + logical_shape=logical_shape, + storage_to_compute_routes=tuple( + _MatrixBlockRoute( + block=block, + source_participants=holders, + destination_participants=(owner,), + ) + for holders, block in storage_blocks + ), + compute_to_storage_routes=tuple( + _MatrixBlockRoute( + block=block, + source_participants=(owner,), + destination_participants=holders, + ) + for holders, block in storage_blocks + ), + ) + + +@dataclass(frozen=True, slots=True) +class _PackedSpan: + """Physical packed-buffer location for a logical matrix block.""" + + block: _MatrixBlock + buffer_offset: int + + @property + def numel(self) -> int: + return self.block.numel + + +class _CommunicationSchedule: + """Physical execution strategy produced from redistribution routes.""" + + __slots__ = () + participants: tuple[int, ...] + local_participant: int + input_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] + output_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] + input_buffer_numel: int + output_buffer_numel: int + + def execute( + self, output: Tensor, input: Tensor + ) -> tuple[dist.Work, ...]: + raise NotImplementedError + + +@dataclass(frozen=True, slots=True) +class _PackedAllToAllSchedule(_CommunicationSchedule): + process_group: dist.ProcessGroup + participants: tuple[int, ...] + local_participant: int + input_split_sizes: tuple[int, ...] + output_split_sizes: tuple[int, ...] + input_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] + output_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] + + @property + def input_buffer_numel(self) -> int: + return sum(self.input_split_sizes) + + @property + def output_buffer_numel(self) -> int: + return sum(self.output_split_sizes) + + def execute( + self, output: Tensor, input: Tensor + ) -> tuple[dist.Work, ...]: + dist.all_to_all_single( + output[: self.output_buffer_numel], + input[: self.input_buffer_numel], + output_split_sizes=list(self.output_split_sizes), + input_split_sizes=list(self.input_split_sizes), + group=self.process_group, + ) + return () + + +@dataclass(frozen=True, slots=True) +class _AllGatherSchedule(_CommunicationSchedule): + process_group: dist.ProcessGroup + participants: tuple[int, ...] + local_participant: int + input_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] + output_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] + input_buffer_numel: int + output_buffer_numel: int + + def execute( + self, output: Tensor, input: Tensor + ) -> tuple[dist.Work, ...]: + dist.all_gather_into_tensor( + output[: self.output_buffer_numel], + input[: self.input_buffer_numel], + group=self.process_group, + ) + return () + + +@dataclass(frozen=True, slots=True) +class _ReduceScatterSchedule(_CommunicationSchedule): + process_group: dist.ProcessGroup + participants: tuple[int, ...] + local_participant: int + input_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] + output_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] + input_buffer_numel: int + output_buffer_numel: int + reduce_op: dist.ReduceOp + + def execute( + self, output: Tensor, input: Tensor + ) -> tuple[dist.Work, ...]: + dist.reduce_scatter_tensor( + output[: self.output_buffer_numel], + input[: self.input_buffer_numel], + op=self.reduce_op, + group=self.process_group, + ) + return () + + +@dataclass(frozen=True, slots=True) +class _PackedP2PTransfer: + peer: int + buffer_offset: int + numel: int + + +@dataclass(frozen=True, slots=True) +class _P2PSchedule(_CommunicationSchedule): + process_group: dist.ProcessGroup + participants: tuple[int, ...] + local_participant: int + input_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] + output_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] + sends: tuple[_PackedP2PTransfer, ...] + receives: tuple[_PackedP2PTransfer, ...] + input_buffer_numel: int + output_buffer_numel: int + + def execute( + self, output: Tensor, input: Tensor + ) -> tuple[dist.Work, ...]: + operations = [ + dist.P2POp( + dist.isend, + input.narrow(0, transfer.buffer_offset, transfer.numel), + transfer.peer, + self.process_group, + ) + for transfer in self.sends + ] + operations.extend( + dist.P2POp( + dist.irecv, + output.narrow(0, transfer.buffer_offset, transfer.numel), + transfer.peer, + self.process_group, + ) + for transfer in self.receives + ) + return tuple(dist.batch_isend_irecv(operations)) if operations else () + + +@dataclass(frozen=True, slots=True) +class _LocalSchedule(_CommunicationSchedule): + participants: tuple[int, ...] = () + local_participant: int = -1 + input_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] = () + output_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] = () + input_buffer_numel: int = 0 + output_buffer_numel: int = 0 + + def execute( + self, output: Tensor, input: Tensor + ) -> tuple[dist.Work, ...]: + if self.input_buffer_numel != self.output_buffer_numel: + raise ValueError("local schedules require equal buffer sizes") + output[: self.output_buffer_numel].copy_( + input[: self.input_buffer_numel] + ) + return () + + +@dataclass(slots=True) +class _BucketPlan(Generic[_ItemT]): + local_items: tuple[_ItemT, ...] + redistributed_items: tuple[_ItemT, ...] + redistribution_plans: tuple[_RedistributionPlan, ...] + group: _RedistributionGroup + storage_to_compute_schedule: _CommunicationSchedule + compute_to_storage_schedule: _CommunicationSchedule + dtype: torch.dtype + device: torch.device + + +@dataclass(frozen=True, slots=True) +class _RedistributionGroup: + process_group: dist.ProcessGroup + participants: tuple[int, ...] + local_participant: int + + +@dataclass(frozen=True, slots=True) +class _BucketPlanningResult(Generic[_ItemT]): + plans: tuple[_BucketPlan[_ItemT], ...] + ordered_items: tuple[_ItemT, ...] + + +@dataclass(slots=True) +class _BucketWork(Generic[_ItemT]): + plan: _BucketPlan[_ItemT] + storage_buffer: Tensor + compute_fragment_buffer: Tensor + forward_ready: torch.Event | None = None + compute_done: torch.Event | None = None + done: torch.Event | None = None + storage_to_compute_works: tuple[dist.Work, ...] = () + compute_to_storage_works: tuple[dist.Work, ...] = () + + +@dataclass(slots=True) +class _BufferSlot: + storage_exchange_storage: dict[ + tuple[torch.device, torch.dtype], Tensor + ] = field(default_factory=dict) + compute_exchange_storage: dict[ + tuple[torch.device, torch.dtype], Tensor + ] = field(default_factory=dict) + compute_storage: dict[tuple[torch.device, torch.dtype], Tensor] = field( + default_factory=dict + ) + + @staticmethod + def _ensure_capacity( + storage: dict[tuple[torch.device, torch.dtype], Tensor], + *, + numel: int, + dtype: torch.dtype, + device: torch.device, + ) -> Tensor: + key = (device, dtype) + buffer = storage.get(key) + if buffer is None or buffer.numel() < numel: + buffer = torch.empty(numel, dtype=dtype, device=device) + storage[key] = buffer + return buffer[:numel] + + def communication_buffers( + self, plan: _BucketPlan[Any] + ) -> tuple[Tensor, Tensor]: + to_compute = plan.storage_to_compute_schedule + to_storage = plan.compute_to_storage_schedule + return ( + self._ensure_capacity( + self.storage_exchange_storage, + numel=max( + to_compute.input_buffer_numel, + to_storage.output_buffer_numel, + ), + dtype=plan.dtype, + device=plan.device, + ), + self._ensure_capacity( + self.compute_exchange_storage, + numel=max( + to_compute.output_buffer_numel, + to_storage.input_buffer_numel, + ), + dtype=plan.dtype, + device=plan.device, + ), + ) + + def compute_buffer( + self, + shape: torch.Size | tuple[int, ...], + *, + dtype: torch.dtype, + device: torch.device, + ) -> Tensor: + return self._ensure_capacity( + self.compute_storage, + numel=math.prod(shape), + dtype=dtype, + device=device, + ).view(shape) + + +@dataclass(slots=True) +class _CommunicationContext: + device_handle: ModuleType + transfer_stream: torch.Stream + slots: tuple[_BufferSlot, _BufferSlot] + + @classmethod + def create(cls, device: torch.device) -> _CommunicationContext: + device_handle = torch.get_device_module(device) + transfer_stream = device_handle.Stream(device=device, priority=0) + return cls( + device_handle=device_handle, + transfer_stream=transfer_stream, + slots=(_BufferSlot(), _BufferSlot()), + ) + + +class _BucketedRedistributionRuntime(Generic[_ItemT]): + """Execute bucket plans with one-bucket-ahead communication prefetch. + + Callbacks run under the stream selected by the runtime. They must enqueue + work without synchronizing or calling ``Tensor.record_stream()``. + """ + + def __init__(self, device: torch.device) -> None: + self._device = device + self._context: _CommunicationContext | None = None + + def run( + self, + plans: Sequence[_BucketPlan[_ItemT]], + *, + local_tensor_spec: Callable[ + [_ItemT], tuple[torch.Size, torch.dtype, torch.device] + ], + compute_shape: Callable[[_ItemT], torch.Size], + prepare: Callable[[_ItemT, Tensor], None], + compute: Callable[[_ItemT, Tensor], None], + finalize: Callable[[_ItemT, Tensor], None], + ) -> None: + if self._context is None: + self._context = _CommunicationContext.create(self._device) + context = self._context + handle = context.device_handle + caller = handle.current_stream(self._device) + context.transfer_stream.wait_stream(caller) + + pending: list[_BucketWork[_ItemT]] = [] + redistributed_index = 0 + try: + for plan in plans: + slot = context.slots[redistributed_index % 2] + if not plan.redistributed_items: + with handle.stream(caller): + self._compute_local( + plan, + slot, + local_tensor_spec=local_tensor_spec, + prepare=prepare, + compute=compute, + finalize=finalize, + ) + continue + work = self._begin( + plan, + slot, + caller, + context, + local_tensor_spec=local_tensor_spec, + compute_shape=compute_shape, + prepare=prepare, + compute=compute, + finalize=finalize, + ) + redistributed_index += 1 + pending.append(work) + if len(pending) == 2: + oldest = pending.pop(0) + self._complete(oldest, context, finalize=finalize) + self._release(oldest, caller) + for work in pending: + self._complete(work, context, finalize=finalize) + self._release(work, caller) + except Exception: + # Preserve allocator lifetime ordering for work already enqueued on + # either stream. This is an error-path drain, not synchronization. + context.transfer_stream.wait_stream(caller) + caller.wait_stream(context.transfer_stream) + raise + + @staticmethod + def _begin( + plan: _BucketPlan[_ItemT], + slot: _BufferSlot, + caller_stream: torch.Stream, + context: _CommunicationContext, + *, + local_tensor_spec: Callable[ + [_ItemT], tuple[torch.Size, torch.dtype, torch.device] + ], + compute_shape: Callable[[_ItemT], torch.Size], + prepare: Callable[[_ItemT, Tensor], None], + compute: Callable[[_ItemT, Tensor], None], + finalize: Callable[[_ItemT, Tensor], None], + ) -> _BucketWork[_ItemT]: + handle = context.device_handle + transfer = context.transfer_stream + with handle.stream(transfer): + storage_buffer, compute_fragment_buffer = slot.communication_buffers( + plan + ) + work = _BucketWork(plan, storage_buffer, compute_fragment_buffer) + _prepare_redistributed(plan, storage_buffer, prepare=prepare) + work.storage_to_compute_works = _execute_schedule( + plan.storage_to_compute_schedule, + output=compute_fragment_buffer, + input=storage_buffer, + ) + work.forward_ready = handle.Event() + work.forward_ready.record(transfer) + + with handle.stream(caller_stream): + _BucketedRedistributionRuntime._compute_local( + plan, + slot, + local_tensor_spec=local_tensor_spec, + prepare=prepare, + compute=compute, + finalize=finalize, + ) + caller_stream.wait_event(work.forward_ready) + _compute_redistributed( + work, + slot, + compute_shape=compute_shape, + compute=compute, + ) + work.compute_done = handle.Event() + work.compute_done.record(caller_stream) + return work + + @staticmethod + def _complete( + work: _BucketWork[_ItemT], + context: _CommunicationContext, + *, + finalize: Callable[[_ItemT, Tensor], None], + ) -> None: + assert work.compute_done is not None + handle = context.device_handle + transfer = context.transfer_stream + with handle.stream(transfer): + transfer.wait_event(work.compute_done) + work.compute_to_storage_works = _execute_schedule( + work.plan.compute_to_storage_schedule, + output=work.storage_buffer, + input=work.compute_fragment_buffer, + ) + _finalize_redistributed(work, finalize=finalize) + work.done = handle.Event() + work.done.record(transfer) + + @staticmethod + def _release( + work: _BucketWork[_ItemT], caller_stream: torch.Stream + ) -> None: + assert work.done is not None + caller_stream.wait_event(work.done) + + @staticmethod + def _compute_local( + plan: _BucketPlan[_ItemT], + slot: _BufferSlot, + *, + local_tensor_spec: Callable[ + [_ItemT], tuple[torch.Size, torch.dtype, torch.device] + ], + prepare: Callable[[_ItemT, Tensor], None], + compute: Callable[[_ItemT, Tensor], None], + finalize: Callable[[_ItemT, Tensor], None], + ) -> None: + for item in plan.local_items: + shape, dtype, device = local_tensor_spec(item) + prepared = slot.compute_buffer(shape, dtype=dtype, device=device) + prepare(item, prepared) + compute(item, prepared) + finalize(item, prepared) + + +def _prepare_redistributed( + plan: _BucketPlan[_ItemT], + storage_buffer: Tensor, + *, + prepare: Callable[[_ItemT, Tensor], None], +) -> None: + schedule = plan.storage_to_compute_schedule + for index, item in enumerate(plan.redistributed_items): + spans = schedule.input_spans_by_parameter[index] + assert len(spans) == 1 + span = spans[0] + out = storage_buffer[ + span.buffer_offset : span.buffer_offset + span.numel + ].view(span.block.shape) + prepare(item, out) + + +def _execute_schedule( + schedule: _CommunicationSchedule, + *, + output: Tensor, + input: Tensor, +) -> tuple[dist.Work, ...]: + works = schedule.execute(output, input) + for work in works: + work.wait() + return works + + +def _compute_redistributed( + work: _BucketWork[_ItemT], + slot: _BufferSlot, + *, + compute_shape: Callable[[_ItemT], torch.Size], + compute: Callable[[_ItemT, Tensor], None], +) -> None: + plan = work.plan + to_compute = plan.storage_to_compute_schedule + to_storage = plan.compute_to_storage_schedule + for index, item in enumerate(plan.redistributed_items): + received_spans = to_compute.output_spans_by_parameter[index] + if not received_spans: + continue + compute_tensor = slot.compute_buffer( + compute_shape(item), + dtype=plan.dtype, + device=plan.device, + ) + for span in received_spans: + received = work.compute_fragment_buffer[ + span.buffer_offset : span.buffer_offset + span.numel + ] + _matrix_block_view(compute_tensor, span.block).copy_( + received.view(span.block.shape) + ) + + compute(item, compute_tensor) + + for span in to_storage.input_spans_by_parameter[index]: + packed = work.compute_fragment_buffer[ + span.buffer_offset : span.buffer_offset + span.numel + ] + packed.view(span.block.shape).copy_( + _matrix_block_view(compute_tensor, span.block) + ) + + +def _finalize_redistributed( + work: _BucketWork[_ItemT], + *, + finalize: Callable[[_ItemT, Tensor], None], +) -> None: + schedule = work.plan.compute_to_storage_schedule + for index, item in enumerate(work.plan.redistributed_items): + spans = schedule.output_spans_by_parameter[index] + assert len(spans) == 1 + span = spans[0] + update = work.storage_buffer[ + span.buffer_offset : span.buffer_offset + span.numel + ].view(span.block.shape) + finalize(item, update) + + +def _copy_transfers( + routes: tuple[_MatrixBlockRoute, ...], participants: tuple[int, ...] +) -> tuple[tuple[int, int, _MatrixBlock], ...]: + participant_order = { + participant: index for index, participant in enumerate(participants) + } + transfers = [] + for route in routes: + if route.reduce_op is not None: + raise ValueError("packed all-to-all cannot lower reduction routes") + sources = tuple( + sorted(route.source_participants, key=participant_order.__getitem__) + ) + for destination in route.destination_participants: + source = destination if destination in sources else sources[0] + transfers.append((source, destination, route.block)) + return tuple(transfers) + + +def _packed_spans_by_parameter( + indexed_spans: list[tuple[int, _PackedSpan]], parameter_count: int +) -> tuple[tuple[_PackedSpan, ...], ...]: + return tuple( + tuple( + span + for span_parameter_index, span in indexed_spans + if span_parameter_index == parameter_index + ) + for parameter_index in range(parameter_count) + ) + + +def _lower_packed_all_to_all( + redistribution_plans: tuple[_RedistributionPlan, ...], + *, + direction: str, + process_group: dist.ProcessGroup, + local_participant: int, +) -> _PackedAllToAllSchedule: + participants = redistribution_plans[0].participants + if any(plan.participants != participants for plan in redistribution_plans): + raise ValueError("one all-to-all schedule requires one participant order") + if tuple(dist.get_process_group_ranks(process_group)) != participants: + raise ValueError( + "redistribution participants must match process-group rank order" + ) + if local_participant not in participants: + raise ValueError("local rank is not a redistribution participant") + if direction == "storage_to_compute": + routes_by_parameter = tuple( + plan.storage_to_compute_routes for plan in redistribution_plans + ) + elif direction == "compute_to_storage": + routes_by_parameter = tuple( + plan.compute_to_storage_routes for plan in redistribution_plans + ) + else: + raise ValueError(f"unsupported redistribution direction {direction!r}") + transfers_by_parameter = tuple( + _copy_transfers(routes, participants) for routes in routes_by_parameter + ) + + input_split_sizes = [] + input_spans = [] + input_cursor = 0 + for destination in participants: + split_start = input_cursor + for parameter_index, transfers in enumerate(transfers_by_parameter): + for source, transfer_destination, block in transfers: + if source != local_participant or transfer_destination != destination: + continue + input_spans.append( + (parameter_index, _PackedSpan(block, input_cursor)) + ) + input_cursor += block.numel + input_split_sizes.append(input_cursor - split_start) + + output_split_sizes = [] + output_spans = [] + output_cursor = 0 + for source in participants: + split_start = output_cursor + for parameter_index, transfers in enumerate(transfers_by_parameter): + for transfer_source, destination, block in transfers: + if transfer_source != source or destination != local_participant: + continue + output_spans.append( + (parameter_index, _PackedSpan(block, output_cursor)) + ) + output_cursor += block.numel + output_split_sizes.append(output_cursor - split_start) + + return _PackedAllToAllSchedule( + process_group=process_group, + participants=participants, + local_participant=local_participant, + input_split_sizes=tuple(input_split_sizes), + output_split_sizes=tuple(output_split_sizes), + input_spans_by_parameter=_packed_spans_by_parameter( + input_spans, len(redistribution_plans) + ), + output_spans_by_parameter=_packed_spans_by_parameter( + output_spans, len(redistribution_plans) + ), + ) + + +def _device_mesh_ranks(mesh: DeviceMesh) -> tuple[int, ...]: + if mesh.ndim == 1: + return tuple(dist.get_process_group_ranks(mesh.get_group())) + return tuple(mesh.mesh.flatten().tolist()) + + +def _redistribution_group(mesh: DeviceMesh) -> _RedistributionGroup: + if mesh.ndim != 1: + raise ValueError("optimizer redistribution mesh must be one-dimensional") + process_group = mesh.get_group() + participants = tuple(dist.get_process_group_ranks(process_group)) + return _RedistributionGroup( + process_group=process_group, + participants=participants, + local_participant=participants[dist.get_rank(process_group)], + ) + + +def _normalize_dim(dim: int, ndim: int) -> int: + normalized = dim if dim >= 0 else dim + ndim + if normalized < 0 or normalized >= ndim: + raise ValueError(f"dimension {dim} is invalid for a rank-{ndim} tensor") + return normalized + + +def _dtensor_storage_block_for_participant( + tensor: DTensor, + participant: int, +) -> _MatrixBlock: + mesh_shape = tuple(tensor.device_mesh.shape) + mesh_rank = _device_mesh_ranks(tensor.device_mesh).index(participant) + coordinate = [0] * len(mesh_shape) + for mesh_dim in range(len(mesh_shape) - 1, -1, -1): + mesh_rank, coordinate[mesh_dim] = divmod(mesh_rank, mesh_shape[mesh_dim]) + + local_shape = list(tensor.shape) + global_offsets = [0] * tensor.ndim + for mesh_dim, placement in enumerate(tensor.placements): + if type(placement) is not Shard: + raise ValueError( + "redistributed optimizer storage requires exact Shard placements" + ) + tensor_dim = _normalize_dim(placement.dim, tensor.ndim) + local_size, global_offset = Shard.local_shard_size_and_offset( + tensor.shape[tensor_dim], + mesh_shape[mesh_dim], + coordinate[mesh_dim], + ) + local_shape[tensor_dim] = local_size + global_offsets[tensor_dim] = global_offset + return _MatrixBlock( + offsets=tuple(global_offsets), + shape=tuple(local_shape), + ) + + +def _dtensor_storage_blocks( + tensor: DTensor, + participants: tuple[int, ...], +) -> tuple[tuple[tuple[int, ...], _MatrixBlock], ...]: + storage_participants = _device_mesh_ranks(tensor.device_mesh) + if storage_participants != participants: + raise ValueError( + "bucket mesh participants must match redistributed DTensor storage" + ) + return tuple( + ( + (participant,), + _dtensor_storage_block_for_participant(tensor, participant), + ) + for participant in participants + ) + + +def _build_bucket_plans( + items: Sequence[_ItemT], + specs: Sequence[BucketSpec], + *, + fqn: Callable[[_ItemT], str], + compute_locally: Callable[[_ItemT], bool], + storage_dtensor: Callable[[_ItemT], DTensor], +) -> _BucketPlanningResult[_ItemT]: + resolved = _resolve_buckets(items, specs, fqn=fqn) + plans = [] + ordered_items = [] + for spec, bucket in zip(specs, resolved, strict=True): + if not bucket: + continue + group = _redistribution_group(spec.mesh) + local_items = tuple( + sorted( + (item for item in bucket if compute_locally(item)), + key=fqn, + ) + ) + redistributed_items = tuple( + sorted( + (item for item in bucket if not compute_locally(item)), + key=fqn, + ) + ) + expected_owners = {fqn(item) for item in redistributed_items} + provided_owners = set(spec.owner_rank_by_fqn) + if provided_owners != expected_owners: + raise ValueError( + f"bucket {spec.name!r} owner assignment must exactly cover " + "whole-tensor-owned parameters; " + f"missing={sorted(expected_owners - provided_owners)}, " + f"extra={sorted(provided_owners - expected_owners)}" + ) + ordered_items.extend(local_items) + ordered_items.extend(redistributed_items) + + if not redistributed_items: + tensor = storage_dtensor(local_items[0]).to_local() + plans.append( + _BucketPlan( + local_items=local_items, + redistributed_items=(), + redistribution_plans=(), + group=group, + storage_to_compute_schedule=_LocalSchedule(), + compute_to_storage_schedule=_LocalSchedule(), + dtype=tensor.dtype, + device=tensor.device, + ) + ) + continue + + owner_ranks = [ + spec.owner_rank_by_fqn[fqn(item)] for item in redistributed_items + ] + if any(rank not in range(len(group.participants)) for rank in owner_ranks): + raise ValueError( + f"bucket {spec.name!r} has owner outside its process group" + ) + + storage_dtensors = [storage_dtensor(item) for item in redistributed_items] + local_tensors = [tensor.to_local() for tensor in storage_dtensors] + dtype = local_tensors[0].dtype + device = local_tensors[0].device + if any( + tensor.dtype != dtype or tensor.device != device + for tensor in local_tensors + ): + raise ValueError(f"bucket {spec.name!r} mixes dtype or device") + + blocks_by_item = tuple( + _dtensor_storage_blocks(tensor, group.participants) + for tensor in storage_dtensors + ) + for tensor, blocks in zip(local_tensors, blocks_by_item, strict=True): + local_blocks = [ + block + for holders, block in blocks + if group.local_participant in holders + ] + if len(local_blocks) != 1 or tuple(tensor.shape) != local_blocks[0].shape: + raise ValueError( + f"bucket {spec.name!r} storage block does not match its mesh" + ) + + redistribution_plans = tuple( + _build_owned_redistribution_plan( + blocks, + participants=group.participants, + owner=group.participants[owner_rank], + logical_shape=tuple(tensor.shape), + ) + for tensor, blocks, owner_rank in zip( + storage_dtensors, blocks_by_item, owner_ranks, strict=True + ) + ) + plans.append( + _BucketPlan( + local_items=local_items, + redistributed_items=redistributed_items, + redistribution_plans=redistribution_plans, + group=group, + storage_to_compute_schedule=_lower_packed_all_to_all( + redistribution_plans, + direction="storage_to_compute", + process_group=group.process_group, + local_participant=group.local_participant, + ), + compute_to_storage_schedule=_lower_packed_all_to_all( + redistribution_plans, + direction="compute_to_storage", + process_group=group.process_group, + local_participant=group.local_participant, + ), + dtype=dtype, + device=device, + ) + ) + + return _BucketPlanningResult( + plans=tuple(plans), + ordered_items=tuple(ordered_items), + ) + + +def _validate_bucket_plans_across_ranks( + plans: Sequence[_BucketPlan[_ItemT]], + *, + item_signature: Callable[[_ItemT], tuple[Any, ...]], +) -> None: + for plan in plans: + description = ( + str(plan.dtype), + plan.device.type, + tuple( + _redistribution_plan_key(redistribution_plan) + for redistribution_plan in plan.redistribution_plans + ), + [ + item_signature(item) + for item in plan.local_items + plan.redistributed_items + ], + ) + digest = hashlib.sha256(repr(description).encode()).digest() + plan_hash = int.from_bytes(digest[:7], "little") + local_hash = torch.tensor(plan_hash, dtype=torch.int64, device=plan.device) + process_group = plan.group.process_group + gathered = [ + torch.empty_like(local_hash) + for _ in range(dist.get_world_size(process_group)) + ] + dist.all_gather(gathered, local_hash, group=process_group) + if any(value.item() != plan_hash for value in gathered): + raise RuntimeError("optimizer bucket plans differ across ranks") + + +def _matrix_block_view(tensor: Tensor, block: _MatrixBlock) -> Tensor: + view = tensor[ + tuple( + slice(offset, offset + size) + for offset, size in zip(block.offsets, block.shape, strict=True) + ) + ] + assert tuple(view.shape) == block.shape + return view + + +def _redistribution_plan_key(plan: _RedistributionPlan) -> tuple[Any, ...]: + def route_key(route: _MatrixBlockRoute) -> tuple[Any, ...]: + return ( + route.block.offsets, + route.block.shape, + route.source_participants, + route.destination_participants, + str(route.reduce_op), + ) + + return ( + plan.participants, + plan.logical_shape, + tuple(map(route_key, plan.storage_to_compute_routes)), + tuple(map(route_key, plan.compute_to_storage_routes)), + ) diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py new file mode 100644 index 0000000000..497017d34a --- /dev/null +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -0,0 +1,650 @@ +# 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. + +"""Standalone bucketed Distributed Muon optimizer.""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any, cast + +import torch +from torch import Tensor +from torch.distributed.tensor import DTensor, Replicate, Shard +from torch.distributed.tensor.placement_types import _StridedShard +from torch.optim import Optimizer +from .bucketed_redistribution import ( + _BucketedRedistributionRuntime, + _BucketPlan, + _build_bucket_plans, + _device_mesh_ranks, + _validate_bucket_plans_across_ranks, + assign_balanced_owners, + BucketSpec, +) + + +__all__ = ["BucketSpec", "assign_balanced_owners", "Owned"] + + + +@dataclass(frozen=True, slots=True) +class Owned: + """Require a complete matrix; sharded storage uses a ``BucketSpec`` owner.""" + + +@dataclass(frozen=True, slots=True) +class _PreparedParameterComputeView: + global_compute_shape: torch.Size + local_compute_tensor: Tensor + + +class DistributedMuon(Optimizer): + """Internal runtime constructed through ``build_distributed_muon``.""" + + def __init__( + self, + params: Iterable[Tensor] | Iterable[dict[str, Any]], + *, + bucket_spec: Sequence[BucketSpec], + _prepared_compute_views: Mapping[ + str, _PreparedParameterComputeView + ], + lr: float = 1e-3, + weight_decay: float = 0.1, + momentum: float = 0.95, + nesterov: bool = True, + ns_coefficients: tuple[float, float, float] = (3.4445, -4.7750, 2.0315), + eps: float = 1e-7, + ns_steps: int = 5, + adjust_lr_fn: str | None = None, + ) -> None: + defaults = { + "lr": lr, + "weight_decay": weight_decay, + "momentum": momentum, + "nesterov": nesterov, + "ns_coefficients": ns_coefficients, + "eps": eps, + "ns_steps": ns_steps, + "adjust_lr_fn": adjust_lr_fn, + } + params = [ + dict(param_or_group) + if isinstance(param_or_group, dict) + else param_or_group + for param_or_group in params + ] + self._first_step_validated = False + self._prepared_compute_views = dict(_prepared_compute_views) + super().__init__(params, defaults) + assert all( + isinstance(param, DTensor) and param.device.type == "cuda" + for group in self.param_groups + for param in group["params"] + ), "DistributedMuon requires CUDA DTensor parameters" + group_compute_placements = [] + for group in self.param_groups: + compute_placement = group.pop("_compute_placement", None) + group_compute_placements.append(compute_placement) + self._group_compute_placements = tuple(group_compute_placements) + + self._specs = tuple(bucket_spec) + self._validate_groups() + self._initialize_plan() + self._validate_plan_across_ranks() + self._redistribution_runtime = _BucketedRedistributionRuntime[ + _ParameterComputeLayout + ](self._tensor_device) + self._frozen_param_names = tuple( + tuple(group.get("param_names", ())) for group in self.param_groups + ) + + @torch.no_grad() + def step( + self, closure: Callable[[], float] | None = None + ) -> float | None: + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + self._preflight_step() + self._redistribution_runtime.run( + self._plans, + local_tensor_spec=self._local_tensor_spec, + compute_shape=self._compute_shape, + prepare=self._prepare_local, + compute=self._compute_update, + finalize=self._apply_update, + ) + return loss + + def add_param_group(self, param_group: dict[str, Any]) -> None: + if hasattr(self, "_plans"): + raise RuntimeError( + "DistributedMuon parameter groups are frozen" + ) + super().add_param_group(param_group) + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + # Compute layout is intentionally not duplicated in optimizer state. TorchTitan + # must reconstruct it from the same model and optimizer config before resume. + saved_groups = state_dict.get("param_groups", ()) + if len(saved_groups) != len(self._frozen_param_names) or any( + "param_names" in saved and tuple(saved["param_names"]) != names + for saved, names in zip( + saved_groups, self._frozen_param_names, strict=True + ) + ): + raise ValueError("checkpoint changed DistributedMuon's parameter groups") + super().load_state_dict(state_dict) + self._validate_plan_across_ranks() + self._first_step_validated = False + + def _validate_groups(self) -> None: + for group_index, group in enumerate(self.param_groups): + if group.get("fused") or group.get("foreach"): + raise NotImplementedError( + "DistributedMuon does not support fused or foreach" + ) + ns_steps = group["ns_steps"] + coefficients = group["ns_coefficients"] + if ( + any( + group[name] < 0 + for name in ("lr", "weight_decay", "momentum", "eps") + ) + or not isinstance(ns_steps, int) + or not 0 <= ns_steps < 100 + or len(coefficients) != 3 + or not all(isinstance(value, (int, float)) for value in coefficients) + or group["adjust_lr_fn"] + not in (None, "original", "match_rms_adamw", "spectral_unclamped") + ): + raise ValueError(f"invalid DistributedMuon group {group_index}") + + def _build_parameter_compute_layouts( + self, + ) -> tuple[_ParameterComputeLayout, ...]: + parameters = [] + seen_names = set() + seen_params = set() + for group_index, group in enumerate(self.param_groups): + params = group["params"] + names = group.get("param_names") + if names is None or len(names) != len(params): + raise ValueError( + "DistributedMuon requires param_names aligned with params" + ) + for fqn, param in zip(names, params, strict=True): + if fqn in seen_names or id(param) in seen_params: + raise ValueError(f"duplicate Muon parameter {fqn!r}") + seen_names.add(fqn) + seen_params.add(id(param)) + parameters.append((group_index, fqn, param)) + + prepared_fqns = self._prepared_compute_views.keys() + if prepared_fqns != seen_names: + raise ValueError( + "prepared compute views must exactly cover parameter FQNs; " + f"missing={sorted(seen_names - prepared_fqns)}, " + f"extra={sorted(prepared_fqns - seen_names)}" + ) + compute_layouts = [] + for group_index, fqn, param in parameters: + compute_placement = self._group_compute_placements[group_index] + prepared = self._prepared_compute_views[fqn] + if not isinstance(prepared, _PreparedParameterComputeView): + raise TypeError( + f"invalid prepared compute view for parameter {fqn!r}" + ) + global_compute_shape = torch.Size(prepared.global_compute_shape) + local_compute_tensor = prepared.local_compute_tensor + compute_locally = _validate_muon_parameter( + fqn, + param, + global_compute_shape, + local_compute_tensor, + compute_placement, + ) + compute_layouts.append( + _ParameterComputeLayout( + fqn=fqn, + param=param, + group_index=group_index, + global_compute_shape=global_compute_shape, + local_compute_tensor=local_compute_tensor, + compute_placement=compute_placement, + compute_locally=compute_locally, + ) + ) + return tuple(compute_layouts) + + def _initialize_plan(self) -> None: + compute_layouts = self._build_parameter_compute_layouts() + result = _build_bucket_plans( + compute_layouts, + self._specs, + fqn=lambda item: item.fqn, + compute_locally=lambda item: item.compute_locally, + storage_dtensor=lambda item: item.param, + ) + self._plans = result.plans + self._parameter_compute_layouts = result.ordered_items + self._tensor_device = self._plans[0].device + + def _validate_plan_across_ranks(self) -> None: + _validate_bucket_plans_across_ranks( + self._plans, + item_signature=self._plan_item_signature, + ) + + def _plan_item_signature( + self, compute_layout: _ParameterComputeLayout + ) -> tuple[Any, ...]: + return ( + compute_layout.fqn, + compute_layout.group_index, + tuple(compute_layout.param.shape), + tuple(compute_layout.param.stride()), + str(compute_layout.param.dtype), + compute_layout.param.to_local().device.type, + tuple(compute_layout.global_compute_shape), + compute_layout.compute_locally, + _compute_placement_key(compute_layout.compute_placement), + _device_mesh_ranks(compute_layout.param.device_mesh), + tuple(map(str, compute_layout.param.placements)), + self._group_signature(compute_layout), + ) + + def _group(self, compute_layout: _ParameterComputeLayout) -> dict[str, Any]: + return self.param_groups[compute_layout.group_index] + + def _group_signature( + self, compute_layout: _ParameterComputeLayout + ) -> tuple[Any, ...]: + group = self._group(compute_layout) + return tuple( + group[key] + for key in ( + "lr", + "weight_decay", + "momentum", + "nesterov", + "ns_coefficients", + "eps", + "ns_steps", + "adjust_lr_fn", + ) + ) + + def _preflight_step(self) -> None: + initialize_state = not self._first_step_validated + for compute_layout in self._parameter_compute_layouts: + if ( + compute_layout.compute_locally + and compute_layout.param.to_local().untyped_storage().data_ptr() + != compute_layout.local_compute_tensor.untyped_storage().data_ptr() + ): + raise RuntimeError( + f"parameter local storage changed for {compute_layout.fqn!r}; " + "rebuild DistributedMuon" + ) + if initialize_state and any( + compute_layout.param.grad is None + for compute_layout in self._parameter_compute_layouts + ): + raise RuntimeError("DistributedMuon requires every configured gradient") + + gradients = [] + for compute_layout in self._parameter_compute_layouts: + grad = self._gradient(compute_layout) + gradients.append((compute_layout, grad)) + if initialize_state: + self._validate_momentum(compute_layout) + + # State creation happens only after every gradient and existing state + # tensor has passed validation, so a deterministic input error cannot + # partially update an earlier bucket. + if initialize_state: + for compute_layout, grad in gradients: + self._momentum(compute_layout, grad) + self._first_step_validated = True + + @staticmethod + def _has_storage_layout( + tensor: DTensor, compute_layout: _ParameterComputeLayout + ) -> bool: + local = tensor.to_local() + param_local = compute_layout.param.to_local() + return ( + tensor.shape == compute_layout.param.shape + and tensor.stride() == compute_layout.param.stride() + and _device_mesh_ranks(tensor.device_mesh) + == _device_mesh_ranks(compute_layout.param.device_mesh) + and tensor.placements == compute_layout.param.placements + and local.shape == param_local.shape + and local.stride() == param_local.stride() + and local.dtype == param_local.dtype + and local.device == param_local.device + and local.is_contiguous() + ) + + def _gradient(self, compute_layout: _ParameterComputeLayout) -> DTensor: + grad = compute_layout.param.grad + if not isinstance(grad, DTensor) or not self._has_storage_layout( + grad, compute_layout + ): + raise RuntimeError( + f"gradient storage layout changed for {compute_layout.fqn!r}" + ) + return grad + + def _validate_momentum(self, compute_layout: _ParameterComputeLayout) -> None: + momentum = self.state.get(compute_layout.param, {}).get("momentum_buffer") + if momentum is None: + return + if not isinstance(momentum, DTensor) or not self._has_storage_layout( + momentum, compute_layout + ): + raise RuntimeError( + f"momentum storage layout changed for {compute_layout.fqn!r}" + ) + + def _momentum( + self, compute_layout: _ParameterComputeLayout, grad: DTensor + ) -> DTensor: + state = self.state[compute_layout.param] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like( + grad, memory_format=torch.preserve_format + ) + return state["momentum_buffer"] + + def _update_local_momentum( + self, compute_layout: _ParameterComputeLayout + ) -> tuple[Tensor, Tensor, dict[str, Any]]: + grad = cast(DTensor, compute_layout.param.grad) + momentum = cast(DTensor, self.state[compute_layout.param]["momentum_buffer"]) + local_grad = grad.to_local().view_as(compute_layout.local_compute_tensor) + local_momentum = momentum.to_local().view_as( + compute_layout.local_compute_tensor + ) + group = self._group(compute_layout) + local_momentum.lerp_(local_grad, 1 - group["momentum"]) + torch.autograd.graph.increment_version(momentum) + return local_grad, local_momentum, group + + @staticmethod + def _write_prepared( + group: dict[str, Any], grad: Tensor, momentum: Tensor, out: Tensor + ) -> None: + if group["nesterov"]: + torch.lerp( + grad, + momentum, + group["momentum"], + out=out, + ) + else: + out.copy_(momentum) + + def _prepare_local( + self, compute_layout: _ParameterComputeLayout, out: Tensor + ) -> None: + grad, momentum, group = self._update_local_momentum(compute_layout) + self._write_prepared(group, grad, momentum, out) + + def _compute_update( + self, compute_layout: _ParameterComputeLayout, compute: Tensor + ) -> None: + group = self._group(compute_layout) + _compute_muon_update( + compute, + lr=group["lr"], + ns_coefficients=group["ns_coefficients"], + ns_steps=group["ns_steps"], + eps=group["eps"], + adjust_lr_fn=group["adjust_lr_fn"], + out=compute, + ) + + def _apply_update( + self, compute_layout: _ParameterComputeLayout, update: Tensor + ) -> None: + group = self._group(compute_layout) + local_param = ( + compute_layout.local_compute_tensor + if compute_layout.compute_locally + else compute_layout.param.to_local() + ) + local_param.mul_(1 - group["lr"] * group["weight_decay"]) + local_param.add_(update) + torch.autograd.graph.increment_version(compute_layout.param) + + @staticmethod + def _local_tensor_spec( + compute_layout: _ParameterComputeLayout, + ) -> tuple[torch.Size, torch.dtype, torch.device]: + tensor = compute_layout.local_compute_tensor + return tensor.shape, tensor.dtype, tensor.device + + @staticmethod + def _compute_shape( + compute_layout: _ParameterComputeLayout, + ) -> torch.Size: + return compute_layout.global_compute_shape + +@dataclass(frozen=True, slots=True) +class _ParameterComputeLayout: + fqn: str + param: DTensor + group_index: int + global_compute_shape: torch.Size + local_compute_tensor: Tensor + compute_placement: Owned | Shard + compute_locally: bool + + +def _has_replicated_storage(param: DTensor) -> bool: + return all(type(placement) is Replicate for placement in param.placements) + + +def _has_dim0_sharded_storage(param: DTensor) -> bool: + has_shard = False + for placement in param.placements: + # FSDP2 emits _StridedShard when a later TP/EP axis already shards + # this dimension. Keep the allowlist exact so new placements fail closed. + if type(placement) in (Shard, _StridedShard): + if getattr(placement, "dim") % param.ndim != 0: + return False + has_shard = True + elif type(placement) is not Replicate: + return False + return has_shard + + +def _validate_muon_parameter( + fqn: str, + param: DTensor, + global_compute_shape: torch.Size, + local_compute_tensor: Tensor, + compute_placement: object, +) -> bool: + local = param.to_local() + if ( + torch.is_complex(param) + or param.ndim < 2 + or not local.is_contiguous() + or tuple(param.stride()) + != tuple(torch.empty(param.shape, device="meta").stride()) + ): + raise ValueError( + f"Muon parameter {fqn!r} has unsupported shape or storage" + ) + + if ( + len(global_compute_shape) < 2 + or local_compute_tensor.ndim < 2 + or math.prod(global_compute_shape) != param.numel() + or local_compute_tensor.numel() != local.numel() + or local_compute_tensor.dtype != local.dtype + or local_compute_tensor.device != local.device + or not local_compute_tensor.is_contiguous() + or local_compute_tensor.data_ptr() != local.data_ptr() + ): + raise ValueError( + f"invalid prepared compute view for parameter {fqn!r}" + ) + + if compute_placement is None: + raise ValueError( + f"Muon parameter {fqn!r} requires explicit compute_placement" + ) + + replicated_storage = _has_replicated_storage(param) + if isinstance(compute_placement, Shard): + if len(global_compute_shape) < 3: + raise ValueError( + "compute Shard requires a batch of complete Muon matrices" + ) + compute_dim = _normalize_dim( + compute_placement.dim, len(global_compute_shape) + ) + if compute_dim != 0: + raise ValueError("DistributedMuon currently supports compute Shard(0)") + if local_compute_tensor.ndim != len(global_compute_shape): + raise ValueError( + f"compute Shard(0) for {fqn!r} must keep complete matrices local" + ) + if replicated_storage: + if local_compute_tensor.shape != global_compute_shape: + raise ValueError( + f"replicated storage for {fqn!r} must contain the complete " + "compute tensor" + ) + elif ( + local_compute_tensor.shape[1:] != global_compute_shape[1:] + or not _has_dim0_sharded_storage(param) + ): + raise ValueError( + f"compute Shard(0) for {fqn!r} must already match storage sharding" + ) + return True + elif not isinstance(compute_placement, Owned): + raise TypeError(f"unsupported compute placement {compute_placement!r}") + elif len(global_compute_shape) != 2 or param.ndim != 2: + raise ValueError( + f"owned Muon parameter {fqn!r} requires matrix storage" + ) + elif replicated_storage: + if local_compute_tensor.shape != global_compute_shape: + raise ValueError( + f"replicated storage for {fqn!r} must contain the complete " + "compute tensor" + ) + return True + elif ( + param.device_mesh.ndim != 1 + or len(param.placements) != 1 + or type(param.placements[0]) is not Shard + ): + raise ValueError( + f"owned Muon parameter {fqn!r} requires replicated or 1D Shard " + "matrix storage" + ) + return False + + +def _normalize_dim(dim: int, ndim: int) -> int: + normalized = dim if dim >= 0 else dim + ndim + if normalized < 0 or normalized >= ndim: + raise ValueError(f"dimension {dim} is invalid for a rank-{ndim} tensor") + return normalized + + +def _compute_placement_key( + placement: Owned | Shard, +) -> tuple[Any, ...]: + if isinstance(placement, Owned): + return ("owned",) + return ("shard", placement.dim) + + +# Keep the functional math aligned with torch.optim.Muon while owning the +# implementation here so the distributed runtime has no Muon dependency. +def _zeropower_via_newtonschulz( + update: Tensor, + *, + ns_coefficients: tuple[float, float, float], + ns_steps: int, + eps: float, +) -> Tensor: + """Compute Muon's approximate polar factor without using torch.optim.Muon.""" + a, b, c = ns_coefficients + result = update.to(dtype=torch.bfloat16, copy=True) + transposed = result.shape[-2] > result.shape[-1] + if transposed: + result = result.transpose(-2, -1) + result.div_(result.norm(dim=(-2, -1), keepdim=True).clamp_min(eps)) + + if result.ndim == 2: + for _ in range(ns_steps): + gram = result @ result.T + gram_update = torch.addmm(gram, gram, gram, beta=b, alpha=c) + result = torch.addmm(result, gram_update, result, beta=a) + else: + original_shape = result.shape + matrices = result.reshape(-1, *original_shape[-2:]) + for _ in range(ns_steps): + gram = matrices @ matrices.transpose(-2, -1) + gram_update = torch.baddbmm(gram, gram, gram, beta=b, alpha=c) + matrices = torch.baddbmm(matrices, gram_update, matrices, beta=a) + result = matrices.reshape(original_shape) + + return result.transpose(-2, -1) if transposed else result + + +def _adjust_learning_rate( + lr: float, + adjust_lr_fn: str | None, + compute_matrix_shape: torch.Size, +) -> float: + rows, columns = compute_matrix_shape[-2:] + if adjust_lr_fn is None or adjust_lr_fn == "original": + ratio = math.sqrt(max(1, rows / columns)) + elif adjust_lr_fn == "match_rms_adamw": + ratio = 0.2 * math.sqrt(max(rows, columns)) + elif adjust_lr_fn == "spectral_unclamped": + ratio = math.sqrt(rows / columns) + else: + raise ValueError(f"unsupported adjust_lr_fn {adjust_lr_fn!r}") + return lr * ratio + + +def _compute_muon_update( + prepared: Tensor, + *, + lr: float, + ns_coefficients: tuple[float, float, float], + ns_steps: int, + eps: float, + adjust_lr_fn: str | None, + out: Tensor, +) -> Tensor: + direction = _zeropower_via_newtonschulz( + prepared, + ns_coefficients=ns_coefficients, + ns_steps=ns_steps, + eps=eps, + ) + adjusted_lr = _adjust_learning_rate(lr, adjust_lr_fn, prepared.shape) + out.zero_() + out.add_(direction, alpha=-adjusted_lr) + return out diff --git a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py new file mode 100644 index 0000000000..4cbd3aa201 --- /dev/null +++ b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py @@ -0,0 +1,208 @@ +# 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. + +"""Muon parameter views and pre-construction layout preparation.""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Any + +import torch +from torch import Tensor +from torch.distributed.tensor import DTensor, Replicate, Shard +from .bucketed_redistribution import ( + _bind_bucket_configs, + BucketConfig, + BucketSpec, +) +from .muon import ( + _PreparedParameterComputeView, + DistributedMuon, + Owned, +) + + +__all__ = [ + "BatchedMatrixComputeView", + "build_distributed_muon", + "MuonComputeSharding", +] + + +@dataclass(frozen=True, slots=True) +class BatchedMatrixComputeView: + """Unflatten a storage dimension into a batch of matrices.""" + + num_matrices: int + matrices_flattened_into_dim: int = 0 + + def __post_init__(self) -> None: + if ( + isinstance(self.num_matrices, bool) + or not isinstance(self.num_matrices, int) + or self.num_matrices <= 0 + ): + raise ValueError("num_matrices must be a positive integer") + if isinstance(self.matrices_flattened_into_dim, bool) or not isinstance( + self.matrices_flattened_into_dim, int + ): + raise ValueError("matrices_flattened_into_dim must be an integer") + if self.matrices_flattened_into_dim != 0: + raise ValueError("only matrices_flattened_into_dim=0 is supported") + + def _resolve(self, storage_shape: torch.Size) -> _ResolvedBatchedMatrixView: + if len(storage_shape) != 2: + raise ValueError("BatchedMatrixComputeView requires rank-2 storage") + flattened_extent = storage_shape[self.matrices_flattened_into_dim] + if flattened_extent == 0 or flattened_extent % self.num_matrices: + raise ValueError( + f"storage shape {tuple(storage_shape)} is not divisible into " + f"{self.num_matrices} matrices" + ) + return _ResolvedBatchedMatrixView( + matrix_rows=flattened_extent // self.num_matrices, + matrix_columns=storage_shape[1], + ) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class MuonComputeSharding: + """Define the logical Muon compute tensor and its required placement.""" + + # Applied before compute placement, so placement dimensions refer to the + # viewed tensor. A future view_after_placement mode can apply a local view + # after redistribution; that ordering is not supported yet. + view_before_placement: BatchedMatrixComputeView | None = None + placement: Owned | Shard + + def __post_init__(self) -> None: + if not isinstance(self.placement, (Owned, Shard)): + raise TypeError("placement must be Owned or Shard") + if self.view_before_placement is not None and not isinstance( + self.view_before_placement, BatchedMatrixComputeView + ): + raise TypeError( + "view_before_placement must be a BatchedMatrixComputeView or None" + ) + + +@dataclass(frozen=True, slots=True) +class _ResolvedBatchedMatrixView: + matrix_rows: int + matrix_columns: int + + def compute_shape(self, storage_shape: torch.Size) -> torch.Size: + if len(storage_shape) != 2: + raise ValueError("batched-matrix compute view requires rank-2 storage") + if ( + storage_shape[0] % self.matrix_rows + or storage_shape[1] != self.matrix_columns + ): + raise ValueError( + f"storage shape {tuple(storage_shape)} is not aligned to " + f"matrix shape {(self.matrix_rows, self.matrix_columns)}" + ) + return torch.Size( + ( + storage_shape[0] // self.matrix_rows, + self.matrix_rows, + self.matrix_columns, + ) + ) + + +def build_distributed_muon( + params: Iterable[Tensor] | Iterable[dict[str, Any]], + *, + bucket_spec: Sequence[BucketSpec] | None = None, + bucket_configs: Sequence[BucketConfig] | None = None, + **kwargs: Any, +) -> DistributedMuon: + """Prepare parameter views and construct the DistributedMuon runtime.""" + if (bucket_spec is None) == (bucket_configs is None): + raise ValueError("provide exactly one of bucket_spec or bucket_configs") + + prepared_params = [] + parameters_to_prepare = [] + for param_or_group in params: + if not isinstance(param_or_group, dict): + prepared_params.append(param_or_group) + continue + group = dict(param_or_group) + compute_sharding = group.pop("compute_sharding", None) + if not isinstance(compute_sharding, MuonComputeSharding): + raise TypeError("compute_sharding must be a MuonComputeSharding") + compute_view = compute_sharding.view_before_placement + group["_compute_placement"] = compute_sharding.placement + raw_params = group.get("params", ()) + group_params = ( + (raw_params,) if isinstance(raw_params, Tensor) else tuple(raw_params) + ) + raw_param_names = group.get("param_names") + param_names = ( + () if raw_param_names is None else tuple(raw_param_names) + ) + if raw_param_names is None or len(group_params) != len(param_names): + raise ValueError("params and param_names must be aligned") + group["params"] = group_params + group["param_names"] = param_names + + for param, fqn in zip(group_params, param_names, strict=True): + parameters_to_prepare.append((param, fqn, compute_view)) + prepared_params.append(group) + + if bucket_configs is not None: + storage_by_fqn = { + fqn: param + for param, fqn, _compute_view in parameters_to_prepare + if isinstance(param, DTensor) + } + if len(storage_by_fqn) != len(parameters_to_prepare): + raise TypeError("bucket_configs require named DTensor parameters") + bucket_spec = _bind_bucket_configs(bucket_configs, storage_by_fqn) + assert bucket_spec is not None + bucket_spec = tuple(bucket_spec) + + prepared_compute_views = {} + for param, fqn, compute_view in parameters_to_prepare: + global_storage_shape = torch.Size(param.shape) + if compute_view is not None and any( + type(placement) not in (Shard, Replicate) + for placement in getattr(param, "placements", ()) + ): + raise ValueError( + f"batched-matrix Muon parameter {fqn!r} requires exact " + "Shard or Replicate storage placements" + ) + local_storage = param.to_local() if isinstance(param, DTensor) else param + compute_storage = ( + local_storage.detach() if isinstance(param, DTensor) else local_storage + ) + local_storage_shape = torch.Size(local_storage.shape) + if compute_view is None: + global_compute_shape = global_storage_shape + local_compute_tensor = compute_storage + else: + resolved_view = compute_view._resolve(global_storage_shape) + global_compute_shape = resolved_view.compute_shape( + global_storage_shape + ) + local_compute_tensor = compute_storage.view( + resolved_view.compute_shape(local_storage_shape) + ) + prepared_compute_views[fqn] = _PreparedParameterComputeView( + global_compute_shape=global_compute_shape, + local_compute_tensor=local_compute_tensor, + ) + + return DistributedMuon( + prepared_params, + bucket_spec=bucket_spec, + _prepared_compute_views=prepared_compute_views, + **kwargs, + ) diff --git a/torchtitan/components/optimizer.py b/torchtitan/components/optimizer.py index 1a7afaf9a6..03f5cf7f60 100644 --- a/torchtitan/components/optimizer.py +++ b/torchtitan/components/optimizer.py @@ -8,11 +8,12 @@ from collections import defaultdict from collections.abc import Callable, Iterator from dataclasses import dataclass, field -from typing import Any, cast, Generic, Literal, overload, Protocol, TypeVar +from typing import Annotated, Any, cast, Generic, Literal, overload, Protocol, TypeVar import torch import torch.distributed.tensor import torch.nn as nn +import tyro from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl from torch.distributed.checkpoint.stateful import Stateful from torch.distributed.tensor import Replicate @@ -23,6 +24,9 @@ init_optim_state, load_flat_optim_state_dict, ) +from torchtitan.components.distributed_optimizers.muon_parameter_prep import ( + build_distributed_muon, +) from torchtitan.config import Configurable from torchtitan.distributed import ParallelDims from torchtitan.tools.logging import logger @@ -40,8 +44,8 @@ class ParamGroupConfig: """Configuration for a parameter group with its own optimizer. Each entry specifies a regex pattern matching parameter FQNs and a - self-contained optimizer setup. ``optimizer_name`` and ``optimizer_kwargs`` - fully define the optimizer for matched parameters — no implicit inheritance. + self-contained parameter-group setup. ``optimizer_name`` and + ``optimizer_kwargs`` fully define the group — no implicit inheritance. Patterns are checked in order; first match wins. Place specific patterns before broad ones, and use ``r".*"`` as the last entry to catch all @@ -105,6 +109,15 @@ class Config(Configurable.Config): regex pattern and a self-contained optimizer setup. Patterns are checked in order; first match wins.""" + optimizer_init_kwargs: Annotated[ + dict[str, dict[str, Any]], tyro.conf.Suppress + ] = field(default_factory=dict) + """Programmatic optimizer-wide constructor arguments keyed by name. + + Use this for instance-wide objects such as communication bucket specs; + parameter-group hyperparameters belong in ``ParamGroupConfig``. + """ + implementation: Literal[ "for-loop", "foreach", "fused", "fused_opt_states_bf16" ] = "fused" @@ -127,14 +140,15 @@ class Config(Configurable.Config): model_parts: list[nn.Module] @staticmethod - def _resolve_optimizer_cls(name: str) -> type: - optimizer_classes = { + def _resolve_optimizer_factory(name: str) -> Callable[..., Optimizer]: + optimizer_factories: dict[str, Callable[..., Optimizer]] = { "Adam": torch.optim.Adam, "AdamW": torch.optim.AdamW, + "DistributedMuon": build_distributed_muon, } - if name not in optimizer_classes: + if name not in optimizer_factories: raise NotImplementedError(f"Optimizer {name} not added.") - return optimizer_classes[name] + return optimizer_factories[name] @staticmethod def _build_impl_kwargs(config: Config) -> dict[str, Any]: @@ -205,6 +219,14 @@ def _build_param_groups( def __init__(self, config: Config, *, model_parts: list[nn.Module]) -> None: impl_kwargs = self._build_impl_kwargs(config) param_group_configs = config.param_groups + unknown_init_kwargs = config.optimizer_init_kwargs.keys() - { + group.optimizer_name for group in param_group_configs + } + if unknown_init_kwargs: + raise ValueError( + "optimizer_init_kwargs contains unconfigured optimizers: " + f"{sorted(unknown_init_kwargs)}" + ) all_params = [] self.optimizers = [] self.model_parts = model_parts @@ -214,7 +236,10 @@ def __init__(self, config: Config, *, model_parts: list[nn.Module]) -> None: model, param_group_configs, impl_kwargs ) for opt_name, opt_param_groups in groups_by_opt_name.items(): - optimizer = self._resolve_optimizer_cls(opt_name)(opt_param_groups) + optimizer = self._resolve_optimizer_factory(opt_name)( + opt_param_groups, + **config.optimizer_init_kwargs.get(opt_name, {}), + ) self.optimizers.append(optimizer) self._log_optimizer(optimizer, part_idx, patterns_by_opt_name[opt_name]) for group in opt_param_groups: diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index 2a144a6914..104522723f 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -4,11 +4,25 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from torch.distributed.tensor import Shard +from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( + assign_balanced_owners, + BucketConfig, +) from torchtitan.components.checkpoint import CheckpointManager +from torchtitan.components.distributed_optimizers.muon import Owned from torchtitan.components.loss import ChunkedLossWrapper, CrossEntropyLoss from torchtitan.components.lr_scheduler import LRSchedulersContainer from torchtitan.components.metrics import MetricsProcessor -from torchtitan.components.optimizer import default_adamw +from torchtitan.components.distributed_optimizers.muon_parameter_prep import ( + BatchedMatrixComputeView, + MuonComputeSharding, +) +from torchtitan.components.optimizer import ( + default_adamw, + OptimizersContainer, + ParamGroupConfig, +) from torchtitan.components.quantization import ( Float8GroupedExpertsConverter, Float8LinearConverter, @@ -160,6 +174,157 @@ def deepseek_v3_16b() -> Trainer.Config: ) +def deepseek_v3_16b_distributed_muon() -> Trainer.Config: + """DSV3-16B with local-block and bucketed owner-compute Muon.""" + config = deepseek_v3_16b() + owner_group_size = 8 + config.optimizer = _deepseek_v3_distributed_muon_optimizer( + n_layers=27, + num_matrices=16, + wkv_a_matrix_shape=(576, 2048), + owner_group_size=owner_group_size, + lr=2.2e-4, + ) + config.parallelism = ParallelismConfig( + data_parallel_replicate_degree=1, + data_parallel_shard_degree=owner_group_size, + tensor_parallel_degree=1, + context_parallel_degree=1, + pipeline_parallel_degree=1, + expert_parallel_degree=4, + enable_sequence_parallel=False, + spmd_backend="spmd_types", + ) + return config + + +def _deepseek_v3_distributed_muon_optimizer( + *, + n_layers: int, + num_matrices: int, + wkv_a_matrix_shape: tuple[int, int], + owner_group_size: int, + lr: float, +) -> OptimizersContainer.Config: + muon_kwargs = { + "lr": lr, + "weight_decay": 0.1, + "fused": False, + "foreach": False, + } + adamw_kwargs = { + "lr": lr, + "betas": (0.9, 0.95), + "eps": 1e-8, + "weight_decay": 0.1, + "fused": False, + "foreach": True, + } + param_groups = [ + ParamGroupConfig( + pattern=r"attention\.wq\.weight$", + optimizer_name="DistributedMuon", + optimizer_kwargs={ + **muon_kwargs, + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=num_matrices, + matrices_flattened_into_dim=0, + ), + placement=Shard(0), + ), + }, + ), + ParamGroupConfig( + pattern=r"attention\.wkv_a\.weight$", + optimizer_name="DistributedMuon", + optimizer_kwargs={ + **muon_kwargs, + "compute_sharding": MuonComputeSharding(placement=Owned()), + }, + ), + ParamGroupConfig( + pattern=r"attention\.wkv_b\.weight$", + optimizer_name="DistributedMuon", + optimizer_kwargs={ + **muon_kwargs, + "compute_sharding": MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=num_matrices, + matrices_flattened_into_dim=0, + ), + placement=Shard(0), + ), + }, + ), + ] + for projection in ("w1_EFD", "w2_EDF", "w3_EFD"): + param_groups.append( + ParamGroupConfig( + pattern=rf"routed_experts\.inner_experts\.{projection}$", + optimizer_name="DistributedMuon", + optimizer_kwargs={ + **muon_kwargs, + "compute_sharding": MuonComputeSharding( + placement=Shard(0) + ), + }, + ) + ) + param_groups.append( + ParamGroupConfig( + pattern=r".*", + optimizer_name="AdamW", + optimizer_kwargs=adamw_kwargs.copy(), + ) + ) + + def layer_fqns(layer_id: int) -> tuple[str, ...]: + prefix = f"layers.{layer_id}" + fqns = tuple( + f"{prefix}.attention.{projection}.weight" + for projection in ("wq", "wkv_a", "wkv_b") + ) + if layer_id: + fqns += tuple( + f"{prefix}.moe.routed_experts.inner_experts.{projection}" + for projection in ("w1_EFD", "w2_EDF", "w3_EFD") + ) + return fqns + + layer_bucket_fqns = tuple(layer_fqns(layer_id) for layer_id in range(n_layers)) + owner_rank_by_bucket = assign_balanced_owners( + layer_bucket_fqns, + { + f"layers.{layer_id}.attention.wkv_a.weight": ( + wkv_a_matrix_shape[0] * wkv_a_matrix_shape[1] + ) + for layer_id in range(n_layers) + }, + num_ranks=owner_group_size, + ) + bucket_configs = tuple( + BucketConfig( + name=f"layers.{layer_id}", + patterns=fqns, + owner_rank_by_fqn=owners, + mesh_axis="dp_shard", + ) + for layer_id, (fqns, owners) in enumerate( + zip(layer_bucket_fqns, owner_rank_by_bucket, strict=True) + ) + ) + return OptimizersContainer.Config( + implementation="foreach", + param_groups=param_groups, + optimizer_init_kwargs={ + "DistributedMuon": { + "bucket_configs": bucket_configs, + } + }, + ) + + def deepseek_v3_16b_hybridep() -> Trainer.Config: config = deepseek_v3_16b() config.model_spec = model_registry(