Add built binary [ci skip]

build/torch27-cxx11-cu118-x86_64-linux/optimizer/_ops.py

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

build/torch27-cxx11-cu118-x86_64-linux/optimizer/{_optimizer_9c21645_dirty.abi3.so → _optimizer_fceb334_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bf8b97161714dff91953d26ae0bf59ebc9f3653ce57a3998723cc08aa97b71e6
+oid sha256:be3c6bfa9c13eae63d8c208954dfb399939bfe64065d0187464db5ff2dce5c30
 size 1787368

build/torch27-cxx11-cu118-x86_64-linux/optimizer/matmul_transpose_triton.py

@@ -0,0 +1,106 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def get_autotune_config():
+    return [
+        triton.Config(
+            {
+                'BLOCK_SIZE_M': blk_m,
+                'BLOCK_SIZE_K': blk_k,
+                'GROUP_SIZE_M': grp_sz
+            },
+            num_stages=n_stages,
+            num_warps=n_warps) for blk_m in [32, 64, 128]
+        for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5]
+        for n_warps in [4, 8]
+    ]
+
+
+@triton.autotune(
+    configs=get_autotune_config(),
+    key=['M', 'K'],
+)
+@triton.jit
+def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn,
+               BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+               GROUP_SIZE_M: tl.constexpr):
+    """
+    Core kernel jit function of matmul_transpose that computes y = x @ x.T
+    The code is a simple adaptation from the triton `matmul` tutorial:
+    https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+    """
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    if pid_m > pid_n:
+        return
+
+    offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # we use a & b ptrs to denote different rows of x.
+    a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+    b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        b = tl.load(b_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator)
+        a_ptrs += BLOCK_SIZE_K * stride_xk
+        b_ptrs += BLOCK_SIZE_K * stride_xk
+    # use dtype.element_ty to accommodate different input datatypes as in cpp templates
+    # https://github.com/triton-lang/triton/issues/2252
+    c = accumulator.to(x.dtype.element_ty)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+    # transpose and copy
+    if pid_m < pid_n:
+        ct_ptrs = y + stride_ym * offs_cn[:,
+                                          None] + stride_yn * offs_cm[None, :]
+        ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+        tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask)
+
+
+def matmul_transpose_assign(d_in, d_out):
+    assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor"
+    assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor"
+    assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device"
+    assert d_in.dtype == d_out.dtype, "Inputs must have the same data type"
+    assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor"
+    assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor"
+    assert d_in.size(0) == d_out.size(0) == d_out.size(0), \
+            "First dimension of `d_in` must match first and second dimension of `d_out`"
+
+    d_in = d_in.contiguous()
+    M, K = d_in.shape
+    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(
+        M, META['BLOCK_SIZE_M']), )
+    with torch.cuda.device(d_in.device.index):
+        mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1),
+                         d_out.stride(0), d_out.stride(1))
+
+
+def matmul_transpose(d_in):
+    M, _ = d_in.shape
+    d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype)
+    matmul_transpose_assign(d_in, d_out)
+    return d_out

build/torch27-cxx11-cu118-x86_64-linux/optimizer/muon.py

@@ -8,14 +8,19 @@ import torch
 import torch.distributed as dist
 from torch.distributed._tensor import DTensor, Replicate, Shard
 
+from .matmul_transpose_triton import matmul_transpose_assign
+
 logger = logging.getLogger(__name__)
 
+COMM_DTYPE = torch.bfloat16
+
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
 # Muon's Newton–Schulz iteration causes high variance in singular values
 # Idea: give each iteration its own 3 coefficients and optimize them via gradient descent.
 @torch.no_grad()
+# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon
 def _zeropower_via_newtonschulz5(G, steps):
     """
     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
@@ -27,13 +32,15 @@ def _zeropower_via_newtonschulz5(G, steps):
     performance at all relative to UV^T, where USV^T = G is the SVD.
     """
     assert len(G.shape) == 2
-    assert G.dtype == torch.bfloat16
+    assert G.dtype == COMM_DTYPE
     X = G  # no manual typecast
 
     if G.size(0) > G.size(1):
         X = X.T
     # Ensure spectral norm is at most 1
     X = X / (X.norm() + 1e-7)
+    buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
+    buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
     for a, b, c in [
         (4.0848, -6.8946, 2.9270),
@@ -42,13 +49,10 @@ def _zeropower_via_newtonschulz5(G, steps):
         (2.8769, -3.1427, 1.2046),
         (2.8366, -3.0525, 1.2012),
     ]:
-        A = X @ X.T
-        # B = (
-        #    b * A + c * A @ A
-        # )
-        B = torch.addmm(A, A, A, alpha=c, beta=b)
-        # X = a * X + B @ X
-        X = torch.addmm(X, B, X, alpha=1.0, beta=a)
+        matmul_transpose_assign(X, buf1)
+        matmul_transpose_assign(buf1, buf2)
+        buf1.mul_(b).add_(buf2, alpha=c)
+        X = torch.addmm(X, buf1, X, alpha=1.0, beta=a)
 
     if G.size(0) > G.size(1):
         X = X.T
@@ -69,51 +73,131 @@ class _muon_state:
     qk_clip_state = None
 
 
+def split_elems_for_src(param, state, 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
+
+
 @torch.no_grad()
-def _gather(p, state, rank, comm_stream, none_grad):
-    """
-    Gather the gradients to worker_rank.
-    If none_grad is True, free p.grad after the gather.
-    """
+def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(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(),
+                                                  dtype=COMM_DTYPE,
+                                                  device="cuda")
+            else:
+                state.gathered_grad = None
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
+
+
+@torch.no_grad()
+def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
+                    alloc_event):
     with torch.cuda.stream(comm_stream):
-        g = p.grad
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            gather_list = [
-                torch.empty_like(g.to_local(), dtype=torch.bfloat16)
-                for _ in range(num_ranks)
-            ]
-        else:
-            gather_list = None
-
-        g = g.to(torch.bfloat16)
-        torch.distributed.gather(
-            g.to_local(),
-            dst=state.worker_rank,
-            gather_list=gather_list,
-            group=state.process_group,
+        # Calculate sending tensors
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
+        for p in params:
+            state = param_to_state[id(p)]
+            dst = state.worker_rank
+            assert dst < num_ranks
+            shard_elems = split_elems_for_src(p, state, rank, num_ranks)
+            g = p.grad
+            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            assert g.numel() == shard_elems
+            per_dst[dst].append(g)
+            send_counts[dst] += shard_elems
+
+        assert all(
+            len(v) > 0
+            for v in per_dst), "all params should be sharded to all devices"
+
+        send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
+
+        # Calculate receiving tensors
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                total += split_elems_for_src(p, state, src, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        if rank == state.worker_rank:
-            if state.gathered_grad is not None:
-                raise RuntimeError(
-                    "Gather event already exists, which should not happen.")
-            state.gathered_grad = torch.cat(gather_list, dim=0)
-            state.gather_event = torch.cuda.Event()
-            state.gather_event.record()
-        else:
-            state.gathered_grad = None
-            state.gather_event = None
-        gather_list = None
-        if none_grad:
-            # We can safely free p.grad without calling record_stream:
-            #   p.grad.to_local().record_stream(comm_stream)
-            # Explanation:
-            # 1. p.grad is created on the default stream, but the default stream
-            #    is synchronized with the comm stream later.
-            # 2. There is no further activity on the default stream before the optimizer finishes.
-            # Therefore, it is safe to free p.grad directly on the comm stream.
-            p.grad = None
+
+        # Reconstructs gathered grad from the received buffer
+        #
+        #                  recv_buf (num ranks = 3)
+        #
+        #      From rank 0        From rank 1        From rank 2
+        # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # p1_n -> p2_n -> p3_n
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        write_offsets = {id(p): 0 for p in owner_params}
+        for src in range(num_ranks):
+            if recv_counts[src] == 0:
+                continue
+
+            block = recv_counts[src]
+            inner_off = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                n = split_elems_for_src(p, state, src, num_ranks)
+                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)
+                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:
+                state.gathered_grad = None
+                state.gather_event = None
+            if none_grad:
+                p.grad = None
 
 
 @torch.no_grad()
@@ -127,45 +211,119 @@ def _compute_u(p, state, steps, rank, compute_stream):
                 raise RuntimeError("Gather event must be set before compute.")
             compute_stream.wait_event(state.gather_event)
             u = _zeropower_via_newtonschulz5(state.gathered_grad, steps)
+            state.gathered_grad = None
             state.computed_u = u
-        state.scattered_u = torch.empty_like(p.to_local(),
-                                             dtype=torch.bfloat16)
-        state.compute_event = torch.cuda.Event()
-        state.compute_event.record()
-        u = None
+            state.compute_event = torch.cuda.Event()
+            state.compute_event.record()
+        else:
+            state.computed_u = None
+            state.compute_event = None
 
 
 @torch.no_grad()
-def _scatter(p, state, rank, comm_stream):
-    """
-    Scatter the computed_u from worker_rank to all ranks.
-    """
+def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(compute_stream):
+        for p in params:
+            state = param_to_state[id(p)]
+            state.scattered_u = torch.empty_like(p.to_local(),
+                                                 dtype=COMM_DTYPE)
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
 
+
+def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
     with torch.cuda.stream(comm_stream):
-        if state.compute_event is None:
-            raise RuntimeError("Compute event must be set before scatter.")
-        comm_stream.wait_event(state.compute_event)
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            # Clear the gathered gradient to free memory
-            state.gathered_grad = None
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
 
-            u = state.computed_u
-            scatter_list = list(torch.split(u, p.size(0) // num_ranks, dim=0))
-            scatter_list = [s.contiguous() for s in scatter_list]
-        else:
-            scatter_list = None
+        if owner_params:
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                if state.compute_event is None:
+                    raise RuntimeError(
+                        "Compute event must be set before scatter.")
+                comm_stream.wait_event(state.compute_event)
+                state.gathered_grad = None
 
-        torch.distributed.scatter(
-            state.scattered_u,
-            scatter_list=scatter_list,
-            src=state.worker_rank,
-            group=state.process_group,
+                assert state.computed_u is not None
+
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+
+                offset = 0
+                for dst in range(num_ranks):
+                    n = split_elems_for_src(p, state, dst, num_ranks)
+                    assert n > 0
+
+                    su = u_full.narrow(0, offset, n)
+                    per_dst[dst].append(su)
+                    send_counts[dst] += n
+                    offset += n
+
+                assert offset == u_full.numel()
+
+        if any(len(v) > 0 for v in per_dst):
+            send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        else:
+            # all_to_all requires participation from all ranks
+            # Even non-owner ranks must join the collective call
+            send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
+
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                total += split_elems_for_src(p, state, rank, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        assert recv_total > 0
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        state.scatter_event = torch.cuda.Event()
-        state.scatter_event.record()
-        scatter_list = None
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        for src in range(num_ranks):
+            block = recv_counts[src]
+            if block == 0:
+                continue
+
+            inner_off = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                n = split_elems_for_src(p, state, rank, num_ranks)
+                assert n > 0
+
+                flat_local = recv_buf.narrow(0, off + inner_off,
+                                             n).view_as(p.to_local())
+                state.scattered_u.copy_(flat_local)
+
+                state.scatter_event = torch.cuda.Event()
+                state.scatter_event.record(comm_stream)
+                inner_off += n
+
+            assert inner_off == block
+            off += block
 
 
 def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
@@ -444,11 +602,14 @@ class Muon(torch.optim.Optimizer):
             if mesh is None:
                 mesh = p.device_mesh
                 shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
+                self.rank = dist.get_rank(group=process_group)
             elif 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()
+            param_to_state[id(
+                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
             param_to_state[id(p)].process_group = process_group
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
             param_to_state[id(p)].qk_clip_state = qk_clip_state
@@ -478,7 +639,7 @@ class Muon(torch.optim.Optimizer):
             else:
                 g = buf
 
-            u = _zeropower_via_newtonschulz5(g.bfloat16(),
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
                                              steps=group["ns_steps"])
 
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
@@ -493,15 +654,12 @@ class Muon(torch.optim.Optimizer):
     def _update_g(self, p, g, group, momentum):
         # 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)
+        buf = state.setdefault("momentum_buffer", torch.zeros_like(g))
+        torch.add(g, buf, alpha=momentum, out=buf)
         if group["nesterov"]:
-            g = g.add(buf, alpha=momentum)
-        else:
-            g = buf
-        return g
+            g.add_(buf, alpha=momentum)
+            return g
+        return buf
 
     @staticmethod
     def _update_p(p, u, lr, adjusted_lr, weight_decay):
@@ -585,11 +743,15 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
-        def enqueue_gathers(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _gather(p, state, self.rank, self.comm_stream,
-                        group["none_grad"])
+        def enqueue_all2all_gather(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_gathered_grad(target_params,
+                                                   param_to_state, self.rank,
+                                                   self.compute_stream)
+                _all2all_gather(target_params, param_to_state, self.rank,
+                                self.comm_stream, group["none_grad"],
+                                alloc_event)
 
         def enqueue_computes(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -597,10 +759,14 @@ class Muon(torch.optim.Optimizer):
                 _compute_u(p, state, group["ns_steps"], self.rank,
                            self.compute_stream)
 
-        def enqueue_scatters(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _scatter(p, state, self.rank, self.comm_stream)
+        def enqueue_all2all_scatter(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_scattered_u(target_params, param_to_state,
+                                                 self.rank,
+                                                 self.compute_stream)
+                _all2all_scatter(target_params, param_to_state, self.rank,
+                                 self.comm_stream, alloc_event)
 
         def enqueue_update_param(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -615,14 +781,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        enqueue_gathers(0, chunk_size)
+        PRE_STEP = 5
+        for i in range(0, PRE_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_computes(i, chunk_size)
-            if i > 0:
-                enqueue_update_param(i - chunk_size, chunk_size)
-            enqueue_gathers(i + chunk_size, chunk_size)
-            enqueue_scatters(i, chunk_size)
-        enqueue_update_param(i, chunk_size)
+            enqueue_all2all_scatter(i, chunk_size)
+            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_update_param(i, chunk_size)
+            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)

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

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

build/{torch28-cxx11-cu126-x86_64-linux/optimizer/_optimizer_9c21645_dirty.abi3.so → torch27-cxx11-cu126-x86_64-linux/optimizer/_optimizer_fceb334_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:42ae6ac1cf967d7d23cac7930c8db635105f60631220a60b9cee060d082f40ae
+oid sha256:c31efe0c8bad71c6fbc002f269eff7a30754e1583bb8c892b8ce5cdbdfc5121c
 size 1824256

build/torch27-cxx11-cu126-x86_64-linux/optimizer/matmul_transpose_triton.py

@@ -0,0 +1,106 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def get_autotune_config():
+    return [
+        triton.Config(
+            {
+                'BLOCK_SIZE_M': blk_m,
+                'BLOCK_SIZE_K': blk_k,
+                'GROUP_SIZE_M': grp_sz
+            },
+            num_stages=n_stages,
+            num_warps=n_warps) for blk_m in [32, 64, 128]
+        for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5]
+        for n_warps in [4, 8]
+    ]
+
+
+@triton.autotune(
+    configs=get_autotune_config(),
+    key=['M', 'K'],
+)
+@triton.jit
+def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn,
+               BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+               GROUP_SIZE_M: tl.constexpr):
+    """
+    Core kernel jit function of matmul_transpose that computes y = x @ x.T
+    The code is a simple adaptation from the triton `matmul` tutorial:
+    https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+    """
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    if pid_m > pid_n:
+        return
+
+    offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # we use a & b ptrs to denote different rows of x.
+    a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+    b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        b = tl.load(b_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator)
+        a_ptrs += BLOCK_SIZE_K * stride_xk
+        b_ptrs += BLOCK_SIZE_K * stride_xk
+    # use dtype.element_ty to accommodate different input datatypes as in cpp templates
+    # https://github.com/triton-lang/triton/issues/2252
+    c = accumulator.to(x.dtype.element_ty)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+    # transpose and copy
+    if pid_m < pid_n:
+        ct_ptrs = y + stride_ym * offs_cn[:,
+                                          None] + stride_yn * offs_cm[None, :]
+        ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+        tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask)
+
+
+def matmul_transpose_assign(d_in, d_out):
+    assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor"
+    assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor"
+    assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device"
+    assert d_in.dtype == d_out.dtype, "Inputs must have the same data type"
+    assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor"
+    assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor"
+    assert d_in.size(0) == d_out.size(0) == d_out.size(0), \
+            "First dimension of `d_in` must match first and second dimension of `d_out`"
+
+    d_in = d_in.contiguous()
+    M, K = d_in.shape
+    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(
+        M, META['BLOCK_SIZE_M']), )
+    with torch.cuda.device(d_in.device.index):
+        mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1),
+                         d_out.stride(0), d_out.stride(1))
+
+
+def matmul_transpose(d_in):
+    M, _ = d_in.shape
+    d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype)
+    matmul_transpose_assign(d_in, d_out)
+    return d_out

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

@@ -8,14 +8,19 @@ import torch
 import torch.distributed as dist
 from torch.distributed._tensor import DTensor, Replicate, Shard
 
+from .matmul_transpose_triton import matmul_transpose_assign
+
 logger = logging.getLogger(__name__)
 
+COMM_DTYPE = torch.bfloat16
+
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
 # Muon's Newton–Schulz iteration causes high variance in singular values
 # Idea: give each iteration its own 3 coefficients and optimize them via gradient descent.
 @torch.no_grad()
+# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon
 def _zeropower_via_newtonschulz5(G, steps):
     """
     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
@@ -27,13 +32,15 @@ def _zeropower_via_newtonschulz5(G, steps):
     performance at all relative to UV^T, where USV^T = G is the SVD.
     """
     assert len(G.shape) == 2
-    assert G.dtype == torch.bfloat16
+    assert G.dtype == COMM_DTYPE
     X = G  # no manual typecast
 
     if G.size(0) > G.size(1):
         X = X.T
     # Ensure spectral norm is at most 1
     X = X / (X.norm() + 1e-7)
+    buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
+    buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
     for a, b, c in [
         (4.0848, -6.8946, 2.9270),
@@ -42,13 +49,10 @@ def _zeropower_via_newtonschulz5(G, steps):
         (2.8769, -3.1427, 1.2046),
         (2.8366, -3.0525, 1.2012),
     ]:
-        A = X @ X.T
-        # B = (
-        #    b * A + c * A @ A
-        # )
-        B = torch.addmm(A, A, A, alpha=c, beta=b)
-        # X = a * X + B @ X
-        X = torch.addmm(X, B, X, alpha=1.0, beta=a)
+        matmul_transpose_assign(X, buf1)
+        matmul_transpose_assign(buf1, buf2)
+        buf1.mul_(b).add_(buf2, alpha=c)
+        X = torch.addmm(X, buf1, X, alpha=1.0, beta=a)
 
     if G.size(0) > G.size(1):
         X = X.T
@@ -69,51 +73,131 @@ class _muon_state:
     qk_clip_state = None
 
 
+def split_elems_for_src(param, state, 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
+
+
 @torch.no_grad()
-def _gather(p, state, rank, comm_stream, none_grad):
-    """
-    Gather the gradients to worker_rank.
-    If none_grad is True, free p.grad after the gather.
-    """
+def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(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(),
+                                                  dtype=COMM_DTYPE,
+                                                  device="cuda")
+            else:
+                state.gathered_grad = None
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
+
+
+@torch.no_grad()
+def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
+                    alloc_event):
     with torch.cuda.stream(comm_stream):
-        g = p.grad
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            gather_list = [
-                torch.empty_like(g.to_local(), dtype=torch.bfloat16)
-                for _ in range(num_ranks)
-            ]
-        else:
-            gather_list = None
-
-        g = g.to(torch.bfloat16)
-        torch.distributed.gather(
-            g.to_local(),
-            dst=state.worker_rank,
-            gather_list=gather_list,
-            group=state.process_group,
+        # Calculate sending tensors
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
+        for p in params:
+            state = param_to_state[id(p)]
+            dst = state.worker_rank
+            assert dst < num_ranks
+            shard_elems = split_elems_for_src(p, state, rank, num_ranks)
+            g = p.grad
+            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            assert g.numel() == shard_elems
+            per_dst[dst].append(g)
+            send_counts[dst] += shard_elems
+
+        assert all(
+            len(v) > 0
+            for v in per_dst), "all params should be sharded to all devices"
+
+        send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
+
+        # Calculate receiving tensors
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                total += split_elems_for_src(p, state, src, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        if rank == state.worker_rank:
-            if state.gathered_grad is not None:
-                raise RuntimeError(
-                    "Gather event already exists, which should not happen.")
-            state.gathered_grad = torch.cat(gather_list, dim=0)
-            state.gather_event = torch.cuda.Event()
-            state.gather_event.record()
-        else:
-            state.gathered_grad = None
-            state.gather_event = None
-        gather_list = None
-        if none_grad:
-            # We can safely free p.grad without calling record_stream:
-            #   p.grad.to_local().record_stream(comm_stream)
-            # Explanation:
-            # 1. p.grad is created on the default stream, but the default stream
-            #    is synchronized with the comm stream later.
-            # 2. There is no further activity on the default stream before the optimizer finishes.
-            # Therefore, it is safe to free p.grad directly on the comm stream.
-            p.grad = None
+
+        # Reconstructs gathered grad from the received buffer
+        #
+        #                  recv_buf (num ranks = 3)
+        #
+        #      From rank 0        From rank 1        From rank 2
+        # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # p1_n -> p2_n -> p3_n
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        write_offsets = {id(p): 0 for p in owner_params}
+        for src in range(num_ranks):
+            if recv_counts[src] == 0:
+                continue
+
+            block = recv_counts[src]
+            inner_off = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                n = split_elems_for_src(p, state, src, num_ranks)
+                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)
+                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:
+                state.gathered_grad = None
+                state.gather_event = None
+            if none_grad:
+                p.grad = None
 
 
 @torch.no_grad()
@@ -127,45 +211,119 @@ def _compute_u(p, state, steps, rank, compute_stream):
                 raise RuntimeError("Gather event must be set before compute.")
             compute_stream.wait_event(state.gather_event)
             u = _zeropower_via_newtonschulz5(state.gathered_grad, steps)
+            state.gathered_grad = None
             state.computed_u = u
-        state.scattered_u = torch.empty_like(p.to_local(),
-                                             dtype=torch.bfloat16)
-        state.compute_event = torch.cuda.Event()
-        state.compute_event.record()
-        u = None
+            state.compute_event = torch.cuda.Event()
+            state.compute_event.record()
+        else:
+            state.computed_u = None
+            state.compute_event = None
 
 
 @torch.no_grad()
-def _scatter(p, state, rank, comm_stream):
-    """
-    Scatter the computed_u from worker_rank to all ranks.
-    """
+def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(compute_stream):
+        for p in params:
+            state = param_to_state[id(p)]
+            state.scattered_u = torch.empty_like(p.to_local(),
+                                                 dtype=COMM_DTYPE)
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
 
+
+def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
     with torch.cuda.stream(comm_stream):
-        if state.compute_event is None:
-            raise RuntimeError("Compute event must be set before scatter.")
-        comm_stream.wait_event(state.compute_event)
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            # Clear the gathered gradient to free memory
-            state.gathered_grad = None
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
 
-            u = state.computed_u
-            scatter_list = list(torch.split(u, p.size(0) // num_ranks, dim=0))
-            scatter_list = [s.contiguous() for s in scatter_list]
-        else:
-            scatter_list = None
+        if owner_params:
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                if state.compute_event is None:
+                    raise RuntimeError(
+                        "Compute event must be set before scatter.")
+                comm_stream.wait_event(state.compute_event)
+                state.gathered_grad = None
 
-        torch.distributed.scatter(
-            state.scattered_u,
-            scatter_list=scatter_list,
-            src=state.worker_rank,
-            group=state.process_group,
+                assert state.computed_u is not None
+
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+
+                offset = 0
+                for dst in range(num_ranks):
+                    n = split_elems_for_src(p, state, dst, num_ranks)
+                    assert n > 0
+
+                    su = u_full.narrow(0, offset, n)
+                    per_dst[dst].append(su)
+                    send_counts[dst] += n
+                    offset += n
+
+                assert offset == u_full.numel()
+
+        if any(len(v) > 0 for v in per_dst):
+            send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        else:
+            # all_to_all requires participation from all ranks
+            # Even non-owner ranks must join the collective call
+            send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
+
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                total += split_elems_for_src(p, state, rank, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        assert recv_total > 0
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        state.scatter_event = torch.cuda.Event()
-        state.scatter_event.record()
-        scatter_list = None
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        for src in range(num_ranks):
+            block = recv_counts[src]
+            if block == 0:
+                continue
+
+            inner_off = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                n = split_elems_for_src(p, state, rank, num_ranks)
+                assert n > 0
+
+                flat_local = recv_buf.narrow(0, off + inner_off,
+                                             n).view_as(p.to_local())
+                state.scattered_u.copy_(flat_local)
+
+                state.scatter_event = torch.cuda.Event()
+                state.scatter_event.record(comm_stream)
+                inner_off += n
+
+            assert inner_off == block
+            off += block
 
 
 def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
@@ -444,11 +602,14 @@ class Muon(torch.optim.Optimizer):
             if mesh is None:
                 mesh = p.device_mesh
                 shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
+                self.rank = dist.get_rank(group=process_group)
             elif 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()
+            param_to_state[id(
+                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
             param_to_state[id(p)].process_group = process_group
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
             param_to_state[id(p)].qk_clip_state = qk_clip_state
@@ -478,7 +639,7 @@ class Muon(torch.optim.Optimizer):
             else:
                 g = buf
 
-            u = _zeropower_via_newtonschulz5(g.bfloat16(),
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
                                              steps=group["ns_steps"])
 
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
@@ -493,15 +654,12 @@ class Muon(torch.optim.Optimizer):
     def _update_g(self, p, g, group, momentum):
         # 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)
+        buf = state.setdefault("momentum_buffer", torch.zeros_like(g))
+        torch.add(g, buf, alpha=momentum, out=buf)
         if group["nesterov"]:
-            g = g.add(buf, alpha=momentum)
-        else:
-            g = buf
-        return g
+            g.add_(buf, alpha=momentum)
+            return g
+        return buf
 
     @staticmethod
     def _update_p(p, u, lr, adjusted_lr, weight_decay):
@@ -585,11 +743,15 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
-        def enqueue_gathers(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _gather(p, state, self.rank, self.comm_stream,
-                        group["none_grad"])
+        def enqueue_all2all_gather(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_gathered_grad(target_params,
+                                                   param_to_state, self.rank,
+                                                   self.compute_stream)
+                _all2all_gather(target_params, param_to_state, self.rank,
+                                self.comm_stream, group["none_grad"],
+                                alloc_event)
 
         def enqueue_computes(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -597,10 +759,14 @@ class Muon(torch.optim.Optimizer):
                 _compute_u(p, state, group["ns_steps"], self.rank,
                            self.compute_stream)
 
-        def enqueue_scatters(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _scatter(p, state, self.rank, self.comm_stream)
+        def enqueue_all2all_scatter(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_scattered_u(target_params, param_to_state,
+                                                 self.rank,
+                                                 self.compute_stream)
+                _all2all_scatter(target_params, param_to_state, self.rank,
+                                 self.comm_stream, alloc_event)
 
         def enqueue_update_param(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -615,14 +781,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        enqueue_gathers(0, chunk_size)
+        PRE_STEP = 5
+        for i in range(0, PRE_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_computes(i, chunk_size)
-            if i > 0:
-                enqueue_update_param(i - chunk_size, chunk_size)
-            enqueue_gathers(i + chunk_size, chunk_size)
-            enqueue_scatters(i, chunk_size)
-        enqueue_update_param(i, chunk_size)
+            enqueue_all2all_scatter(i, chunk_size)
+            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_update_param(i, chunk_size)
+            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)

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

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

build/{torch28-cxx11-cu128-x86_64-linux/optimizer/_optimizer_9c21645_dirty.abi3.so → torch27-cxx11-cu128-x86_64-linux/optimizer/_optimizer_fceb334_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dae71b7e998e72130093a86f8c983c3379510e23525e3cdcd4afe5c21bf4d3db
+oid sha256:5b7313ff1257a692e4e5cfdec474b10a7a795f2e87edb0099c7e5d21f588634f
 size 1883344

build/torch27-cxx11-cu128-x86_64-linux/optimizer/matmul_transpose_triton.py

@@ -0,0 +1,106 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def get_autotune_config():
+    return [
+        triton.Config(
+            {
+                'BLOCK_SIZE_M': blk_m,
+                'BLOCK_SIZE_K': blk_k,
+                'GROUP_SIZE_M': grp_sz
+            },
+            num_stages=n_stages,
+            num_warps=n_warps) for blk_m in [32, 64, 128]
+        for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5]
+        for n_warps in [4, 8]
+    ]
+
+
+@triton.autotune(
+    configs=get_autotune_config(),
+    key=['M', 'K'],
+)
+@triton.jit
+def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn,
+               BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+               GROUP_SIZE_M: tl.constexpr):
+    """
+    Core kernel jit function of matmul_transpose that computes y = x @ x.T
+    The code is a simple adaptation from the triton `matmul` tutorial:
+    https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+    """
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    if pid_m > pid_n:
+        return
+
+    offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # we use a & b ptrs to denote different rows of x.
+    a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+    b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        b = tl.load(b_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator)
+        a_ptrs += BLOCK_SIZE_K * stride_xk
+        b_ptrs += BLOCK_SIZE_K * stride_xk
+    # use dtype.element_ty to accommodate different input datatypes as in cpp templates
+    # https://github.com/triton-lang/triton/issues/2252
+    c = accumulator.to(x.dtype.element_ty)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+    # transpose and copy
+    if pid_m < pid_n:
+        ct_ptrs = y + stride_ym * offs_cn[:,
+                                          None] + stride_yn * offs_cm[None, :]
+        ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+        tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask)
+
+
+def matmul_transpose_assign(d_in, d_out):
+    assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor"
+    assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor"
+    assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device"
+    assert d_in.dtype == d_out.dtype, "Inputs must have the same data type"
+    assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor"
+    assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor"
+    assert d_in.size(0) == d_out.size(0) == d_out.size(0), \
+            "First dimension of `d_in` must match first and second dimension of `d_out`"
+
+    d_in = d_in.contiguous()
+    M, K = d_in.shape
+    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(
+        M, META['BLOCK_SIZE_M']), )
+    with torch.cuda.device(d_in.device.index):
+        mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1),
+                         d_out.stride(0), d_out.stride(1))
+
+
+def matmul_transpose(d_in):
+    M, _ = d_in.shape
+    d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype)
+    matmul_transpose_assign(d_in, d_out)
+    return d_out

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

@@ -8,14 +8,19 @@ import torch
 import torch.distributed as dist
 from torch.distributed._tensor import DTensor, Replicate, Shard
 
+from .matmul_transpose_triton import matmul_transpose_assign
+
 logger = logging.getLogger(__name__)
 
+COMM_DTYPE = torch.bfloat16
+
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
 # Muon's Newton–Schulz iteration causes high variance in singular values
 # Idea: give each iteration its own 3 coefficients and optimize them via gradient descent.
 @torch.no_grad()
+# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon
 def _zeropower_via_newtonschulz5(G, steps):
     """
     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
@@ -27,13 +32,15 @@ def _zeropower_via_newtonschulz5(G, steps):
     performance at all relative to UV^T, where USV^T = G is the SVD.
     """
     assert len(G.shape) == 2
-    assert G.dtype == torch.bfloat16
+    assert G.dtype == COMM_DTYPE
     X = G  # no manual typecast
 
     if G.size(0) > G.size(1):
         X = X.T
     # Ensure spectral norm is at most 1
     X = X / (X.norm() + 1e-7)
+    buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
+    buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
     for a, b, c in [
         (4.0848, -6.8946, 2.9270),
@@ -42,13 +49,10 @@ def _zeropower_via_newtonschulz5(G, steps):
         (2.8769, -3.1427, 1.2046),
         (2.8366, -3.0525, 1.2012),
     ]:
-        A = X @ X.T
-        # B = (
-        #    b * A + c * A @ A
-        # )
-        B = torch.addmm(A, A, A, alpha=c, beta=b)
-        # X = a * X + B @ X
-        X = torch.addmm(X, B, X, alpha=1.0, beta=a)
+        matmul_transpose_assign(X, buf1)
+        matmul_transpose_assign(buf1, buf2)
+        buf1.mul_(b).add_(buf2, alpha=c)
+        X = torch.addmm(X, buf1, X, alpha=1.0, beta=a)
 
     if G.size(0) > G.size(1):
         X = X.T
@@ -69,51 +73,131 @@ class _muon_state:
     qk_clip_state = None
 
 
+def split_elems_for_src(param, state, 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
+
+
 @torch.no_grad()
-def _gather(p, state, rank, comm_stream, none_grad):
-    """
-    Gather the gradients to worker_rank.
-    If none_grad is True, free p.grad after the gather.
-    """
+def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(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(),
+                                                  dtype=COMM_DTYPE,
+                                                  device="cuda")
+            else:
+                state.gathered_grad = None
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
+
+
+@torch.no_grad()
+def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
+                    alloc_event):
     with torch.cuda.stream(comm_stream):
-        g = p.grad
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            gather_list = [
-                torch.empty_like(g.to_local(), dtype=torch.bfloat16)
-                for _ in range(num_ranks)
-            ]
-        else:
-            gather_list = None
-
-        g = g.to(torch.bfloat16)
-        torch.distributed.gather(
-            g.to_local(),
-            dst=state.worker_rank,
-            gather_list=gather_list,
-            group=state.process_group,
+        # Calculate sending tensors
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
+        for p in params:
+            state = param_to_state[id(p)]
+            dst = state.worker_rank
+            assert dst < num_ranks
+            shard_elems = split_elems_for_src(p, state, rank, num_ranks)
+            g = p.grad
+            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            assert g.numel() == shard_elems
+            per_dst[dst].append(g)
+            send_counts[dst] += shard_elems
+
+        assert all(
+            len(v) > 0
+            for v in per_dst), "all params should be sharded to all devices"
+
+        send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
+
+        # Calculate receiving tensors
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                total += split_elems_for_src(p, state, src, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        if rank == state.worker_rank:
-            if state.gathered_grad is not None:
-                raise RuntimeError(
-                    "Gather event already exists, which should not happen.")
-            state.gathered_grad = torch.cat(gather_list, dim=0)
-            state.gather_event = torch.cuda.Event()
-            state.gather_event.record()
-        else:
-            state.gathered_grad = None
-            state.gather_event = None
-        gather_list = None
-        if none_grad:
-            # We can safely free p.grad without calling record_stream:
-            #   p.grad.to_local().record_stream(comm_stream)
-            # Explanation:
-            # 1. p.grad is created on the default stream, but the default stream
-            #    is synchronized with the comm stream later.
-            # 2. There is no further activity on the default stream before the optimizer finishes.
-            # Therefore, it is safe to free p.grad directly on the comm stream.
-            p.grad = None
+
+        # Reconstructs gathered grad from the received buffer
+        #
+        #                  recv_buf (num ranks = 3)
+        #
+        #      From rank 0        From rank 1        From rank 2
+        # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # p1_n -> p2_n -> p3_n
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        write_offsets = {id(p): 0 for p in owner_params}
+        for src in range(num_ranks):
+            if recv_counts[src] == 0:
+                continue
+
+            block = recv_counts[src]
+            inner_off = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                n = split_elems_for_src(p, state, src, num_ranks)
+                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)
+                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:
+                state.gathered_grad = None
+                state.gather_event = None
+            if none_grad:
+                p.grad = None
 
 
 @torch.no_grad()
@@ -127,45 +211,119 @@ def _compute_u(p, state, steps, rank, compute_stream):
                 raise RuntimeError("Gather event must be set before compute.")
             compute_stream.wait_event(state.gather_event)
             u = _zeropower_via_newtonschulz5(state.gathered_grad, steps)
+            state.gathered_grad = None
             state.computed_u = u
-        state.scattered_u = torch.empty_like(p.to_local(),
-                                             dtype=torch.bfloat16)
-        state.compute_event = torch.cuda.Event()
-        state.compute_event.record()
-        u = None
+            state.compute_event = torch.cuda.Event()
+            state.compute_event.record()
+        else:
+            state.computed_u = None
+            state.compute_event = None
 
 
 @torch.no_grad()
-def _scatter(p, state, rank, comm_stream):
-    """
-    Scatter the computed_u from worker_rank to all ranks.
-    """
+def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(compute_stream):
+        for p in params:
+            state = param_to_state[id(p)]
+            state.scattered_u = torch.empty_like(p.to_local(),
+                                                 dtype=COMM_DTYPE)
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
 
+
+def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
     with torch.cuda.stream(comm_stream):
-        if state.compute_event is None:
-            raise RuntimeError("Compute event must be set before scatter.")
-        comm_stream.wait_event(state.compute_event)
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            # Clear the gathered gradient to free memory
-            state.gathered_grad = None
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
 
-            u = state.computed_u
-            scatter_list = list(torch.split(u, p.size(0) // num_ranks, dim=0))
-            scatter_list = [s.contiguous() for s in scatter_list]
-        else:
-            scatter_list = None
+        if owner_params:
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                if state.compute_event is None:
+                    raise RuntimeError(
+                        "Compute event must be set before scatter.")
+                comm_stream.wait_event(state.compute_event)
+                state.gathered_grad = None
 
-        torch.distributed.scatter(
-            state.scattered_u,
-            scatter_list=scatter_list,
-            src=state.worker_rank,
-            group=state.process_group,
+                assert state.computed_u is not None
+
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+
+                offset = 0
+                for dst in range(num_ranks):
+                    n = split_elems_for_src(p, state, dst, num_ranks)
+                    assert n > 0
+
+                    su = u_full.narrow(0, offset, n)
+                    per_dst[dst].append(su)
+                    send_counts[dst] += n
+                    offset += n
+
+                assert offset == u_full.numel()
+
+        if any(len(v) > 0 for v in per_dst):
+            send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        else:
+            # all_to_all requires participation from all ranks
+            # Even non-owner ranks must join the collective call
+            send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
+
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                total += split_elems_for_src(p, state, rank, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        assert recv_total > 0
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        state.scatter_event = torch.cuda.Event()
-        state.scatter_event.record()
-        scatter_list = None
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        for src in range(num_ranks):
+            block = recv_counts[src]
+            if block == 0:
+                continue
+
+            inner_off = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                n = split_elems_for_src(p, state, rank, num_ranks)
+                assert n > 0
+
+                flat_local = recv_buf.narrow(0, off + inner_off,
+                                             n).view_as(p.to_local())
+                state.scattered_u.copy_(flat_local)
+
+                state.scatter_event = torch.cuda.Event()
+                state.scatter_event.record(comm_stream)
+                inner_off += n
+
+            assert inner_off == block
+            off += block
 
 
 def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
@@ -444,11 +602,14 @@ class Muon(torch.optim.Optimizer):
             if mesh is None:
                 mesh = p.device_mesh
                 shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
+                self.rank = dist.get_rank(group=process_group)
             elif 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()
+            param_to_state[id(
+                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
             param_to_state[id(p)].process_group = process_group
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
             param_to_state[id(p)].qk_clip_state = qk_clip_state
@@ -478,7 +639,7 @@ class Muon(torch.optim.Optimizer):
             else:
                 g = buf
 
-            u = _zeropower_via_newtonschulz5(g.bfloat16(),
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
                                              steps=group["ns_steps"])
 
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
@@ -493,15 +654,12 @@ class Muon(torch.optim.Optimizer):
     def _update_g(self, p, g, group, momentum):
         # 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)
+        buf = state.setdefault("momentum_buffer", torch.zeros_like(g))
+        torch.add(g, buf, alpha=momentum, out=buf)
         if group["nesterov"]:
-            g = g.add(buf, alpha=momentum)
-        else:
-            g = buf
-        return g
+            g.add_(buf, alpha=momentum)
+            return g
+        return buf
 
     @staticmethod
     def _update_p(p, u, lr, adjusted_lr, weight_decay):
@@ -585,11 +743,15 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
-        def enqueue_gathers(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _gather(p, state, self.rank, self.comm_stream,
-                        group["none_grad"])
+        def enqueue_all2all_gather(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_gathered_grad(target_params,
+                                                   param_to_state, self.rank,
+                                                   self.compute_stream)
+                _all2all_gather(target_params, param_to_state, self.rank,
+                                self.comm_stream, group["none_grad"],
+                                alloc_event)
 
         def enqueue_computes(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -597,10 +759,14 @@ class Muon(torch.optim.Optimizer):
                 _compute_u(p, state, group["ns_steps"], self.rank,
                            self.compute_stream)
 
-        def enqueue_scatters(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _scatter(p, state, self.rank, self.comm_stream)
+        def enqueue_all2all_scatter(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_scattered_u(target_params, param_to_state,
+                                                 self.rank,
+                                                 self.compute_stream)
+                _all2all_scatter(target_params, param_to_state, self.rank,
+                                 self.comm_stream, alloc_event)
 
         def enqueue_update_param(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -615,14 +781,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        enqueue_gathers(0, chunk_size)
+        PRE_STEP = 5
+        for i in range(0, PRE_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_computes(i, chunk_size)
-            if i > 0:
-                enqueue_update_param(i - chunk_size, chunk_size)
-            enqueue_gathers(i + chunk_size, chunk_size)
-            enqueue_scatters(i, chunk_size)
-        enqueue_update_param(i, chunk_size)
+            enqueue_all2all_scatter(i, chunk_size)
+            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_update_param(i, chunk_size)
+            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)

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

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

build/torch27-cxx11-rocm63-x86_64-linux/optimizer/{_optimizer_9c21645_dirty.abi3.so → _optimizer_fceb334_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:41492cb1479920b654768a5597d88670dd0caeedbdcd73fd63afa31ffc6961d6
+oid sha256:083283a8bd6037ed40c50acc478806786547298abe51a9494bcac0e468ca09fd
 size 1749776

build/torch27-cxx11-rocm63-x86_64-linux/optimizer/matmul_transpose_triton.py

@@ -0,0 +1,106 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def get_autotune_config():
+    return [
+        triton.Config(
+            {
+                'BLOCK_SIZE_M': blk_m,
+                'BLOCK_SIZE_K': blk_k,
+                'GROUP_SIZE_M': grp_sz
+            },
+            num_stages=n_stages,
+            num_warps=n_warps) for blk_m in [32, 64, 128]
+        for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5]
+        for n_warps in [4, 8]
+    ]
+
+
+@triton.autotune(
+    configs=get_autotune_config(),
+    key=['M', 'K'],
+)
+@triton.jit
+def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn,
+               BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+               GROUP_SIZE_M: tl.constexpr):
+    """
+    Core kernel jit function of matmul_transpose that computes y = x @ x.T
+    The code is a simple adaptation from the triton `matmul` tutorial:
+    https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+    """
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    if pid_m > pid_n:
+        return
+
+    offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # we use a & b ptrs to denote different rows of x.
+    a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+    b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        b = tl.load(b_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator)
+        a_ptrs += BLOCK_SIZE_K * stride_xk
+        b_ptrs += BLOCK_SIZE_K * stride_xk
+    # use dtype.element_ty to accommodate different input datatypes as in cpp templates
+    # https://github.com/triton-lang/triton/issues/2252
+    c = accumulator.to(x.dtype.element_ty)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+    # transpose and copy
+    if pid_m < pid_n:
+        ct_ptrs = y + stride_ym * offs_cn[:,
+                                          None] + stride_yn * offs_cm[None, :]
+        ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+        tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask)
+
+
+def matmul_transpose_assign(d_in, d_out):
+    assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor"
+    assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor"
+    assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device"
+    assert d_in.dtype == d_out.dtype, "Inputs must have the same data type"
+    assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor"
+    assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor"
+    assert d_in.size(0) == d_out.size(0) == d_out.size(0), \
+            "First dimension of `d_in` must match first and second dimension of `d_out`"
+
+    d_in = d_in.contiguous()
+    M, K = d_in.shape
+    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(
+        M, META['BLOCK_SIZE_M']), )
+    with torch.cuda.device(d_in.device.index):
+        mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1),
+                         d_out.stride(0), d_out.stride(1))
+
+
+def matmul_transpose(d_in):
+    M, _ = d_in.shape
+    d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype)
+    matmul_transpose_assign(d_in, d_out)
+    return d_out

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

@@ -8,14 +8,19 @@ import torch
 import torch.distributed as dist
 from torch.distributed._tensor import DTensor, Replicate, Shard
 
+from .matmul_transpose_triton import matmul_transpose_assign
+
 logger = logging.getLogger(__name__)
 
+COMM_DTYPE = torch.bfloat16
+
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
 # Muon's Newton–Schulz iteration causes high variance in singular values
 # Idea: give each iteration its own 3 coefficients and optimize them via gradient descent.
 @torch.no_grad()
+# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon
 def _zeropower_via_newtonschulz5(G, steps):
     """
     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
@@ -27,13 +32,15 @@ def _zeropower_via_newtonschulz5(G, steps):
     performance at all relative to UV^T, where USV^T = G is the SVD.
     """
     assert len(G.shape) == 2
-    assert G.dtype == torch.bfloat16
+    assert G.dtype == COMM_DTYPE
     X = G  # no manual typecast
 
     if G.size(0) > G.size(1):
         X = X.T
     # Ensure spectral norm is at most 1
     X = X / (X.norm() + 1e-7)
+    buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
+    buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
     for a, b, c in [
         (4.0848, -6.8946, 2.9270),
@@ -42,13 +49,10 @@ def _zeropower_via_newtonschulz5(G, steps):
         (2.8769, -3.1427, 1.2046),
         (2.8366, -3.0525, 1.2012),
     ]:
-        A = X @ X.T
-        # B = (
-        #    b * A + c * A @ A
-        # )
-        B = torch.addmm(A, A, A, alpha=c, beta=b)
-        # X = a * X + B @ X
-        X = torch.addmm(X, B, X, alpha=1.0, beta=a)
+        matmul_transpose_assign(X, buf1)
+        matmul_transpose_assign(buf1, buf2)
+        buf1.mul_(b).add_(buf2, alpha=c)
+        X = torch.addmm(X, buf1, X, alpha=1.0, beta=a)
 
     if G.size(0) > G.size(1):
         X = X.T
@@ -69,51 +73,131 @@ class _muon_state:
     qk_clip_state = None
 
 
+def split_elems_for_src(param, state, 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
+
+
 @torch.no_grad()
-def _gather(p, state, rank, comm_stream, none_grad):
-    """
-    Gather the gradients to worker_rank.
-    If none_grad is True, free p.grad after the gather.
-    """
+def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(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(),
+                                                  dtype=COMM_DTYPE,
+                                                  device="cuda")
+            else:
+                state.gathered_grad = None
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
+
+
+@torch.no_grad()
+def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
+                    alloc_event):
     with torch.cuda.stream(comm_stream):
-        g = p.grad
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            gather_list = [
-                torch.empty_like(g.to_local(), dtype=torch.bfloat16)
-                for _ in range(num_ranks)
-            ]
-        else:
-            gather_list = None
-
-        g = g.to(torch.bfloat16)
-        torch.distributed.gather(
-            g.to_local(),
-            dst=state.worker_rank,
-            gather_list=gather_list,
-            group=state.process_group,
+        # Calculate sending tensors
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
+        for p in params:
+            state = param_to_state[id(p)]
+            dst = state.worker_rank
+            assert dst < num_ranks
+            shard_elems = split_elems_for_src(p, state, rank, num_ranks)
+            g = p.grad
+            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            assert g.numel() == shard_elems
+            per_dst[dst].append(g)
+            send_counts[dst] += shard_elems
+
+        assert all(
+            len(v) > 0
+            for v in per_dst), "all params should be sharded to all devices"
+
+        send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
+
+        # Calculate receiving tensors
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                total += split_elems_for_src(p, state, src, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        if rank == state.worker_rank:
-            if state.gathered_grad is not None:
-                raise RuntimeError(
-                    "Gather event already exists, which should not happen.")
-            state.gathered_grad = torch.cat(gather_list, dim=0)
-            state.gather_event = torch.cuda.Event()
-            state.gather_event.record()
-        else:
-            state.gathered_grad = None
-            state.gather_event = None
-        gather_list = None
-        if none_grad:
-            # We can safely free p.grad without calling record_stream:
-            #   p.grad.to_local().record_stream(comm_stream)
-            # Explanation:
-            # 1. p.grad is created on the default stream, but the default stream
-            #    is synchronized with the comm stream later.
-            # 2. There is no further activity on the default stream before the optimizer finishes.
-            # Therefore, it is safe to free p.grad directly on the comm stream.
-            p.grad = None
+
+        # Reconstructs gathered grad from the received buffer
+        #
+        #                  recv_buf (num ranks = 3)
+        #
+        #      From rank 0        From rank 1        From rank 2
+        # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # p1_n -> p2_n -> p3_n
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        write_offsets = {id(p): 0 for p in owner_params}
+        for src in range(num_ranks):
+            if recv_counts[src] == 0:
+                continue
+
+            block = recv_counts[src]
+            inner_off = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                n = split_elems_for_src(p, state, src, num_ranks)
+                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)
+                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:
+                state.gathered_grad = None
+                state.gather_event = None
+            if none_grad:
+                p.grad = None
 
 
 @torch.no_grad()
@@ -127,45 +211,119 @@ def _compute_u(p, state, steps, rank, compute_stream):
                 raise RuntimeError("Gather event must be set before compute.")
             compute_stream.wait_event(state.gather_event)
             u = _zeropower_via_newtonschulz5(state.gathered_grad, steps)
+            state.gathered_grad = None
             state.computed_u = u
-        state.scattered_u = torch.empty_like(p.to_local(),
-                                             dtype=torch.bfloat16)
-        state.compute_event = torch.cuda.Event()
-        state.compute_event.record()
-        u = None
+            state.compute_event = torch.cuda.Event()
+            state.compute_event.record()
+        else:
+            state.computed_u = None
+            state.compute_event = None
 
 
 @torch.no_grad()
-def _scatter(p, state, rank, comm_stream):
-    """
-    Scatter the computed_u from worker_rank to all ranks.
-    """
+def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(compute_stream):
+        for p in params:
+            state = param_to_state[id(p)]
+            state.scattered_u = torch.empty_like(p.to_local(),
+                                                 dtype=COMM_DTYPE)
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
 
+
+def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
     with torch.cuda.stream(comm_stream):
-        if state.compute_event is None:
-            raise RuntimeError("Compute event must be set before scatter.")
-        comm_stream.wait_event(state.compute_event)
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            # Clear the gathered gradient to free memory
-            state.gathered_grad = None
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
 
-            u = state.computed_u
-            scatter_list = list(torch.split(u, p.size(0) // num_ranks, dim=0))
-            scatter_list = [s.contiguous() for s in scatter_list]
-        else:
-            scatter_list = None
+        if owner_params:
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                if state.compute_event is None:
+                    raise RuntimeError(
+                        "Compute event must be set before scatter.")
+                comm_stream.wait_event(state.compute_event)
+                state.gathered_grad = None
 
-        torch.distributed.scatter(
-            state.scattered_u,
-            scatter_list=scatter_list,
-            src=state.worker_rank,
-            group=state.process_group,
+                assert state.computed_u is not None
+
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+
+                offset = 0
+                for dst in range(num_ranks):
+                    n = split_elems_for_src(p, state, dst, num_ranks)
+                    assert n > 0
+
+                    su = u_full.narrow(0, offset, n)
+                    per_dst[dst].append(su)
+                    send_counts[dst] += n
+                    offset += n
+
+                assert offset == u_full.numel()
+
+        if any(len(v) > 0 for v in per_dst):
+            send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        else:
+            # all_to_all requires participation from all ranks
+            # Even non-owner ranks must join the collective call
+            send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
+
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                total += split_elems_for_src(p, state, rank, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        assert recv_total > 0
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        state.scatter_event = torch.cuda.Event()
-        state.scatter_event.record()
-        scatter_list = None
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        for src in range(num_ranks):
+            block = recv_counts[src]
+            if block == 0:
+                continue
+
+            inner_off = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                n = split_elems_for_src(p, state, rank, num_ranks)
+                assert n > 0
+
+                flat_local = recv_buf.narrow(0, off + inner_off,
+                                             n).view_as(p.to_local())
+                state.scattered_u.copy_(flat_local)
+
+                state.scatter_event = torch.cuda.Event()
+                state.scatter_event.record(comm_stream)
+                inner_off += n
+
+            assert inner_off == block
+            off += block
 
 
 def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
@@ -444,11 +602,14 @@ class Muon(torch.optim.Optimizer):
             if mesh is None:
                 mesh = p.device_mesh
                 shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
+                self.rank = dist.get_rank(group=process_group)
             elif 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()
+            param_to_state[id(
+                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
             param_to_state[id(p)].process_group = process_group
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
             param_to_state[id(p)].qk_clip_state = qk_clip_state
@@ -478,7 +639,7 @@ class Muon(torch.optim.Optimizer):
             else:
                 g = buf
 
-            u = _zeropower_via_newtonschulz5(g.bfloat16(),
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
                                              steps=group["ns_steps"])
 
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
@@ -493,15 +654,12 @@ class Muon(torch.optim.Optimizer):
     def _update_g(self, p, g, group, momentum):
         # 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)
+        buf = state.setdefault("momentum_buffer", torch.zeros_like(g))
+        torch.add(g, buf, alpha=momentum, out=buf)
         if group["nesterov"]:
-            g = g.add(buf, alpha=momentum)
-        else:
-            g = buf
-        return g
+            g.add_(buf, alpha=momentum)
+            return g
+        return buf
 
     @staticmethod
     def _update_p(p, u, lr, adjusted_lr, weight_decay):
@@ -585,11 +743,15 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
-        def enqueue_gathers(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _gather(p, state, self.rank, self.comm_stream,
-                        group["none_grad"])
+        def enqueue_all2all_gather(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_gathered_grad(target_params,
+                                                   param_to_state, self.rank,
+                                                   self.compute_stream)
+                _all2all_gather(target_params, param_to_state, self.rank,
+                                self.comm_stream, group["none_grad"],
+                                alloc_event)
 
         def enqueue_computes(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -597,10 +759,14 @@ class Muon(torch.optim.Optimizer):
                 _compute_u(p, state, group["ns_steps"], self.rank,
                            self.compute_stream)
 
-        def enqueue_scatters(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _scatter(p, state, self.rank, self.comm_stream)
+        def enqueue_all2all_scatter(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_scattered_u(target_params, param_to_state,
+                                                 self.rank,
+                                                 self.compute_stream)
+                _all2all_scatter(target_params, param_to_state, self.rank,
+                                 self.comm_stream, alloc_event)
 
         def enqueue_update_param(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -615,14 +781,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        enqueue_gathers(0, chunk_size)
+        PRE_STEP = 5
+        for i in range(0, PRE_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_computes(i, chunk_size)
-            if i > 0:
-                enqueue_update_param(i - chunk_size, chunk_size)
-            enqueue_gathers(i + chunk_size, chunk_size)
-            enqueue_scatters(i, chunk_size)
-        enqueue_update_param(i, chunk_size)
+            enqueue_all2all_scatter(i, chunk_size)
+            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_update_param(i, chunk_size)
+            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)

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

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

build/{torch27-cxx11-cu126-x86_64-linux/optimizer/_optimizer_9c21645_dirty.abi3.so → torch28-cxx11-cu126-x86_64-linux/optimizer/_optimizer_fceb334_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:42ae6ac1cf967d7d23cac7930c8db635105f60631220a60b9cee060d082f40ae
+oid sha256:c31efe0c8bad71c6fbc002f269eff7a30754e1583bb8c892b8ce5cdbdfc5121c
 size 1824256

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

@@ -0,0 +1,106 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def get_autotune_config():
+    return [
+        triton.Config(
+            {
+                'BLOCK_SIZE_M': blk_m,
+                'BLOCK_SIZE_K': blk_k,
+                'GROUP_SIZE_M': grp_sz
+            },
+            num_stages=n_stages,
+            num_warps=n_warps) for blk_m in [32, 64, 128]
+        for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5]
+        for n_warps in [4, 8]
+    ]
+
+
+@triton.autotune(
+    configs=get_autotune_config(),
+    key=['M', 'K'],
+)
+@triton.jit
+def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn,
+               BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+               GROUP_SIZE_M: tl.constexpr):
+    """
+    Core kernel jit function of matmul_transpose that computes y = x @ x.T
+    The code is a simple adaptation from the triton `matmul` tutorial:
+    https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+    """
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    if pid_m > pid_n:
+        return
+
+    offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # we use a & b ptrs to denote different rows of x.
+    a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+    b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        b = tl.load(b_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator)
+        a_ptrs += BLOCK_SIZE_K * stride_xk
+        b_ptrs += BLOCK_SIZE_K * stride_xk
+    # use dtype.element_ty to accommodate different input datatypes as in cpp templates
+    # https://github.com/triton-lang/triton/issues/2252
+    c = accumulator.to(x.dtype.element_ty)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+    # transpose and copy
+    if pid_m < pid_n:
+        ct_ptrs = y + stride_ym * offs_cn[:,
+                                          None] + stride_yn * offs_cm[None, :]
+        ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+        tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask)
+
+
+def matmul_transpose_assign(d_in, d_out):
+    assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor"
+    assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor"
+    assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device"
+    assert d_in.dtype == d_out.dtype, "Inputs must have the same data type"
+    assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor"
+    assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor"
+    assert d_in.size(0) == d_out.size(0) == d_out.size(0), \
+            "First dimension of `d_in` must match first and second dimension of `d_out`"
+
+    d_in = d_in.contiguous()
+    M, K = d_in.shape
+    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(
+        M, META['BLOCK_SIZE_M']), )
+    with torch.cuda.device(d_in.device.index):
+        mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1),
+                         d_out.stride(0), d_out.stride(1))
+
+
+def matmul_transpose(d_in):
+    M, _ = d_in.shape
+    d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype)
+    matmul_transpose_assign(d_in, d_out)
+    return d_out

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

@@ -8,14 +8,19 @@ import torch
 import torch.distributed as dist
 from torch.distributed._tensor import DTensor, Replicate, Shard
 
+from .matmul_transpose_triton import matmul_transpose_assign
+
 logger = logging.getLogger(__name__)
 
+COMM_DTYPE = torch.bfloat16
+
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
 # Muon's Newton–Schulz iteration causes high variance in singular values
 # Idea: give each iteration its own 3 coefficients and optimize them via gradient descent.
 @torch.no_grad()
+# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon
 def _zeropower_via_newtonschulz5(G, steps):
     """
     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
@@ -27,13 +32,15 @@ def _zeropower_via_newtonschulz5(G, steps):
     performance at all relative to UV^T, where USV^T = G is the SVD.
     """
     assert len(G.shape) == 2
-    assert G.dtype == torch.bfloat16
+    assert G.dtype == COMM_DTYPE
     X = G  # no manual typecast
 
     if G.size(0) > G.size(1):
         X = X.T
     # Ensure spectral norm is at most 1
     X = X / (X.norm() + 1e-7)
+    buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
+    buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
     for a, b, c in [
         (4.0848, -6.8946, 2.9270),
@@ -42,13 +49,10 @@ def _zeropower_via_newtonschulz5(G, steps):
         (2.8769, -3.1427, 1.2046),
         (2.8366, -3.0525, 1.2012),
     ]:
-        A = X @ X.T
-        # B = (
-        #    b * A + c * A @ A
-        # )
-        B = torch.addmm(A, A, A, alpha=c, beta=b)
-        # X = a * X + B @ X
-        X = torch.addmm(X, B, X, alpha=1.0, beta=a)
+        matmul_transpose_assign(X, buf1)
+        matmul_transpose_assign(buf1, buf2)
+        buf1.mul_(b).add_(buf2, alpha=c)
+        X = torch.addmm(X, buf1, X, alpha=1.0, beta=a)
 
     if G.size(0) > G.size(1):
         X = X.T
@@ -69,51 +73,131 @@ class _muon_state:
     qk_clip_state = None
 
 
+def split_elems_for_src(param, state, 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
+
+
 @torch.no_grad()
-def _gather(p, state, rank, comm_stream, none_grad):
-    """
-    Gather the gradients to worker_rank.
-    If none_grad is True, free p.grad after the gather.
-    """
+def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(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(),
+                                                  dtype=COMM_DTYPE,
+                                                  device="cuda")
+            else:
+                state.gathered_grad = None
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
+
+
+@torch.no_grad()
+def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
+                    alloc_event):
     with torch.cuda.stream(comm_stream):
-        g = p.grad
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            gather_list = [
-                torch.empty_like(g.to_local(), dtype=torch.bfloat16)
-                for _ in range(num_ranks)
-            ]
-        else:
-            gather_list = None
-
-        g = g.to(torch.bfloat16)
-        torch.distributed.gather(
-            g.to_local(),
-            dst=state.worker_rank,
-            gather_list=gather_list,
-            group=state.process_group,
+        # Calculate sending tensors
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
+        for p in params:
+            state = param_to_state[id(p)]
+            dst = state.worker_rank
+            assert dst < num_ranks
+            shard_elems = split_elems_for_src(p, state, rank, num_ranks)
+            g = p.grad
+            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            assert g.numel() == shard_elems
+            per_dst[dst].append(g)
+            send_counts[dst] += shard_elems
+
+        assert all(
+            len(v) > 0
+            for v in per_dst), "all params should be sharded to all devices"
+
+        send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
+
+        # Calculate receiving tensors
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                total += split_elems_for_src(p, state, src, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        if rank == state.worker_rank:
-            if state.gathered_grad is not None:
-                raise RuntimeError(
-                    "Gather event already exists, which should not happen.")
-            state.gathered_grad = torch.cat(gather_list, dim=0)
-            state.gather_event = torch.cuda.Event()
-            state.gather_event.record()
-        else:
-            state.gathered_grad = None
-            state.gather_event = None
-        gather_list = None
-        if none_grad:
-            # We can safely free p.grad without calling record_stream:
-            #   p.grad.to_local().record_stream(comm_stream)
-            # Explanation:
-            # 1. p.grad is created on the default stream, but the default stream
-            #    is synchronized with the comm stream later.
-            # 2. There is no further activity on the default stream before the optimizer finishes.
-            # Therefore, it is safe to free p.grad directly on the comm stream.
-            p.grad = None
+
+        # Reconstructs gathered grad from the received buffer
+        #
+        #                  recv_buf (num ranks = 3)
+        #
+        #      From rank 0        From rank 1        From rank 2
+        # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # p1_n -> p2_n -> p3_n
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        write_offsets = {id(p): 0 for p in owner_params}
+        for src in range(num_ranks):
+            if recv_counts[src] == 0:
+                continue
+
+            block = recv_counts[src]
+            inner_off = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                n = split_elems_for_src(p, state, src, num_ranks)
+                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)
+                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:
+                state.gathered_grad = None
+                state.gather_event = None
+            if none_grad:
+                p.grad = None
 
 
 @torch.no_grad()
@@ -127,45 +211,119 @@ def _compute_u(p, state, steps, rank, compute_stream):
                 raise RuntimeError("Gather event must be set before compute.")
             compute_stream.wait_event(state.gather_event)
             u = _zeropower_via_newtonschulz5(state.gathered_grad, steps)
+            state.gathered_grad = None
             state.computed_u = u
-        state.scattered_u = torch.empty_like(p.to_local(),
-                                             dtype=torch.bfloat16)
-        state.compute_event = torch.cuda.Event()
-        state.compute_event.record()
-        u = None
+            state.compute_event = torch.cuda.Event()
+            state.compute_event.record()
+        else:
+            state.computed_u = None
+            state.compute_event = None
 
 
 @torch.no_grad()
-def _scatter(p, state, rank, comm_stream):
-    """
-    Scatter the computed_u from worker_rank to all ranks.
-    """
+def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(compute_stream):
+        for p in params:
+            state = param_to_state[id(p)]
+            state.scattered_u = torch.empty_like(p.to_local(),
+                                                 dtype=COMM_DTYPE)
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
 
+
+def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
     with torch.cuda.stream(comm_stream):
-        if state.compute_event is None:
-            raise RuntimeError("Compute event must be set before scatter.")
-        comm_stream.wait_event(state.compute_event)
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            # Clear the gathered gradient to free memory
-            state.gathered_grad = None
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
 
-            u = state.computed_u
-            scatter_list = list(torch.split(u, p.size(0) // num_ranks, dim=0))
-            scatter_list = [s.contiguous() for s in scatter_list]
-        else:
-            scatter_list = None
+        if owner_params:
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                if state.compute_event is None:
+                    raise RuntimeError(
+                        "Compute event must be set before scatter.")
+                comm_stream.wait_event(state.compute_event)
+                state.gathered_grad = None
 
-        torch.distributed.scatter(
-            state.scattered_u,
-            scatter_list=scatter_list,
-            src=state.worker_rank,
-            group=state.process_group,
+                assert state.computed_u is not None
+
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+
+                offset = 0
+                for dst in range(num_ranks):
+                    n = split_elems_for_src(p, state, dst, num_ranks)
+                    assert n > 0
+
+                    su = u_full.narrow(0, offset, n)
+                    per_dst[dst].append(su)
+                    send_counts[dst] += n
+                    offset += n
+
+                assert offset == u_full.numel()
+
+        if any(len(v) > 0 for v in per_dst):
+            send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        else:
+            # all_to_all requires participation from all ranks
+            # Even non-owner ranks must join the collective call
+            send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
+
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                total += split_elems_for_src(p, state, rank, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        assert recv_total > 0
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        state.scatter_event = torch.cuda.Event()
-        state.scatter_event.record()
-        scatter_list = None
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        for src in range(num_ranks):
+            block = recv_counts[src]
+            if block == 0:
+                continue
+
+            inner_off = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                n = split_elems_for_src(p, state, rank, num_ranks)
+                assert n > 0
+
+                flat_local = recv_buf.narrow(0, off + inner_off,
+                                             n).view_as(p.to_local())
+                state.scattered_u.copy_(flat_local)
+
+                state.scatter_event = torch.cuda.Event()
+                state.scatter_event.record(comm_stream)
+                inner_off += n
+
+            assert inner_off == block
+            off += block
 
 
 def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
@@ -444,11 +602,14 @@ class Muon(torch.optim.Optimizer):
             if mesh is None:
                 mesh = p.device_mesh
                 shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
+                self.rank = dist.get_rank(group=process_group)
             elif 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()
+            param_to_state[id(
+                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
             param_to_state[id(p)].process_group = process_group
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
             param_to_state[id(p)].qk_clip_state = qk_clip_state
@@ -478,7 +639,7 @@ class Muon(torch.optim.Optimizer):
             else:
                 g = buf
 
-            u = _zeropower_via_newtonschulz5(g.bfloat16(),
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
                                              steps=group["ns_steps"])
 
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
@@ -493,15 +654,12 @@ class Muon(torch.optim.Optimizer):
     def _update_g(self, p, g, group, momentum):
         # 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)
+        buf = state.setdefault("momentum_buffer", torch.zeros_like(g))
+        torch.add(g, buf, alpha=momentum, out=buf)
         if group["nesterov"]:
-            g = g.add(buf, alpha=momentum)
-        else:
-            g = buf
-        return g
+            g.add_(buf, alpha=momentum)
+            return g
+        return buf
 
     @staticmethod
     def _update_p(p, u, lr, adjusted_lr, weight_decay):
@@ -585,11 +743,15 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
-        def enqueue_gathers(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _gather(p, state, self.rank, self.comm_stream,
-                        group["none_grad"])
+        def enqueue_all2all_gather(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_gathered_grad(target_params,
+                                                   param_to_state, self.rank,
+                                                   self.compute_stream)
+                _all2all_gather(target_params, param_to_state, self.rank,
+                                self.comm_stream, group["none_grad"],
+                                alloc_event)
 
         def enqueue_computes(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -597,10 +759,14 @@ class Muon(torch.optim.Optimizer):
                 _compute_u(p, state, group["ns_steps"], self.rank,
                            self.compute_stream)
 
-        def enqueue_scatters(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _scatter(p, state, self.rank, self.comm_stream)
+        def enqueue_all2all_scatter(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_scattered_u(target_params, param_to_state,
+                                                 self.rank,
+                                                 self.compute_stream)
+                _all2all_scatter(target_params, param_to_state, self.rank,
+                                 self.comm_stream, alloc_event)
 
         def enqueue_update_param(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -615,14 +781,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        enqueue_gathers(0, chunk_size)
+        PRE_STEP = 5
+        for i in range(0, PRE_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_computes(i, chunk_size)
-            if i > 0:
-                enqueue_update_param(i - chunk_size, chunk_size)
-            enqueue_gathers(i + chunk_size, chunk_size)
-            enqueue_scatters(i, chunk_size)
-        enqueue_update_param(i, chunk_size)
+            enqueue_all2all_scatter(i, chunk_size)
+            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_update_param(i, chunk_size)
+            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)

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

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

build/{torch28-cxx11-cu129-x86_64-linux/optimizer/_optimizer_9c21645_dirty.abi3.so → torch28-cxx11-cu128-x86_64-linux/optimizer/_optimizer_fceb334_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eb40a06623bb3668b82ff248b5a3c1bcf41e7f3f860888b261505b3a71257bc7
+oid sha256:5b7313ff1257a692e4e5cfdec474b10a7a795f2e87edb0099c7e5d21f588634f
 size 1883344

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

@@ -0,0 +1,106 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def get_autotune_config():
+    return [
+        triton.Config(
+            {
+                'BLOCK_SIZE_M': blk_m,
+                'BLOCK_SIZE_K': blk_k,
+                'GROUP_SIZE_M': grp_sz
+            },
+            num_stages=n_stages,
+            num_warps=n_warps) for blk_m in [32, 64, 128]
+        for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5]
+        for n_warps in [4, 8]
+    ]
+
+
+@triton.autotune(
+    configs=get_autotune_config(),
+    key=['M', 'K'],
+)
+@triton.jit
+def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn,
+               BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+               GROUP_SIZE_M: tl.constexpr):
+    """
+    Core kernel jit function of matmul_transpose that computes y = x @ x.T
+    The code is a simple adaptation from the triton `matmul` tutorial:
+    https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+    """
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    if pid_m > pid_n:
+        return
+
+    offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # we use a & b ptrs to denote different rows of x.
+    a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+    b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        b = tl.load(b_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator)
+        a_ptrs += BLOCK_SIZE_K * stride_xk
+        b_ptrs += BLOCK_SIZE_K * stride_xk
+    # use dtype.element_ty to accommodate different input datatypes as in cpp templates
+    # https://github.com/triton-lang/triton/issues/2252
+    c = accumulator.to(x.dtype.element_ty)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+    # transpose and copy
+    if pid_m < pid_n:
+        ct_ptrs = y + stride_ym * offs_cn[:,
+                                          None] + stride_yn * offs_cm[None, :]
+        ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+        tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask)
+
+
+def matmul_transpose_assign(d_in, d_out):
+    assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor"
+    assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor"
+    assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device"
+    assert d_in.dtype == d_out.dtype, "Inputs must have the same data type"
+    assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor"
+    assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor"
+    assert d_in.size(0) == d_out.size(0) == d_out.size(0), \
+            "First dimension of `d_in` must match first and second dimension of `d_out`"
+
+    d_in = d_in.contiguous()
+    M, K = d_in.shape
+    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(
+        M, META['BLOCK_SIZE_M']), )
+    with torch.cuda.device(d_in.device.index):
+        mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1),
+                         d_out.stride(0), d_out.stride(1))
+
+
+def matmul_transpose(d_in):
+    M, _ = d_in.shape
+    d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype)
+    matmul_transpose_assign(d_in, d_out)
+    return d_out

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

@@ -8,14 +8,19 @@ import torch
 import torch.distributed as dist
 from torch.distributed._tensor import DTensor, Replicate, Shard
 
+from .matmul_transpose_triton import matmul_transpose_assign
+
 logger = logging.getLogger(__name__)
 
+COMM_DTYPE = torch.bfloat16
+
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
 # Muon's Newton–Schulz iteration causes high variance in singular values
 # Idea: give each iteration its own 3 coefficients and optimize them via gradient descent.
 @torch.no_grad()
+# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon
 def _zeropower_via_newtonschulz5(G, steps):
     """
     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
@@ -27,13 +32,15 @@ def _zeropower_via_newtonschulz5(G, steps):
     performance at all relative to UV^T, where USV^T = G is the SVD.
     """
     assert len(G.shape) == 2
-    assert G.dtype == torch.bfloat16
+    assert G.dtype == COMM_DTYPE
     X = G  # no manual typecast
 
     if G.size(0) > G.size(1):
         X = X.T
     # Ensure spectral norm is at most 1
     X = X / (X.norm() + 1e-7)
+    buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
+    buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
     for a, b, c in [
         (4.0848, -6.8946, 2.9270),
@@ -42,13 +49,10 @@ def _zeropower_via_newtonschulz5(G, steps):
         (2.8769, -3.1427, 1.2046),
         (2.8366, -3.0525, 1.2012),
     ]:
-        A = X @ X.T
-        # B = (
-        #    b * A + c * A @ A
-        # )
-        B = torch.addmm(A, A, A, alpha=c, beta=b)
-        # X = a * X + B @ X
-        X = torch.addmm(X, B, X, alpha=1.0, beta=a)
+        matmul_transpose_assign(X, buf1)
+        matmul_transpose_assign(buf1, buf2)
+        buf1.mul_(b).add_(buf2, alpha=c)
+        X = torch.addmm(X, buf1, X, alpha=1.0, beta=a)
 
     if G.size(0) > G.size(1):
         X = X.T
@@ -69,51 +73,131 @@ class _muon_state:
     qk_clip_state = None
 
 
+def split_elems_for_src(param, state, 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
+
+
 @torch.no_grad()
-def _gather(p, state, rank, comm_stream, none_grad):
-    """
-    Gather the gradients to worker_rank.
-    If none_grad is True, free p.grad after the gather.
-    """
+def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(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(),
+                                                  dtype=COMM_DTYPE,
+                                                  device="cuda")
+            else:
+                state.gathered_grad = None
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
+
+
+@torch.no_grad()
+def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
+                    alloc_event):
     with torch.cuda.stream(comm_stream):
-        g = p.grad
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            gather_list = [
-                torch.empty_like(g.to_local(), dtype=torch.bfloat16)
-                for _ in range(num_ranks)
-            ]
-        else:
-            gather_list = None
-
-        g = g.to(torch.bfloat16)
-        torch.distributed.gather(
-            g.to_local(),
-            dst=state.worker_rank,
-            gather_list=gather_list,
-            group=state.process_group,
+        # Calculate sending tensors
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
+        for p in params:
+            state = param_to_state[id(p)]
+            dst = state.worker_rank
+            assert dst < num_ranks
+            shard_elems = split_elems_for_src(p, state, rank, num_ranks)
+            g = p.grad
+            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            assert g.numel() == shard_elems
+            per_dst[dst].append(g)
+            send_counts[dst] += shard_elems
+
+        assert all(
+            len(v) > 0
+            for v in per_dst), "all params should be sharded to all devices"
+
+        send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
+
+        # Calculate receiving tensors
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                total += split_elems_for_src(p, state, src, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        if rank == state.worker_rank:
-            if state.gathered_grad is not None:
-                raise RuntimeError(
-                    "Gather event already exists, which should not happen.")
-            state.gathered_grad = torch.cat(gather_list, dim=0)
-            state.gather_event = torch.cuda.Event()
-            state.gather_event.record()
-        else:
-            state.gathered_grad = None
-            state.gather_event = None
-        gather_list = None
-        if none_grad:
-            # We can safely free p.grad without calling record_stream:
-            #   p.grad.to_local().record_stream(comm_stream)
-            # Explanation:
-            # 1. p.grad is created on the default stream, but the default stream
-            #    is synchronized with the comm stream later.
-            # 2. There is no further activity on the default stream before the optimizer finishes.
-            # Therefore, it is safe to free p.grad directly on the comm stream.
-            p.grad = None
+
+        # Reconstructs gathered grad from the received buffer
+        #
+        #                  recv_buf (num ranks = 3)
+        #
+        #      From rank 0        From rank 1        From rank 2
+        # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # p1_n -> p2_n -> p3_n
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        write_offsets = {id(p): 0 for p in owner_params}
+        for src in range(num_ranks):
+            if recv_counts[src] == 0:
+                continue
+
+            block = recv_counts[src]
+            inner_off = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                n = split_elems_for_src(p, state, src, num_ranks)
+                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)
+                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:
+                state.gathered_grad = None
+                state.gather_event = None
+            if none_grad:
+                p.grad = None
 
 
 @torch.no_grad()
@@ -127,45 +211,119 @@ def _compute_u(p, state, steps, rank, compute_stream):
                 raise RuntimeError("Gather event must be set before compute.")
             compute_stream.wait_event(state.gather_event)
             u = _zeropower_via_newtonschulz5(state.gathered_grad, steps)
+            state.gathered_grad = None
             state.computed_u = u
-        state.scattered_u = torch.empty_like(p.to_local(),
-                                             dtype=torch.bfloat16)
-        state.compute_event = torch.cuda.Event()
-        state.compute_event.record()
-        u = None
+            state.compute_event = torch.cuda.Event()
+            state.compute_event.record()
+        else:
+            state.computed_u = None
+            state.compute_event = None
 
 
 @torch.no_grad()
-def _scatter(p, state, rank, comm_stream):
-    """
-    Scatter the computed_u from worker_rank to all ranks.
-    """
+def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(compute_stream):
+        for p in params:
+            state = param_to_state[id(p)]
+            state.scattered_u = torch.empty_like(p.to_local(),
+                                                 dtype=COMM_DTYPE)
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
 
+
+def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
     with torch.cuda.stream(comm_stream):
-        if state.compute_event is None:
-            raise RuntimeError("Compute event must be set before scatter.")
-        comm_stream.wait_event(state.compute_event)
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            # Clear the gathered gradient to free memory
-            state.gathered_grad = None
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
 
-            u = state.computed_u
-            scatter_list = list(torch.split(u, p.size(0) // num_ranks, dim=0))
-            scatter_list = [s.contiguous() for s in scatter_list]
-        else:
-            scatter_list = None
+        if owner_params:
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                if state.compute_event is None:
+                    raise RuntimeError(
+                        "Compute event must be set before scatter.")
+                comm_stream.wait_event(state.compute_event)
+                state.gathered_grad = None
 
-        torch.distributed.scatter(
-            state.scattered_u,
-            scatter_list=scatter_list,
-            src=state.worker_rank,
-            group=state.process_group,
+                assert state.computed_u is not None
+
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+
+                offset = 0
+                for dst in range(num_ranks):
+                    n = split_elems_for_src(p, state, dst, num_ranks)
+                    assert n > 0
+
+                    su = u_full.narrow(0, offset, n)
+                    per_dst[dst].append(su)
+                    send_counts[dst] += n
+                    offset += n
+
+                assert offset == u_full.numel()
+
+        if any(len(v) > 0 for v in per_dst):
+            send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        else:
+            # all_to_all requires participation from all ranks
+            # Even non-owner ranks must join the collective call
+            send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
+
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                total += split_elems_for_src(p, state, rank, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        assert recv_total > 0
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        state.scatter_event = torch.cuda.Event()
-        state.scatter_event.record()
-        scatter_list = None
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        for src in range(num_ranks):
+            block = recv_counts[src]
+            if block == 0:
+                continue
+
+            inner_off = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                n = split_elems_for_src(p, state, rank, num_ranks)
+                assert n > 0
+
+                flat_local = recv_buf.narrow(0, off + inner_off,
+                                             n).view_as(p.to_local())
+                state.scattered_u.copy_(flat_local)
+
+                state.scatter_event = torch.cuda.Event()
+                state.scatter_event.record(comm_stream)
+                inner_off += n
+
+            assert inner_off == block
+            off += block
 
 
 def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
@@ -444,11 +602,14 @@ class Muon(torch.optim.Optimizer):
             if mesh is None:
                 mesh = p.device_mesh
                 shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
+                self.rank = dist.get_rank(group=process_group)
             elif 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()
+            param_to_state[id(
+                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
             param_to_state[id(p)].process_group = process_group
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
             param_to_state[id(p)].qk_clip_state = qk_clip_state
@@ -478,7 +639,7 @@ class Muon(torch.optim.Optimizer):
             else:
                 g = buf
 
-            u = _zeropower_via_newtonschulz5(g.bfloat16(),
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
                                              steps=group["ns_steps"])
 
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
@@ -493,15 +654,12 @@ class Muon(torch.optim.Optimizer):
     def _update_g(self, p, g, group, momentum):
         # 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)
+        buf = state.setdefault("momentum_buffer", torch.zeros_like(g))
+        torch.add(g, buf, alpha=momentum, out=buf)
         if group["nesterov"]:
-            g = g.add(buf, alpha=momentum)
-        else:
-            g = buf
-        return g
+            g.add_(buf, alpha=momentum)
+            return g
+        return buf
 
     @staticmethod
     def _update_p(p, u, lr, adjusted_lr, weight_decay):
@@ -585,11 +743,15 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
-        def enqueue_gathers(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _gather(p, state, self.rank, self.comm_stream,
-                        group["none_grad"])
+        def enqueue_all2all_gather(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_gathered_grad(target_params,
+                                                   param_to_state, self.rank,
+                                                   self.compute_stream)
+                _all2all_gather(target_params, param_to_state, self.rank,
+                                self.comm_stream, group["none_grad"],
+                                alloc_event)
 
         def enqueue_computes(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -597,10 +759,14 @@ class Muon(torch.optim.Optimizer):
                 _compute_u(p, state, group["ns_steps"], self.rank,
                            self.compute_stream)
 
-        def enqueue_scatters(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _scatter(p, state, self.rank, self.comm_stream)
+        def enqueue_all2all_scatter(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_scattered_u(target_params, param_to_state,
+                                                 self.rank,
+                                                 self.compute_stream)
+                _all2all_scatter(target_params, param_to_state, self.rank,
+                                 self.comm_stream, alloc_event)
 
         def enqueue_update_param(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -615,14 +781,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        enqueue_gathers(0, chunk_size)
+        PRE_STEP = 5
+        for i in range(0, PRE_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_computes(i, chunk_size)
-            if i > 0:
-                enqueue_update_param(i - chunk_size, chunk_size)
-            enqueue_gathers(i + chunk_size, chunk_size)
-            enqueue_scatters(i, chunk_size)
-        enqueue_update_param(i, chunk_size)
+            enqueue_all2all_scatter(i, chunk_size)
+            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_update_param(i, chunk_size)
+            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)

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

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

build/{torch27-cxx11-cu128-x86_64-linux/optimizer/_optimizer_9c21645_dirty.abi3.so → torch28-cxx11-cu129-x86_64-linux/optimizer/_optimizer_fceb334_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dae71b7e998e72130093a86f8c983c3379510e23525e3cdcd4afe5c21bf4d3db
+oid sha256:afde41795cbfff495808e4a5da7c9309795bf9b1e49e08f3cb459b5115c9bf98
 size 1883344

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

@@ -0,0 +1,106 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def get_autotune_config():
+    return [
+        triton.Config(
+            {
+                'BLOCK_SIZE_M': blk_m,
+                'BLOCK_SIZE_K': blk_k,
+                'GROUP_SIZE_M': grp_sz
+            },
+            num_stages=n_stages,
+            num_warps=n_warps) for blk_m in [32, 64, 128]
+        for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5]
+        for n_warps in [4, 8]
+    ]
+
+
+@triton.autotune(
+    configs=get_autotune_config(),
+    key=['M', 'K'],
+)
+@triton.jit
+def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn,
+               BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+               GROUP_SIZE_M: tl.constexpr):
+    """
+    Core kernel jit function of matmul_transpose that computes y = x @ x.T
+    The code is a simple adaptation from the triton `matmul` tutorial:
+    https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+    """
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    if pid_m > pid_n:
+        return
+
+    offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # we use a & b ptrs to denote different rows of x.
+    a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+    b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        b = tl.load(b_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator)
+        a_ptrs += BLOCK_SIZE_K * stride_xk
+        b_ptrs += BLOCK_SIZE_K * stride_xk
+    # use dtype.element_ty to accommodate different input datatypes as in cpp templates
+    # https://github.com/triton-lang/triton/issues/2252
+    c = accumulator.to(x.dtype.element_ty)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+    # transpose and copy
+    if pid_m < pid_n:
+        ct_ptrs = y + stride_ym * offs_cn[:,
+                                          None] + stride_yn * offs_cm[None, :]
+        ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+        tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask)
+
+
+def matmul_transpose_assign(d_in, d_out):
+    assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor"
+    assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor"
+    assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device"
+    assert d_in.dtype == d_out.dtype, "Inputs must have the same data type"
+    assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor"
+    assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor"
+    assert d_in.size(0) == d_out.size(0) == d_out.size(0), \
+            "First dimension of `d_in` must match first and second dimension of `d_out`"
+
+    d_in = d_in.contiguous()
+    M, K = d_in.shape
+    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(
+        M, META['BLOCK_SIZE_M']), )
+    with torch.cuda.device(d_in.device.index):
+        mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1),
+                         d_out.stride(0), d_out.stride(1))
+
+
+def matmul_transpose(d_in):
+    M, _ = d_in.shape
+    d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype)
+    matmul_transpose_assign(d_in, d_out)
+    return d_out

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

@@ -8,14 +8,19 @@ import torch
 import torch.distributed as dist
 from torch.distributed._tensor import DTensor, Replicate, Shard
 
+from .matmul_transpose_triton import matmul_transpose_assign
+
 logger = logging.getLogger(__name__)
 
+COMM_DTYPE = torch.bfloat16
+
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
 # Muon's Newton–Schulz iteration causes high variance in singular values
 # Idea: give each iteration its own 3 coefficients and optimize them via gradient descent.
 @torch.no_grad()
+# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon
 def _zeropower_via_newtonschulz5(G, steps):
     """
     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
@@ -27,13 +32,15 @@ def _zeropower_via_newtonschulz5(G, steps):
     performance at all relative to UV^T, where USV^T = G is the SVD.
     """
     assert len(G.shape) == 2
-    assert G.dtype == torch.bfloat16
+    assert G.dtype == COMM_DTYPE
     X = G  # no manual typecast
 
     if G.size(0) > G.size(1):
         X = X.T
     # Ensure spectral norm is at most 1
     X = X / (X.norm() + 1e-7)
+    buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
+    buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
     for a, b, c in [
         (4.0848, -6.8946, 2.9270),
@@ -42,13 +49,10 @@ def _zeropower_via_newtonschulz5(G, steps):
         (2.8769, -3.1427, 1.2046),
         (2.8366, -3.0525, 1.2012),
     ]:
-        A = X @ X.T
-        # B = (
-        #    b * A + c * A @ A
-        # )
-        B = torch.addmm(A, A, A, alpha=c, beta=b)
-        # X = a * X + B @ X
-        X = torch.addmm(X, B, X, alpha=1.0, beta=a)
+        matmul_transpose_assign(X, buf1)
+        matmul_transpose_assign(buf1, buf2)
+        buf1.mul_(b).add_(buf2, alpha=c)
+        X = torch.addmm(X, buf1, X, alpha=1.0, beta=a)
 
     if G.size(0) > G.size(1):
         X = X.T
@@ -69,51 +73,131 @@ class _muon_state:
     qk_clip_state = None
 
 
+def split_elems_for_src(param, state, 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
+
+
 @torch.no_grad()
-def _gather(p, state, rank, comm_stream, none_grad):
-    """
-    Gather the gradients to worker_rank.
-    If none_grad is True, free p.grad after the gather.
-    """
+def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(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(),
+                                                  dtype=COMM_DTYPE,
+                                                  device="cuda")
+            else:
+                state.gathered_grad = None
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
+
+
+@torch.no_grad()
+def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
+                    alloc_event):
     with torch.cuda.stream(comm_stream):
-        g = p.grad
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            gather_list = [
-                torch.empty_like(g.to_local(), dtype=torch.bfloat16)
-                for _ in range(num_ranks)
-            ]
-        else:
-            gather_list = None
-
-        g = g.to(torch.bfloat16)
-        torch.distributed.gather(
-            g.to_local(),
-            dst=state.worker_rank,
-            gather_list=gather_list,
-            group=state.process_group,
+        # Calculate sending tensors
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
+        for p in params:
+            state = param_to_state[id(p)]
+            dst = state.worker_rank
+            assert dst < num_ranks
+            shard_elems = split_elems_for_src(p, state, rank, num_ranks)
+            g = p.grad
+            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            assert g.numel() == shard_elems
+            per_dst[dst].append(g)
+            send_counts[dst] += shard_elems
+
+        assert all(
+            len(v) > 0
+            for v in per_dst), "all params should be sharded to all devices"
+
+        send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
+
+        # Calculate receiving tensors
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                total += split_elems_for_src(p, state, src, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        if rank == state.worker_rank:
-            if state.gathered_grad is not None:
-                raise RuntimeError(
-                    "Gather event already exists, which should not happen.")
-            state.gathered_grad = torch.cat(gather_list, dim=0)
-            state.gather_event = torch.cuda.Event()
-            state.gather_event.record()
-        else:
-            state.gathered_grad = None
-            state.gather_event = None
-        gather_list = None
-        if none_grad:
-            # We can safely free p.grad without calling record_stream:
-            #   p.grad.to_local().record_stream(comm_stream)
-            # Explanation:
-            # 1. p.grad is created on the default stream, but the default stream
-            #    is synchronized with the comm stream later.
-            # 2. There is no further activity on the default stream before the optimizer finishes.
-            # Therefore, it is safe to free p.grad directly on the comm stream.
-            p.grad = None
+
+        # Reconstructs gathered grad from the received buffer
+        #
+        #                  recv_buf (num ranks = 3)
+        #
+        #      From rank 0        From rank 1        From rank 2
+        # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # p1_n -> p2_n -> p3_n
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        write_offsets = {id(p): 0 for p in owner_params}
+        for src in range(num_ranks):
+            if recv_counts[src] == 0:
+                continue
+
+            block = recv_counts[src]
+            inner_off = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                n = split_elems_for_src(p, state, src, num_ranks)
+                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)
+                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:
+                state.gathered_grad = None
+                state.gather_event = None
+            if none_grad:
+                p.grad = None
 
 
 @torch.no_grad()
@@ -127,45 +211,119 @@ def _compute_u(p, state, steps, rank, compute_stream):
                 raise RuntimeError("Gather event must be set before compute.")
             compute_stream.wait_event(state.gather_event)
             u = _zeropower_via_newtonschulz5(state.gathered_grad, steps)
+            state.gathered_grad = None
             state.computed_u = u
-        state.scattered_u = torch.empty_like(p.to_local(),
-                                             dtype=torch.bfloat16)
-        state.compute_event = torch.cuda.Event()
-        state.compute_event.record()
-        u = None
+            state.compute_event = torch.cuda.Event()
+            state.compute_event.record()
+        else:
+            state.computed_u = None
+            state.compute_event = None
 
 
 @torch.no_grad()
-def _scatter(p, state, rank, comm_stream):
-    """
-    Scatter the computed_u from worker_rank to all ranks.
-    """
+def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(compute_stream):
+        for p in params:
+            state = param_to_state[id(p)]
+            state.scattered_u = torch.empty_like(p.to_local(),
+                                                 dtype=COMM_DTYPE)
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
 
+
+def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
     with torch.cuda.stream(comm_stream):
-        if state.compute_event is None:
-            raise RuntimeError("Compute event must be set before scatter.")
-        comm_stream.wait_event(state.compute_event)
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            # Clear the gathered gradient to free memory
-            state.gathered_grad = None
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
 
-            u = state.computed_u
-            scatter_list = list(torch.split(u, p.size(0) // num_ranks, dim=0))
-            scatter_list = [s.contiguous() for s in scatter_list]
-        else:
-            scatter_list = None
+        if owner_params:
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                if state.compute_event is None:
+                    raise RuntimeError(
+                        "Compute event must be set before scatter.")
+                comm_stream.wait_event(state.compute_event)
+                state.gathered_grad = None
 
-        torch.distributed.scatter(
-            state.scattered_u,
-            scatter_list=scatter_list,
-            src=state.worker_rank,
-            group=state.process_group,
+                assert state.computed_u is not None
+
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+
+                offset = 0
+                for dst in range(num_ranks):
+                    n = split_elems_for_src(p, state, dst, num_ranks)
+                    assert n > 0
+
+                    su = u_full.narrow(0, offset, n)
+                    per_dst[dst].append(su)
+                    send_counts[dst] += n
+                    offset += n
+
+                assert offset == u_full.numel()
+
+        if any(len(v) > 0 for v in per_dst):
+            send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        else:
+            # all_to_all requires participation from all ranks
+            # Even non-owner ranks must join the collective call
+            send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
+
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                total += split_elems_for_src(p, state, rank, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        assert recv_total > 0
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        state.scatter_event = torch.cuda.Event()
-        state.scatter_event.record()
-        scatter_list = None
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        for src in range(num_ranks):
+            block = recv_counts[src]
+            if block == 0:
+                continue
+
+            inner_off = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                n = split_elems_for_src(p, state, rank, num_ranks)
+                assert n > 0
+
+                flat_local = recv_buf.narrow(0, off + inner_off,
+                                             n).view_as(p.to_local())
+                state.scattered_u.copy_(flat_local)
+
+                state.scatter_event = torch.cuda.Event()
+                state.scatter_event.record(comm_stream)
+                inner_off += n
+
+            assert inner_off == block
+            off += block
 
 
 def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
@@ -444,11 +602,14 @@ class Muon(torch.optim.Optimizer):
             if mesh is None:
                 mesh = p.device_mesh
                 shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
+                self.rank = dist.get_rank(group=process_group)
             elif 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()
+            param_to_state[id(
+                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
             param_to_state[id(p)].process_group = process_group
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
             param_to_state[id(p)].qk_clip_state = qk_clip_state
@@ -478,7 +639,7 @@ class Muon(torch.optim.Optimizer):
             else:
                 g = buf
 
-            u = _zeropower_via_newtonschulz5(g.bfloat16(),
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
                                              steps=group["ns_steps"])
 
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
@@ -493,15 +654,12 @@ class Muon(torch.optim.Optimizer):
     def _update_g(self, p, g, group, momentum):
         # 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)
+        buf = state.setdefault("momentum_buffer", torch.zeros_like(g))
+        torch.add(g, buf, alpha=momentum, out=buf)
         if group["nesterov"]:
-            g = g.add(buf, alpha=momentum)
-        else:
-            g = buf
-        return g
+            g.add_(buf, alpha=momentum)
+            return g
+        return buf
 
     @staticmethod
     def _update_p(p, u, lr, adjusted_lr, weight_decay):
@@ -585,11 +743,15 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
-        def enqueue_gathers(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _gather(p, state, self.rank, self.comm_stream,
-                        group["none_grad"])
+        def enqueue_all2all_gather(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_gathered_grad(target_params,
+                                                   param_to_state, self.rank,
+                                                   self.compute_stream)
+                _all2all_gather(target_params, param_to_state, self.rank,
+                                self.comm_stream, group["none_grad"],
+                                alloc_event)
 
         def enqueue_computes(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -597,10 +759,14 @@ class Muon(torch.optim.Optimizer):
                 _compute_u(p, state, group["ns_steps"], self.rank,
                            self.compute_stream)
 
-        def enqueue_scatters(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _scatter(p, state, self.rank, self.comm_stream)
+        def enqueue_all2all_scatter(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_scattered_u(target_params, param_to_state,
+                                                 self.rank,
+                                                 self.compute_stream)
+                _all2all_scatter(target_params, param_to_state, self.rank,
+                                 self.comm_stream, alloc_event)
 
         def enqueue_update_param(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -615,14 +781,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        enqueue_gathers(0, chunk_size)
+        PRE_STEP = 5
+        for i in range(0, PRE_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_computes(i, chunk_size)
-            if i > 0:
-                enqueue_update_param(i - chunk_size, chunk_size)
-            enqueue_gathers(i + chunk_size, chunk_size)
-            enqueue_scatters(i, chunk_size)
-        enqueue_update_param(i, chunk_size)
+            enqueue_all2all_scatter(i, chunk_size)
+            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_update_param(i, chunk_size)
+            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)

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

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

build/torch28-cxx11-rocm63-x86_64-linux/optimizer/{_optimizer_9c21645_dirty.abi3.so → _optimizer_fceb334_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d8f845b8df6426eb5db57e4525b8dd3c80004c44759b01a3e39cc37a817813b5
+oid sha256:d45bbfe1dbab50f52d15893e65ec125af6101cf36f41576f9647326fdedc328b
 size 1749936

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

@@ -0,0 +1,106 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def get_autotune_config():
+    return [
+        triton.Config(
+            {
+                'BLOCK_SIZE_M': blk_m,
+                'BLOCK_SIZE_K': blk_k,
+                'GROUP_SIZE_M': grp_sz
+            },
+            num_stages=n_stages,
+            num_warps=n_warps) for blk_m in [32, 64, 128]
+        for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5]
+        for n_warps in [4, 8]
+    ]
+
+
+@triton.autotune(
+    configs=get_autotune_config(),
+    key=['M', 'K'],
+)
+@triton.jit
+def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn,
+               BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+               GROUP_SIZE_M: tl.constexpr):
+    """
+    Core kernel jit function of matmul_transpose that computes y = x @ x.T
+    The code is a simple adaptation from the triton `matmul` tutorial:
+    https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+    """
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    if pid_m > pid_n:
+        return
+
+    offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # we use a & b ptrs to denote different rows of x.
+    a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+    b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        b = tl.load(b_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator)
+        a_ptrs += BLOCK_SIZE_K * stride_xk
+        b_ptrs += BLOCK_SIZE_K * stride_xk
+    # use dtype.element_ty to accommodate different input datatypes as in cpp templates
+    # https://github.com/triton-lang/triton/issues/2252
+    c = accumulator.to(x.dtype.element_ty)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+    # transpose and copy
+    if pid_m < pid_n:
+        ct_ptrs = y + stride_ym * offs_cn[:,
+                                          None] + stride_yn * offs_cm[None, :]
+        ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+        tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask)
+
+
+def matmul_transpose_assign(d_in, d_out):
+    assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor"
+    assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor"
+    assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device"
+    assert d_in.dtype == d_out.dtype, "Inputs must have the same data type"
+    assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor"
+    assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor"
+    assert d_in.size(0) == d_out.size(0) == d_out.size(0), \
+            "First dimension of `d_in` must match first and second dimension of `d_out`"
+
+    d_in = d_in.contiguous()
+    M, K = d_in.shape
+    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(
+        M, META['BLOCK_SIZE_M']), )
+    with torch.cuda.device(d_in.device.index):
+        mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1),
+                         d_out.stride(0), d_out.stride(1))
+
+
+def matmul_transpose(d_in):
+    M, _ = d_in.shape
+    d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype)
+    matmul_transpose_assign(d_in, d_out)
+    return d_out

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

@@ -8,14 +8,19 @@ import torch
 import torch.distributed as dist
 from torch.distributed._tensor import DTensor, Replicate, Shard
 
+from .matmul_transpose_triton import matmul_transpose_assign
+
 logger = logging.getLogger(__name__)
 
+COMM_DTYPE = torch.bfloat16
+
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
 # Muon's Newton–Schulz iteration causes high variance in singular values
 # Idea: give each iteration its own 3 coefficients and optimize them via gradient descent.
 @torch.no_grad()
+# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon
 def _zeropower_via_newtonschulz5(G, steps):
     """
     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
@@ -27,13 +32,15 @@ def _zeropower_via_newtonschulz5(G, steps):
     performance at all relative to UV^T, where USV^T = G is the SVD.
     """
     assert len(G.shape) == 2
-    assert G.dtype == torch.bfloat16
+    assert G.dtype == COMM_DTYPE
     X = G  # no manual typecast
 
     if G.size(0) > G.size(1):
         X = X.T
     # Ensure spectral norm is at most 1
     X = X / (X.norm() + 1e-7)
+    buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
+    buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
     for a, b, c in [
         (4.0848, -6.8946, 2.9270),
@@ -42,13 +49,10 @@ def _zeropower_via_newtonschulz5(G, steps):
         (2.8769, -3.1427, 1.2046),
         (2.8366, -3.0525, 1.2012),
     ]:
-        A = X @ X.T
-        # B = (
-        #    b * A + c * A @ A
-        # )
-        B = torch.addmm(A, A, A, alpha=c, beta=b)
-        # X = a * X + B @ X
-        X = torch.addmm(X, B, X, alpha=1.0, beta=a)
+        matmul_transpose_assign(X, buf1)
+        matmul_transpose_assign(buf1, buf2)
+        buf1.mul_(b).add_(buf2, alpha=c)
+        X = torch.addmm(X, buf1, X, alpha=1.0, beta=a)
 
     if G.size(0) > G.size(1):
         X = X.T
@@ -69,51 +73,131 @@ class _muon_state:
     qk_clip_state = None
 
 
+def split_elems_for_src(param, state, 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
+
+
 @torch.no_grad()
-def _gather(p, state, rank, comm_stream, none_grad):
-    """
-    Gather the gradients to worker_rank.
-    If none_grad is True, free p.grad after the gather.
-    """
+def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(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(),
+                                                  dtype=COMM_DTYPE,
+                                                  device="cuda")
+            else:
+                state.gathered_grad = None
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
+
+
+@torch.no_grad()
+def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
+                    alloc_event):
     with torch.cuda.stream(comm_stream):
-        g = p.grad
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            gather_list = [
-                torch.empty_like(g.to_local(), dtype=torch.bfloat16)
-                for _ in range(num_ranks)
-            ]
-        else:
-            gather_list = None
-
-        g = g.to(torch.bfloat16)
-        torch.distributed.gather(
-            g.to_local(),
-            dst=state.worker_rank,
-            gather_list=gather_list,
-            group=state.process_group,
+        # Calculate sending tensors
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
+        for p in params:
+            state = param_to_state[id(p)]
+            dst = state.worker_rank
+            assert dst < num_ranks
+            shard_elems = split_elems_for_src(p, state, rank, num_ranks)
+            g = p.grad
+            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            assert g.numel() == shard_elems
+            per_dst[dst].append(g)
+            send_counts[dst] += shard_elems
+
+        assert all(
+            len(v) > 0
+            for v in per_dst), "all params should be sharded to all devices"
+
+        send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
+
+        # Calculate receiving tensors
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                total += split_elems_for_src(p, state, src, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        if rank == state.worker_rank:
-            if state.gathered_grad is not None:
-                raise RuntimeError(
-                    "Gather event already exists, which should not happen.")
-            state.gathered_grad = torch.cat(gather_list, dim=0)
-            state.gather_event = torch.cuda.Event()
-            state.gather_event.record()
-        else:
-            state.gathered_grad = None
-            state.gather_event = None
-        gather_list = None
-        if none_grad:
-            # We can safely free p.grad without calling record_stream:
-            #   p.grad.to_local().record_stream(comm_stream)
-            # Explanation:
-            # 1. p.grad is created on the default stream, but the default stream
-            #    is synchronized with the comm stream later.
-            # 2. There is no further activity on the default stream before the optimizer finishes.
-            # Therefore, it is safe to free p.grad directly on the comm stream.
-            p.grad = None
+
+        # Reconstructs gathered grad from the received buffer
+        #
+        #                  recv_buf (num ranks = 3)
+        #
+        #      From rank 0        From rank 1        From rank 2
+        # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # p1_n -> p2_n -> p3_n
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        write_offsets = {id(p): 0 for p in owner_params}
+        for src in range(num_ranks):
+            if recv_counts[src] == 0:
+                continue
+
+            block = recv_counts[src]
+            inner_off = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                n = split_elems_for_src(p, state, src, num_ranks)
+                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)
+                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:
+                state.gathered_grad = None
+                state.gather_event = None
+            if none_grad:
+                p.grad = None
 
 
 @torch.no_grad()
@@ -127,45 +211,119 @@ def _compute_u(p, state, steps, rank, compute_stream):
                 raise RuntimeError("Gather event must be set before compute.")
             compute_stream.wait_event(state.gather_event)
             u = _zeropower_via_newtonschulz5(state.gathered_grad, steps)
+            state.gathered_grad = None
             state.computed_u = u
-        state.scattered_u = torch.empty_like(p.to_local(),
-                                             dtype=torch.bfloat16)
-        state.compute_event = torch.cuda.Event()
-        state.compute_event.record()
-        u = None
+            state.compute_event = torch.cuda.Event()
+            state.compute_event.record()
+        else:
+            state.computed_u = None
+            state.compute_event = None
 
 
 @torch.no_grad()
-def _scatter(p, state, rank, comm_stream):
-    """
-    Scatter the computed_u from worker_rank to all ranks.
-    """
+def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(compute_stream):
+        for p in params:
+            state = param_to_state[id(p)]
+            state.scattered_u = torch.empty_like(p.to_local(),
+                                                 dtype=COMM_DTYPE)
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
 
+
+def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
     with torch.cuda.stream(comm_stream):
-        if state.compute_event is None:
-            raise RuntimeError("Compute event must be set before scatter.")
-        comm_stream.wait_event(state.compute_event)
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            # Clear the gathered gradient to free memory
-            state.gathered_grad = None
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
 
-            u = state.computed_u
-            scatter_list = list(torch.split(u, p.size(0) // num_ranks, dim=0))
-            scatter_list = [s.contiguous() for s in scatter_list]
-        else:
-            scatter_list = None
+        if owner_params:
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                if state.compute_event is None:
+                    raise RuntimeError(
+                        "Compute event must be set before scatter.")
+                comm_stream.wait_event(state.compute_event)
+                state.gathered_grad = None
 
-        torch.distributed.scatter(
-            state.scattered_u,
-            scatter_list=scatter_list,
-            src=state.worker_rank,
-            group=state.process_group,
+                assert state.computed_u is not None
+
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+
+                offset = 0
+                for dst in range(num_ranks):
+                    n = split_elems_for_src(p, state, dst, num_ranks)
+                    assert n > 0
+
+                    su = u_full.narrow(0, offset, n)
+                    per_dst[dst].append(su)
+                    send_counts[dst] += n
+                    offset += n
+
+                assert offset == u_full.numel()
+
+        if any(len(v) > 0 for v in per_dst):
+            send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        else:
+            # all_to_all requires participation from all ranks
+            # Even non-owner ranks must join the collective call
+            send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
+
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                total += split_elems_for_src(p, state, rank, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        assert recv_total > 0
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        state.scatter_event = torch.cuda.Event()
-        state.scatter_event.record()
-        scatter_list = None
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        for src in range(num_ranks):
+            block = recv_counts[src]
+            if block == 0:
+                continue
+
+            inner_off = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                n = split_elems_for_src(p, state, rank, num_ranks)
+                assert n > 0
+
+                flat_local = recv_buf.narrow(0, off + inner_off,
+                                             n).view_as(p.to_local())
+                state.scattered_u.copy_(flat_local)
+
+                state.scatter_event = torch.cuda.Event()
+                state.scatter_event.record(comm_stream)
+                inner_off += n
+
+            assert inner_off == block
+            off += block
 
 
 def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
@@ -444,11 +602,14 @@ class Muon(torch.optim.Optimizer):
             if mesh is None:
                 mesh = p.device_mesh
                 shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
+                self.rank = dist.get_rank(group=process_group)
             elif 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()
+            param_to_state[id(
+                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
             param_to_state[id(p)].process_group = process_group
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
             param_to_state[id(p)].qk_clip_state = qk_clip_state
@@ -478,7 +639,7 @@ class Muon(torch.optim.Optimizer):
             else:
                 g = buf
 
-            u = _zeropower_via_newtonschulz5(g.bfloat16(),
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
                                              steps=group["ns_steps"])
 
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
@@ -493,15 +654,12 @@ class Muon(torch.optim.Optimizer):
     def _update_g(self, p, g, group, momentum):
         # 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)
+        buf = state.setdefault("momentum_buffer", torch.zeros_like(g))
+        torch.add(g, buf, alpha=momentum, out=buf)
         if group["nesterov"]:
-            g = g.add(buf, alpha=momentum)
-        else:
-            g = buf
-        return g
+            g.add_(buf, alpha=momentum)
+            return g
+        return buf
 
     @staticmethod
     def _update_p(p, u, lr, adjusted_lr, weight_decay):
@@ -585,11 +743,15 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
-        def enqueue_gathers(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _gather(p, state, self.rank, self.comm_stream,
-                        group["none_grad"])
+        def enqueue_all2all_gather(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_gathered_grad(target_params,
+                                                   param_to_state, self.rank,
+                                                   self.compute_stream)
+                _all2all_gather(target_params, param_to_state, self.rank,
+                                self.comm_stream, group["none_grad"],
+                                alloc_event)
 
         def enqueue_computes(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -597,10 +759,14 @@ class Muon(torch.optim.Optimizer):
                 _compute_u(p, state, group["ns_steps"], self.rank,
                            self.compute_stream)
 
-        def enqueue_scatters(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _scatter(p, state, self.rank, self.comm_stream)
+        def enqueue_all2all_scatter(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_scattered_u(target_params, param_to_state,
+                                                 self.rank,
+                                                 self.compute_stream)
+                _all2all_scatter(target_params, param_to_state, self.rank,
+                                 self.comm_stream, alloc_event)
 
         def enqueue_update_param(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -615,14 +781,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        enqueue_gathers(0, chunk_size)
+        PRE_STEP = 5
+        for i in range(0, PRE_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_computes(i, chunk_size)
-            if i > 0:
-                enqueue_update_param(i - chunk_size, chunk_size)
-            enqueue_gathers(i + chunk_size, chunk_size)
-            enqueue_scatters(i, chunk_size)
-        enqueue_update_param(i, chunk_size)
+            enqueue_all2all_scatter(i, chunk_size)
+            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_update_param(i, chunk_size)
+            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)

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

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

build/torch28-cxx11-rocm64-x86_64-linux/optimizer/{_optimizer_9c21645_dirty.abi3.so → _optimizer_fceb334_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9a477575e3cc30e54d355b3e778240dc25fb0dab30362f3540dc5f925ac03ba1
+oid sha256:7344e0b7800f75f7c80488e7096e23a4ecf22609fee805613ac77685e7f0893c
 size 1750024

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

@@ -0,0 +1,106 @@
+import torch
+import triton
+import triton.language as tl
+
+
+def get_autotune_config():
+    return [
+        triton.Config(
+            {
+                'BLOCK_SIZE_M': blk_m,
+                'BLOCK_SIZE_K': blk_k,
+                'GROUP_SIZE_M': grp_sz
+            },
+            num_stages=n_stages,
+            num_warps=n_warps) for blk_m in [32, 64, 128]
+        for blk_k in [32, 64] for grp_sz in [8] for n_stages in [3, 4, 5]
+        for n_warps in [4, 8]
+    ]
+
+
+@triton.autotune(
+    configs=get_autotune_config(),
+    key=['M', 'K'],
+)
+@triton.jit
+def mmt_kernel(x, y, M, K, stride_xm, stride_xk, stride_ym, stride_yn,
+               BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+               GROUP_SIZE_M: tl.constexpr):
+    """
+    Core kernel jit function of matmul_transpose that computes y = x @ x.T
+    The code is a simple adaptation from the triton `matmul` tutorial:
+    https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+    """
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    if pid_m > pid_n:
+        return
+
+    offs_xm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_xn = (pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # we use a & b ptrs to denote different rows of x.
+    a_ptrs = x + (offs_xm[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+    b_ptrs = x + (offs_xn[:, None] * stride_xm + offs_k[None, :] * stride_xk)
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_M), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = tl.load(a_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        b = tl.load(b_ptrs,
+                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
+                    other=0.0)
+        accumulator = tl.dot(a, tl.permute(b, (1, 0)), accumulator)
+        a_ptrs += BLOCK_SIZE_K * stride_xk
+        b_ptrs += BLOCK_SIZE_K * stride_xk
+    # use dtype.element_ty to accommodate different input datatypes as in cpp templates
+    # https://github.com/triton-lang/triton/issues/2252
+    c = accumulator.to(x.dtype.element_ty)
+
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    c_ptrs = y + stride_ym * offs_cm[:, None] + stride_yn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < M)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+    # transpose and copy
+    if pid_m < pid_n:
+        ct_ptrs = y + stride_ym * offs_cn[:,
+                                          None] + stride_yn * offs_cm[None, :]
+        ct_mask = (offs_cn[:, None] < M) & (offs_cm[None, :] < M)
+        tl.store(ct_ptrs, tl.permute(c, (1, 0)), mask=ct_mask)
+
+
+def matmul_transpose_assign(d_in, d_out):
+    assert d_in.is_cuda, "Input `d_in` must be a CUDA tensor"
+    assert d_out.is_cuda, "Input `d_out` must be a CUDA tensor"
+    assert d_in.device == d_out.device, "Inputs `d_in` and `d_out` must be on the same CUDA device"
+    assert d_in.dtype == d_out.dtype, "Inputs must have the same data type"
+    assert d_in.ndim == 2, "Input `d_in` must be a 2D tensor"
+    assert d_out.ndim == 2, "Input `d_out` must be a 2D tensor"
+    assert d_in.size(0) == d_out.size(0) == d_out.size(0), \
+            "First dimension of `d_in` must match first and second dimension of `d_out`"
+
+    d_in = d_in.contiguous()
+    M, K = d_in.shape
+    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(
+        M, META['BLOCK_SIZE_M']), )
+    with torch.cuda.device(d_in.device.index):
+        mmt_kernel[grid](d_in, d_out, M, K, d_in.stride(0), d_in.stride(1),
+                         d_out.stride(0), d_out.stride(1))
+
+
+def matmul_transpose(d_in):
+    M, _ = d_in.shape
+    d_out = torch.empty((M, M), device=d_in.device, dtype=d_in.dtype)
+    matmul_transpose_assign(d_in, d_out)
+    return d_out

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

@@ -8,14 +8,19 @@ import torch
 import torch.distributed as dist
 from torch.distributed._tensor import DTensor, Replicate, Shard
 
+from .matmul_transpose_triton import matmul_transpose_assign
+
 logger = logging.getLogger(__name__)
 
+COMM_DTYPE = torch.bfloat16
+
 
 # This code snippet is a modified version adapted from the following GitHub repositories:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
 # Muon's Newton–Schulz iteration causes high variance in singular values
 # Idea: give each iteration its own 3 coefficients and optimize them via gradient descent.
 @torch.no_grad()
+# matmul_transpose_assign from : https://github.com/nil0x9/flash-muon
 def _zeropower_via_newtonschulz5(G, steps):
     """
     Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
@@ -27,13 +32,15 @@ def _zeropower_via_newtonschulz5(G, steps):
     performance at all relative to UV^T, where USV^T = G is the SVD.
     """
     assert len(G.shape) == 2
-    assert G.dtype == torch.bfloat16
+    assert G.dtype == COMM_DTYPE
     X = G  # no manual typecast
 
     if G.size(0) > G.size(1):
         X = X.T
     # Ensure spectral norm is at most 1
     X = X / (X.norm() + 1e-7)
+    buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
+    buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
     for a, b, c in [
         (4.0848, -6.8946, 2.9270),
@@ -42,13 +49,10 @@ def _zeropower_via_newtonschulz5(G, steps):
         (2.8769, -3.1427, 1.2046),
         (2.8366, -3.0525, 1.2012),
     ]:
-        A = X @ X.T
-        # B = (
-        #    b * A + c * A @ A
-        # )
-        B = torch.addmm(A, A, A, alpha=c, beta=b)
-        # X = a * X + B @ X
-        X = torch.addmm(X, B, X, alpha=1.0, beta=a)
+        matmul_transpose_assign(X, buf1)
+        matmul_transpose_assign(buf1, buf2)
+        buf1.mul_(b).add_(buf2, alpha=c)
+        X = torch.addmm(X, buf1, X, alpha=1.0, beta=a)
 
     if G.size(0) > G.size(1):
         X = X.T
@@ -69,51 +73,131 @@ class _muon_state:
     qk_clip_state = None
 
 
+def split_elems_for_src(param, state, 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
+
+
 @torch.no_grad()
-def _gather(p, state, rank, comm_stream, none_grad):
-    """
-    Gather the gradients to worker_rank.
-    If none_grad is True, free p.grad after the gather.
-    """
+def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(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(),
+                                                  dtype=COMM_DTYPE,
+                                                  device="cuda")
+            else:
+                state.gathered_grad = None
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
+
+
+@torch.no_grad()
+def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
+                    alloc_event):
     with torch.cuda.stream(comm_stream):
-        g = p.grad
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            gather_list = [
-                torch.empty_like(g.to_local(), dtype=torch.bfloat16)
-                for _ in range(num_ranks)
-            ]
-        else:
-            gather_list = None
-
-        g = g.to(torch.bfloat16)
-        torch.distributed.gather(
-            g.to_local(),
-            dst=state.worker_rank,
-            gather_list=gather_list,
-            group=state.process_group,
+        # Calculate sending tensors
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
+        for p in params:
+            state = param_to_state[id(p)]
+            dst = state.worker_rank
+            assert dst < num_ranks
+            shard_elems = split_elems_for_src(p, state, rank, num_ranks)
+            g = p.grad
+            g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
+            assert g.numel() == shard_elems
+            per_dst[dst].append(g)
+            send_counts[dst] += shard_elems
+
+        assert all(
+            len(v) > 0
+            for v in per_dst), "all params should be sharded to all devices"
+
+        send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
+
+        # Calculate receiving tensors
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                total += split_elems_for_src(p, state, src, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        if rank == state.worker_rank:
-            if state.gathered_grad is not None:
-                raise RuntimeError(
-                    "Gather event already exists, which should not happen.")
-            state.gathered_grad = torch.cat(gather_list, dim=0)
-            state.gather_event = torch.cuda.Event()
-            state.gather_event.record()
-        else:
-            state.gathered_grad = None
-            state.gather_event = None
-        gather_list = None
-        if none_grad:
-            # We can safely free p.grad without calling record_stream:
-            #   p.grad.to_local().record_stream(comm_stream)
-            # Explanation:
-            # 1. p.grad is created on the default stream, but the default stream
-            #    is synchronized with the comm stream later.
-            # 2. There is no further activity on the default stream before the optimizer finishes.
-            # Therefore, it is safe to free p.grad directly on the comm stream.
-            p.grad = None
+
+        # Reconstructs gathered grad from the received buffer
+        #
+        #                  recv_buf (num ranks = 3)
+        #
+        #      From rank 0        From rank 1        From rank 2
+        # | p1_0, p2_0, p3_0 | p1_1, p2_1, p3_1 | p1_2, p2_2, p3_2 |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # p1_n -> p2_n -> p3_n
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        write_offsets = {id(p): 0 for p in owner_params}
+        for src in range(num_ranks):
+            if recv_counts[src] == 0:
+                continue
+
+            block = recv_counts[src]
+            inner_off = 0
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                assert state.worker_rank == rank
+                n = split_elems_for_src(p, state, src, num_ranks)
+                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)
+                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:
+                state.gathered_grad = None
+                state.gather_event = None
+            if none_grad:
+                p.grad = None
 
 
 @torch.no_grad()
@@ -127,45 +211,119 @@ def _compute_u(p, state, steps, rank, compute_stream):
                 raise RuntimeError("Gather event must be set before compute.")
             compute_stream.wait_event(state.gather_event)
             u = _zeropower_via_newtonschulz5(state.gathered_grad, steps)
+            state.gathered_grad = None
             state.computed_u = u
-        state.scattered_u = torch.empty_like(p.to_local(),
-                                             dtype=torch.bfloat16)
-        state.compute_event = torch.cuda.Event()
-        state.compute_event.record()
-        u = None
+            state.compute_event = torch.cuda.Event()
+            state.compute_event.record()
+        else:
+            state.computed_u = None
+            state.compute_event = None
 
 
 @torch.no_grad()
-def _scatter(p, state, rank, comm_stream):
-    """
-    Scatter the computed_u from worker_rank to all ranks.
-    """
+def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    with torch.cuda.stream(compute_stream):
+        for p in params:
+            state = param_to_state[id(p)]
+            state.scattered_u = torch.empty_like(p.to_local(),
+                                                 dtype=COMM_DTYPE)
+
+        alloc_event = torch.cuda.Event()
+        alloc_event.record(compute_stream)
+        return alloc_event
 
+
+def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
     with torch.cuda.stream(comm_stream):
-        if state.compute_event is None:
-            raise RuntimeError("Compute event must be set before scatter.")
-        comm_stream.wait_event(state.compute_event)
+        process_group = param_to_state[id(params[0])].process_group
+        num_ranks = dist.get_world_size(group=process_group)
+        owner_params = [
+            p for p in params if param_to_state[id(p)].worker_rank == rank
+        ]
 
-        if rank == state.worker_rank:
-            num_ranks = dist.get_world_size(group=state.process_group)
-            # Clear the gathered gradient to free memory
-            state.gathered_grad = None
+        per_dst = [[] for _ in range(num_ranks)]
+        send_counts = [0] * num_ranks
 
-            u = state.computed_u
-            scatter_list = list(torch.split(u, p.size(0) // num_ranks, dim=0))
-            scatter_list = [s.contiguous() for s in scatter_list]
-        else:
-            scatter_list = None
+        if owner_params:
+            for p in owner_params:
+                state = param_to_state[id(p)]
+                if state.compute_event is None:
+                    raise RuntimeError(
+                        "Compute event must be set before scatter.")
+                comm_stream.wait_event(state.compute_event)
+                state.gathered_grad = None
 
-        torch.distributed.scatter(
-            state.scattered_u,
-            scatter_list=scatter_list,
-            src=state.worker_rank,
-            group=state.process_group,
+                assert state.computed_u is not None
+
+                u_full = state.computed_u.to(COMM_DTYPE).contiguous().view(-1)
+
+                offset = 0
+                for dst in range(num_ranks):
+                    n = split_elems_for_src(p, state, dst, num_ranks)
+                    assert n > 0
+
+                    su = u_full.narrow(0, offset, n)
+                    per_dst[dst].append(su)
+                    send_counts[dst] += n
+                    offset += n
+
+                assert offset == u_full.numel()
+
+        if any(len(v) > 0 for v in per_dst):
+            send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        else:
+            # all_to_all requires participation from all ranks
+            # Even non-owner ranks must join the collective call
+            send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
+
+        recv_counts = [0] * num_ranks
+        for src in range(num_ranks):
+            total = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                total += split_elems_for_src(p, state, rank, num_ranks)
+            recv_counts[src] = total
+
+        recv_total = sum(recv_counts)
+        assert recv_total > 0
+        recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        dist.all_to_all_single(
+            recv_buf,
+            send_buf,
+            output_split_sizes=recv_counts,
+            input_split_sizes=send_counts,
+            group=process_group,
         )
-        state.scatter_event = torch.cuda.Event()
-        state.scatter_event.record()
-        scatter_list = None
+
+        comm_stream.wait_event(alloc_event)
+
+        off = 0
+        for src in range(num_ranks):
+            block = recv_counts[src]
+            if block == 0:
+                continue
+
+            inner_off = 0
+            for p in params:
+                state = param_to_state[id(p)]
+                if state.worker_rank != src:
+                    continue
+                n = split_elems_for_src(p, state, rank, num_ranks)
+                assert n > 0
+
+                flat_local = recv_buf.narrow(0, off + inner_off,
+                                             n).view_as(p.to_local())
+                state.scattered_u.copy_(flat_local)
+
+                state.scatter_event = torch.cuda.Event()
+                state.scatter_event.record(comm_stream)
+                inner_off += n
+
+            assert inner_off == block
+            off += block
 
 
 def _update_param(p, state, lr, adjusted_lr, weight_decay, rank,
@@ -444,11 +602,14 @@ class Muon(torch.optim.Optimizer):
             if mesh is None:
                 mesh = p.device_mesh
                 shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
+                self.rank = dist.get_rank(group=process_group)
             elif 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()
+            param_to_state[id(
+                p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
             param_to_state[id(p)].process_group = process_group
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
             param_to_state[id(p)].qk_clip_state = qk_clip_state
@@ -478,7 +639,7 @@ class Muon(torch.optim.Optimizer):
             else:
                 g = buf
 
-            u = _zeropower_via_newtonschulz5(g.bfloat16(),
+            u = _zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
                                              steps=group["ns_steps"])
 
             adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
@@ -493,15 +654,12 @@ class Muon(torch.optim.Optimizer):
     def _update_g(self, p, g, group, momentum):
         # 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)
+        buf = state.setdefault("momentum_buffer", torch.zeros_like(g))
+        torch.add(g, buf, alpha=momentum, out=buf)
         if group["nesterov"]:
-            g = g.add(buf, alpha=momentum)
-        else:
-            g = buf
-        return g
+            g.add_(buf, alpha=momentum)
+            return g
+        return buf
 
     @staticmethod
     def _update_p(p, u, lr, adjusted_lr, weight_decay):
@@ -585,11 +743,15 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
-        def enqueue_gathers(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _gather(p, state, self.rank, self.comm_stream,
-                        group["none_grad"])
+        def enqueue_all2all_gather(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_gathered_grad(target_params,
+                                                   param_to_state, self.rank,
+                                                   self.compute_stream)
+                _all2all_gather(target_params, param_to_state, self.rank,
+                                self.comm_stream, group["none_grad"],
+                                alloc_event)
 
         def enqueue_computes(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -597,10 +759,14 @@ class Muon(torch.optim.Optimizer):
                 _compute_u(p, state, group["ns_steps"], self.rank,
                            self.compute_stream)
 
-        def enqueue_scatters(start_idx, chunk_size):
-            for p in ordered_params[start_idx:start_idx + chunk_size]:
-                state = param_to_state[id(p)]
-                _scatter(p, state, self.rank, self.comm_stream)
+        def enqueue_all2all_scatter(start_idx, chunk_size):
+            target_params = ordered_params[start_idx:start_idx + chunk_size]
+            if target_params:
+                alloc_event = _alloc_scattered_u(target_params, param_to_state,
+                                                 self.rank,
+                                                 self.compute_stream)
+                _all2all_scatter(target_params, param_to_state, self.rank,
+                                 self.comm_stream, alloc_event)
 
         def enqueue_update_param(start_idx, chunk_size):
             for p in ordered_params[start_idx:start_idx + chunk_size]:
@@ -615,14 +781,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        enqueue_gathers(0, chunk_size)
+        PRE_STEP = 5
+        for i in range(0, PRE_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_computes(i, chunk_size)
-            if i > 0:
-                enqueue_update_param(i - chunk_size, chunk_size)
-            enqueue_gathers(i + chunk_size, chunk_size)
-            enqueue_scatters(i, chunk_size)
-        enqueue_update_param(i, chunk_size)
+            enqueue_all2all_scatter(i, chunk_size)
+            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_update_param(i, chunk_size)
+            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)