From cdad326a2f2b314d5411633a707e77fde4563d45 Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Fri, 31 Jul 2026 11:48:32 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- ...est_deepseek_v3_distributed_muon_config.py | 167 ++ tests/unit_tests/test_distributed_muon.py | 731 +++++++++ .../unit_tests/test_distributed_muon_math.py | 109 ++ torchtitan/components/distributed_muon.py | 1364 +++++++++++++++++ torchtitan/components/optimizer.py | 31 +- .../models/deepseek_v3/config_registry.py | 151 +- 6 files changed, 2548 insertions(+), 5 deletions(-) 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 torchtitan/components/distributed_muon.py 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..df28de5350 --- /dev/null +++ b/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py @@ -0,0 +1,167 @@ +# 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_muon import ( + BucketSpec, + assign_balanced_owners, +) +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 = BucketSpec(patterns=("a",), owner_rank_by_fqn=owners) + 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), + (".attention.wo.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), 186) + 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), 191) + self.assertEqual(adamw_names, model_names - expected_muon_names) + self.assertEqual(len(muon_names | adamw_names), 377) + self.assertFalse(muon_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_b.weight", + ".attention.wo.weight", + ) + } + self.assertEqual( + groups_by_suffix[".attention.wq.weight"]["matrix_shape"], + (192, 2048), + ) + self.assertEqual( + groups_by_suffix[".attention.wkv_b.weight"]["matrix_shape"], + (256, 512), + ) + self.assertEqual( + groups_by_suffix[".attention.wo.weight"]["matrix_shape"], + (2048, 128), + ) + self.assertEqual( + groups_by_suffix[".attention.wo.weight"]["matrix_block_dim"], + 1, + ) + + def test_bucket_and_parallelism_config(self): + optimizer_config = self.config.optimizer + bucket_specs = optimizer_config.optimizer_init_kwargs["DistributedMuon"][ + "bucket_spec" + ] + self.assertEqual(len(bucket_specs), 27) + self.assertTrue(all(isinstance(spec, BucketSpec) for spec in bucket_specs)) + self.assertEqual( + [spec.name for spec in bucket_specs], + [f"layers.{layer_id}" for layer_id in range(27)], + ) + for layer_id, spec in enumerate(bucket_specs): + prefix = f"layers.{layer_id}" + expected = tuple( + f"{prefix}.attention.{projection}.weight" + for projection in ("wq", "wkv_a", "wkv_b", "wo") + ) + if layer_id: + expected += tuple( + f"{prefix}.moe.routed_experts.inner_experts.{projection}" + for projection in ("w1_EFD", "w2_EDF", "w3_EFD") + ) + self.assertEqual(spec.patterns, expected) + self.assertEqual( + spec.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..ef3d0c0d5f --- /dev/null +++ b/tests/unit_tests/test_distributed_muon.py @@ -0,0 +1,731 @@ +# 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, Shard +from torch.testing._internal.distributed._tensor.common_dtensor import ( + DTensorTestBase, + with_comms, +) +from torchtitan.components.distributed_muon import ( + BucketSpec, + DistributedMuon, +) +from torchtitan.components.checkpoint_utils import ( + get_flat_optim_state_dict, + init_optim_state, +) + + +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,)) + 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 DistributedMuon( + [ + { + "params": [redistributed], + "param_names": ["layers.0.redistributed"], + }, + { + "params": [local_blocks], + "param_names": ["layers.0.local_blocks"], + "matrix_shape": (2, 3), + }, + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.redistributed": 1}, + 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_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) + self.assertEqual( + optimizer._plans[0].distributed_bindings[0].owner_rank, + 1, + ) + 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"): + DistributedMuon( + [ + { + "params": [redistributed, local_blocks], + "param_names": [ + "layers.0.redistributed", + "layers.0.local_blocks", + ], + "matrix_shape": (2, 3), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("*.redistributed",), + owner_rank_by_fqn={}, + ) + ], + ) + + with self.assertRaisesRegex(ValueError, "must match one bucket"): + DistributedMuon( + [ + { + "params": [redistributed, local_blocks], + "param_names": [ + "layers.0.redistributed", + "layers.0.local_blocks", + ], + "matrix_shape": (2, 3), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={}, + ), + BucketSpec( + patterns=("*.local_blocks",), + owner_rank_by_fqn={}, + ), + ], + ) + + 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_requires_valid_owner_assignments(self): + 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"], + } + ] + + with self.assertRaisesRegex(ValueError, "exactly cover"): + DistributedMuon( + params, + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.first": 0}, + ) + ], + ) + + with self.assertRaisesRegex(ValueError, "outside its process group"): + DistributedMuon( + params, + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={ + "layers.0.first": 0, + "layers.0.second": self.world_size, + }, + ) + ], + ) + + with self.assertRaisesRegex(ValueError, "invalid DistributedMuon group"): + DistributedMuon( + [ + { + "params": [first], + "param_names": ["layers.0.first"], + "matrix_block_dim": 1, + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={}, + ) + ], + ) + + with self.assertRaisesRegex(ValueError, "whole-matrix-owned"): + DistributedMuon( + [ + { + "params": [first], + "param_names": ["layers.0.first"], + "matrix_shape": (4, 1), + "matrix_block_dim": 1, + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.first": 0}, + ) + ], + ) + + @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"): + DistributedMuon( + [ + { + "params": [redistributed], + "param_names": ["layers.0.redistributed"], + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.redistributed": self.rank}, + ) + ], + ) + + with self.assertRaisesRegex(RuntimeError, "plans differ across ranks"): + DistributedMuon( + [ + { + "params": [redistributed], + "param_names": ["layers.0.redistributed"], + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.redistributed": 0}, + ) + ], + lr=0.01 if self.rank == 0 else 0.02, + ) + + @with_comms + def test_mixed_owned_and_head_sharded_bucket_matches_plain_muon(self): + owned_value = ( + torch.arange(12, device=self.device).reshape(4, 3).float().div_(10) + ) + wo_value = ( + torch.arange(24, device=self.device).reshape(4, 6).float().div_(13) + ) + owned = self._parameter(owned_value) + wo = self._parameter(wo_value) + + def make_optimizer( + owned_param: torch.nn.Parameter, wo_param: torch.nn.Parameter + ) -> DistributedMuon: + return DistributedMuon( + [ + { + "params": [owned_param], + "param_names": ["layers.0.owned"], + }, + { + "params": [wo_param], + "param_names": ["layers.0.wo"], + "matrix_shape": (4, 2), + "matrix_block_dim": 1, + }, + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.owned": 1}, + ) + ], + lr=0.03, + weight_decay=0.2, + momentum=0.8, + nesterov=True, + ns_steps=2, + ) + + optimizer = make_optimizer(owned, wo) + plan = optimizer._plans[0] + self.assertEqual(plan.input_split_sizes, [8, 10]) + self.assertEqual( + plan.output_split_sizes, + [8, 8] if self.mesh.get_local_rank() == 0 else [10, 10], + ) + + reference_owned = torch.nn.Parameter(owned_value.clone()) + reference_heads = tuple( + torch.nn.Parameter(head.clone()) + for head in wo_value.unflatten(1, (3, 2)).movedim(1, 0) + ) + reference = torch.optim.Muon( + [reference_owned, *reference_heads], + lr=0.03, + weight_decay=0.2, + momentum=0.8, + nesterov=True, + ns_steps=2, + ) + + first_owned_grad = ( + torch.arange(1, 13, device=self.device) + .reshape(4, 3) + .float() + .div_(17) + ) + first_wo_grad = ( + torch.arange(1, 25, device=self.device) + .reshape(4, 6) + .float() + .div_(19) + ) + all_to_all_single = dist.all_to_all_single + rank = self.mesh.get_local_rank() + buffer_signature = None + for owned_grad, wo_grad in ( + (first_owned_grad, first_wo_grad), + ( + first_owned_grad.flip(0).contiguous(), + first_wo_grad.flip(0).contiguous(), + ), + ): + owned.grad = distribute_tensor( + owned_grad.clone(), self.mesh, (Shard(0),) + ) + wo.grad = distribute_tensor(wo_grad.clone(), self.mesh, (Shard(0),)) + reference_owned.grad = owned_grad.clone() + for head, head_grad in zip( + reference_heads, + wo_grad.unflatten(1, (3, 2)).movedim(1, 0), + strict=True, + ): + head.grad = head_grad.clone() + + with patch( + "torchtitan.components.distributed_muon.dist.all_to_all_single", + wraps=all_to_all_single, + ) as collective: + optimizer.step() + self.assertEqual(collective.call_count, 2) + reference.step() + + expected_owned = reference_owned.detach().chunk( + self.world_size, dim=0 + )[rank] + expected_wo = ( + torch.stack([head.detach() for head in reference_heads]) + .movedim(0, 1) + .contiguous() + .view(4, 6) + .chunk(self.world_size, dim=0)[rank] + ) + torch.testing.assert_close(owned.to_local(), expected_owned) + torch.testing.assert_close(wo.to_local(), expected_wo) + + expected_owned_momentum = reference.state[reference_owned][ + "momentum_buffer" + ].chunk(self.world_size, dim=0)[rank] + expected_wo_momentum = ( + torch.stack( + [ + reference.state[head]["momentum_buffer"] + for head in reference_heads + ] + ) + .movedim(0, 1) + .contiguous() + .view(4, 6) + .chunk(self.world_size, dim=0)[rank] + ) + torch.testing.assert_close( + optimizer.state[owned]["momentum_buffer"].to_local(), + expected_owned_momentum, + ) + torch.testing.assert_close( + optimizer.state[wo]["momentum_buffer"].to_local(), + expected_wo_momentum, + ) + for param in (owned, wo): + momentum = optimizer.state[param]["momentum_buffer"] + self.assertIsInstance(param, DTensor) + self.assertIsInstance(momentum, DTensor) + self.assertEqual(param.placements, (Shard(0),)) + self.assertEqual(momentum.placements, (Shard(0),)) + self.assertTrue(param.to_local().is_contiguous()) + self.assertTrue(momentum.to_local().is_contiguous()) + + context = optimizer._communication_context + assert context is not None + current_signature = tuple( + tuple( + ( + storage_name, + tuple( + (str(key), tensor.data_ptr(), tensor.numel()) + for key, tensor in storage.items() + ), + ) + for storage_name, storage in ( + ("local", slot.local_storage), + ("routed", slot.routed_storage), + ("compute", slot.compute_storage), + ) + ) + for slot in context.slots + ) + if buffer_signature is None: + buffer_signature = current_signature + else: + self.assertEqual(current_signature, buffer_signature) + + state_dict = optimizer.state_dict() + resumed_owned = self._parameter(reference_owned.detach()) + resumed_wo_value = ( + torch.stack([head.detach() for head in reference_heads]) + .movedim(0, 1) + .contiguous() + .view(4, 6) + ) + resumed_wo = self._parameter(resumed_wo_value) + resumed = make_optimizer(resumed_owned, resumed_wo) + resumed.load_state_dict(state_dict) + self.assertEqual(resumed.param_groups[1]["matrix_block_dim"], 1) + torch.testing.assert_close( + resumed.state[resumed_owned]["momentum_buffer"].to_local(), + optimizer.state[owned]["momentum_buffer"].to_local(), + ) + torch.testing.assert_close( + resumed.state[resumed_wo]["momentum_buffer"].to_local(), + optimizer.state[wo]["momentum_buffer"].to_local(), + ) + + @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_muon.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() + 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_muon.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() >= 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 = DistributedMuon( + [ + { + "params": [distributed_0, distributed_2], + "param_names": [ + "layers.0.redistributed", + "layers.2.redistributed", + ], + }, + { + "params": [local_blocks], + "param_names": ["layers.1.local_blocks"], + "matrix_shape": (2, 3), + }, + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={"layers.0.redistributed": 0}, + ), + BucketSpec( + patterns=("layers.1.*",), + owner_rank_by_fqn={}, + ), + BucketSpec( + patterns=("layers.2.*",), + owner_rank_by_fqn={"layers.2.redistributed": 0}, + ), + ], + 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_muon.dist.all_to_all_single", + wraps=all_to_all_single, + ) as collective: + optimizer.step() + reference.step() + + self.assertEqual(collective.call_count, 4) + 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..a7c0a2d171 --- /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_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/torchtitan/components/distributed_muon.py b/torchtitan/components/distributed_muon.py new file mode 100644 index 0000000000..82f99db101 --- /dev/null +++ b/torchtitan/components/distributed_muon.py @@ -0,0 +1,1364 @@ +# 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 fnmatch +import hashlib +import heapq +import math +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field, replace +from types import ModuleType +from typing import Any, cast + +import torch +import torch.distributed as dist +import torch.distributed.tensor.placement_types as placement_types +from torch import Tensor +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor, Replicate, Shard +from torch.distributed.tensor.placement_types import Placement +from torch.optim import Optimizer + + +__all__ = ["BucketSpec", "assign_balanced_owners", "DistributedMuon"] + +_DEFAULT_NS_COEFFICIENTS = (3.4445, -4.7750, 2.0315) + + +@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. + ``owner_rank_by_fqn`` must exactly cover the bucket's whole-matrix-owned + parameters using process-group-local ranks. Local and matrix-batch-sharded + parameters have no owner entry. ``name`` is diagnostic metadata only. + """ + + patterns: tuple[str, ...] + owner_rank_by_fqn: Mapping[str, int] + 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 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) + + +class DistributedMuon(Optimizer): + """CUDA Muon with bucketed storage-to-matrix-compute routing.""" + + def __init__( + self, + params: Iterable[Tensor] | Iterable[dict[str, Any]], + *, + bucket_spec: Sequence[BucketSpec], + lr: float = 1e-3, + weight_decay: float = 0.1, + momentum: float = 0.95, + nesterov: bool = True, + ns_coefficients: tuple[float, float, float] = _DEFAULT_NS_COEFFICIENTS, + 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, + } + self._communication_context: _CommunicationContext | None = None + self._first_step_validated = False + 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" + + self._control_group = self._infer_control_group() + self._control_group_ranks = tuple( + dist.get_process_group_ranks(self._control_group) + if self._control_group is not None + else range(dist.get_world_size()) + ) + + self._specs = tuple(bucket_spec) + + setup_error: Exception | None = None + try: + self._validate_groups() + self._initialize_plan() + except Exception as error: + setup_error = error + self._synchronize_setup_error(setup_error) + self._validate_plan_across_ranks() + self._frozen_group_metadata = self._group_metadata() + + @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._pipelined_step() + 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: + saved_groups = state_dict.get("param_groups", ()) + if len(saved_groups) != len(self._frozen_group_metadata) or any( + ("param_names" in saved and tuple(saved["param_names"]) != names) + or saved.get("matrix_shape") != matrix_shape + or _effective_matrix_block_dim(saved) != matrix_block_dim + for saved, (names, matrix_shape, matrix_block_dim) in zip( + saved_groups, self._frozen_group_metadata, strict=True + ) + ): + raise ValueError("checkpoint changed DistributedMuon's static plan") + super().load_state_dict(state_dict) + self._validate_plan_across_ranks() + self._first_step_validated = False + + def _infer_control_group(self) -> dist.ProcessGroup | None: + for group in self.param_groups: + for param in group["params"]: + if isinstance(param, DTensor) and param.device_mesh.ndim == 1: + return param.device_mesh.get_group() + return None + + 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"] + matrix_block_dim = group.get("matrix_block_dim") + 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") + or matrix_block_dim not in (None, 0, 1) + or ( + matrix_block_dim is not None + and group.get("matrix_shape") is None + ) + ): + raise ValueError(f"invalid DistributedMuon group {group_index}") + + def _unassigned_bindings(self) -> tuple[_ParamPlan, ...]: + bindings = [] + 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" + ) + matrix_shape = group.get("matrix_shape") + matrix_block_dim = _effective_matrix_block_dim(group) + 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)) + 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" + ) + _validate_matrix_shape(param, matrix_shape) + local_blocks = _local_block_layout( + param, matrix_shape, matrix_block_dim + ) + sharded_blocks = _matrix_batch_shard_layout( + param, matrix_shape, matrix_block_dim + ) + if not local_blocks and not sharded_blocks: + if ( + param.ndim != 2 + or param.device_mesh.ndim != 1 + or len(param.placements) != 1 + or type(param.placements[0]) is not Shard + or param.placements[0].dim % param.ndim != 0 + ): + raise ValueError( + f"Muon parameter {fqn!r} is neither complete local " + "matrix blocks nor 1D Shard(0)" + ) + if not local_blocks: + world_size = param.device_mesh.size() + if sharded_blocks: + assert matrix_shape is not None + matrix_count = param.shape[1] // matrix_shape[1] + if matrix_count < world_size: + raise ValueError( + f"Muon parameter {fqn!r} has fewer matrix blocks " + "than compute ranks" + ) + bindings.append( + _ParamPlan( + fqn=fqn, + param=param, + group_index=group_index, + matrix_shape=matrix_shape, + local_blocks=local_blocks, + sharded_blocks=sharded_blocks, + global_shape=torch.Size(param.shape), + global_stride=tuple(param.stride()), + local_shape=torch.Size(local.shape), + local_stride=tuple(local.stride()), + mesh_ranks=_storage_mesh_ranks(param.device_mesh), + placements=tuple(param.placements), + ) + ) + return tuple(bindings) + + def _resolve_buckets( + self, bindings: tuple[_ParamPlan, ...] + ) -> list[list[_ParamPlan]]: + resolved = [[] for _ in self._specs] + for binding in bindings: + matches = [ + index + for index, spec in enumerate(self._specs) + if any( + fnmatch.fnmatchcase(binding.fqn, pattern) + for pattern in spec.patterns + ) + ] + if len(matches) != 1: + raise ValueError( + f"Muon parameter {binding.fqn!r} must match one bucket" + ) + resolved[matches[0]].append(binding) + return resolved + + def _initialize_plan(self) -> None: + bindings = self._unassigned_bindings() + resolved = self._resolve_buckets(bindings) + plans = [] + planned_bindings = [] + expected_distributed_ranks: tuple[int, ...] | None = None + for spec, bucket in zip(self._specs, resolved, strict=True): + if not bucket: + continue + local_bindings = tuple( + sorted( + (binding for binding in bucket if binding.local_blocks), + key=lambda item: item.fqn, + ) + ) + distributed = tuple( + sorted( + (binding for binding in bucket if not binding.local_blocks), + key=lambda item: item.fqn, + ) + ) + expected_owners = { + binding.fqn + for binding in distributed + if not binding.sharded_blocks + } + 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-matrix-owned parameters; " + f"missing={sorted(expected_owners - provided_owners)}, " + f"extra={sorted(provided_owners - expected_owners)}" + ) + distributed_bindings = tuple( + replace( + binding, + owner_rank=spec.owner_rank_by_fqn.get(binding.fqn, -1), + ) + for binding in distributed + ) + planned_bindings.extend(local_bindings) + planned_bindings.extend(distributed_bindings) + + if not distributed_bindings: + local_tensor = local_bindings[0].param.to_local() + plans.append( + _BucketPlan( + local_bindings=local_bindings, + distributed_bindings=(), + process_group=None, + group_rank=-1, + world_size=0, + input_split_sizes=[], + output_split_sizes=[], + send_segments_by_binding=(), + receive_segments_by_binding=(), + dtype=local_tensor.dtype, + device=local_tensor.device, + local_buffer_numel=0, + routed_buffer_numel=0, + ) + ) + continue + + ranks = distributed_bindings[0].mesh_ranks + if any( + binding.mesh_ranks != ranks for binding in distributed_bindings + ) or ( + expected_distributed_ranks is not None + and ranks != expected_distributed_ranks + ): + raise ValueError( + "redistributed Muon parameters must use one process group" + ) + expected_distributed_ranks = ranks + mesh = distributed_bindings[0].param.device_mesh + process_group = mesh.get_group() + group_rank = mesh.get_local_rank() + world_size = mesh.size() + owner_ranks = [ + binding.owner_rank + for binding in distributed_bindings + if not binding.sharded_blocks + ] + if any(rank not in range(world_size) for rank in owner_ranks): + raise ValueError( + f"bucket {spec.name!r} has owner outside its process group" + ) + local_tensors = [binding.param.to_local() for binding in distributed_bindings] + 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") + ( + input_splits, + output_splits, + send_segments, + receive_segments, + ) = _routing_metadata( + distributed_bindings, group_rank, world_size + ) + plans.append( + _BucketPlan( + local_bindings=local_bindings, + distributed_bindings=distributed_bindings, + process_group=process_group, + group_rank=group_rank, + world_size=world_size, + input_split_sizes=input_splits, + output_split_sizes=output_splits, + send_segments_by_binding=_segments_by_binding( + send_segments, len(distributed_bindings) + ), + receive_segments_by_binding=_segments_by_binding( + receive_segments, len(distributed_bindings) + ), + dtype=dtype, + device=device, + local_buffer_numel=sum(input_splits), + routed_buffer_numel=sum(output_splits), + ) + ) + + self._plans = tuple(plans) + self._bindings = tuple(planned_bindings) + self._tensor_device = self._plans[0].device + if ( + expected_distributed_ranks is not None + and expected_distributed_ranks != self._control_group_ranks + ): + raise ValueError( + "redistributed Muon parameters must use the optimizer control group" + ) + + def _synchronize_setup_error(self, error: Exception | None) -> None: + first_param = cast(DTensor, self.param_groups[0]["params"][0]) + device = first_param.to_local().device + status = torch.tensor(int(error is not None), dtype=torch.int32, device=device) + dist.all_reduce(status, group=self._control_group) + if status.item(): + if error is not None: + raise error + raise RuntimeError("DistributedMuon setup failed on another rank") + + def _validate_plan_across_ranks(self) -> None: + description = [ + ( + str(plan.dtype), + plan.device.type, + plan.world_size, + tuple( + _routing_metadata( + plan.distributed_bindings, rank, plan.world_size + ) + for rank in range(plan.world_size) + ), + [ + ( + binding.fqn, + binding.group_index, + tuple(binding.global_shape), + binding.global_stride, + tuple(binding.local_shape), + binding.local_stride, + str(binding.param.dtype), + binding.param.to_local().device.type, + binding.matrix_shape, + binding.local_blocks, + binding.sharded_blocks, + binding.owner_rank, + binding.mesh_ranks, + tuple(map(str, binding.placements)), + self._group_signature(binding), + ) + for binding in plan.local_bindings + plan.distributed_bindings + ], + ) + for plan in self._plans + ] + 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=self._tensor_device) + gathered = [ + torch.empty_like(local_hash) + for _ in range(dist.get_world_size(self._control_group)) + ] + dist.all_gather(gathered, local_hash, group=self._control_group) + if any(value.item() != plan_hash for value in gathered): + raise RuntimeError("DistributedMuon plans differ across ranks") + + def _group(self, binding: _ParamPlan) -> dict[str, Any]: + return self.param_groups[binding.group_index] + + def _group_signature(self, binding: _ParamPlan) -> tuple[Any, ...]: + group = self._group(binding) + return tuple( + group[key] + for key in ( + "lr", + "weight_decay", + "momentum", + "nesterov", + "ns_coefficients", + "eps", + "ns_steps", + "adjust_lr_fn", + ) + ) + + def _group_metadata(self) -> tuple[tuple[Any, ...], ...]: + return tuple( + ( + tuple(group.get("param_names", ())), + group.get("matrix_shape"), + _effective_matrix_block_dim(group), + ) + for group in self.param_groups + ) + + def _preflight_step(self) -> None: + initialize_state = not self._first_step_validated + if initialize_state: + missing = sum(binding.param.grad is None for binding in self._bindings) + status = torch.tensor( + missing, dtype=torch.int32, device=self._tensor_device + ) + dist.all_reduce(status, group=self._control_group) + if status.item(): + raise RuntimeError( + "DistributedMuon requires every configured gradient" + ) + + gradients = [] + for binding in self._bindings: + grad = self._gradient(binding) + gradients.append((binding, grad)) + if initialize_state: + self._validate_momentum(binding) + + # 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 binding, grad in gradients: + self._momentum(binding, grad) + self._first_step_validated = True + + @staticmethod + def _has_storage_layout(tensor: DTensor, binding: _ParamPlan) -> bool: + local = tensor.to_local() + param_local = binding.param.to_local() + return ( + torch.Size(tensor.shape) == binding.global_shape + and tuple(tensor.stride()) == binding.global_stride + and _storage_mesh_ranks(tensor.device_mesh) == binding.mesh_ranks + and tuple(tensor.placements) == binding.placements + and local.shape == binding.local_shape + and tuple(local.stride()) == binding.local_stride + and local.dtype == param_local.dtype + and local.device == param_local.device + and local.is_contiguous() + ) + + def _gradient(self, binding: _ParamPlan) -> DTensor: + grad = binding.param.grad + if not isinstance(grad, DTensor) or not self._has_storage_layout( + grad, binding + ): + raise RuntimeError(f"gradient layout changed for {binding.fqn!r}") + return grad + + def _validate_momentum(self, binding: _ParamPlan) -> None: + momentum = self.state.get(binding.param, {}).get("momentum_buffer") + if momentum is None: + return + if not isinstance(momentum, DTensor) or not self._has_storage_layout( + momentum, binding + ): + raise RuntimeError(f"momentum layout changed for {binding.fqn!r}") + + def _momentum(self, binding: _ParamPlan, grad: DTensor) -> DTensor: + state = self.state[binding.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, binding: _ParamPlan + ) -> tuple[Tensor, Tensor, dict[str, Any]]: + grad = cast(DTensor, binding.param.grad) + momentum = cast(DTensor, self.state[binding.param]["momentum_buffer"]) + local_grad = grad.to_local() + local_momentum = momentum.to_local() + group = self._group(binding) + 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, binding: _ParamPlan, out: Tensor) -> None: + grad, momentum, group = self._update_local_momentum(binding) + self._write_prepared(group, grad, momentum, out) + + def _compute_update(self, binding: _ParamPlan, prepared: Tensor) -> Tensor: + group = self._group(binding) + logical_prepared = _matrix_view(prepared, binding.matrix_shape) + update = _compute_muon_update( + logical_prepared, + lr=group["lr"], + ns_coefficients=group["ns_coefficients"], + ns_steps=group["ns_steps"], + eps=group["eps"], + adjust_lr_fn=group["adjust_lr_fn"], + out=logical_prepared, + ) + return update.view(prepared.shape) + + def _apply_update(self, binding: _ParamPlan, update: Tensor) -> None: + group = self._group(binding) + local_param = binding.param.to_local() + local_param.mul_(1 - group["lr"] * group["weight_decay"]) + local_param.add_(update) + torch.autograd.graph.increment_version(binding.param) + + def _compute_local_bindings( + self, plan: _BucketPlan, slot: _BufferSlot + ) -> None: + for binding in plan.local_bindings: + local_param = binding.param.to_local() + prepared = slot.compute_buffer( + local_param.shape, + dtype=local_param.dtype, + device=local_param.device, + ) + self._prepare_local(binding, prepared) + self._apply_update(binding, self._compute_update(binding, prepared)) + + def _prepare_distributed( + self, plan: _BucketPlan, local_buffer: Tensor + ) -> None: + for index, binding in enumerate(plan.distributed_bindings): + grad, momentum, group = self._update_local_momentum(binding) + for segment in plan.send_segments_by_binding[index]: + out = local_buffer[ + segment.buffer_offset : segment.buffer_offset + segment.numel + ] + if binding.sharded_blocks: + assert binding.matrix_shape is not None + matrix_columns = binding.matrix_shape[1] + matrix_count = binding.global_shape[1] // matrix_columns + grad_blocks = grad.view( + segment.storage_row_count, + matrix_count, + matrix_columns, + ) + momentum_blocks = momentum.view_as(grad_blocks) + block_slice = slice( + segment.matrix_block_offset, + segment.matrix_block_offset + segment.matrix_block_count, + ) + grad_piece = grad_blocks[:, block_slice, :].movedim(1, 0) + momentum_piece = momentum_blocks[:, block_slice, :].movedim( + 1, 0 + ) + out = out.view( + segment.matrix_block_count, + segment.storage_row_count, + matrix_columns, + ) + self._write_prepared( + group, grad_piece, momentum_piece, out + ) + else: + self._write_prepared( + group, grad, momentum, out.view(grad.shape) + ) + + @staticmethod + def _forward(work: _BucketWork) -> None: + plan = work.plan + assert plan.process_group is not None + dist.all_to_all_single( + work.routed_buffer, + work.local_buffer, + output_split_sizes=plan.output_split_sizes, + input_split_sizes=plan.input_split_sizes, + group=plan.process_group, + ) + + def _compute_redistributed( + self, work: _BucketWork, slot: _BufferSlot + ) -> None: + plan = work.plan + for index, binding in enumerate(plan.distributed_bindings): + compute_shape = _redistributed_compute_shape( + binding, plan.group_rank, plan.world_size + ) + if compute_shape is None: + continue + compute = slot.compute_buffer( + compute_shape, dtype=plan.dtype, device=plan.device + ) + segments = plan.receive_segments_by_binding[index] + for segment in segments: + received = work.routed_buffer[ + segment.buffer_offset : segment.buffer_offset + segment.numel + ] + if binding.sharded_blocks: + assert binding.matrix_shape is not None + matrix_columns = binding.matrix_shape[1] + compute[ + :, + segment.storage_row_offset : ( + segment.storage_row_offset + + segment.storage_row_count + ), + :, + ].copy_( + received.view( + segment.matrix_block_count, + segment.storage_row_count, + matrix_columns, + ) + ) + else: + compute[ + segment.storage_row_offset : ( + segment.storage_row_offset + + segment.storage_row_count + ) + ].copy_(received.view(segment.storage_row_count, -1)) + + self._compute_update(binding, compute) + + for segment in segments: + routed = work.routed_buffer[ + segment.buffer_offset : segment.buffer_offset + segment.numel + ] + if binding.sharded_blocks: + assert binding.matrix_shape is not None + routed.view( + segment.matrix_block_count, + segment.storage_row_count, + binding.matrix_shape[1], + ).copy_( + compute[ + :, + segment.storage_row_offset : ( + segment.storage_row_offset + + segment.storage_row_count + ), + :, + ] + ) + else: + routed.view( + segment.storage_row_count, *binding.global_shape[1:] + ).copy_( + compute[ + segment.storage_row_offset : ( + segment.storage_row_offset + + segment.storage_row_count + ) + ] + ) + + @staticmethod + def _reverse(work: _BucketWork) -> None: + plan = work.plan + assert plan.process_group is not None + dist.all_to_all_single( + work.local_buffer, + work.routed_buffer, + output_split_sizes=plan.input_split_sizes, + input_split_sizes=plan.output_split_sizes, + group=plan.process_group, + ) + + def _finalize_distributed(self, work: _BucketWork) -> None: + plan = work.plan + for index, binding in enumerate(plan.distributed_bindings): + local_param = binding.param.to_local() + segments = plan.send_segments_by_binding[index] + if not binding.sharded_blocks: + assert len(segments) == 1 + segment = segments[0] + update = work.local_buffer[ + segment.buffer_offset : segment.buffer_offset + segment.numel + ].view(local_param.shape) + self._apply_update(binding, update) + continue + + assert binding.matrix_shape is not None + group = self._group(binding) + matrix_columns = binding.matrix_shape[1] + matrix_count = binding.global_shape[1] // matrix_columns + local_blocks = local_param.view( + local_param.shape[0], matrix_count, matrix_columns + ) + local_blocks.mul_(1 - group["lr"] * group["weight_decay"]) + for segment in segments: + update = work.local_buffer[ + segment.buffer_offset : segment.buffer_offset + segment.numel + ].view( + segment.matrix_block_count, + segment.storage_row_count, + matrix_columns, + ) + block_slice = slice( + segment.matrix_block_offset, + segment.matrix_block_offset + segment.matrix_block_count, + ) + local_blocks[:, block_slice, :].add_(update.movedim(0, 1)) + torch.autograd.graph.increment_version(binding.param) + + def _begin_pipelined( + self, + plan: _BucketPlan, + slot: _BufferSlot, + caller_stream: torch.Stream, + context: _CommunicationContext, + ) -> _BucketWork: + handle = context.device_handle + transfer = context.transfer_stream + with handle.stream(transfer): + local_buffer, routed_buffer = slot.communication_buffers(plan) + work = _BucketWork(plan, local_buffer, routed_buffer) + self._prepare_distributed(plan, local_buffer) + self._forward(work) + work.forward_ready = handle.Event() + work.forward_ready.record(transfer) + + with handle.stream(caller_stream): + self._compute_local_bindings(plan, slot) + caller_stream.wait_event(work.forward_ready) + self._compute_redistributed(work, slot) + work.compute_done = handle.Event() + work.compute_done.record(caller_stream) + return work + + def _complete_pipelined( + self, work: _BucketWork, context: _CommunicationContext + ) -> 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) + self._reverse(work) + self._finalize_distributed(work) + work.done = handle.Event() + work.done.record(transfer) + + @staticmethod + def _release_pipelined(work: _BucketWork, caller_stream: torch.Stream) -> None: + assert work.done is not None + caller_stream.wait_event(work.done) + + def _pipelined_step(self) -> None: + if self._communication_context is None: + self._communication_context = _CommunicationContext.create( + self._tensor_device + ) + context = self._communication_context + handle = context.device_handle + caller = handle.current_stream(self._tensor_device) + context.transfer_stream.wait_stream(caller) + + pending: list[_BucketWork] = [] + distributed_index = 0 + try: + for plan in self._plans: + slot = context.slots[distributed_index % 2] + if not plan.distributed_bindings: + with handle.stream(caller): + self._compute_local_bindings(plan, slot) + continue + work = self._begin_pipelined(plan, slot, caller, context) + distributed_index += 1 + pending.append(work) + if len(pending) == 2: + oldest = pending.pop(0) + self._complete_pipelined(oldest, context) + self._release_pipelined(oldest, caller) + for work in pending: + self._complete_pipelined(work, context) + self._release_pipelined(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 + +@dataclass(frozen=True, slots=True) +class _ParamPlan: + fqn: str + param: DTensor + group_index: int + matrix_shape: tuple[int, int] | None + local_blocks: bool + sharded_blocks: bool + global_shape: torch.Size + global_stride: tuple[int, ...] + local_shape: torch.Size + local_stride: tuple[int, ...] + mesh_ranks: tuple[int, ...] + placements: tuple[Placement, ...] + owner_rank: int = -1 + + +@dataclass(frozen=True, slots=True) +class _RouteSegment: + binding_index: int + buffer_offset: int + numel: int + storage_row_offset: int + storage_row_count: int + matrix_block_offset: int = 0 + matrix_block_count: int = 0 + + +@dataclass(slots=True) +class _BucketPlan: + local_bindings: tuple[_ParamPlan, ...] + distributed_bindings: tuple[_ParamPlan, ...] + process_group: dist.ProcessGroup | None + group_rank: int + world_size: int + input_split_sizes: list[int] + output_split_sizes: list[int] + send_segments_by_binding: tuple[tuple[_RouteSegment, ...], ...] + receive_segments_by_binding: tuple[tuple[_RouteSegment, ...], ...] + dtype: torch.dtype + device: torch.device + local_buffer_numel: int + routed_buffer_numel: int + + +@dataclass(slots=True) +class _BucketWork: + plan: _BucketPlan + local_buffer: Tensor + routed_buffer: Tensor + forward_ready: torch.Event | None = None + compute_done: torch.Event | None = None + done: torch.Event | None = None + + +@dataclass(slots=True) +class _BufferSlot: + local_storage: dict[tuple[torch.device, torch.dtype], Tensor] = field( + default_factory=dict + ) + routed_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) -> tuple[Tensor, Tensor]: + return ( + self._ensure_capacity( + self.local_storage, + numel=plan.local_buffer_numel, + dtype=plan.dtype, + device=plan.device, + ), + self._ensure_capacity( + self.routed_storage, + numel=plan.routed_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()), + ) + + +def _is_shard_like(placement: Placement) -> bool: + predicate = getattr(placement_types, "_is_shard_like", None) + if predicate is not None: + return predicate(placement) + strided_shard_type = getattr(placement_types, "_StridedShard", None) + return isinstance(placement, Shard) or ( + strided_shard_type is not None and isinstance(placement, strided_shard_type) + ) + + +def _storage_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 _effective_matrix_block_dim(group: Mapping[str, Any]) -> int | None: + if group.get("matrix_shape") is None: + return None + matrix_block_dim = group.get("matrix_block_dim") + return 0 if matrix_block_dim is None else matrix_block_dim + + +def _validate_matrix_shape( + tensor: Tensor, matrix_shape: tuple[int, int] | None +) -> None: + if matrix_shape is None: + return + if ( + not isinstance(matrix_shape, tuple) + or len(matrix_shape) != 2 + or not all(isinstance(dim, int) and dim > 0 for dim in matrix_shape) + or tensor.numel() % math.prod(matrix_shape) + ): + raise ValueError( + f"invalid matrix_shape {matrix_shape!r} for {tuple(tensor.shape)}" + ) + + +def _matrix_view(tensor: Tensor, matrix_shape: tuple[int, int] | None) -> Tensor: + if matrix_shape is None: + return tensor + return tensor.view(-1, *matrix_shape) + + +# 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, + matrix_shape: torch.Size, +) -> float: + rows, columns = 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 + + +def _local_block_layout( + param: DTensor, + matrix_shape: tuple[int, int] | None, + matrix_block_dim: int | None, +) -> bool: + if matrix_block_dim == 1: + return False + local = param.to_local() + shard_placements = [] + for placement in param.placements: + if _is_shard_like(placement): + shard_dim = getattr(placement, "dim") % param.ndim + if shard_dim != 0: + return False + shard_placements.append(placement) + elif not isinstance(placement, Replicate): + return False + if not shard_placements: + return False + + if matrix_shape is not None: + rows, columns = matrix_shape + if param.shape[-1] != columns or local.shape[-1] != columns: + return False + if local.numel() % (rows * columns): + return False + if param.device_mesh.ndim == 1 and len(param.placements) == 1: + rank = param.device_mesh.get_local_rank() + local_rows, row_offset = Shard.local_shard_size_and_offset( + param.shape[0], param.device_mesh.size(), rank + ) + if row_offset % rows or local_rows % rows: + return False + return True + + return ( + param.ndim > 2 + and tuple(local.shape[-2:]) == tuple(param.shape[-2:]) + ) + + +def _matrix_batch_shard_layout( + param: DTensor, + matrix_shape: tuple[int, int] | None, + matrix_block_dim: int | None, +) -> bool: + if matrix_block_dim != 1: + return False + matrix_rows, matrix_columns = cast(tuple[int, int], matrix_shape) + if ( + param.ndim != 2 + or param.device_mesh.ndim != 1 + or len(param.placements) != 1 + or type(param.placements[0]) is not Shard + or param.placements[0].dim % param.ndim != 0 + or param.shape[0] != matrix_rows + or param.shape[1] % matrix_columns + ): + raise ValueError( + "matrix_block_dim=1 requires a rank-2 1D Shard(0) matrix batch" + ) + return True + + +def _redistributed_compute_shape( + binding: _ParamPlan, group_rank: int, world_size: int +) -> torch.Size | None: + if not binding.sharded_blocks: + return binding.global_shape if binding.owner_rank == group_rank else None + assert binding.matrix_shape is not None + matrix_count = binding.global_shape[1] // binding.matrix_shape[1] + local_count, _ = Shard.local_shard_size_and_offset( + matrix_count, world_size, group_rank + ) + return torch.Size((local_count, *binding.matrix_shape)) + + +def _segments_by_binding( + segments: tuple[_RouteSegment, ...], binding_count: int +) -> tuple[tuple[_RouteSegment, ...], ...]: + return tuple( + tuple( + segment + for segment in segments + if segment.binding_index == binding_index + ) + for binding_index in range(binding_count) + ) + + +def _route_segment( + binding: _ParamPlan, + binding_index: int, + source_rank: int, + destination_rank: int, + world_size: int, + buffer_offset: int, +) -> _RouteSegment | None: + storage_row_count, storage_row_offset = Shard.local_shard_size_and_offset( + binding.global_shape[0], world_size, source_rank + ) + matrix_block_count = 0 + matrix_block_offset = 0 + if binding.sharded_blocks: + assert binding.matrix_shape is not None + matrix_count = binding.global_shape[1] // binding.matrix_shape[1] + matrix_block_count, matrix_block_offset = ( + Shard.local_shard_size_and_offset( + matrix_count, world_size, destination_rank + ) + ) + numel = ( + storage_row_count * matrix_block_count * binding.matrix_shape[1] + ) + elif binding.owner_rank == destination_rank: + numel = storage_row_count * math.prod(binding.global_shape[1:]) + else: + return None + return _RouteSegment( + binding_index=binding_index, + buffer_offset=buffer_offset, + numel=numel, + storage_row_offset=storage_row_offset, + storage_row_count=storage_row_count, + matrix_block_offset=matrix_block_offset, + matrix_block_count=matrix_block_count, + ) + + +def _routing_metadata( + bindings: tuple[_ParamPlan, ...], group_rank: int, world_size: int +) -> tuple[ + list[int], + list[int], + tuple[_RouteSegment, ...], + tuple[_RouteSegment, ...], +]: + input_split_sizes = [] + send_segments = [] + send_cursor = 0 + for destination_rank in range(world_size): + split_start = send_cursor + for binding_index, binding in enumerate(bindings): + segment = _route_segment( + binding, + binding_index, + group_rank, + destination_rank, + world_size, + send_cursor, + ) + if segment is None: + continue + send_segments.append(segment) + send_cursor += segment.numel + input_split_sizes.append(send_cursor - split_start) + + output_split_sizes = [] + receive_segments = [] + receive_cursor = 0 + for source_rank in range(world_size): + split_start = receive_cursor + for binding_index, binding in enumerate(bindings): + segment = _route_segment( + binding, + binding_index, + source_rank, + group_rank, + world_size, + receive_cursor, + ) + if segment is None: + continue + receive_segments.append(segment) + receive_cursor += segment.numel + output_split_sizes.append(receive_cursor - split_start) + expected_send_numel = sum( + Shard.local_shard_size_and_offset( + binding.global_shape[0], world_size, group_rank + )[0] + * math.prod(binding.global_shape[1:]) + for binding in bindings + ) + expected_receive_numel = sum( + math.prod(compute_shape) + for binding in bindings + if ( + compute_shape := _redistributed_compute_shape( + binding, group_rank, world_size + ) + ) + is not None + ) + assert send_cursor == expected_send_numel + assert receive_cursor == expected_receive_numel + return ( + input_split_sizes, + output_split_sizes, + tuple(send_segments), + tuple(receive_segments), + ) diff --git a/torchtitan/components/optimizer.py b/torchtitan/components/optimizer.py index 1a7afaf9a6..644ee319df 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,7 @@ init_optim_state, load_flat_optim_state_dict, ) +from torchtitan.components.distributed_muon import DistributedMuon from torchtitan.config import Configurable from torchtitan.distributed import ParallelDims from torchtitan.tools.logging import logger @@ -40,8 +42,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 +107,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" @@ -131,6 +142,7 @@ def _resolve_optimizer_cls(name: str) -> type: optimizer_classes = { "Adam": torch.optim.Adam, "AdamW": torch.optim.AdamW, + "DistributedMuon": DistributedMuon, } if name not in optimizer_classes: raise NotImplementedError(f"Optimizer {name} not added.") @@ -205,6 +217,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 +234,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_cls(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..2dba97388c 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -5,10 +5,18 @@ # LICENSE file in the root directory of this source tree. from torchtitan.components.checkpoint import CheckpointManager +from torchtitan.components.distributed_muon import ( + BucketSpec, + assign_balanced_owners, +) 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.optimizer import ( + default_adamw, + OptimizersContainer, + ParamGroupConfig, +) from torchtitan.components.quantization import ( Float8GroupedExpertsConverter, Float8LinearConverter, @@ -160,6 +168,147 @@ 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, + wq_matrix_shape=(192, 2048), + wkv_a_matrix_shape=(576, 2048), + wkv_b_matrix_shape=(256, 512), + wo_matrix_shape=(2048, 128), + 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, + wq_matrix_shape: tuple[int, int], + wkv_a_matrix_shape: tuple[int, int], + wkv_b_matrix_shape: tuple[int, int], + wo_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, + "matrix_shape": wq_matrix_shape, + }, + ), + ParamGroupConfig( + pattern=r"attention\.wkv_a\.weight$", + optimizer_name="DistributedMuon", + optimizer_kwargs=muon_kwargs.copy(), + ), + ParamGroupConfig( + pattern=r"attention\.wkv_b\.weight$", + optimizer_name="DistributedMuon", + optimizer_kwargs={ + **muon_kwargs, + "matrix_shape": wkv_b_matrix_shape, + }, + ), + ParamGroupConfig( + pattern=r"attention\.wo\.weight$", + optimizer_name="DistributedMuon", + optimizer_kwargs={ + **muon_kwargs, + "matrix_shape": wo_matrix_shape, + "matrix_block_dim": 1, + }, + ), + ] + 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.copy(), + ) + ) + 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", "wo") + ) + 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_spec = tuple( + BucketSpec( + name=f"layers.{layer_id}", + patterns=fqns, + owner_rank_by_fqn=owners, + ) + 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_spec": bucket_spec} + }, + ) + + def deepseek_v3_16b_hybridep() -> Trainer.Config: config = deepseek_v3_16b() config.model_spec = model_registry(