Support param group with various placements (#13)

* feat(muon): group parameters by placements for parallel Muon execution

* refactor(muon): refactor step func and group params with it's placement

* feat(muon): support general mesh

* refactor(muon): refactor state init

* refactor(muon): refactor test

* fix(muon): fix general mesh, add chunk_size argument

* refactor(muon): change overlap_step to warmup_step

* refactor(muon-test): rewrite README, add conftest.py and use explicit flags

* chore(muon): clarify N-D sharding support and add test reference

* fix: use device_mesh as key to group params

* Add built binary [skip-build]

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

.pre-commit-config.yaml

@@ -31,7 +31,3 @@ repos:
   hooks:
   - id: pymarkdown
     args: [fix]
-- repo: https://github.com/rhysd/actionlint
-  rev: v1.7.7
-  hooks:
-  - id: actionlint

README.md

@@ -11,7 +11,13 @@ Optimizer is a python package that provides:
 - with support for parallelism techniques for efficient large-scale training.
 
 ## Currently implemented
-- [Parallel Muon with FSDP2](./docs/muon/parallel_muon.pdf)
+- Parallel Muon with N-D sharding
+  - arxiv URL: (TBW)
+  - Supports **general N-D sharding configurations**
+    - The implementation is not tied to any specific parallel strategy.
+    - Verified from basic FSDP2 setups up to hybrid configurations such as
+      **(2 TP + 2 DP-Replicate + 2 DP-Shard)**.
+    - Verified configurations can be found in [test_muon.py](./test/test_muon.py)
 
 ## Usage
 
@@ -39,6 +45,9 @@ optim = optimizer.Muon(
 )
 ```
 
+## Test
+- Check [test/README.md](./test/README.md) for how to run the tests.
+
 ## Pre-commit Hooks
 
 This project uses [pre-commit](https://pre-commit.com/) to automatically check and format code before commits.

build/torch28-cxx11-cu126-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_811726c_dirty
-ops = torch.ops._optimizer_811726c_dirty
+from . import _optimizer_23d68bb_dirty
+ops = torch.ops._optimizer_23d68bb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_811726c_dirty::{op_name}"
+    return f"_optimizer_23d68bb_dirty::{op_name}"

build/torch28-cxx11-cu126-x86_64-linux/optimizer/{_optimizer_811726c_dirty.abi3.so → _optimizer_23d68bb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:511199ac2ae46febc8aeeb96e843a748da7d6fdea4922572ccf27ee5eabe312d
-size 1816064
+oid sha256:35708a107d9ac807fa3e63bbacfc6234fd7622a689a79eae3e43fce11f85d3da
+size 1924376

build/torch28-cxx11-cu126-x86_64-linux/optimizer/distributed/utils.py

@@ -0,0 +1,174 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import (Placement, Shard,
+                                                      _StridedShard)
+
+
+def get_slices_of_dtensor(
+    target: DTensor | torch.Tensor,
+    local_rank: int,
+    shard_mesh: DeviceMesh,
+    shard_placements: tuple[Placement],
+) -> tuple[slice]:
+    """
+    Get the slice of local tensor for a given rank from a tensor.
+    Args:
+        target (DTensor | torch.Tensor): The target tensor.
+        rank (int): The local rank of the shard group.
+        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
+        shard_placements (tuple[Placement]): The shard placements.
+    """
+
+    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
+
+    # find the global rank of the local rank in the shard mesh
+    rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank]
+
+    rank_coords = (shard_mesh.mesh == rank).nonzero()
+
+    assert len(rank_coords) == 1
+    rank_coords = tuple(rank_coords[0].tolist())
+
+    assert len(rank_coords) == len(shard_placements)
+
+    # Caution: Assuming replicate-to-shard of the shard mesh goes with
+    # left-to-right sharding. This is ensured by the sorting logic of
+    # construct_shard_mesh function.
+    for i, (rank_coord,
+            placement) in enumerate(zip(rank_coords, shard_placements)):
+        assert isinstance(placement, Shard)
+
+        num_ranks = shard_mesh.mesh.shape[i]
+
+        dim = placement.dim
+        dim_size = (slices[dim].stop - slices[dim].start)
+
+        if dim_size % num_ranks != 0:
+            raise NotImplementedError(
+                f"Dimension size {dim_size} is not divisible "
+                f"by number of ranks {num_ranks} for shard "
+                f"placement on dim {dim}.")
+
+        shard_size = dim_size // num_ranks
+
+        start = slices[dim].start + rank_coord * shard_size
+        end = start + shard_size
+
+        assert start < end <= slices[dim].stop
+
+        slices[dim] = slice(start, end)
+
+    return tuple(slices)
+
+
+_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, ProcessGroup]] = dict()
+
+
+def construct_shard_mesh(
+    placements: tuple[Placement],
+    mesh: DeviceMesh,
+) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
+    """
+    Construct Shard Mesh and Placements for unsharding.
+    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
+    """
+    my_rank = dist.get_rank()
+
+    assert mesh.mesh.device.type == 'cpu'
+
+    # Copy mesh to avoid modifying the original mesh
+    mesh = mesh.mesh.clone()
+
+    # 1. Sort placements. Replicate first, then Shard by dim ascending.
+
+    # For Shard, strided shard comes after regular shard on the same dim
+    # to preserve left-to-right order of replicate-to-shard.
+    # This is because that strided shard is using stride to represent
+    # more fine-grained sharding on the same dim.
+    # Please check the URL below for _StridedShard.
+    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
+
+    def placement_sort_key(
+        placement_with_index: tuple[float, Placement]
+    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
+        index, placement = placement_with_index
+        is_replicate = placement.is_replicate()
+        is_shard = placement.is_shard()
+        is_partial = placement.is_partial()
+
+        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
+        assert not is_partial, "Partial placement is not supported."
+
+        if is_replicate:
+            return (-1.0, 0, index)
+        elif is_shard:
+            if isinstance(placement, _StridedShard):
+                return (placement.dim, 1 / placement.split_factor, index)
+            return (placement.dim, 0, index)
+        else:
+            raise TypeError(f"Unknown placement type: {type(placement)}")
+
+    placements_with_index: list[tuple[int,
+                                      Placement]] = list(enumerate(placements))
+    placements_with_index = sorted(placements_with_index,
+                                   key=placement_sort_key)
+
+    sorted_indices, sorted_placements = zip(*placements_with_index)
+
+    # 2. Permute mesh according to sorted placements.
+    sorted_mesh = mesh.permute(sorted_indices)
+
+    # 3. Collect list of shard meshes by removing replicate dims
+    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
+    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
+    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
+
+    # merge replicate dims
+    # shard_meshes became a list of shard meshes with a length of replicate degree
+    if num_replicates > 0:
+        sorted_mesh = sorted_mesh.flatten(
+            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
+        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
+    else:
+        shard_meshes = [sorted_mesh]
+    shard_placements = sorted_placements[num_replicates:]
+
+    # assume all shard placements are different
+    assert len(shard_placements) == len(set(shard_placements))
+
+    # 4. Construct ProcessGroups
+    # Caution: all groups should be created in the same order in all processes,
+    # even though each process only needs its own group.
+
+    # To use tensor as dict key, convert it to tuple
+    def tensor_to_tuple(t):
+        if isinstance(t, torch.Tensor):
+            t = t.tolist()
+        if isinstance(t, list):
+            return tuple(tensor_to_tuple(x) for x in t)
+        return t
+
+    my_shard_mesh_as_tuple = None
+    for shard_mesh in shard_meshes:
+        assert isinstance(shard_mesh, torch.Tensor)
+        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
+
+        if (my_rank == shard_mesh).any().item():
+            assert my_shard_mesh_as_tuple is None
+            my_shard_mesh_as_tuple = shard_mesh_as_tuple
+
+        # update global cache
+        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
+            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
+            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
+                DeviceMesh(device_type="cuda", mesh=shard_mesh),
+                shard_process_group,
+            )
+
+    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
+        my_shard_mesh_as_tuple]
+
+    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch28-cxx11-cu126-x86_64-linux/optimizer/muon.py

@@ -1,18 +1,24 @@
 import logging
 import math
 import types
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor, Replicate, Shard
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor.placement_types import Placement
 
+from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
 from .matmul_transpose_triton import matmul_transpose_assign
 
 logger = logging.getLogger(__name__)
 
 COMM_DTYPE = torch.bfloat16
+DEFAULT_CHUNK_SIZE_RATIO = 4
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -62,23 +68,39 @@ def _zeropower_via_newtonschulz5(G, steps):
 @dataclass
 class _muon_state:
     # TODO: use Optional
-    worker_rank: int | None = None
+    worker_rank: int
+    process_group: ProcessGroup
+    shard_mesh: DeviceMesh
+    shard_placements: tuple[Placement, ...]
+    name: str
+    qk_clip_state: torch.Tensor | None = None
     gathered_grad: torch.Tensor | None = None
     scattered_u: DTensor | None = None
     computed_u: torch.Tensor | None = None
     gather_event: torch.cuda.Event | None = None
     compute_event: torch.cuda.Event | None = None
     scatter_event: torch.cuda.Event | None = None
-    process_group = None
-    qk_clip_state = None
 
 
-def split_elems_for_src(param, src_rank, num_ranks) -> int:
-    rows = param.shape[0]
-    cols = int(param.numel() // rows)
-    base, rem = divmod(rows, num_ranks)
-    my_rows = base + (1 if src_rank < rem else 0)
-    return my_rows * cols
+def numel_for_rank(
+    param: DTensor,
+    local_rank: int,
+    state: _muon_state,
+) -> int:
+    slices = get_slices_of_dtensor(
+        param,
+        local_rank,
+        state.shard_mesh,
+        state.shard_placements,
+    )
+
+    numel = 1
+    for s, dim in zip(slices, param.shape):
+        start, stop, step = s.indices(dim)
+        length = max(0, (stop - start + (step - 1)) // step)
+        numel *= length
+
+    return numel
 
 
 @torch.no_grad()
@@ -91,8 +113,7 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
         for p in params:
             state = param_to_state[id(p)]
             if rank == state.worker_rank:
-                num_ranks = dist.get_world_size(group=state.process_group)
-                state.gathered_grad = torch.empty(p.grad.numel(),
+                state.gathered_grad = torch.empty(p.shape,
                                                   dtype=COMM_DTYPE,
                                                   device="cuda")
             else:
@@ -121,11 +142,11 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             state = param_to_state[id(p)]
             dst = state.worker_rank
             assert dst < num_ranks
-            shard_elems = split_elems_for_src(p, rank, num_ranks)
+            shard_elems = numel_for_rank(p, rank, state)
             g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            g = g.to_local().to(COMM_DTYPE).contiguous()
             assert g.numel() == shard_elems
-            per_dst[dst].append(g)
+            per_dst[dst].append(g.view(-1))
             send_counts[dst] += shard_elems
 
         assert any(
@@ -148,13 +169,18 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                total += split_elems_for_src(p, src, num_ranks)
+                total += numel_for_rank(p, src, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
         #All2All
+        logger.debug(f"send_buf size: {send_buf.numel()}, "
+                     f"recv_buf size: {recv_buf.numel()}, "
+                     f"recv_counts: {recv_counts}, "
+                     f"send_counts: {send_counts}, "
+                     f"process_group: {str(process_group)}")
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -179,7 +205,6 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
@@ -189,22 +214,28 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                n = split_elems_for_src(p, src, num_ranks)
+
+                # get the slice of the full dtensor corresponding to rank src.
+                slices = get_slices_of_dtensor(state.gathered_grad, src,
+                                               state.shard_mesh,
+                                               state.shard_placements)
+
+                dst = state.gathered_grad[slices]
+                assert dst._base is state.gathered_grad
+
+                n = dst.numel()
                 assert n > 0
 
                 sg = recv_buf.narrow(0, off + inner_off, n)
-                woff = write_offsets[id(p)]
-                dst = state.gathered_grad.narrow(0, woff, n)
+                sg = sg.reshape_as(dst)
                 dst.copy_(sg)
 
-                write_offsets[id(p)] += n
                 inner_off += n
             off += block
 
         for p in params:
             state = param_to_state[id(p)]
             if state.worker_rank == rank:
-                state.gathered_grad = state.gathered_grad.view_as(p)
                 state.gather_event = torch.cuda.Event()
                 state.gather_event.record(comm_stream)
             else:
@@ -277,14 +308,19 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert state.computed_u is not None
 
-                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous()
 
                 offset = 0
                 for dst in range(num_ranks):
-                    n = split_elems_for_src(p, dst, num_ranks)
+                    # get the slice of the full tensor corresponding to rank dst.
+                    slices = get_slices_of_dtensor(u_full, dst,
+                                                   state.shard_mesh,
+                                                   state.shard_placements)
+                    su = u_full[slices].flatten()
+
+                    n = su.numel()
                     assert n > 0
 
-                    su = u_full.narrow(0, offset, n)
                     per_dst[dst].append(su)
                     send_counts[dst] += n
                     offset += n
@@ -313,7 +349,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                total += split_elems_for_src(p, rank, num_ranks)
+                total += numel_for_rank(p, rank, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
@@ -357,7 +393,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                n = split_elems_for_src(p, rank, num_ranks)
+                n = numel_for_rank(p, rank, state)
                 assert n > 0
 
                 flat_local = recv_buf.narrow(0, off + inner_off,
@@ -398,11 +434,23 @@ def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
         state.scattered_u = None
         u_dtensor = None
 
-        scales_full = Muon._compute_scales(p, state.qk_clip_state)
+        scales_full = Muon._compute_scales(
+            p,
+            state.qk_clip_state) if state.qk_clip_state is not None else None
         if scales_full is not None:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            local_rank = dist.get_rank(group=state.process_group)
-            scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank]
+            # Have to slice scales_full among dim 0
+            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
+                                                  state.shard_placements)
+            ratio = p.shape[0] // scales_full.shape[0]
+            scales_slice = slice(
+                None if weight_slices[0].start is None else
+                weight_slices[0].start // ratio,
+                None if weight_slices[0].stop is None else
+                weight_slices[0].stop // ratio,
+                None,
+            )
+
+            scales_local = scales_full[scales_slice]
             scales_local = DTensor.from_local(
                 scales_local,
                 placements=p.placements,
@@ -478,11 +526,11 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
 @dataclass
 class QKClipInfo:
     """Per-parameter dynamic info computed from config + runtime logits."""
-    kind: Optional[str]  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
-    indices: List[int]  # which heads to consider for clipping
+    kind: str | None  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
+    indices: list[int]  # which heads to consider for clipping
     head_dim: int  # from config
     threshold: float  # from config
-    logit: Optional[torch.Tensor]
+    logit: torch.Tensor | None
 
 
 class Muon(torch.optim.Optimizer):
@@ -525,11 +573,16 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
-        overlap_step : How many all2all gather, compute operations are launched in advance
-                       before the corresponding all2all scatter steps begin.
-                       A higher overlap_step increases memory usage but can improve
-                       performance by overlapping communication.
-                       Parallel muon only. 
+        warmup_step : How many all2all gather, compute operations are launched in advance
+                      before the corresponding all2all scatter steps begin.
+                      A higher warmup_step increases memory usage but can improve
+                      performance by overlapping communication.
+                      Parallel muon only.
+        chunk_size : Batch size of parameters to process in each
+                     all2all gather/compute/scatter step.
+                     Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
+        use_distributed_muon: Use distributed muon by Liu et al. (2024).
+                              For testing purpose only.
     """
 
     def __init__(self,
@@ -549,7 +602,9 @@ class Muon(torch.optim.Optimizer):
                      "head_dim": 128,
                      "threshold": 100
                  },
-                 overlap_step=5):
+                 warmup_step=5,
+                 chunk_size=-1,
+                 use_distributed_muon=False):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -579,7 +634,9 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
-        self.overlap_step = overlap_step
+        self.warmup_step = warmup_step
+        self.chunk_size = chunk_size
+        self.use_distributed_muon = use_distributed_muon
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -597,6 +654,12 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
+    def set_rank_once(self, rank):
+        if self.rank is None:
+            self.rank = rank
+        else:
+            assert self.rank == rank
+
     def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
@@ -604,26 +667,13 @@ class Muon(torch.optim.Optimizer):
         assert isinstance(
             p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0), ):
-            # Case for FSDP
-            process_group = p.device_mesh.get_group(mesh_dim=0)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
-        elif p.placements == (Replicate(), Shard(dim=0)):
-            # Case for HSDP
-            process_group = p.device_mesh.get_group(mesh_dim=1)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if self.rank in shard_mesh:
-                    return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
-        else:
-            raise ValueError(f"Unsupported placements ({p.placements}).")
+        shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+            p.placements, p.device_mesh)
+
+        # set rank with the local rank in the shard process group
+        self.set_rank_once(dist.get_rank(group=shard_pg))
+
+        return shard_mesh, shard_pg, shard_placements
 
     def init_state_and_assign_params(self, names, params, group, qk_logits):
         param_to_state = {}
@@ -655,23 +705,32 @@ class Muon(torch.optim.Optimizer):
         ordered_params = list(params_sorted)
 
         round_robin = 0
-        mesh = None
-        shard_mesh = None
-        process_group = None
+        mesh = ordered_params[0].device_mesh
+        placements = ordered_params[0].placements
+
+        shard_mesh, shard_pg, shard_placements = self.get_shard_mesh(
+            ordered_params[0])
+        shard_mesh_flattened = shard_mesh.mesh.flatten()
+        num_ranks = dist.get_world_size(group=shard_pg)
+
         for n, p in zip(ordered_names, ordered_params):
-            if mesh is None:
-                mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p)
-            elif mesh != p.device_mesh:
+            if mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
-            num_ranks = dist.get_world_size(group=process_group)
-            param_to_state[id(p)] = _muon_state()
-            param_to_state[id(
-                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
-            param_to_state[id(p)].process_group = process_group
+            if placements != p.placements:
+                raise ValueError("All parameters must have same placements.")
+
+            worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
+            round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-            param_to_state[id(p)].qk_clip_state = qk_clip_state
-            round_robin = (round_robin + 1) % len(shard_mesh)
+
+            param_to_state[id(p)] = _muon_state(
+                worker_rank=worker_rank,
+                process_group=shard_pg,
+                shard_mesh=shard_mesh,
+                shard_placements=shard_placements,
+                name=n,
+                qk_clip_state=qk_clip_state,
+            )
 
         return param_to_state, ordered_params
 
@@ -705,10 +764,73 @@ class Muon(torch.optim.Optimizer):
 
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
 
-            scales_full = self._compute_scales(p, qk_clip_state)
+            scales_full = self._compute_scales(
+                p, qk_clip_state) if qk_clip_state is not None else None
             if scales_full is not None:
                 Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
 
+    def distributed_muon(
+        self,
+        names: list[str],
+        params: list[torch.nn.Parameter],
+        group: dict[str, Any],
+        lr: float,
+        weight_decay: float,
+        momentum: float,
+        qk_logits: list[torch.Tensor | DTensor] | None,
+    ):
+        """ Implementation of Distributed Muon by Liu et al. """
+        if qk_logits is not None:
+            raise NotImplementedError("QK clipping is not supported yet")
+
+        if isinstance(params[0], DTensor):
+            shard_mesh, _, shard_placements = construct_shard_mesh(
+                placements=params[0].placements,
+                mesh=params[0].device_mesh,
+            )
+
+        for n, p in zip(names, params):
+            g = p.grad
+            if g is None:
+                continue
+            if g.ndim > 2:
+                g = g.view(g.size(0), -1)
+            assert g is not None
+
+            # calc update
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(momentum).add_(g)
+            if group["nesterov"]:
+                g = g.add(buf, alpha=momentum)
+            else:
+                g = buf
+
+            # Gather G
+            if isinstance(p.data, DTensor):
+                g = g.full_tensor()
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                             steps=group["ns_steps"])
+
+            if isinstance(p.data, DTensor):
+                slices = get_slices_of_dtensor(
+                    target=p,
+                    local_rank=dist.get_rank(),
+                    shard_mesh=shard_mesh,
+                    shard_placements=shard_placements,
+                )
+                u_shard = u[slices]
+                u = DTensor.from_local(
+                    u_shard,
+                    device_mesh=p.device_mesh,
+                    placements=p.placements,
+                )
+
+            adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+            Muon._update_p(p, u, lr, adjusted_lr, weight_decay)
+
     def _update_g(self, p, g, group, momentum):
         # calc update
         state = self.state[p]
@@ -727,6 +849,9 @@ class Muon(torch.optim.Optimizer):
         p.data.add_(u, alpha=-adjusted_lr)
 
     def get_qk_clip_info(self, n, qk_logits):
+        if self.clip_config is None:
+            return None
+
         head_dim = self.clip_config.get('head_dim')
         threshold = self.clip_config.get('threshold')
         kind, layer_idx = parse_qk_layer(n)
@@ -737,6 +862,11 @@ class Muon(torch.optim.Optimizer):
             indices_key = 'q_indices' if 'q' in kind else 'k_indices'
             indices = self.clip_config.get(indices_key, []) or []
 
+            if isinstance(logit, DTensor):
+                # In TP settings, qk_logits may be DTensor
+                # We convert it to full tensor here for simplicity
+                logit = logit.full_tensor()
+
         return QKClipInfo(
             kind=kind,
             indices=indices,
@@ -835,22 +965,28 @@ class Muon(torch.optim.Optimizer):
                 _update_param(p, state, lr, adjusted_lr, weight_decay,
                               self.rank, self.compute_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(
-            params[0])].process_group)
+        if self.chunk_size == -1:
+            shard_ranks = dist.get_world_size(param_to_state[id(
+                params[0])].process_group)
+            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+        elif self.chunk_size > 0:
+            chunk_size = self.chunk_size
+        else:
+            raise ValueError("chunk_size must be -1 or a positive integer.")
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        overlap_step = self.overlap_step
-        for i in range(0, overlap_step):
+        warmup_step = self.warmup_step
+        for i in range(0, warmup_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
+            enqueue_computes(i + warmup_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)
@@ -866,7 +1002,7 @@ class Muon(torch.optim.Optimizer):
         amsgrad: bool,
         beta1: float,
         beta2: float,
-        lr: Union[float, torch.Tensor],
+        lr: float | torch.Tensor,
         weight_decay: float,
         eps: float,
         maximize: bool,
@@ -876,10 +1012,10 @@ class Muon(torch.optim.Optimizer):
 
         # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer
         # treating it as a scalar.
-        lr_dict: Optional[DeviceDict] = ({
+        lr_dict: DeviceDict | None = ({
             lr.device: lr
         } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else
-                                         None)
+                                      None)
         grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype(
             [
                 params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
@@ -926,6 +1062,159 @@ class Muon(torch.optim.Optimizer):
                 maximize=maximize,
             )
 
+    def _step_muon(self, group, qk_logits=None):
+        params = group["params"]
+        lr = group["lr"]
+        weight_decay = group["weight_decay"]
+        momentum = group["momentum"]
+        names = group["names"]
+
+        param_dtensors = []
+        param_tensors = []
+        name_dtensors = []
+        name_tensors = []
+
+        if self.use_distributed_muon:
+            self.distributed_muon(names=names,
+                                  params=params,
+                                  group=group,
+                                  lr=lr,
+                                  weight_decay=weight_decay,
+                                  momentum=momentum,
+                                  qk_logits=qk_logits)
+            return
+
+        for n, p in zip(names, params):
+            if p is None or p.grad is None:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(
+                        isinstance(placement, Replicate)
+                        for placement in p.placements):
+                    param_tensors.append(p)
+                    name_tensors.append(n)
+                else:
+                    param_dtensors.append(p)
+                    name_dtensors.append(n)
+            elif isinstance(p.data, torch.Tensor):
+                param_tensors.append(p)
+                name_tensors.append(n)
+            else:
+                raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+        logger.debug(
+            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors"
+        )
+
+        if len(param_dtensors) > 0:
+            if not dist.is_initialized():
+                raise RuntimeError(
+                    "Parallel Muon requires torch.distributed to be initialized."
+                )
+
+            # To support different placements, we group parameters by placements
+            # and run parallel Muon on each group.
+
+            placement_to_params = defaultdict(lambda: ([], []))
+            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
+
+            assert len(name_dtensors) == len(param_dtensors)
+            for n, p in zip(name_dtensors, param_dtensors):
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][0].append(n)
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][1].append(p)
+
+            for _, (names, params) in placement_to_params.items():
+                self.parallel(
+                    names,
+                    params,
+                    group,
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    momentum=momentum,
+                    qk_logits=qk_logits,
+                )
+
+        if len(param_tensors) > 0:
+            self.base(
+                name_tensors,
+                param_tensors,
+                group,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                qk_logits=qk_logits,
+            )
+
+    def _step_adamw_params(self, params, group):
+        params_with_grads = []
+        grads = []
+        moment1 = []
+        moment2 = []
+        max_exp_avg_sqs = []
+        state_steps = []
+        lr = group["lr"]
+        beta1, beta2 = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        weight_decay = group["weight_decay"]
+
+        for p in params:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            params_with_grads.append(p)
+            grads.append(g)
+            if "step" not in state:
+                state["step"] = (torch.zeros((),
+                                             dtype=torch.float32,
+                                             device=p.device))
+                state["moment1"] = torch.zeros_like(g)
+                state["moment2"] = torch.zeros_like(g)
+            moment1.append(state["moment1"])
+            moment2.append(state["moment2"])
+            if not isinstance(state["step"], torch.Tensor):
+                step_tensor = torch.tensor(state["step"],
+                                           dtype=torch.float32,
+                                           device=p.device)
+            else:
+                step_tensor = state["step"]
+            state_steps.append(step_tensor)
+
+        self._fused_adamw(
+            params_with_grads,
+            grads,
+            moment1,
+            moment2,
+            max_exp_avg_sqs,
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+    def _step_adamw(self, group):
+        params = group["params"]
+
+        # group params with it's type and placement
+        placement_to_params: dict[tuple[Placement | type,
+                                        DeviceMesh | None]] = defaultdict(list)
+        for p in params:
+            match p:
+                case DTensor():
+                    placement_to_params[tuple([p.placements,
+                                               p.device_mesh])].append(p)
+                case torch.Tensor():
+                    placement_to_params[tuple([torch.Tensor, None])].append(p)
+
+        for params in placement_to_params.values():
+            self._step_adamw_params(params, group)
+
     def step(self, closure=None, qk_logits=None):
         """Perform a single optimization step.
 
@@ -943,127 +1232,9 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         for group in self.param_groups:
-            params = group["params"]
-
             if group["use_muon"]:
-                ############################
-                #           Muon           #
-                ############################
-                lr = group["lr"]
-                weight_decay = group["weight_decay"]
-                momentum = group["momentum"]
-                names = group["names"]
-
-                param_dtensors = []
-                param_tensors = []
-                name_dtensors = []
-                name_tensors = []
-
-                for n, p in zip(names, params):
-                    if p is None or p.grad is None:
-                        continue
-                    if isinstance(p.data, DTensor):
-                        if all(
-                                isinstance(placement, Replicate)
-                                for placement in p.placements):
-                            param_tensors.append(p)
-                            name_tensors.append(n)
-                        else:
-                            param_dtensors.append(p)
-                            name_dtensors.append(n)
-                    elif isinstance(p.data, torch.Tensor):
-                        param_tensors.append(p)
-                        name_tensors.append(n)
-                    else:
-                        raise TypeError(
-                            f"Unsupported parameter type: {type(p.data)}")
-
-                if self.debug:
-                    print(
-                        f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
-                        flush=True,
-                    )
-
-                if len(param_dtensors) > 0:
-                    if not dist.is_initialized():
-                        raise RuntimeError(
-                            "Parallel Muon requires torch.distributed to be initialized."
-                        )
-
-                    self.parallel(
-                        name_dtensors,
-                        param_dtensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
-                if len(param_tensors) > 0:
-                    self.base(
-                        name_tensors,
-                        param_tensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
+                self._step_muon(group, qk_logits=qk_logits)
             else:
-                ############################
-                #       AdamW backup       #
-                ############################
-
-                params_with_grads = []
-                grads = []
-                moment1 = []
-                moment2 = []
-                max_exp_avg_sqs = []
-                state_steps = []
-                lr = group["lr"]
-                beta1, beta2 = group["adamw_betas"]
-                eps = group["adamw_eps"]
-                weight_decay = group["weight_decay"]
-
-                for p in params:
-                    g = p.grad
-                    if g is None:
-                        continue
-                    state = self.state[p]
-                    params_with_grads.append(p)
-                    grads.append(g)
-                    if "step" not in state:
-                        state["step"] = (torch.zeros((),
-                                                     dtype=torch.float32,
-                                                     device=p.device))
-                        state["moment1"] = torch.zeros_like(g)
-                        state["moment2"] = torch.zeros_like(g)
-                    moment1.append(state["moment1"])
-                    moment2.append(state["moment2"])
-                    if not isinstance(state["step"], torch.Tensor):
-                        step_tensor = torch.tensor(state["step"],
-                                                   dtype=torch.float32,
-                                                   device=p.device)
-                    else:
-                        step_tensor = state["step"]
-                    state_steps.append(step_tensor)
-
-                self._fused_adamw(
-                    params_with_grads,
-                    grads,
-                    moment1,
-                    moment2,
-                    max_exp_avg_sqs,
-                    state_steps,
-                    amsgrad=False,
-                    beta1=beta1,
-                    beta2=beta2,
-                    lr=lr,
-                    weight_decay=weight_decay,
-                    eps=eps,
-                    maximize=False,
-                )
+                self._step_adamw(group)
 
         return loss

build/torch28-cxx11-cu128-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_811726c_dirty
-ops = torch.ops._optimizer_811726c_dirty
+from . import _optimizer_23d68bb_dirty
+ops = torch.ops._optimizer_23d68bb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_811726c_dirty::{op_name}"
+    return f"_optimizer_23d68bb_dirty::{op_name}"

build/torch28-cxx11-cu128-x86_64-linux/optimizer/{_optimizer_811726c_dirty.abi3.so → _optimizer_23d68bb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b3cdb515b6c56204224cc307b66d34fcee1cd5e27b4117197a71b784d34fadc5
-size 1871056
+oid sha256:03c3bbbbc5c4ceb5cebfe3a2e411f155bebb390f1921c14d59fcf791dd556da1
+size 1983488

build/torch28-cxx11-cu128-x86_64-linux/optimizer/distributed/utils.py

@@ -0,0 +1,174 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import (Placement, Shard,
+                                                      _StridedShard)
+
+
+def get_slices_of_dtensor(
+    target: DTensor | torch.Tensor,
+    local_rank: int,
+    shard_mesh: DeviceMesh,
+    shard_placements: tuple[Placement],
+) -> tuple[slice]:
+    """
+    Get the slice of local tensor for a given rank from a tensor.
+    Args:
+        target (DTensor | torch.Tensor): The target tensor.
+        rank (int): The local rank of the shard group.
+        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
+        shard_placements (tuple[Placement]): The shard placements.
+    """
+
+    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
+
+    # find the global rank of the local rank in the shard mesh
+    rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank]
+
+    rank_coords = (shard_mesh.mesh == rank).nonzero()
+
+    assert len(rank_coords) == 1
+    rank_coords = tuple(rank_coords[0].tolist())
+
+    assert len(rank_coords) == len(shard_placements)
+
+    # Caution: Assuming replicate-to-shard of the shard mesh goes with
+    # left-to-right sharding. This is ensured by the sorting logic of
+    # construct_shard_mesh function.
+    for i, (rank_coord,
+            placement) in enumerate(zip(rank_coords, shard_placements)):
+        assert isinstance(placement, Shard)
+
+        num_ranks = shard_mesh.mesh.shape[i]
+
+        dim = placement.dim
+        dim_size = (slices[dim].stop - slices[dim].start)
+
+        if dim_size % num_ranks != 0:
+            raise NotImplementedError(
+                f"Dimension size {dim_size} is not divisible "
+                f"by number of ranks {num_ranks} for shard "
+                f"placement on dim {dim}.")
+
+        shard_size = dim_size // num_ranks
+
+        start = slices[dim].start + rank_coord * shard_size
+        end = start + shard_size
+
+        assert start < end <= slices[dim].stop
+
+        slices[dim] = slice(start, end)
+
+    return tuple(slices)
+
+
+_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, ProcessGroup]] = dict()
+
+
+def construct_shard_mesh(
+    placements: tuple[Placement],
+    mesh: DeviceMesh,
+) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
+    """
+    Construct Shard Mesh and Placements for unsharding.
+    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
+    """
+    my_rank = dist.get_rank()
+
+    assert mesh.mesh.device.type == 'cpu'
+
+    # Copy mesh to avoid modifying the original mesh
+    mesh = mesh.mesh.clone()
+
+    # 1. Sort placements. Replicate first, then Shard by dim ascending.
+
+    # For Shard, strided shard comes after regular shard on the same dim
+    # to preserve left-to-right order of replicate-to-shard.
+    # This is because that strided shard is using stride to represent
+    # more fine-grained sharding on the same dim.
+    # Please check the URL below for _StridedShard.
+    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
+
+    def placement_sort_key(
+        placement_with_index: tuple[float, Placement]
+    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
+        index, placement = placement_with_index
+        is_replicate = placement.is_replicate()
+        is_shard = placement.is_shard()
+        is_partial = placement.is_partial()
+
+        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
+        assert not is_partial, "Partial placement is not supported."
+
+        if is_replicate:
+            return (-1.0, 0, index)
+        elif is_shard:
+            if isinstance(placement, _StridedShard):
+                return (placement.dim, 1 / placement.split_factor, index)
+            return (placement.dim, 0, index)
+        else:
+            raise TypeError(f"Unknown placement type: {type(placement)}")
+
+    placements_with_index: list[tuple[int,
+                                      Placement]] = list(enumerate(placements))
+    placements_with_index = sorted(placements_with_index,
+                                   key=placement_sort_key)
+
+    sorted_indices, sorted_placements = zip(*placements_with_index)
+
+    # 2. Permute mesh according to sorted placements.
+    sorted_mesh = mesh.permute(sorted_indices)
+
+    # 3. Collect list of shard meshes by removing replicate dims
+    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
+    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
+    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
+
+    # merge replicate dims
+    # shard_meshes became a list of shard meshes with a length of replicate degree
+    if num_replicates > 0:
+        sorted_mesh = sorted_mesh.flatten(
+            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
+        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
+    else:
+        shard_meshes = [sorted_mesh]
+    shard_placements = sorted_placements[num_replicates:]
+
+    # assume all shard placements are different
+    assert len(shard_placements) == len(set(shard_placements))
+
+    # 4. Construct ProcessGroups
+    # Caution: all groups should be created in the same order in all processes,
+    # even though each process only needs its own group.
+
+    # To use tensor as dict key, convert it to tuple
+    def tensor_to_tuple(t):
+        if isinstance(t, torch.Tensor):
+            t = t.tolist()
+        if isinstance(t, list):
+            return tuple(tensor_to_tuple(x) for x in t)
+        return t
+
+    my_shard_mesh_as_tuple = None
+    for shard_mesh in shard_meshes:
+        assert isinstance(shard_mesh, torch.Tensor)
+        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
+
+        if (my_rank == shard_mesh).any().item():
+            assert my_shard_mesh_as_tuple is None
+            my_shard_mesh_as_tuple = shard_mesh_as_tuple
+
+        # update global cache
+        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
+            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
+            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
+                DeviceMesh(device_type="cuda", mesh=shard_mesh),
+                shard_process_group,
+            )
+
+    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
+        my_shard_mesh_as_tuple]
+
+    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch28-cxx11-cu128-x86_64-linux/optimizer/muon.py

@@ -1,18 +1,24 @@
 import logging
 import math
 import types
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor, Replicate, Shard
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor.placement_types import Placement
 
+from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
 from .matmul_transpose_triton import matmul_transpose_assign
 
 logger = logging.getLogger(__name__)
 
 COMM_DTYPE = torch.bfloat16
+DEFAULT_CHUNK_SIZE_RATIO = 4
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -62,23 +68,39 @@ def _zeropower_via_newtonschulz5(G, steps):
 @dataclass
 class _muon_state:
     # TODO: use Optional
-    worker_rank: int | None = None
+    worker_rank: int
+    process_group: ProcessGroup
+    shard_mesh: DeviceMesh
+    shard_placements: tuple[Placement, ...]
+    name: str
+    qk_clip_state: torch.Tensor | None = None
     gathered_grad: torch.Tensor | None = None
     scattered_u: DTensor | None = None
     computed_u: torch.Tensor | None = None
     gather_event: torch.cuda.Event | None = None
     compute_event: torch.cuda.Event | None = None
     scatter_event: torch.cuda.Event | None = None
-    process_group = None
-    qk_clip_state = None
 
 
-def split_elems_for_src(param, src_rank, num_ranks) -> int:
-    rows = param.shape[0]
-    cols = int(param.numel() // rows)
-    base, rem = divmod(rows, num_ranks)
-    my_rows = base + (1 if src_rank < rem else 0)
-    return my_rows * cols
+def numel_for_rank(
+    param: DTensor,
+    local_rank: int,
+    state: _muon_state,
+) -> int:
+    slices = get_slices_of_dtensor(
+        param,
+        local_rank,
+        state.shard_mesh,
+        state.shard_placements,
+    )
+
+    numel = 1
+    for s, dim in zip(slices, param.shape):
+        start, stop, step = s.indices(dim)
+        length = max(0, (stop - start + (step - 1)) // step)
+        numel *= length
+
+    return numel
 
 
 @torch.no_grad()
@@ -91,8 +113,7 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
         for p in params:
             state = param_to_state[id(p)]
             if rank == state.worker_rank:
-                num_ranks = dist.get_world_size(group=state.process_group)
-                state.gathered_grad = torch.empty(p.grad.numel(),
+                state.gathered_grad = torch.empty(p.shape,
                                                   dtype=COMM_DTYPE,
                                                   device="cuda")
             else:
@@ -121,11 +142,11 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             state = param_to_state[id(p)]
             dst = state.worker_rank
             assert dst < num_ranks
-            shard_elems = split_elems_for_src(p, rank, num_ranks)
+            shard_elems = numel_for_rank(p, rank, state)
             g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            g = g.to_local().to(COMM_DTYPE).contiguous()
             assert g.numel() == shard_elems
-            per_dst[dst].append(g)
+            per_dst[dst].append(g.view(-1))
             send_counts[dst] += shard_elems
 
         assert any(
@@ -148,13 +169,18 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                total += split_elems_for_src(p, src, num_ranks)
+                total += numel_for_rank(p, src, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
         #All2All
+        logger.debug(f"send_buf size: {send_buf.numel()}, "
+                     f"recv_buf size: {recv_buf.numel()}, "
+                     f"recv_counts: {recv_counts}, "
+                     f"send_counts: {send_counts}, "
+                     f"process_group: {str(process_group)}")
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -179,7 +205,6 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
@@ -189,22 +214,28 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                n = split_elems_for_src(p, src, num_ranks)
+
+                # get the slice of the full dtensor corresponding to rank src.
+                slices = get_slices_of_dtensor(state.gathered_grad, src,
+                                               state.shard_mesh,
+                                               state.shard_placements)
+
+                dst = state.gathered_grad[slices]
+                assert dst._base is state.gathered_grad
+
+                n = dst.numel()
                 assert n > 0
 
                 sg = recv_buf.narrow(0, off + inner_off, n)
-                woff = write_offsets[id(p)]
-                dst = state.gathered_grad.narrow(0, woff, n)
+                sg = sg.reshape_as(dst)
                 dst.copy_(sg)
 
-                write_offsets[id(p)] += n
                 inner_off += n
             off += block
 
         for p in params:
             state = param_to_state[id(p)]
             if state.worker_rank == rank:
-                state.gathered_grad = state.gathered_grad.view_as(p)
                 state.gather_event = torch.cuda.Event()
                 state.gather_event.record(comm_stream)
             else:
@@ -277,14 +308,19 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert state.computed_u is not None
 
-                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous()
 
                 offset = 0
                 for dst in range(num_ranks):
-                    n = split_elems_for_src(p, dst, num_ranks)
+                    # get the slice of the full tensor corresponding to rank dst.
+                    slices = get_slices_of_dtensor(u_full, dst,
+                                                   state.shard_mesh,
+                                                   state.shard_placements)
+                    su = u_full[slices].flatten()
+
+                    n = su.numel()
                     assert n > 0
 
-                    su = u_full.narrow(0, offset, n)
                     per_dst[dst].append(su)
                     send_counts[dst] += n
                     offset += n
@@ -313,7 +349,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                total += split_elems_for_src(p, rank, num_ranks)
+                total += numel_for_rank(p, rank, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
@@ -357,7 +393,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                n = split_elems_for_src(p, rank, num_ranks)
+                n = numel_for_rank(p, rank, state)
                 assert n > 0
 
                 flat_local = recv_buf.narrow(0, off + inner_off,
@@ -398,11 +434,23 @@ def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
         state.scattered_u = None
         u_dtensor = None
 
-        scales_full = Muon._compute_scales(p, state.qk_clip_state)
+        scales_full = Muon._compute_scales(
+            p,
+            state.qk_clip_state) if state.qk_clip_state is not None else None
         if scales_full is not None:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            local_rank = dist.get_rank(group=state.process_group)
-            scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank]
+            # Have to slice scales_full among dim 0
+            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
+                                                  state.shard_placements)
+            ratio = p.shape[0] // scales_full.shape[0]
+            scales_slice = slice(
+                None if weight_slices[0].start is None else
+                weight_slices[0].start // ratio,
+                None if weight_slices[0].stop is None else
+                weight_slices[0].stop // ratio,
+                None,
+            )
+
+            scales_local = scales_full[scales_slice]
             scales_local = DTensor.from_local(
                 scales_local,
                 placements=p.placements,
@@ -478,11 +526,11 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
 @dataclass
 class QKClipInfo:
     """Per-parameter dynamic info computed from config + runtime logits."""
-    kind: Optional[str]  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
-    indices: List[int]  # which heads to consider for clipping
+    kind: str | None  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
+    indices: list[int]  # which heads to consider for clipping
     head_dim: int  # from config
     threshold: float  # from config
-    logit: Optional[torch.Tensor]
+    logit: torch.Tensor | None
 
 
 class Muon(torch.optim.Optimizer):
@@ -525,11 +573,16 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
-        overlap_step : How many all2all gather, compute operations are launched in advance
-                       before the corresponding all2all scatter steps begin.
-                       A higher overlap_step increases memory usage but can improve
-                       performance by overlapping communication.
-                       Parallel muon only. 
+        warmup_step : How many all2all gather, compute operations are launched in advance
+                      before the corresponding all2all scatter steps begin.
+                      A higher warmup_step increases memory usage but can improve
+                      performance by overlapping communication.
+                      Parallel muon only.
+        chunk_size : Batch size of parameters to process in each
+                     all2all gather/compute/scatter step.
+                     Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
+        use_distributed_muon: Use distributed muon by Liu et al. (2024).
+                              For testing purpose only.
     """
 
     def __init__(self,
@@ -549,7 +602,9 @@ class Muon(torch.optim.Optimizer):
                      "head_dim": 128,
                      "threshold": 100
                  },
-                 overlap_step=5):
+                 warmup_step=5,
+                 chunk_size=-1,
+                 use_distributed_muon=False):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -579,7 +634,9 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
-        self.overlap_step = overlap_step
+        self.warmup_step = warmup_step
+        self.chunk_size = chunk_size
+        self.use_distributed_muon = use_distributed_muon
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -597,6 +654,12 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
+    def set_rank_once(self, rank):
+        if self.rank is None:
+            self.rank = rank
+        else:
+            assert self.rank == rank
+
     def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
@@ -604,26 +667,13 @@ class Muon(torch.optim.Optimizer):
         assert isinstance(
             p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0), ):
-            # Case for FSDP
-            process_group = p.device_mesh.get_group(mesh_dim=0)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
-        elif p.placements == (Replicate(), Shard(dim=0)):
-            # Case for HSDP
-            process_group = p.device_mesh.get_group(mesh_dim=1)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if self.rank in shard_mesh:
-                    return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
-        else:
-            raise ValueError(f"Unsupported placements ({p.placements}).")
+        shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+            p.placements, p.device_mesh)
+
+        # set rank with the local rank in the shard process group
+        self.set_rank_once(dist.get_rank(group=shard_pg))
+
+        return shard_mesh, shard_pg, shard_placements
 
     def init_state_and_assign_params(self, names, params, group, qk_logits):
         param_to_state = {}
@@ -655,23 +705,32 @@ class Muon(torch.optim.Optimizer):
         ordered_params = list(params_sorted)
 
         round_robin = 0
-        mesh = None
-        shard_mesh = None
-        process_group = None
+        mesh = ordered_params[0].device_mesh
+        placements = ordered_params[0].placements
+
+        shard_mesh, shard_pg, shard_placements = self.get_shard_mesh(
+            ordered_params[0])
+        shard_mesh_flattened = shard_mesh.mesh.flatten()
+        num_ranks = dist.get_world_size(group=shard_pg)
+
         for n, p in zip(ordered_names, ordered_params):
-            if mesh is None:
-                mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p)
-            elif mesh != p.device_mesh:
+            if mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
-            num_ranks = dist.get_world_size(group=process_group)
-            param_to_state[id(p)] = _muon_state()
-            param_to_state[id(
-                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
-            param_to_state[id(p)].process_group = process_group
+            if placements != p.placements:
+                raise ValueError("All parameters must have same placements.")
+
+            worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
+            round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-            param_to_state[id(p)].qk_clip_state = qk_clip_state
-            round_robin = (round_robin + 1) % len(shard_mesh)
+
+            param_to_state[id(p)] = _muon_state(
+                worker_rank=worker_rank,
+                process_group=shard_pg,
+                shard_mesh=shard_mesh,
+                shard_placements=shard_placements,
+                name=n,
+                qk_clip_state=qk_clip_state,
+            )
 
         return param_to_state, ordered_params
 
@@ -705,10 +764,73 @@ class Muon(torch.optim.Optimizer):
 
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
 
-            scales_full = self._compute_scales(p, qk_clip_state)
+            scales_full = self._compute_scales(
+                p, qk_clip_state) if qk_clip_state is not None else None
             if scales_full is not None:
                 Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
 
+    def distributed_muon(
+        self,
+        names: list[str],
+        params: list[torch.nn.Parameter],
+        group: dict[str, Any],
+        lr: float,
+        weight_decay: float,
+        momentum: float,
+        qk_logits: list[torch.Tensor | DTensor] | None,
+    ):
+        """ Implementation of Distributed Muon by Liu et al. """
+        if qk_logits is not None:
+            raise NotImplementedError("QK clipping is not supported yet")
+
+        if isinstance(params[0], DTensor):
+            shard_mesh, _, shard_placements = construct_shard_mesh(
+                placements=params[0].placements,
+                mesh=params[0].device_mesh,
+            )
+
+        for n, p in zip(names, params):
+            g = p.grad
+            if g is None:
+                continue
+            if g.ndim > 2:
+                g = g.view(g.size(0), -1)
+            assert g is not None
+
+            # calc update
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(momentum).add_(g)
+            if group["nesterov"]:
+                g = g.add(buf, alpha=momentum)
+            else:
+                g = buf
+
+            # Gather G
+            if isinstance(p.data, DTensor):
+                g = g.full_tensor()
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                             steps=group["ns_steps"])
+
+            if isinstance(p.data, DTensor):
+                slices = get_slices_of_dtensor(
+                    target=p,
+                    local_rank=dist.get_rank(),
+                    shard_mesh=shard_mesh,
+                    shard_placements=shard_placements,
+                )
+                u_shard = u[slices]
+                u = DTensor.from_local(
+                    u_shard,
+                    device_mesh=p.device_mesh,
+                    placements=p.placements,
+                )
+
+            adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+            Muon._update_p(p, u, lr, adjusted_lr, weight_decay)
+
     def _update_g(self, p, g, group, momentum):
         # calc update
         state = self.state[p]
@@ -727,6 +849,9 @@ class Muon(torch.optim.Optimizer):
         p.data.add_(u, alpha=-adjusted_lr)
 
     def get_qk_clip_info(self, n, qk_logits):
+        if self.clip_config is None:
+            return None
+
         head_dim = self.clip_config.get('head_dim')
         threshold = self.clip_config.get('threshold')
         kind, layer_idx = parse_qk_layer(n)
@@ -737,6 +862,11 @@ class Muon(torch.optim.Optimizer):
             indices_key = 'q_indices' if 'q' in kind else 'k_indices'
             indices = self.clip_config.get(indices_key, []) or []
 
+            if isinstance(logit, DTensor):
+                # In TP settings, qk_logits may be DTensor
+                # We convert it to full tensor here for simplicity
+                logit = logit.full_tensor()
+
         return QKClipInfo(
             kind=kind,
             indices=indices,
@@ -835,22 +965,28 @@ class Muon(torch.optim.Optimizer):
                 _update_param(p, state, lr, adjusted_lr, weight_decay,
                               self.rank, self.compute_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(
-            params[0])].process_group)
+        if self.chunk_size == -1:
+            shard_ranks = dist.get_world_size(param_to_state[id(
+                params[0])].process_group)
+            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+        elif self.chunk_size > 0:
+            chunk_size = self.chunk_size
+        else:
+            raise ValueError("chunk_size must be -1 or a positive integer.")
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        overlap_step = self.overlap_step
-        for i in range(0, overlap_step):
+        warmup_step = self.warmup_step
+        for i in range(0, warmup_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
+            enqueue_computes(i + warmup_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)
@@ -866,7 +1002,7 @@ class Muon(torch.optim.Optimizer):
         amsgrad: bool,
         beta1: float,
         beta2: float,
-        lr: Union[float, torch.Tensor],
+        lr: float | torch.Tensor,
         weight_decay: float,
         eps: float,
         maximize: bool,
@@ -876,10 +1012,10 @@ class Muon(torch.optim.Optimizer):
 
         # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer
         # treating it as a scalar.
-        lr_dict: Optional[DeviceDict] = ({
+        lr_dict: DeviceDict | None = ({
             lr.device: lr
         } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else
-                                         None)
+                                      None)
         grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype(
             [
                 params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
@@ -926,6 +1062,159 @@ class Muon(torch.optim.Optimizer):
                 maximize=maximize,
             )
 
+    def _step_muon(self, group, qk_logits=None):
+        params = group["params"]
+        lr = group["lr"]
+        weight_decay = group["weight_decay"]
+        momentum = group["momentum"]
+        names = group["names"]
+
+        param_dtensors = []
+        param_tensors = []
+        name_dtensors = []
+        name_tensors = []
+
+        if self.use_distributed_muon:
+            self.distributed_muon(names=names,
+                                  params=params,
+                                  group=group,
+                                  lr=lr,
+                                  weight_decay=weight_decay,
+                                  momentum=momentum,
+                                  qk_logits=qk_logits)
+            return
+
+        for n, p in zip(names, params):
+            if p is None or p.grad is None:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(
+                        isinstance(placement, Replicate)
+                        for placement in p.placements):
+                    param_tensors.append(p)
+                    name_tensors.append(n)
+                else:
+                    param_dtensors.append(p)
+                    name_dtensors.append(n)
+            elif isinstance(p.data, torch.Tensor):
+                param_tensors.append(p)
+                name_tensors.append(n)
+            else:
+                raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+        logger.debug(
+            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors"
+        )
+
+        if len(param_dtensors) > 0:
+            if not dist.is_initialized():
+                raise RuntimeError(
+                    "Parallel Muon requires torch.distributed to be initialized."
+                )
+
+            # To support different placements, we group parameters by placements
+            # and run parallel Muon on each group.
+
+            placement_to_params = defaultdict(lambda: ([], []))
+            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
+
+            assert len(name_dtensors) == len(param_dtensors)
+            for n, p in zip(name_dtensors, param_dtensors):
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][0].append(n)
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][1].append(p)
+
+            for _, (names, params) in placement_to_params.items():
+                self.parallel(
+                    names,
+                    params,
+                    group,
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    momentum=momentum,
+                    qk_logits=qk_logits,
+                )
+
+        if len(param_tensors) > 0:
+            self.base(
+                name_tensors,
+                param_tensors,
+                group,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                qk_logits=qk_logits,
+            )
+
+    def _step_adamw_params(self, params, group):
+        params_with_grads = []
+        grads = []
+        moment1 = []
+        moment2 = []
+        max_exp_avg_sqs = []
+        state_steps = []
+        lr = group["lr"]
+        beta1, beta2 = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        weight_decay = group["weight_decay"]
+
+        for p in params:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            params_with_grads.append(p)
+            grads.append(g)
+            if "step" not in state:
+                state["step"] = (torch.zeros((),
+                                             dtype=torch.float32,
+                                             device=p.device))
+                state["moment1"] = torch.zeros_like(g)
+                state["moment2"] = torch.zeros_like(g)
+            moment1.append(state["moment1"])
+            moment2.append(state["moment2"])
+            if not isinstance(state["step"], torch.Tensor):
+                step_tensor = torch.tensor(state["step"],
+                                           dtype=torch.float32,
+                                           device=p.device)
+            else:
+                step_tensor = state["step"]
+            state_steps.append(step_tensor)
+
+        self._fused_adamw(
+            params_with_grads,
+            grads,
+            moment1,
+            moment2,
+            max_exp_avg_sqs,
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+    def _step_adamw(self, group):
+        params = group["params"]
+
+        # group params with it's type and placement
+        placement_to_params: dict[tuple[Placement | type,
+                                        DeviceMesh | None]] = defaultdict(list)
+        for p in params:
+            match p:
+                case DTensor():
+                    placement_to_params[tuple([p.placements,
+                                               p.device_mesh])].append(p)
+                case torch.Tensor():
+                    placement_to_params[tuple([torch.Tensor, None])].append(p)
+
+        for params in placement_to_params.values():
+            self._step_adamw_params(params, group)
+
     def step(self, closure=None, qk_logits=None):
         """Perform a single optimization step.
 
@@ -943,127 +1232,9 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         for group in self.param_groups:
-            params = group["params"]
-
             if group["use_muon"]:
-                ############################
-                #           Muon           #
-                ############################
-                lr = group["lr"]
-                weight_decay = group["weight_decay"]
-                momentum = group["momentum"]
-                names = group["names"]
-
-                param_dtensors = []
-                param_tensors = []
-                name_dtensors = []
-                name_tensors = []
-
-                for n, p in zip(names, params):
-                    if p is None or p.grad is None:
-                        continue
-                    if isinstance(p.data, DTensor):
-                        if all(
-                                isinstance(placement, Replicate)
-                                for placement in p.placements):
-                            param_tensors.append(p)
-                            name_tensors.append(n)
-                        else:
-                            param_dtensors.append(p)
-                            name_dtensors.append(n)
-                    elif isinstance(p.data, torch.Tensor):
-                        param_tensors.append(p)
-                        name_tensors.append(n)
-                    else:
-                        raise TypeError(
-                            f"Unsupported parameter type: {type(p.data)}")
-
-                if self.debug:
-                    print(
-                        f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
-                        flush=True,
-                    )
-
-                if len(param_dtensors) > 0:
-                    if not dist.is_initialized():
-                        raise RuntimeError(
-                            "Parallel Muon requires torch.distributed to be initialized."
-                        )
-
-                    self.parallel(
-                        name_dtensors,
-                        param_dtensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
-                if len(param_tensors) > 0:
-                    self.base(
-                        name_tensors,
-                        param_tensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
+                self._step_muon(group, qk_logits=qk_logits)
             else:
-                ############################
-                #       AdamW backup       #
-                ############################
-
-                params_with_grads = []
-                grads = []
-                moment1 = []
-                moment2 = []
-                max_exp_avg_sqs = []
-                state_steps = []
-                lr = group["lr"]
-                beta1, beta2 = group["adamw_betas"]
-                eps = group["adamw_eps"]
-                weight_decay = group["weight_decay"]
-
-                for p in params:
-                    g = p.grad
-                    if g is None:
-                        continue
-                    state = self.state[p]
-                    params_with_grads.append(p)
-                    grads.append(g)
-                    if "step" not in state:
-                        state["step"] = (torch.zeros((),
-                                                     dtype=torch.float32,
-                                                     device=p.device))
-                        state["moment1"] = torch.zeros_like(g)
-                        state["moment2"] = torch.zeros_like(g)
-                    moment1.append(state["moment1"])
-                    moment2.append(state["moment2"])
-                    if not isinstance(state["step"], torch.Tensor):
-                        step_tensor = torch.tensor(state["step"],
-                                                   dtype=torch.float32,
-                                                   device=p.device)
-                    else:
-                        step_tensor = state["step"]
-                    state_steps.append(step_tensor)
-
-                self._fused_adamw(
-                    params_with_grads,
-                    grads,
-                    moment1,
-                    moment2,
-                    max_exp_avg_sqs,
-                    state_steps,
-                    amsgrad=False,
-                    beta1=beta1,
-                    beta2=beta2,
-                    lr=lr,
-                    weight_decay=weight_decay,
-                    eps=eps,
-                    maximize=False,
-                )
+                self._step_adamw(group)
 
         return loss

build/torch28-cxx11-cu129-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_811726c_dirty
-ops = torch.ops._optimizer_811726c_dirty
+from . import _optimizer_23d68bb_dirty
+ops = torch.ops._optimizer_23d68bb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_811726c_dirty::{op_name}"
+    return f"_optimizer_23d68bb_dirty::{op_name}"

build/torch28-cxx11-cu129-x86_64-linux/optimizer/{_optimizer_811726c_dirty.abi3.so → _optimizer_23d68bb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b957f60eab442d3ff5a5525d16a1b4b71e8c6be32edb874d9a5681953c61f0c2
-size 1871056
+oid sha256:1cbcd3df518412314d547a86b947998802e488e8aec0f22bf8b59fbc2d1c91e8
+size 1983488

build/torch28-cxx11-cu129-x86_64-linux/optimizer/distributed/utils.py

@@ -0,0 +1,174 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import (Placement, Shard,
+                                                      _StridedShard)
+
+
+def get_slices_of_dtensor(
+    target: DTensor | torch.Tensor,
+    local_rank: int,
+    shard_mesh: DeviceMesh,
+    shard_placements: tuple[Placement],
+) -> tuple[slice]:
+    """
+    Get the slice of local tensor for a given rank from a tensor.
+    Args:
+        target (DTensor | torch.Tensor): The target tensor.
+        rank (int): The local rank of the shard group.
+        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
+        shard_placements (tuple[Placement]): The shard placements.
+    """
+
+    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
+
+    # find the global rank of the local rank in the shard mesh
+    rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank]
+
+    rank_coords = (shard_mesh.mesh == rank).nonzero()
+
+    assert len(rank_coords) == 1
+    rank_coords = tuple(rank_coords[0].tolist())
+
+    assert len(rank_coords) == len(shard_placements)
+
+    # Caution: Assuming replicate-to-shard of the shard mesh goes with
+    # left-to-right sharding. This is ensured by the sorting logic of
+    # construct_shard_mesh function.
+    for i, (rank_coord,
+            placement) in enumerate(zip(rank_coords, shard_placements)):
+        assert isinstance(placement, Shard)
+
+        num_ranks = shard_mesh.mesh.shape[i]
+
+        dim = placement.dim
+        dim_size = (slices[dim].stop - slices[dim].start)
+
+        if dim_size % num_ranks != 0:
+            raise NotImplementedError(
+                f"Dimension size {dim_size} is not divisible "
+                f"by number of ranks {num_ranks} for shard "
+                f"placement on dim {dim}.")
+
+        shard_size = dim_size // num_ranks
+
+        start = slices[dim].start + rank_coord * shard_size
+        end = start + shard_size
+
+        assert start < end <= slices[dim].stop
+
+        slices[dim] = slice(start, end)
+
+    return tuple(slices)
+
+
+_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, ProcessGroup]] = dict()
+
+
+def construct_shard_mesh(
+    placements: tuple[Placement],
+    mesh: DeviceMesh,
+) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
+    """
+    Construct Shard Mesh and Placements for unsharding.
+    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
+    """
+    my_rank = dist.get_rank()
+
+    assert mesh.mesh.device.type == 'cpu'
+
+    # Copy mesh to avoid modifying the original mesh
+    mesh = mesh.mesh.clone()
+
+    # 1. Sort placements. Replicate first, then Shard by dim ascending.
+
+    # For Shard, strided shard comes after regular shard on the same dim
+    # to preserve left-to-right order of replicate-to-shard.
+    # This is because that strided shard is using stride to represent
+    # more fine-grained sharding on the same dim.
+    # Please check the URL below for _StridedShard.
+    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
+
+    def placement_sort_key(
+        placement_with_index: tuple[float, Placement]
+    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
+        index, placement = placement_with_index
+        is_replicate = placement.is_replicate()
+        is_shard = placement.is_shard()
+        is_partial = placement.is_partial()
+
+        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
+        assert not is_partial, "Partial placement is not supported."
+
+        if is_replicate:
+            return (-1.0, 0, index)
+        elif is_shard:
+            if isinstance(placement, _StridedShard):
+                return (placement.dim, 1 / placement.split_factor, index)
+            return (placement.dim, 0, index)
+        else:
+            raise TypeError(f"Unknown placement type: {type(placement)}")
+
+    placements_with_index: list[tuple[int,
+                                      Placement]] = list(enumerate(placements))
+    placements_with_index = sorted(placements_with_index,
+                                   key=placement_sort_key)
+
+    sorted_indices, sorted_placements = zip(*placements_with_index)
+
+    # 2. Permute mesh according to sorted placements.
+    sorted_mesh = mesh.permute(sorted_indices)
+
+    # 3. Collect list of shard meshes by removing replicate dims
+    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
+    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
+    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
+
+    # merge replicate dims
+    # shard_meshes became a list of shard meshes with a length of replicate degree
+    if num_replicates > 0:
+        sorted_mesh = sorted_mesh.flatten(
+            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
+        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
+    else:
+        shard_meshes = [sorted_mesh]
+    shard_placements = sorted_placements[num_replicates:]
+
+    # assume all shard placements are different
+    assert len(shard_placements) == len(set(shard_placements))
+
+    # 4. Construct ProcessGroups
+    # Caution: all groups should be created in the same order in all processes,
+    # even though each process only needs its own group.
+
+    # To use tensor as dict key, convert it to tuple
+    def tensor_to_tuple(t):
+        if isinstance(t, torch.Tensor):
+            t = t.tolist()
+        if isinstance(t, list):
+            return tuple(tensor_to_tuple(x) for x in t)
+        return t
+
+    my_shard_mesh_as_tuple = None
+    for shard_mesh in shard_meshes:
+        assert isinstance(shard_mesh, torch.Tensor)
+        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
+
+        if (my_rank == shard_mesh).any().item():
+            assert my_shard_mesh_as_tuple is None
+            my_shard_mesh_as_tuple = shard_mesh_as_tuple
+
+        # update global cache
+        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
+            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
+            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
+                DeviceMesh(device_type="cuda", mesh=shard_mesh),
+                shard_process_group,
+            )
+
+    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
+        my_shard_mesh_as_tuple]
+
+    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch28-cxx11-cu129-x86_64-linux/optimizer/muon.py

@@ -1,18 +1,24 @@
 import logging
 import math
 import types
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor, Replicate, Shard
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor.placement_types import Placement
 
+from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
 from .matmul_transpose_triton import matmul_transpose_assign
 
 logger = logging.getLogger(__name__)
 
 COMM_DTYPE = torch.bfloat16
+DEFAULT_CHUNK_SIZE_RATIO = 4
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -62,23 +68,39 @@ def _zeropower_via_newtonschulz5(G, steps):
 @dataclass
 class _muon_state:
     # TODO: use Optional
-    worker_rank: int | None = None
+    worker_rank: int
+    process_group: ProcessGroup
+    shard_mesh: DeviceMesh
+    shard_placements: tuple[Placement, ...]
+    name: str
+    qk_clip_state: torch.Tensor | None = None
     gathered_grad: torch.Tensor | None = None
     scattered_u: DTensor | None = None
     computed_u: torch.Tensor | None = None
     gather_event: torch.cuda.Event | None = None
     compute_event: torch.cuda.Event | None = None
     scatter_event: torch.cuda.Event | None = None
-    process_group = None
-    qk_clip_state = None
 
 
-def split_elems_for_src(param, src_rank, num_ranks) -> int:
-    rows = param.shape[0]
-    cols = int(param.numel() // rows)
-    base, rem = divmod(rows, num_ranks)
-    my_rows = base + (1 if src_rank < rem else 0)
-    return my_rows * cols
+def numel_for_rank(
+    param: DTensor,
+    local_rank: int,
+    state: _muon_state,
+) -> int:
+    slices = get_slices_of_dtensor(
+        param,
+        local_rank,
+        state.shard_mesh,
+        state.shard_placements,
+    )
+
+    numel = 1
+    for s, dim in zip(slices, param.shape):
+        start, stop, step = s.indices(dim)
+        length = max(0, (stop - start + (step - 1)) // step)
+        numel *= length
+
+    return numel
 
 
 @torch.no_grad()
@@ -91,8 +113,7 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
         for p in params:
             state = param_to_state[id(p)]
             if rank == state.worker_rank:
-                num_ranks = dist.get_world_size(group=state.process_group)
-                state.gathered_grad = torch.empty(p.grad.numel(),
+                state.gathered_grad = torch.empty(p.shape,
                                                   dtype=COMM_DTYPE,
                                                   device="cuda")
             else:
@@ -121,11 +142,11 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             state = param_to_state[id(p)]
             dst = state.worker_rank
             assert dst < num_ranks
-            shard_elems = split_elems_for_src(p, rank, num_ranks)
+            shard_elems = numel_for_rank(p, rank, state)
             g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            g = g.to_local().to(COMM_DTYPE).contiguous()
             assert g.numel() == shard_elems
-            per_dst[dst].append(g)
+            per_dst[dst].append(g.view(-1))
             send_counts[dst] += shard_elems
 
         assert any(
@@ -148,13 +169,18 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                total += split_elems_for_src(p, src, num_ranks)
+                total += numel_for_rank(p, src, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
         #All2All
+        logger.debug(f"send_buf size: {send_buf.numel()}, "
+                     f"recv_buf size: {recv_buf.numel()}, "
+                     f"recv_counts: {recv_counts}, "
+                     f"send_counts: {send_counts}, "
+                     f"process_group: {str(process_group)}")
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -179,7 +205,6 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
@@ -189,22 +214,28 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                n = split_elems_for_src(p, src, num_ranks)
+
+                # get the slice of the full dtensor corresponding to rank src.
+                slices = get_slices_of_dtensor(state.gathered_grad, src,
+                                               state.shard_mesh,
+                                               state.shard_placements)
+
+                dst = state.gathered_grad[slices]
+                assert dst._base is state.gathered_grad
+
+                n = dst.numel()
                 assert n > 0
 
                 sg = recv_buf.narrow(0, off + inner_off, n)
-                woff = write_offsets[id(p)]
-                dst = state.gathered_grad.narrow(0, woff, n)
+                sg = sg.reshape_as(dst)
                 dst.copy_(sg)
 
-                write_offsets[id(p)] += n
                 inner_off += n
             off += block
 
         for p in params:
             state = param_to_state[id(p)]
             if state.worker_rank == rank:
-                state.gathered_grad = state.gathered_grad.view_as(p)
                 state.gather_event = torch.cuda.Event()
                 state.gather_event.record(comm_stream)
             else:
@@ -277,14 +308,19 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert state.computed_u is not None
 
-                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous()
 
                 offset = 0
                 for dst in range(num_ranks):
-                    n = split_elems_for_src(p, dst, num_ranks)
+                    # get the slice of the full tensor corresponding to rank dst.
+                    slices = get_slices_of_dtensor(u_full, dst,
+                                                   state.shard_mesh,
+                                                   state.shard_placements)
+                    su = u_full[slices].flatten()
+
+                    n = su.numel()
                     assert n > 0
 
-                    su = u_full.narrow(0, offset, n)
                     per_dst[dst].append(su)
                     send_counts[dst] += n
                     offset += n
@@ -313,7 +349,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                total += split_elems_for_src(p, rank, num_ranks)
+                total += numel_for_rank(p, rank, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
@@ -357,7 +393,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                n = split_elems_for_src(p, rank, num_ranks)
+                n = numel_for_rank(p, rank, state)
                 assert n > 0
 
                 flat_local = recv_buf.narrow(0, off + inner_off,
@@ -398,11 +434,23 @@ def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
         state.scattered_u = None
         u_dtensor = None
 
-        scales_full = Muon._compute_scales(p, state.qk_clip_state)
+        scales_full = Muon._compute_scales(
+            p,
+            state.qk_clip_state) if state.qk_clip_state is not None else None
         if scales_full is not None:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            local_rank = dist.get_rank(group=state.process_group)
-            scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank]
+            # Have to slice scales_full among dim 0
+            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
+                                                  state.shard_placements)
+            ratio = p.shape[0] // scales_full.shape[0]
+            scales_slice = slice(
+                None if weight_slices[0].start is None else
+                weight_slices[0].start // ratio,
+                None if weight_slices[0].stop is None else
+                weight_slices[0].stop // ratio,
+                None,
+            )
+
+            scales_local = scales_full[scales_slice]
             scales_local = DTensor.from_local(
                 scales_local,
                 placements=p.placements,
@@ -478,11 +526,11 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
 @dataclass
 class QKClipInfo:
     """Per-parameter dynamic info computed from config + runtime logits."""
-    kind: Optional[str]  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
-    indices: List[int]  # which heads to consider for clipping
+    kind: str | None  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
+    indices: list[int]  # which heads to consider for clipping
     head_dim: int  # from config
     threshold: float  # from config
-    logit: Optional[torch.Tensor]
+    logit: torch.Tensor | None
 
 
 class Muon(torch.optim.Optimizer):
@@ -525,11 +573,16 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
-        overlap_step : How many all2all gather, compute operations are launched in advance
-                       before the corresponding all2all scatter steps begin.
-                       A higher overlap_step increases memory usage but can improve
-                       performance by overlapping communication.
-                       Parallel muon only. 
+        warmup_step : How many all2all gather, compute operations are launched in advance
+                      before the corresponding all2all scatter steps begin.
+                      A higher warmup_step increases memory usage but can improve
+                      performance by overlapping communication.
+                      Parallel muon only.
+        chunk_size : Batch size of parameters to process in each
+                     all2all gather/compute/scatter step.
+                     Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
+        use_distributed_muon: Use distributed muon by Liu et al. (2024).
+                              For testing purpose only.
     """
 
     def __init__(self,
@@ -549,7 +602,9 @@ class Muon(torch.optim.Optimizer):
                      "head_dim": 128,
                      "threshold": 100
                  },
-                 overlap_step=5):
+                 warmup_step=5,
+                 chunk_size=-1,
+                 use_distributed_muon=False):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -579,7 +634,9 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
-        self.overlap_step = overlap_step
+        self.warmup_step = warmup_step
+        self.chunk_size = chunk_size
+        self.use_distributed_muon = use_distributed_muon
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -597,6 +654,12 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
+    def set_rank_once(self, rank):
+        if self.rank is None:
+            self.rank = rank
+        else:
+            assert self.rank == rank
+
     def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
@@ -604,26 +667,13 @@ class Muon(torch.optim.Optimizer):
         assert isinstance(
             p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0), ):
-            # Case for FSDP
-            process_group = p.device_mesh.get_group(mesh_dim=0)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
-        elif p.placements == (Replicate(), Shard(dim=0)):
-            # Case for HSDP
-            process_group = p.device_mesh.get_group(mesh_dim=1)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if self.rank in shard_mesh:
-                    return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
-        else:
-            raise ValueError(f"Unsupported placements ({p.placements}).")
+        shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+            p.placements, p.device_mesh)
+
+        # set rank with the local rank in the shard process group
+        self.set_rank_once(dist.get_rank(group=shard_pg))
+
+        return shard_mesh, shard_pg, shard_placements
 
     def init_state_and_assign_params(self, names, params, group, qk_logits):
         param_to_state = {}
@@ -655,23 +705,32 @@ class Muon(torch.optim.Optimizer):
         ordered_params = list(params_sorted)
 
         round_robin = 0
-        mesh = None
-        shard_mesh = None
-        process_group = None
+        mesh = ordered_params[0].device_mesh
+        placements = ordered_params[0].placements
+
+        shard_mesh, shard_pg, shard_placements = self.get_shard_mesh(
+            ordered_params[0])
+        shard_mesh_flattened = shard_mesh.mesh.flatten()
+        num_ranks = dist.get_world_size(group=shard_pg)
+
         for n, p in zip(ordered_names, ordered_params):
-            if mesh is None:
-                mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p)
-            elif mesh != p.device_mesh:
+            if mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
-            num_ranks = dist.get_world_size(group=process_group)
-            param_to_state[id(p)] = _muon_state()
-            param_to_state[id(
-                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
-            param_to_state[id(p)].process_group = process_group
+            if placements != p.placements:
+                raise ValueError("All parameters must have same placements.")
+
+            worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
+            round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-            param_to_state[id(p)].qk_clip_state = qk_clip_state
-            round_robin = (round_robin + 1) % len(shard_mesh)
+
+            param_to_state[id(p)] = _muon_state(
+                worker_rank=worker_rank,
+                process_group=shard_pg,
+                shard_mesh=shard_mesh,
+                shard_placements=shard_placements,
+                name=n,
+                qk_clip_state=qk_clip_state,
+            )
 
         return param_to_state, ordered_params
 
@@ -705,10 +764,73 @@ class Muon(torch.optim.Optimizer):
 
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
 
-            scales_full = self._compute_scales(p, qk_clip_state)
+            scales_full = self._compute_scales(
+                p, qk_clip_state) if qk_clip_state is not None else None
             if scales_full is not None:
                 Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
 
+    def distributed_muon(
+        self,
+        names: list[str],
+        params: list[torch.nn.Parameter],
+        group: dict[str, Any],
+        lr: float,
+        weight_decay: float,
+        momentum: float,
+        qk_logits: list[torch.Tensor | DTensor] | None,
+    ):
+        """ Implementation of Distributed Muon by Liu et al. """
+        if qk_logits is not None:
+            raise NotImplementedError("QK clipping is not supported yet")
+
+        if isinstance(params[0], DTensor):
+            shard_mesh, _, shard_placements = construct_shard_mesh(
+                placements=params[0].placements,
+                mesh=params[0].device_mesh,
+            )
+
+        for n, p in zip(names, params):
+            g = p.grad
+            if g is None:
+                continue
+            if g.ndim > 2:
+                g = g.view(g.size(0), -1)
+            assert g is not None
+
+            # calc update
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(momentum).add_(g)
+            if group["nesterov"]:
+                g = g.add(buf, alpha=momentum)
+            else:
+                g = buf
+
+            # Gather G
+            if isinstance(p.data, DTensor):
+                g = g.full_tensor()
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                             steps=group["ns_steps"])
+
+            if isinstance(p.data, DTensor):
+                slices = get_slices_of_dtensor(
+                    target=p,
+                    local_rank=dist.get_rank(),
+                    shard_mesh=shard_mesh,
+                    shard_placements=shard_placements,
+                )
+                u_shard = u[slices]
+                u = DTensor.from_local(
+                    u_shard,
+                    device_mesh=p.device_mesh,
+                    placements=p.placements,
+                )
+
+            adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+            Muon._update_p(p, u, lr, adjusted_lr, weight_decay)
+
     def _update_g(self, p, g, group, momentum):
         # calc update
         state = self.state[p]
@@ -727,6 +849,9 @@ class Muon(torch.optim.Optimizer):
         p.data.add_(u, alpha=-adjusted_lr)
 
     def get_qk_clip_info(self, n, qk_logits):
+        if self.clip_config is None:
+            return None
+
         head_dim = self.clip_config.get('head_dim')
         threshold = self.clip_config.get('threshold')
         kind, layer_idx = parse_qk_layer(n)
@@ -737,6 +862,11 @@ class Muon(torch.optim.Optimizer):
             indices_key = 'q_indices' if 'q' in kind else 'k_indices'
             indices = self.clip_config.get(indices_key, []) or []
 
+            if isinstance(logit, DTensor):
+                # In TP settings, qk_logits may be DTensor
+                # We convert it to full tensor here for simplicity
+                logit = logit.full_tensor()
+
         return QKClipInfo(
             kind=kind,
             indices=indices,
@@ -835,22 +965,28 @@ class Muon(torch.optim.Optimizer):
                 _update_param(p, state, lr, adjusted_lr, weight_decay,
                               self.rank, self.compute_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(
-            params[0])].process_group)
+        if self.chunk_size == -1:
+            shard_ranks = dist.get_world_size(param_to_state[id(
+                params[0])].process_group)
+            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+        elif self.chunk_size > 0:
+            chunk_size = self.chunk_size
+        else:
+            raise ValueError("chunk_size must be -1 or a positive integer.")
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        overlap_step = self.overlap_step
-        for i in range(0, overlap_step):
+        warmup_step = self.warmup_step
+        for i in range(0, warmup_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
+            enqueue_computes(i + warmup_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)
@@ -866,7 +1002,7 @@ class Muon(torch.optim.Optimizer):
         amsgrad: bool,
         beta1: float,
         beta2: float,
-        lr: Union[float, torch.Tensor],
+        lr: float | torch.Tensor,
         weight_decay: float,
         eps: float,
         maximize: bool,
@@ -876,10 +1012,10 @@ class Muon(torch.optim.Optimizer):
 
         # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer
         # treating it as a scalar.
-        lr_dict: Optional[DeviceDict] = ({
+        lr_dict: DeviceDict | None = ({
             lr.device: lr
         } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else
-                                         None)
+                                      None)
         grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype(
             [
                 params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
@@ -926,6 +1062,159 @@ class Muon(torch.optim.Optimizer):
                 maximize=maximize,
             )
 
+    def _step_muon(self, group, qk_logits=None):
+        params = group["params"]
+        lr = group["lr"]
+        weight_decay = group["weight_decay"]
+        momentum = group["momentum"]
+        names = group["names"]
+
+        param_dtensors = []
+        param_tensors = []
+        name_dtensors = []
+        name_tensors = []
+
+        if self.use_distributed_muon:
+            self.distributed_muon(names=names,
+                                  params=params,
+                                  group=group,
+                                  lr=lr,
+                                  weight_decay=weight_decay,
+                                  momentum=momentum,
+                                  qk_logits=qk_logits)
+            return
+
+        for n, p in zip(names, params):
+            if p is None or p.grad is None:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(
+                        isinstance(placement, Replicate)
+                        for placement in p.placements):
+                    param_tensors.append(p)
+                    name_tensors.append(n)
+                else:
+                    param_dtensors.append(p)
+                    name_dtensors.append(n)
+            elif isinstance(p.data, torch.Tensor):
+                param_tensors.append(p)
+                name_tensors.append(n)
+            else:
+                raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+        logger.debug(
+            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors"
+        )
+
+        if len(param_dtensors) > 0:
+            if not dist.is_initialized():
+                raise RuntimeError(
+                    "Parallel Muon requires torch.distributed to be initialized."
+                )
+
+            # To support different placements, we group parameters by placements
+            # and run parallel Muon on each group.
+
+            placement_to_params = defaultdict(lambda: ([], []))
+            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
+
+            assert len(name_dtensors) == len(param_dtensors)
+            for n, p in zip(name_dtensors, param_dtensors):
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][0].append(n)
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][1].append(p)
+
+            for _, (names, params) in placement_to_params.items():
+                self.parallel(
+                    names,
+                    params,
+                    group,
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    momentum=momentum,
+                    qk_logits=qk_logits,
+                )
+
+        if len(param_tensors) > 0:
+            self.base(
+                name_tensors,
+                param_tensors,
+                group,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                qk_logits=qk_logits,
+            )
+
+    def _step_adamw_params(self, params, group):
+        params_with_grads = []
+        grads = []
+        moment1 = []
+        moment2 = []
+        max_exp_avg_sqs = []
+        state_steps = []
+        lr = group["lr"]
+        beta1, beta2 = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        weight_decay = group["weight_decay"]
+
+        for p in params:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            params_with_grads.append(p)
+            grads.append(g)
+            if "step" not in state:
+                state["step"] = (torch.zeros((),
+                                             dtype=torch.float32,
+                                             device=p.device))
+                state["moment1"] = torch.zeros_like(g)
+                state["moment2"] = torch.zeros_like(g)
+            moment1.append(state["moment1"])
+            moment2.append(state["moment2"])
+            if not isinstance(state["step"], torch.Tensor):
+                step_tensor = torch.tensor(state["step"],
+                                           dtype=torch.float32,
+                                           device=p.device)
+            else:
+                step_tensor = state["step"]
+            state_steps.append(step_tensor)
+
+        self._fused_adamw(
+            params_with_grads,
+            grads,
+            moment1,
+            moment2,
+            max_exp_avg_sqs,
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+    def _step_adamw(self, group):
+        params = group["params"]
+
+        # group params with it's type and placement
+        placement_to_params: dict[tuple[Placement | type,
+                                        DeviceMesh | None]] = defaultdict(list)
+        for p in params:
+            match p:
+                case DTensor():
+                    placement_to_params[tuple([p.placements,
+                                               p.device_mesh])].append(p)
+                case torch.Tensor():
+                    placement_to_params[tuple([torch.Tensor, None])].append(p)
+
+        for params in placement_to_params.values():
+            self._step_adamw_params(params, group)
+
     def step(self, closure=None, qk_logits=None):
         """Perform a single optimization step.
 
@@ -943,127 +1232,9 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         for group in self.param_groups:
-            params = group["params"]
-
             if group["use_muon"]:
-                ############################
-                #           Muon           #
-                ############################
-                lr = group["lr"]
-                weight_decay = group["weight_decay"]
-                momentum = group["momentum"]
-                names = group["names"]
-
-                param_dtensors = []
-                param_tensors = []
-                name_dtensors = []
-                name_tensors = []
-
-                for n, p in zip(names, params):
-                    if p is None or p.grad is None:
-                        continue
-                    if isinstance(p.data, DTensor):
-                        if all(
-                                isinstance(placement, Replicate)
-                                for placement in p.placements):
-                            param_tensors.append(p)
-                            name_tensors.append(n)
-                        else:
-                            param_dtensors.append(p)
-                            name_dtensors.append(n)
-                    elif isinstance(p.data, torch.Tensor):
-                        param_tensors.append(p)
-                        name_tensors.append(n)
-                    else:
-                        raise TypeError(
-                            f"Unsupported parameter type: {type(p.data)}")
-
-                if self.debug:
-                    print(
-                        f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
-                        flush=True,
-                    )
-
-                if len(param_dtensors) > 0:
-                    if not dist.is_initialized():
-                        raise RuntimeError(
-                            "Parallel Muon requires torch.distributed to be initialized."
-                        )
-
-                    self.parallel(
-                        name_dtensors,
-                        param_dtensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
-                if len(param_tensors) > 0:
-                    self.base(
-                        name_tensors,
-                        param_tensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
+                self._step_muon(group, qk_logits=qk_logits)
             else:
-                ############################
-                #       AdamW backup       #
-                ############################
-
-                params_with_grads = []
-                grads = []
-                moment1 = []
-                moment2 = []
-                max_exp_avg_sqs = []
-                state_steps = []
-                lr = group["lr"]
-                beta1, beta2 = group["adamw_betas"]
-                eps = group["adamw_eps"]
-                weight_decay = group["weight_decay"]
-
-                for p in params:
-                    g = p.grad
-                    if g is None:
-                        continue
-                    state = self.state[p]
-                    params_with_grads.append(p)
-                    grads.append(g)
-                    if "step" not in state:
-                        state["step"] = (torch.zeros((),
-                                                     dtype=torch.float32,
-                                                     device=p.device))
-                        state["moment1"] = torch.zeros_like(g)
-                        state["moment2"] = torch.zeros_like(g)
-                    moment1.append(state["moment1"])
-                    moment2.append(state["moment2"])
-                    if not isinstance(state["step"], torch.Tensor):
-                        step_tensor = torch.tensor(state["step"],
-                                                   dtype=torch.float32,
-                                                   device=p.device)
-                    else:
-                        step_tensor = state["step"]
-                    state_steps.append(step_tensor)
-
-                self._fused_adamw(
-                    params_with_grads,
-                    grads,
-                    moment1,
-                    moment2,
-                    max_exp_avg_sqs,
-                    state_steps,
-                    amsgrad=False,
-                    beta1=beta1,
-                    beta2=beta2,
-                    lr=lr,
-                    weight_decay=weight_decay,
-                    eps=eps,
-                    maximize=False,
-                )
+                self._step_adamw(group)
 
         return loss

build/torch28-cxx11-rocm63-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_811726c_dirty
-ops = torch.ops._optimizer_811726c_dirty
+from . import _optimizer_23d68bb_dirty
+ops = torch.ops._optimizer_23d68bb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_811726c_dirty::{op_name}"
+    return f"_optimizer_23d68bb_dirty::{op_name}"

build/torch28-cxx11-rocm63-x86_64-linux/optimizer/{_optimizer_811726c_dirty.abi3.so → _optimizer_23d68bb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:898ff08457f77c2f6ef504c73570cc87c5c5fd9a144528dbf8af4c03ffc21049
-size 1749232
+oid sha256:8a2999010ee158e13e3ef247e877dfab073b5bde7babefe2b2b5273b760c7ddf
+size 1852152

build/torch28-cxx11-rocm63-x86_64-linux/optimizer/distributed/utils.py

@@ -0,0 +1,174 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import (Placement, Shard,
+                                                      _StridedShard)
+
+
+def get_slices_of_dtensor(
+    target: DTensor | torch.Tensor,
+    local_rank: int,
+    shard_mesh: DeviceMesh,
+    shard_placements: tuple[Placement],
+) -> tuple[slice]:
+    """
+    Get the slice of local tensor for a given rank from a tensor.
+    Args:
+        target (DTensor | torch.Tensor): The target tensor.
+        rank (int): The local rank of the shard group.
+        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
+        shard_placements (tuple[Placement]): The shard placements.
+    """
+
+    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
+
+    # find the global rank of the local rank in the shard mesh
+    rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank]
+
+    rank_coords = (shard_mesh.mesh == rank).nonzero()
+
+    assert len(rank_coords) == 1
+    rank_coords = tuple(rank_coords[0].tolist())
+
+    assert len(rank_coords) == len(shard_placements)
+
+    # Caution: Assuming replicate-to-shard of the shard mesh goes with
+    # left-to-right sharding. This is ensured by the sorting logic of
+    # construct_shard_mesh function.
+    for i, (rank_coord,
+            placement) in enumerate(zip(rank_coords, shard_placements)):
+        assert isinstance(placement, Shard)
+
+        num_ranks = shard_mesh.mesh.shape[i]
+
+        dim = placement.dim
+        dim_size = (slices[dim].stop - slices[dim].start)
+
+        if dim_size % num_ranks != 0:
+            raise NotImplementedError(
+                f"Dimension size {dim_size} is not divisible "
+                f"by number of ranks {num_ranks} for shard "
+                f"placement on dim {dim}.")
+
+        shard_size = dim_size // num_ranks
+
+        start = slices[dim].start + rank_coord * shard_size
+        end = start + shard_size
+
+        assert start < end <= slices[dim].stop
+
+        slices[dim] = slice(start, end)
+
+    return tuple(slices)
+
+
+_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, ProcessGroup]] = dict()
+
+
+def construct_shard_mesh(
+    placements: tuple[Placement],
+    mesh: DeviceMesh,
+) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
+    """
+    Construct Shard Mesh and Placements for unsharding.
+    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
+    """
+    my_rank = dist.get_rank()
+
+    assert mesh.mesh.device.type == 'cpu'
+
+    # Copy mesh to avoid modifying the original mesh
+    mesh = mesh.mesh.clone()
+
+    # 1. Sort placements. Replicate first, then Shard by dim ascending.
+
+    # For Shard, strided shard comes after regular shard on the same dim
+    # to preserve left-to-right order of replicate-to-shard.
+    # This is because that strided shard is using stride to represent
+    # more fine-grained sharding on the same dim.
+    # Please check the URL below for _StridedShard.
+    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
+
+    def placement_sort_key(
+        placement_with_index: tuple[float, Placement]
+    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
+        index, placement = placement_with_index
+        is_replicate = placement.is_replicate()
+        is_shard = placement.is_shard()
+        is_partial = placement.is_partial()
+
+        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
+        assert not is_partial, "Partial placement is not supported."
+
+        if is_replicate:
+            return (-1.0, 0, index)
+        elif is_shard:
+            if isinstance(placement, _StridedShard):
+                return (placement.dim, 1 / placement.split_factor, index)
+            return (placement.dim, 0, index)
+        else:
+            raise TypeError(f"Unknown placement type: {type(placement)}")
+
+    placements_with_index: list[tuple[int,
+                                      Placement]] = list(enumerate(placements))
+    placements_with_index = sorted(placements_with_index,
+                                   key=placement_sort_key)
+
+    sorted_indices, sorted_placements = zip(*placements_with_index)
+
+    # 2. Permute mesh according to sorted placements.
+    sorted_mesh = mesh.permute(sorted_indices)
+
+    # 3. Collect list of shard meshes by removing replicate dims
+    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
+    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
+    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
+
+    # merge replicate dims
+    # shard_meshes became a list of shard meshes with a length of replicate degree
+    if num_replicates > 0:
+        sorted_mesh = sorted_mesh.flatten(
+            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
+        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
+    else:
+        shard_meshes = [sorted_mesh]
+    shard_placements = sorted_placements[num_replicates:]
+
+    # assume all shard placements are different
+    assert len(shard_placements) == len(set(shard_placements))
+
+    # 4. Construct ProcessGroups
+    # Caution: all groups should be created in the same order in all processes,
+    # even though each process only needs its own group.
+
+    # To use tensor as dict key, convert it to tuple
+    def tensor_to_tuple(t):
+        if isinstance(t, torch.Tensor):
+            t = t.tolist()
+        if isinstance(t, list):
+            return tuple(tensor_to_tuple(x) for x in t)
+        return t
+
+    my_shard_mesh_as_tuple = None
+    for shard_mesh in shard_meshes:
+        assert isinstance(shard_mesh, torch.Tensor)
+        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
+
+        if (my_rank == shard_mesh).any().item():
+            assert my_shard_mesh_as_tuple is None
+            my_shard_mesh_as_tuple = shard_mesh_as_tuple
+
+        # update global cache
+        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
+            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
+            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
+                DeviceMesh(device_type="cuda", mesh=shard_mesh),
+                shard_process_group,
+            )
+
+    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
+        my_shard_mesh_as_tuple]
+
+    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch28-cxx11-rocm63-x86_64-linux/optimizer/muon.py

@@ -1,18 +1,24 @@
 import logging
 import math
 import types
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor, Replicate, Shard
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor.placement_types import Placement
 
+from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
 from .matmul_transpose_triton import matmul_transpose_assign
 
 logger = logging.getLogger(__name__)
 
 COMM_DTYPE = torch.bfloat16
+DEFAULT_CHUNK_SIZE_RATIO = 4
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -62,23 +68,39 @@ def _zeropower_via_newtonschulz5(G, steps):
 @dataclass
 class _muon_state:
     # TODO: use Optional
-    worker_rank: int | None = None
+    worker_rank: int
+    process_group: ProcessGroup
+    shard_mesh: DeviceMesh
+    shard_placements: tuple[Placement, ...]
+    name: str
+    qk_clip_state: torch.Tensor | None = None
     gathered_grad: torch.Tensor | None = None
     scattered_u: DTensor | None = None
     computed_u: torch.Tensor | None = None
     gather_event: torch.cuda.Event | None = None
     compute_event: torch.cuda.Event | None = None
     scatter_event: torch.cuda.Event | None = None
-    process_group = None
-    qk_clip_state = None
 
 
-def split_elems_for_src(param, src_rank, num_ranks) -> int:
-    rows = param.shape[0]
-    cols = int(param.numel() // rows)
-    base, rem = divmod(rows, num_ranks)
-    my_rows = base + (1 if src_rank < rem else 0)
-    return my_rows * cols
+def numel_for_rank(
+    param: DTensor,
+    local_rank: int,
+    state: _muon_state,
+) -> int:
+    slices = get_slices_of_dtensor(
+        param,
+        local_rank,
+        state.shard_mesh,
+        state.shard_placements,
+    )
+
+    numel = 1
+    for s, dim in zip(slices, param.shape):
+        start, stop, step = s.indices(dim)
+        length = max(0, (stop - start + (step - 1)) // step)
+        numel *= length
+
+    return numel
 
 
 @torch.no_grad()
@@ -91,8 +113,7 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
         for p in params:
             state = param_to_state[id(p)]
             if rank == state.worker_rank:
-                num_ranks = dist.get_world_size(group=state.process_group)
-                state.gathered_grad = torch.empty(p.grad.numel(),
+                state.gathered_grad = torch.empty(p.shape,
                                                   dtype=COMM_DTYPE,
                                                   device="cuda")
             else:
@@ -121,11 +142,11 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             state = param_to_state[id(p)]
             dst = state.worker_rank
             assert dst < num_ranks
-            shard_elems = split_elems_for_src(p, rank, num_ranks)
+            shard_elems = numel_for_rank(p, rank, state)
             g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            g = g.to_local().to(COMM_DTYPE).contiguous()
             assert g.numel() == shard_elems
-            per_dst[dst].append(g)
+            per_dst[dst].append(g.view(-1))
             send_counts[dst] += shard_elems
 
         assert any(
@@ -148,13 +169,18 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                total += split_elems_for_src(p, src, num_ranks)
+                total += numel_for_rank(p, src, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
         #All2All
+        logger.debug(f"send_buf size: {send_buf.numel()}, "
+                     f"recv_buf size: {recv_buf.numel()}, "
+                     f"recv_counts: {recv_counts}, "
+                     f"send_counts: {send_counts}, "
+                     f"process_group: {str(process_group)}")
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -179,7 +205,6 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
@@ -189,22 +214,28 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                n = split_elems_for_src(p, src, num_ranks)
+
+                # get the slice of the full dtensor corresponding to rank src.
+                slices = get_slices_of_dtensor(state.gathered_grad, src,
+                                               state.shard_mesh,
+                                               state.shard_placements)
+
+                dst = state.gathered_grad[slices]
+                assert dst._base is state.gathered_grad
+
+                n = dst.numel()
                 assert n > 0
 
                 sg = recv_buf.narrow(0, off + inner_off, n)
-                woff = write_offsets[id(p)]
-                dst = state.gathered_grad.narrow(0, woff, n)
+                sg = sg.reshape_as(dst)
                 dst.copy_(sg)
 
-                write_offsets[id(p)] += n
                 inner_off += n
             off += block
 
         for p in params:
             state = param_to_state[id(p)]
             if state.worker_rank == rank:
-                state.gathered_grad = state.gathered_grad.view_as(p)
                 state.gather_event = torch.cuda.Event()
                 state.gather_event.record(comm_stream)
             else:
@@ -277,14 +308,19 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert state.computed_u is not None
 
-                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous()
 
                 offset = 0
                 for dst in range(num_ranks):
-                    n = split_elems_for_src(p, dst, num_ranks)
+                    # get the slice of the full tensor corresponding to rank dst.
+                    slices = get_slices_of_dtensor(u_full, dst,
+                                                   state.shard_mesh,
+                                                   state.shard_placements)
+                    su = u_full[slices].flatten()
+
+                    n = su.numel()
                     assert n > 0
 
-                    su = u_full.narrow(0, offset, n)
                     per_dst[dst].append(su)
                     send_counts[dst] += n
                     offset += n
@@ -313,7 +349,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                total += split_elems_for_src(p, rank, num_ranks)
+                total += numel_for_rank(p, rank, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
@@ -357,7 +393,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                n = split_elems_for_src(p, rank, num_ranks)
+                n = numel_for_rank(p, rank, state)
                 assert n > 0
 
                 flat_local = recv_buf.narrow(0, off + inner_off,
@@ -398,11 +434,23 @@ def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
         state.scattered_u = None
         u_dtensor = None
 
-        scales_full = Muon._compute_scales(p, state.qk_clip_state)
+        scales_full = Muon._compute_scales(
+            p,
+            state.qk_clip_state) if state.qk_clip_state is not None else None
         if scales_full is not None:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            local_rank = dist.get_rank(group=state.process_group)
-            scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank]
+            # Have to slice scales_full among dim 0
+            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
+                                                  state.shard_placements)
+            ratio = p.shape[0] // scales_full.shape[0]
+            scales_slice = slice(
+                None if weight_slices[0].start is None else
+                weight_slices[0].start // ratio,
+                None if weight_slices[0].stop is None else
+                weight_slices[0].stop // ratio,
+                None,
+            )
+
+            scales_local = scales_full[scales_slice]
             scales_local = DTensor.from_local(
                 scales_local,
                 placements=p.placements,
@@ -478,11 +526,11 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
 @dataclass
 class QKClipInfo:
     """Per-parameter dynamic info computed from config + runtime logits."""
-    kind: Optional[str]  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
-    indices: List[int]  # which heads to consider for clipping
+    kind: str | None  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
+    indices: list[int]  # which heads to consider for clipping
     head_dim: int  # from config
     threshold: float  # from config
-    logit: Optional[torch.Tensor]
+    logit: torch.Tensor | None
 
 
 class Muon(torch.optim.Optimizer):
@@ -525,11 +573,16 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
-        overlap_step : How many all2all gather, compute operations are launched in advance
-                       before the corresponding all2all scatter steps begin.
-                       A higher overlap_step increases memory usage but can improve
-                       performance by overlapping communication.
-                       Parallel muon only. 
+        warmup_step : How many all2all gather, compute operations are launched in advance
+                      before the corresponding all2all scatter steps begin.
+                      A higher warmup_step increases memory usage but can improve
+                      performance by overlapping communication.
+                      Parallel muon only.
+        chunk_size : Batch size of parameters to process in each
+                     all2all gather/compute/scatter step.
+                     Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
+        use_distributed_muon: Use distributed muon by Liu et al. (2024).
+                              For testing purpose only.
     """
 
     def __init__(self,
@@ -549,7 +602,9 @@ class Muon(torch.optim.Optimizer):
                      "head_dim": 128,
                      "threshold": 100
                  },
-                 overlap_step=5):
+                 warmup_step=5,
+                 chunk_size=-1,
+                 use_distributed_muon=False):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -579,7 +634,9 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
-        self.overlap_step = overlap_step
+        self.warmup_step = warmup_step
+        self.chunk_size = chunk_size
+        self.use_distributed_muon = use_distributed_muon
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -597,6 +654,12 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
+    def set_rank_once(self, rank):
+        if self.rank is None:
+            self.rank = rank
+        else:
+            assert self.rank == rank
+
     def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
@@ -604,26 +667,13 @@ class Muon(torch.optim.Optimizer):
         assert isinstance(
             p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0), ):
-            # Case for FSDP
-            process_group = p.device_mesh.get_group(mesh_dim=0)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
-        elif p.placements == (Replicate(), Shard(dim=0)):
-            # Case for HSDP
-            process_group = p.device_mesh.get_group(mesh_dim=1)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if self.rank in shard_mesh:
-                    return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
-        else:
-            raise ValueError(f"Unsupported placements ({p.placements}).")
+        shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+            p.placements, p.device_mesh)
+
+        # set rank with the local rank in the shard process group
+        self.set_rank_once(dist.get_rank(group=shard_pg))
+
+        return shard_mesh, shard_pg, shard_placements
 
     def init_state_and_assign_params(self, names, params, group, qk_logits):
         param_to_state = {}
@@ -655,23 +705,32 @@ class Muon(torch.optim.Optimizer):
         ordered_params = list(params_sorted)
 
         round_robin = 0
-        mesh = None
-        shard_mesh = None
-        process_group = None
+        mesh = ordered_params[0].device_mesh
+        placements = ordered_params[0].placements
+
+        shard_mesh, shard_pg, shard_placements = self.get_shard_mesh(
+            ordered_params[0])
+        shard_mesh_flattened = shard_mesh.mesh.flatten()
+        num_ranks = dist.get_world_size(group=shard_pg)
+
         for n, p in zip(ordered_names, ordered_params):
-            if mesh is None:
-                mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p)
-            elif mesh != p.device_mesh:
+            if mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
-            num_ranks = dist.get_world_size(group=process_group)
-            param_to_state[id(p)] = _muon_state()
-            param_to_state[id(
-                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
-            param_to_state[id(p)].process_group = process_group
+            if placements != p.placements:
+                raise ValueError("All parameters must have same placements.")
+
+            worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
+            round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-            param_to_state[id(p)].qk_clip_state = qk_clip_state
-            round_robin = (round_robin + 1) % len(shard_mesh)
+
+            param_to_state[id(p)] = _muon_state(
+                worker_rank=worker_rank,
+                process_group=shard_pg,
+                shard_mesh=shard_mesh,
+                shard_placements=shard_placements,
+                name=n,
+                qk_clip_state=qk_clip_state,
+            )
 
         return param_to_state, ordered_params
 
@@ -705,10 +764,73 @@ class Muon(torch.optim.Optimizer):
 
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
 
-            scales_full = self._compute_scales(p, qk_clip_state)
+            scales_full = self._compute_scales(
+                p, qk_clip_state) if qk_clip_state is not None else None
             if scales_full is not None:
                 Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
 
+    def distributed_muon(
+        self,
+        names: list[str],
+        params: list[torch.nn.Parameter],
+        group: dict[str, Any],
+        lr: float,
+        weight_decay: float,
+        momentum: float,
+        qk_logits: list[torch.Tensor | DTensor] | None,
+    ):
+        """ Implementation of Distributed Muon by Liu et al. """
+        if qk_logits is not None:
+            raise NotImplementedError("QK clipping is not supported yet")
+
+        if isinstance(params[0], DTensor):
+            shard_mesh, _, shard_placements = construct_shard_mesh(
+                placements=params[0].placements,
+                mesh=params[0].device_mesh,
+            )
+
+        for n, p in zip(names, params):
+            g = p.grad
+            if g is None:
+                continue
+            if g.ndim > 2:
+                g = g.view(g.size(0), -1)
+            assert g is not None
+
+            # calc update
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(momentum).add_(g)
+            if group["nesterov"]:
+                g = g.add(buf, alpha=momentum)
+            else:
+                g = buf
+
+            # Gather G
+            if isinstance(p.data, DTensor):
+                g = g.full_tensor()
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                             steps=group["ns_steps"])
+
+            if isinstance(p.data, DTensor):
+                slices = get_slices_of_dtensor(
+                    target=p,
+                    local_rank=dist.get_rank(),
+                    shard_mesh=shard_mesh,
+                    shard_placements=shard_placements,
+                )
+                u_shard = u[slices]
+                u = DTensor.from_local(
+                    u_shard,
+                    device_mesh=p.device_mesh,
+                    placements=p.placements,
+                )
+
+            adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+            Muon._update_p(p, u, lr, adjusted_lr, weight_decay)
+
     def _update_g(self, p, g, group, momentum):
         # calc update
         state = self.state[p]
@@ -727,6 +849,9 @@ class Muon(torch.optim.Optimizer):
         p.data.add_(u, alpha=-adjusted_lr)
 
     def get_qk_clip_info(self, n, qk_logits):
+        if self.clip_config is None:
+            return None
+
         head_dim = self.clip_config.get('head_dim')
         threshold = self.clip_config.get('threshold')
         kind, layer_idx = parse_qk_layer(n)
@@ -737,6 +862,11 @@ class Muon(torch.optim.Optimizer):
             indices_key = 'q_indices' if 'q' in kind else 'k_indices'
             indices = self.clip_config.get(indices_key, []) or []
 
+            if isinstance(logit, DTensor):
+                # In TP settings, qk_logits may be DTensor
+                # We convert it to full tensor here for simplicity
+                logit = logit.full_tensor()
+
         return QKClipInfo(
             kind=kind,
             indices=indices,
@@ -835,22 +965,28 @@ class Muon(torch.optim.Optimizer):
                 _update_param(p, state, lr, adjusted_lr, weight_decay,
                               self.rank, self.compute_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(
-            params[0])].process_group)
+        if self.chunk_size == -1:
+            shard_ranks = dist.get_world_size(param_to_state[id(
+                params[0])].process_group)
+            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+        elif self.chunk_size > 0:
+            chunk_size = self.chunk_size
+        else:
+            raise ValueError("chunk_size must be -1 or a positive integer.")
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        overlap_step = self.overlap_step
-        for i in range(0, overlap_step):
+        warmup_step = self.warmup_step
+        for i in range(0, warmup_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
+            enqueue_computes(i + warmup_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)
@@ -866,7 +1002,7 @@ class Muon(torch.optim.Optimizer):
         amsgrad: bool,
         beta1: float,
         beta2: float,
-        lr: Union[float, torch.Tensor],
+        lr: float | torch.Tensor,
         weight_decay: float,
         eps: float,
         maximize: bool,
@@ -876,10 +1012,10 @@ class Muon(torch.optim.Optimizer):
 
         # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer
         # treating it as a scalar.
-        lr_dict: Optional[DeviceDict] = ({
+        lr_dict: DeviceDict | None = ({
             lr.device: lr
         } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else
-                                         None)
+                                      None)
         grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype(
             [
                 params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
@@ -926,6 +1062,159 @@ class Muon(torch.optim.Optimizer):
                 maximize=maximize,
             )
 
+    def _step_muon(self, group, qk_logits=None):
+        params = group["params"]
+        lr = group["lr"]
+        weight_decay = group["weight_decay"]
+        momentum = group["momentum"]
+        names = group["names"]
+
+        param_dtensors = []
+        param_tensors = []
+        name_dtensors = []
+        name_tensors = []
+
+        if self.use_distributed_muon:
+            self.distributed_muon(names=names,
+                                  params=params,
+                                  group=group,
+                                  lr=lr,
+                                  weight_decay=weight_decay,
+                                  momentum=momentum,
+                                  qk_logits=qk_logits)
+            return
+
+        for n, p in zip(names, params):
+            if p is None or p.grad is None:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(
+                        isinstance(placement, Replicate)
+                        for placement in p.placements):
+                    param_tensors.append(p)
+                    name_tensors.append(n)
+                else:
+                    param_dtensors.append(p)
+                    name_dtensors.append(n)
+            elif isinstance(p.data, torch.Tensor):
+                param_tensors.append(p)
+                name_tensors.append(n)
+            else:
+                raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+        logger.debug(
+            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors"
+        )
+
+        if len(param_dtensors) > 0:
+            if not dist.is_initialized():
+                raise RuntimeError(
+                    "Parallel Muon requires torch.distributed to be initialized."
+                )
+
+            # To support different placements, we group parameters by placements
+            # and run parallel Muon on each group.
+
+            placement_to_params = defaultdict(lambda: ([], []))
+            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
+
+            assert len(name_dtensors) == len(param_dtensors)
+            for n, p in zip(name_dtensors, param_dtensors):
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][0].append(n)
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][1].append(p)
+
+            for _, (names, params) in placement_to_params.items():
+                self.parallel(
+                    names,
+                    params,
+                    group,
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    momentum=momentum,
+                    qk_logits=qk_logits,
+                )
+
+        if len(param_tensors) > 0:
+            self.base(
+                name_tensors,
+                param_tensors,
+                group,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                qk_logits=qk_logits,
+            )
+
+    def _step_adamw_params(self, params, group):
+        params_with_grads = []
+        grads = []
+        moment1 = []
+        moment2 = []
+        max_exp_avg_sqs = []
+        state_steps = []
+        lr = group["lr"]
+        beta1, beta2 = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        weight_decay = group["weight_decay"]
+
+        for p in params:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            params_with_grads.append(p)
+            grads.append(g)
+            if "step" not in state:
+                state["step"] = (torch.zeros((),
+                                             dtype=torch.float32,
+                                             device=p.device))
+                state["moment1"] = torch.zeros_like(g)
+                state["moment2"] = torch.zeros_like(g)
+            moment1.append(state["moment1"])
+            moment2.append(state["moment2"])
+            if not isinstance(state["step"], torch.Tensor):
+                step_tensor = torch.tensor(state["step"],
+                                           dtype=torch.float32,
+                                           device=p.device)
+            else:
+                step_tensor = state["step"]
+            state_steps.append(step_tensor)
+
+        self._fused_adamw(
+            params_with_grads,
+            grads,
+            moment1,
+            moment2,
+            max_exp_avg_sqs,
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+    def _step_adamw(self, group):
+        params = group["params"]
+
+        # group params with it's type and placement
+        placement_to_params: dict[tuple[Placement | type,
+                                        DeviceMesh | None]] = defaultdict(list)
+        for p in params:
+            match p:
+                case DTensor():
+                    placement_to_params[tuple([p.placements,
+                                               p.device_mesh])].append(p)
+                case torch.Tensor():
+                    placement_to_params[tuple([torch.Tensor, None])].append(p)
+
+        for params in placement_to_params.values():
+            self._step_adamw_params(params, group)
+
     def step(self, closure=None, qk_logits=None):
         """Perform a single optimization step.
 
@@ -943,127 +1232,9 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         for group in self.param_groups:
-            params = group["params"]
-
             if group["use_muon"]:
-                ############################
-                #           Muon           #
-                ############################
-                lr = group["lr"]
-                weight_decay = group["weight_decay"]
-                momentum = group["momentum"]
-                names = group["names"]
-
-                param_dtensors = []
-                param_tensors = []
-                name_dtensors = []
-                name_tensors = []
-
-                for n, p in zip(names, params):
-                    if p is None or p.grad is None:
-                        continue
-                    if isinstance(p.data, DTensor):
-                        if all(
-                                isinstance(placement, Replicate)
-                                for placement in p.placements):
-                            param_tensors.append(p)
-                            name_tensors.append(n)
-                        else:
-                            param_dtensors.append(p)
-                            name_dtensors.append(n)
-                    elif isinstance(p.data, torch.Tensor):
-                        param_tensors.append(p)
-                        name_tensors.append(n)
-                    else:
-                        raise TypeError(
-                            f"Unsupported parameter type: {type(p.data)}")
-
-                if self.debug:
-                    print(
-                        f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
-                        flush=True,
-                    )
-
-                if len(param_dtensors) > 0:
-                    if not dist.is_initialized():
-                        raise RuntimeError(
-                            "Parallel Muon requires torch.distributed to be initialized."
-                        )
-
-                    self.parallel(
-                        name_dtensors,
-                        param_dtensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
-                if len(param_tensors) > 0:
-                    self.base(
-                        name_tensors,
-                        param_tensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
+                self._step_muon(group, qk_logits=qk_logits)
             else:
-                ############################
-                #       AdamW backup       #
-                ############################
-
-                params_with_grads = []
-                grads = []
-                moment1 = []
-                moment2 = []
-                max_exp_avg_sqs = []
-                state_steps = []
-                lr = group["lr"]
-                beta1, beta2 = group["adamw_betas"]
-                eps = group["adamw_eps"]
-                weight_decay = group["weight_decay"]
-
-                for p in params:
-                    g = p.grad
-                    if g is None:
-                        continue
-                    state = self.state[p]
-                    params_with_grads.append(p)
-                    grads.append(g)
-                    if "step" not in state:
-                        state["step"] = (torch.zeros((),
-                                                     dtype=torch.float32,
-                                                     device=p.device))
-                        state["moment1"] = torch.zeros_like(g)
-                        state["moment2"] = torch.zeros_like(g)
-                    moment1.append(state["moment1"])
-                    moment2.append(state["moment2"])
-                    if not isinstance(state["step"], torch.Tensor):
-                        step_tensor = torch.tensor(state["step"],
-                                                   dtype=torch.float32,
-                                                   device=p.device)
-                    else:
-                        step_tensor = state["step"]
-                    state_steps.append(step_tensor)
-
-                self._fused_adamw(
-                    params_with_grads,
-                    grads,
-                    moment1,
-                    moment2,
-                    max_exp_avg_sqs,
-                    state_steps,
-                    amsgrad=False,
-                    beta1=beta1,
-                    beta2=beta2,
-                    lr=lr,
-                    weight_decay=weight_decay,
-                    eps=eps,
-                    maximize=False,
-                )
+                self._step_adamw(group)
 
         return loss

build/torch28-cxx11-rocm64-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_811726c_dirty
-ops = torch.ops._optimizer_811726c_dirty
+from . import _optimizer_23d68bb_dirty
+ops = torch.ops._optimizer_23d68bb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_811726c_dirty::{op_name}"
+    return f"_optimizer_23d68bb_dirty::{op_name}"

build/torch28-cxx11-rocm64-x86_64-linux/optimizer/_optimizer_23d68bb_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55f869cf4220f2033d4e499da522da46794a682495c2b688dbcac0ec89135cf4
+size 1852240

build/torch28-cxx11-rocm64-x86_64-linux/optimizer/_optimizer_811726c_dirty.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:72d100180fd73094f7b1c6e765eb4a77f103ad392fdee571687cb0c66d304177
-size 1749320

build/torch28-cxx11-rocm64-x86_64-linux/optimizer/distributed/utils.py

@@ -0,0 +1,174 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import (Placement, Shard,
+                                                      _StridedShard)
+
+
+def get_slices_of_dtensor(
+    target: DTensor | torch.Tensor,
+    local_rank: int,
+    shard_mesh: DeviceMesh,
+    shard_placements: tuple[Placement],
+) -> tuple[slice]:
+    """
+    Get the slice of local tensor for a given rank from a tensor.
+    Args:
+        target (DTensor | torch.Tensor): The target tensor.
+        rank (int): The local rank of the shard group.
+        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
+        shard_placements (tuple[Placement]): The shard placements.
+    """
+
+    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
+
+    # find the global rank of the local rank in the shard mesh
+    rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank]
+
+    rank_coords = (shard_mesh.mesh == rank).nonzero()
+
+    assert len(rank_coords) == 1
+    rank_coords = tuple(rank_coords[0].tolist())
+
+    assert len(rank_coords) == len(shard_placements)
+
+    # Caution: Assuming replicate-to-shard of the shard mesh goes with
+    # left-to-right sharding. This is ensured by the sorting logic of
+    # construct_shard_mesh function.
+    for i, (rank_coord,
+            placement) in enumerate(zip(rank_coords, shard_placements)):
+        assert isinstance(placement, Shard)
+
+        num_ranks = shard_mesh.mesh.shape[i]
+
+        dim = placement.dim
+        dim_size = (slices[dim].stop - slices[dim].start)
+
+        if dim_size % num_ranks != 0:
+            raise NotImplementedError(
+                f"Dimension size {dim_size} is not divisible "
+                f"by number of ranks {num_ranks} for shard "
+                f"placement on dim {dim}.")
+
+        shard_size = dim_size // num_ranks
+
+        start = slices[dim].start + rank_coord * shard_size
+        end = start + shard_size
+
+        assert start < end <= slices[dim].stop
+
+        slices[dim] = slice(start, end)
+
+    return tuple(slices)
+
+
+_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, ProcessGroup]] = dict()
+
+
+def construct_shard_mesh(
+    placements: tuple[Placement],
+    mesh: DeviceMesh,
+) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
+    """
+    Construct Shard Mesh and Placements for unsharding.
+    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
+    """
+    my_rank = dist.get_rank()
+
+    assert mesh.mesh.device.type == 'cpu'
+
+    # Copy mesh to avoid modifying the original mesh
+    mesh = mesh.mesh.clone()
+
+    # 1. Sort placements. Replicate first, then Shard by dim ascending.
+
+    # For Shard, strided shard comes after regular shard on the same dim
+    # to preserve left-to-right order of replicate-to-shard.
+    # This is because that strided shard is using stride to represent
+    # more fine-grained sharding on the same dim.
+    # Please check the URL below for _StridedShard.
+    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
+
+    def placement_sort_key(
+        placement_with_index: tuple[float, Placement]
+    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
+        index, placement = placement_with_index
+        is_replicate = placement.is_replicate()
+        is_shard = placement.is_shard()
+        is_partial = placement.is_partial()
+
+        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
+        assert not is_partial, "Partial placement is not supported."
+
+        if is_replicate:
+            return (-1.0, 0, index)
+        elif is_shard:
+            if isinstance(placement, _StridedShard):
+                return (placement.dim, 1 / placement.split_factor, index)
+            return (placement.dim, 0, index)
+        else:
+            raise TypeError(f"Unknown placement type: {type(placement)}")
+
+    placements_with_index: list[tuple[int,
+                                      Placement]] = list(enumerate(placements))
+    placements_with_index = sorted(placements_with_index,
+                                   key=placement_sort_key)
+
+    sorted_indices, sorted_placements = zip(*placements_with_index)
+
+    # 2. Permute mesh according to sorted placements.
+    sorted_mesh = mesh.permute(sorted_indices)
+
+    # 3. Collect list of shard meshes by removing replicate dims
+    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
+    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
+    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
+
+    # merge replicate dims
+    # shard_meshes became a list of shard meshes with a length of replicate degree
+    if num_replicates > 0:
+        sorted_mesh = sorted_mesh.flatten(
+            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
+        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
+    else:
+        shard_meshes = [sorted_mesh]
+    shard_placements = sorted_placements[num_replicates:]
+
+    # assume all shard placements are different
+    assert len(shard_placements) == len(set(shard_placements))
+
+    # 4. Construct ProcessGroups
+    # Caution: all groups should be created in the same order in all processes,
+    # even though each process only needs its own group.
+
+    # To use tensor as dict key, convert it to tuple
+    def tensor_to_tuple(t):
+        if isinstance(t, torch.Tensor):
+            t = t.tolist()
+        if isinstance(t, list):
+            return tuple(tensor_to_tuple(x) for x in t)
+        return t
+
+    my_shard_mesh_as_tuple = None
+    for shard_mesh in shard_meshes:
+        assert isinstance(shard_mesh, torch.Tensor)
+        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
+
+        if (my_rank == shard_mesh).any().item():
+            assert my_shard_mesh_as_tuple is None
+            my_shard_mesh_as_tuple = shard_mesh_as_tuple
+
+        # update global cache
+        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
+            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
+            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
+                DeviceMesh(device_type="cuda", mesh=shard_mesh),
+                shard_process_group,
+            )
+
+    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
+        my_shard_mesh_as_tuple]
+
+    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch28-cxx11-rocm64-x86_64-linux/optimizer/muon.py

@@ -1,18 +1,24 @@
 import logging
 import math
 import types
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor, Replicate, Shard
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor.placement_types import Placement
 
+from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
 from .matmul_transpose_triton import matmul_transpose_assign
 
 logger = logging.getLogger(__name__)
 
 COMM_DTYPE = torch.bfloat16
+DEFAULT_CHUNK_SIZE_RATIO = 4
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -62,23 +68,39 @@ def _zeropower_via_newtonschulz5(G, steps):
 @dataclass
 class _muon_state:
     # TODO: use Optional
-    worker_rank: int | None = None
+    worker_rank: int
+    process_group: ProcessGroup
+    shard_mesh: DeviceMesh
+    shard_placements: tuple[Placement, ...]
+    name: str
+    qk_clip_state: torch.Tensor | None = None
     gathered_grad: torch.Tensor | None = None
     scattered_u: DTensor | None = None
     computed_u: torch.Tensor | None = None
     gather_event: torch.cuda.Event | None = None
     compute_event: torch.cuda.Event | None = None
     scatter_event: torch.cuda.Event | None = None
-    process_group = None
-    qk_clip_state = None
 
 
-def split_elems_for_src(param, src_rank, num_ranks) -> int:
-    rows = param.shape[0]
-    cols = int(param.numel() // rows)
-    base, rem = divmod(rows, num_ranks)
-    my_rows = base + (1 if src_rank < rem else 0)
-    return my_rows * cols
+def numel_for_rank(
+    param: DTensor,
+    local_rank: int,
+    state: _muon_state,
+) -> int:
+    slices = get_slices_of_dtensor(
+        param,
+        local_rank,
+        state.shard_mesh,
+        state.shard_placements,
+    )
+
+    numel = 1
+    for s, dim in zip(slices, param.shape):
+        start, stop, step = s.indices(dim)
+        length = max(0, (stop - start + (step - 1)) // step)
+        numel *= length
+
+    return numel
 
 
 @torch.no_grad()
@@ -91,8 +113,7 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
         for p in params:
             state = param_to_state[id(p)]
             if rank == state.worker_rank:
-                num_ranks = dist.get_world_size(group=state.process_group)
-                state.gathered_grad = torch.empty(p.grad.numel(),
+                state.gathered_grad = torch.empty(p.shape,
                                                   dtype=COMM_DTYPE,
                                                   device="cuda")
             else:
@@ -121,11 +142,11 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             state = param_to_state[id(p)]
             dst = state.worker_rank
             assert dst < num_ranks
-            shard_elems = split_elems_for_src(p, rank, num_ranks)
+            shard_elems = numel_for_rank(p, rank, state)
             g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            g = g.to_local().to(COMM_DTYPE).contiguous()
             assert g.numel() == shard_elems
-            per_dst[dst].append(g)
+            per_dst[dst].append(g.view(-1))
             send_counts[dst] += shard_elems
 
         assert any(
@@ -148,13 +169,18 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                total += split_elems_for_src(p, src, num_ranks)
+                total += numel_for_rank(p, src, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
         #All2All
+        logger.debug(f"send_buf size: {send_buf.numel()}, "
+                     f"recv_buf size: {recv_buf.numel()}, "
+                     f"recv_counts: {recv_counts}, "
+                     f"send_counts: {send_counts}, "
+                     f"process_group: {str(process_group)}")
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -179,7 +205,6 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
@@ -189,22 +214,28 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                n = split_elems_for_src(p, src, num_ranks)
+
+                # get the slice of the full dtensor corresponding to rank src.
+                slices = get_slices_of_dtensor(state.gathered_grad, src,
+                                               state.shard_mesh,
+                                               state.shard_placements)
+
+                dst = state.gathered_grad[slices]
+                assert dst._base is state.gathered_grad
+
+                n = dst.numel()
                 assert n > 0
 
                 sg = recv_buf.narrow(0, off + inner_off, n)
-                woff = write_offsets[id(p)]
-                dst = state.gathered_grad.narrow(0, woff, n)
+                sg = sg.reshape_as(dst)
                 dst.copy_(sg)
 
-                write_offsets[id(p)] += n
                 inner_off += n
             off += block
 
         for p in params:
             state = param_to_state[id(p)]
             if state.worker_rank == rank:
-                state.gathered_grad = state.gathered_grad.view_as(p)
                 state.gather_event = torch.cuda.Event()
                 state.gather_event.record(comm_stream)
             else:
@@ -277,14 +308,19 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert state.computed_u is not None
 
-                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous()
 
                 offset = 0
                 for dst in range(num_ranks):
-                    n = split_elems_for_src(p, dst, num_ranks)
+                    # get the slice of the full tensor corresponding to rank dst.
+                    slices = get_slices_of_dtensor(u_full, dst,
+                                                   state.shard_mesh,
+                                                   state.shard_placements)
+                    su = u_full[slices].flatten()
+
+                    n = su.numel()
                     assert n > 0
 
-                    su = u_full.narrow(0, offset, n)
                     per_dst[dst].append(su)
                     send_counts[dst] += n
                     offset += n
@@ -313,7 +349,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                total += split_elems_for_src(p, rank, num_ranks)
+                total += numel_for_rank(p, rank, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
@@ -357,7 +393,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                n = split_elems_for_src(p, rank, num_ranks)
+                n = numel_for_rank(p, rank, state)
                 assert n > 0
 
                 flat_local = recv_buf.narrow(0, off + inner_off,
@@ -398,11 +434,23 @@ def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
         state.scattered_u = None
         u_dtensor = None
 
-        scales_full = Muon._compute_scales(p, state.qk_clip_state)
+        scales_full = Muon._compute_scales(
+            p,
+            state.qk_clip_state) if state.qk_clip_state is not None else None
         if scales_full is not None:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            local_rank = dist.get_rank(group=state.process_group)
-            scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank]
+            # Have to slice scales_full among dim 0
+            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
+                                                  state.shard_placements)
+            ratio = p.shape[0] // scales_full.shape[0]
+            scales_slice = slice(
+                None if weight_slices[0].start is None else
+                weight_slices[0].start // ratio,
+                None if weight_slices[0].stop is None else
+                weight_slices[0].stop // ratio,
+                None,
+            )
+
+            scales_local = scales_full[scales_slice]
             scales_local = DTensor.from_local(
                 scales_local,
                 placements=p.placements,
@@ -478,11 +526,11 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
 @dataclass
 class QKClipInfo:
     """Per-parameter dynamic info computed from config + runtime logits."""
-    kind: Optional[str]  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
-    indices: List[int]  # which heads to consider for clipping
+    kind: str | None  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
+    indices: list[int]  # which heads to consider for clipping
     head_dim: int  # from config
     threshold: float  # from config
-    logit: Optional[torch.Tensor]
+    logit: torch.Tensor | None
 
 
 class Muon(torch.optim.Optimizer):
@@ -525,11 +573,16 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
-        overlap_step : How many all2all gather, compute operations are launched in advance
-                       before the corresponding all2all scatter steps begin.
-                       A higher overlap_step increases memory usage but can improve
-                       performance by overlapping communication.
-                       Parallel muon only. 
+        warmup_step : How many all2all gather, compute operations are launched in advance
+                      before the corresponding all2all scatter steps begin.
+                      A higher warmup_step increases memory usage but can improve
+                      performance by overlapping communication.
+                      Parallel muon only.
+        chunk_size : Batch size of parameters to process in each
+                     all2all gather/compute/scatter step.
+                     Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
+        use_distributed_muon: Use distributed muon by Liu et al. (2024).
+                              For testing purpose only.
     """
 
     def __init__(self,
@@ -549,7 +602,9 @@ class Muon(torch.optim.Optimizer):
                      "head_dim": 128,
                      "threshold": 100
                  },
-                 overlap_step=5):
+                 warmup_step=5,
+                 chunk_size=-1,
+                 use_distributed_muon=False):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -579,7 +634,9 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
-        self.overlap_step = overlap_step
+        self.warmup_step = warmup_step
+        self.chunk_size = chunk_size
+        self.use_distributed_muon = use_distributed_muon
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -597,6 +654,12 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
+    def set_rank_once(self, rank):
+        if self.rank is None:
+            self.rank = rank
+        else:
+            assert self.rank == rank
+
     def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
@@ -604,26 +667,13 @@ class Muon(torch.optim.Optimizer):
         assert isinstance(
             p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0), ):
-            # Case for FSDP
-            process_group = p.device_mesh.get_group(mesh_dim=0)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
-        elif p.placements == (Replicate(), Shard(dim=0)):
-            # Case for HSDP
-            process_group = p.device_mesh.get_group(mesh_dim=1)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if self.rank in shard_mesh:
-                    return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
-        else:
-            raise ValueError(f"Unsupported placements ({p.placements}).")
+        shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+            p.placements, p.device_mesh)
+
+        # set rank with the local rank in the shard process group
+        self.set_rank_once(dist.get_rank(group=shard_pg))
+
+        return shard_mesh, shard_pg, shard_placements
 
     def init_state_and_assign_params(self, names, params, group, qk_logits):
         param_to_state = {}
@@ -655,23 +705,32 @@ class Muon(torch.optim.Optimizer):
         ordered_params = list(params_sorted)
 
         round_robin = 0
-        mesh = None
-        shard_mesh = None
-        process_group = None
+        mesh = ordered_params[0].device_mesh
+        placements = ordered_params[0].placements
+
+        shard_mesh, shard_pg, shard_placements = self.get_shard_mesh(
+            ordered_params[0])
+        shard_mesh_flattened = shard_mesh.mesh.flatten()
+        num_ranks = dist.get_world_size(group=shard_pg)
+
         for n, p in zip(ordered_names, ordered_params):
-            if mesh is None:
-                mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p)
-            elif mesh != p.device_mesh:
+            if mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
-            num_ranks = dist.get_world_size(group=process_group)
-            param_to_state[id(p)] = _muon_state()
-            param_to_state[id(
-                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
-            param_to_state[id(p)].process_group = process_group
+            if placements != p.placements:
+                raise ValueError("All parameters must have same placements.")
+
+            worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
+            round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-            param_to_state[id(p)].qk_clip_state = qk_clip_state
-            round_robin = (round_robin + 1) % len(shard_mesh)
+
+            param_to_state[id(p)] = _muon_state(
+                worker_rank=worker_rank,
+                process_group=shard_pg,
+                shard_mesh=shard_mesh,
+                shard_placements=shard_placements,
+                name=n,
+                qk_clip_state=qk_clip_state,
+            )
 
         return param_to_state, ordered_params
 
@@ -705,10 +764,73 @@ class Muon(torch.optim.Optimizer):
 
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
 
-            scales_full = self._compute_scales(p, qk_clip_state)
+            scales_full = self._compute_scales(
+                p, qk_clip_state) if qk_clip_state is not None else None
             if scales_full is not None:
                 Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
 
+    def distributed_muon(
+        self,
+        names: list[str],
+        params: list[torch.nn.Parameter],
+        group: dict[str, Any],
+        lr: float,
+        weight_decay: float,
+        momentum: float,
+        qk_logits: list[torch.Tensor | DTensor] | None,
+    ):
+        """ Implementation of Distributed Muon by Liu et al. """
+        if qk_logits is not None:
+            raise NotImplementedError("QK clipping is not supported yet")
+
+        if isinstance(params[0], DTensor):
+            shard_mesh, _, shard_placements = construct_shard_mesh(
+                placements=params[0].placements,
+                mesh=params[0].device_mesh,
+            )
+
+        for n, p in zip(names, params):
+            g = p.grad
+            if g is None:
+                continue
+            if g.ndim > 2:
+                g = g.view(g.size(0), -1)
+            assert g is not None
+
+            # calc update
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(momentum).add_(g)
+            if group["nesterov"]:
+                g = g.add(buf, alpha=momentum)
+            else:
+                g = buf
+
+            # Gather G
+            if isinstance(p.data, DTensor):
+                g = g.full_tensor()
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                             steps=group["ns_steps"])
+
+            if isinstance(p.data, DTensor):
+                slices = get_slices_of_dtensor(
+                    target=p,
+                    local_rank=dist.get_rank(),
+                    shard_mesh=shard_mesh,
+                    shard_placements=shard_placements,
+                )
+                u_shard = u[slices]
+                u = DTensor.from_local(
+                    u_shard,
+                    device_mesh=p.device_mesh,
+                    placements=p.placements,
+                )
+
+            adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+            Muon._update_p(p, u, lr, adjusted_lr, weight_decay)
+
     def _update_g(self, p, g, group, momentum):
         # calc update
         state = self.state[p]
@@ -727,6 +849,9 @@ class Muon(torch.optim.Optimizer):
         p.data.add_(u, alpha=-adjusted_lr)
 
     def get_qk_clip_info(self, n, qk_logits):
+        if self.clip_config is None:
+            return None
+
         head_dim = self.clip_config.get('head_dim')
         threshold = self.clip_config.get('threshold')
         kind, layer_idx = parse_qk_layer(n)
@@ -737,6 +862,11 @@ class Muon(torch.optim.Optimizer):
             indices_key = 'q_indices' if 'q' in kind else 'k_indices'
             indices = self.clip_config.get(indices_key, []) or []
 
+            if isinstance(logit, DTensor):
+                # In TP settings, qk_logits may be DTensor
+                # We convert it to full tensor here for simplicity
+                logit = logit.full_tensor()
+
         return QKClipInfo(
             kind=kind,
             indices=indices,
@@ -835,22 +965,28 @@ class Muon(torch.optim.Optimizer):
                 _update_param(p, state, lr, adjusted_lr, weight_decay,
                               self.rank, self.compute_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(
-            params[0])].process_group)
+        if self.chunk_size == -1:
+            shard_ranks = dist.get_world_size(param_to_state[id(
+                params[0])].process_group)
+            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+        elif self.chunk_size > 0:
+            chunk_size = self.chunk_size
+        else:
+            raise ValueError("chunk_size must be -1 or a positive integer.")
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        overlap_step = self.overlap_step
-        for i in range(0, overlap_step):
+        warmup_step = self.warmup_step
+        for i in range(0, warmup_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
+            enqueue_computes(i + warmup_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)
@@ -866,7 +1002,7 @@ class Muon(torch.optim.Optimizer):
         amsgrad: bool,
         beta1: float,
         beta2: float,
-        lr: Union[float, torch.Tensor],
+        lr: float | torch.Tensor,
         weight_decay: float,
         eps: float,
         maximize: bool,
@@ -876,10 +1012,10 @@ class Muon(torch.optim.Optimizer):
 
         # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer
         # treating it as a scalar.
-        lr_dict: Optional[DeviceDict] = ({
+        lr_dict: DeviceDict | None = ({
             lr.device: lr
         } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else
-                                         None)
+                                      None)
         grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype(
             [
                 params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
@@ -926,6 +1062,159 @@ class Muon(torch.optim.Optimizer):
                 maximize=maximize,
             )
 
+    def _step_muon(self, group, qk_logits=None):
+        params = group["params"]
+        lr = group["lr"]
+        weight_decay = group["weight_decay"]
+        momentum = group["momentum"]
+        names = group["names"]
+
+        param_dtensors = []
+        param_tensors = []
+        name_dtensors = []
+        name_tensors = []
+
+        if self.use_distributed_muon:
+            self.distributed_muon(names=names,
+                                  params=params,
+                                  group=group,
+                                  lr=lr,
+                                  weight_decay=weight_decay,
+                                  momentum=momentum,
+                                  qk_logits=qk_logits)
+            return
+
+        for n, p in zip(names, params):
+            if p is None or p.grad is None:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(
+                        isinstance(placement, Replicate)
+                        for placement in p.placements):
+                    param_tensors.append(p)
+                    name_tensors.append(n)
+                else:
+                    param_dtensors.append(p)
+                    name_dtensors.append(n)
+            elif isinstance(p.data, torch.Tensor):
+                param_tensors.append(p)
+                name_tensors.append(n)
+            else:
+                raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+        logger.debug(
+            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors"
+        )
+
+        if len(param_dtensors) > 0:
+            if not dist.is_initialized():
+                raise RuntimeError(
+                    "Parallel Muon requires torch.distributed to be initialized."
+                )
+
+            # To support different placements, we group parameters by placements
+            # and run parallel Muon on each group.
+
+            placement_to_params = defaultdict(lambda: ([], []))
+            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
+
+            assert len(name_dtensors) == len(param_dtensors)
+            for n, p in zip(name_dtensors, param_dtensors):
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][0].append(n)
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][1].append(p)
+
+            for _, (names, params) in placement_to_params.items():
+                self.parallel(
+                    names,
+                    params,
+                    group,
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    momentum=momentum,
+                    qk_logits=qk_logits,
+                )
+
+        if len(param_tensors) > 0:
+            self.base(
+                name_tensors,
+                param_tensors,
+                group,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                qk_logits=qk_logits,
+            )
+
+    def _step_adamw_params(self, params, group):
+        params_with_grads = []
+        grads = []
+        moment1 = []
+        moment2 = []
+        max_exp_avg_sqs = []
+        state_steps = []
+        lr = group["lr"]
+        beta1, beta2 = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        weight_decay = group["weight_decay"]
+
+        for p in params:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            params_with_grads.append(p)
+            grads.append(g)
+            if "step" not in state:
+                state["step"] = (torch.zeros((),
+                                             dtype=torch.float32,
+                                             device=p.device))
+                state["moment1"] = torch.zeros_like(g)
+                state["moment2"] = torch.zeros_like(g)
+            moment1.append(state["moment1"])
+            moment2.append(state["moment2"])
+            if not isinstance(state["step"], torch.Tensor):
+                step_tensor = torch.tensor(state["step"],
+                                           dtype=torch.float32,
+                                           device=p.device)
+            else:
+                step_tensor = state["step"]
+            state_steps.append(step_tensor)
+
+        self._fused_adamw(
+            params_with_grads,
+            grads,
+            moment1,
+            moment2,
+            max_exp_avg_sqs,
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+    def _step_adamw(self, group):
+        params = group["params"]
+
+        # group params with it's type and placement
+        placement_to_params: dict[tuple[Placement | type,
+                                        DeviceMesh | None]] = defaultdict(list)
+        for p in params:
+            match p:
+                case DTensor():
+                    placement_to_params[tuple([p.placements,
+                                               p.device_mesh])].append(p)
+                case torch.Tensor():
+                    placement_to_params[tuple([torch.Tensor, None])].append(p)
+
+        for params in placement_to_params.values():
+            self._step_adamw_params(params, group)
+
     def step(self, closure=None, qk_logits=None):
         """Perform a single optimization step.
 
@@ -943,127 +1232,9 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         for group in self.param_groups:
-            params = group["params"]
-
             if group["use_muon"]:
-                ############################
-                #           Muon           #
-                ############################
-                lr = group["lr"]
-                weight_decay = group["weight_decay"]
-                momentum = group["momentum"]
-                names = group["names"]
-
-                param_dtensors = []
-                param_tensors = []
-                name_dtensors = []
-                name_tensors = []
-
-                for n, p in zip(names, params):
-                    if p is None or p.grad is None:
-                        continue
-                    if isinstance(p.data, DTensor):
-                        if all(
-                                isinstance(placement, Replicate)
-                                for placement in p.placements):
-                            param_tensors.append(p)
-                            name_tensors.append(n)
-                        else:
-                            param_dtensors.append(p)
-                            name_dtensors.append(n)
-                    elif isinstance(p.data, torch.Tensor):
-                        param_tensors.append(p)
-                        name_tensors.append(n)
-                    else:
-                        raise TypeError(
-                            f"Unsupported parameter type: {type(p.data)}")
-
-                if self.debug:
-                    print(
-                        f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
-                        flush=True,
-                    )
-
-                if len(param_dtensors) > 0:
-                    if not dist.is_initialized():
-                        raise RuntimeError(
-                            "Parallel Muon requires torch.distributed to be initialized."
-                        )
-
-                    self.parallel(
-                        name_dtensors,
-                        param_dtensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
-                if len(param_tensors) > 0:
-                    self.base(
-                        name_tensors,
-                        param_tensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
+                self._step_muon(group, qk_logits=qk_logits)
             else:
-                ############################
-                #       AdamW backup       #
-                ############################
-
-                params_with_grads = []
-                grads = []
-                moment1 = []
-                moment2 = []
-                max_exp_avg_sqs = []
-                state_steps = []
-                lr = group["lr"]
-                beta1, beta2 = group["adamw_betas"]
-                eps = group["adamw_eps"]
-                weight_decay = group["weight_decay"]
-
-                for p in params:
-                    g = p.grad
-                    if g is None:
-                        continue
-                    state = self.state[p]
-                    params_with_grads.append(p)
-                    grads.append(g)
-                    if "step" not in state:
-                        state["step"] = (torch.zeros((),
-                                                     dtype=torch.float32,
-                                                     device=p.device))
-                        state["moment1"] = torch.zeros_like(g)
-                        state["moment2"] = torch.zeros_like(g)
-                    moment1.append(state["moment1"])
-                    moment2.append(state["moment2"])
-                    if not isinstance(state["step"], torch.Tensor):
-                        step_tensor = torch.tensor(state["step"],
-                                                   dtype=torch.float32,
-                                                   device=p.device)
-                    else:
-                        step_tensor = state["step"]
-                    state_steps.append(step_tensor)
-
-                self._fused_adamw(
-                    params_with_grads,
-                    grads,
-                    moment1,
-                    moment2,
-                    max_exp_avg_sqs,
-                    state_steps,
-                    amsgrad=False,
-                    beta1=beta1,
-                    beta2=beta2,
-                    lr=lr,
-                    weight_decay=weight_decay,
-                    eps=eps,
-                    maximize=False,
-                )
+                self._step_adamw(group)
 
         return loss

build/torch29-cxx11-cu126-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_811726c_dirty
-ops = torch.ops._optimizer_811726c_dirty
+from . import _optimizer_23d68bb_dirty
+ops = torch.ops._optimizer_23d68bb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_811726c_dirty::{op_name}"
+    return f"_optimizer_23d68bb_dirty::{op_name}"

build/torch29-cxx11-cu126-x86_64-linux/optimizer/_optimizer_23d68bb_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca847c77875fc19f211a4c8ac217e9664b46c6862aa3234c270aacfea519d0f5
+size 1924376

build/torch29-cxx11-cu126-x86_64-linux/optimizer/_optimizer_811726c_dirty.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:87c8e75ead1c831dabfce1abbd7c100aa72c9b2988dfc0e1554216ca8005267c
-size 1816064

build/torch29-cxx11-cu126-x86_64-linux/optimizer/distributed/utils.py

@@ -0,0 +1,174 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import (Placement, Shard,
+                                                      _StridedShard)
+
+
+def get_slices_of_dtensor(
+    target: DTensor | torch.Tensor,
+    local_rank: int,
+    shard_mesh: DeviceMesh,
+    shard_placements: tuple[Placement],
+) -> tuple[slice]:
+    """
+    Get the slice of local tensor for a given rank from a tensor.
+    Args:
+        target (DTensor | torch.Tensor): The target tensor.
+        rank (int): The local rank of the shard group.
+        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
+        shard_placements (tuple[Placement]): The shard placements.
+    """
+
+    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
+
+    # find the global rank of the local rank in the shard mesh
+    rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank]
+
+    rank_coords = (shard_mesh.mesh == rank).nonzero()
+
+    assert len(rank_coords) == 1
+    rank_coords = tuple(rank_coords[0].tolist())
+
+    assert len(rank_coords) == len(shard_placements)
+
+    # Caution: Assuming replicate-to-shard of the shard mesh goes with
+    # left-to-right sharding. This is ensured by the sorting logic of
+    # construct_shard_mesh function.
+    for i, (rank_coord,
+            placement) in enumerate(zip(rank_coords, shard_placements)):
+        assert isinstance(placement, Shard)
+
+        num_ranks = shard_mesh.mesh.shape[i]
+
+        dim = placement.dim
+        dim_size = (slices[dim].stop - slices[dim].start)
+
+        if dim_size % num_ranks != 0:
+            raise NotImplementedError(
+                f"Dimension size {dim_size} is not divisible "
+                f"by number of ranks {num_ranks} for shard "
+                f"placement on dim {dim}.")
+
+        shard_size = dim_size // num_ranks
+
+        start = slices[dim].start + rank_coord * shard_size
+        end = start + shard_size
+
+        assert start < end <= slices[dim].stop
+
+        slices[dim] = slice(start, end)
+
+    return tuple(slices)
+
+
+_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, ProcessGroup]] = dict()
+
+
+def construct_shard_mesh(
+    placements: tuple[Placement],
+    mesh: DeviceMesh,
+) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
+    """
+    Construct Shard Mesh and Placements for unsharding.
+    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
+    """
+    my_rank = dist.get_rank()
+
+    assert mesh.mesh.device.type == 'cpu'
+
+    # Copy mesh to avoid modifying the original mesh
+    mesh = mesh.mesh.clone()
+
+    # 1. Sort placements. Replicate first, then Shard by dim ascending.
+
+    # For Shard, strided shard comes after regular shard on the same dim
+    # to preserve left-to-right order of replicate-to-shard.
+    # This is because that strided shard is using stride to represent
+    # more fine-grained sharding on the same dim.
+    # Please check the URL below for _StridedShard.
+    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
+
+    def placement_sort_key(
+        placement_with_index: tuple[float, Placement]
+    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
+        index, placement = placement_with_index
+        is_replicate = placement.is_replicate()
+        is_shard = placement.is_shard()
+        is_partial = placement.is_partial()
+
+        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
+        assert not is_partial, "Partial placement is not supported."
+
+        if is_replicate:
+            return (-1.0, 0, index)
+        elif is_shard:
+            if isinstance(placement, _StridedShard):
+                return (placement.dim, 1 / placement.split_factor, index)
+            return (placement.dim, 0, index)
+        else:
+            raise TypeError(f"Unknown placement type: {type(placement)}")
+
+    placements_with_index: list[tuple[int,
+                                      Placement]] = list(enumerate(placements))
+    placements_with_index = sorted(placements_with_index,
+                                   key=placement_sort_key)
+
+    sorted_indices, sorted_placements = zip(*placements_with_index)
+
+    # 2. Permute mesh according to sorted placements.
+    sorted_mesh = mesh.permute(sorted_indices)
+
+    # 3. Collect list of shard meshes by removing replicate dims
+    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
+    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
+    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
+
+    # merge replicate dims
+    # shard_meshes became a list of shard meshes with a length of replicate degree
+    if num_replicates > 0:
+        sorted_mesh = sorted_mesh.flatten(
+            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
+        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
+    else:
+        shard_meshes = [sorted_mesh]
+    shard_placements = sorted_placements[num_replicates:]
+
+    # assume all shard placements are different
+    assert len(shard_placements) == len(set(shard_placements))
+
+    # 4. Construct ProcessGroups
+    # Caution: all groups should be created in the same order in all processes,
+    # even though each process only needs its own group.
+
+    # To use tensor as dict key, convert it to tuple
+    def tensor_to_tuple(t):
+        if isinstance(t, torch.Tensor):
+            t = t.tolist()
+        if isinstance(t, list):
+            return tuple(tensor_to_tuple(x) for x in t)
+        return t
+
+    my_shard_mesh_as_tuple = None
+    for shard_mesh in shard_meshes:
+        assert isinstance(shard_mesh, torch.Tensor)
+        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
+
+        if (my_rank == shard_mesh).any().item():
+            assert my_shard_mesh_as_tuple is None
+            my_shard_mesh_as_tuple = shard_mesh_as_tuple
+
+        # update global cache
+        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
+            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
+            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
+                DeviceMesh(device_type="cuda", mesh=shard_mesh),
+                shard_process_group,
+            )
+
+    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
+        my_shard_mesh_as_tuple]
+
+    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch29-cxx11-cu126-x86_64-linux/optimizer/muon.py

@@ -1,18 +1,24 @@
 import logging
 import math
 import types
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor, Replicate, Shard
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor.placement_types import Placement
 
+from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
 from .matmul_transpose_triton import matmul_transpose_assign
 
 logger = logging.getLogger(__name__)
 
 COMM_DTYPE = torch.bfloat16
+DEFAULT_CHUNK_SIZE_RATIO = 4
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -62,23 +68,39 @@ def _zeropower_via_newtonschulz5(G, steps):
 @dataclass
 class _muon_state:
     # TODO: use Optional
-    worker_rank: int | None = None
+    worker_rank: int
+    process_group: ProcessGroup
+    shard_mesh: DeviceMesh
+    shard_placements: tuple[Placement, ...]
+    name: str
+    qk_clip_state: torch.Tensor | None = None
     gathered_grad: torch.Tensor | None = None
     scattered_u: DTensor | None = None
     computed_u: torch.Tensor | None = None
     gather_event: torch.cuda.Event | None = None
     compute_event: torch.cuda.Event | None = None
     scatter_event: torch.cuda.Event | None = None
-    process_group = None
-    qk_clip_state = None
 
 
-def split_elems_for_src(param, src_rank, num_ranks) -> int:
-    rows = param.shape[0]
-    cols = int(param.numel() // rows)
-    base, rem = divmod(rows, num_ranks)
-    my_rows = base + (1 if src_rank < rem else 0)
-    return my_rows * cols
+def numel_for_rank(
+    param: DTensor,
+    local_rank: int,
+    state: _muon_state,
+) -> int:
+    slices = get_slices_of_dtensor(
+        param,
+        local_rank,
+        state.shard_mesh,
+        state.shard_placements,
+    )
+
+    numel = 1
+    for s, dim in zip(slices, param.shape):
+        start, stop, step = s.indices(dim)
+        length = max(0, (stop - start + (step - 1)) // step)
+        numel *= length
+
+    return numel
 
 
 @torch.no_grad()
@@ -91,8 +113,7 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
         for p in params:
             state = param_to_state[id(p)]
             if rank == state.worker_rank:
-                num_ranks = dist.get_world_size(group=state.process_group)
-                state.gathered_grad = torch.empty(p.grad.numel(),
+                state.gathered_grad = torch.empty(p.shape,
                                                   dtype=COMM_DTYPE,
                                                   device="cuda")
             else:
@@ -121,11 +142,11 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             state = param_to_state[id(p)]
             dst = state.worker_rank
             assert dst < num_ranks
-            shard_elems = split_elems_for_src(p, rank, num_ranks)
+            shard_elems = numel_for_rank(p, rank, state)
             g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            g = g.to_local().to(COMM_DTYPE).contiguous()
             assert g.numel() == shard_elems
-            per_dst[dst].append(g)
+            per_dst[dst].append(g.view(-1))
             send_counts[dst] += shard_elems
 
         assert any(
@@ -148,13 +169,18 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                total += split_elems_for_src(p, src, num_ranks)
+                total += numel_for_rank(p, src, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
         #All2All
+        logger.debug(f"send_buf size: {send_buf.numel()}, "
+                     f"recv_buf size: {recv_buf.numel()}, "
+                     f"recv_counts: {recv_counts}, "
+                     f"send_counts: {send_counts}, "
+                     f"process_group: {str(process_group)}")
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -179,7 +205,6 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
@@ -189,22 +214,28 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                n = split_elems_for_src(p, src, num_ranks)
+
+                # get the slice of the full dtensor corresponding to rank src.
+                slices = get_slices_of_dtensor(state.gathered_grad, src,
+                                               state.shard_mesh,
+                                               state.shard_placements)
+
+                dst = state.gathered_grad[slices]
+                assert dst._base is state.gathered_grad
+
+                n = dst.numel()
                 assert n > 0
 
                 sg = recv_buf.narrow(0, off + inner_off, n)
-                woff = write_offsets[id(p)]
-                dst = state.gathered_grad.narrow(0, woff, n)
+                sg = sg.reshape_as(dst)
                 dst.copy_(sg)
 
-                write_offsets[id(p)] += n
                 inner_off += n
             off += block
 
         for p in params:
             state = param_to_state[id(p)]
             if state.worker_rank == rank:
-                state.gathered_grad = state.gathered_grad.view_as(p)
                 state.gather_event = torch.cuda.Event()
                 state.gather_event.record(comm_stream)
             else:
@@ -277,14 +308,19 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert state.computed_u is not None
 
-                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous()
 
                 offset = 0
                 for dst in range(num_ranks):
-                    n = split_elems_for_src(p, dst, num_ranks)
+                    # get the slice of the full tensor corresponding to rank dst.
+                    slices = get_slices_of_dtensor(u_full, dst,
+                                                   state.shard_mesh,
+                                                   state.shard_placements)
+                    su = u_full[slices].flatten()
+
+                    n = su.numel()
                     assert n > 0
 
-                    su = u_full.narrow(0, offset, n)
                     per_dst[dst].append(su)
                     send_counts[dst] += n
                     offset += n
@@ -313,7 +349,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                total += split_elems_for_src(p, rank, num_ranks)
+                total += numel_for_rank(p, rank, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
@@ -357,7 +393,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                n = split_elems_for_src(p, rank, num_ranks)
+                n = numel_for_rank(p, rank, state)
                 assert n > 0
 
                 flat_local = recv_buf.narrow(0, off + inner_off,
@@ -398,11 +434,23 @@ def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
         state.scattered_u = None
         u_dtensor = None
 
-        scales_full = Muon._compute_scales(p, state.qk_clip_state)
+        scales_full = Muon._compute_scales(
+            p,
+            state.qk_clip_state) if state.qk_clip_state is not None else None
         if scales_full is not None:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            local_rank = dist.get_rank(group=state.process_group)
-            scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank]
+            # Have to slice scales_full among dim 0
+            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
+                                                  state.shard_placements)
+            ratio = p.shape[0] // scales_full.shape[0]
+            scales_slice = slice(
+                None if weight_slices[0].start is None else
+                weight_slices[0].start // ratio,
+                None if weight_slices[0].stop is None else
+                weight_slices[0].stop // ratio,
+                None,
+            )
+
+            scales_local = scales_full[scales_slice]
             scales_local = DTensor.from_local(
                 scales_local,
                 placements=p.placements,
@@ -478,11 +526,11 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
 @dataclass
 class QKClipInfo:
     """Per-parameter dynamic info computed from config + runtime logits."""
-    kind: Optional[str]  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
-    indices: List[int]  # which heads to consider for clipping
+    kind: str | None  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
+    indices: list[int]  # which heads to consider for clipping
     head_dim: int  # from config
     threshold: float  # from config
-    logit: Optional[torch.Tensor]
+    logit: torch.Tensor | None
 
 
 class Muon(torch.optim.Optimizer):
@@ -525,11 +573,16 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
-        overlap_step : How many all2all gather, compute operations are launched in advance
-                       before the corresponding all2all scatter steps begin.
-                       A higher overlap_step increases memory usage but can improve
-                       performance by overlapping communication.
-                       Parallel muon only. 
+        warmup_step : How many all2all gather, compute operations are launched in advance
+                      before the corresponding all2all scatter steps begin.
+                      A higher warmup_step increases memory usage but can improve
+                      performance by overlapping communication.
+                      Parallel muon only.
+        chunk_size : Batch size of parameters to process in each
+                     all2all gather/compute/scatter step.
+                     Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
+        use_distributed_muon: Use distributed muon by Liu et al. (2024).
+                              For testing purpose only.
     """
 
     def __init__(self,
@@ -549,7 +602,9 @@ class Muon(torch.optim.Optimizer):
                      "head_dim": 128,
                      "threshold": 100
                  },
-                 overlap_step=5):
+                 warmup_step=5,
+                 chunk_size=-1,
+                 use_distributed_muon=False):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -579,7 +634,9 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
-        self.overlap_step = overlap_step
+        self.warmup_step = warmup_step
+        self.chunk_size = chunk_size
+        self.use_distributed_muon = use_distributed_muon
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -597,6 +654,12 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
+    def set_rank_once(self, rank):
+        if self.rank is None:
+            self.rank = rank
+        else:
+            assert self.rank == rank
+
     def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
@@ -604,26 +667,13 @@ class Muon(torch.optim.Optimizer):
         assert isinstance(
             p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0), ):
-            # Case for FSDP
-            process_group = p.device_mesh.get_group(mesh_dim=0)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
-        elif p.placements == (Replicate(), Shard(dim=0)):
-            # Case for HSDP
-            process_group = p.device_mesh.get_group(mesh_dim=1)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if self.rank in shard_mesh:
-                    return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
-        else:
-            raise ValueError(f"Unsupported placements ({p.placements}).")
+        shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+            p.placements, p.device_mesh)
+
+        # set rank with the local rank in the shard process group
+        self.set_rank_once(dist.get_rank(group=shard_pg))
+
+        return shard_mesh, shard_pg, shard_placements
 
     def init_state_and_assign_params(self, names, params, group, qk_logits):
         param_to_state = {}
@@ -655,23 +705,32 @@ class Muon(torch.optim.Optimizer):
         ordered_params = list(params_sorted)
 
         round_robin = 0
-        mesh = None
-        shard_mesh = None
-        process_group = None
+        mesh = ordered_params[0].device_mesh
+        placements = ordered_params[0].placements
+
+        shard_mesh, shard_pg, shard_placements = self.get_shard_mesh(
+            ordered_params[0])
+        shard_mesh_flattened = shard_mesh.mesh.flatten()
+        num_ranks = dist.get_world_size(group=shard_pg)
+
         for n, p in zip(ordered_names, ordered_params):
-            if mesh is None:
-                mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p)
-            elif mesh != p.device_mesh:
+            if mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
-            num_ranks = dist.get_world_size(group=process_group)
-            param_to_state[id(p)] = _muon_state()
-            param_to_state[id(
-                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
-            param_to_state[id(p)].process_group = process_group
+            if placements != p.placements:
+                raise ValueError("All parameters must have same placements.")
+
+            worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
+            round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-            param_to_state[id(p)].qk_clip_state = qk_clip_state
-            round_robin = (round_robin + 1) % len(shard_mesh)
+
+            param_to_state[id(p)] = _muon_state(
+                worker_rank=worker_rank,
+                process_group=shard_pg,
+                shard_mesh=shard_mesh,
+                shard_placements=shard_placements,
+                name=n,
+                qk_clip_state=qk_clip_state,
+            )
 
         return param_to_state, ordered_params
 
@@ -705,10 +764,73 @@ class Muon(torch.optim.Optimizer):
 
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
 
-            scales_full = self._compute_scales(p, qk_clip_state)
+            scales_full = self._compute_scales(
+                p, qk_clip_state) if qk_clip_state is not None else None
             if scales_full is not None:
                 Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
 
+    def distributed_muon(
+        self,
+        names: list[str],
+        params: list[torch.nn.Parameter],
+        group: dict[str, Any],
+        lr: float,
+        weight_decay: float,
+        momentum: float,
+        qk_logits: list[torch.Tensor | DTensor] | None,
+    ):
+        """ Implementation of Distributed Muon by Liu et al. """
+        if qk_logits is not None:
+            raise NotImplementedError("QK clipping is not supported yet")
+
+        if isinstance(params[0], DTensor):
+            shard_mesh, _, shard_placements = construct_shard_mesh(
+                placements=params[0].placements,
+                mesh=params[0].device_mesh,
+            )
+
+        for n, p in zip(names, params):
+            g = p.grad
+            if g is None:
+                continue
+            if g.ndim > 2:
+                g = g.view(g.size(0), -1)
+            assert g is not None
+
+            # calc update
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(momentum).add_(g)
+            if group["nesterov"]:
+                g = g.add(buf, alpha=momentum)
+            else:
+                g = buf
+
+            # Gather G
+            if isinstance(p.data, DTensor):
+                g = g.full_tensor()
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                             steps=group["ns_steps"])
+
+            if isinstance(p.data, DTensor):
+                slices = get_slices_of_dtensor(
+                    target=p,
+                    local_rank=dist.get_rank(),
+                    shard_mesh=shard_mesh,
+                    shard_placements=shard_placements,
+                )
+                u_shard = u[slices]
+                u = DTensor.from_local(
+                    u_shard,
+                    device_mesh=p.device_mesh,
+                    placements=p.placements,
+                )
+
+            adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+            Muon._update_p(p, u, lr, adjusted_lr, weight_decay)
+
     def _update_g(self, p, g, group, momentum):
         # calc update
         state = self.state[p]
@@ -727,6 +849,9 @@ class Muon(torch.optim.Optimizer):
         p.data.add_(u, alpha=-adjusted_lr)
 
     def get_qk_clip_info(self, n, qk_logits):
+        if self.clip_config is None:
+            return None
+
         head_dim = self.clip_config.get('head_dim')
         threshold = self.clip_config.get('threshold')
         kind, layer_idx = parse_qk_layer(n)
@@ -737,6 +862,11 @@ class Muon(torch.optim.Optimizer):
             indices_key = 'q_indices' if 'q' in kind else 'k_indices'
             indices = self.clip_config.get(indices_key, []) or []
 
+            if isinstance(logit, DTensor):
+                # In TP settings, qk_logits may be DTensor
+                # We convert it to full tensor here for simplicity
+                logit = logit.full_tensor()
+
         return QKClipInfo(
             kind=kind,
             indices=indices,
@@ -835,22 +965,28 @@ class Muon(torch.optim.Optimizer):
                 _update_param(p, state, lr, adjusted_lr, weight_decay,
                               self.rank, self.compute_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(
-            params[0])].process_group)
+        if self.chunk_size == -1:
+            shard_ranks = dist.get_world_size(param_to_state[id(
+                params[0])].process_group)
+            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+        elif self.chunk_size > 0:
+            chunk_size = self.chunk_size
+        else:
+            raise ValueError("chunk_size must be -1 or a positive integer.")
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        overlap_step = self.overlap_step
-        for i in range(0, overlap_step):
+        warmup_step = self.warmup_step
+        for i in range(0, warmup_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
+            enqueue_computes(i + warmup_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)
@@ -866,7 +1002,7 @@ class Muon(torch.optim.Optimizer):
         amsgrad: bool,
         beta1: float,
         beta2: float,
-        lr: Union[float, torch.Tensor],
+        lr: float | torch.Tensor,
         weight_decay: float,
         eps: float,
         maximize: bool,
@@ -876,10 +1012,10 @@ class Muon(torch.optim.Optimizer):
 
         # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer
         # treating it as a scalar.
-        lr_dict: Optional[DeviceDict] = ({
+        lr_dict: DeviceDict | None = ({
             lr.device: lr
         } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else
-                                         None)
+                                      None)
         grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype(
             [
                 params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
@@ -926,6 +1062,159 @@ class Muon(torch.optim.Optimizer):
                 maximize=maximize,
             )
 
+    def _step_muon(self, group, qk_logits=None):
+        params = group["params"]
+        lr = group["lr"]
+        weight_decay = group["weight_decay"]
+        momentum = group["momentum"]
+        names = group["names"]
+
+        param_dtensors = []
+        param_tensors = []
+        name_dtensors = []
+        name_tensors = []
+
+        if self.use_distributed_muon:
+            self.distributed_muon(names=names,
+                                  params=params,
+                                  group=group,
+                                  lr=lr,
+                                  weight_decay=weight_decay,
+                                  momentum=momentum,
+                                  qk_logits=qk_logits)
+            return
+
+        for n, p in zip(names, params):
+            if p is None or p.grad is None:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(
+                        isinstance(placement, Replicate)
+                        for placement in p.placements):
+                    param_tensors.append(p)
+                    name_tensors.append(n)
+                else:
+                    param_dtensors.append(p)
+                    name_dtensors.append(n)
+            elif isinstance(p.data, torch.Tensor):
+                param_tensors.append(p)
+                name_tensors.append(n)
+            else:
+                raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+        logger.debug(
+            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors"
+        )
+
+        if len(param_dtensors) > 0:
+            if not dist.is_initialized():
+                raise RuntimeError(
+                    "Parallel Muon requires torch.distributed to be initialized."
+                )
+
+            # To support different placements, we group parameters by placements
+            # and run parallel Muon on each group.
+
+            placement_to_params = defaultdict(lambda: ([], []))
+            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
+
+            assert len(name_dtensors) == len(param_dtensors)
+            for n, p in zip(name_dtensors, param_dtensors):
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][0].append(n)
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][1].append(p)
+
+            for _, (names, params) in placement_to_params.items():
+                self.parallel(
+                    names,
+                    params,
+                    group,
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    momentum=momentum,
+                    qk_logits=qk_logits,
+                )
+
+        if len(param_tensors) > 0:
+            self.base(
+                name_tensors,
+                param_tensors,
+                group,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                qk_logits=qk_logits,
+            )
+
+    def _step_adamw_params(self, params, group):
+        params_with_grads = []
+        grads = []
+        moment1 = []
+        moment2 = []
+        max_exp_avg_sqs = []
+        state_steps = []
+        lr = group["lr"]
+        beta1, beta2 = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        weight_decay = group["weight_decay"]
+
+        for p in params:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            params_with_grads.append(p)
+            grads.append(g)
+            if "step" not in state:
+                state["step"] = (torch.zeros((),
+                                             dtype=torch.float32,
+                                             device=p.device))
+                state["moment1"] = torch.zeros_like(g)
+                state["moment2"] = torch.zeros_like(g)
+            moment1.append(state["moment1"])
+            moment2.append(state["moment2"])
+            if not isinstance(state["step"], torch.Tensor):
+                step_tensor = torch.tensor(state["step"],
+                                           dtype=torch.float32,
+                                           device=p.device)
+            else:
+                step_tensor = state["step"]
+            state_steps.append(step_tensor)
+
+        self._fused_adamw(
+            params_with_grads,
+            grads,
+            moment1,
+            moment2,
+            max_exp_avg_sqs,
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+    def _step_adamw(self, group):
+        params = group["params"]
+
+        # group params with it's type and placement
+        placement_to_params: dict[tuple[Placement | type,
+                                        DeviceMesh | None]] = defaultdict(list)
+        for p in params:
+            match p:
+                case DTensor():
+                    placement_to_params[tuple([p.placements,
+                                               p.device_mesh])].append(p)
+                case torch.Tensor():
+                    placement_to_params[tuple([torch.Tensor, None])].append(p)
+
+        for params in placement_to_params.values():
+            self._step_adamw_params(params, group)
+
     def step(self, closure=None, qk_logits=None):
         """Perform a single optimization step.
 
@@ -943,127 +1232,9 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         for group in self.param_groups:
-            params = group["params"]
-
             if group["use_muon"]:
-                ############################
-                #           Muon           #
-                ############################
-                lr = group["lr"]
-                weight_decay = group["weight_decay"]
-                momentum = group["momentum"]
-                names = group["names"]
-
-                param_dtensors = []
-                param_tensors = []
-                name_dtensors = []
-                name_tensors = []
-
-                for n, p in zip(names, params):
-                    if p is None or p.grad is None:
-                        continue
-                    if isinstance(p.data, DTensor):
-                        if all(
-                                isinstance(placement, Replicate)
-                                for placement in p.placements):
-                            param_tensors.append(p)
-                            name_tensors.append(n)
-                        else:
-                            param_dtensors.append(p)
-                            name_dtensors.append(n)
-                    elif isinstance(p.data, torch.Tensor):
-                        param_tensors.append(p)
-                        name_tensors.append(n)
-                    else:
-                        raise TypeError(
-                            f"Unsupported parameter type: {type(p.data)}")
-
-                if self.debug:
-                    print(
-                        f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
-                        flush=True,
-                    )
-
-                if len(param_dtensors) > 0:
-                    if not dist.is_initialized():
-                        raise RuntimeError(
-                            "Parallel Muon requires torch.distributed to be initialized."
-                        )
-
-                    self.parallel(
-                        name_dtensors,
-                        param_dtensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
-                if len(param_tensors) > 0:
-                    self.base(
-                        name_tensors,
-                        param_tensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
+                self._step_muon(group, qk_logits=qk_logits)
             else:
-                ############################
-                #       AdamW backup       #
-                ############################
-
-                params_with_grads = []
-                grads = []
-                moment1 = []
-                moment2 = []
-                max_exp_avg_sqs = []
-                state_steps = []
-                lr = group["lr"]
-                beta1, beta2 = group["adamw_betas"]
-                eps = group["adamw_eps"]
-                weight_decay = group["weight_decay"]
-
-                for p in params:
-                    g = p.grad
-                    if g is None:
-                        continue
-                    state = self.state[p]
-                    params_with_grads.append(p)
-                    grads.append(g)
-                    if "step" not in state:
-                        state["step"] = (torch.zeros((),
-                                                     dtype=torch.float32,
-                                                     device=p.device))
-                        state["moment1"] = torch.zeros_like(g)
-                        state["moment2"] = torch.zeros_like(g)
-                    moment1.append(state["moment1"])
-                    moment2.append(state["moment2"])
-                    if not isinstance(state["step"], torch.Tensor):
-                        step_tensor = torch.tensor(state["step"],
-                                                   dtype=torch.float32,
-                                                   device=p.device)
-                    else:
-                        step_tensor = state["step"]
-                    state_steps.append(step_tensor)
-
-                self._fused_adamw(
-                    params_with_grads,
-                    grads,
-                    moment1,
-                    moment2,
-                    max_exp_avg_sqs,
-                    state_steps,
-                    amsgrad=False,
-                    beta1=beta1,
-                    beta2=beta2,
-                    lr=lr,
-                    weight_decay=weight_decay,
-                    eps=eps,
-                    maximize=False,
-                )
+                self._step_adamw(group)
 
         return loss

build/torch29-cxx11-cu128-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_811726c_dirty
-ops = torch.ops._optimizer_811726c_dirty
+from . import _optimizer_23d68bb_dirty
+ops = torch.ops._optimizer_23d68bb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_811726c_dirty::{op_name}"
+    return f"_optimizer_23d68bb_dirty::{op_name}"

build/torch29-cxx11-cu128-x86_64-linux/optimizer/_optimizer_23d68bb_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc97ff00a3255d5eb363958b1e619eadbc4315f1930d0fb59cfc9560c3951721
+size 1983488

build/torch29-cxx11-cu128-x86_64-linux/optimizer/_optimizer_811726c_dirty.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ab1875be65811d88c407f36077aced58056a4feeb9946d7cd40ec55c7e1025c8
-size 1871056

build/torch29-cxx11-cu128-x86_64-linux/optimizer/distributed/utils.py

@@ -0,0 +1,174 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import (Placement, Shard,
+                                                      _StridedShard)
+
+
+def get_slices_of_dtensor(
+    target: DTensor | torch.Tensor,
+    local_rank: int,
+    shard_mesh: DeviceMesh,
+    shard_placements: tuple[Placement],
+) -> tuple[slice]:
+    """
+    Get the slice of local tensor for a given rank from a tensor.
+    Args:
+        target (DTensor | torch.Tensor): The target tensor.
+        rank (int): The local rank of the shard group.
+        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
+        shard_placements (tuple[Placement]): The shard placements.
+    """
+
+    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
+
+    # find the global rank of the local rank in the shard mesh
+    rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank]
+
+    rank_coords = (shard_mesh.mesh == rank).nonzero()
+
+    assert len(rank_coords) == 1
+    rank_coords = tuple(rank_coords[0].tolist())
+
+    assert len(rank_coords) == len(shard_placements)
+
+    # Caution: Assuming replicate-to-shard of the shard mesh goes with
+    # left-to-right sharding. This is ensured by the sorting logic of
+    # construct_shard_mesh function.
+    for i, (rank_coord,
+            placement) in enumerate(zip(rank_coords, shard_placements)):
+        assert isinstance(placement, Shard)
+
+        num_ranks = shard_mesh.mesh.shape[i]
+
+        dim = placement.dim
+        dim_size = (slices[dim].stop - slices[dim].start)
+
+        if dim_size % num_ranks != 0:
+            raise NotImplementedError(
+                f"Dimension size {dim_size} is not divisible "
+                f"by number of ranks {num_ranks} for shard "
+                f"placement on dim {dim}.")
+
+        shard_size = dim_size // num_ranks
+
+        start = slices[dim].start + rank_coord * shard_size
+        end = start + shard_size
+
+        assert start < end <= slices[dim].stop
+
+        slices[dim] = slice(start, end)
+
+    return tuple(slices)
+
+
+_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, ProcessGroup]] = dict()
+
+
+def construct_shard_mesh(
+    placements: tuple[Placement],
+    mesh: DeviceMesh,
+) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
+    """
+    Construct Shard Mesh and Placements for unsharding.
+    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
+    """
+    my_rank = dist.get_rank()
+
+    assert mesh.mesh.device.type == 'cpu'
+
+    # Copy mesh to avoid modifying the original mesh
+    mesh = mesh.mesh.clone()
+
+    # 1. Sort placements. Replicate first, then Shard by dim ascending.
+
+    # For Shard, strided shard comes after regular shard on the same dim
+    # to preserve left-to-right order of replicate-to-shard.
+    # This is because that strided shard is using stride to represent
+    # more fine-grained sharding on the same dim.
+    # Please check the URL below for _StridedShard.
+    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
+
+    def placement_sort_key(
+        placement_with_index: tuple[float, Placement]
+    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
+        index, placement = placement_with_index
+        is_replicate = placement.is_replicate()
+        is_shard = placement.is_shard()
+        is_partial = placement.is_partial()
+
+        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
+        assert not is_partial, "Partial placement is not supported."
+
+        if is_replicate:
+            return (-1.0, 0, index)
+        elif is_shard:
+            if isinstance(placement, _StridedShard):
+                return (placement.dim, 1 / placement.split_factor, index)
+            return (placement.dim, 0, index)
+        else:
+            raise TypeError(f"Unknown placement type: {type(placement)}")
+
+    placements_with_index: list[tuple[int,
+                                      Placement]] = list(enumerate(placements))
+    placements_with_index = sorted(placements_with_index,
+                                   key=placement_sort_key)
+
+    sorted_indices, sorted_placements = zip(*placements_with_index)
+
+    # 2. Permute mesh according to sorted placements.
+    sorted_mesh = mesh.permute(sorted_indices)
+
+    # 3. Collect list of shard meshes by removing replicate dims
+    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
+    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
+    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
+
+    # merge replicate dims
+    # shard_meshes became a list of shard meshes with a length of replicate degree
+    if num_replicates > 0:
+        sorted_mesh = sorted_mesh.flatten(
+            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
+        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
+    else:
+        shard_meshes = [sorted_mesh]
+    shard_placements = sorted_placements[num_replicates:]
+
+    # assume all shard placements are different
+    assert len(shard_placements) == len(set(shard_placements))
+
+    # 4. Construct ProcessGroups
+    # Caution: all groups should be created in the same order in all processes,
+    # even though each process only needs its own group.
+
+    # To use tensor as dict key, convert it to tuple
+    def tensor_to_tuple(t):
+        if isinstance(t, torch.Tensor):
+            t = t.tolist()
+        if isinstance(t, list):
+            return tuple(tensor_to_tuple(x) for x in t)
+        return t
+
+    my_shard_mesh_as_tuple = None
+    for shard_mesh in shard_meshes:
+        assert isinstance(shard_mesh, torch.Tensor)
+        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
+
+        if (my_rank == shard_mesh).any().item():
+            assert my_shard_mesh_as_tuple is None
+            my_shard_mesh_as_tuple = shard_mesh_as_tuple
+
+        # update global cache
+        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
+            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
+            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
+                DeviceMesh(device_type="cuda", mesh=shard_mesh),
+                shard_process_group,
+            )
+
+    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
+        my_shard_mesh_as_tuple]
+
+    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch29-cxx11-cu128-x86_64-linux/optimizer/muon.py

@@ -1,18 +1,24 @@
 import logging
 import math
 import types
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor, Replicate, Shard
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor.placement_types import Placement
 
+from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
 from .matmul_transpose_triton import matmul_transpose_assign
 
 logger = logging.getLogger(__name__)
 
 COMM_DTYPE = torch.bfloat16
+DEFAULT_CHUNK_SIZE_RATIO = 4
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -62,23 +68,39 @@ def _zeropower_via_newtonschulz5(G, steps):
 @dataclass
 class _muon_state:
     # TODO: use Optional
-    worker_rank: int | None = None
+    worker_rank: int
+    process_group: ProcessGroup
+    shard_mesh: DeviceMesh
+    shard_placements: tuple[Placement, ...]
+    name: str
+    qk_clip_state: torch.Tensor | None = None
     gathered_grad: torch.Tensor | None = None
     scattered_u: DTensor | None = None
     computed_u: torch.Tensor | None = None
     gather_event: torch.cuda.Event | None = None
     compute_event: torch.cuda.Event | None = None
     scatter_event: torch.cuda.Event | None = None
-    process_group = None
-    qk_clip_state = None
 
 
-def split_elems_for_src(param, src_rank, num_ranks) -> int:
-    rows = param.shape[0]
-    cols = int(param.numel() // rows)
-    base, rem = divmod(rows, num_ranks)
-    my_rows = base + (1 if src_rank < rem else 0)
-    return my_rows * cols
+def numel_for_rank(
+    param: DTensor,
+    local_rank: int,
+    state: _muon_state,
+) -> int:
+    slices = get_slices_of_dtensor(
+        param,
+        local_rank,
+        state.shard_mesh,
+        state.shard_placements,
+    )
+
+    numel = 1
+    for s, dim in zip(slices, param.shape):
+        start, stop, step = s.indices(dim)
+        length = max(0, (stop - start + (step - 1)) // step)
+        numel *= length
+
+    return numel
 
 
 @torch.no_grad()
@@ -91,8 +113,7 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
         for p in params:
             state = param_to_state[id(p)]
             if rank == state.worker_rank:
-                num_ranks = dist.get_world_size(group=state.process_group)
-                state.gathered_grad = torch.empty(p.grad.numel(),
+                state.gathered_grad = torch.empty(p.shape,
                                                   dtype=COMM_DTYPE,
                                                   device="cuda")
             else:
@@ -121,11 +142,11 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             state = param_to_state[id(p)]
             dst = state.worker_rank
             assert dst < num_ranks
-            shard_elems = split_elems_for_src(p, rank, num_ranks)
+            shard_elems = numel_for_rank(p, rank, state)
             g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            g = g.to_local().to(COMM_DTYPE).contiguous()
             assert g.numel() == shard_elems
-            per_dst[dst].append(g)
+            per_dst[dst].append(g.view(-1))
             send_counts[dst] += shard_elems
 
         assert any(
@@ -148,13 +169,18 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                total += split_elems_for_src(p, src, num_ranks)
+                total += numel_for_rank(p, src, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
         #All2All
+        logger.debug(f"send_buf size: {send_buf.numel()}, "
+                     f"recv_buf size: {recv_buf.numel()}, "
+                     f"recv_counts: {recv_counts}, "
+                     f"send_counts: {send_counts}, "
+                     f"process_group: {str(process_group)}")
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -179,7 +205,6 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
@@ -189,22 +214,28 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                n = split_elems_for_src(p, src, num_ranks)
+
+                # get the slice of the full dtensor corresponding to rank src.
+                slices = get_slices_of_dtensor(state.gathered_grad, src,
+                                               state.shard_mesh,
+                                               state.shard_placements)
+
+                dst = state.gathered_grad[slices]
+                assert dst._base is state.gathered_grad
+
+                n = dst.numel()
                 assert n > 0
 
                 sg = recv_buf.narrow(0, off + inner_off, n)
-                woff = write_offsets[id(p)]
-                dst = state.gathered_grad.narrow(0, woff, n)
+                sg = sg.reshape_as(dst)
                 dst.copy_(sg)
 
-                write_offsets[id(p)] += n
                 inner_off += n
             off += block
 
         for p in params:
             state = param_to_state[id(p)]
             if state.worker_rank == rank:
-                state.gathered_grad = state.gathered_grad.view_as(p)
                 state.gather_event = torch.cuda.Event()
                 state.gather_event.record(comm_stream)
             else:
@@ -277,14 +308,19 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert state.computed_u is not None
 
-                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous()
 
                 offset = 0
                 for dst in range(num_ranks):
-                    n = split_elems_for_src(p, dst, num_ranks)
+                    # get the slice of the full tensor corresponding to rank dst.
+                    slices = get_slices_of_dtensor(u_full, dst,
+                                                   state.shard_mesh,
+                                                   state.shard_placements)
+                    su = u_full[slices].flatten()
+
+                    n = su.numel()
                     assert n > 0
 
-                    su = u_full.narrow(0, offset, n)
                     per_dst[dst].append(su)
                     send_counts[dst] += n
                     offset += n
@@ -313,7 +349,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                total += split_elems_for_src(p, rank, num_ranks)
+                total += numel_for_rank(p, rank, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
@@ -357,7 +393,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                n = split_elems_for_src(p, rank, num_ranks)
+                n = numel_for_rank(p, rank, state)
                 assert n > 0
 
                 flat_local = recv_buf.narrow(0, off + inner_off,
@@ -398,11 +434,23 @@ def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
         state.scattered_u = None
         u_dtensor = None
 
-        scales_full = Muon._compute_scales(p, state.qk_clip_state)
+        scales_full = Muon._compute_scales(
+            p,
+            state.qk_clip_state) if state.qk_clip_state is not None else None
         if scales_full is not None:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            local_rank = dist.get_rank(group=state.process_group)
-            scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank]
+            # Have to slice scales_full among dim 0
+            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
+                                                  state.shard_placements)
+            ratio = p.shape[0] // scales_full.shape[0]
+            scales_slice = slice(
+                None if weight_slices[0].start is None else
+                weight_slices[0].start // ratio,
+                None if weight_slices[0].stop is None else
+                weight_slices[0].stop // ratio,
+                None,
+            )
+
+            scales_local = scales_full[scales_slice]
             scales_local = DTensor.from_local(
                 scales_local,
                 placements=p.placements,
@@ -478,11 +526,11 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
 @dataclass
 class QKClipInfo:
     """Per-parameter dynamic info computed from config + runtime logits."""
-    kind: Optional[str]  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
-    indices: List[int]  # which heads to consider for clipping
+    kind: str | None  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
+    indices: list[int]  # which heads to consider for clipping
     head_dim: int  # from config
     threshold: float  # from config
-    logit: Optional[torch.Tensor]
+    logit: torch.Tensor | None
 
 
 class Muon(torch.optim.Optimizer):
@@ -525,11 +573,16 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
-        overlap_step : How many all2all gather, compute operations are launched in advance
-                       before the corresponding all2all scatter steps begin.
-                       A higher overlap_step increases memory usage but can improve
-                       performance by overlapping communication.
-                       Parallel muon only. 
+        warmup_step : How many all2all gather, compute operations are launched in advance
+                      before the corresponding all2all scatter steps begin.
+                      A higher warmup_step increases memory usage but can improve
+                      performance by overlapping communication.
+                      Parallel muon only.
+        chunk_size : Batch size of parameters to process in each
+                     all2all gather/compute/scatter step.
+                     Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
+        use_distributed_muon: Use distributed muon by Liu et al. (2024).
+                              For testing purpose only.
     """
 
     def __init__(self,
@@ -549,7 +602,9 @@ class Muon(torch.optim.Optimizer):
                      "head_dim": 128,
                      "threshold": 100
                  },
-                 overlap_step=5):
+                 warmup_step=5,
+                 chunk_size=-1,
+                 use_distributed_muon=False):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -579,7 +634,9 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
-        self.overlap_step = overlap_step
+        self.warmup_step = warmup_step
+        self.chunk_size = chunk_size
+        self.use_distributed_muon = use_distributed_muon
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -597,6 +654,12 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
+    def set_rank_once(self, rank):
+        if self.rank is None:
+            self.rank = rank
+        else:
+            assert self.rank == rank
+
     def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
@@ -604,26 +667,13 @@ class Muon(torch.optim.Optimizer):
         assert isinstance(
             p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0), ):
-            # Case for FSDP
-            process_group = p.device_mesh.get_group(mesh_dim=0)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
-        elif p.placements == (Replicate(), Shard(dim=0)):
-            # Case for HSDP
-            process_group = p.device_mesh.get_group(mesh_dim=1)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if self.rank in shard_mesh:
-                    return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
-        else:
-            raise ValueError(f"Unsupported placements ({p.placements}).")
+        shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+            p.placements, p.device_mesh)
+
+        # set rank with the local rank in the shard process group
+        self.set_rank_once(dist.get_rank(group=shard_pg))
+
+        return shard_mesh, shard_pg, shard_placements
 
     def init_state_and_assign_params(self, names, params, group, qk_logits):
         param_to_state = {}
@@ -655,23 +705,32 @@ class Muon(torch.optim.Optimizer):
         ordered_params = list(params_sorted)
 
         round_robin = 0
-        mesh = None
-        shard_mesh = None
-        process_group = None
+        mesh = ordered_params[0].device_mesh
+        placements = ordered_params[0].placements
+
+        shard_mesh, shard_pg, shard_placements = self.get_shard_mesh(
+            ordered_params[0])
+        shard_mesh_flattened = shard_mesh.mesh.flatten()
+        num_ranks = dist.get_world_size(group=shard_pg)
+
         for n, p in zip(ordered_names, ordered_params):
-            if mesh is None:
-                mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p)
-            elif mesh != p.device_mesh:
+            if mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
-            num_ranks = dist.get_world_size(group=process_group)
-            param_to_state[id(p)] = _muon_state()
-            param_to_state[id(
-                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
-            param_to_state[id(p)].process_group = process_group
+            if placements != p.placements:
+                raise ValueError("All parameters must have same placements.")
+
+            worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
+            round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-            param_to_state[id(p)].qk_clip_state = qk_clip_state
-            round_robin = (round_robin + 1) % len(shard_mesh)
+
+            param_to_state[id(p)] = _muon_state(
+                worker_rank=worker_rank,
+                process_group=shard_pg,
+                shard_mesh=shard_mesh,
+                shard_placements=shard_placements,
+                name=n,
+                qk_clip_state=qk_clip_state,
+            )
 
         return param_to_state, ordered_params
 
@@ -705,10 +764,73 @@ class Muon(torch.optim.Optimizer):
 
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
 
-            scales_full = self._compute_scales(p, qk_clip_state)
+            scales_full = self._compute_scales(
+                p, qk_clip_state) if qk_clip_state is not None else None
             if scales_full is not None:
                 Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
 
+    def distributed_muon(
+        self,
+        names: list[str],
+        params: list[torch.nn.Parameter],
+        group: dict[str, Any],
+        lr: float,
+        weight_decay: float,
+        momentum: float,
+        qk_logits: list[torch.Tensor | DTensor] | None,
+    ):
+        """ Implementation of Distributed Muon by Liu et al. """
+        if qk_logits is not None:
+            raise NotImplementedError("QK clipping is not supported yet")
+
+        if isinstance(params[0], DTensor):
+            shard_mesh, _, shard_placements = construct_shard_mesh(
+                placements=params[0].placements,
+                mesh=params[0].device_mesh,
+            )
+
+        for n, p in zip(names, params):
+            g = p.grad
+            if g is None:
+                continue
+            if g.ndim > 2:
+                g = g.view(g.size(0), -1)
+            assert g is not None
+
+            # calc update
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(momentum).add_(g)
+            if group["nesterov"]:
+                g = g.add(buf, alpha=momentum)
+            else:
+                g = buf
+
+            # Gather G
+            if isinstance(p.data, DTensor):
+                g = g.full_tensor()
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                             steps=group["ns_steps"])
+
+            if isinstance(p.data, DTensor):
+                slices = get_slices_of_dtensor(
+                    target=p,
+                    local_rank=dist.get_rank(),
+                    shard_mesh=shard_mesh,
+                    shard_placements=shard_placements,
+                )
+                u_shard = u[slices]
+                u = DTensor.from_local(
+                    u_shard,
+                    device_mesh=p.device_mesh,
+                    placements=p.placements,
+                )
+
+            adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+            Muon._update_p(p, u, lr, adjusted_lr, weight_decay)
+
     def _update_g(self, p, g, group, momentum):
         # calc update
         state = self.state[p]
@@ -727,6 +849,9 @@ class Muon(torch.optim.Optimizer):
         p.data.add_(u, alpha=-adjusted_lr)
 
     def get_qk_clip_info(self, n, qk_logits):
+        if self.clip_config is None:
+            return None
+
         head_dim = self.clip_config.get('head_dim')
         threshold = self.clip_config.get('threshold')
         kind, layer_idx = parse_qk_layer(n)
@@ -737,6 +862,11 @@ class Muon(torch.optim.Optimizer):
             indices_key = 'q_indices' if 'q' in kind else 'k_indices'
             indices = self.clip_config.get(indices_key, []) or []
 
+            if isinstance(logit, DTensor):
+                # In TP settings, qk_logits may be DTensor
+                # We convert it to full tensor here for simplicity
+                logit = logit.full_tensor()
+
         return QKClipInfo(
             kind=kind,
             indices=indices,
@@ -835,22 +965,28 @@ class Muon(torch.optim.Optimizer):
                 _update_param(p, state, lr, adjusted_lr, weight_decay,
                               self.rank, self.compute_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(
-            params[0])].process_group)
+        if self.chunk_size == -1:
+            shard_ranks = dist.get_world_size(param_to_state[id(
+                params[0])].process_group)
+            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+        elif self.chunk_size > 0:
+            chunk_size = self.chunk_size
+        else:
+            raise ValueError("chunk_size must be -1 or a positive integer.")
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        overlap_step = self.overlap_step
-        for i in range(0, overlap_step):
+        warmup_step = self.warmup_step
+        for i in range(0, warmup_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
+            enqueue_computes(i + warmup_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)
@@ -866,7 +1002,7 @@ class Muon(torch.optim.Optimizer):
         amsgrad: bool,
         beta1: float,
         beta2: float,
-        lr: Union[float, torch.Tensor],
+        lr: float | torch.Tensor,
         weight_decay: float,
         eps: float,
         maximize: bool,
@@ -876,10 +1012,10 @@ class Muon(torch.optim.Optimizer):
 
         # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer
         # treating it as a scalar.
-        lr_dict: Optional[DeviceDict] = ({
+        lr_dict: DeviceDict | None = ({
             lr.device: lr
         } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else
-                                         None)
+                                      None)
         grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype(
             [
                 params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
@@ -926,6 +1062,159 @@ class Muon(torch.optim.Optimizer):
                 maximize=maximize,
             )
 
+    def _step_muon(self, group, qk_logits=None):
+        params = group["params"]
+        lr = group["lr"]
+        weight_decay = group["weight_decay"]
+        momentum = group["momentum"]
+        names = group["names"]
+
+        param_dtensors = []
+        param_tensors = []
+        name_dtensors = []
+        name_tensors = []
+
+        if self.use_distributed_muon:
+            self.distributed_muon(names=names,
+                                  params=params,
+                                  group=group,
+                                  lr=lr,
+                                  weight_decay=weight_decay,
+                                  momentum=momentum,
+                                  qk_logits=qk_logits)
+            return
+
+        for n, p in zip(names, params):
+            if p is None or p.grad is None:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(
+                        isinstance(placement, Replicate)
+                        for placement in p.placements):
+                    param_tensors.append(p)
+                    name_tensors.append(n)
+                else:
+                    param_dtensors.append(p)
+                    name_dtensors.append(n)
+            elif isinstance(p.data, torch.Tensor):
+                param_tensors.append(p)
+                name_tensors.append(n)
+            else:
+                raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+        logger.debug(
+            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors"
+        )
+
+        if len(param_dtensors) > 0:
+            if not dist.is_initialized():
+                raise RuntimeError(
+                    "Parallel Muon requires torch.distributed to be initialized."
+                )
+
+            # To support different placements, we group parameters by placements
+            # and run parallel Muon on each group.
+
+            placement_to_params = defaultdict(lambda: ([], []))
+            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
+
+            assert len(name_dtensors) == len(param_dtensors)
+            for n, p in zip(name_dtensors, param_dtensors):
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][0].append(n)
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][1].append(p)
+
+            for _, (names, params) in placement_to_params.items():
+                self.parallel(
+                    names,
+                    params,
+                    group,
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    momentum=momentum,
+                    qk_logits=qk_logits,
+                )
+
+        if len(param_tensors) > 0:
+            self.base(
+                name_tensors,
+                param_tensors,
+                group,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                qk_logits=qk_logits,
+            )
+
+    def _step_adamw_params(self, params, group):
+        params_with_grads = []
+        grads = []
+        moment1 = []
+        moment2 = []
+        max_exp_avg_sqs = []
+        state_steps = []
+        lr = group["lr"]
+        beta1, beta2 = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        weight_decay = group["weight_decay"]
+
+        for p in params:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            params_with_grads.append(p)
+            grads.append(g)
+            if "step" not in state:
+                state["step"] = (torch.zeros((),
+                                             dtype=torch.float32,
+                                             device=p.device))
+                state["moment1"] = torch.zeros_like(g)
+                state["moment2"] = torch.zeros_like(g)
+            moment1.append(state["moment1"])
+            moment2.append(state["moment2"])
+            if not isinstance(state["step"], torch.Tensor):
+                step_tensor = torch.tensor(state["step"],
+                                           dtype=torch.float32,
+                                           device=p.device)
+            else:
+                step_tensor = state["step"]
+            state_steps.append(step_tensor)
+
+        self._fused_adamw(
+            params_with_grads,
+            grads,
+            moment1,
+            moment2,
+            max_exp_avg_sqs,
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+    def _step_adamw(self, group):
+        params = group["params"]
+
+        # group params with it's type and placement
+        placement_to_params: dict[tuple[Placement | type,
+                                        DeviceMesh | None]] = defaultdict(list)
+        for p in params:
+            match p:
+                case DTensor():
+                    placement_to_params[tuple([p.placements,
+                                               p.device_mesh])].append(p)
+                case torch.Tensor():
+                    placement_to_params[tuple([torch.Tensor, None])].append(p)
+
+        for params in placement_to_params.values():
+            self._step_adamw_params(params, group)
+
     def step(self, closure=None, qk_logits=None):
         """Perform a single optimization step.
 
@@ -943,127 +1232,9 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         for group in self.param_groups:
-            params = group["params"]
-
             if group["use_muon"]:
-                ############################
-                #           Muon           #
-                ############################
-                lr = group["lr"]
-                weight_decay = group["weight_decay"]
-                momentum = group["momentum"]
-                names = group["names"]
-
-                param_dtensors = []
-                param_tensors = []
-                name_dtensors = []
-                name_tensors = []
-
-                for n, p in zip(names, params):
-                    if p is None or p.grad is None:
-                        continue
-                    if isinstance(p.data, DTensor):
-                        if all(
-                                isinstance(placement, Replicate)
-                                for placement in p.placements):
-                            param_tensors.append(p)
-                            name_tensors.append(n)
-                        else:
-                            param_dtensors.append(p)
-                            name_dtensors.append(n)
-                    elif isinstance(p.data, torch.Tensor):
-                        param_tensors.append(p)
-                        name_tensors.append(n)
-                    else:
-                        raise TypeError(
-                            f"Unsupported parameter type: {type(p.data)}")
-
-                if self.debug:
-                    print(
-                        f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
-                        flush=True,
-                    )
-
-                if len(param_dtensors) > 0:
-                    if not dist.is_initialized():
-                        raise RuntimeError(
-                            "Parallel Muon requires torch.distributed to be initialized."
-                        )
-
-                    self.parallel(
-                        name_dtensors,
-                        param_dtensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
-                if len(param_tensors) > 0:
-                    self.base(
-                        name_tensors,
-                        param_tensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
+                self._step_muon(group, qk_logits=qk_logits)
             else:
-                ############################
-                #       AdamW backup       #
-                ############################
-
-                params_with_grads = []
-                grads = []
-                moment1 = []
-                moment2 = []
-                max_exp_avg_sqs = []
-                state_steps = []
-                lr = group["lr"]
-                beta1, beta2 = group["adamw_betas"]
-                eps = group["adamw_eps"]
-                weight_decay = group["weight_decay"]
-
-                for p in params:
-                    g = p.grad
-                    if g is None:
-                        continue
-                    state = self.state[p]
-                    params_with_grads.append(p)
-                    grads.append(g)
-                    if "step" not in state:
-                        state["step"] = (torch.zeros((),
-                                                     dtype=torch.float32,
-                                                     device=p.device))
-                        state["moment1"] = torch.zeros_like(g)
-                        state["moment2"] = torch.zeros_like(g)
-                    moment1.append(state["moment1"])
-                    moment2.append(state["moment2"])
-                    if not isinstance(state["step"], torch.Tensor):
-                        step_tensor = torch.tensor(state["step"],
-                                                   dtype=torch.float32,
-                                                   device=p.device)
-                    else:
-                        step_tensor = state["step"]
-                    state_steps.append(step_tensor)
-
-                self._fused_adamw(
-                    params_with_grads,
-                    grads,
-                    moment1,
-                    moment2,
-                    max_exp_avg_sqs,
-                    state_steps,
-                    amsgrad=False,
-                    beta1=beta1,
-                    beta2=beta2,
-                    lr=lr,
-                    weight_decay=weight_decay,
-                    eps=eps,
-                    maximize=False,
-                )
+                self._step_adamw(group)
 
         return loss

build/torch29-cxx11-cu130-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_811726c_dirty
-ops = torch.ops._optimizer_811726c_dirty
+from . import _optimizer_23d68bb_dirty
+ops = torch.ops._optimizer_23d68bb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_811726c_dirty::{op_name}"
+    return f"_optimizer_23d68bb_dirty::{op_name}"

build/torch29-cxx11-cu130-x86_64-linux/optimizer/_optimizer_23d68bb_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa394498c52692c29094cbd2cc3da6c4c37aefaa4454c97487f8e91827fbd814
+size 1988672

build/torch29-cxx11-cu130-x86_64-linux/optimizer/_optimizer_811726c_dirty.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:52a744cf30c60fe1e8fc35ebb0d3421d679bb2047fbb4602846bd6902cfa9e52
-size 1872152

build/torch29-cxx11-cu130-x86_64-linux/optimizer/distributed/utils.py

@@ -0,0 +1,174 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import (Placement, Shard,
+                                                      _StridedShard)
+
+
+def get_slices_of_dtensor(
+    target: DTensor | torch.Tensor,
+    local_rank: int,
+    shard_mesh: DeviceMesh,
+    shard_placements: tuple[Placement],
+) -> tuple[slice]:
+    """
+    Get the slice of local tensor for a given rank from a tensor.
+    Args:
+        target (DTensor | torch.Tensor): The target tensor.
+        rank (int): The local rank of the shard group.
+        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
+        shard_placements (tuple[Placement]): The shard placements.
+    """
+
+    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
+
+    # find the global rank of the local rank in the shard mesh
+    rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank]
+
+    rank_coords = (shard_mesh.mesh == rank).nonzero()
+
+    assert len(rank_coords) == 1
+    rank_coords = tuple(rank_coords[0].tolist())
+
+    assert len(rank_coords) == len(shard_placements)
+
+    # Caution: Assuming replicate-to-shard of the shard mesh goes with
+    # left-to-right sharding. This is ensured by the sorting logic of
+    # construct_shard_mesh function.
+    for i, (rank_coord,
+            placement) in enumerate(zip(rank_coords, shard_placements)):
+        assert isinstance(placement, Shard)
+
+        num_ranks = shard_mesh.mesh.shape[i]
+
+        dim = placement.dim
+        dim_size = (slices[dim].stop - slices[dim].start)
+
+        if dim_size % num_ranks != 0:
+            raise NotImplementedError(
+                f"Dimension size {dim_size} is not divisible "
+                f"by number of ranks {num_ranks} for shard "
+                f"placement on dim {dim}.")
+
+        shard_size = dim_size // num_ranks
+
+        start = slices[dim].start + rank_coord * shard_size
+        end = start + shard_size
+
+        assert start < end <= slices[dim].stop
+
+        slices[dim] = slice(start, end)
+
+    return tuple(slices)
+
+
+_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, ProcessGroup]] = dict()
+
+
+def construct_shard_mesh(
+    placements: tuple[Placement],
+    mesh: DeviceMesh,
+) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
+    """
+    Construct Shard Mesh and Placements for unsharding.
+    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
+    """
+    my_rank = dist.get_rank()
+
+    assert mesh.mesh.device.type == 'cpu'
+
+    # Copy mesh to avoid modifying the original mesh
+    mesh = mesh.mesh.clone()
+
+    # 1. Sort placements. Replicate first, then Shard by dim ascending.
+
+    # For Shard, strided shard comes after regular shard on the same dim
+    # to preserve left-to-right order of replicate-to-shard.
+    # This is because that strided shard is using stride to represent
+    # more fine-grained sharding on the same dim.
+    # Please check the URL below for _StridedShard.
+    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
+
+    def placement_sort_key(
+        placement_with_index: tuple[float, Placement]
+    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
+        index, placement = placement_with_index
+        is_replicate = placement.is_replicate()
+        is_shard = placement.is_shard()
+        is_partial = placement.is_partial()
+
+        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
+        assert not is_partial, "Partial placement is not supported."
+
+        if is_replicate:
+            return (-1.0, 0, index)
+        elif is_shard:
+            if isinstance(placement, _StridedShard):
+                return (placement.dim, 1 / placement.split_factor, index)
+            return (placement.dim, 0, index)
+        else:
+            raise TypeError(f"Unknown placement type: {type(placement)}")
+
+    placements_with_index: list[tuple[int,
+                                      Placement]] = list(enumerate(placements))
+    placements_with_index = sorted(placements_with_index,
+                                   key=placement_sort_key)
+
+    sorted_indices, sorted_placements = zip(*placements_with_index)
+
+    # 2. Permute mesh according to sorted placements.
+    sorted_mesh = mesh.permute(sorted_indices)
+
+    # 3. Collect list of shard meshes by removing replicate dims
+    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
+    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
+    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
+
+    # merge replicate dims
+    # shard_meshes became a list of shard meshes with a length of replicate degree
+    if num_replicates > 0:
+        sorted_mesh = sorted_mesh.flatten(
+            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
+        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
+    else:
+        shard_meshes = [sorted_mesh]
+    shard_placements = sorted_placements[num_replicates:]
+
+    # assume all shard placements are different
+    assert len(shard_placements) == len(set(shard_placements))
+
+    # 4. Construct ProcessGroups
+    # Caution: all groups should be created in the same order in all processes,
+    # even though each process only needs its own group.
+
+    # To use tensor as dict key, convert it to tuple
+    def tensor_to_tuple(t):
+        if isinstance(t, torch.Tensor):
+            t = t.tolist()
+        if isinstance(t, list):
+            return tuple(tensor_to_tuple(x) for x in t)
+        return t
+
+    my_shard_mesh_as_tuple = None
+    for shard_mesh in shard_meshes:
+        assert isinstance(shard_mesh, torch.Tensor)
+        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
+
+        if (my_rank == shard_mesh).any().item():
+            assert my_shard_mesh_as_tuple is None
+            my_shard_mesh_as_tuple = shard_mesh_as_tuple
+
+        # update global cache
+        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
+            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
+            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
+                DeviceMesh(device_type="cuda", mesh=shard_mesh),
+                shard_process_group,
+            )
+
+    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
+        my_shard_mesh_as_tuple]
+
+    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch29-cxx11-cu130-x86_64-linux/optimizer/muon.py

@@ -1,18 +1,24 @@
 import logging
 import math
 import types
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor, Replicate, Shard
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor.placement_types import Placement
 
+from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
 from .matmul_transpose_triton import matmul_transpose_assign
 
 logger = logging.getLogger(__name__)
 
 COMM_DTYPE = torch.bfloat16
+DEFAULT_CHUNK_SIZE_RATIO = 4
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -62,23 +68,39 @@ def _zeropower_via_newtonschulz5(G, steps):
 @dataclass
 class _muon_state:
     # TODO: use Optional
-    worker_rank: int | None = None
+    worker_rank: int
+    process_group: ProcessGroup
+    shard_mesh: DeviceMesh
+    shard_placements: tuple[Placement, ...]
+    name: str
+    qk_clip_state: torch.Tensor | None = None
     gathered_grad: torch.Tensor | None = None
     scattered_u: DTensor | None = None
     computed_u: torch.Tensor | None = None
     gather_event: torch.cuda.Event | None = None
     compute_event: torch.cuda.Event | None = None
     scatter_event: torch.cuda.Event | None = None
-    process_group = None
-    qk_clip_state = None
 
 
-def split_elems_for_src(param, src_rank, num_ranks) -> int:
-    rows = param.shape[0]
-    cols = int(param.numel() // rows)
-    base, rem = divmod(rows, num_ranks)
-    my_rows = base + (1 if src_rank < rem else 0)
-    return my_rows * cols
+def numel_for_rank(
+    param: DTensor,
+    local_rank: int,
+    state: _muon_state,
+) -> int:
+    slices = get_slices_of_dtensor(
+        param,
+        local_rank,
+        state.shard_mesh,
+        state.shard_placements,
+    )
+
+    numel = 1
+    for s, dim in zip(slices, param.shape):
+        start, stop, step = s.indices(dim)
+        length = max(0, (stop - start + (step - 1)) // step)
+        numel *= length
+
+    return numel
 
 
 @torch.no_grad()
@@ -91,8 +113,7 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
         for p in params:
             state = param_to_state[id(p)]
             if rank == state.worker_rank:
-                num_ranks = dist.get_world_size(group=state.process_group)
-                state.gathered_grad = torch.empty(p.grad.numel(),
+                state.gathered_grad = torch.empty(p.shape,
                                                   dtype=COMM_DTYPE,
                                                   device="cuda")
             else:
@@ -121,11 +142,11 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             state = param_to_state[id(p)]
             dst = state.worker_rank
             assert dst < num_ranks
-            shard_elems = split_elems_for_src(p, rank, num_ranks)
+            shard_elems = numel_for_rank(p, rank, state)
             g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            g = g.to_local().to(COMM_DTYPE).contiguous()
             assert g.numel() == shard_elems
-            per_dst[dst].append(g)
+            per_dst[dst].append(g.view(-1))
             send_counts[dst] += shard_elems
 
         assert any(
@@ -148,13 +169,18 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                total += split_elems_for_src(p, src, num_ranks)
+                total += numel_for_rank(p, src, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
         #All2All
+        logger.debug(f"send_buf size: {send_buf.numel()}, "
+                     f"recv_buf size: {recv_buf.numel()}, "
+                     f"recv_counts: {recv_counts}, "
+                     f"send_counts: {send_counts}, "
+                     f"process_group: {str(process_group)}")
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -179,7 +205,6 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
@@ -189,22 +214,28 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                n = split_elems_for_src(p, src, num_ranks)
+
+                # get the slice of the full dtensor corresponding to rank src.
+                slices = get_slices_of_dtensor(state.gathered_grad, src,
+                                               state.shard_mesh,
+                                               state.shard_placements)
+
+                dst = state.gathered_grad[slices]
+                assert dst._base is state.gathered_grad
+
+                n = dst.numel()
                 assert n > 0
 
                 sg = recv_buf.narrow(0, off + inner_off, n)
-                woff = write_offsets[id(p)]
-                dst = state.gathered_grad.narrow(0, woff, n)
+                sg = sg.reshape_as(dst)
                 dst.copy_(sg)
 
-                write_offsets[id(p)] += n
                 inner_off += n
             off += block
 
         for p in params:
             state = param_to_state[id(p)]
             if state.worker_rank == rank:
-                state.gathered_grad = state.gathered_grad.view_as(p)
                 state.gather_event = torch.cuda.Event()
                 state.gather_event.record(comm_stream)
             else:
@@ -277,14 +308,19 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert state.computed_u is not None
 
-                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous()
 
                 offset = 0
                 for dst in range(num_ranks):
-                    n = split_elems_for_src(p, dst, num_ranks)
+                    # get the slice of the full tensor corresponding to rank dst.
+                    slices = get_slices_of_dtensor(u_full, dst,
+                                                   state.shard_mesh,
+                                                   state.shard_placements)
+                    su = u_full[slices].flatten()
+
+                    n = su.numel()
                     assert n > 0
 
-                    su = u_full.narrow(0, offset, n)
                     per_dst[dst].append(su)
                     send_counts[dst] += n
                     offset += n
@@ -313,7 +349,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                total += split_elems_for_src(p, rank, num_ranks)
+                total += numel_for_rank(p, rank, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
@@ -357,7 +393,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                n = split_elems_for_src(p, rank, num_ranks)
+                n = numel_for_rank(p, rank, state)
                 assert n > 0
 
                 flat_local = recv_buf.narrow(0, off + inner_off,
@@ -398,11 +434,23 @@ def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
         state.scattered_u = None
         u_dtensor = None
 
-        scales_full = Muon._compute_scales(p, state.qk_clip_state)
+        scales_full = Muon._compute_scales(
+            p,
+            state.qk_clip_state) if state.qk_clip_state is not None else None
         if scales_full is not None:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            local_rank = dist.get_rank(group=state.process_group)
-            scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank]
+            # Have to slice scales_full among dim 0
+            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
+                                                  state.shard_placements)
+            ratio = p.shape[0] // scales_full.shape[0]
+            scales_slice = slice(
+                None if weight_slices[0].start is None else
+                weight_slices[0].start // ratio,
+                None if weight_slices[0].stop is None else
+                weight_slices[0].stop // ratio,
+                None,
+            )
+
+            scales_local = scales_full[scales_slice]
             scales_local = DTensor.from_local(
                 scales_local,
                 placements=p.placements,
@@ -478,11 +526,11 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
 @dataclass
 class QKClipInfo:
     """Per-parameter dynamic info computed from config + runtime logits."""
-    kind: Optional[str]  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
-    indices: List[int]  # which heads to consider for clipping
+    kind: str | None  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
+    indices: list[int]  # which heads to consider for clipping
     head_dim: int  # from config
     threshold: float  # from config
-    logit: Optional[torch.Tensor]
+    logit: torch.Tensor | None
 
 
 class Muon(torch.optim.Optimizer):
@@ -525,11 +573,16 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
-        overlap_step : How many all2all gather, compute operations are launched in advance
-                       before the corresponding all2all scatter steps begin.
-                       A higher overlap_step increases memory usage but can improve
-                       performance by overlapping communication.
-                       Parallel muon only. 
+        warmup_step : How many all2all gather, compute operations are launched in advance
+                      before the corresponding all2all scatter steps begin.
+                      A higher warmup_step increases memory usage but can improve
+                      performance by overlapping communication.
+                      Parallel muon only.
+        chunk_size : Batch size of parameters to process in each
+                     all2all gather/compute/scatter step.
+                     Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
+        use_distributed_muon: Use distributed muon by Liu et al. (2024).
+                              For testing purpose only.
     """
 
     def __init__(self,
@@ -549,7 +602,9 @@ class Muon(torch.optim.Optimizer):
                      "head_dim": 128,
                      "threshold": 100
                  },
-                 overlap_step=5):
+                 warmup_step=5,
+                 chunk_size=-1,
+                 use_distributed_muon=False):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -579,7 +634,9 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
-        self.overlap_step = overlap_step
+        self.warmup_step = warmup_step
+        self.chunk_size = chunk_size
+        self.use_distributed_muon = use_distributed_muon
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -597,6 +654,12 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
+    def set_rank_once(self, rank):
+        if self.rank is None:
+            self.rank = rank
+        else:
+            assert self.rank == rank
+
     def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
@@ -604,26 +667,13 @@ class Muon(torch.optim.Optimizer):
         assert isinstance(
             p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0), ):
-            # Case for FSDP
-            process_group = p.device_mesh.get_group(mesh_dim=0)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
-        elif p.placements == (Replicate(), Shard(dim=0)):
-            # Case for HSDP
-            process_group = p.device_mesh.get_group(mesh_dim=1)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if self.rank in shard_mesh:
-                    return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
-        else:
-            raise ValueError(f"Unsupported placements ({p.placements}).")
+        shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+            p.placements, p.device_mesh)
+
+        # set rank with the local rank in the shard process group
+        self.set_rank_once(dist.get_rank(group=shard_pg))
+
+        return shard_mesh, shard_pg, shard_placements
 
     def init_state_and_assign_params(self, names, params, group, qk_logits):
         param_to_state = {}
@@ -655,23 +705,32 @@ class Muon(torch.optim.Optimizer):
         ordered_params = list(params_sorted)
 
         round_robin = 0
-        mesh = None
-        shard_mesh = None
-        process_group = None
+        mesh = ordered_params[0].device_mesh
+        placements = ordered_params[0].placements
+
+        shard_mesh, shard_pg, shard_placements = self.get_shard_mesh(
+            ordered_params[0])
+        shard_mesh_flattened = shard_mesh.mesh.flatten()
+        num_ranks = dist.get_world_size(group=shard_pg)
+
         for n, p in zip(ordered_names, ordered_params):
-            if mesh is None:
-                mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p)
-            elif mesh != p.device_mesh:
+            if mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
-            num_ranks = dist.get_world_size(group=process_group)
-            param_to_state[id(p)] = _muon_state()
-            param_to_state[id(
-                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
-            param_to_state[id(p)].process_group = process_group
+            if placements != p.placements:
+                raise ValueError("All parameters must have same placements.")
+
+            worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
+            round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-            param_to_state[id(p)].qk_clip_state = qk_clip_state
-            round_robin = (round_robin + 1) % len(shard_mesh)
+
+            param_to_state[id(p)] = _muon_state(
+                worker_rank=worker_rank,
+                process_group=shard_pg,
+                shard_mesh=shard_mesh,
+                shard_placements=shard_placements,
+                name=n,
+                qk_clip_state=qk_clip_state,
+            )
 
         return param_to_state, ordered_params
 
@@ -705,10 +764,73 @@ class Muon(torch.optim.Optimizer):
 
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
 
-            scales_full = self._compute_scales(p, qk_clip_state)
+            scales_full = self._compute_scales(
+                p, qk_clip_state) if qk_clip_state is not None else None
             if scales_full is not None:
                 Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
 
+    def distributed_muon(
+        self,
+        names: list[str],
+        params: list[torch.nn.Parameter],
+        group: dict[str, Any],
+        lr: float,
+        weight_decay: float,
+        momentum: float,
+        qk_logits: list[torch.Tensor | DTensor] | None,
+    ):
+        """ Implementation of Distributed Muon by Liu et al. """
+        if qk_logits is not None:
+            raise NotImplementedError("QK clipping is not supported yet")
+
+        if isinstance(params[0], DTensor):
+            shard_mesh, _, shard_placements = construct_shard_mesh(
+                placements=params[0].placements,
+                mesh=params[0].device_mesh,
+            )
+
+        for n, p in zip(names, params):
+            g = p.grad
+            if g is None:
+                continue
+            if g.ndim > 2:
+                g = g.view(g.size(0), -1)
+            assert g is not None
+
+            # calc update
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(momentum).add_(g)
+            if group["nesterov"]:
+                g = g.add(buf, alpha=momentum)
+            else:
+                g = buf
+
+            # Gather G
+            if isinstance(p.data, DTensor):
+                g = g.full_tensor()
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                             steps=group["ns_steps"])
+
+            if isinstance(p.data, DTensor):
+                slices = get_slices_of_dtensor(
+                    target=p,
+                    local_rank=dist.get_rank(),
+                    shard_mesh=shard_mesh,
+                    shard_placements=shard_placements,
+                )
+                u_shard = u[slices]
+                u = DTensor.from_local(
+                    u_shard,
+                    device_mesh=p.device_mesh,
+                    placements=p.placements,
+                )
+
+            adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+            Muon._update_p(p, u, lr, adjusted_lr, weight_decay)
+
     def _update_g(self, p, g, group, momentum):
         # calc update
         state = self.state[p]
@@ -727,6 +849,9 @@ class Muon(torch.optim.Optimizer):
         p.data.add_(u, alpha=-adjusted_lr)
 
     def get_qk_clip_info(self, n, qk_logits):
+        if self.clip_config is None:
+            return None
+
         head_dim = self.clip_config.get('head_dim')
         threshold = self.clip_config.get('threshold')
         kind, layer_idx = parse_qk_layer(n)
@@ -737,6 +862,11 @@ class Muon(torch.optim.Optimizer):
             indices_key = 'q_indices' if 'q' in kind else 'k_indices'
             indices = self.clip_config.get(indices_key, []) or []
 
+            if isinstance(logit, DTensor):
+                # In TP settings, qk_logits may be DTensor
+                # We convert it to full tensor here for simplicity
+                logit = logit.full_tensor()
+
         return QKClipInfo(
             kind=kind,
             indices=indices,
@@ -835,22 +965,28 @@ class Muon(torch.optim.Optimizer):
                 _update_param(p, state, lr, adjusted_lr, weight_decay,
                               self.rank, self.compute_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(
-            params[0])].process_group)
+        if self.chunk_size == -1:
+            shard_ranks = dist.get_world_size(param_to_state[id(
+                params[0])].process_group)
+            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+        elif self.chunk_size > 0:
+            chunk_size = self.chunk_size
+        else:
+            raise ValueError("chunk_size must be -1 or a positive integer.")
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        overlap_step = self.overlap_step
-        for i in range(0, overlap_step):
+        warmup_step = self.warmup_step
+        for i in range(0, warmup_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
+            enqueue_computes(i + warmup_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)
@@ -866,7 +1002,7 @@ class Muon(torch.optim.Optimizer):
         amsgrad: bool,
         beta1: float,
         beta2: float,
-        lr: Union[float, torch.Tensor],
+        lr: float | torch.Tensor,
         weight_decay: float,
         eps: float,
         maximize: bool,
@@ -876,10 +1012,10 @@ class Muon(torch.optim.Optimizer):
 
         # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer
         # treating it as a scalar.
-        lr_dict: Optional[DeviceDict] = ({
+        lr_dict: DeviceDict | None = ({
             lr.device: lr
         } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else
-                                         None)
+                                      None)
         grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype(
             [
                 params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
@@ -926,6 +1062,159 @@ class Muon(torch.optim.Optimizer):
                 maximize=maximize,
             )
 
+    def _step_muon(self, group, qk_logits=None):
+        params = group["params"]
+        lr = group["lr"]
+        weight_decay = group["weight_decay"]
+        momentum = group["momentum"]
+        names = group["names"]
+
+        param_dtensors = []
+        param_tensors = []
+        name_dtensors = []
+        name_tensors = []
+
+        if self.use_distributed_muon:
+            self.distributed_muon(names=names,
+                                  params=params,
+                                  group=group,
+                                  lr=lr,
+                                  weight_decay=weight_decay,
+                                  momentum=momentum,
+                                  qk_logits=qk_logits)
+            return
+
+        for n, p in zip(names, params):
+            if p is None or p.grad is None:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(
+                        isinstance(placement, Replicate)
+                        for placement in p.placements):
+                    param_tensors.append(p)
+                    name_tensors.append(n)
+                else:
+                    param_dtensors.append(p)
+                    name_dtensors.append(n)
+            elif isinstance(p.data, torch.Tensor):
+                param_tensors.append(p)
+                name_tensors.append(n)
+            else:
+                raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+        logger.debug(
+            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors"
+        )
+
+        if len(param_dtensors) > 0:
+            if not dist.is_initialized():
+                raise RuntimeError(
+                    "Parallel Muon requires torch.distributed to be initialized."
+                )
+
+            # To support different placements, we group parameters by placements
+            # and run parallel Muon on each group.
+
+            placement_to_params = defaultdict(lambda: ([], []))
+            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
+
+            assert len(name_dtensors) == len(param_dtensors)
+            for n, p in zip(name_dtensors, param_dtensors):
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][0].append(n)
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][1].append(p)
+
+            for _, (names, params) in placement_to_params.items():
+                self.parallel(
+                    names,
+                    params,
+                    group,
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    momentum=momentum,
+                    qk_logits=qk_logits,
+                )
+
+        if len(param_tensors) > 0:
+            self.base(
+                name_tensors,
+                param_tensors,
+                group,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                qk_logits=qk_logits,
+            )
+
+    def _step_adamw_params(self, params, group):
+        params_with_grads = []
+        grads = []
+        moment1 = []
+        moment2 = []
+        max_exp_avg_sqs = []
+        state_steps = []
+        lr = group["lr"]
+        beta1, beta2 = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        weight_decay = group["weight_decay"]
+
+        for p in params:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            params_with_grads.append(p)
+            grads.append(g)
+            if "step" not in state:
+                state["step"] = (torch.zeros((),
+                                             dtype=torch.float32,
+                                             device=p.device))
+                state["moment1"] = torch.zeros_like(g)
+                state["moment2"] = torch.zeros_like(g)
+            moment1.append(state["moment1"])
+            moment2.append(state["moment2"])
+            if not isinstance(state["step"], torch.Tensor):
+                step_tensor = torch.tensor(state["step"],
+                                           dtype=torch.float32,
+                                           device=p.device)
+            else:
+                step_tensor = state["step"]
+            state_steps.append(step_tensor)
+
+        self._fused_adamw(
+            params_with_grads,
+            grads,
+            moment1,
+            moment2,
+            max_exp_avg_sqs,
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+    def _step_adamw(self, group):
+        params = group["params"]
+
+        # group params with it's type and placement
+        placement_to_params: dict[tuple[Placement | type,
+                                        DeviceMesh | None]] = defaultdict(list)
+        for p in params:
+            match p:
+                case DTensor():
+                    placement_to_params[tuple([p.placements,
+                                               p.device_mesh])].append(p)
+                case torch.Tensor():
+                    placement_to_params[tuple([torch.Tensor, None])].append(p)
+
+        for params in placement_to_params.values():
+            self._step_adamw_params(params, group)
+
     def step(self, closure=None, qk_logits=None):
         """Perform a single optimization step.
 
@@ -943,127 +1232,9 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         for group in self.param_groups:
-            params = group["params"]
-
             if group["use_muon"]:
-                ############################
-                #           Muon           #
-                ############################
-                lr = group["lr"]
-                weight_decay = group["weight_decay"]
-                momentum = group["momentum"]
-                names = group["names"]
-
-                param_dtensors = []
-                param_tensors = []
-                name_dtensors = []
-                name_tensors = []
-
-                for n, p in zip(names, params):
-                    if p is None or p.grad is None:
-                        continue
-                    if isinstance(p.data, DTensor):
-                        if all(
-                                isinstance(placement, Replicate)
-                                for placement in p.placements):
-                            param_tensors.append(p)
-                            name_tensors.append(n)
-                        else:
-                            param_dtensors.append(p)
-                            name_dtensors.append(n)
-                    elif isinstance(p.data, torch.Tensor):
-                        param_tensors.append(p)
-                        name_tensors.append(n)
-                    else:
-                        raise TypeError(
-                            f"Unsupported parameter type: {type(p.data)}")
-
-                if self.debug:
-                    print(
-                        f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
-                        flush=True,
-                    )
-
-                if len(param_dtensors) > 0:
-                    if not dist.is_initialized():
-                        raise RuntimeError(
-                            "Parallel Muon requires torch.distributed to be initialized."
-                        )
-
-                    self.parallel(
-                        name_dtensors,
-                        param_dtensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
-                if len(param_tensors) > 0:
-                    self.base(
-                        name_tensors,
-                        param_tensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
+                self._step_muon(group, qk_logits=qk_logits)
             else:
-                ############################
-                #       AdamW backup       #
-                ############################
-
-                params_with_grads = []
-                grads = []
-                moment1 = []
-                moment2 = []
-                max_exp_avg_sqs = []
-                state_steps = []
-                lr = group["lr"]
-                beta1, beta2 = group["adamw_betas"]
-                eps = group["adamw_eps"]
-                weight_decay = group["weight_decay"]
-
-                for p in params:
-                    g = p.grad
-                    if g is None:
-                        continue
-                    state = self.state[p]
-                    params_with_grads.append(p)
-                    grads.append(g)
-                    if "step" not in state:
-                        state["step"] = (torch.zeros((),
-                                                     dtype=torch.float32,
-                                                     device=p.device))
-                        state["moment1"] = torch.zeros_like(g)
-                        state["moment2"] = torch.zeros_like(g)
-                    moment1.append(state["moment1"])
-                    moment2.append(state["moment2"])
-                    if not isinstance(state["step"], torch.Tensor):
-                        step_tensor = torch.tensor(state["step"],
-                                                   dtype=torch.float32,
-                                                   device=p.device)
-                    else:
-                        step_tensor = state["step"]
-                    state_steps.append(step_tensor)
-
-                self._fused_adamw(
-                    params_with_grads,
-                    grads,
-                    moment1,
-                    moment2,
-                    max_exp_avg_sqs,
-                    state_steps,
-                    amsgrad=False,
-                    beta1=beta1,
-                    beta2=beta2,
-                    lr=lr,
-                    weight_decay=weight_decay,
-                    eps=eps,
-                    maximize=False,
-                )
+                self._step_adamw(group)
 
         return loss

build/torch29-cxx11-rocm63-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_811726c_dirty
-ops = torch.ops._optimizer_811726c_dirty
+from . import _optimizer_23d68bb_dirty
+ops = torch.ops._optimizer_23d68bb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_811726c_dirty::{op_name}"
+    return f"_optimizer_23d68bb_dirty::{op_name}"

build/torch29-cxx11-rocm63-x86_64-linux/optimizer/_optimizer_23d68bb_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d297c32252c7f030f3ec60ab1cc908cf145c8ecc710a25690a528d06115ab998
+size 1852184

build/torch29-cxx11-rocm63-x86_64-linux/optimizer/_optimizer_811726c_dirty.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0661740cd0f97ca56ef83979c5a5fa059bcba411148f89d836e9305065578e73
-size 1749264

build/torch29-cxx11-rocm63-x86_64-linux/optimizer/distributed/utils.py

@@ -0,0 +1,174 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import (Placement, Shard,
+                                                      _StridedShard)
+
+
+def get_slices_of_dtensor(
+    target: DTensor | torch.Tensor,
+    local_rank: int,
+    shard_mesh: DeviceMesh,
+    shard_placements: tuple[Placement],
+) -> tuple[slice]:
+    """
+    Get the slice of local tensor for a given rank from a tensor.
+    Args:
+        target (DTensor | torch.Tensor): The target tensor.
+        rank (int): The local rank of the shard group.
+        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
+        shard_placements (tuple[Placement]): The shard placements.
+    """
+
+    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
+
+    # find the global rank of the local rank in the shard mesh
+    rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank]
+
+    rank_coords = (shard_mesh.mesh == rank).nonzero()
+
+    assert len(rank_coords) == 1
+    rank_coords = tuple(rank_coords[0].tolist())
+
+    assert len(rank_coords) == len(shard_placements)
+
+    # Caution: Assuming replicate-to-shard of the shard mesh goes with
+    # left-to-right sharding. This is ensured by the sorting logic of
+    # construct_shard_mesh function.
+    for i, (rank_coord,
+            placement) in enumerate(zip(rank_coords, shard_placements)):
+        assert isinstance(placement, Shard)
+
+        num_ranks = shard_mesh.mesh.shape[i]
+
+        dim = placement.dim
+        dim_size = (slices[dim].stop - slices[dim].start)
+
+        if dim_size % num_ranks != 0:
+            raise NotImplementedError(
+                f"Dimension size {dim_size} is not divisible "
+                f"by number of ranks {num_ranks} for shard "
+                f"placement on dim {dim}.")
+
+        shard_size = dim_size // num_ranks
+
+        start = slices[dim].start + rank_coord * shard_size
+        end = start + shard_size
+
+        assert start < end <= slices[dim].stop
+
+        slices[dim] = slice(start, end)
+
+    return tuple(slices)
+
+
+_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, ProcessGroup]] = dict()
+
+
+def construct_shard_mesh(
+    placements: tuple[Placement],
+    mesh: DeviceMesh,
+) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
+    """
+    Construct Shard Mesh and Placements for unsharding.
+    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
+    """
+    my_rank = dist.get_rank()
+
+    assert mesh.mesh.device.type == 'cpu'
+
+    # Copy mesh to avoid modifying the original mesh
+    mesh = mesh.mesh.clone()
+
+    # 1. Sort placements. Replicate first, then Shard by dim ascending.
+
+    # For Shard, strided shard comes after regular shard on the same dim
+    # to preserve left-to-right order of replicate-to-shard.
+    # This is because that strided shard is using stride to represent
+    # more fine-grained sharding on the same dim.
+    # Please check the URL below for _StridedShard.
+    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
+
+    def placement_sort_key(
+        placement_with_index: tuple[float, Placement]
+    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
+        index, placement = placement_with_index
+        is_replicate = placement.is_replicate()
+        is_shard = placement.is_shard()
+        is_partial = placement.is_partial()
+
+        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
+        assert not is_partial, "Partial placement is not supported."
+
+        if is_replicate:
+            return (-1.0, 0, index)
+        elif is_shard:
+            if isinstance(placement, _StridedShard):
+                return (placement.dim, 1 / placement.split_factor, index)
+            return (placement.dim, 0, index)
+        else:
+            raise TypeError(f"Unknown placement type: {type(placement)}")
+
+    placements_with_index: list[tuple[int,
+                                      Placement]] = list(enumerate(placements))
+    placements_with_index = sorted(placements_with_index,
+                                   key=placement_sort_key)
+
+    sorted_indices, sorted_placements = zip(*placements_with_index)
+
+    # 2. Permute mesh according to sorted placements.
+    sorted_mesh = mesh.permute(sorted_indices)
+
+    # 3. Collect list of shard meshes by removing replicate dims
+    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
+    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
+    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
+
+    # merge replicate dims
+    # shard_meshes became a list of shard meshes with a length of replicate degree
+    if num_replicates > 0:
+        sorted_mesh = sorted_mesh.flatten(
+            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
+        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
+    else:
+        shard_meshes = [sorted_mesh]
+    shard_placements = sorted_placements[num_replicates:]
+
+    # assume all shard placements are different
+    assert len(shard_placements) == len(set(shard_placements))
+
+    # 4. Construct ProcessGroups
+    # Caution: all groups should be created in the same order in all processes,
+    # even though each process only needs its own group.
+
+    # To use tensor as dict key, convert it to tuple
+    def tensor_to_tuple(t):
+        if isinstance(t, torch.Tensor):
+            t = t.tolist()
+        if isinstance(t, list):
+            return tuple(tensor_to_tuple(x) for x in t)
+        return t
+
+    my_shard_mesh_as_tuple = None
+    for shard_mesh in shard_meshes:
+        assert isinstance(shard_mesh, torch.Tensor)
+        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
+
+        if (my_rank == shard_mesh).any().item():
+            assert my_shard_mesh_as_tuple is None
+            my_shard_mesh_as_tuple = shard_mesh_as_tuple
+
+        # update global cache
+        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
+            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
+            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
+                DeviceMesh(device_type="cuda", mesh=shard_mesh),
+                shard_process_group,
+            )
+
+    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
+        my_shard_mesh_as_tuple]
+
+    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch29-cxx11-rocm63-x86_64-linux/optimizer/muon.py

@@ -1,18 +1,24 @@
 import logging
 import math
 import types
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor, Replicate, Shard
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor.placement_types import Placement
 
+from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
 from .matmul_transpose_triton import matmul_transpose_assign
 
 logger = logging.getLogger(__name__)
 
 COMM_DTYPE = torch.bfloat16
+DEFAULT_CHUNK_SIZE_RATIO = 4
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -62,23 +68,39 @@ def _zeropower_via_newtonschulz5(G, steps):
 @dataclass
 class _muon_state:
     # TODO: use Optional
-    worker_rank: int | None = None
+    worker_rank: int
+    process_group: ProcessGroup
+    shard_mesh: DeviceMesh
+    shard_placements: tuple[Placement, ...]
+    name: str
+    qk_clip_state: torch.Tensor | None = None
     gathered_grad: torch.Tensor | None = None
     scattered_u: DTensor | None = None
     computed_u: torch.Tensor | None = None
     gather_event: torch.cuda.Event | None = None
     compute_event: torch.cuda.Event | None = None
     scatter_event: torch.cuda.Event | None = None
-    process_group = None
-    qk_clip_state = None
 
 
-def split_elems_for_src(param, src_rank, num_ranks) -> int:
-    rows = param.shape[0]
-    cols = int(param.numel() // rows)
-    base, rem = divmod(rows, num_ranks)
-    my_rows = base + (1 if src_rank < rem else 0)
-    return my_rows * cols
+def numel_for_rank(
+    param: DTensor,
+    local_rank: int,
+    state: _muon_state,
+) -> int:
+    slices = get_slices_of_dtensor(
+        param,
+        local_rank,
+        state.shard_mesh,
+        state.shard_placements,
+    )
+
+    numel = 1
+    for s, dim in zip(slices, param.shape):
+        start, stop, step = s.indices(dim)
+        length = max(0, (stop - start + (step - 1)) // step)
+        numel *= length
+
+    return numel
 
 
 @torch.no_grad()
@@ -91,8 +113,7 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
         for p in params:
             state = param_to_state[id(p)]
             if rank == state.worker_rank:
-                num_ranks = dist.get_world_size(group=state.process_group)
-                state.gathered_grad = torch.empty(p.grad.numel(),
+                state.gathered_grad = torch.empty(p.shape,
                                                   dtype=COMM_DTYPE,
                                                   device="cuda")
             else:
@@ -121,11 +142,11 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             state = param_to_state[id(p)]
             dst = state.worker_rank
             assert dst < num_ranks
-            shard_elems = split_elems_for_src(p, rank, num_ranks)
+            shard_elems = numel_for_rank(p, rank, state)
             g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            g = g.to_local().to(COMM_DTYPE).contiguous()
             assert g.numel() == shard_elems
-            per_dst[dst].append(g)
+            per_dst[dst].append(g.view(-1))
             send_counts[dst] += shard_elems
 
         assert any(
@@ -148,13 +169,18 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                total += split_elems_for_src(p, src, num_ranks)
+                total += numel_for_rank(p, src, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
         #All2All
+        logger.debug(f"send_buf size: {send_buf.numel()}, "
+                     f"recv_buf size: {recv_buf.numel()}, "
+                     f"recv_counts: {recv_counts}, "
+                     f"send_counts: {send_counts}, "
+                     f"process_group: {str(process_group)}")
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -179,7 +205,6 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
@@ -189,22 +214,28 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                n = split_elems_for_src(p, src, num_ranks)
+
+                # get the slice of the full dtensor corresponding to rank src.
+                slices = get_slices_of_dtensor(state.gathered_grad, src,
+                                               state.shard_mesh,
+                                               state.shard_placements)
+
+                dst = state.gathered_grad[slices]
+                assert dst._base is state.gathered_grad
+
+                n = dst.numel()
                 assert n > 0
 
                 sg = recv_buf.narrow(0, off + inner_off, n)
-                woff = write_offsets[id(p)]
-                dst = state.gathered_grad.narrow(0, woff, n)
+                sg = sg.reshape_as(dst)
                 dst.copy_(sg)
 
-                write_offsets[id(p)] += n
                 inner_off += n
             off += block
 
         for p in params:
             state = param_to_state[id(p)]
             if state.worker_rank == rank:
-                state.gathered_grad = state.gathered_grad.view_as(p)
                 state.gather_event = torch.cuda.Event()
                 state.gather_event.record(comm_stream)
             else:
@@ -277,14 +308,19 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert state.computed_u is not None
 
-                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous()
 
                 offset = 0
                 for dst in range(num_ranks):
-                    n = split_elems_for_src(p, dst, num_ranks)
+                    # get the slice of the full tensor corresponding to rank dst.
+                    slices = get_slices_of_dtensor(u_full, dst,
+                                                   state.shard_mesh,
+                                                   state.shard_placements)
+                    su = u_full[slices].flatten()
+
+                    n = su.numel()
                     assert n > 0
 
-                    su = u_full.narrow(0, offset, n)
                     per_dst[dst].append(su)
                     send_counts[dst] += n
                     offset += n
@@ -313,7 +349,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                total += split_elems_for_src(p, rank, num_ranks)
+                total += numel_for_rank(p, rank, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
@@ -357,7 +393,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                n = split_elems_for_src(p, rank, num_ranks)
+                n = numel_for_rank(p, rank, state)
                 assert n > 0
 
                 flat_local = recv_buf.narrow(0, off + inner_off,
@@ -398,11 +434,23 @@ def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
         state.scattered_u = None
         u_dtensor = None
 
-        scales_full = Muon._compute_scales(p, state.qk_clip_state)
+        scales_full = Muon._compute_scales(
+            p,
+            state.qk_clip_state) if state.qk_clip_state is not None else None
         if scales_full is not None:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            local_rank = dist.get_rank(group=state.process_group)
-            scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank]
+            # Have to slice scales_full among dim 0
+            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
+                                                  state.shard_placements)
+            ratio = p.shape[0] // scales_full.shape[0]
+            scales_slice = slice(
+                None if weight_slices[0].start is None else
+                weight_slices[0].start // ratio,
+                None if weight_slices[0].stop is None else
+                weight_slices[0].stop // ratio,
+                None,
+            )
+
+            scales_local = scales_full[scales_slice]
             scales_local = DTensor.from_local(
                 scales_local,
                 placements=p.placements,
@@ -478,11 +526,11 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
 @dataclass
 class QKClipInfo:
     """Per-parameter dynamic info computed from config + runtime logits."""
-    kind: Optional[str]  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
-    indices: List[int]  # which heads to consider for clipping
+    kind: str | None  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
+    indices: list[int]  # which heads to consider for clipping
     head_dim: int  # from config
     threshold: float  # from config
-    logit: Optional[torch.Tensor]
+    logit: torch.Tensor | None
 
 
 class Muon(torch.optim.Optimizer):
@@ -525,11 +573,16 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
-        overlap_step : How many all2all gather, compute operations are launched in advance
-                       before the corresponding all2all scatter steps begin.
-                       A higher overlap_step increases memory usage but can improve
-                       performance by overlapping communication.
-                       Parallel muon only. 
+        warmup_step : How many all2all gather, compute operations are launched in advance
+                      before the corresponding all2all scatter steps begin.
+                      A higher warmup_step increases memory usage but can improve
+                      performance by overlapping communication.
+                      Parallel muon only.
+        chunk_size : Batch size of parameters to process in each
+                     all2all gather/compute/scatter step.
+                     Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
+        use_distributed_muon: Use distributed muon by Liu et al. (2024).
+                              For testing purpose only.
     """
 
     def __init__(self,
@@ -549,7 +602,9 @@ class Muon(torch.optim.Optimizer):
                      "head_dim": 128,
                      "threshold": 100
                  },
-                 overlap_step=5):
+                 warmup_step=5,
+                 chunk_size=-1,
+                 use_distributed_muon=False):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -579,7 +634,9 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
-        self.overlap_step = overlap_step
+        self.warmup_step = warmup_step
+        self.chunk_size = chunk_size
+        self.use_distributed_muon = use_distributed_muon
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -597,6 +654,12 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
+    def set_rank_once(self, rank):
+        if self.rank is None:
+            self.rank = rank
+        else:
+            assert self.rank == rank
+
     def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
@@ -604,26 +667,13 @@ class Muon(torch.optim.Optimizer):
         assert isinstance(
             p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0), ):
-            # Case for FSDP
-            process_group = p.device_mesh.get_group(mesh_dim=0)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
-        elif p.placements == (Replicate(), Shard(dim=0)):
-            # Case for HSDP
-            process_group = p.device_mesh.get_group(mesh_dim=1)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if self.rank in shard_mesh:
-                    return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
-        else:
-            raise ValueError(f"Unsupported placements ({p.placements}).")
+        shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+            p.placements, p.device_mesh)
+
+        # set rank with the local rank in the shard process group
+        self.set_rank_once(dist.get_rank(group=shard_pg))
+
+        return shard_mesh, shard_pg, shard_placements
 
     def init_state_and_assign_params(self, names, params, group, qk_logits):
         param_to_state = {}
@@ -655,23 +705,32 @@ class Muon(torch.optim.Optimizer):
         ordered_params = list(params_sorted)
 
         round_robin = 0
-        mesh = None
-        shard_mesh = None
-        process_group = None
+        mesh = ordered_params[0].device_mesh
+        placements = ordered_params[0].placements
+
+        shard_mesh, shard_pg, shard_placements = self.get_shard_mesh(
+            ordered_params[0])
+        shard_mesh_flattened = shard_mesh.mesh.flatten()
+        num_ranks = dist.get_world_size(group=shard_pg)
+
         for n, p in zip(ordered_names, ordered_params):
-            if mesh is None:
-                mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p)
-            elif mesh != p.device_mesh:
+            if mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
-            num_ranks = dist.get_world_size(group=process_group)
-            param_to_state[id(p)] = _muon_state()
-            param_to_state[id(
-                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
-            param_to_state[id(p)].process_group = process_group
+            if placements != p.placements:
+                raise ValueError("All parameters must have same placements.")
+
+            worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
+            round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-            param_to_state[id(p)].qk_clip_state = qk_clip_state
-            round_robin = (round_robin + 1) % len(shard_mesh)
+
+            param_to_state[id(p)] = _muon_state(
+                worker_rank=worker_rank,
+                process_group=shard_pg,
+                shard_mesh=shard_mesh,
+                shard_placements=shard_placements,
+                name=n,
+                qk_clip_state=qk_clip_state,
+            )
 
         return param_to_state, ordered_params
 
@@ -705,10 +764,73 @@ class Muon(torch.optim.Optimizer):
 
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
 
-            scales_full = self._compute_scales(p, qk_clip_state)
+            scales_full = self._compute_scales(
+                p, qk_clip_state) if qk_clip_state is not None else None
             if scales_full is not None:
                 Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
 
+    def distributed_muon(
+        self,
+        names: list[str],
+        params: list[torch.nn.Parameter],
+        group: dict[str, Any],
+        lr: float,
+        weight_decay: float,
+        momentum: float,
+        qk_logits: list[torch.Tensor | DTensor] | None,
+    ):
+        """ Implementation of Distributed Muon by Liu et al. """
+        if qk_logits is not None:
+            raise NotImplementedError("QK clipping is not supported yet")
+
+        if isinstance(params[0], DTensor):
+            shard_mesh, _, shard_placements = construct_shard_mesh(
+                placements=params[0].placements,
+                mesh=params[0].device_mesh,
+            )
+
+        for n, p in zip(names, params):
+            g = p.grad
+            if g is None:
+                continue
+            if g.ndim > 2:
+                g = g.view(g.size(0), -1)
+            assert g is not None
+
+            # calc update
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(momentum).add_(g)
+            if group["nesterov"]:
+                g = g.add(buf, alpha=momentum)
+            else:
+                g = buf
+
+            # Gather G
+            if isinstance(p.data, DTensor):
+                g = g.full_tensor()
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                             steps=group["ns_steps"])
+
+            if isinstance(p.data, DTensor):
+                slices = get_slices_of_dtensor(
+                    target=p,
+                    local_rank=dist.get_rank(),
+                    shard_mesh=shard_mesh,
+                    shard_placements=shard_placements,
+                )
+                u_shard = u[slices]
+                u = DTensor.from_local(
+                    u_shard,
+                    device_mesh=p.device_mesh,
+                    placements=p.placements,
+                )
+
+            adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+            Muon._update_p(p, u, lr, adjusted_lr, weight_decay)
+
     def _update_g(self, p, g, group, momentum):
         # calc update
         state = self.state[p]
@@ -727,6 +849,9 @@ class Muon(torch.optim.Optimizer):
         p.data.add_(u, alpha=-adjusted_lr)
 
     def get_qk_clip_info(self, n, qk_logits):
+        if self.clip_config is None:
+            return None
+
         head_dim = self.clip_config.get('head_dim')
         threshold = self.clip_config.get('threshold')
         kind, layer_idx = parse_qk_layer(n)
@@ -737,6 +862,11 @@ class Muon(torch.optim.Optimizer):
             indices_key = 'q_indices' if 'q' in kind else 'k_indices'
             indices = self.clip_config.get(indices_key, []) or []
 
+            if isinstance(logit, DTensor):
+                # In TP settings, qk_logits may be DTensor
+                # We convert it to full tensor here for simplicity
+                logit = logit.full_tensor()
+
         return QKClipInfo(
             kind=kind,
             indices=indices,
@@ -835,22 +965,28 @@ class Muon(torch.optim.Optimizer):
                 _update_param(p, state, lr, adjusted_lr, weight_decay,
                               self.rank, self.compute_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(
-            params[0])].process_group)
+        if self.chunk_size == -1:
+            shard_ranks = dist.get_world_size(param_to_state[id(
+                params[0])].process_group)
+            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+        elif self.chunk_size > 0:
+            chunk_size = self.chunk_size
+        else:
+            raise ValueError("chunk_size must be -1 or a positive integer.")
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        overlap_step = self.overlap_step
-        for i in range(0, overlap_step):
+        warmup_step = self.warmup_step
+        for i in range(0, warmup_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
+            enqueue_computes(i + warmup_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)
@@ -866,7 +1002,7 @@ class Muon(torch.optim.Optimizer):
         amsgrad: bool,
         beta1: float,
         beta2: float,
-        lr: Union[float, torch.Tensor],
+        lr: float | torch.Tensor,
         weight_decay: float,
         eps: float,
         maximize: bool,
@@ -876,10 +1012,10 @@ class Muon(torch.optim.Optimizer):
 
         # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer
         # treating it as a scalar.
-        lr_dict: Optional[DeviceDict] = ({
+        lr_dict: DeviceDict | None = ({
             lr.device: lr
         } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else
-                                         None)
+                                      None)
         grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype(
             [
                 params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
@@ -926,6 +1062,159 @@ class Muon(torch.optim.Optimizer):
                 maximize=maximize,
             )
 
+    def _step_muon(self, group, qk_logits=None):
+        params = group["params"]
+        lr = group["lr"]
+        weight_decay = group["weight_decay"]
+        momentum = group["momentum"]
+        names = group["names"]
+
+        param_dtensors = []
+        param_tensors = []
+        name_dtensors = []
+        name_tensors = []
+
+        if self.use_distributed_muon:
+            self.distributed_muon(names=names,
+                                  params=params,
+                                  group=group,
+                                  lr=lr,
+                                  weight_decay=weight_decay,
+                                  momentum=momentum,
+                                  qk_logits=qk_logits)
+            return
+
+        for n, p in zip(names, params):
+            if p is None or p.grad is None:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(
+                        isinstance(placement, Replicate)
+                        for placement in p.placements):
+                    param_tensors.append(p)
+                    name_tensors.append(n)
+                else:
+                    param_dtensors.append(p)
+                    name_dtensors.append(n)
+            elif isinstance(p.data, torch.Tensor):
+                param_tensors.append(p)
+                name_tensors.append(n)
+            else:
+                raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+        logger.debug(
+            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors"
+        )
+
+        if len(param_dtensors) > 0:
+            if not dist.is_initialized():
+                raise RuntimeError(
+                    "Parallel Muon requires torch.distributed to be initialized."
+                )
+
+            # To support different placements, we group parameters by placements
+            # and run parallel Muon on each group.
+
+            placement_to_params = defaultdict(lambda: ([], []))
+            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
+
+            assert len(name_dtensors) == len(param_dtensors)
+            for n, p in zip(name_dtensors, param_dtensors):
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][0].append(n)
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][1].append(p)
+
+            for _, (names, params) in placement_to_params.items():
+                self.parallel(
+                    names,
+                    params,
+                    group,
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    momentum=momentum,
+                    qk_logits=qk_logits,
+                )
+
+        if len(param_tensors) > 0:
+            self.base(
+                name_tensors,
+                param_tensors,
+                group,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                qk_logits=qk_logits,
+            )
+
+    def _step_adamw_params(self, params, group):
+        params_with_grads = []
+        grads = []
+        moment1 = []
+        moment2 = []
+        max_exp_avg_sqs = []
+        state_steps = []
+        lr = group["lr"]
+        beta1, beta2 = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        weight_decay = group["weight_decay"]
+
+        for p in params:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            params_with_grads.append(p)
+            grads.append(g)
+            if "step" not in state:
+                state["step"] = (torch.zeros((),
+                                             dtype=torch.float32,
+                                             device=p.device))
+                state["moment1"] = torch.zeros_like(g)
+                state["moment2"] = torch.zeros_like(g)
+            moment1.append(state["moment1"])
+            moment2.append(state["moment2"])
+            if not isinstance(state["step"], torch.Tensor):
+                step_tensor = torch.tensor(state["step"],
+                                           dtype=torch.float32,
+                                           device=p.device)
+            else:
+                step_tensor = state["step"]
+            state_steps.append(step_tensor)
+
+        self._fused_adamw(
+            params_with_grads,
+            grads,
+            moment1,
+            moment2,
+            max_exp_avg_sqs,
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+    def _step_adamw(self, group):
+        params = group["params"]
+
+        # group params with it's type and placement
+        placement_to_params: dict[tuple[Placement | type,
+                                        DeviceMesh | None]] = defaultdict(list)
+        for p in params:
+            match p:
+                case DTensor():
+                    placement_to_params[tuple([p.placements,
+                                               p.device_mesh])].append(p)
+                case torch.Tensor():
+                    placement_to_params[tuple([torch.Tensor, None])].append(p)
+
+        for params in placement_to_params.values():
+            self._step_adamw_params(params, group)
+
     def step(self, closure=None, qk_logits=None):
         """Perform a single optimization step.
 
@@ -943,127 +1232,9 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         for group in self.param_groups:
-            params = group["params"]
-
             if group["use_muon"]:
-                ############################
-                #           Muon           #
-                ############################
-                lr = group["lr"]
-                weight_decay = group["weight_decay"]
-                momentum = group["momentum"]
-                names = group["names"]
-
-                param_dtensors = []
-                param_tensors = []
-                name_dtensors = []
-                name_tensors = []
-
-                for n, p in zip(names, params):
-                    if p is None or p.grad is None:
-                        continue
-                    if isinstance(p.data, DTensor):
-                        if all(
-                                isinstance(placement, Replicate)
-                                for placement in p.placements):
-                            param_tensors.append(p)
-                            name_tensors.append(n)
-                        else:
-                            param_dtensors.append(p)
-                            name_dtensors.append(n)
-                    elif isinstance(p.data, torch.Tensor):
-                        param_tensors.append(p)
-                        name_tensors.append(n)
-                    else:
-                        raise TypeError(
-                            f"Unsupported parameter type: {type(p.data)}")
-
-                if self.debug:
-                    print(
-                        f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
-                        flush=True,
-                    )
-
-                if len(param_dtensors) > 0:
-                    if not dist.is_initialized():
-                        raise RuntimeError(
-                            "Parallel Muon requires torch.distributed to be initialized."
-                        )
-
-                    self.parallel(
-                        name_dtensors,
-                        param_dtensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
-                if len(param_tensors) > 0:
-                    self.base(
-                        name_tensors,
-                        param_tensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
+                self._step_muon(group, qk_logits=qk_logits)
             else:
-                ############################
-                #       AdamW backup       #
-                ############################
-
-                params_with_grads = []
-                grads = []
-                moment1 = []
-                moment2 = []
-                max_exp_avg_sqs = []
-                state_steps = []
-                lr = group["lr"]
-                beta1, beta2 = group["adamw_betas"]
-                eps = group["adamw_eps"]
-                weight_decay = group["weight_decay"]
-
-                for p in params:
-                    g = p.grad
-                    if g is None:
-                        continue
-                    state = self.state[p]
-                    params_with_grads.append(p)
-                    grads.append(g)
-                    if "step" not in state:
-                        state["step"] = (torch.zeros((),
-                                                     dtype=torch.float32,
-                                                     device=p.device))
-                        state["moment1"] = torch.zeros_like(g)
-                        state["moment2"] = torch.zeros_like(g)
-                    moment1.append(state["moment1"])
-                    moment2.append(state["moment2"])
-                    if not isinstance(state["step"], torch.Tensor):
-                        step_tensor = torch.tensor(state["step"],
-                                                   dtype=torch.float32,
-                                                   device=p.device)
-                    else:
-                        step_tensor = state["step"]
-                    state_steps.append(step_tensor)
-
-                self._fused_adamw(
-                    params_with_grads,
-                    grads,
-                    moment1,
-                    moment2,
-                    max_exp_avg_sqs,
-                    state_steps,
-                    amsgrad=False,
-                    beta1=beta1,
-                    beta2=beta2,
-                    lr=lr,
-                    weight_decay=weight_decay,
-                    eps=eps,
-                    maximize=False,
-                )
+                self._step_adamw(group)
 
         return loss

build/torch29-cxx11-rocm64-x86_64-linux/optimizer/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _optimizer_811726c_dirty
-ops = torch.ops._optimizer_811726c_dirty
+from . import _optimizer_23d68bb_dirty
+ops = torch.ops._optimizer_23d68bb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_811726c_dirty::{op_name}"
+    return f"_optimizer_23d68bb_dirty::{op_name}"

build/torch29-cxx11-rocm64-x86_64-linux/optimizer/_optimizer_23d68bb_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8de22742ad0d387021a7b812ee3b7d0c8c54191914c8c0469886f6d2c082e9e3
+size 1852272

build/torch29-cxx11-rocm64-x86_64-linux/optimizer/_optimizer_811726c_dirty.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:08b55491319446b12d0d890926506639640414edcba945e0f71afef0fac369d5
-size 1749352

build/torch29-cxx11-rocm64-x86_64-linux/optimizer/distributed/utils.py

@@ -0,0 +1,174 @@
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import (Placement, Shard,
+                                                      _StridedShard)
+
+
+def get_slices_of_dtensor(
+    target: DTensor | torch.Tensor,
+    local_rank: int,
+    shard_mesh: DeviceMesh,
+    shard_placements: tuple[Placement],
+) -> tuple[slice]:
+    """
+    Get the slice of local tensor for a given rank from a tensor.
+    Args:
+        target (DTensor | torch.Tensor): The target tensor.
+        rank (int): The local rank of the shard group.
+        shard_mesh (DeviceMesh): The shard mesh. It consists of global ranks.
+        shard_placements (tuple[Placement]): The shard placements.
+    """
+
+    slices: list[slice] = [slice(0, dim_size) for dim_size in target.size()]
+
+    # find the global rank of the local rank in the shard mesh
+    rank = sorted(shard_mesh.mesh.flatten().tolist())[local_rank]
+
+    rank_coords = (shard_mesh.mesh == rank).nonzero()
+
+    assert len(rank_coords) == 1
+    rank_coords = tuple(rank_coords[0].tolist())
+
+    assert len(rank_coords) == len(shard_placements)
+
+    # Caution: Assuming replicate-to-shard of the shard mesh goes with
+    # left-to-right sharding. This is ensured by the sorting logic of
+    # construct_shard_mesh function.
+    for i, (rank_coord,
+            placement) in enumerate(zip(rank_coords, shard_placements)):
+        assert isinstance(placement, Shard)
+
+        num_ranks = shard_mesh.mesh.shape[i]
+
+        dim = placement.dim
+        dim_size = (slices[dim].stop - slices[dim].start)
+
+        if dim_size % num_ranks != 0:
+            raise NotImplementedError(
+                f"Dimension size {dim_size} is not divisible "
+                f"by number of ranks {num_ranks} for shard "
+                f"placement on dim {dim}.")
+
+        shard_size = dim_size // num_ranks
+
+        start = slices[dim].start + rank_coord * shard_size
+        end = start + shard_size
+
+        assert start < end <= slices[dim].stop
+
+        slices[dim] = slice(start, end)
+
+    return tuple(slices)
+
+
+_ranks_to_dist_cache: dict[tuple[int, ...], tuple[DeviceMesh, ProcessGroup]] = dict()
+
+
+def construct_shard_mesh(
+    placements: tuple[Placement],
+    mesh: DeviceMesh,
+) -> (DeviceMesh, ProcessGroup, tuple[Placement]):
+    """
+    Construct Shard Mesh and Placements for unsharding.
+    It removes Replicate placements and constructs a new Mesh and ProcessGroup.
+    """
+    my_rank = dist.get_rank()
+
+    assert mesh.mesh.device.type == 'cpu'
+
+    # Copy mesh to avoid modifying the original mesh
+    mesh = mesh.mesh.clone()
+
+    # 1. Sort placements. Replicate first, then Shard by dim ascending.
+
+    # For Shard, strided shard comes after regular shard on the same dim
+    # to preserve left-to-right order of replicate-to-shard.
+    # This is because that strided shard is using stride to represent
+    # more fine-grained sharding on the same dim.
+    # Please check the URL below for _StridedShard.
+    # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/tensor/placement_types.py#L366
+
+    def placement_sort_key(
+        placement_with_index: tuple[float, Placement]
+    ) -> tuple[int, float, int]:  # (dim, split factor, original index)
+        index, placement = placement_with_index
+        is_replicate = placement.is_replicate()
+        is_shard = placement.is_shard()
+        is_partial = placement.is_partial()
+
+        assert is_replicate or is_shard, f"Unsupported placement type: {type(placement)}"
+        assert not is_partial, "Partial placement is not supported."
+
+        if is_replicate:
+            return (-1.0, 0, index)
+        elif is_shard:
+            if isinstance(placement, _StridedShard):
+                return (placement.dim, 1 / placement.split_factor, index)
+            return (placement.dim, 0, index)
+        else:
+            raise TypeError(f"Unknown placement type: {type(placement)}")
+
+    placements_with_index: list[tuple[int,
+                                      Placement]] = list(enumerate(placements))
+    placements_with_index = sorted(placements_with_index,
+                                   key=placement_sort_key)
+
+    sorted_indices, sorted_placements = zip(*placements_with_index)
+
+    # 2. Permute mesh according to sorted placements.
+    sorted_mesh = mesh.permute(sorted_indices)
+
+    # 3. Collect list of shard meshes by removing replicate dims
+    # For example, (2, 3, 4, 4) with placements [R, R, S(0), S(1)]
+    # shard_meshes should be list with 2 * 3 = 6 shard meshes of shape (4, 4)
+    num_replicates = sum(1 for p in sorted_placements if p.is_replicate())
+
+    # merge replicate dims
+    # shard_meshes became a list of shard meshes with a length of replicate degree
+    if num_replicates > 0:
+        sorted_mesh = sorted_mesh.flatten(
+            0, num_replicates - 1) if num_replicates > 1 else sorted_mesh
+        shard_meshes = list(torch.unbind(sorted_mesh, dim=0))
+    else:
+        shard_meshes = [sorted_mesh]
+    shard_placements = sorted_placements[num_replicates:]
+
+    # assume all shard placements are different
+    assert len(shard_placements) == len(set(shard_placements))
+
+    # 4. Construct ProcessGroups
+    # Caution: all groups should be created in the same order in all processes,
+    # even though each process only needs its own group.
+
+    # To use tensor as dict key, convert it to tuple
+    def tensor_to_tuple(t):
+        if isinstance(t, torch.Tensor):
+            t = t.tolist()
+        if isinstance(t, list):
+            return tuple(tensor_to_tuple(x) for x in t)
+        return t
+
+    my_shard_mesh_as_tuple = None
+    for shard_mesh in shard_meshes:
+        assert isinstance(shard_mesh, torch.Tensor)
+        shard_mesh_as_tuple = tensor_to_tuple(shard_mesh)
+
+        if (my_rank == shard_mesh).any().item():
+            assert my_shard_mesh_as_tuple is None
+            my_shard_mesh_as_tuple = shard_mesh_as_tuple
+
+        # update global cache
+        if shard_mesh_as_tuple not in _ranks_to_dist_cache:
+            shard_process_group = dist.new_group(shard_mesh.flatten().tolist())
+            _ranks_to_dist_cache[shard_mesh_as_tuple] = (
+                DeviceMesh(device_type="cuda", mesh=shard_mesh),
+                shard_process_group,
+            )
+
+    my_shard_mesh, my_shard_process_group = _ranks_to_dist_cache[
+        my_shard_mesh_as_tuple]
+
+    return my_shard_mesh, my_shard_process_group, shard_placements

build/torch29-cxx11-rocm64-x86_64-linux/optimizer/muon.py

@@ -1,18 +1,24 @@
 import logging
 import math
 import types
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union, cast
+from typing import Any, cast
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import DTensor, Replicate, Shard
+from torch.distributed import ProcessGroup
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor.placement_types import Placement
 
+from .distributed.utils import construct_shard_mesh, get_slices_of_dtensor
 from .matmul_transpose_triton import matmul_transpose_assign
 
 logger = logging.getLogger(__name__)
 
 COMM_DTYPE = torch.bfloat16
+DEFAULT_CHUNK_SIZE_RATIO = 4
 
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
@@ -62,23 +68,39 @@ def _zeropower_via_newtonschulz5(G, steps):
 @dataclass
 class _muon_state:
     # TODO: use Optional
-    worker_rank: int | None = None
+    worker_rank: int
+    process_group: ProcessGroup
+    shard_mesh: DeviceMesh
+    shard_placements: tuple[Placement, ...]
+    name: str
+    qk_clip_state: torch.Tensor | None = None
     gathered_grad: torch.Tensor | None = None
     scattered_u: DTensor | None = None
     computed_u: torch.Tensor | None = None
     gather_event: torch.cuda.Event | None = None
     compute_event: torch.cuda.Event | None = None
     scatter_event: torch.cuda.Event | None = None
-    process_group = None
-    qk_clip_state = None
 
 
-def split_elems_for_src(param, src_rank, num_ranks) -> int:
-    rows = param.shape[0]
-    cols = int(param.numel() // rows)
-    base, rem = divmod(rows, num_ranks)
-    my_rows = base + (1 if src_rank < rem else 0)
-    return my_rows * cols
+def numel_for_rank(
+    param: DTensor,
+    local_rank: int,
+    state: _muon_state,
+) -> int:
+    slices = get_slices_of_dtensor(
+        param,
+        local_rank,
+        state.shard_mesh,
+        state.shard_placements,
+    )
+
+    numel = 1
+    for s, dim in zip(slices, param.shape):
+        start, stop, step = s.indices(dim)
+        length = max(0, (stop - start + (step - 1)) // step)
+        numel *= length
+
+    return numel
 
 
 @torch.no_grad()
@@ -91,8 +113,7 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
         for p in params:
             state = param_to_state[id(p)]
             if rank == state.worker_rank:
-                num_ranks = dist.get_world_size(group=state.process_group)
-                state.gathered_grad = torch.empty(p.grad.numel(),
+                state.gathered_grad = torch.empty(p.shape,
                                                   dtype=COMM_DTYPE,
                                                   device="cuda")
             else:
@@ -121,11 +142,11 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             state = param_to_state[id(p)]
             dst = state.worker_rank
             assert dst < num_ranks
-            shard_elems = split_elems_for_src(p, rank, num_ranks)
+            shard_elems = numel_for_rank(p, rank, state)
             g = p.grad
-            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            g = g.to_local().to(COMM_DTYPE).contiguous()
             assert g.numel() == shard_elems
-            per_dst[dst].append(g)
+            per_dst[dst].append(g.view(-1))
             send_counts[dst] += shard_elems
 
         assert any(
@@ -148,13 +169,18 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                total += split_elems_for_src(p, src, num_ranks)
+                total += numel_for_rank(p, src, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
         #All2All
+        logger.debug(f"send_buf size: {send_buf.numel()}, "
+                     f"recv_buf size: {recv_buf.numel()}, "
+                     f"recv_counts: {recv_counts}, "
+                     f"send_counts: {send_counts}, "
+                     f"process_group: {str(process_group)}")
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -179,7 +205,6 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owned_params}
         for src in range(num_ranks):
             if recv_counts[src] == 0:
                 continue
@@ -189,22 +214,28 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
-                n = split_elems_for_src(p, src, num_ranks)
+
+                # get the slice of the full dtensor corresponding to rank src.
+                slices = get_slices_of_dtensor(state.gathered_grad, src,
+                                               state.shard_mesh,
+                                               state.shard_placements)
+
+                dst = state.gathered_grad[slices]
+                assert dst._base is state.gathered_grad
+
+                n = dst.numel()
                 assert n > 0
 
                 sg = recv_buf.narrow(0, off + inner_off, n)
-                woff = write_offsets[id(p)]
-                dst = state.gathered_grad.narrow(0, woff, n)
+                sg = sg.reshape_as(dst)
                 dst.copy_(sg)
 
-                write_offsets[id(p)] += n
                 inner_off += n
             off += block
 
         for p in params:
             state = param_to_state[id(p)]
             if state.worker_rank == rank:
-                state.gathered_grad = state.gathered_grad.view_as(p)
                 state.gather_event = torch.cuda.Event()
                 state.gather_event.record(comm_stream)
             else:
@@ -277,14 +308,19 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert state.computed_u is not None
 
-                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous()
 
                 offset = 0
                 for dst in range(num_ranks):
-                    n = split_elems_for_src(p, dst, num_ranks)
+                    # get the slice of the full tensor corresponding to rank dst.
+                    slices = get_slices_of_dtensor(u_full, dst,
+                                                   state.shard_mesh,
+                                                   state.shard_placements)
+                    su = u_full[slices].flatten()
+
+                    n = su.numel()
                     assert n > 0
 
-                    su = u_full.narrow(0, offset, n)
                     per_dst[dst].append(su)
                     send_counts[dst] += n
                     offset += n
@@ -313,7 +349,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                total += split_elems_for_src(p, rank, num_ranks)
+                total += numel_for_rank(p, rank, state)
             recv_counts[src] = total
 
         recv_total = sum(recv_counts)
@@ -357,7 +393,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
                 state = param_to_state[id(p)]
                 if state.worker_rank != src:
                     continue
-                n = split_elems_for_src(p, rank, num_ranks)
+                n = numel_for_rank(p, rank, state)
                 assert n > 0
 
                 flat_local = recv_buf.narrow(0, off + inner_off,
@@ -398,11 +434,23 @@ def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
         state.scattered_u = None
         u_dtensor = None
 
-        scales_full = Muon._compute_scales(p, state.qk_clip_state)
+        scales_full = Muon._compute_scales(
+            p,
+            state.qk_clip_state) if state.qk_clip_state is not None else None
         if scales_full is not None:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            local_rank = dist.get_rank(group=state.process_group)
-            scales_local = scales_full.chunk(num_ranks, dim=0)[local_rank]
+            # Have to slice scales_full among dim 0
+            weight_slices = get_slices_of_dtensor(p, rank, state.shard_mesh,
+                                                  state.shard_placements)
+            ratio = p.shape[0] // scales_full.shape[0]
+            scales_slice = slice(
+                None if weight_slices[0].start is None else
+                weight_slices[0].start // ratio,
+                None if weight_slices[0].stop is None else
+                weight_slices[0].stop // ratio,
+                None,
+            )
+
+            scales_local = scales_full[scales_slice]
             scales_local = DTensor.from_local(
                 scales_local,
                 placements=p.placements,
@@ -478,11 +526,11 @@ def parse_qk_layer(name: str) -> tuple[str | None, int]:
 @dataclass
 class QKClipInfo:
     """Per-parameter dynamic info computed from config + runtime logits."""
-    kind: Optional[str]  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
-    indices: List[int]  # which heads to consider for clipping
+    kind: str | None  # 'wq'/'q_proj' or 'wk'/'k_proj' or None
+    indices: list[int]  # which heads to consider for clipping
     head_dim: int  # from config
     threshold: float  # from config
-    logit: Optional[torch.Tensor]
+    logit: torch.Tensor | None
 
 
 class Muon(torch.optim.Optimizer):
@@ -525,11 +573,16 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
-        overlap_step : How many all2all gather, compute operations are launched in advance
-                       before the corresponding all2all scatter steps begin.
-                       A higher overlap_step increases memory usage but can improve
-                       performance by overlapping communication.
-                       Parallel muon only. 
+        warmup_step : How many all2all gather, compute operations are launched in advance
+                      before the corresponding all2all scatter steps begin.
+                      A higher warmup_step increases memory usage but can improve
+                      performance by overlapping communication.
+                      Parallel muon only.
+        chunk_size : Batch size of parameters to process in each
+                     all2all gather/compute/scatter step.
+                     Use shard ranks * DEFAULT_CHUNK_SIZE_RATIO when -1 is specified.
+        use_distributed_muon: Use distributed muon by Liu et al. (2024).
+                              For testing purpose only.
     """
 
     def __init__(self,
@@ -549,7 +602,9 @@ class Muon(torch.optim.Optimizer):
                      "head_dim": 128,
                      "threshold": 100
                  },
-                 overlap_step=5):
+                 warmup_step=5,
+                 chunk_size=-1,
+                 use_distributed_muon=False):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -579,7 +634,9 @@ class Muon(torch.optim.Optimizer):
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
-        self.overlap_step = overlap_step
+        self.warmup_step = warmup_step
+        self.chunk_size = chunk_size
+        self.use_distributed_muon = use_distributed_muon
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -597,6 +654,12 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
+    def set_rank_once(self, rank):
+        if self.rank is None:
+            self.rank = rank
+        else:
+            assert self.rank == rank
+
     def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
@@ -604,26 +667,13 @@ class Muon(torch.optim.Optimizer):
         assert isinstance(
             p, DTensor), "Parallel Muon only supports DTensor parameters."
 
-        if p.placements == (Shard(dim=0), ):
-            # Case for FSDP
-            process_group = p.device_mesh.get_group(mesh_dim=0)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
-        elif p.placements == (Replicate(), Shard(dim=0)):
-            # Case for HSDP
-            process_group = p.device_mesh.get_group(mesh_dim=1)
-            if self.rank is None:
-                self.rank = dist.get_rank(group=process_group)
-            else:
-                assert self.rank == dist.get_rank(group=process_group)
-            for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if self.rank in shard_mesh:
-                    return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
-        else:
-            raise ValueError(f"Unsupported placements ({p.placements}).")
+        shard_mesh, shard_pg, shard_placements = construct_shard_mesh(
+            p.placements, p.device_mesh)
+
+        # set rank with the local rank in the shard process group
+        self.set_rank_once(dist.get_rank(group=shard_pg))
+
+        return shard_mesh, shard_pg, shard_placements
 
     def init_state_and_assign_params(self, names, params, group, qk_logits):
         param_to_state = {}
@@ -655,23 +705,32 @@ class Muon(torch.optim.Optimizer):
         ordered_params = list(params_sorted)
 
         round_robin = 0
-        mesh = None
-        shard_mesh = None
-        process_group = None
+        mesh = ordered_params[0].device_mesh
+        placements = ordered_params[0].placements
+
+        shard_mesh, shard_pg, shard_placements = self.get_shard_mesh(
+            ordered_params[0])
+        shard_mesh_flattened = shard_mesh.mesh.flatten()
+        num_ranks = dist.get_world_size(group=shard_pg)
+
         for n, p in zip(ordered_names, ordered_params):
-            if mesh is None:
-                mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p)
-            elif mesh != p.device_mesh:
+            if mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
-            num_ranks = dist.get_world_size(group=process_group)
-            param_to_state[id(p)] = _muon_state()
-            param_to_state[id(
-                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
-            param_to_state[id(p)].process_group = process_group
+            if placements != p.placements:
+                raise ValueError("All parameters must have same placements.")
+
+            worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
+            round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
-            param_to_state[id(p)].qk_clip_state = qk_clip_state
-            round_robin = (round_robin + 1) % len(shard_mesh)
+
+            param_to_state[id(p)] = _muon_state(
+                worker_rank=worker_rank,
+                process_group=shard_pg,
+                shard_mesh=shard_mesh,
+                shard_placements=shard_placements,
+                name=n,
+                qk_clip_state=qk_clip_state,
+            )
 
         return param_to_state, ordered_params
 
@@ -705,10 +764,73 @@ class Muon(torch.optim.Optimizer):
 
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
 
-            scales_full = self._compute_scales(p, qk_clip_state)
+            scales_full = self._compute_scales(
+                p, qk_clip_state) if qk_clip_state is not None else None
             if scales_full is not None:
                 Muon._qk_clip(p, scales_full, qk_clip_state.head_dim)
 
+    def distributed_muon(
+        self,
+        names: list[str],
+        params: list[torch.nn.Parameter],
+        group: dict[str, Any],
+        lr: float,
+        weight_decay: float,
+        momentum: float,
+        qk_logits: list[torch.Tensor | DTensor] | None,
+    ):
+        """ Implementation of Distributed Muon by Liu et al. """
+        if qk_logits is not None:
+            raise NotImplementedError("QK clipping is not supported yet")
+
+        if isinstance(params[0], DTensor):
+            shard_mesh, _, shard_placements = construct_shard_mesh(
+                placements=params[0].placements,
+                mesh=params[0].device_mesh,
+            )
+
+        for n, p in zip(names, params):
+            g = p.grad
+            if g is None:
+                continue
+            if g.ndim > 2:
+                g = g.view(g.size(0), -1)
+            assert g is not None
+
+            # calc update
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(momentum).add_(g)
+            if group["nesterov"]:
+                g = g.add(buf, alpha=momentum)
+            else:
+                g = buf
+
+            # Gather G
+            if isinstance(p.data, DTensor):
+                g = g.full_tensor()
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                             steps=group["ns_steps"])
+
+            if isinstance(p.data, DTensor):
+                slices = get_slices_of_dtensor(
+                    target=p,
+                    local_rank=dist.get_rank(),
+                    shard_mesh=shard_mesh,
+                    shard_placements=shard_placements,
+                )
+                u_shard = u[slices]
+                u = DTensor.from_local(
+                    u_shard,
+                    device_mesh=p.device_mesh,
+                    placements=p.placements,
+                )
+
+            adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+            Muon._update_p(p, u, lr, adjusted_lr, weight_decay)
+
     def _update_g(self, p, g, group, momentum):
         # calc update
         state = self.state[p]
@@ -727,6 +849,9 @@ class Muon(torch.optim.Optimizer):
         p.data.add_(u, alpha=-adjusted_lr)
 
     def get_qk_clip_info(self, n, qk_logits):
+        if self.clip_config is None:
+            return None
+
         head_dim = self.clip_config.get('head_dim')
         threshold = self.clip_config.get('threshold')
         kind, layer_idx = parse_qk_layer(n)
@@ -737,6 +862,11 @@ class Muon(torch.optim.Optimizer):
             indices_key = 'q_indices' if 'q' in kind else 'k_indices'
             indices = self.clip_config.get(indices_key, []) or []
 
+            if isinstance(logit, DTensor):
+                # In TP settings, qk_logits may be DTensor
+                # We convert it to full tensor here for simplicity
+                logit = logit.full_tensor()
+
         return QKClipInfo(
             kind=kind,
             indices=indices,
@@ -835,22 +965,28 @@ class Muon(torch.optim.Optimizer):
                 _update_param(p, state, lr, adjusted_lr, weight_decay,
                               self.rank, self.compute_stream)
 
-        chunk_size = dist.get_world_size(param_to_state[id(
-            params[0])].process_group)
+        if self.chunk_size == -1:
+            shard_ranks = dist.get_world_size(param_to_state[id(
+                params[0])].process_group)
+            chunk_size = shard_ranks * DEFAULT_CHUNK_SIZE_RATIO
+        elif self.chunk_size > 0:
+            chunk_size = self.chunk_size
+        else:
+            raise ValueError("chunk_size must be -1 or a positive integer.")
 
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        overlap_step = self.overlap_step
-        for i in range(0, overlap_step):
+        warmup_step = self.warmup_step
+        for i in range(0, warmup_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + warmup_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
+            enqueue_computes(i + warmup_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)
@@ -866,7 +1002,7 @@ class Muon(torch.optim.Optimizer):
         amsgrad: bool,
         beta1: float,
         beta2: float,
-        lr: Union[float, torch.Tensor],
+        lr: float | torch.Tensor,
         weight_decay: float,
         eps: float,
         maximize: bool,
@@ -876,10 +1012,10 @@ class Muon(torch.optim.Optimizer):
 
         # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer
         # treating it as a scalar.
-        lr_dict: Optional[DeviceDict] = ({
+        lr_dict: DeviceDict | None = ({
             lr.device: lr
         } if isinstance(lr, torch.Tensor) and str(lr.device) != "cpu" else
-                                         None)
+                                      None)
         grouped_tensors = torch.optim.Optimizer._group_tensors_by_device_and_dtype(
             [
                 params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs,
@@ -926,6 +1062,159 @@ class Muon(torch.optim.Optimizer):
                 maximize=maximize,
             )
 
+    def _step_muon(self, group, qk_logits=None):
+        params = group["params"]
+        lr = group["lr"]
+        weight_decay = group["weight_decay"]
+        momentum = group["momentum"]
+        names = group["names"]
+
+        param_dtensors = []
+        param_tensors = []
+        name_dtensors = []
+        name_tensors = []
+
+        if self.use_distributed_muon:
+            self.distributed_muon(names=names,
+                                  params=params,
+                                  group=group,
+                                  lr=lr,
+                                  weight_decay=weight_decay,
+                                  momentum=momentum,
+                                  qk_logits=qk_logits)
+            return
+
+        for n, p in zip(names, params):
+            if p is None or p.grad is None:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(
+                        isinstance(placement, Replicate)
+                        for placement in p.placements):
+                    param_tensors.append(p)
+                    name_tensors.append(n)
+                else:
+                    param_dtensors.append(p)
+                    name_dtensors.append(n)
+            elif isinstance(p.data, torch.Tensor):
+                param_tensors.append(p)
+                name_tensors.append(n)
+            else:
+                raise TypeError(f"Unsupported parameter type: {type(p.data)}")
+
+        logger.debug(
+            f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors"
+        )
+
+        if len(param_dtensors) > 0:
+            if not dist.is_initialized():
+                raise RuntimeError(
+                    "Parallel Muon requires torch.distributed to be initialized."
+                )
+
+            # To support different placements, we group parameters by placements
+            # and run parallel Muon on each group.
+
+            placement_to_params = defaultdict(lambda: ([], []))
+            # type: dict[tuple[Placement, DeviceMesh], tuple[list[str], list[DTensor]]]
+
+            assert len(name_dtensors) == len(param_dtensors)
+            for n, p in zip(name_dtensors, param_dtensors):
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][0].append(n)
+                placement_to_params[tuple([p.placements,
+                                           p.device_mesh])][1].append(p)
+
+            for _, (names, params) in placement_to_params.items():
+                self.parallel(
+                    names,
+                    params,
+                    group,
+                    lr=lr,
+                    weight_decay=weight_decay,
+                    momentum=momentum,
+                    qk_logits=qk_logits,
+                )
+
+        if len(param_tensors) > 0:
+            self.base(
+                name_tensors,
+                param_tensors,
+                group,
+                lr=lr,
+                weight_decay=weight_decay,
+                momentum=momentum,
+                qk_logits=qk_logits,
+            )
+
+    def _step_adamw_params(self, params, group):
+        params_with_grads = []
+        grads = []
+        moment1 = []
+        moment2 = []
+        max_exp_avg_sqs = []
+        state_steps = []
+        lr = group["lr"]
+        beta1, beta2 = group["adamw_betas"]
+        eps = group["adamw_eps"]
+        weight_decay = group["weight_decay"]
+
+        for p in params:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            params_with_grads.append(p)
+            grads.append(g)
+            if "step" not in state:
+                state["step"] = (torch.zeros((),
+                                             dtype=torch.float32,
+                                             device=p.device))
+                state["moment1"] = torch.zeros_like(g)
+                state["moment2"] = torch.zeros_like(g)
+            moment1.append(state["moment1"])
+            moment2.append(state["moment2"])
+            if not isinstance(state["step"], torch.Tensor):
+                step_tensor = torch.tensor(state["step"],
+                                           dtype=torch.float32,
+                                           device=p.device)
+            else:
+                step_tensor = state["step"]
+            state_steps.append(step_tensor)
+
+        self._fused_adamw(
+            params_with_grads,
+            grads,
+            moment1,
+            moment2,
+            max_exp_avg_sqs,
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+    def _step_adamw(self, group):
+        params = group["params"]
+
+        # group params with it's type and placement
+        placement_to_params: dict[tuple[Placement | type,
+                                        DeviceMesh | None]] = defaultdict(list)
+        for p in params:
+            match p:
+                case DTensor():
+                    placement_to_params[tuple([p.placements,
+                                               p.device_mesh])].append(p)
+                case torch.Tensor():
+                    placement_to_params[tuple([torch.Tensor, None])].append(p)
+
+        for params in placement_to_params.values():
+            self._step_adamw_params(params, group)
+
     def step(self, closure=None, qk_logits=None):
         """Perform a single optimization step.
 
@@ -943,127 +1232,9 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         for group in self.param_groups:
-            params = group["params"]
-
             if group["use_muon"]:
-                ############################
-                #           Muon           #
-                ############################
-                lr = group["lr"]
-                weight_decay = group["weight_decay"]
-                momentum = group["momentum"]
-                names = group["names"]
-
-                param_dtensors = []
-                param_tensors = []
-                name_dtensors = []
-                name_tensors = []
-
-                for n, p in zip(names, params):
-                    if p is None or p.grad is None:
-                        continue
-                    if isinstance(p.data, DTensor):
-                        if all(
-                                isinstance(placement, Replicate)
-                                for placement in p.placements):
-                            param_tensors.append(p)
-                            name_tensors.append(n)
-                        else:
-                            param_dtensors.append(p)
-                            name_dtensors.append(n)
-                    elif isinstance(p.data, torch.Tensor):
-                        param_tensors.append(p)
-                        name_tensors.append(n)
-                    else:
-                        raise TypeError(
-                            f"Unsupported parameter type: {type(p.data)}")
-
-                if self.debug:
-                    print(
-                        f"[Muon] {len(param_dtensors)} DTensors, {len(param_tensors)} Tensors",
-                        flush=True,
-                    )
-
-                if len(param_dtensors) > 0:
-                    if not dist.is_initialized():
-                        raise RuntimeError(
-                            "Parallel Muon requires torch.distributed to be initialized."
-                        )
-
-                    self.parallel(
-                        name_dtensors,
-                        param_dtensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
-                if len(param_tensors) > 0:
-                    self.base(
-                        name_tensors,
-                        param_tensors,
-                        group,
-                        lr=lr,
-                        weight_decay=weight_decay,
-                        momentum=momentum,
-                        qk_logits=qk_logits,
-                    )
-
+                self._step_muon(group, qk_logits=qk_logits)
             else:
-                ############################
-                #       AdamW backup       #
-                ############################
-
-                params_with_grads = []
-                grads = []
-                moment1 = []
-                moment2 = []
-                max_exp_avg_sqs = []
-                state_steps = []
-                lr = group["lr"]
-                beta1, beta2 = group["adamw_betas"]
-                eps = group["adamw_eps"]
-                weight_decay = group["weight_decay"]
-
-                for p in params:
-                    g = p.grad
-                    if g is None:
-                        continue
-                    state = self.state[p]
-                    params_with_grads.append(p)
-                    grads.append(g)
-                    if "step" not in state:
-                        state["step"] = (torch.zeros((),
-                                                     dtype=torch.float32,
-                                                     device=p.device))
-                        state["moment1"] = torch.zeros_like(g)
-                        state["moment2"] = torch.zeros_like(g)
-                    moment1.append(state["moment1"])
-                    moment2.append(state["moment2"])
-                    if not isinstance(state["step"], torch.Tensor):
-                        step_tensor = torch.tensor(state["step"],
-                                                   dtype=torch.float32,
-                                                   device=p.device)
-                    else:
-                        step_tensor = state["step"]
-                    state_steps.append(step_tensor)
-
-                self._fused_adamw(
-                    params_with_grads,
-                    grads,
-                    moment1,
-                    moment2,
-                    max_exp_avg_sqs,
-                    state_steps,
-                    amsgrad=False,
-                    beta1=beta1,
-                    beta2=beta2,
-                    lr=lr,
-                    weight_decay=weight_decay,
-                    eps=eps,
-                    maximize=False,
-                )
+                self._step_adamw(group)
 
         return loss