From d0afbfb220f899eaa12080b8550cfd1fde7ec8f3 Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Mon, 3 Aug 2026 21:55:12 -0700 Subject: [PATCH 01/20] Update [ghstack-poisoned] --- .../test_bucketed_optimizer_redistribution.py | 259 ++++ ...est_deepseek_v3_distributed_muon_config.py | 194 +++ tests/unit_tests/test_distributed_muon.py | 1031 ++++++++++++++ .../unit_tests/test_distributed_muon_math.py | 106 ++ tests/unit_tests/test_muon_parameter_prep.py | 189 +++ .../distributed_optimizers/__init__.py | 1 + .../bucketed_redistribution.py | 1226 +++++++++++++++++ .../components/distributed_optimizers/muon.py | 660 +++++++++ .../muon_parameter_prep.py | 208 +++ torchtitan/components/optimizer.py | 41 +- .../models/deepseek_v3/config_registry.py | 167 ++- 11 files changed, 4073 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..a45ef76f83 --- /dev/null +++ b/tests/unit_tests/test_distributed_muon.py @@ -0,0 +1,1031 @@ +# 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, "layers.0.local_blocks"): + 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, + ) + + local_blocks.grad = None + with self.assertRaisesRegex(RuntimeError, "layers.0.local_blocks"): + optimizer.step() + + 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..a1d34d26b3 --- /dev/null +++ b/tests/unit_tests/test_distributed_muon_math.py @@ -0,0 +1,106 @@ +# 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 itertools +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 dtype, shape in itertools.product( + (torch.bfloat16, torch.float32), ((3, 5), (5, 3)) + ): + with self.subTest(dtype=dtype, shape=shape): + generator = torch.Generator().manual_seed(4) + initial = torch.randn(shape, generator=generator, dtype=dtype) + gradients = [ + torch.randn(shape, generator=generator, dtype=dtype) + 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, adjusted_lr = _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, alpha=-adjusted_lr) + + self.assertTrue(torch.equal(actual_param, reference_param)) + self.assertTrue( + torch.equal( + actual_momentum, + reference.state[reference_param]["momentum_buffer"], + ) + ) + + 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 + )[0] + for matrix in prepared + ] + ) + + self.assertTrue(torch.equal(batched, independent)) 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..1d2e66a9d7 --- /dev/null +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -0,0 +1,660 @@ +# 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 + missing_gradients = [ + compute_layout.fqn + for compute_layout in self._parameter_compute_layouts + if compute_layout.param.grad is None + ] + if missing_gradients: + raise RuntimeError( + "DistributedMuon requires every configured gradient before " + f"step(); missing gradients: {missing_gradients}" + ) + + 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" + ) + 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, direction: 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() + ) + adjusted_lr = _adjust_learning_rate( + group["lr"], + group["adjust_lr_fn"], + compute_layout.global_compute_shape, + ) + local_param.mul_(1 - group["lr"] * group["weight_decay"]) + local_param.add_(direction, alpha=-adjusted_lr) + 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, +) -> tuple[Tensor, float]: + 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) + # Pre-scaling the direction can change FP32 rounding versus Muon's add_. + out.copy_(direction) + return out, adjusted_lr 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( From 3c3e8c0443c26c219fb26d16baa686e7ceb866fc Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Mon, 3 Aug 2026 22:09:54 -0700 Subject: [PATCH 02/20] Update [ghstack-poisoned] --- tests/unit_tests/test_distributed_muon.py | 6 ++++-- torchtitan/components/distributed_optimizers/muon.py | 5 +++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/test_distributed_muon.py b/tests/unit_tests/test_distributed_muon.py index a45ef76f83..15f5ad892d 100644 --- a/tests/unit_tests/test_distributed_muon.py +++ b/tests/unit_tests/test_distributed_muon.py @@ -281,8 +281,10 @@ def test_constructor_requires_exact_bucket_coverage_without_creating_state(self) redistributed.grad = distribute_tensor( torch.ones(4, 3, device=self.device), self.mesh, (Shard(0),) ) - with self.assertRaisesRegex(RuntimeError, "layers.0.local_blocks"): - optimizer.step() + with patch("torch.distributed.all_reduce") as validation_collective: + with self.assertRaisesRegex(RuntimeError, "layers.0.local_blocks"): + optimizer.step() + validation_collective.assert_not_called() self.assertEqual(len(optimizer.state), 0) torch.testing.assert_close(redistributed.to_local(), redistributed_before) redistributed.grad = None diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index 1d2e66a9d7..35a9c0cbf3 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -285,6 +285,11 @@ def _group_signature( ) def _preflight_step(self) -> None: + """Fail the local worker before bucket communication on invalid input. + + TorchTitan's elastic launcher terminates peer workers after this error + escapes. Do not add a validation collective to the optimizer hot path. + """ initialize_state = not self._first_step_validated missing_gradients = [ compute_layout.fqn From 48bbd4b5006eb1100260fa5a2e908c8f93effbb0 Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Mon, 3 Aug 2026 22:24:29 -0700 Subject: [PATCH 03/20] Update [ghstack-poisoned] --- tests/unit_tests/test_distributed_muon.py | 121 +++++++++++------- .../components/distributed_optimizers/muon.py | 41 +++--- .../muon_parameter_prep.py | 6 +- 3 files changed, 97 insertions(+), 71 deletions(-) diff --git a/tests/unit_tests/test_distributed_muon.py b/tests/unit_tests/test_distributed_muon.py index 15f5ad892d..d4870260db 100644 --- a/tests/unit_tests/test_distributed_muon.py +++ b/tests/unit_tests/test_distributed_muon.py @@ -628,37 +628,25 @@ def test_shard1_owned_matches_plain_muon(self): ) @with_comms - def test_replicated_storage_matches_plain_muon_without_redistribution(self): - values = [ - torch.arange(offset, offset + 12, device=self.device) + def test_replicated_storage_and_compute_match_plain_muon(self): + value = ( + torch.arange(1, 13, 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 + ) + parameter = torch.nn.Parameter( + distribute_tensor(value.clone(), self.mesh, (Replicate(),)) ) optimizer = build_distributed_muon( [ { - "params": [owned], - "param_names": ["layers.0.owned"], - "compute_sharding": MuonComputeSharding(placement=Owned()), - }, - { - "params": [batched], - "param_names": ["layers.0.batched"], + "params": [parameter], + "param_names": ["layers.0.matrix"], "compute_sharding": MuonComputeSharding( - view_before_placement=BatchedMatrixComputeView( - num_matrices=2, matrices_flattened_into_dim=0 - ), - placement=Shard(0), + placement=Replicate() ), - }, + } ], bucket_spec=[ BucketSpec( @@ -674,25 +662,18 @@ def test_replicated_storage_matches_plain_muon_without_redistribution(self): 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()), - ] + grad = value.flip(0).contiguous() + parameter.grad = distribute_tensor(grad, self.mesh, (Replicate(),)) + reference = torch.nn.Parameter(value.clone()) + reference.grad = grad.clone() reference_optimizer = torch.optim.Muon( - references, + [reference], 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." @@ -703,22 +684,64 @@ def test_replicated_storage_matches_plain_muon_without_redistribution(self): reference_optimizer.step() collective.assert_not_called() - for param, reference in zip( - (owned, batched), references, strict=True + self.assertEqual(parameter.placements, (Replicate(),)) + self.assertEqual(parameter.grad.placements, (Replicate(),)) + momentum = optimizer.state[parameter]["momentum_buffer"] + self.assertEqual(momentum.placements, (Replicate(),)) + torch.testing.assert_close(parameter.to_local(), reference) + torch.testing.assert_close( + momentum.to_local(), + reference_optimizer.state[reference]["momentum_buffer"], + ) + + @with_comms + def test_constructor_rejects_mismatched_replicated_compute_placement(self): + value = torch.arange(12, device=self.device).reshape(4, 3).float() + replicated = torch.nn.Parameter( + distribute_tensor(value.clone(), self.mesh, (Replicate(),)) + ) + sharded = self._parameter(value) + + for storage, placement, compute_view, owners, message in ( + ( + replicated, + Owned(), + None, + {"layers.0.weight": 0}, + "requires compute Replicate", + ), + ( + replicated, + Shard(0), + BatchedMatrixComputeView(num_matrices=2), + {}, + "requires compute Replicate", + ), + (sharded, Replicate(), None, {}, "requires replicated storage"), ): - 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 self.subTest( + storage=storage.placements, placement=placement + ): + with self.assertRaisesRegex(ValueError, message): + build_distributed_muon( + [ + { + "params": [storage], + "param_names": ["layers.0.weight"], + "compute_sharding": MuonComputeSharding( + view_before_placement=compute_view, + placement=placement, + ), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn=owners, + mesh=self.mesh, + ) + ], + ) @with_comms def test_step_matches_plain_muon_and_continues_from_state_dict(self): diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index 35a9c0cbf3..a1278e6a26 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -463,7 +463,7 @@ class _ParameterComputeLayout: group_index: int global_compute_shape: torch.Size local_compute_tensor: Tensor - compute_placement: Owned | Shard + compute_placement: Owned | Replicate | Shard compute_locally: bool @@ -524,7 +524,22 @@ def _validate_muon_parameter( ) replicated_storage = _has_replicated_storage(param) - if isinstance(compute_placement, Shard): + if isinstance(compute_placement, Replicate): + if not replicated_storage: + raise ValueError( + f"compute Replicate for {fqn!r} requires 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 replicated_storage: + raise ValueError( + f"replicated storage for {fqn!r} requires compute Replicate" + ) + elif isinstance(compute_placement, Shard): if len(global_compute_shape) < 3: raise ValueError( "compute Shard requires a batch of complete Muon matrices" @@ -538,13 +553,7 @@ def _validate_muon_parameter( 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 ( + if ( local_compute_tensor.shape[1:] != global_compute_shape[1:] or not _has_dim0_sharded_storage(param) ): @@ -558,21 +567,13 @@ def _validate_muon_parameter( 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" + f"owned Muon parameter {fqn!r} requires 1D Shard matrix storage" ) return False @@ -585,10 +586,12 @@ def _normalize_dim(dim: int, ndim: int) -> int: def _compute_placement_key( - placement: Owned | Shard, + placement: Owned | Replicate | Shard, ) -> tuple[Any, ...]: if isinstance(placement, Owned): return ("owned",) + if isinstance(placement, Replicate): + return ("replicate",) return ("shard", placement.dim) diff --git a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py index 4cbd3aa201..9c83f43b03 100644 --- a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py +++ b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py @@ -78,11 +78,11 @@ class MuonComputeSharding: # 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 + placement: Owned | Replicate | Shard def __post_init__(self) -> None: - if not isinstance(self.placement, (Owned, Shard)): - raise TypeError("placement must be Owned or Shard") + if not isinstance(self.placement, (Owned, Replicate, Shard)): + raise TypeError("placement must be Owned, Replicate, or Shard") if self.view_before_placement is not None and not isinstance( self.view_before_placement, BatchedMatrixComputeView ): From 7e53368995de84f85ca9e0e94370fee96317a54a Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Mon, 3 Aug 2026 22:32:50 -0700 Subject: [PATCH 04/20] Update [ghstack-poisoned] --- .../test_bucketed_optimizer_redistribution.py | 4 ++-- .../distributed_optimizers/bucketed_redistribution.py | 11 +++++++++-- torchtitan/components/distributed_optimizers/muon.py | 4 ++-- 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/tests/unit_tests/test_bucketed_optimizer_redistribution.py b/tests/unit_tests/test_bucketed_optimizer_redistribution.py index 4ebc0a4ed2..f41653069f 100644 --- a/tests/unit_tests/test_bucketed_optimizer_redistribution.py +++ b/tests/unit_tests/test_bucketed_optimizer_redistribution.py @@ -13,7 +13,7 @@ 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_bucket_plans, _build_owned_redistribution_plan, _lower_packed_all_to_all, _MatrixBlock, @@ -61,7 +61,7 @@ class Item: "_dtensor_storage_blocks", return_value=blocks, ): - result = _build_bucket_plans( + result = _build_owned_bucket_plans( (item,), ( BucketSpec( diff --git a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py index 3357800f51..cb847375cc 100644 --- a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py +++ b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py @@ -4,7 +4,7 @@ # 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.""" +"""Private bucketed storage-to-compute runtime for DistributedMuon.""" from __future__ import annotations @@ -1038,7 +1038,7 @@ def _dtensor_storage_blocks( ) -def _build_bucket_plans( +def _build_owned_bucket_plans( items: Sequence[_ItemT], specs: Sequence[BucketSpec], *, @@ -1046,6 +1046,13 @@ def _build_bucket_plans( compute_locally: Callable[[_ItemT], bool], storage_dtensor: Callable[[_ItemT], DTensor], ) -> _BucketPlanningResult[_ItemT]: + """Build the local and whole-matrix-owned DistributedMuon plans. + + The active planner supports Replicate -> Replicate and Shard(0) matrix + batches as local compute, plus Shard(...) -> Owned through packed + all-to-all and Owned -> Shard(...) through reverse packed all-to-all. + Other placement transitions are intentionally unsupported. + """ resolved = _resolve_buckets(items, specs, fqn=fqn) plans = [] ordered_items = [] diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index a1278e6a26..1cf5db066d 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -21,7 +21,7 @@ from .bucketed_redistribution import ( _BucketedRedistributionRuntime, _BucketPlan, - _build_bucket_plans, + _build_owned_bucket_plans, _device_mesh_ranks, _validate_bucket_plans_across_ranks, assign_balanced_owners, @@ -228,7 +228,7 @@ def _build_parameter_compute_layouts( def _initialize_plan(self) -> None: compute_layouts = self._build_parameter_compute_layouts() - result = _build_bucket_plans( + result = _build_owned_bucket_plans( compute_layouts, self._specs, fqn=lambda item: item.fqn, From 50a429ffab9cc4fde4df8a7c70e7ed49b192f3ce Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Mon, 3 Aug 2026 22:48:38 -0700 Subject: [PATCH 05/20] Update [ghstack-poisoned] --- tests/unit_tests/test_distributed_muon.py | 56 +++++++++++++++++-- .../components/distributed_optimizers/muon.py | 56 +++++++++++++++++-- .../muon_parameter_prep.py | 7 +++ 3 files changed, 110 insertions(+), 9 deletions(-) diff --git a/tests/unit_tests/test_distributed_muon.py b/tests/unit_tests/test_distributed_muon.py index d4870260db..546b460d2b 100644 --- a/tests/unit_tests/test_distributed_muon.py +++ b/tests/unit_tests/test_distributed_muon.py @@ -29,6 +29,7 @@ from torchtitan.components.checkpoint_utils import ( get_flat_optim_state_dict, init_optim_state, + load_flat_optim_state_dict, ) from torchtitan.components.distributed_optimizers.muon_parameter_prep import ( build_distributed_muon, @@ -86,20 +87,27 @@ def _optimizer( self, redistributed: torch.nn.Parameter, local_blocks: torch.nn.Parameter, + *, + local_num_matrices: int = 2, + owner_rank: int = 1, + redistributed_compute_placement: Owned | Replicate = Owned(), ) -> DistributedMuon: return build_distributed_muon( [ { "params": [redistributed], "param_names": ["layers.0.redistributed"], - "compute_sharding": MuonComputeSharding(placement=Owned()), + "compute_sharding": MuonComputeSharding( + placement=redistributed_compute_placement + ), }, { "params": [local_blocks], "param_names": ["layers.0.local_blocks"], "compute_sharding": MuonComputeSharding( view_before_placement=BatchedMatrixComputeView( - num_matrices=2, matrices_flattened_into_dim=0 + num_matrices=local_num_matrices, + matrices_flattened_into_dim=0, ), placement=Shard(0), ), @@ -108,7 +116,11 @@ def _optimizer( bucket_configs=[ BucketConfig( patterns=("layers.0.*",), - owner_rank_by_fqn={"layers.0.redistributed": 1}, + owner_rank_by_fqn=( + {"layers.0.redistributed": owner_rank} + if isinstance(redistributed_compute_placement, Owned) + else {} + ), mesh_axis="dp_shard", name="layers.0", ) @@ -825,6 +837,7 @@ def test_step_matches_plain_muon_and_continues_from_state_dict(self): optimizer.step() state_dict = optimizer.state_dict() + flat_state_dict = get_flat_optim_state_dict(optimizer) self.assertTrue( all( "compute_sharding" not in group @@ -836,10 +849,43 @@ def test_step_matches_plain_muon_and_continues_from_state_dict(self): resumed_local_blocks = self._parameter( torch.cat([parameter.detach() for parameter in reference_local_blocks]) ) + changed_view_optimizer = self._optimizer( + self._parameter(reference_redistributed.detach()), + self._parameter( + torch.cat( + [parameter.detach() for parameter in reference_local_blocks] + ) + ), + local_num_matrices=4, + ) + with self.assertRaisesRegex(ValueError, "compute layout"): + changed_view_optimizer.load_state_dict(state_dict) + + changed_placement_optimizer = self._optimizer( + torch.nn.Parameter( + distribute_tensor( + reference_redistributed.detach().clone(), + self.mesh, + (Replicate(),), + ) + ), + self._parameter( + torch.cat( + [parameter.detach() for parameter in reference_local_blocks] + ) + ), + redistributed_compute_placement=Replicate(), + ) + with self.assertRaisesRegex(ValueError, "compute layout"): + changed_placement_optimizer.load_state_dict(state_dict) + resumed_optimizer = self._optimizer( - resumed_redistributed, resumed_local_blocks + resumed_redistributed, + resumed_local_blocks, + owner_rank=0, ) - resumed_optimizer.load_state_dict(state_dict) + init_optim_state(resumed_optimizer) + load_flat_optim_state_dict(resumed_optimizer, flat_state_dict) second_redistributed_grad = first_redistributed_grad.flip(0).contiguous() second_local_blocks_grad = first_local_blocks_grad.flip(0).contiguous() diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index 1cf5db066d..5178dc1d80 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -8,6 +8,7 @@ from __future__ import annotations +import hashlib import math from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass @@ -32,6 +33,9 @@ __all__ = ["BucketSpec", "assign_balanced_owners", "Owned"] +_LAYOUT_FINGERPRINT_KEY = "_distributed_muon_layout_fingerprint" +_LAYOUT_FINGERPRINT_VERSION = 1 + @dataclass(frozen=True, slots=True) class Owned: @@ -40,6 +44,7 @@ class Owned: @dataclass(frozen=True, slots=True) class _PreparedParameterComputeView: + compute_view_key: tuple[Any, ...] global_compute_shape: torch.Size local_compute_tensor: Tensor @@ -104,6 +109,7 @@ def __init__( self._frozen_param_names = tuple( tuple(group.get("param_names", ())) for group in self.param_groups ) + self._set_checkpoint_layout_fingerprints() @torch.no_grad() def step( @@ -133,8 +139,6 @@ def add_param_group(self, param_group: dict[str, Any]) -> None: 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 @@ -143,6 +147,14 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> None: ) ): raise ValueError("checkpoint changed DistributedMuon's parameter groups") + if any( + saved.get(_LAYOUT_FINGERPRINT_KEY) + != current[_LAYOUT_FINGERPRINT_KEY] + for saved, current in zip( + saved_groups, self.param_groups, strict=True + ) + ): + raise ValueError("checkpoint changed DistributedMuon's compute layout") super().load_state_dict(state_dict) self._validate_plan_across_ranks() self._first_step_validated = False @@ -218,6 +230,7 @@ def _build_parameter_compute_layouts( fqn=fqn, param=param, group_index=group_index, + compute_view_key=prepared.compute_view_key, global_compute_shape=global_compute_shape, local_compute_tensor=local_compute_tensor, compute_placement=compute_placement, @@ -239,6 +252,35 @@ def _initialize_plan(self) -> None: self._parameter_compute_layouts = result.ordered_items self._tensor_device = self._plans[0].device + def _set_checkpoint_layout_fingerprints(self) -> None: + layouts_by_fqn = { + layout.fqn: layout for layout in self._parameter_compute_layouts + } + for group, names in zip( + self.param_groups, self._frozen_param_names, strict=True + ): + entries = [] + for fqn in names: + layout = layouts_by_fqn[fqn] + entries.append( + ( + fqn, + tuple(layout.param.shape), + layout.compute_view_key, + tuple(layout.global_compute_shape), + _compute_placement_key( + layout.compute_placement, + len(layout.global_compute_shape), + ), + ) + ) + # Flat optimizer checkpoints repeat group fields for every FQN, so + # store a fixed-size digest rather than the full group descriptor. + group[_LAYOUT_FINGERPRINT_KEY] = ( + _LAYOUT_FINGERPRINT_VERSION, + hashlib.sha256(repr(tuple(entries)).encode()).hexdigest(), + ) + def _validate_plan_across_ranks(self) -> None: _validate_bucket_plans_across_ranks( self._plans, @@ -257,7 +299,11 @@ def _plan_item_signature( compute_layout.param.to_local().device.type, tuple(compute_layout.global_compute_shape), compute_layout.compute_locally, - _compute_placement_key(compute_layout.compute_placement), + compute_layout.compute_view_key, + _compute_placement_key( + compute_layout.compute_placement, + len(compute_layout.global_compute_shape), + ), _device_mesh_ranks(compute_layout.param.device_mesh), tuple(map(str, compute_layout.param.placements)), self._group_signature(compute_layout), @@ -461,6 +507,7 @@ class _ParameterComputeLayout: fqn: str param: DTensor group_index: int + compute_view_key: tuple[Any, ...] global_compute_shape: torch.Size local_compute_tensor: Tensor compute_placement: Owned | Replicate | Shard @@ -587,12 +634,13 @@ def _normalize_dim(dim: int, ndim: int) -> int: def _compute_placement_key( placement: Owned | Replicate | Shard, + ndim: int, ) -> tuple[Any, ...]: if isinstance(placement, Owned): return ("owned",) if isinstance(placement, Replicate): return ("replicate",) - return ("shard", placement.dim) + return ("shard", _normalize_dim(placement.dim, ndim)) # Keep the functional math aligned with torch.optim.Muon while owning the diff --git a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py index 9c83f43b03..3d4965ad60 100644 --- a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py +++ b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py @@ -185,9 +185,15 @@ def build_distributed_muon( ) local_storage_shape = torch.Size(local_storage.shape) if compute_view is None: + compute_view_key = ("identity",) global_compute_shape = global_storage_shape local_compute_tensor = compute_storage else: + compute_view_key = ( + "batched_matrix", + compute_view.num_matrices, + compute_view.matrices_flattened_into_dim, + ) resolved_view = compute_view._resolve(global_storage_shape) global_compute_shape = resolved_view.compute_shape( global_storage_shape @@ -196,6 +202,7 @@ def build_distributed_muon( resolved_view.compute_shape(local_storage_shape) ) prepared_compute_views[fqn] = _PreparedParameterComputeView( + compute_view_key=compute_view_key, global_compute_shape=global_compute_shape, local_compute_tensor=local_compute_tensor, ) From ce7f26383c5b9403adade32dbc134aa209eb2e69 Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Mon, 3 Aug 2026 23:20:32 -0700 Subject: [PATCH 06/20] Update [ghstack-poisoned] --- ...est_deepseek_v3_distributed_muon_config.py | 4 +- tests/unit_tests/test_distributed_muon.py | 2 +- .../bucketed_redistribution.py | 37 +++++++++++++++---- .../models/deepseek_v3/config_registry.py | 2 +- 4 files changed, 34 insertions(+), 11 deletions(-) diff --git a/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py b/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py index b1f93a80bf..beb66cf421 100644 --- a/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py +++ b/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py @@ -53,7 +53,7 @@ def test_balanced_owner_assignment(self): spec = BucketConfig( patterns=("a",), owner_rank_by_fqn=owners, - mesh_axis="dp_shard", + mesh_axes=("dp_shard",), ) owners["a"] = 1 self.assertEqual(spec.owner_rank_by_fqn, {"a": 0}) @@ -169,7 +169,7 @@ def test_bucket_and_parallelism_config(self): for projection in ("w1_EFD", "w2_EDF", "w3_EFD") ) self.assertEqual(config.patterns, expected) - self.assertEqual(config.mesh_axis, "dp_shard") + self.assertEqual(config.mesh_axes, ("dp_shard",)) self.assertEqual( config.owner_rank_by_fqn, {f"{prefix}.attention.wkv_a.weight": layer_id % 8}, diff --git a/tests/unit_tests/test_distributed_muon.py b/tests/unit_tests/test_distributed_muon.py index 546b460d2b..362660c0a4 100644 --- a/tests/unit_tests/test_distributed_muon.py +++ b/tests/unit_tests/test_distributed_muon.py @@ -121,7 +121,7 @@ def _optimizer( if isinstance(redistributed_compute_placement, Owned) else {} ), - mesh_axis="dp_shard", + mesh_axes=("dp_shard",), name="layers.0", ) ], diff --git a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py index cb847375cc..bbcf00a899 100644 --- a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py +++ b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py @@ -29,16 +29,23 @@ @dataclass(frozen=True, slots=True) class BucketConfig: - """Static bucket configuration resolved after runtime meshes exist.""" + """Static bucket configuration resolved after runtime meshes exist. + + ``mesh_axes`` selects an ordered storage submesh. Multiple axes are + flattened into the one-dimensional communication mesh used by the bucket. + """ patterns: tuple[str, ...] owner_rank_by_fqn: Mapping[str, int] - mesh_axis: str + mesh_axes: tuple[str, ...] name: str = "" def __post_init__(self) -> None: + if isinstance(self.mesh_axes, str) or not self.mesh_axes: + raise ValueError("mesh_axes must be a non-empty sequence of axis names") object.__setattr__(self, "patterns", tuple(self.patterns)) object.__setattr__(self, "owner_rank_by_fqn", dict(self.owner_rank_by_fqn)) + object.__setattr__(self, "mesh_axes", tuple(self.mesh_axes)) def bind(self, mesh: DeviceMesh) -> BucketSpec: return BucketSpec( @@ -92,14 +99,30 @@ def _bind_bucket_configs( 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 + if storage_mesh.mesh_dim_names is None or any( + axis not in storage_mesh.mesh_dim_names + for axis in config.mesh_axes ): raise ValueError( - f"bucket {config.name!r} mesh axis {config.mesh_axis!r} " - f"is not present on storage for {fqn!r}" + f"bucket {config.name!r} mesh axes {config.mesh_axes!r} " + f"are not present on storage for {fqn!r}" + ) + storage_axis_order = tuple( + axis + for axis in storage_mesh.mesh_dim_names + if axis in config.mesh_axes + ) + if storage_axis_order != config.mesh_axes: + raise ValueError( + f"bucket {config.name!r} mesh axes must follow storage " + f"order {storage_mesh.mesh_dim_names!r}" ) - meshes.append(storage_mesh[config.mesh_axis]) + selected_mesh = storage_mesh[config.mesh_axes] + meshes.append( + selected_mesh._flatten() + if selected_mesh.ndim > 1 + else selected_mesh + ) mesh = meshes[0] if any(not torch.equal(candidate.mesh, mesh.mesh) for candidate in meshes[1:]): diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index 104522723f..f8a1f4cebe 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -308,7 +308,7 @@ def layer_fqns(layer_id: int) -> tuple[str, ...]: name=f"layers.{layer_id}", patterns=fqns, owner_rank_by_fqn=owners, - mesh_axis="dp_shard", + mesh_axes=("dp_shard",), ) for layer_id, (fqns, owners) in enumerate( zip(layer_bucket_fqns, owner_rank_by_bucket, strict=True) From 7486252b31c40ab816de4b0823bc70acffe677fe Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Mon, 3 Aug 2026 23:30:05 -0700 Subject: [PATCH 07/20] Update [ghstack-poisoned] --- tests/unit_tests/test_distributed_muon.py | 29 +++++++++++++++++++ .../components/distributed_optimizers/muon.py | 19 ++++++++++-- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/test_distributed_muon.py b/tests/unit_tests/test_distributed_muon.py index 362660c0a4..66d06e712d 100644 --- a/tests/unit_tests/test_distributed_muon.py +++ b/tests/unit_tests/test_distributed_muon.py @@ -255,6 +255,35 @@ def build(parameter, name, compute_placement, owner_rank=None): with self.assertRaisesRegex(RuntimeError, "local storage changed"): optimizer.step() + backing = torch.arange(24, device=self.device).float() + first_region = backing[:12].view(2, 2, 3) + second_region = backing[12:].view(2, 2, 3) + shared_storage = torch.nn.Parameter( + DTensor.from_local( + first_region, + self.mesh, + (Shard(0),), + run_check=False, + ) + ) + shared_optimizer = build(shared_storage, "shared_storage", Shard(0)) + shared_storage.grad = DTensor.from_local( + torch.ones_like(first_region), + self.mesh, + (Shard(0),), + run_check=False, + ) + self.assertEqual( + shared_storage.to_local().untyped_storage().data_ptr(), + second_region.untyped_storage().data_ptr(), + ) + self.assertNotEqual( + shared_storage.to_local().data_ptr(), second_region.data_ptr() + ) + shared_storage._local_tensor = second_region + with self.assertRaisesRegex(RuntimeError, "local storage changed"): + shared_optimizer.step() + dim1_sharded = make_parameter( torch.arange(24, device=self.device).reshape(2, 4, 3).float(), 1 ) diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index 5178dc1d80..08b9cd2e40 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -233,6 +233,9 @@ def _build_parameter_compute_layouts( compute_view_key=prepared.compute_view_key, global_compute_shape=global_compute_shape, local_compute_tensor=local_compute_tensor, + local_storage_signature=_local_storage_signature( + param.to_local() + ), compute_placement=compute_placement, compute_locally=compute_locally, ) @@ -351,8 +354,8 @@ def _preflight_step(self) -> None: 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() + and _local_storage_signature(compute_layout.param.to_local()) + != compute_layout.local_storage_signature ): raise RuntimeError( f"parameter local storage changed for {compute_layout.fqn!r}; " @@ -510,10 +513,22 @@ class _ParameterComputeLayout: compute_view_key: tuple[Any, ...] global_compute_shape: torch.Size local_compute_tensor: Tensor + local_storage_signature: tuple[Any, ...] compute_placement: Owned | Replicate | Shard compute_locally: bool +def _local_storage_signature(tensor: Tensor) -> tuple[Any, ...]: + return ( + tensor.data_ptr(), + tensor.storage_offset(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device, + ) + + def _has_replicated_storage(param: DTensor) -> bool: return all(type(placement) is Replicate for placement in param.placements) From 6633c3e0efbcfb88997ad3ecb49fa31057dd034d Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Mon, 3 Aug 2026 23:45:09 -0700 Subject: [PATCH 08/20] Update [ghstack-poisoned] --- tests/unit_tests/test_muon_parameter_prep.py | 17 +++++++++ .../components/distributed_optimizers/muon.py | 35 ++++++++++++------- .../muon_parameter_prep.py | 11 +++--- 3 files changed, 44 insertions(+), 19 deletions(-) diff --git a/tests/unit_tests/test_muon_parameter_prep.py b/tests/unit_tests/test_muon_parameter_prep.py index cea29a7a6d..3f9d05e190 100644 --- a/tests/unit_tests/test_muon_parameter_prep.py +++ b/tests/unit_tests/test_muon_parameter_prep.py @@ -155,6 +155,23 @@ def test_builder_validates_global_shape_and_aligned_names(self): ) def test_builder_requires_compute_sharding(self): + with self.assertRaisesRegex(TypeError, "named parameter groups"): + build_distributed_muon([torch.empty(2, 2)], bucket_spec=()) + + with self.assertRaisesRegex(TypeError, "DTensor parameters"): + build_distributed_muon( + [ + { + "params": [torch.empty(2, 2)], + "param_names": ["weight"], + "compute_sharding": MuonComputeSharding( + placement=Owned() + ), + } + ], + bucket_spec=(), + ) + with self.assertRaisesRegex(TypeError, "must be a MuonComputeSharding"): build_distributed_muon( [{"params": [], "param_names": [], "compute_sharding": object()}], diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index 08b9cd2e40..4307f2bbf7 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -54,7 +54,7 @@ class DistributedMuon(Optimizer): def __init__( self, - params: Iterable[Tensor] | Iterable[dict[str, Any]], + params: Iterable[dict[str, Any]], *, bucket_spec: Sequence[BucketSpec], _prepared_compute_views: Mapping[ @@ -79,20 +79,14 @@ def __init__( "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 - ] + params = list(params) + if any(not isinstance(param_group, dict) for param_group in params): + raise TypeError("DistributedMuon requires named parameter groups") + params = [dict(param_group) for param_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" + self._tensor_device = self._validate_parameter_storage() group_compute_placements = [] for group in self.param_groups: compute_placement = group.pop("_compute_placement", None) @@ -181,6 +175,22 @@ def _validate_groups(self) -> None: ): raise ValueError(f"invalid DistributedMuon group {group_index}") + def _validate_parameter_storage(self) -> torch.device: + local_devices = set() + for group in self.param_groups: + for param in group["params"]: + if not isinstance(param, DTensor): + raise TypeError("DistributedMuon requires DTensor parameters") + local_device = param.to_local().device + if local_device.type != "cuda": + raise ValueError("DistributedMuon requires CUDA parameters") + local_devices.add(local_device) + if len(local_devices) != 1: + raise ValueError( + "DistributedMuon requires one CUDA device per process" + ) + return local_devices.pop() + def _build_parameter_compute_layouts( self, ) -> tuple[_ParameterComputeLayout, ...]: @@ -253,7 +263,6 @@ def _initialize_plan(self) -> None: ) self._plans = result.plans self._parameter_compute_layouts = result.ordered_items - self._tensor_device = self._plans[0].device def _set_checkpoint_layout_fingerprints(self) -> None: layouts_by_fqn = { diff --git a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py index 3d4965ad60..76e71e2766 100644 --- a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py +++ b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py @@ -117,7 +117,7 @@ def compute_shape(self, storage_shape: torch.Size) -> torch.Size: def build_distributed_muon( - params: Iterable[Tensor] | Iterable[dict[str, Any]], + params: Iterable[dict[str, Any]], *, bucket_spec: Sequence[BucketSpec] | None = None, bucket_configs: Sequence[BucketConfig] | None = None, @@ -129,11 +129,10 @@ def build_distributed_muon( 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) + for param_group in params: + if not isinstance(param_group, dict): + raise TypeError("DistributedMuon requires named parameter groups") + group = dict(param_group) compute_sharding = group.pop("compute_sharding", None) if not isinstance(compute_sharding, MuonComputeSharding): raise TypeError("compute_sharding must be a MuonComputeSharding") From 7a7697459296e5da6984302a1297b41c1047deef Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Mon, 3 Aug 2026 23:50:26 -0700 Subject: [PATCH 09/20] Update [ghstack-poisoned] --- torchtitan/components/optimizer.py | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/torchtitan/components/optimizer.py b/torchtitan/components/optimizer.py index 03f5cf7f60..663b995b10 100644 --- a/torchtitan/components/optimizer.py +++ b/torchtitan/components/optimizer.py @@ -8,12 +8,11 @@ from collections import defaultdict from collections.abc import Callable, Iterator from dataclasses import dataclass, field -from typing import Annotated, Any, cast, Generic, Literal, overload, Protocol, TypeVar +from typing import 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 @@ -44,8 +43,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 parameter-group setup. ``optimizer_name`` and - ``optimizer_kwargs`` fully define the group — no implicit inheritance. + self-contained optimizer setup. ``optimizer_name`` and ``optimizer_kwargs`` + fully define the optimizer for matched parameters — 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 @@ -109,9 +108,7 @@ 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) + optimizer_init_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict) """Programmatic optimizer-wide constructor arguments keyed by name. Use this for instance-wide objects such as communication bucket specs; @@ -219,14 +216,6 @@ 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 From 4889977c7f80cd3904ca103cababa6b0220072c8 Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Tue, 4 Aug 2026 09:51:57 -0700 Subject: [PATCH 10/20] Update [ghstack-poisoned] --- .../unit_tests/test_kimi_k2_7_muon_config.py | 155 ++++++++++++++++++ .../models/kimi_k2_7/config_registry.py | 154 ++++++++++++++++- 2 files changed, 308 insertions(+), 1 deletion(-) create mode 100644 tests/unit_tests/test_kimi_k2_7_muon_config.py diff --git a/tests/unit_tests/test_kimi_k2_7_muon_config.py b/tests/unit_tests/test_kimi_k2_7_muon_config.py new file mode 100644 index 0000000000..d0bca99993 --- /dev/null +++ b/tests/unit_tests/test_kimi_k2_7_muon_config.py @@ -0,0 +1,155 @@ +# 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.muon import Owned +from torchtitan.components.distributed_optimizers.muon_parameter_prep import ( + BatchedMatrixComputeView, + MuonComputeSharding, +) +from torchtitan.components.optimizer import OptimizersContainer +from torchtitan.models.kimi_k2_7.config_registry import ( + kimi_k2_5_muon, +) + + +class TestKimiK25MuonConfig(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.config = kimi_k2_5_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_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())) + expected_muon_names = set() + for suffix, count in ( + (".attention.wq_a.weight", 61), + (".attention.wq_b.weight", 61), + (".attention.wkv_a.weight", 61), + (".attention.wkv_b.weight", 61), + (".moe.routed_experts.inner_experts.w1_EFD", 60), + (".moe.routed_experts.inner_experts.w2_EDF", 60), + (".moe.routed_experts.inner_experts.w3_EFD", 60), + ): + 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"] + } + adamw_names = { + name + for group in groups_by_optimizer["AdamW"] + for name in group["param_names"] + } + self.assertEqual(len(muon_names), 424) + self.assertEqual(muon_names, expected_muon_names) + self.assertEqual(adamw_names, model_names - expected_muon_names) + self.assertFalse(muon_names & adamw_names) + self.assertTrue( + { + name + for name in model_names + if name.endswith(".attention.wo.weight") + } + <= adamw_names + ) + + group_by_suffix = { + suffix: next( + group + for group in muon_groups + if group["param_names"][0].endswith(suffix) + ) + for suffix in ( + ".attention.wq_a.weight", + ".attention.wq_b.weight", + ".attention.wkv_a.weight", + ".attention.wkv_b.weight", + ".moe.routed_experts.inner_experts.w1_EFD", + ) + } + per_head = MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=64, + matrices_flattened_into_dim=0, + ), + placement=Shard(0), + ) + expected_sharding = { + ".attention.wq_a.weight": MuonComputeSharding(placement=Owned()), + ".attention.wq_b.weight": per_head, + ".attention.wkv_a.weight": MuonComputeSharding(placement=Owned()), + ".attention.wkv_b.weight": per_head, + ".moe.routed_experts.inner_experts.w1_EFD": MuonComputeSharding( + placement=Shard(0) + ), + } + for suffix, group in group_by_suffix.items(): + self.assertEqual(group["compute_sharding"], expected_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(len(bucket_configs), 61) + for layer_id, bucket in enumerate(bucket_configs): + prefix = f"layers.{layer_id}" + expected = tuple( + f"{prefix}.attention.{projection}.weight" + for projection in ("wq_a", "wq_b", "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(bucket.name, prefix) + self.assertEqual(bucket.patterns, expected) + self.assertEqual(bucket.mesh_axes, ("dp_shard",)) + self.assertEqual( + set(bucket.owner_rank_by_fqn), + { + f"{prefix}.attention.wq_a.weight", + f"{prefix}.attention.wkv_a.weight", + }, + ) + self.assertTrue( + all(rank in range(64) for rank in bucket.owner_rank_by_fqn.values()) + ) + + parallelism = self.config.parallelism + self.assertEqual(parallelism.data_parallel_replicate_degree, 1) + self.assertEqual(parallelism.data_parallel_shard_degree, 64) + self.assertEqual(parallelism.expert_parallel_degree, 8) + 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") + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/models/kimi_k2_7/config_registry.py b/torchtitan/models/kimi_k2_7/config_registry.py index 5dae9b381b..2c4939575d 100644 --- a/torchtitan/models/kimi_k2_7/config_registry.py +++ b/torchtitan/models/kimi_k2_7/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.distributed_optimizers.muon import Owned +from torchtitan.components.distributed_optimizers.muon_parameter_prep import ( + BatchedMatrixComputeView, + MuonComputeSharding, +) from torchtitan.components.checkpoint import CheckpointManager 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.tokenizer import MultiModalTokenizer from torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig from torchtitan.distributed.activation_checkpoint import FullAC, SelectiveAC @@ -40,6 +54,121 @@ def _mm_dataloader(dataset: str, **kwargs) -> MMDataLoader.Config: ) +def _kimi_k2_5_distributed_muon_optimizer( + *, + n_layers: int, + num_heads: 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, + } + per_head = MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=num_heads, + matrices_flattened_into_dim=0, + ), + placement=Shard(0), + ) + attention_shardings = { + "wq_a": MuonComputeSharding(placement=Owned()), + "wq_b": per_head, + "wkv_a": MuonComputeSharding(placement=Owned()), + "wkv_b": per_head, + } + param_groups = [ + ParamGroupConfig( + pattern=rf"attention\.{projection}\.weight$", + optimizer_name="DistributedMuon", + optimizer_kwargs={ + **muon_kwargs, + "compute_sharding": compute_sharding, + }, + ) + for projection, compute_sharding in attention_shardings.items() + ] + 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, + ) + ) + + def layer_fqns(layer_id: int) -> tuple[str, ...]: + prefix = f"layers.{layer_id}" + fqns = tuple( + f"{prefix}.attention.{projection}.weight" + for projection in attention_shardings + ) + 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)) + owned_projection_numel = { + "wq_a": 1536 * 7168, + "wkv_a": 576 * 7168, + } + owner_rank_by_bucket = assign_balanced_owners( + layer_bucket_fqns, + { + f"layers.{layer_id}.attention.{projection}.weight": numel + for layer_id in range(n_layers) + for projection, numel in owned_projection_numel.items() + }, + num_ranks=owner_group_size, + ) + bucket_configs = tuple( + BucketConfig( + name=f"layers.{layer_id}", + patterns=fqns, + owner_rank_by_fqn=owners, + mesh_axes=("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 kimi_k2_5_debugmodel() -> Trainer.Config: model_spec = model_registry("debugmodel") return Trainer.Config( @@ -181,3 +310,26 @@ def kimi_k2_5() -> Trainer.Config: activation_checkpoint=FullAC.Config(), compile=compile_config, ) + + +def kimi_k2_5_muon() -> Trainer.Config: + """Full Kimi K2.5 with projection and per-expert DistributedMuon.""" + config = kimi_k2_5() + owner_group_size = 64 + config.optimizer = _kimi_k2_5_distributed_muon_optimizer( + n_layers=61, + num_heads=64, + 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=8, + enable_sequence_parallel=False, + spmd_backend="spmd_types", + ) + return config From 7109a6368ad312a83094c2d57c22bc44bc982a9c Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Tue, 4 Aug 2026 10:04:04 -0700 Subject: [PATCH 11/20] Update [ghstack-poisoned] --- .../test_bucketed_optimizer_redistribution.py | 22 ++ ...est_deepseek_v3_distributed_muon_config.py | 194 ------------------ .../models/deepseek_v3/config_registry.py | 167 +-------------- 3 files changed, 23 insertions(+), 360 deletions(-) delete mode 100644 tests/unit_tests/test_deepseek_v3_distributed_muon_config.py diff --git a/tests/unit_tests/test_bucketed_optimizer_redistribution.py b/tests/unit_tests/test_bucketed_optimizer_redistribution.py index f41653069f..03f9857cfd 100644 --- a/tests/unit_tests/test_bucketed_optimizer_redistribution.py +++ b/tests/unit_tests/test_bucketed_optimizer_redistribution.py @@ -21,11 +21,33 @@ _PackedAllToAllSchedule, _RedistributionGroup, _RedistributionPlan, + assign_balanced_owners, + BucketConfig, BucketSpec, ) class TestBucketedOptimizerRedistribution(unittest.TestCase): + 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} + config = BucketConfig( + patterns=("a",), + owner_rank_by_fqn=owners, + mesh_axes=("dp_shard",), + ) + owners["a"] = 1 + self.assertEqual(config.owner_rank_by_fqn, {"a": 0}) + def test_bucket_planner_preserves_empty_local_storage_block(self): @dataclass(frozen=True) class Item: diff --git a/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py b/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py deleted file mode 100644 index beb66cf421..0000000000 --- a/tests/unit_tests/test_deepseek_v3_distributed_muon_config.py +++ /dev/null @@ -1,194 +0,0 @@ -# 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_axes=("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_axes, ("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/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index f8a1f4cebe..2a144a6914 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -4,25 +4,11 @@ # 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.distributed_optimizers.muon_parameter_prep import ( - BatchedMatrixComputeView, - MuonComputeSharding, -) -from torchtitan.components.optimizer import ( - default_adamw, - OptimizersContainer, - ParamGroupConfig, -) +from torchtitan.components.optimizer import default_adamw from torchtitan.components.quantization import ( Float8GroupedExpertsConverter, Float8LinearConverter, @@ -174,157 +160,6 @@ 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_axes=("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( From c0e890eb72102e59a3dc770575141b5fe992323c Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Tue, 4 Aug 2026 10:34:02 -0700 Subject: [PATCH 12/20] Update [ghstack-poisoned] --- .../test_bucketed_optimizer_redistribution.py | 51 ------- .../unit_tests/test_distributed_muon_math.py | 20 +-- .../unit_tests/test_kimi_k2_7_muon_config.py | 10 -- .../bucketed_redistribution.py | 125 +----------------- .../components/distributed_optimizers/muon.py | 10 +- .../models/kimi_k2_7/config_registry.py | 8 +- 6 files changed, 23 insertions(+), 201 deletions(-) diff --git a/tests/unit_tests/test_bucketed_optimizer_redistribution.py b/tests/unit_tests/test_bucketed_optimizer_redistribution.py index 03f9857cfd..83c6f64ab0 100644 --- a/tests/unit_tests/test_bucketed_optimizer_redistribution.py +++ b/tests/unit_tests/test_bucketed_optimizer_redistribution.py @@ -9,7 +9,6 @@ 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 ( @@ -182,56 +181,6 @@ def test_equivalent_replicas_prefer_local_copy_source(self): 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( diff --git a/tests/unit_tests/test_distributed_muon_math.py b/tests/unit_tests/test_distributed_muon_math.py index a1d34d26b3..555a257f1e 100644 --- a/tests/unit_tests/test_distributed_muon_math.py +++ b/tests/unit_tests/test_distributed_muon_math.py @@ -8,7 +8,10 @@ import unittest import torch -from torchtitan.components.distributed_optimizers.muon import _compute_muon_update +from torchtitan.components.distributed_optimizers.muon import ( + _adjust_learning_rate, + _compute_muon_update, +) class TestDistributedMuonMath(unittest.TestCase): @@ -54,14 +57,17 @@ def test_two_steps_match_torch_muon(self): actual_momentum, optimizer_kwargs["momentum"], ) - update, adjusted_lr = _compute_muon_update( + 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"], + ) + adjusted_lr = _adjust_learning_rate( + optimizer_kwargs["lr"], + optimizer_kwargs["adjust_lr_fn"], + prepared.shape, ) actual_param.mul_( 1 @@ -80,25 +86,23 @@ def test_two_steps_match_torch_muon(self): 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( + batched = _compute_muon_update( prepared, out=torch.empty_like(prepared), **kwargs ) independent = torch.stack( [ _compute_muon_update( matrix, out=torch.empty_like(matrix), **kwargs - )[0] + ) for matrix in prepared ] ) diff --git a/tests/unit_tests/test_kimi_k2_7_muon_config.py b/tests/unit_tests/test_kimi_k2_7_muon_config.py index d0bca99993..bef9f37efe 100644 --- a/tests/unit_tests/test_kimi_k2_7_muon_config.py +++ b/tests/unit_tests/test_kimi_k2_7_muon_config.py @@ -63,18 +63,8 @@ def test_parameter_routing(self): for group in groups_by_optimizer["AdamW"] for name in group["param_names"] } - self.assertEqual(len(muon_names), 424) self.assertEqual(muon_names, expected_muon_names) self.assertEqual(adamw_names, model_names - expected_muon_names) - self.assertFalse(muon_names & adamw_names) - self.assertTrue( - { - name - for name in model_names - if name.endswith(".attention.wo.weight") - } - <= adamw_names - ) group_by_suffix = { suffix: next( diff --git a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py index bbcf00a899..15a7cd89d0 100644 --- a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py +++ b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py @@ -199,16 +199,11 @@ def numel(self) -> int: @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. - """ + """Map one logical block from storage holders to compute holders.""" block: _MatrixBlock source_participants: tuple[int, ...] destination_participants: tuple[int, ...] - reduce_op: dist.ReduceOp | None = None @dataclass(frozen=True, slots=True) @@ -361,7 +356,7 @@ class _CommunicationSchedule: def execute( self, output: Tensor, input: Tensor - ) -> tuple[dist.Work, ...]: + ) -> None: raise NotImplementedError @@ -385,7 +380,7 @@ def output_buffer_numel(self) -> int: def execute( self, output: Tensor, input: Tensor - ) -> tuple[dist.Work, ...]: + ) -> None: dist.all_to_all_single( output[: self.output_buffer_numel], input[: self.input_buffer_numel], @@ -393,94 +388,6 @@ def execute( 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) @@ -494,13 +401,12 @@ class _LocalSchedule(_CommunicationSchedule): def execute( self, output: Tensor, input: Tensor - ) -> tuple[dist.Work, ...]: + ) -> None: 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) @@ -536,8 +442,6 @@ class _BucketWork(Generic[_ItemT]): 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) @@ -721,8 +625,7 @@ def _begin( ) 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, + plan.storage_to_compute_schedule.execute( output=compute_fragment_buffer, input=storage_buffer, ) @@ -761,8 +664,7 @@ def _complete( 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, + work.plan.compute_to_storage_schedule.execute( output=work.storage_buffer, input=work.compute_fragment_buffer, ) @@ -814,18 +716,6 @@ def _prepare_redistributed( 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, @@ -888,8 +778,6 @@ def _copy_transfers( } 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__) ) @@ -1245,7 +1133,6 @@ def route_key(route: _MatrixBlockRoute) -> tuple[Any, ...]: route.block.shape, route.source_participants, route.destination_participants, - str(route.reduce_op), ) return ( diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index 4307f2bbf7..fb7da735f6 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -475,11 +475,9 @@ def _compute_update( 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, ) @@ -721,20 +719,16 @@ def _adjust_learning_rate( 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, -) -> tuple[Tensor, float]: +) -> 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) - # Pre-scaling the direction can change FP32 rounding versus Muon's add_. out.copy_(direction) - return out, adjusted_lr + return out diff --git a/torchtitan/models/kimi_k2_7/config_registry.py b/torchtitan/models/kimi_k2_7/config_registry.py index 2c4939575d..bfb3a82798 100644 --- a/torchtitan/models/kimi_k2_7/config_registry.py +++ b/torchtitan/models/kimi_k2_7/config_registry.py @@ -64,7 +64,6 @@ def _kimi_k2_5_distributed_muon_optimizer( muon_kwargs = { "lr": lr, "weight_decay": 0.1, - "fused": False, "foreach": False, } adamw_kwargs = { @@ -72,8 +71,6 @@ def _kimi_k2_5_distributed_muon_optimizer( "betas": (0.9, 0.95), "eps": 1e-8, "weight_decay": 0.1, - "fused": False, - "foreach": True, } per_head = MuonComputeSharding( view_before_placement=BatchedMatrixComputeView( @@ -88,6 +85,7 @@ def _kimi_k2_5_distributed_muon_optimizer( "wkv_a": MuonComputeSharding(placement=Owned()), "wkv_b": per_head, } + expert_projections = ("w1_EFD", "w2_EDF", "w3_EFD") param_groups = [ ParamGroupConfig( pattern=rf"attention\.{projection}\.weight$", @@ -99,7 +97,7 @@ def _kimi_k2_5_distributed_muon_optimizer( ) for projection, compute_sharding in attention_shardings.items() ] - for projection in ("w1_EFD", "w2_EDF", "w3_EFD"): + for projection in expert_projections: param_groups.append( ParamGroupConfig( pattern=rf"routed_experts\.inner_experts\.{projection}$", @@ -129,7 +127,7 @@ def layer_fqns(layer_id: int) -> tuple[str, ...]: if layer_id: fqns += tuple( f"{prefix}.moe.routed_experts.inner_experts.{projection}" - for projection in ("w1_EFD", "w2_EDF", "w3_EFD") + for projection in expert_projections ) return fqns From 2da360ec9f0b783a902fe3d0c31961411404f25c Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Tue, 4 Aug 2026 10:51:06 -0700 Subject: [PATCH 13/20] Update [ghstack-poisoned] --- .../test_bucketed_optimizer_redistribution.py | 33 +---- tests/unit_tests/test_distributed_muon.py | 63 +------- tests/unit_tests/test_muon_parameter_prep.py | 30 +--- .../bucketed_redistribution.py | 100 +++---------- .../components/distributed_optimizers/muon.py | 138 ++++-------------- .../muon_parameter_prep.py | 51 +++---- 6 files changed, 88 insertions(+), 327 deletions(-) diff --git a/tests/unit_tests/test_bucketed_optimizer_redistribution.py b/tests/unit_tests/test_bucketed_optimizer_redistribution.py index 83c6f64ab0..a78e865b43 100644 --- a/tests/unit_tests/test_bucketed_optimizer_redistribution.py +++ b/tests/unit_tests/test_bucketed_optimizer_redistribution.py @@ -13,7 +13,6 @@ from torch.distributed.tensor import DTensor from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( _build_owned_bucket_plans, - _build_owned_redistribution_plan, _lower_packed_all_to_all, _MatrixBlock, _MatrixBlockRoute, @@ -130,13 +129,13 @@ def test_transport_neutral_routes_lower_to_packed_all_to_all(self): ): forward = _lower_packed_all_to_all( (plan,), - direction="storage_to_compute", + storage_to_compute=True, process_group=object(), local_participant=7, ) reverse = _lower_packed_all_to_all( (plan,), - direction="compute_to_storage", + storage_to_compute=False, process_group=object(), local_participant=7, ) @@ -155,32 +154,6 @@ def test_transport_neutral_routes_lower_to_packed_all_to_all(self): (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_routes_require_an_exact_nonoverlapping_partition(self): def plan(blocks): routes = tuple( @@ -197,7 +170,7 @@ def plan(blocks): ( (_MatrixBlock((0, 0), (3, 3)),), ValueError, - "outside", + "in bounds", ), ( ( diff --git a/tests/unit_tests/test_distributed_muon.py b/tests/unit_tests/test_distributed_muon.py index 66d06e712d..3723d1bc0f 100644 --- a/tests/unit_tests/test_distributed_muon.py +++ b/tests/unit_tests/test_distributed_muon.py @@ -287,17 +287,13 @@ def build(parameter, name, compute_placement, owner_rank=None): 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" - ): + with self.assertRaisesRegex(ValueError, "storage-to-compute layout"): 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" - ): + with self.assertRaisesRegex(ValueError, "storage-to-compute layout"): build(owned, "owned", Owned(), owner_rank=0) @with_comms @@ -428,9 +424,6 @@ def test_constructor_rejects_storage_shards_that_split_matrices(self): @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() ) @@ -445,39 +438,7 @@ def test_constructor_requires_valid_owner_assignments(self): } ] - 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"): + with self.assertRaisesRegex(ValueError, "storage-to-compute layout"): build_distributed_muon( [ { @@ -743,34 +704,26 @@ def test_constructor_rejects_mismatched_replicated_compute_placement(self): ) sharded = self._parameter(value) - for storage, placement, compute_view, owners, message in ( + for storage, placement, owners in ( ( replicated, Owned(), - None, {"layers.0.weight": 0}, - "requires compute Replicate", - ), - ( - replicated, - Shard(0), - BatchedMatrixComputeView(num_matrices=2), - {}, - "requires compute Replicate", ), - (sharded, Replicate(), None, {}, "requires replicated storage"), + (sharded, Replicate(), {}), ): with self.subTest( storage=storage.placements, placement=placement ): - with self.assertRaisesRegex(ValueError, message): + with self.assertRaisesRegex( + ValueError, "storage-to-compute layout" + ): build_distributed_muon( [ { "params": [storage], "param_names": ["layers.0.weight"], "compute_sharding": MuonComputeSharding( - view_before_placement=compute_view, placement=placement, ), } diff --git a/tests/unit_tests/test_muon_parameter_prep.py b/tests/unit_tests/test_muon_parameter_prep.py index 3f9d05e190..3f799c4a4f 100644 --- a/tests/unit_tests/test_muon_parameter_prep.py +++ b/tests/unit_tests/test_muon_parameter_prep.py @@ -20,16 +20,8 @@ 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, "positive integer"): + BatchedMatrixComputeView(0) with self.assertRaisesRegex( ValueError, "only matrices_flattened_into_dim=0" ): @@ -118,12 +110,9 @@ def test_builder_compiles_layout_without_mutating_caller_group(self): ) 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"), - ): + for shape in ((2, 3, 4), (5, 4)): with self.subTest(shape=shape): - with self.assertRaisesRegex(ValueError, message): + with self.assertRaisesRegex(ValueError, "cannot be viewed"): build_distributed_muon( [ { @@ -154,10 +143,7 @@ def test_builder_validates_global_shape_and_aligned_names(self): bucket_spec=(), ) - def test_builder_requires_compute_sharding(self): - with self.assertRaisesRegex(TypeError, "named parameter groups"): - build_distributed_muon([torch.empty(2, 2)], bucket_spec=()) - + def test_builder_requires_dtensor_storage(self): with self.assertRaisesRegex(TypeError, "DTensor parameters"): build_distributed_muon( [ @@ -172,12 +158,6 @@ def test_builder_requires_compute_sharding(self): bucket_spec=(), ) - 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)) diff --git a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py index 15a7cd89d0..e6cea1588c 100644 --- a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py +++ b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py @@ -41,8 +41,6 @@ class BucketConfig: name: str = "" def __post_init__(self) -> None: - if isinstance(self.mesh_axes, str) or not self.mesh_axes: - raise ValueError("mesh_axes must be a non-empty sequence of axis names") object.__setattr__(self, "patterns", tuple(self.patterns)) object.__setattr__(self, "owner_rank_by_fqn", dict(self.owner_rank_by_fqn)) object.__setattr__(self, "mesh_axes", tuple(self.mesh_axes)) @@ -92,31 +90,11 @@ def _bind_bucket_configs( 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") + continue 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 any( - axis not in storage_mesh.mesh_dim_names - for axis in config.mesh_axes - ): - raise ValueError( - f"bucket {config.name!r} mesh axes {config.mesh_axes!r} " - f"are not present on storage for {fqn!r}" - ) - storage_axis_order = tuple( - axis - for axis in storage_mesh.mesh_dim_names - if axis in config.mesh_axes - ) - if storage_axis_order != config.mesh_axes: - raise ValueError( - f"bucket {config.name!r} mesh axes must follow storage " - f"order {storage_mesh.mesh_dim_names!r}" - ) selected_mesh = storage_mesh[config.mesh_axes] meshes.append( selected_mesh._flatten() @@ -261,20 +239,18 @@ def _validate_matrix_block_partition( *, 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( + if any(size < 0 for size in logical_shape) or any( + len(block.offsets) != len(logical_shape) + or len(block.shape) != len(logical_shape) + or 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") + ) + for block in blocks + ): + raise ValueError(f"{direction} blocks must be in bounds") positive_blocks = tuple(block for block in blocks if block.numel) for index, first in enumerate(positive_blocks): @@ -343,25 +319,8 @@ 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 - ) -> None: - raise NotImplementedError - - @dataclass(frozen=True, slots=True) -class _PackedAllToAllSchedule(_CommunicationSchedule): +class _PackedAllToAllSchedule: process_group: dist.ProcessGroup participants: tuple[int, ...] local_participant: int @@ -391,7 +350,7 @@ def execute( @dataclass(frozen=True, slots=True) -class _LocalSchedule(_CommunicationSchedule): +class _LocalSchedule: participants: tuple[int, ...] = () local_participant: int = -1 input_spans_by_parameter: tuple[tuple[_PackedSpan, ...], ...] = () @@ -402,13 +361,14 @@ class _LocalSchedule(_CommunicationSchedule): def execute( self, output: Tensor, input: Tensor ) -> None: - 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] ) +_CommunicationSchedule = _PackedAllToAllSchedule | _LocalSchedule + + @dataclass(slots=True) class _BucketPlan(Generic[_ItemT]): local_items: tuple[_ItemT, ...] @@ -803,29 +763,19 @@ def _packed_spans_by_parameter( def _lower_packed_all_to_all( redistribution_plans: tuple[_RedistributionPlan, ...], *, - direction: str, + storage_to_compute: bool, 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": + if storage_to_compute: routes_by_parameter = tuple( plan.storage_to_compute_routes for plan in redistribution_plans ) - elif direction == "compute_to_storage": + else: 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 ) @@ -882,8 +832,6 @@ def _device_mesh_ranks(mesh: DeviceMesh) -> tuple[int, ...]: 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( @@ -893,13 +841,6 @@ def _redistribution_group(mesh: DeviceMesh) -> _RedistributionGroup: ) -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, @@ -917,7 +858,7 @@ def _dtensor_storage_block_for_participant( raise ValueError( "redistributed optimizer storage requires exact Shard placements" ) - tensor_dim = _normalize_dim(placement.dim, tensor.ndim) + tensor_dim = placement.dim % tensor.ndim local_size, global_offset = Shard.local_shard_size_and_offset( tensor.shape[tensor_dim], mesh_shape[mesh_dim], @@ -1063,13 +1004,13 @@ def _build_owned_bucket_plans( group=group, storage_to_compute_schedule=_lower_packed_all_to_all( redistribution_plans, - direction="storage_to_compute", + storage_to_compute=True, 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", + storage_to_compute=False, process_group=group.process_group, local_participant=group.local_participant, ), @@ -1122,7 +1063,6 @@ def _matrix_block_view(tensor: Tensor, block: _MatrixBlock) -> Tensor: for offset, size in zip(block.offsets, block.shape, strict=True) ) ] - assert tuple(view.shape) == block.shape return view diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index fb7da735f6..0764c3752d 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -79,10 +79,6 @@ def __init__( "ns_steps": ns_steps, "adjust_lr_fn": adjust_lr_fn, } - params = list(params) - if any(not isinstance(param_group, dict) for param_group in params): - raise TypeError("DistributedMuon requires named parameter groups") - params = [dict(param_group) for param_group in params] self._first_step_validated = False self._prepared_compute_views = dict(_prepared_compute_views) super().__init__(params, defaults) @@ -100,9 +96,6 @@ def __init__( self._redistribution_runtime = _BucketedRedistributionRuntime[ _ParameterComputeLayout ](self._tensor_device) - self._frozen_param_names = tuple( - tuple(group.get("param_names", ())) for group in self.param_groups - ) self._set_checkpoint_layout_fingerprints() @torch.no_grad() @@ -134,13 +127,6 @@ def add_param_group(self, param_group: dict[str, Any]) -> None: 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_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") if any( saved.get(_LAYOUT_FINGERPRINT_KEY) != current[_LAYOUT_FINGERPRINT_KEY] @@ -155,14 +141,12 @@ def load_state_dict(self, state_dict: dict[str, Any]) -> 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"] if ( - any( + group.get("fused") + or group.get("foreach") + or any( group[name] < 0 for name in ("lr", "weight_decay", "momentum", "eps") ) @@ -173,7 +157,7 @@ def _validate_groups(self) -> None: or group["adjust_lr_fn"] not in (None, "original", "match_rms_adamw", "spectral_unclamped") ): - raise ValueError(f"invalid DistributedMuon group {group_index}") + raise ValueError(f"unsupported DistributedMuon group {group_index}") def _validate_parameter_storage(self) -> torch.device: local_devices = set() @@ -182,10 +166,8 @@ def _validate_parameter_storage(self) -> torch.device: if not isinstance(param, DTensor): raise TypeError("DistributedMuon requires DTensor parameters") local_device = param.to_local().device - if local_device.type != "cuda": - raise ValueError("DistributedMuon requires CUDA parameters") local_devices.add(local_device) - if len(local_devices) != 1: + if len(local_devices) != 1 or next(iter(local_devices)).type != "cuda": raise ValueError( "DistributedMuon requires one CUDA device per process" ) @@ -199,11 +181,7 @@ def _build_parameter_compute_layouts( 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" - ) + names = group["param_names"] 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}") @@ -211,21 +189,10 @@ def _build_parameter_compute_layouts( 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( @@ -268,11 +235,9 @@ def _set_checkpoint_layout_fingerprints(self) -> None: layouts_by_fqn = { layout.fqn: layout for layout in self._parameter_compute_layouts } - for group, names in zip( - self.param_groups, self._frozen_param_names, strict=True - ): + for group in self.param_groups: entries = [] - for fqn in names: + for fqn in group["param_names"]: layout = layouts_by_fqn[fqn] entries.append( ( @@ -554,6 +519,14 @@ def _has_dim0_sharded_storage(param: DTensor) -> bool: return has_shard +def _has_owned_sharded_storage(param: DTensor) -> bool: + return ( + param.device_mesh.ndim == 1 + and len(param.placements) == 1 + and type(param.placements[0]) is Shard + ) + + def _validate_muon_parameter( fqn: str, param: DTensor, @@ -566,85 +539,30 @@ def _validate_muon_parameter( 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, Replicate): - if not replicated_storage: - raise ValueError( - f"compute Replicate for {fqn!r} requires replicated storage" - ) - if local_compute_tensor.shape != global_compute_shape: - raise ValueError( - f"replicated storage for {fqn!r} must contain the complete " - "compute tensor" - ) + if isinstance(compute_placement, Replicate) and replicated_storage: return True - elif replicated_storage: - raise ValueError( - f"replicated storage for {fqn!r} requires compute Replicate" - ) elif 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 ( - local_compute_tensor.shape[1:] != global_compute_shape[1:] - or not _has_dim0_sharded_storage(param) + len(global_compute_shape) >= 3 + and _normalize_dim(compute_placement.dim, len(global_compute_shape)) == 0 + and local_compute_tensor.shape[1:] == global_compute_shape[1:] + and _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" - ) + return True elif ( - param.device_mesh.ndim != 1 - or len(param.placements) != 1 - or type(param.placements[0]) is not Shard + isinstance(compute_placement, Owned) + and len(global_compute_shape) == 2 + and param.ndim == 2 + and _has_owned_sharded_storage(param) ): - raise ValueError( - f"owned Muon parameter {fqn!r} requires 1D Shard matrix storage" - ) - return False + return False + raise ValueError(f"unsupported storage-to-compute layout for {fqn!r}") def _normalize_dim(dim: int, ndim: int) -> int: diff --git a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py index 76e71e2766..aa0055b945 100644 --- a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py +++ b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py @@ -48,24 +48,21 @@ def __post_init__(self) -> None: 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: + if ( + len(storage_shape) != 2 + or storage_shape[0] == 0 + or storage_shape[0] % self.num_matrices + ): raise ValueError( - f"storage shape {tuple(storage_shape)} is not divisible into " + f"storage shape {tuple(storage_shape)} cannot be viewed as " f"{self.num_matrices} matrices" ) return _ResolvedBatchedMatrixView( - matrix_rows=flattened_extent // self.num_matrices, + matrix_rows=storage_shape[0] // self.num_matrices, matrix_columns=storage_shape[1], ) @@ -81,13 +78,14 @@ class MuonComputeSharding: placement: Owned | Replicate | Shard def __post_init__(self) -> None: - if not isinstance(self.placement, (Owned, Replicate, Shard)): - raise TypeError("placement must be Owned, Replicate, or Shard") - if self.view_before_placement is not None and not isinstance( - self.view_before_placement, BatchedMatrixComputeView + if not isinstance(self.placement, (Owned, Replicate, Shard)) or ( + 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" + "MuonComputeSharding requires a supported view and placement" ) @@ -97,10 +95,9 @@ class _ResolvedBatchedMatrixView: 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 + len(storage_shape) != 2 + or storage_shape[0] % self.matrix_rows or storage_shape[1] != self.matrix_columns ): raise ValueError( @@ -130,12 +127,8 @@ def build_distributed_muon( prepared_params = [] parameters_to_prepare = [] for param_group in params: - if not isinstance(param_group, dict): - raise TypeError("DistributedMuon requires named parameter groups") group = dict(param_group) - compute_sharding = group.pop("compute_sharding", None) - if not isinstance(compute_sharding, MuonComputeSharding): - raise TypeError("compute_sharding must be a MuonComputeSharding") + compute_sharding = group.pop("compute_sharding") compute_view = compute_sharding.view_before_placement group["_compute_placement"] = compute_sharding.placement raw_params = group.get("params", ()) @@ -164,8 +157,8 @@ def build_distributed_muon( 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) + else: + bucket_spec = tuple(bucket_spec) prepared_compute_views = {} for param, fqn, compute_view in parameters_to_prepare: @@ -194,8 +187,12 @@ def build_distributed_muon( compute_view.matrices_flattened_into_dim, ) resolved_view = compute_view._resolve(global_storage_shape) - global_compute_shape = resolved_view.compute_shape( - global_storage_shape + global_compute_shape = torch.Size( + ( + compute_view.num_matrices, + resolved_view.matrix_rows, + resolved_view.matrix_columns, + ) ) local_compute_tensor = compute_storage.view( resolved_view.compute_shape(local_storage_shape) From e74b3b3de0273e3154d5dbf6ee216a6bdf0e024e Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Tue, 4 Aug 2026 11:45:05 -0700 Subject: [PATCH 14/20] Update [ghstack-poisoned] --- .../unit_tests/test_kimi_k2_7_muon_config.py | 45 +++++++++++++--- .../models/kimi_k2_7/config_registry.py | 53 ++++++++++++++++--- 2 files changed, 84 insertions(+), 14 deletions(-) diff --git a/tests/unit_tests/test_kimi_k2_7_muon_config.py b/tests/unit_tests/test_kimi_k2_7_muon_config.py index bef9f37efe..1791b2e4ad 100644 --- a/tests/unit_tests/test_kimi_k2_7_muon_config.py +++ b/tests/unit_tests/test_kimi_k2_7_muon_config.py @@ -46,9 +46,17 @@ def test_parameter_routing(self): (".attention.wq_b.weight", 61), (".attention.wkv_a.weight", 61), (".attention.wkv_b.weight", 61), + (".attention.wo.weight", 61), + (".feed_forward.w1.weight", 1), + (".feed_forward.w2.weight", 1), + (".feed_forward.w3.weight", 1), (".moe.routed_experts.inner_experts.w1_EFD", 60), (".moe.routed_experts.inner_experts.w2_EDF", 60), (".moe.routed_experts.inner_experts.w3_EFD", 60), + (".moe.router.gate.weight", 60), + (".moe.shared_experts.w1.weight", 60), + (".moe.shared_experts.w2.weight", 60), + (".moe.shared_experts.w3.weight", 60), ): names = {name for name in model_names if name.endswith(suffix)} self.assertEqual(len(names), count, suffix) @@ -77,6 +85,10 @@ def test_parameter_routing(self): ".attention.wq_b.weight", ".attention.wkv_a.weight", ".attention.wkv_b.weight", + ".attention.wo.weight", + ".feed_forward.w1.weight", + ".moe.router.gate.weight", + ".moe.shared_experts.w1.weight", ".moe.routed_experts.inner_experts.w1_EFD", ) } @@ -92,6 +104,12 @@ def test_parameter_routing(self): ".attention.wq_b.weight": per_head, ".attention.wkv_a.weight": MuonComputeSharding(placement=Owned()), ".attention.wkv_b.weight": per_head, + ".attention.wo.weight": MuonComputeSharding(placement=Owned()), + ".feed_forward.w1.weight": MuonComputeSharding(placement=Owned()), + ".moe.router.gate.weight": MuonComputeSharding(placement=Owned()), + ".moe.shared_experts.w1.weight": MuonComputeSharding( + placement=Owned() + ), ".moe.routed_experts.inner_experts.w1_EFD": MuonComputeSharding( placement=Shard(0) ), @@ -109,22 +127,37 @@ def test_bucket_and_parallelism_config(self): prefix = f"layers.{layer_id}" expected = tuple( f"{prefix}.attention.{projection}.weight" - for projection in ("wq_a", "wq_b", "wkv_a", "wkv_b") + for projection in ("wq_a", "wq_b", "wkv_a", "wkv_b", "wo") ) - if layer_id: + expected_owners = { + f"{prefix}.attention.{projection}.weight" + for projection in ("wq_a", "wkv_a", "wo") + } + if not layer_id: + dense_fqns = tuple( + f"{prefix}.feed_forward.{projection}.weight" + for projection in ("w1", "w2", "w3") + ) + expected += dense_fqns + expected_owners.update(dense_fqns) + else: expected += tuple( f"{prefix}.moe.routed_experts.inner_experts.{projection}" for projection in ("w1_EFD", "w2_EDF", "w3_EFD") ) + router_fqn = f"{prefix}.moe.router.gate.weight" + shared_fqns = tuple( + f"{prefix}.moe.shared_experts.{projection}.weight" + for projection in ("w1", "w2", "w3") + ) + expected += (router_fqn,) + shared_fqns + expected_owners.update((router_fqn, *shared_fqns)) self.assertEqual(bucket.name, prefix) self.assertEqual(bucket.patterns, expected) self.assertEqual(bucket.mesh_axes, ("dp_shard",)) self.assertEqual( set(bucket.owner_rank_by_fqn), - { - f"{prefix}.attention.wq_a.weight", - f"{prefix}.attention.wkv_a.weight", - }, + expected_owners, ) self.assertTrue( all(rank in range(64) for rank in bucket.owner_rank_by_fqn.values()) diff --git a/torchtitan/models/kimi_k2_7/config_registry.py b/torchtitan/models/kimi_k2_7/config_registry.py index bfb3a82798..42770ea727 100644 --- a/torchtitan/models/kimi_k2_7/config_registry.py +++ b/torchtitan/models/kimi_k2_7/config_registry.py @@ -84,6 +84,7 @@ def _kimi_k2_5_distributed_muon_optimizer( "wq_b": per_head, "wkv_a": MuonComputeSharding(placement=Owned()), "wkv_b": per_head, + "wo": MuonComputeSharding(placement=Owned()), } expert_projections = ("w1_EFD", "w2_EDF", "w3_EFD") param_groups = [ @@ -110,6 +111,23 @@ def _kimi_k2_5_distributed_muon_optimizer( }, ) ) + for pattern in ( + r"feed_forward\.w[123]\.weight$", + r"moe\.router\.gate\.weight$", + r"moe\.shared_experts\.w[123]\.weight$", + ): + param_groups.append( + ParamGroupConfig( + pattern=pattern, + optimizer_name="DistributedMuon", + optimizer_kwargs={ + **muon_kwargs, + "compute_sharding": MuonComputeSharding( + placement=Owned() + ), + }, + ) + ) param_groups.append( ParamGroupConfig( pattern=r".*", @@ -124,24 +142,43 @@ def layer_fqns(layer_id: int) -> tuple[str, ...]: f"{prefix}.attention.{projection}.weight" for projection in attention_shardings ) - if layer_id: + if not layer_id: + fqns += tuple( + f"{prefix}.feed_forward.{projection}.weight" + for projection in ("w1", "w2", "w3") + ) + else: fqns += tuple( f"{prefix}.moe.routed_experts.inner_experts.{projection}" for projection in expert_projections ) + fqns += (f"{prefix}.moe.router.gate.weight",) + fqns += tuple( + f"{prefix}.moe.shared_experts.{projection}.weight" + for projection in ("w1", "w2", "w3") + ) return fqns layer_bucket_fqns = tuple(layer_fqns(layer_id) for layer_id in range(n_layers)) - owned_projection_numel = { - "wq_a": 1536 * 7168, - "wkv_a": 576 * 7168, + owned_parameter_numel_by_suffix = { + "attention.wq_a.weight": 1536 * 7168, + "attention.wkv_a.weight": 576 * 7168, + "attention.wo.weight": 7168 * 8192, + "feed_forward.w1.weight": 18432 * 7168, + "feed_forward.w2.weight": 7168 * 18432, + "feed_forward.w3.weight": 18432 * 7168, + "moe.router.gate.weight": 384 * 7168, + "moe.shared_experts.w1.weight": 2048 * 7168, + "moe.shared_experts.w2.weight": 7168 * 2048, + "moe.shared_experts.w3.weight": 2048 * 7168, } owner_rank_by_bucket = assign_balanced_owners( layer_bucket_fqns, { - f"layers.{layer_id}.attention.{projection}.weight": numel - for layer_id in range(n_layers) - for projection, numel in owned_projection_numel.items() + f"layers.{layer_id}.{suffix}": numel + for layer_id, fqns in enumerate(layer_bucket_fqns) + for suffix, numel in owned_parameter_numel_by_suffix.items() + if f"layers.{layer_id}.{suffix}" in fqns }, num_ranks=owner_group_size, ) @@ -311,7 +348,7 @@ def kimi_k2_5() -> Trainer.Config: def kimi_k2_5_muon() -> Trainer.Config: - """Full Kimi K2.5 with projection and per-expert DistributedMuon.""" + """Full Kimi K2.5 with DistributedMuon for text-tower matrices.""" config = kimi_k2_5() owner_group_size = 64 config.optimizer = _kimi_k2_5_distributed_muon_optimizer( From 0bef542bbade0a5f2205bf6dae71aad8ebc7b88b Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Wed, 5 Aug 2026 18:50:41 -0700 Subject: [PATCH 15/20] Update [ghstack-poisoned] --- tests/unit_tests/test_distributed_muon.py | 128 +++++++++++++++++- .../bucketed_redistribution.py | 4 +- .../components/distributed_optimizers/muon.py | 81 +++++++---- 3 files changed, 182 insertions(+), 31 deletions(-) diff --git a/tests/unit_tests/test_distributed_muon.py b/tests/unit_tests/test_distributed_muon.py index 99ffafac25..62ad2062d2 100644 --- a/tests/unit_tests/test_distributed_muon.py +++ b/tests/unit_tests/test_distributed_muon.py @@ -9,7 +9,7 @@ import torch import torch.distributed as dist -from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.device_mesh import DeviceMesh, 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 ( @@ -411,6 +411,80 @@ def test_constructor_requires_exact_bucket_coverage_without_creating_state(self) 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_flat_state_dict_loads_after_group_membership_changes(self): + values = { + "layers.0.a": torch.arange(12, device=self.device).reshape(4, 3).float(), + "layers.0.b": torch.arange(12, 24, device=self.device) + .reshape(4, 3) + .float(), + } + + def build(names, compute_view=None): + parameters = [ + torch.nn.Parameter( + distribute_tensor(values[name].clone(), self.mesh, (Replicate(),)) + ) + for name in names + ] + optimizer = build_distributed_muon( + [ + { + "params": parameters, + "param_names": names, + "compute_sharding": MuonComputeSharding( + view_before_placement=compute_view, + placement=Replicate(), + ), + } + ], + bucket_spec=[ + BucketSpec( + patterns=("layers.0.*",), + owner_rank_by_fqn={}, + mesh=self.mesh, + ) + ], + ns_steps=2, + ) + return parameters, optimizer + + source_parameters, source_optimizer = build(("layers.0.a", "layers.0.b")) + for parameter in source_parameters: + parameter.grad = distribute_tensor( + torch.ones_like(parameter.to_local()), + self.mesh, + (Replicate(),), + ) + source_optimizer.step() + flat_state_dict = get_flat_optim_state_dict(source_optimizer) + self.assertIn( + "state.layers.0.a._distributed_muon_layout_fingerprint", + flat_state_dict, + ) + + target_parameters, target_optimizer = build(("layers.0.a",)) + init_optim_state(target_optimizer) + load_flat_optim_state_dict(target_optimizer, flat_state_dict) + torch.testing.assert_close( + target_optimizer.state[target_parameters[0]]["momentum_buffer"].to_local(), + source_optimizer.state[source_parameters[0]]["momentum_buffer"].to_local(), + ) + + _, changed_layout_optimizer = build( + ("layers.0.a",), + BatchedMatrixComputeView( + num_matrices=2, + matrices_flattened_into_dim=0, + ), + ) + init_optim_state(changed_layout_optimizer) + with self.assertRaisesRegex(ValueError, "compute layout"): + load_flat_optim_state_dict( + changed_layout_optimizer, + flat_state_dict, + ) + @with_comms def test_constructor_rejects_storage_shards_that_split_matrices(self): parameter = self._parameter( @@ -733,6 +807,58 @@ def test_constructor_rejects_mismatched_replicated_compute_placement(self): ], ) + @with_comms + def test_step_rejects_gradient_with_reordered_mesh(self): + value = torch.arange(12, device=self.device).reshape(4, 3).float() + parameter = self._parameter(value) + 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": 0}, + mesh=self.mesh, + ) + ], + ns_steps=1, + ) + reversed_mesh = DeviceMesh( + self.device_type, + torch.arange(self.world_size - 1, -1, -1), + mesh_dim_names=("dp_shard",), + ) + self.assertEqual( + tuple(dist.get_process_group_ranks(reversed_mesh.get_group())), + tuple(dist.get_process_group_ranks(self.mesh.get_group())), + ) + self.assertNotEqual( + tuple(reversed_mesh.mesh.flatten().tolist()), + tuple(self.mesh.mesh.flatten().tolist()), + ) + parameter.grad = distribute_tensor( + value.flip(0).contiguous(), reversed_mesh, (Shard(0),) + ) + parameter_before = parameter.to_local().clone() + + with patch( + "torchtitan.components.distributed_optimizers.bucketed_redistribution." + "dist.all_to_all_single" + ) as collective: + with self.assertRaisesRegex( + RuntimeError, "gradient storage layout changed" + ): + optimizer.step() + + collective.assert_not_called() + self.assertEqual(len(optimizer.state), 0) + torch.testing.assert_close(parameter.to_local(), parameter_before) + @with_comms def test_step_matches_plain_muon_and_continues_from_state_dict(self): redistributed_value = ( diff --git a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py index 3f56a306a1..9a55c3ee99 100644 --- a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py +++ b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py @@ -812,8 +812,8 @@ def _lower_packed_all_to_all( def _device_mesh_ranks(mesh: DeviceMesh) -> tuple[int, ...]: - if mesh.ndim == 1: - return tuple(dist.get_process_group_ranks(mesh.get_group())) + # Process groups can canonicalize rank order, but DeviceMesh order defines + # which logical shard each global rank holds. return tuple(mesh.mesh.flatten().tolist()) diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index 2b1bad2cc7..8a5627dbec 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -127,13 +127,33 @@ def add_param_group(self, param_group: dict[str, Any]) -> None: raise RuntimeError("DistributedMuon parameter groups are frozen") super().add_param_group(param_group) + def state_dict(self) -> dict[str, Any]: + state_dict = super().state_dict() + for saved_group, current_group in zip( + state_dict["param_groups"], self.param_groups, strict=True + ): + for param_id, fqn in zip( + saved_group["params"], current_group["param_names"], strict=True + ): + state_dict["state"].setdefault(param_id, {})[ + _LAYOUT_FINGERPRINT_KEY + ] = self._layout_fingerprints_by_fqn[fqn] + return state_dict + def load_state_dict(self, state_dict: dict[str, Any]) -> None: saved_groups = state_dict.get("param_groups", ()) - if any( - saved.get(_LAYOUT_FINGERPRINT_KEY) != current[_LAYOUT_FINGERPRINT_KEY] - for saved, current in zip(saved_groups, self.param_groups, strict=True) + saved_state = state_dict.get("state", {}) + for saved_group, current_group in zip( + saved_groups, self.param_groups, strict=True ): - raise ValueError("checkpoint changed DistributedMuon's compute layout") + for param_id, fqn in zip( + saved_group["params"], current_group["param_names"], strict=True + ): + fingerprint = saved_state.get(param_id, {}).get(_LAYOUT_FINGERPRINT_KEY) + if fingerprint != self._layout_fingerprints_by_fqn[fqn]: + raise ValueError( + "checkpoint changed DistributedMuon's compute layout" + ) super().load_state_dict(state_dict) self._validate_plan_across_ranks() self._first_step_validated = False @@ -227,30 +247,23 @@ def _initialize_plan(self) -> None: self._parameter_compute_layouts = result.ordered_items def _set_checkpoint_layout_fingerprints(self) -> None: - layouts_by_fqn = { - layout.fqn: layout for layout in self._parameter_compute_layouts - } - for group in self.param_groups: - entries = [] - for fqn in group["param_names"]: - layout = layouts_by_fqn[fqn] - entries.append( - ( - fqn, - tuple(layout.param.shape), - layout.compute_view_key, - tuple(layout.global_compute_shape), - _compute_placement_key( - layout.compute_placement, - len(layout.global_compute_shape), - ), - ) - ) - # Flat optimizer checkpoints repeat group fields for every FQN, so - # store a fixed-size digest rather than the full group descriptor. - group[_LAYOUT_FINGERPRINT_KEY] = ( + self._layout_fingerprints_by_fqn = {} + for layout in self._parameter_compute_layouts: + descriptor = ( + layout.fqn, + tuple(layout.param.shape), + layout.compute_view_key, + tuple(layout.global_compute_shape), + _compute_placement_key( + layout.compute_placement, + len(layout.global_compute_shape), + ), + ) + self._layout_fingerprints_by_fqn[layout.fqn] = ( _LAYOUT_FINGERPRINT_VERSION, - hashlib.sha256(repr(tuple(entries)).encode()).hexdigest(), + # Optimizer.load_state_dict rebuilds iterable state values via + # type(value)(generator), which round-trips bytes but not strings. + hashlib.sha256(repr(descriptor).encode()).digest(), ) def _validate_plan_across_ranks(self) -> None: @@ -375,7 +388,16 @@ def _gradient(self, compute_layout: _ParameterComputeLayout) -> DTensor: return grad def _validate_momentum(self, compute_layout: _ParameterComputeLayout) -> None: - momentum = self.state.get(compute_layout.param, {}).get("momentum_buffer") + state = self.state.get(compute_layout.param, {}) + fingerprint = state.get(_LAYOUT_FINGERPRINT_KEY) + if ( + fingerprint is not None + and fingerprint != self._layout_fingerprints_by_fqn[compute_layout.fqn] + ): + raise RuntimeError( + f"optimizer state layout changed for {compute_layout.fqn!r}" + ) + momentum = state.get("momentum_buffer") if momentum is None: return if not isinstance(momentum, DTensor) or not self._has_storage_layout( @@ -389,6 +411,9 @@ def _momentum( self, compute_layout: _ParameterComputeLayout, grad: DTensor ) -> DTensor: state = self.state[compute_layout.param] + state[_LAYOUT_FINGERPRINT_KEY] = self._layout_fingerprints_by_fqn[ + compute_layout.fqn + ] if "momentum_buffer" not in state: state["momentum_buffer"] = torch.zeros_like( grad, memory_format=torch.preserve_format From df25de0239d6dce1eafd5b85b9df1a39b7b93e0e Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Wed, 5 Aug 2026 22:32:51 -0700 Subject: [PATCH 16/20] Update [ghstack-poisoned] --- tests/unit_tests/test_kimi_k2_7_muon_config.py | 4 ++++ torchtitan/models/kimi_k2_7/config_registry.py | 7 +++++++ 2 files changed, 11 insertions(+) diff --git a/tests/unit_tests/test_kimi_k2_7_muon_config.py b/tests/unit_tests/test_kimi_k2_7_muon_config.py index 686665b51a..2d8ca5952b 100644 --- a/tests/unit_tests/test_kimi_k2_7_muon_config.py +++ b/tests/unit_tests/test_kimi_k2_7_muon_config.py @@ -70,6 +70,10 @@ def test_parameter_routing(self): } self.assertEqual(muon_names, expected_muon_names) self.assertEqual(adamw_names, model_names - expected_muon_names) + self.assertEqual( + {group["adjust_lr_fn"] for group in muon_groups}, + {"match_rms_adamw"}, + ) group_by_suffix = { suffix: next( diff --git a/torchtitan/models/kimi_k2_7/config_registry.py b/torchtitan/models/kimi_k2_7/config_registry.py index 7e690c3688..0f9fa15866 100644 --- a/torchtitan/models/kimi_k2_7/config_registry.py +++ b/torchtitan/models/kimi_k2_7/config_registry.py @@ -66,6 +66,10 @@ def _kimi_k2_5_distributed_muon_optimizer( "lr": lr, "weight_decay": 0.1, "foreach": False, + # Kimi K2's MuonClip recipe uses 0.2 * sqrt(max(rows, columns)) + # for shape-consistent AdamW-scale updates instead of Muon's original + # aspect-ratio scaling. + "adjust_lr_fn": "match_rms_adamw", } adamw_kwargs = { "lr": lr, @@ -112,6 +116,9 @@ def _kimi_k2_5_distributed_muon_optimizer( ) for pattern in ( r"feed_forward\.w[123]\.weight$", + # Keep the 2D router gate on Muon: this follows the Kimi team's + # matrix-parameter rule, and Moonlight reports a larger SVD-entropy + # gain over AdamW for MoE router weights. r"moe\.router\.gate\.weight$", r"moe\.shared_experts\.w[123]\.weight$", ): From b11cdf0eb64d44c5e1ce9d24e8748d5d17d36427 Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Thu, 6 Aug 2026 09:17:25 -0700 Subject: [PATCH 17/20] Update [ghstack-poisoned] --- .../test_bucketed_optimizer_redistribution.py | 84 +++++++++++++++++++ .../unit_tests/test_kimi_k2_7_muon_config.py | 68 ++++++++------- .../bucketed_redistribution.py | 63 +++++++++----- .../models/kimi_k2_7/config_registry.py | 21 ++++- 4 files changed, 182 insertions(+), 54 deletions(-) diff --git a/tests/unit_tests/test_bucketed_optimizer_redistribution.py b/tests/unit_tests/test_bucketed_optimizer_redistribution.py index b61ca4a973..87363595b5 100644 --- a/tests/unit_tests/test_bucketed_optimizer_redistribution.py +++ b/tests/unit_tests/test_bucketed_optimizer_redistribution.py @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. import unittest +from contextlib import nullcontext from dataclasses import dataclass from unittest.mock import MagicMock, Mock, patch @@ -13,6 +14,7 @@ from torch.distributed.tensor import DTensor from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( _bind_bucket_configs, + _BucketedRedistributionRuntime, _build_owned_bucket_plans, _lower_packed_all_to_all, _MatrixBlock, @@ -27,6 +29,88 @@ class TestBucketedOptimizerRedistribution(unittest.TestCase): + def test_runtime_enqueues_return_before_next_compute(self): + runtime = _BucketedRedistributionRuntime(torch.device("cuda")) + caller_stream = Mock() + context = Mock() + context.device_handle.current_stream.return_value = caller_stream + context.device_handle.stream.side_effect = lambda _stream: nullcontext() + context.device_handle.Event.side_effect = Mock + context.slots = (Mock(), Mock()) + for slot in context.slots: + slot.communication_buffers.return_value = (object(), object()) + runtime._context = context + + plans = tuple( + Mock(redistributed_items=(object(),), local_items=()) for _ in range(3) + ) + plan_names = {plan: f"bucket_{index}" for index, plan in enumerate(plans)} + events = [] + for plan in plans: + plan.storage_to_compute_schedule.execute.side_effect = ( + lambda *, _plan=plan, **_kwargs: events.append( + ("gather", plan_names[_plan]) + ) + ) + plan.compute_to_storage_schedule.execute.side_effect = ( + lambda *, _plan=plan, **_kwargs: events.append( + ("return", plan_names[_plan]) + ) + ) + + def compute_redistributed(work, *_args, **_kwargs): + events.append(("compute", plan_names[work.plan])) + + original_release = runtime._release + + def release(work, caller): + events.append(("release", plan_names[work.plan])) + original_release(work, caller) + + with patch( + "torchtitan.components.distributed_optimizers." + "bucketed_redistribution._prepare_redistributed" + ), patch( + "torchtitan.components.distributed_optimizers." + "bucketed_redistribution._compute_redistributed", + side_effect=compute_redistributed, + ), patch( + "torchtitan.components.distributed_optimizers." + "bucketed_redistribution._finalize_redistributed" + ), patch.object( + runtime, + "_release", + side_effect=release, + ): + runtime.run( + plans, + local_tensor_spec=Mock(), + compute_shape=Mock(), + prepare=Mock(), + compute=Mock(), + finalize=Mock(), + ) + + self.assertEqual( + events, + [ + ("gather", "bucket_0"), + ("compute", "bucket_0"), + ("gather", "bucket_1"), + ("return", "bucket_0"), + ("compute", "bucket_1"), + ("release", "bucket_0"), + ("gather", "bucket_2"), + ("return", "bucket_1"), + ("compute", "bucket_2"), + ("release", "bucket_1"), + ("return", "bucket_2"), + ("release", "bucket_2"), + ], + ) + self.assertEqual(context.slots[0].communication_buffers.call_count, 2) + self.assertEqual(context.slots[1].communication_buffers.call_count, 1) + def test_balanced_owner_assignment(self): self.assertEqual( assign_balanced_owners( diff --git a/tests/unit_tests/test_kimi_k2_7_muon_config.py b/tests/unit_tests/test_kimi_k2_7_muon_config.py index 2d8ca5952b..c5b3c78456 100644 --- a/tests/unit_tests/test_kimi_k2_7_muon_config.py +++ b/tests/unit_tests/test_kimi_k2_7_muon_config.py @@ -121,37 +121,47 @@ def test_bucket_and_parallelism_config(self): bucket_configs = optimizer_config.optimizer_init_kwargs["DistributedMuon"][ "bucket_configs" ] - self.assertEqual(len(bucket_configs), 61) - for layer_id, bucket in enumerate(bucket_configs): - prefix = f"layers.{layer_id}" - expected = tuple( - f"{prefix}.attention.{projection}.weight" - for projection in ("wq_a", "wq_b", "wkv_a", "wkv_b", "wo") - ) - expected_owners = { - f"{prefix}.attention.{projection}.weight" - for projection in ("wq_a", "wkv_a", "wo") - } - if not layer_id: - dense_fqns = tuple( - f"{prefix}.feed_forward.{projection}.weight" - for projection in ("w1", "w2", "w3") - ) - expected += dense_fqns - expected_owners.update(dense_fqns) - else: - expected += tuple( - f"{prefix}.moe.routed_experts.inner_experts.{projection}" - for projection in ("w1_EFD", "w2_EDF", "w3_EFD") + bucket_layer_ids = ((0,),) + tuple( + (first_layer_id, first_layer_id + 1) for first_layer_id in range(1, 61, 2) + ) + self.assertEqual(len(bucket_configs), 31) + for layer_ids, bucket in zip(bucket_layer_ids, bucket_configs, strict=True): + expected = () + expected_owners = set() + for layer_id in layer_ids: + prefix = f"layers.{layer_id}" + attention_fqns = tuple( + f"{prefix}.attention.{projection}.weight" + for projection in ("wq_a", "wq_b", "wkv_a", "wkv_b", "wo") ) - router_fqn = f"{prefix}.moe.router.gate.weight" - shared_fqns = tuple( - f"{prefix}.moe.shared_experts.{projection}.weight" - for projection in ("w1", "w2", "w3") + expected += attention_fqns + expected_owners.update( + f"{prefix}.attention.{projection}.weight" + for projection in ("wq_a", "wkv_a", "wo") ) - expected += (router_fqn,) + shared_fqns - expected_owners.update((router_fqn, *shared_fqns)) - self.assertEqual(bucket.name, prefix) + if not layer_id: + dense_fqns = tuple( + f"{prefix}.feed_forward.{projection}.weight" + for projection in ("w1", "w2", "w3") + ) + expected += dense_fqns + expected_owners.update(dense_fqns) + else: + expert_fqns = tuple( + f"{prefix}.moe.routed_experts.inner_experts.{projection}" + for projection in ("w1_EFD", "w2_EDF", "w3_EFD") + ) + router_fqn = f"{prefix}.moe.router.gate.weight" + shared_fqns = tuple( + f"{prefix}.moe.shared_experts.{projection}.weight" + for projection in ("w1", "w2", "w3") + ) + expected += expert_fqns + (router_fqn,) + shared_fqns + expected_owners.update((router_fqn, *shared_fqns)) + self.assertEqual( + bucket.name, + "layers." + "-".join(map(str, layer_ids)), + ) self.assertEqual(bucket.patterns, expected) self.assertEqual(bucket.mesh_axes, ("dp_shard",)) self.assertEqual( diff --git a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py index 9a55c3ee99..883b56b7f6 100644 --- a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py +++ b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py @@ -511,7 +511,7 @@ def run( caller = handle.current_stream(self._device) context.transfer_stream.wait_stream(caller) - pending: list[_BucketWork[_ItemT]] = [] + previous: _BucketWork[_ItemT] | None = None redistributed_index = 0 try: for plan in plans: @@ -527,9 +527,22 @@ def run( finalize=finalize, ) continue - work = self._begin( + + work = self._begin_communication( plan, slot, + context, + prepare=prepare, + ) + redistributed_index += 1 + # Keep collective launches ahead of owner-dependent work: + # gather(current) -> return(previous) -> compute(current). + if previous is not None: + self._complete(previous, context, finalize=finalize) + + self._compute_bucket( + work, + slot, caller, context, local_tensor_spec=local_tensor_spec, @@ -538,15 +551,13 @@ def run( 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) + if previous is not None: + self._release(previous, caller) + previous = work + + if previous is not None: + self._complete(previous, context, finalize=finalize) + self._release(previous, caller) except Exception: # Preserve allocator lifetime ordering for work already enqueued on # either stream. This is an error-path drain, not synchronization. @@ -555,19 +566,12 @@ def run( raise @staticmethod - def _begin( + def _begin_communication( 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 @@ -581,10 +585,28 @@ def _begin( ) work.forward_ready = handle.Event() work.forward_ready.record(transfer) + return work + @staticmethod + def _compute_bucket( + work: _BucketWork[_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], + ) -> None: + assert work.forward_ready is not None + handle = context.device_handle with handle.stream(caller_stream): _BucketedRedistributionRuntime._compute_local( - plan, + work.plan, slot, local_tensor_spec=local_tensor_spec, prepare=prepare, @@ -600,7 +622,6 @@ def _begin( ) work.compute_done = handle.Event() work.compute_done.record(caller_stream) - return work @staticmethod def _complete( diff --git a/torchtitan/models/kimi_k2_7/config_registry.py b/torchtitan/models/kimi_k2_7/config_registry.py index 0f9fa15866..084752db93 100644 --- a/torchtitan/models/kimi_k2_7/config_registry.py +++ b/torchtitan/models/kimi_k2_7/config_registry.py @@ -164,6 +164,16 @@ def layer_fqns(layer_id: int) -> tuple[str, ...]: return fqns layer_bucket_fqns = tuple(layer_fqns(layer_id) for layer_id in range(n_layers)) + # Layer 0 has a much larger dense MLP, so keep it separate while amortizing + # collective launch overhead across pairs of MoE layers. + bucket_layer_ids = ((0,),) + tuple( + tuple(range(first_layer_id, min(first_layer_id + 2, n_layers))) + for first_layer_id in range(1, n_layers, 2) + ) + bucket_fqns = tuple( + tuple(fqn for layer_id in layer_ids for fqn in layer_bucket_fqns[layer_id]) + for layer_ids in bucket_layer_ids + ) owned_parameter_numel_by_suffix = { "attention.wq_a.weight": 1536 * 7168, "attention.wkv_a.weight": 576 * 7168, @@ -177,7 +187,7 @@ def layer_fqns(layer_id: int) -> tuple[str, ...]: "moe.shared_experts.w3.weight": 2048 * 7168, } owner_rank_by_bucket = assign_balanced_owners( - layer_bucket_fqns, + bucket_fqns, { f"layers.{layer_id}.{suffix}": numel for layer_id, fqns in enumerate(layer_bucket_fqns) @@ -188,13 +198,16 @@ def layer_fqns(layer_id: int) -> tuple[str, ...]: ) bucket_configs = tuple( BucketConfig( - name=f"layers.{layer_id}", + name="layers." + "-".join(map(str, layer_ids)), patterns=fqns, owner_rank_by_fqn=owners, mesh_axes=("dp_shard",), ) - for layer_id, (fqns, owners) in enumerate( - zip(layer_bucket_fqns, owner_rank_by_bucket, strict=True) + for layer_ids, fqns, owners in zip( + bucket_layer_ids, + bucket_fqns, + owner_rank_by_bucket, + strict=True, ) ) return OptimizersContainer.Config( From 0a8d0106632bf687571b68f54761cf573dd0463f Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Thu, 6 Aug 2026 11:17:36 -0700 Subject: [PATCH 18/20] Update [ghstack-poisoned] --- .../integration_test_8gpu_features.yaml | 4 + .../unit_tests/test_kimi_k2_7_muon_config.py | 179 +++++++++++++----- .../models/kimi_k2_7/config_registry.py | 142 ++++++++++---- 3 files changed, 238 insertions(+), 87 deletions(-) diff --git a/.github/workflows/integration_test_8gpu_features.yaml b/.github/workflows/integration_test_8gpu_features.yaml index 2a0abfcc0a..38700df521 100644 --- a/.github/workflows/integration_test_8gpu_features.yaml +++ b/.github/workflows/integration_test_8gpu_features.yaml @@ -89,6 +89,10 @@ jobs: end=$(date +%s) echo "pip install torchao took $((end - start)) seconds" + # Exercise distributed optimizer tests that are skipped in CPU unit-test CI. + python -m pytest tests/unit_tests/test_distributed_muon.py \ + --durations=20 -vv + sudo mkdir -p "$RUNNER_TEMP/artifacts-to-be-uploaded" sudo chown -R $(id -u):$(id -g) "$RUNNER_TEMP/artifacts-to-be-uploaded" diff --git a/tests/unit_tests/test_kimi_k2_7_muon_config.py b/tests/unit_tests/test_kimi_k2_7_muon_config.py index c5b3c78456..3b3092ff76 100644 --- a/tests/unit_tests/test_kimi_k2_7_muon_config.py +++ b/tests/unit_tests/test_kimi_k2_7_muon_config.py @@ -9,19 +9,35 @@ import torch from torch.distributed.tensor import Shard +from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( + 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 -from torchtitan.models.kimi_k2_7.config_registry import kimi_k2_5_muon +from torchtitan.distributed.activation_checkpoint import FullAC +from torchtitan.models.kimi_k2_7.config_registry import ( + kimi_k2_5_muon, + moonlight_16b_a3b_muon, +) -class TestKimiK25MuonConfig(unittest.TestCase): +class _KimiMuonConfigTests: + config_factory = None + num_layers = 0 + num_heads = 0 + num_owner_ranks = 0 + expert_parallel_degree = 0 + attention_projections: tuple[str, ...] = () + owned_attention_projections: frozenset[str] = frozenset() + @classmethod def setUpClass(cls): - cls.config = kimi_k2_5_muon() + assert cls.config_factory is not None + cls.config = cls.config_factory() assert cls.config.model_spec is not None with torch.device("meta"): cls.model = cls.config.model_spec.model.build() @@ -40,23 +56,32 @@ def test_parameter_routing(self): ) model_names = set(dict(self.model.named_parameters())) expected_muon_names = set() - for suffix, count in ( - (".attention.wq_a.weight", 61), - (".attention.wq_b.weight", 61), - (".attention.wkv_a.weight", 61), - (".attention.wkv_b.weight", 61), - (".attention.wo.weight", 61), + suffix_counts = ( + *( + (f".attention.{projection}.weight", self.num_layers) + for projection in self.attention_projections + ), (".feed_forward.w1.weight", 1), (".feed_forward.w2.weight", 1), (".feed_forward.w3.weight", 1), - (".moe.routed_experts.inner_experts.w1_EFD", 60), - (".moe.routed_experts.inner_experts.w2_EDF", 60), - (".moe.routed_experts.inner_experts.w3_EFD", 60), - (".moe.router.gate.weight", 60), - (".moe.shared_experts.w1.weight", 60), - (".moe.shared_experts.w2.weight", 60), - (".moe.shared_experts.w3.weight", 60), - ): + ( + ".moe.routed_experts.inner_experts.w1_EFD", + self.num_layers - 1, + ), + ( + ".moe.routed_experts.inner_experts.w2_EDF", + self.num_layers - 1, + ), + ( + ".moe.routed_experts.inner_experts.w3_EFD", + self.num_layers - 1, + ), + (".moe.router.gate.weight", self.num_layers - 1), + (".moe.shared_experts.w1.weight", self.num_layers - 1), + (".moe.shared_experts.w2.weight", self.num_layers - 1), + (".moe.shared_experts.w3.weight", self.num_layers - 1), + ) + for suffix, count in suffix_counts: names = {name for name in model_names if name.endswith(suffix)} self.assertEqual(len(names), count, suffix) expected_muon_names.update(names) @@ -75,44 +100,49 @@ def test_parameter_routing(self): {"match_rms_adamw"}, ) + representative_suffixes = ( + *( + f".attention.{projection}.weight" + for projection in self.attention_projections + ), + ".feed_forward.w1.weight", + ".moe.router.gate.weight", + ".moe.shared_experts.w1.weight", + ".moe.routed_experts.inner_experts.w1_EFD", + ) group_by_suffix = { suffix: next( group for group in muon_groups if group["param_names"][0].endswith(suffix) ) - for suffix in ( - ".attention.wq_a.weight", - ".attention.wq_b.weight", - ".attention.wkv_a.weight", - ".attention.wkv_b.weight", - ".attention.wo.weight", - ".feed_forward.w1.weight", - ".moe.router.gate.weight", - ".moe.shared_experts.w1.weight", - ".moe.routed_experts.inner_experts.w1_EFD", - ) + for suffix in representative_suffixes } per_head = MuonComputeSharding( view_before_placement=BatchedMatrixComputeView( - num_matrices=64, + num_matrices=self.num_heads, matrices_flattened_into_dim=0, ), placement=Shard(0), ) expected_sharding = { - ".attention.wq_a.weight": MuonComputeSharding(placement=Owned()), - ".attention.wq_b.weight": per_head, - ".attention.wkv_a.weight": MuonComputeSharding(placement=Owned()), - ".attention.wkv_b.weight": per_head, - ".attention.wo.weight": MuonComputeSharding(placement=Owned()), - ".feed_forward.w1.weight": MuonComputeSharding(placement=Owned()), - ".moe.router.gate.weight": MuonComputeSharding(placement=Owned()), - ".moe.shared_experts.w1.weight": MuonComputeSharding(placement=Owned()), - ".moe.routed_experts.inner_experts.w1_EFD": MuonComputeSharding( - placement=Shard(0) - ), + f".attention.{projection}.weight": ( + MuonComputeSharding(placement=Owned()) + if projection in self.owned_attention_projections + else per_head + ) + for projection in self.attention_projections } + expected_sharding.update( + { + ".feed_forward.w1.weight": MuonComputeSharding(placement=Owned()), + ".moe.router.gate.weight": MuonComputeSharding(placement=Owned()), + ".moe.shared_experts.w1.weight": MuonComputeSharding(placement=Owned()), + ".moe.routed_experts.inner_experts.w1_EFD": MuonComputeSharding( + placement=Shard(0) + ), + } + ) for suffix, group in group_by_suffix.items(): self.assertEqual(group["compute_sharding"], expected_sharding[suffix]) @@ -122,9 +152,17 @@ def test_bucket_and_parallelism_config(self): "bucket_configs" ] bucket_layer_ids = ((0,),) + tuple( - (first_layer_id, first_layer_id + 1) for first_layer_id in range(1, 61, 2) + tuple( + range( + first_layer_id, + min(first_layer_id + 2, self.num_layers), + ) + ) + for first_layer_id in range(1, self.num_layers, 2) ) - self.assertEqual(len(bucket_configs), 31) + self.assertEqual(len(bucket_configs), len(bucket_layer_ids)) + expected_bucket_patterns = [] + expected_owned_fqns = set() for layer_ids, bucket in zip(bucket_layer_ids, bucket_configs, strict=True): expected = () expected_owners = set() @@ -132,12 +170,12 @@ def test_bucket_and_parallelism_config(self): prefix = f"layers.{layer_id}" attention_fqns = tuple( f"{prefix}.attention.{projection}.weight" - for projection in ("wq_a", "wq_b", "wkv_a", "wkv_b", "wo") + for projection in self.attention_projections ) expected += attention_fqns expected_owners.update( f"{prefix}.attention.{projection}.weight" - for projection in ("wq_a", "wkv_a", "wo") + for projection in self.owned_attention_projections ) if not layer_id: dense_fqns = tuple( @@ -168,23 +206,70 @@ def test_bucket_and_parallelism_config(self): set(bucket.owner_rank_by_fqn), expected_owners, ) + expected_bucket_patterns.append(expected) + expected_owned_fqns.update(expected_owners) self.assertTrue( - all(rank in range(64) for rank in bucket.owner_rank_by_fqn.values()) + all( + rank in range(self.num_owner_ranks) + for rank in bucket.owner_rank_by_fqn.values() + ) ) + parameter_numel_by_fqn = { + fqn: parameter.numel() + for fqn, parameter in self.model.named_parameters() + if fqn in expected_owned_fqns + } + self.assertEqual(set(parameter_numel_by_fqn), expected_owned_fqns) + self.assertEqual( + tuple(dict(bucket.owner_rank_by_fqn) for bucket in bucket_configs), + assign_balanced_owners( + expected_bucket_patterns, + parameter_numel_by_fqn, + num_ranks=self.num_owner_ranks, + ), + ) + parallelism = self.config.parallelism self.assertEqual(parallelism.data_parallel_replicate_degree, 1) - self.assertEqual(parallelism.data_parallel_shard_degree, 64) - self.assertEqual(parallelism.expert_parallel_degree, 8) + self.assertEqual( + parallelism.data_parallel_shard_degree, + self.num_owner_ranks, + ) + self.assertEqual( + parallelism.expert_parallel_degree, + self.expert_parallel_degree, + ) 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.assertIsInstance(self.config.activation_checkpoint, FullAC.Config) def test_config_is_json_serializable(self): json.dumps(self.config.to_dict()) +class TestKimiK25MuonConfig(_KimiMuonConfigTests, unittest.TestCase): + config_factory = staticmethod(kimi_k2_5_muon) + num_layers = 61 + num_heads = 64 + num_owner_ranks = 64 + expert_parallel_degree = 8 + attention_projections = ("wq_a", "wq_b", "wkv_a", "wkv_b", "wo") + owned_attention_projections = frozenset(("wq_a", "wkv_a", "wo")) + + +class TestMoonlightMuonConfig(_KimiMuonConfigTests, unittest.TestCase): + config_factory = staticmethod(moonlight_16b_a3b_muon) + num_layers = 27 + num_heads = 16 + num_owner_ranks = 8 + expert_parallel_degree = 4 + attention_projections = ("wq", "wkv_a", "wkv_b", "wo") + owned_attention_projections = frozenset(("wkv_a", "wo")) + + if __name__ == "__main__": unittest.main() diff --git a/torchtitan/models/kimi_k2_7/config_registry.py b/torchtitan/models/kimi_k2_7/config_registry.py index 084752db93..863d7a3391 100644 --- a/torchtitan/models/kimi_k2_7/config_registry.py +++ b/torchtitan/models/kimi_k2_7/config_registry.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from collections.abc import Mapping + from torch.distributed.tensor import Shard from torchtitan.components.checkpoint import CheckpointManager @@ -55,12 +57,23 @@ def _mm_dataloader(dataset: str, **kwargs) -> MMDataLoader.Config: ) -def _kimi_k2_5_distributed_muon_optimizer( +def _per_head_muon_sharding(num_heads: int) -> MuonComputeSharding: + return MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView( + num_matrices=num_heads, + matrices_flattened_into_dim=0, + ), + placement=Shard(0), + ) + + +def _kimi_text_distributed_muon_optimizer( *, - n_layers: int, - num_heads: int, - owner_group_size: int, + num_layers: int, + num_owner_ranks: int, lr: float, + attention_shardings: Mapping[str, MuonComputeSharding], + owned_parameter_numel_by_suffix: Mapping[str, int], ) -> OptimizersContainer.Config: muon_kwargs = { "lr": lr, @@ -77,20 +90,6 @@ def _kimi_k2_5_distributed_muon_optimizer( "eps": 1e-8, "weight_decay": 0.1, } - per_head = MuonComputeSharding( - view_before_placement=BatchedMatrixComputeView( - num_matrices=num_heads, - matrices_flattened_into_dim=0, - ), - placement=Shard(0), - ) - attention_shardings = { - "wq_a": MuonComputeSharding(placement=Owned()), - "wq_b": per_head, - "wkv_a": MuonComputeSharding(placement=Owned()), - "wkv_b": per_head, - "wo": MuonComputeSharding(placement=Owned()), - } expert_projections = ("w1_EFD", "w2_EDF", "w3_EFD") param_groups = [ ParamGroupConfig( @@ -163,29 +162,17 @@ def layer_fqns(layer_id: int) -> tuple[str, ...]: ) return fqns - layer_bucket_fqns = tuple(layer_fqns(layer_id) for layer_id in range(n_layers)) + layer_bucket_fqns = tuple(layer_fqns(layer_id) for layer_id in range(num_layers)) # Layer 0 has a much larger dense MLP, so keep it separate while amortizing # collective launch overhead across pairs of MoE layers. bucket_layer_ids = ((0,),) + tuple( - tuple(range(first_layer_id, min(first_layer_id + 2, n_layers))) - for first_layer_id in range(1, n_layers, 2) + tuple(range(first_layer_id, min(first_layer_id + 2, num_layers))) + for first_layer_id in range(1, num_layers, 2) ) bucket_fqns = tuple( tuple(fqn for layer_id in layer_ids for fqn in layer_bucket_fqns[layer_id]) for layer_ids in bucket_layer_ids ) - owned_parameter_numel_by_suffix = { - "attention.wq_a.weight": 1536 * 7168, - "attention.wkv_a.weight": 576 * 7168, - "attention.wo.weight": 7168 * 8192, - "feed_forward.w1.weight": 18432 * 7168, - "feed_forward.w2.weight": 7168 * 18432, - "feed_forward.w3.weight": 18432 * 7168, - "moe.router.gate.weight": 384 * 7168, - "moe.shared_experts.w1.weight": 2048 * 7168, - "moe.shared_experts.w2.weight": 7168 * 2048, - "moe.shared_experts.w3.weight": 2048 * 7168, - } owner_rank_by_bucket = assign_balanced_owners( bucket_fqns, { @@ -194,7 +181,7 @@ def layer_fqns(layer_id: int) -> tuple[str, ...]: for suffix, numel in owned_parameter_numel_by_suffix.items() if f"layers.{layer_id}.{suffix}" in fqns }, - num_ranks=owner_group_size, + num_ranks=num_owner_ranks, ) bucket_configs = tuple( BucketConfig( @@ -221,6 +208,64 @@ def layer_fqns(layer_id: int) -> tuple[str, ...]: ) +def _moonlight_distributed_muon_optimizer( + *, num_owner_ranks: int +) -> OptimizersContainer.Config: + per_head = _per_head_muon_sharding(num_heads=16) + return _kimi_text_distributed_muon_optimizer( + num_layers=27, + num_owner_ranks=num_owner_ranks, + lr=3e-4, + attention_shardings={ + "wq": per_head, + "wkv_a": MuonComputeSharding(placement=Owned()), + "wkv_b": per_head, + "wo": MuonComputeSharding(placement=Owned()), + }, + owned_parameter_numel_by_suffix={ + "attention.wkv_a.weight": 576 * 2048, + "attention.wo.weight": 2048 * 2048, + "feed_forward.w1.weight": 11264 * 2048, + "feed_forward.w2.weight": 2048 * 11264, + "feed_forward.w3.weight": 11264 * 2048, + "moe.router.gate.weight": 64 * 2048, + "moe.shared_experts.w1.weight": 2816 * 2048, + "moe.shared_experts.w2.weight": 2048 * 2816, + "moe.shared_experts.w3.weight": 2816 * 2048, + }, + ) + + +def _kimi_k2_5_distributed_muon_optimizer( + *, num_owner_ranks: int +) -> OptimizersContainer.Config: + per_head = _per_head_muon_sharding(num_heads=64) + return _kimi_text_distributed_muon_optimizer( + num_layers=61, + num_owner_ranks=num_owner_ranks, + lr=2.2e-4, + attention_shardings={ + "wq_a": MuonComputeSharding(placement=Owned()), + "wq_b": per_head, + "wkv_a": MuonComputeSharding(placement=Owned()), + "wkv_b": per_head, + "wo": MuonComputeSharding(placement=Owned()), + }, + owned_parameter_numel_by_suffix={ + "attention.wq_a.weight": 1536 * 7168, + "attention.wkv_a.weight": 576 * 7168, + "attention.wo.weight": 7168 * 8192, + "feed_forward.w1.weight": 18432 * 7168, + "feed_forward.w2.weight": 7168 * 18432, + "feed_forward.w3.weight": 18432 * 7168, + "moe.router.gate.weight": 384 * 7168, + "moe.shared_experts.w1.weight": 2048 * 7168, + "moe.shared_experts.w2.weight": 7168 * 2048, + "moe.shared_experts.w3.weight": 2048 * 7168, + }, + ) + + def kimi_k2_5_debugmodel() -> Trainer.Config: model_spec = model_registry("debugmodel") return Trainer.Config( @@ -289,6 +334,26 @@ def moonlight_16b_a3b() -> Trainer.Config: ) +def moonlight_16b_a3b_muon() -> Trainer.Config: + """Moonlight 16B-A3B with DistributedMuon for matrix parameters.""" + config = moonlight_16b_a3b() + num_owner_ranks = 8 + config.optimizer = _moonlight_distributed_muon_optimizer( + num_owner_ranks=num_owner_ranks + ) + config.parallelism = ParallelismConfig( + data_parallel_replicate_degree=1, + data_parallel_shard_degree=num_owner_ranks, + 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 kimi_vl_a3b() -> Trainer.Config: """Kimi-VL A3B: Moonlight text tower + 2D MoonViT vision (image-text).""" model_spec = model_registry("Kimi-VL-A3B", attn_backend="flex") @@ -367,16 +432,13 @@ def kimi_k2_5() -> Trainer.Config: def kimi_k2_5_muon() -> Trainer.Config: """Full Kimi K2.5 with DistributedMuon for text-tower matrices.""" config = kimi_k2_5() - owner_group_size = 64 + num_owner_ranks = 64 config.optimizer = _kimi_k2_5_distributed_muon_optimizer( - n_layers=61, - num_heads=64, - owner_group_size=owner_group_size, - lr=2.2e-4, + num_owner_ranks=num_owner_ranks, ) config.parallelism = ParallelismConfig( data_parallel_replicate_degree=1, - data_parallel_shard_degree=owner_group_size, + data_parallel_shard_degree=num_owner_ranks, tensor_parallel_degree=1, context_parallel_degree=1, pipeline_parallel_degree=1, From 4ddffff51681a50ff022b49c5b79ffeb444bed95 Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Thu, 6 Aug 2026 11:50:52 -0700 Subject: [PATCH 19/20] Update [ghstack-poisoned] --- tests/unit_tests/test_distributed_muon.py | 233 ++++++++++-------- .../bucketed_redistribution.py | 23 +- .../components/distributed_optimizers/muon.py | 72 +++--- .../muon_parameter_prep.py | 4 +- 4 files changed, 193 insertions(+), 139 deletions(-) diff --git a/tests/unit_tests/test_distributed_muon.py b/tests/unit_tests/test_distributed_muon.py index 62ad2062d2..1a3c86b978 100644 --- a/tests/unit_tests/test_distributed_muon.py +++ b/tests/unit_tests/test_distributed_muon.py @@ -113,16 +113,13 @@ def _optimizer( *, local_num_matrices: int = 2, owner_rank: int = 1, - redistributed_compute_placement: Owned | Replicate = Owned(), ) -> DistributedMuon: return build_distributed_muon( [ { "params": [redistributed], "param_names": ["layers.0.redistributed"], - "compute_sharding": MuonComputeSharding( - placement=redistributed_compute_placement - ), + "compute_sharding": MuonComputeSharding(placement=Owned()), }, { "params": [local_blocks], @@ -139,11 +136,7 @@ def _optimizer( bucket_configs=[ BucketConfig( patterns=("layers.0.*",), - owner_rank_by_fqn=( - {"layers.0.redistributed": owner_rank} - if isinstance(redistributed_compute_placement, Owned) - else {} - ), + owner_rank_by_fqn={"layers.0.redistributed": owner_rank}, mesh_axes=("dp_shard",), name="layers.0", ) @@ -213,6 +206,50 @@ def _assert_matches_reference( ) +@unittest.skipUnless(torch.cuda.device_count() >= 1, "requires one CUDA device") +class TestDistributedMuonSingleRank(_DistributedMuonTestBase): + @property + def world_size(self): + return 1 + + @with_comms + def test_owned_compute_accepts_static_owner_for_replicated_storage(self): + value = torch.arange(12, device=self.device).reshape(4, 3).float() + parameter = torch.nn.Parameter( + distribute_tensor(value.clone(), self.mesh, (Replicate(),)) + ) + 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": 0}, + mesh=self.mesh, + ) + ], + ns_steps=1, + ) + parameter.grad = distribute_tensor( + torch.ones_like(value), self.mesh, (Replicate(),) + ) + + 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() + + collective.assert_not_called() + + @unittest.skipUnless(torch.cuda.device_count() >= 2, "requires two CUDA devices") class TestDistributedMuon(_DistributedMuonTestBase): @with_comms @@ -420,28 +457,35 @@ def test_flat_state_dict_loads_after_group_membership_changes(self): .float(), } - def build(names, compute_view=None): + def build(names, *, compute_locally=False): parameters = [ torch.nn.Parameter( - distribute_tensor(values[name].clone(), self.mesh, (Replicate(),)) + distribute_tensor(values[name].clone(), self.mesh, (Shard(0),)) ) for name in names ] + compute_sharding = ( + MuonComputeSharding( + view_before_placement=BatchedMatrixComputeView(num_matrices=2), + placement=Shard(0), + ) + if compute_locally + else MuonComputeSharding(placement=Owned()) + ) optimizer = build_distributed_muon( [ { "params": parameters, "param_names": names, - "compute_sharding": MuonComputeSharding( - view_before_placement=compute_view, - placement=Replicate(), - ), + "compute_sharding": compute_sharding, } ], bucket_spec=[ BucketSpec( patterns=("layers.0.*",), - owner_rank_by_fqn={}, + owner_rank_by_fqn=( + {} if compute_locally else dict.fromkeys(names, 0) + ), mesh=self.mesh, ) ], @@ -450,11 +494,13 @@ def build(names, compute_view=None): return parameters, optimizer source_parameters, source_optimizer = build(("layers.0.a", "layers.0.b")) - for parameter in source_parameters: + for name, parameter in zip( + ("layers.0.a", "layers.0.b"), source_parameters, strict=True + ): parameter.grad = distribute_tensor( - torch.ones_like(parameter.to_local()), + torch.ones_like(values[name]), self.mesh, - (Replicate(),), + (Shard(0),), ) source_optimizer.step() flat_state_dict = get_flat_optim_state_dict(source_optimizer) @@ -471,13 +517,7 @@ def build(names, compute_view=None): source_optimizer.state[source_parameters[0]]["momentum_buffer"].to_local(), ) - _, changed_layout_optimizer = build( - ("layers.0.a",), - BatchedMatrixComputeView( - num_matrices=2, - matrices_flattened_into_dim=0, - ), - ) + _, changed_layout_optimizer = build(("layers.0.a",), compute_locally=True) init_optim_state(changed_layout_optimizer) with self.assertRaisesRegex(ValueError, "compute layout"): load_flat_optim_state_dict( @@ -711,18 +751,35 @@ def test_shard1_owned_matches_plain_muon(self): ) @with_comms - def test_replicated_storage_and_compute_match_plain_muon(self): - value = torch.arange(1, 13, device=self.device).reshape(4, 3).float().div_(10) - parameter = torch.nn.Parameter( - distribute_tensor(value.clone(), self.mesh, (Replicate(),)) + 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": [parameter], - "param_names": ["layers.0.matrix"], - "compute_sharding": MuonComputeSharding(placement=Replicate()), - } + "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), + placement=Shard(0), + ), + }, ], bucket_spec=[ BucketSpec( @@ -738,18 +795,32 @@ def test_replicated_storage_and_compute_match_plain_muon(self): ns_steps=2, ) - grad = value.flip(0).contiguous() - parameter.grad = distribute_tensor(grad, self.mesh, (Replicate(),)) - reference = torch.nn.Parameter(value.clone()) - reference.grad = grad.clone() + 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(),)) + + owned_reference = torch.nn.Parameter(values[0].clone()) + batched_references = [ + torch.nn.Parameter(matrix.clone()) + for matrix in values[1].view(2, 2, 3).unbind() + ] + references = [owned_reference, *batched_references] reference_optimizer = torch.optim.Muon( - [reference], + references, lr=0.03, weight_decay=0.2, momentum=0.8, nesterov=True, ns_steps=2, ) + owned_reference.grad = grads[0] + for reference, grad in zip( + batched_references, + grads[1].view(2, 2, 3).unbind(), + strict=True, + ): + reference.grad = grad + all_to_all_single = dist.all_to_all_single with patch( "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." @@ -760,52 +831,34 @@ def test_replicated_storage_and_compute_match_plain_muon(self): reference_optimizer.step() collective.assert_not_called() - self.assertEqual(parameter.placements, (Replicate(),)) - self.assertEqual(parameter.grad.placements, (Replicate(),)) - momentum = optimizer.state[parameter]["momentum_buffer"] - self.assertEqual(momentum.placements, (Replicate(),)) - torch.testing.assert_close(parameter.to_local(), reference) - torch.testing.assert_close( - momentum.to_local(), - reference_optimizer.state[reference]["momentum_buffer"], + reference_values = ( + owned_reference, + torch.stack(batched_references).view(batched.shape), ) - - @with_comms - def test_constructor_rejects_mismatched_replicated_compute_placement(self): - value = torch.arange(12, device=self.device).reshape(4, 3).float() - replicated = torch.nn.Parameter( - distribute_tensor(value.clone(), self.mesh, (Replicate(),)) + reference_momenta = ( + reference_optimizer.state[owned_reference]["momentum_buffer"], + torch.stack( + [ + reference_optimizer.state[reference]["momentum_buffer"] + for reference in batched_references + ] + ).view(batched.shape), ) - sharded = self._parameter(value) - - for storage, placement, owners in ( - ( - replicated, - Owned(), - {"layers.0.weight": 0}, - ), - (sharded, Replicate(), {}), + for param, reference, reference_momentum in zip( + (owned, batched), + reference_values, + reference_momenta, + strict=True, ): - with self.subTest(storage=storage.placements, placement=placement): - with self.assertRaisesRegex(ValueError, "storage-to-compute layout"): - build_distributed_muon( - [ - { - "params": [storage], - "param_names": ["layers.0.weight"], - "compute_sharding": MuonComputeSharding( - placement=placement, - ), - } - ], - bucket_spec=[ - BucketSpec( - patterns=("layers.0.*",), - owner_rank_by_fqn=owners, - mesh=self.mesh, - ) - ], - ) + 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) + torch.testing.assert_close( + momentum.to_local(), + reference_momentum, + ) @with_comms def test_step_rejects_gradient_with_reordered_mesh(self): @@ -949,22 +1002,6 @@ def test_step_matches_plain_muon_and_continues_from_state_dict(self): with self.assertRaisesRegex(ValueError, "compute layout"): changed_view_optimizer.load_state_dict(state_dict) - changed_placement_optimizer = self._optimizer( - torch.nn.Parameter( - distribute_tensor( - reference_redistributed.detach().clone(), - self.mesh, - (Replicate(),), - ) - ), - self._parameter( - torch.cat([parameter.detach() for parameter in reference_local_blocks]) - ), - redistributed_compute_placement=Replicate(), - ) - with self.assertRaisesRegex(ValueError, "compute layout"): - changed_placement_optimizer.load_state_dict(state_dict) - resumed_optimizer = self._optimizer( resumed_redistributed, resumed_local_blocks, diff --git a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py index 883b56b7f6..3d7b6ef59c 100644 --- a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py +++ b/torchtitan/components/distributed_optimizers/bucketed_redistribution.py @@ -71,7 +71,9 @@ class BucketSpec: ``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. + owner entry. A rank-0 entry is also accepted for compute-ready parameters + on a one-rank mesh, where sharded storage may normalize to replication. + ``name`` is diagnostic metadata only. """ patterns: tuple[str, ...] @@ -907,8 +909,8 @@ def _build_owned_bucket_plans( ) -> _BucketPlanningResult[_ItemT]: """Build the local and whole-matrix-owned DistributedMuon plans. - The active planner supports Replicate -> Replicate and Shard(0) matrix - batches as local compute, plus Shard(...) -> Owned through packed + The active planner supports replicated storage and dimension-0-sharded + matrix batches as local compute, plus Shard(...) -> Owned through packed all-to-all and Owned -> Shard(...) through reverse packed all-to-all. Other placement transitions are intentionally unsupported. """ @@ -933,12 +935,21 @@ def _build_owned_bucket_plans( ) expected_owners = {fqn(item) for item in redistributed_items} provided_owners = set(spec.owner_rank_by_fqn) - if provided_owners != expected_owners: + # Size-one sharded storage may normalize to Replicate. In that case a + # static rank-0 owner entry is equivalent to the resolved local compute. + redundant_owners = { + fqn(item) + for item in local_items + if len(group.participants) == 1 + and spec.owner_rank_by_fqn.get(fqn(item)) == 0 + } + effective_provided_owners = provided_owners - redundant_owners + if effective_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)}" + f"missing={sorted(expected_owners - effective_provided_owners)}, " + f"extra={sorted(effective_provided_owners - expected_owners)}" ) ordered_items.extend(local_items) ordered_items.extend(redistributed_items) diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index 8a5627dbec..8a7af42feb 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -212,7 +212,7 @@ def _build_parameter_compute_layouts( prepared = self._prepared_compute_views[fqn] global_compute_shape = torch.Size(prepared.global_compute_shape) local_compute_tensor = prepared.local_compute_tensor - compute_locally = _validate_muon_parameter( + resolved_compute_placement = _resolve_compute_placement( fqn, param, global_compute_shape, @@ -228,8 +228,8 @@ def _build_parameter_compute_layouts( global_compute_shape=global_compute_shape, local_compute_tensor=local_compute_tensor, local_storage_signature=_local_storage_signature(param.to_local()), - compute_placement=compute_placement, - compute_locally=compute_locally, + compute_placement_key=resolved_compute_placement.fingerprint_key, + compute_locally=resolved_compute_placement.compute_locally, ) ) return tuple(compute_layouts) @@ -254,10 +254,7 @@ def _set_checkpoint_layout_fingerprints(self) -> None: tuple(layout.param.shape), layout.compute_view_key, tuple(layout.global_compute_shape), - _compute_placement_key( - layout.compute_placement, - len(layout.global_compute_shape), - ), + layout.compute_placement_key, ) self._layout_fingerprints_by_fqn[layout.fqn] = ( _LAYOUT_FINGERPRINT_VERSION, @@ -285,10 +282,7 @@ def _plan_item_signature( tuple(compute_layout.global_compute_shape), compute_layout.compute_locally, compute_layout.compute_view_key, - _compute_placement_key( - compute_layout.compute_placement, - len(compute_layout.global_compute_shape), - ), + compute_layout.compute_placement_key, _device_mesh_ranks(compute_layout.param.device_mesh), tuple(map(str, compute_layout.param.placements)), self._group_signature(compute_layout), @@ -507,7 +501,13 @@ class _ParameterComputeLayout: global_compute_shape: torch.Size local_compute_tensor: Tensor local_storage_signature: tuple[Any, ...] - compute_placement: Owned | Replicate | Shard + compute_placement_key: tuple[Any, ...] + compute_locally: bool + + +@dataclass(frozen=True, slots=True) +class _ResolvedComputePlacement: + fingerprint_key: tuple[Any, ...] compute_locally: bool @@ -549,35 +549,52 @@ def _has_owned_sharded_storage(param: DTensor) -> bool: ) -def _validate_muon_parameter( +def _resolve_compute_placement( fqn: str, param: DTensor, global_compute_shape: torch.Size, local_compute_tensor: Tensor, compute_placement: object, -) -> bool: +) -> _ResolvedComputePlacement: local = param.to_local() if torch.is_complex(param) or param.ndim < 2 or not local.is_contiguous(): raise ValueError(f"Muon parameter {fqn!r} has unsupported shape or storage") replicated_storage = _has_replicated_storage(param) - if isinstance(compute_placement, Replicate) and replicated_storage: - return True - elif isinstance(compute_placement, Shard): + if isinstance(compute_placement, Shard): if ( len(global_compute_shape) >= 3 and _normalize_dim(compute_placement.dim, len(global_compute_shape)) == 0 - and local_compute_tensor.shape[1:] == global_compute_shape[1:] - and _has_dim0_sharded_storage(param) + and ( + ( + replicated_storage + and local_compute_tensor.shape == global_compute_shape + ) + or ( + local_compute_tensor.shape[1:] == global_compute_shape[1:] + and _has_dim0_sharded_storage(param) + ) + ) ): - return True + return _ResolvedComputePlacement( + fingerprint_key=("shard", 0), + compute_locally=True, + ) elif ( isinstance(compute_placement, Owned) and len(global_compute_shape) == 2 and param.ndim == 2 - and _has_owned_sharded_storage(param) ): - return False + if replicated_storage: + return _ResolvedComputePlacement( + fingerprint_key=("owned",), + compute_locally=True, + ) + if _has_owned_sharded_storage(param): + return _ResolvedComputePlacement( + fingerprint_key=("owned",), + compute_locally=False, + ) raise ValueError(f"unsupported storage-to-compute layout for {fqn!r}") @@ -588,17 +605,6 @@ def _normalize_dim(dim: int, ndim: int) -> int: return normalized -def _compute_placement_key( - placement: Owned | Replicate | Shard, - ndim: int, -) -> tuple[Any, ...]: - if isinstance(placement, Owned): - return ("owned",) - if isinstance(placement, Replicate): - return ("replicate",) - return ("shard", _normalize_dim(placement.dim, ndim)) - - # 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( diff --git a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py index ffdd6c68a5..bca6a43f03 100644 --- a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py +++ b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py @@ -70,10 +70,10 @@ class MuonComputeSharding: # 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 | Replicate | Shard + placement: Owned | Shard def __post_init__(self) -> None: - if not isinstance(self.placement, (Owned, Replicate, Shard)) or ( + if not isinstance(self.placement, (Owned, Shard)) or ( self.view_before_placement is not None and not isinstance(self.view_before_placement, BatchedMatrixComputeView) ): From f0e3863b312729caf03d3bf9d0ffa7c6bc9b1007 Mon Sep 17 00:00:00 2001 From: Wei Feng Date: Thu, 6 Aug 2026 13:50:15 -0700 Subject: [PATCH 20/20] Update [ghstack-poisoned] --- tests/unit_tests/test_distributed_muon.py | 20 +- ...tion.py => test_flex_optimizer_reshard.py} | 18 +- .../unit_tests/test_kimi_k2_7_muon_config.py | 2 +- ...tribution.py => flex_optimizer_reshard.py} | 91 +++++---- .../components/distributed_optimizers/muon.py | 186 ++++++++++-------- .../muon_parameter_prep.py | 147 +++++++------- .../models/kimi_k2_7/config_registry.py | 2 +- 7 files changed, 255 insertions(+), 211 deletions(-) rename tests/unit_tests/{test_bucketed_optimizer_redistribution.py => test_flex_optimizer_reshard.py} (94%) rename torchtitan/components/distributed_optimizers/{bucketed_redistribution.py => flex_optimizer_reshard.py} (97%) diff --git a/tests/unit_tests/test_distributed_muon.py b/tests/unit_tests/test_distributed_muon.py index 1a3c86b978..66dd2bb30f 100644 --- a/tests/unit_tests/test_distributed_muon.py +++ b/tests/unit_tests/test_distributed_muon.py @@ -21,7 +21,7 @@ init_optim_state, load_flat_optim_state_dict, ) -from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( +from torchtitan.components.distributed_optimizers.flex_optimizer_reshard import ( BucketConfig, BucketSpec, ) @@ -241,7 +241,7 @@ def test_owned_compute_accepts_static_owner_for_replicated_storage(self): all_to_all_single = dist.all_to_all_single with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard.dist." "all_to_all_single", wraps=all_to_all_single, ) as collective: @@ -298,7 +298,7 @@ def build(parameter, name, compute_placement, owner_rank=None): local_blocks.placements, ) with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard.dist." "all_to_all_single" ) as collective: optimizer.step() @@ -727,7 +727,7 @@ def test_shard1_owned_matches_plain_muon(self): all_to_all_single = dist.all_to_all_single with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard." "dist.all_to_all_single", wraps=all_to_all_single, ) as collective: @@ -823,7 +823,7 @@ def test_replicated_storage_matches_plain_muon_without_redistribution(self): all_to_all_single = dist.all_to_all_single with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard.dist." "all_to_all_single", wraps=all_to_all_single, ) as collective: @@ -900,7 +900,7 @@ def test_step_rejects_gradient_with_reordered_mesh(self): parameter_before = parameter.to_local().clone() with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard." "dist.all_to_all_single" ) as collective: with self.assertRaisesRegex( @@ -960,7 +960,7 @@ def test_step_matches_plain_muon_and_continues_from_state_dict(self): all_to_all_single = dist.all_to_all_single with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard.dist." "all_to_all_single", wraps=all_to_all_single, ) as collective: @@ -1026,7 +1026,7 @@ def test_step_matches_plain_muon_and_continues_from_state_dict(self): parameter.grad = grad.clone() with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard.dist." "all_to_all_single", wraps=all_to_all_single, ) as collective: @@ -1252,7 +1252,7 @@ def test_distinct_bucket_meshes_use_mesh_local_owners(self): all_to_all_single = dist.all_to_all_single with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard.dist." "all_to_all_single", wraps=all_to_all_single, ) as collective: @@ -1346,7 +1346,7 @@ def test_local_only_bucket_does_not_reuse_inflight_slot(self): all_to_all_single = dist.all_to_all_single with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard.dist." "all_to_all_single", wraps=all_to_all_single, ) as collective: diff --git a/tests/unit_tests/test_bucketed_optimizer_redistribution.py b/tests/unit_tests/test_flex_optimizer_reshard.py similarity index 94% rename from tests/unit_tests/test_bucketed_optimizer_redistribution.py rename to tests/unit_tests/test_flex_optimizer_reshard.py index 87363595b5..4734c25232 100644 --- a/tests/unit_tests/test_bucketed_optimizer_redistribution.py +++ b/tests/unit_tests/test_flex_optimizer_reshard.py @@ -12,7 +12,7 @@ import torch from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor -from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( +from torchtitan.components.distributed_optimizers.flex_optimizer_reshard import ( _bind_bucket_configs, _BucketedRedistributionRuntime, _build_owned_bucket_plans, @@ -28,7 +28,7 @@ ) -class TestBucketedOptimizerRedistribution(unittest.TestCase): +class TestFlexOptimizerReshard(unittest.TestCase): def test_runtime_enqueues_return_before_next_compute(self): runtime = _BucketedRedistributionRuntime(torch.device("cuda")) caller_stream = Mock() @@ -69,14 +69,14 @@ def release(work, caller): with patch( "torchtitan.components.distributed_optimizers." - "bucketed_redistribution._prepare_redistributed" + "flex_optimizer_reshard._prepare_redistributed" ), patch( "torchtitan.components.distributed_optimizers." - "bucketed_redistribution._compute_redistributed", + "flex_optimizer_reshard._compute_redistributed", side_effect=compute_redistributed, ), patch( "torchtitan.components.distributed_optimizers." - "bucketed_redistribution._finalize_redistributed" + "flex_optimizer_reshard._finalize_redistributed" ), patch.object( runtime, "_release", @@ -193,15 +193,15 @@ class Item: mesh.ndim = 1 with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard.dist." "get_process_group_ranks", return_value=[3, 7], ), patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard." "_redistribution_group", return_value=group, ), patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard." "_dtensor_storage_blocks", return_value=blocks, ): @@ -247,7 +247,7 @@ def test_transport_neutral_routes_lower_to_packed_all_to_all(self): ) with patch( - "torchtitan.components.distributed_optimizers.bucketed_redistribution.dist." + "torchtitan.components.distributed_optimizers.flex_optimizer_reshard.dist." "get_process_group_ranks", return_value=[3, 7], ): diff --git a/tests/unit_tests/test_kimi_k2_7_muon_config.py b/tests/unit_tests/test_kimi_k2_7_muon_config.py index 3b3092ff76..37a1bbf0d4 100644 --- a/tests/unit_tests/test_kimi_k2_7_muon_config.py +++ b/tests/unit_tests/test_kimi_k2_7_muon_config.py @@ -9,7 +9,7 @@ import torch from torch.distributed.tensor import Shard -from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( +from torchtitan.components.distributed_optimizers.flex_optimizer_reshard import ( assign_balanced_owners, ) from torchtitan.components.distributed_optimizers.muon import Owned diff --git a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py b/torchtitan/components/distributed_optimizers/flex_optimizer_reshard.py similarity index 97% rename from torchtitan/components/distributed_optimizers/bucketed_redistribution.py rename to torchtitan/components/distributed_optimizers/flex_optimizer_reshard.py index 3d7b6ef59c..1bdfc5242f 100644 --- a/torchtitan/components/distributed_optimizers/bucketed_redistribution.py +++ b/torchtitan/components/distributed_optimizers/flex_optimizer_reshard.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Private bucketed storage-to-compute runtime for DistributedMuon.""" +"""Public bucket contracts and private resharding for distributed optimizers.""" from __future__ import annotations @@ -88,10 +88,44 @@ def __post_init__(self) -> None: 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 FQNs across group-local rank indices. + + Only FQNs present in ``memory_estimate_by_fqn`` receive owners. One running + load vector balances cumulatively across buckets; FQN and rank ordering + make equal-load assignments deterministic. + """ + 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 _bind_bucket_configs( configs: Sequence[BucketConfig], storage_by_fqn: Mapping[str, DTensor], ) -> tuple[BucketSpec, ...]: + """Bind static configs to storage meshes after model parallelization.""" specs = [] for config in configs: candidates = tuple(config.owner_rank_by_fqn) or tuple( @@ -116,34 +150,6 @@ def _bind_bucket_configs( 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], @@ -364,6 +370,13 @@ def execute(self, output: Tensor, input: Tensor) -> None: _CommunicationSchedule = _PackedAllToAllSchedule | _LocalSchedule +@dataclass(frozen=True, slots=True) +class _RedistributionGroup: + process_group: dist.ProcessGroup + participants: tuple[int, ...] + local_participant: int + + @dataclass(slots=True) class _BucketPlan(Generic[_ItemT]): local_items: tuple[_ItemT, ...] @@ -376,13 +389,6 @@ class _BucketPlan(Generic[_ItemT]): 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], ...] @@ -485,8 +491,11 @@ def create(cls, device: torch.device) -> _CommunicationContext: 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()``. + ``prepare`` writes a Muon input into runtime-owned scratch, ``compute`` + updates its runtime-owned input in place, and ``finalize`` consumes a + runtime-owned result before reuse. Callbacks run under the stream selected + by the runtime and must not retain tensors, synchronize, or call + ``Tensor.record_stream()``. """ def __init__(self, device: torch.device) -> None: @@ -778,6 +787,7 @@ def _lower_packed_all_to_all( process_group: dist.ProcessGroup, local_participant: int, ) -> _PackedAllToAllSchedule: + """Lower nonempty plans with one shared participant order to packed A2A.""" participants = redistribution_plans[0].participants if storage_to_compute: routes_by_parameter = tuple( @@ -1045,6 +1055,11 @@ def _validate_bucket_plans_across_ranks( *, item_signature: Callable[[_ItemT], tuple[Any, ...]], ) -> None: + """Collectively verify rank-stable plans before runtime communication. + + Every rank must provide the same plan count and process-group order so all + workers enter these validation collectives in the same sequence. + """ for plan in plans: description = ( str(plan.dtype), diff --git a/torchtitan/components/distributed_optimizers/muon.py b/torchtitan/components/distributed_optimizers/muon.py index 8a7af42feb..4c5f451666 100644 --- a/torchtitan/components/distributed_optimizers/muon.py +++ b/torchtitan/components/distributed_optimizers/muon.py @@ -4,7 +4,7 @@ # 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.""" +"""Muon compute placement and the internal DistributedMuon runtime.""" from __future__ import annotations @@ -20,7 +20,7 @@ from torch.distributed.tensor.placement_types import _StridedShard from torch.optim import Optimizer -from .bucketed_redistribution import ( +from .flex_optimizer_reshard import ( _BucketedRedistributionRuntime, _build_owned_bucket_plans, _device_mesh_ranks, @@ -33,24 +33,23 @@ __all__ = ["BucketSpec", "assign_balanced_owners", "Owned"] -_LAYOUT_FINGERPRINT_KEY = "_distributed_muon_layout_fingerprint" -_LAYOUT_FINGERPRINT_VERSION = 1 - - @dataclass(frozen=True, slots=True) class Owned: - """Require a complete matrix; sharded storage uses a ``BucketSpec`` owner.""" + """Require complete 2D matrix compute. - -@dataclass(frozen=True, slots=True) -class _PreparedParameterComputeView: - compute_view_key: tuple[Any, ...] - global_compute_shape: torch.Size - local_compute_tensor: Tensor + This is a Muon compute placement, not a DTensor storage placement. + Replicated storage computes locally; sharded storage uses the parameter's + mesh-local owner from ``BucketSpec.owner_rank_by_fqn``. + """ class DistributedMuon(Optimizer): - """Internal runtime constructed through ``build_distributed_muon``.""" + """Internal runtime constructed through ``build_distributed_muon``. + + Parameter groups, FQNs, storage layouts, compute layouts, and bucket plans + are frozen after construction. Every configured parameter must have a + layout-compatible DTensor gradient before each rank enters ``step()``. + """ def __init__( self, @@ -492,6 +491,17 @@ def _compute_shape( return compute_layout.global_compute_shape +_LAYOUT_FINGERPRINT_KEY = "_distributed_muon_layout_fingerprint" +_LAYOUT_FINGERPRINT_VERSION = 1 + + +@dataclass(frozen=True, slots=True) +class _PreparedParameterComputeView: + compute_view_key: tuple[Any, ...] + global_compute_shape: torch.Size + local_compute_tensor: Tensor + + @dataclass(frozen=True, slots=True) class _ParameterComputeLayout: fqn: str @@ -511,44 +521,6 @@ class _ResolvedComputePlacement: compute_locally: bool -def _local_storage_signature(tensor: Tensor) -> tuple[Any, ...]: - return ( - tensor.data_ptr(), - tensor.storage_offset(), - tuple(tensor.shape), - tuple(tensor.stride()), - tensor.dtype, - tensor.device, - ) - - -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): - shard = cast(Shard | _StridedShard, placement) - if shard.dim % param.ndim != 0: - return False - has_shard = True - elif type(placement) is not Replicate: - return False - return has_shard - - -def _has_owned_sharded_storage(param: DTensor) -> bool: - return ( - param.device_mesh.ndim == 1 - and len(param.placements) == 1 - and type(param.placements[0]) is Shard - ) - - def _resolve_compute_placement( fqn: str, param: DTensor, @@ -556,6 +528,14 @@ def _resolve_compute_placement( local_compute_tensor: Tensor, compute_placement: object, ) -> _ResolvedComputePlacement: + """Validate and canonicalize one storage-to-compute transition. + + ``Owned`` accepts a 2D matrix stored as exact ``Shard`` on a 1D mesh and + redistributes it to its configured owner. ``Shard(0)`` accepts matrix + batches whose storage shards keep each matrix whole. Fully replicated + storage computes locally under either compatible declaration, including + when a size-one mesh axis has normalized sharded storage to replication. + """ local = param.to_local() if torch.is_complex(param) or param.ndim < 2 or not local.is_contiguous(): raise ValueError(f"Muon parameter {fqn!r} has unsupported shape or storage") @@ -598,6 +578,33 @@ def _resolve_compute_placement( raise ValueError(f"unsupported storage-to-compute layout for {fqn!r}") +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): + shard = cast(Shard | _StridedShard, placement) + if shard.dim % param.ndim != 0: + return False + has_shard = True + elif type(placement) is not Replicate: + return False + return has_shard + + +def _has_owned_sharded_storage(param: DTensor) -> bool: + return ( + param.device_mesh.ndim == 1 + and len(param.placements) == 1 + and type(param.placements[0]) is Shard + ) + + def _normalize_dim(dim: int, ndim: int) -> int: normalized = dim if dim >= 0 else dim + ndim if normalized < 0 or normalized >= ndim: @@ -605,40 +612,19 @@ def _normalize_dim(dim: int, ndim: int) -> int: return normalized -# 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 _local_storage_signature(tensor: Tensor) -> tuple[Any, ...]: + return ( + tensor.data_ptr(), + tensor.storage_offset(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device, + ) +# Keep the functional math aligned with torch.optim.Muon while owning the +# implementation here so the distributed runtime has no Muon dependency. def _adjust_learning_rate( lr: float, adjust_lr_fn: str | None, @@ -672,3 +658,35 @@ def _compute_muon_update( ) out.copy_(direction) return out + + +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 diff --git a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py index bca6a43f03..827966157c 100644 --- a/torchtitan/components/distributed_optimizers/muon_parameter_prep.py +++ b/torchtitan/components/distributed_optimizers/muon_parameter_prep.py @@ -4,7 +4,7 @@ # 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.""" +"""Public configuration and construction for DistributedMuon.""" from __future__ import annotations @@ -18,20 +18,20 @@ from torch.distributed.tensor import DTensor, Replicate, Shard from torch.distributed.tensor._utils import _compute_local_shape_and_global_offset -from .bucketed_redistribution import _bind_bucket_configs, BucketConfig, BucketSpec +from .flex_optimizer_reshard import _bind_bucket_configs, BucketConfig, BucketSpec from .muon import _PreparedParameterComputeView, DistributedMuon, Owned __all__ = [ "BatchedMatrixComputeView", - "build_distributed_muon", "MuonComputeSharding", + "build_distributed_muon", ] @dataclass(frozen=True, slots=True) class BatchedMatrixComputeView: - """Unflatten a storage dimension into a batch of matrices.""" + """View 2D storage as matrices with batch and rows flattened into dim 0.""" num_matrices: int matrices_flattened_into_dim: int = 0 @@ -64,7 +64,12 @@ def _resolve(self, storage_shape: torch.Size) -> _ResolvedBatchedMatrixView: @dataclass(frozen=True, slots=True, kw_only=True) class MuonComputeSharding: - """Define the logical Muon compute tensor and its required placement.""" + """Define the logical Muon tensor and its compute placement. + + ``Owned`` requires one rank to compute a complete 2D matrix when storage + is sharded. ``Shard(0)`` computes rank-3-or-higher matrix batches locally + when storage shards preserve complete matrices. + """ # Applied before compute placement, so placement dimensions refer to the # viewed tensor. A future view_after_placement mode can apply a local view @@ -85,68 +90,6 @@ def to_dict(self) -> dict: return {"repr": repr(self)} -@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 - or 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 _validate_batched_matrix_storage_shards( - fqn: str, - param: DTensor, - resolved_view: _ResolvedBatchedMatrixView, -) -> None: - """Validate every storage shard from globally identical DTensor metadata.""" - for placement in param.placements: - if type(placement) is Replicate: - continue - assert type(placement) is Shard - if placement.dim % param.ndim != 0: - raise ValueError( - f"batched-matrix Muon parameter {fqn!r} requires storage " - "shards along tensor dimension 0" - ) - - matrix_rows = resolved_view.matrix_rows - # Every rank must validate all coordinates before DistributedMuon performs - # collectives; checking only the local shard could strand its peers. - coordinates = product( - *(range(mesh_axis_size) for mesh_axis_size in param.device_mesh.shape) - ) - for coordinate in coordinates: - local_shape, global_offset = _compute_local_shape_and_global_offset( - param.shape, - param.device_mesh.shape, - list(coordinate), - param.placements, - ) - if local_shape[0] and ( - local_shape[0] % matrix_rows or global_offset[0] % matrix_rows - ): - raise ValueError( - f"batched-matrix Muon parameter {fqn!r} storage shards are not " - f"aligned to matrix rows of size {matrix_rows}" - ) - - def build_distributed_muon( params: Iterable[dict[str, Any]], *, @@ -154,7 +97,13 @@ def build_distributed_muon( bucket_configs: Sequence[BucketConfig] | None = None, **kwargs: Any, ) -> DistributedMuon: - """Prepare parameter views and construct the DistributedMuon runtime.""" + """Prepare named DTensor parameter groups and construct DistributedMuon. + + Every group must provide aligned ``params`` and ``param_names`` plus one + ``compute_sharding`` contract. Exactly one of ``bucket_spec`` or + ``bucket_configs`` is required. Parameter groups and layouts are frozen + after construction because optimizer state and collectives depend on them. + """ if (bucket_spec is None) == (bucket_configs is None): raise ValueError("provide exactly one of bucket_spec or bucket_configs") @@ -251,3 +200,65 @@ def build_distributed_muon( _prepared_compute_views=prepared_compute_views, **kwargs, ) + + +@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 + or 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 _validate_batched_matrix_storage_shards( + fqn: str, + param: DTensor, + resolved_view: _ResolvedBatchedMatrixView, +) -> None: + """Validate every storage shard from globally identical DTensor metadata.""" + for placement in param.placements: + if type(placement) is Replicate: + continue + assert type(placement) is Shard + if placement.dim % param.ndim != 0: + raise ValueError( + f"batched-matrix Muon parameter {fqn!r} requires storage " + "shards along tensor dimension 0" + ) + + matrix_rows = resolved_view.matrix_rows + # Every rank must validate all coordinates before DistributedMuon performs + # collectives; checking only the local shard could strand its peers. + coordinates = product( + *(range(mesh_axis_size) for mesh_axis_size in param.device_mesh.shape) + ) + for coordinate in coordinates: + local_shape, global_offset = _compute_local_shape_and_global_offset( + param.shape, + param.device_mesh.shape, + list(coordinate), + param.placements, + ) + if local_shape[0] and ( + local_shape[0] % matrix_rows or global_offset[0] % matrix_rows + ): + raise ValueError( + f"batched-matrix Muon parameter {fqn!r} storage shards are not " + f"aligned to matrix rows of size {matrix_rows}" + ) diff --git a/torchtitan/models/kimi_k2_7/config_registry.py b/torchtitan/models/kimi_k2_7/config_registry.py index 863d7a3391..2951262448 100644 --- a/torchtitan/models/kimi_k2_7/config_registry.py +++ b/torchtitan/models/kimi_k2_7/config_registry.py @@ -9,7 +9,7 @@ from torch.distributed.tensor import Shard from torchtitan.components.checkpoint import CheckpointManager -from torchtitan.components.distributed_optimizers.bucketed_redistribution import ( +from torchtitan.components.distributed_optimizers.flex_optimizer_reshard import ( assign_balanced_owners, BucketConfig, )