Add built binary [ci skip]

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

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

build/torch27-cxx11-cu118-x86_64-linux/optimizer/{_optimizer_37be0d0_dirty.abi3.so → _optimizer_6943c45_dirty.abi3.so}

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

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

@@ -73,7 +73,7 @@ class _muon_state:
     qk_clip_state = None
 
 
-def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
+def split_elems_for_src(param, src_rank, num_ranks) -> int:
     rows = param.shape[0]
     cols = int(param.numel() // rows)
     base, rem = divmod(rows, num_ranks)
@@ -83,6 +83,10 @@ def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
 
 @torch.no_grad()
 def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate gathered_grad buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -102,13 +106,17 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
 @torch.no_grad()
 def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
                     alloc_event):
+    """
+    All2all gathers shards so each owner rank reconstructs its full gradient
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
 
-        # Calculate sending tensors
+        # Construct sending buffers
         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
@@ -125,15 +133,16 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             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 = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
-        # Calculate receiving tensors
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 total += split_elems_for_src(p, state, src, num_ranks)
@@ -141,6 +150,8 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -165,14 +176,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owner_params}
+        write_offsets = {id(p): 0 for p in owned_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:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 n = split_elems_for_src(p, state, src, num_ranks)
@@ -222,6 +233,10 @@ def _compute_u(p, state, steps, rank, compute_stream):
 
 @torch.no_grad()
 def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate scattered_u buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -234,18 +249,22 @@ def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
 
 
 def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
+    """
+    All2all scatters full gradients to all ranks
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
+        # Construct sending buffer
         per_dst = [[] for _ in range(num_ranks)]
         send_counts = [0] * num_ranks
 
-        if owner_params:
-            for p in owner_params:
+        if owned_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 if state.compute_event is None:
                     raise RuntimeError(
@@ -276,7 +295,9 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             # Even non-owner ranks must join the collective call
             send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
             for p in params:
@@ -290,6 +311,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
         assert recv_total > 0
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -298,6 +320,21 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             group=process_group,
         )
 
+        # Copy to pre-allocated scattered_u buffer from the received buffer
+        #
+        #                  recv_buf (num ranks = 3, local_rank = 0)
+        #
+        #      From rank 0        From rank 1       From rank 2
+        # | p1_0, p2_0, p3_0 |      p4_0       |    p5_0, p6_0    |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # src(0) :  p1_0 -> p2_0 -> p3_0
+        # src(1) :  p4_0
+        # src(2) :  p5_0 -> p6_0
+
         comm_stream.wait_event(alloc_event)
 
         off = 0
@@ -479,6 +516,11 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
+        overlap_step : How many all2all gather, compute operations are launched in advance
+                       before the corresponding all2all scatter steps begin.
+                       A higher overlap_step increases memory usage but can improve
+                       performance by overlapping communication.
+                       Parallel muon only. 
     """
 
     def __init__(self,
@@ -497,7 +539,8 @@ class Muon(torch.optim.Optimizer):
                      "k_indices": [],
                      "head_dim": 128,
                      "threshold": 100
-                 }):
+                 },
+                 overlap_step=5):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -521,15 +564,13 @@ class Muon(torch.optim.Optimizer):
 
         super().__init__(params, defaults)
 
-        if dist.is_initialized():
-            self.rank = dist.get_rank()
-        else:
-            self.rank = None
+        self.rank = None
 
         self.comm_stream = torch.cuda.Stream()
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
+        self.overlap_step = overlap_step
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -602,7 +643,11 @@ 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)
+                local_rank = dist.get_rank(group=process_group)
+                if self.rank is None:
+                    self.rank = dist.get_rank(group=process_group)
+                else:
+                    assert self.rank == local_rank
             elif mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
 
@@ -743,6 +788,8 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
+        assert self.rank is not None
+
         def enqueue_all2all_gather(start_idx, chunk_size):
             target_params = ordered_params[start_idx:start_idx + chunk_size]
             if target_params:
@@ -781,16 +828,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        PRE_STEP = 5
-        for i in range(0, PRE_STEP):
+        overlap_step = self.overlap_step
+        for i in range(0, overlap_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_computes(i + overlap_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_37be0d0_dirty
-ops = torch.ops._optimizer_37be0d0_dirty
+from . import _optimizer_6943c45_dirty
+ops = torch.ops._optimizer_6943c45_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_37be0d0_dirty::{op_name}"
+    return f"_optimizer_6943c45_dirty::{op_name}"

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

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:10fdf4bbdc0ceb7b6c1504145ea53f837ac237422c43e3b74b12aec789c49dac
+oid sha256:10905a024d5fa31fe31b8370747790b2f0ce60d06881831efb7cc07e1c5e5436
 size 1824256

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

@@ -73,7 +73,7 @@ class _muon_state:
     qk_clip_state = None
 
 
-def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
+def split_elems_for_src(param, src_rank, num_ranks) -> int:
     rows = param.shape[0]
     cols = int(param.numel() // rows)
     base, rem = divmod(rows, num_ranks)
@@ -83,6 +83,10 @@ def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
 
 @torch.no_grad()
 def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate gathered_grad buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -102,13 +106,17 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
 @torch.no_grad()
 def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
                     alloc_event):
+    """
+    All2all gathers shards so each owner rank reconstructs its full gradient
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
 
-        # Calculate sending tensors
+        # Construct sending buffers
         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
@@ -125,15 +133,16 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             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 = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
-        # Calculate receiving tensors
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 total += split_elems_for_src(p, state, src, num_ranks)
@@ -141,6 +150,8 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -165,14 +176,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owner_params}
+        write_offsets = {id(p): 0 for p in owned_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:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 n = split_elems_for_src(p, state, src, num_ranks)
@@ -222,6 +233,10 @@ def _compute_u(p, state, steps, rank, compute_stream):
 
 @torch.no_grad()
 def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate scattered_u buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -234,18 +249,22 @@ def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
 
 
 def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
+    """
+    All2all scatters full gradients to all ranks
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
+        # Construct sending buffer
         per_dst = [[] for _ in range(num_ranks)]
         send_counts = [0] * num_ranks
 
-        if owner_params:
-            for p in owner_params:
+        if owned_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 if state.compute_event is None:
                     raise RuntimeError(
@@ -276,7 +295,9 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             # Even non-owner ranks must join the collective call
             send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
             for p in params:
@@ -290,6 +311,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
         assert recv_total > 0
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -298,6 +320,21 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             group=process_group,
         )
 
+        # Copy to pre-allocated scattered_u buffer from the received buffer
+        #
+        #                  recv_buf (num ranks = 3, local_rank = 0)
+        #
+        #      From rank 0        From rank 1       From rank 2
+        # | p1_0, p2_0, p3_0 |      p4_0       |    p5_0, p6_0    |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # src(0) :  p1_0 -> p2_0 -> p3_0
+        # src(1) :  p4_0
+        # src(2) :  p5_0 -> p6_0
+
         comm_stream.wait_event(alloc_event)
 
         off = 0
@@ -479,6 +516,11 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
+        overlap_step : How many all2all gather, compute operations are launched in advance
+                       before the corresponding all2all scatter steps begin.
+                       A higher overlap_step increases memory usage but can improve
+                       performance by overlapping communication.
+                       Parallel muon only. 
     """
 
     def __init__(self,
@@ -497,7 +539,8 @@ class Muon(torch.optim.Optimizer):
                      "k_indices": [],
                      "head_dim": 128,
                      "threshold": 100
-                 }):
+                 },
+                 overlap_step=5):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -521,15 +564,13 @@ class Muon(torch.optim.Optimizer):
 
         super().__init__(params, defaults)
 
-        if dist.is_initialized():
-            self.rank = dist.get_rank()
-        else:
-            self.rank = None
+        self.rank = None
 
         self.comm_stream = torch.cuda.Stream()
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
+        self.overlap_step = overlap_step
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -602,7 +643,11 @@ 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)
+                local_rank = dist.get_rank(group=process_group)
+                if self.rank is None:
+                    self.rank = dist.get_rank(group=process_group)
+                else:
+                    assert self.rank == local_rank
             elif mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
 
@@ -743,6 +788,8 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
+        assert self.rank is not None
+
         def enqueue_all2all_gather(start_idx, chunk_size):
             target_params = ordered_params[start_idx:start_idx + chunk_size]
             if target_params:
@@ -781,16 +828,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        PRE_STEP = 5
-        for i in range(0, PRE_STEP):
+        overlap_step = self.overlap_step
+        for i in range(0, overlap_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_computes(i + overlap_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_37be0d0_dirty
-ops = torch.ops._optimizer_37be0d0_dirty
+from . import _optimizer_6943c45_dirty
+ops = torch.ops._optimizer_6943c45_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_37be0d0_dirty::{op_name}"
+    return f"_optimizer_6943c45_dirty::{op_name}"

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

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

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

@@ -73,7 +73,7 @@ class _muon_state:
     qk_clip_state = None
 
 
-def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
+def split_elems_for_src(param, src_rank, num_ranks) -> int:
     rows = param.shape[0]
     cols = int(param.numel() // rows)
     base, rem = divmod(rows, num_ranks)
@@ -83,6 +83,10 @@ def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
 
 @torch.no_grad()
 def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate gathered_grad buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -102,13 +106,17 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
 @torch.no_grad()
 def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
                     alloc_event):
+    """
+    All2all gathers shards so each owner rank reconstructs its full gradient
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
 
-        # Calculate sending tensors
+        # Construct sending buffers
         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
@@ -125,15 +133,16 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             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 = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
-        # Calculate receiving tensors
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 total += split_elems_for_src(p, state, src, num_ranks)
@@ -141,6 +150,8 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -165,14 +176,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owner_params}
+        write_offsets = {id(p): 0 for p in owned_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:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 n = split_elems_for_src(p, state, src, num_ranks)
@@ -222,6 +233,10 @@ def _compute_u(p, state, steps, rank, compute_stream):
 
 @torch.no_grad()
 def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate scattered_u buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -234,18 +249,22 @@ def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
 
 
 def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
+    """
+    All2all scatters full gradients to all ranks
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
+        # Construct sending buffer
         per_dst = [[] for _ in range(num_ranks)]
         send_counts = [0] * num_ranks
 
-        if owner_params:
-            for p in owner_params:
+        if owned_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 if state.compute_event is None:
                     raise RuntimeError(
@@ -276,7 +295,9 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             # Even non-owner ranks must join the collective call
             send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
             for p in params:
@@ -290,6 +311,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
         assert recv_total > 0
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -298,6 +320,21 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             group=process_group,
         )
 
+        # Copy to pre-allocated scattered_u buffer from the received buffer
+        #
+        #                  recv_buf (num ranks = 3, local_rank = 0)
+        #
+        #      From rank 0        From rank 1       From rank 2
+        # | p1_0, p2_0, p3_0 |      p4_0       |    p5_0, p6_0    |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # src(0) :  p1_0 -> p2_0 -> p3_0
+        # src(1) :  p4_0
+        # src(2) :  p5_0 -> p6_0
+
         comm_stream.wait_event(alloc_event)
 
         off = 0
@@ -479,6 +516,11 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
+        overlap_step : How many all2all gather, compute operations are launched in advance
+                       before the corresponding all2all scatter steps begin.
+                       A higher overlap_step increases memory usage but can improve
+                       performance by overlapping communication.
+                       Parallel muon only. 
     """
 
     def __init__(self,
@@ -497,7 +539,8 @@ class Muon(torch.optim.Optimizer):
                      "k_indices": [],
                      "head_dim": 128,
                      "threshold": 100
-                 }):
+                 },
+                 overlap_step=5):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -521,15 +564,13 @@ class Muon(torch.optim.Optimizer):
 
         super().__init__(params, defaults)
 
-        if dist.is_initialized():
-            self.rank = dist.get_rank()
-        else:
-            self.rank = None
+        self.rank = None
 
         self.comm_stream = torch.cuda.Stream()
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
+        self.overlap_step = overlap_step
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -602,7 +643,11 @@ 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)
+                local_rank = dist.get_rank(group=process_group)
+                if self.rank is None:
+                    self.rank = dist.get_rank(group=process_group)
+                else:
+                    assert self.rank == local_rank
             elif mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
 
@@ -743,6 +788,8 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
+        assert self.rank is not None
+
         def enqueue_all2all_gather(start_idx, chunk_size):
             target_params = ordered_params[start_idx:start_idx + chunk_size]
             if target_params:
@@ -781,16 +828,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        PRE_STEP = 5
-        for i in range(0, PRE_STEP):
+        overlap_step = self.overlap_step
+        for i in range(0, overlap_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_computes(i + overlap_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_37be0d0_dirty
-ops = torch.ops._optimizer_37be0d0_dirty
+from . import _optimizer_6943c45_dirty
+ops = torch.ops._optimizer_6943c45_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_37be0d0_dirty::{op_name}"
+    return f"_optimizer_6943c45_dirty::{op_name}"

build/torch27-cxx11-rocm63-x86_64-linux/optimizer/{_optimizer_37be0d0_dirty.abi3.so → _optimizer_6943c45_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8f968917ca40eb2ff9b121ea44889d0446a388432a57a01ff8d102ff59695b8a
+oid sha256:f63251067a8472a98f754fc603e51a85b44fb51777a40d229af67b173f2a7b2c
 size 1749776

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

@@ -73,7 +73,7 @@ class _muon_state:
     qk_clip_state = None
 
 
-def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
+def split_elems_for_src(param, src_rank, num_ranks) -> int:
     rows = param.shape[0]
     cols = int(param.numel() // rows)
     base, rem = divmod(rows, num_ranks)
@@ -83,6 +83,10 @@ def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
 
 @torch.no_grad()
 def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate gathered_grad buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -102,13 +106,17 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
 @torch.no_grad()
 def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
                     alloc_event):
+    """
+    All2all gathers shards so each owner rank reconstructs its full gradient
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
 
-        # Calculate sending tensors
+        # Construct sending buffers
         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
@@ -125,15 +133,16 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             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 = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
-        # Calculate receiving tensors
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 total += split_elems_for_src(p, state, src, num_ranks)
@@ -141,6 +150,8 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -165,14 +176,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owner_params}
+        write_offsets = {id(p): 0 for p in owned_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:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 n = split_elems_for_src(p, state, src, num_ranks)
@@ -222,6 +233,10 @@ def _compute_u(p, state, steps, rank, compute_stream):
 
 @torch.no_grad()
 def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate scattered_u buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -234,18 +249,22 @@ def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
 
 
 def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
+    """
+    All2all scatters full gradients to all ranks
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
+        # Construct sending buffer
         per_dst = [[] for _ in range(num_ranks)]
         send_counts = [0] * num_ranks
 
-        if owner_params:
-            for p in owner_params:
+        if owned_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 if state.compute_event is None:
                     raise RuntimeError(
@@ -276,7 +295,9 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             # Even non-owner ranks must join the collective call
             send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
             for p in params:
@@ -290,6 +311,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
         assert recv_total > 0
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -298,6 +320,21 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             group=process_group,
         )
 
+        # Copy to pre-allocated scattered_u buffer from the received buffer
+        #
+        #                  recv_buf (num ranks = 3, local_rank = 0)
+        #
+        #      From rank 0        From rank 1       From rank 2
+        # | p1_0, p2_0, p3_0 |      p4_0       |    p5_0, p6_0    |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # src(0) :  p1_0 -> p2_0 -> p3_0
+        # src(1) :  p4_0
+        # src(2) :  p5_0 -> p6_0
+
         comm_stream.wait_event(alloc_event)
 
         off = 0
@@ -479,6 +516,11 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
+        overlap_step : How many all2all gather, compute operations are launched in advance
+                       before the corresponding all2all scatter steps begin.
+                       A higher overlap_step increases memory usage but can improve
+                       performance by overlapping communication.
+                       Parallel muon only. 
     """
 
     def __init__(self,
@@ -497,7 +539,8 @@ class Muon(torch.optim.Optimizer):
                      "k_indices": [],
                      "head_dim": 128,
                      "threshold": 100
-                 }):
+                 },
+                 overlap_step=5):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -521,15 +564,13 @@ class Muon(torch.optim.Optimizer):
 
         super().__init__(params, defaults)
 
-        if dist.is_initialized():
-            self.rank = dist.get_rank()
-        else:
-            self.rank = None
+        self.rank = None
 
         self.comm_stream = torch.cuda.Stream()
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
+        self.overlap_step = overlap_step
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -602,7 +643,11 @@ 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)
+                local_rank = dist.get_rank(group=process_group)
+                if self.rank is None:
+                    self.rank = dist.get_rank(group=process_group)
+                else:
+                    assert self.rank == local_rank
             elif mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
 
@@ -743,6 +788,8 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
+        assert self.rank is not None
+
         def enqueue_all2all_gather(start_idx, chunk_size):
             target_params = ordered_params[start_idx:start_idx + chunk_size]
             if target_params:
@@ -781,16 +828,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        PRE_STEP = 5
-        for i in range(0, PRE_STEP):
+        overlap_step = self.overlap_step
+        for i in range(0, overlap_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_computes(i + overlap_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_37be0d0_dirty
-ops = torch.ops._optimizer_37be0d0_dirty
+from . import _optimizer_6943c45_dirty
+ops = torch.ops._optimizer_6943c45_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_37be0d0_dirty::{op_name}"
+    return f"_optimizer_6943c45_dirty::{op_name}"

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

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:10fdf4bbdc0ceb7b6c1504145ea53f837ac237422c43e3b74b12aec789c49dac
+oid sha256:10905a024d5fa31fe31b8370747790b2f0ce60d06881831efb7cc07e1c5e5436
 size 1824256

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

@@ -73,7 +73,7 @@ class _muon_state:
     qk_clip_state = None
 
 
-def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
+def split_elems_for_src(param, src_rank, num_ranks) -> int:
     rows = param.shape[0]
     cols = int(param.numel() // rows)
     base, rem = divmod(rows, num_ranks)
@@ -83,6 +83,10 @@ def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
 
 @torch.no_grad()
 def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate gathered_grad buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -102,13 +106,17 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
 @torch.no_grad()
 def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
                     alloc_event):
+    """
+    All2all gathers shards so each owner rank reconstructs its full gradient
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
 
-        # Calculate sending tensors
+        # Construct sending buffers
         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
@@ -125,15 +133,16 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             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 = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
-        # Calculate receiving tensors
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 total += split_elems_for_src(p, state, src, num_ranks)
@@ -141,6 +150,8 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -165,14 +176,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owner_params}
+        write_offsets = {id(p): 0 for p in owned_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:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 n = split_elems_for_src(p, state, src, num_ranks)
@@ -222,6 +233,10 @@ def _compute_u(p, state, steps, rank, compute_stream):
 
 @torch.no_grad()
 def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate scattered_u buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -234,18 +249,22 @@ def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
 
 
 def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
+    """
+    All2all scatters full gradients to all ranks
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
+        # Construct sending buffer
         per_dst = [[] for _ in range(num_ranks)]
         send_counts = [0] * num_ranks
 
-        if owner_params:
-            for p in owner_params:
+        if owned_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 if state.compute_event is None:
                     raise RuntimeError(
@@ -276,7 +295,9 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             # Even non-owner ranks must join the collective call
             send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
             for p in params:
@@ -290,6 +311,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
         assert recv_total > 0
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -298,6 +320,21 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             group=process_group,
         )
 
+        # Copy to pre-allocated scattered_u buffer from the received buffer
+        #
+        #                  recv_buf (num ranks = 3, local_rank = 0)
+        #
+        #      From rank 0        From rank 1       From rank 2
+        # | p1_0, p2_0, p3_0 |      p4_0       |    p5_0, p6_0    |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # src(0) :  p1_0 -> p2_0 -> p3_0
+        # src(1) :  p4_0
+        # src(2) :  p5_0 -> p6_0
+
         comm_stream.wait_event(alloc_event)
 
         off = 0
@@ -479,6 +516,11 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
+        overlap_step : How many all2all gather, compute operations are launched in advance
+                       before the corresponding all2all scatter steps begin.
+                       A higher overlap_step increases memory usage but can improve
+                       performance by overlapping communication.
+                       Parallel muon only. 
     """
 
     def __init__(self,
@@ -497,7 +539,8 @@ class Muon(torch.optim.Optimizer):
                      "k_indices": [],
                      "head_dim": 128,
                      "threshold": 100
-                 }):
+                 },
+                 overlap_step=5):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -521,15 +564,13 @@ class Muon(torch.optim.Optimizer):
 
         super().__init__(params, defaults)
 
-        if dist.is_initialized():
-            self.rank = dist.get_rank()
-        else:
-            self.rank = None
+        self.rank = None
 
         self.comm_stream = torch.cuda.Stream()
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
+        self.overlap_step = overlap_step
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -602,7 +643,11 @@ 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)
+                local_rank = dist.get_rank(group=process_group)
+                if self.rank is None:
+                    self.rank = dist.get_rank(group=process_group)
+                else:
+                    assert self.rank == local_rank
             elif mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
 
@@ -743,6 +788,8 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
+        assert self.rank is not None
+
         def enqueue_all2all_gather(start_idx, chunk_size):
             target_params = ordered_params[start_idx:start_idx + chunk_size]
             if target_params:
@@ -781,16 +828,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        PRE_STEP = 5
-        for i in range(0, PRE_STEP):
+        overlap_step = self.overlap_step
+        for i in range(0, overlap_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_computes(i + overlap_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_37be0d0_dirty
-ops = torch.ops._optimizer_37be0d0_dirty
+from . import _optimizer_6943c45_dirty
+ops = torch.ops._optimizer_6943c45_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_37be0d0_dirty::{op_name}"
+    return f"_optimizer_6943c45_dirty::{op_name}"

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

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4789f8aba6b8d3ff6557b8d44b657551469ddd8857b1e315eb29b09a4c86da1a
+oid sha256:058aa5ff9f1e974cb5b52e4a7af074cef2092c457b5498356fc0fbdd86adf5f3
 size 1883344

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

@@ -73,7 +73,7 @@ class _muon_state:
     qk_clip_state = None
 
 
-def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
+def split_elems_for_src(param, src_rank, num_ranks) -> int:
     rows = param.shape[0]
     cols = int(param.numel() // rows)
     base, rem = divmod(rows, num_ranks)
@@ -83,6 +83,10 @@ def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
 
 @torch.no_grad()
 def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate gathered_grad buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -102,13 +106,17 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
 @torch.no_grad()
 def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
                     alloc_event):
+    """
+    All2all gathers shards so each owner rank reconstructs its full gradient
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
 
-        # Calculate sending tensors
+        # Construct sending buffers
         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
@@ -125,15 +133,16 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             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 = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
-        # Calculate receiving tensors
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 total += split_elems_for_src(p, state, src, num_ranks)
@@ -141,6 +150,8 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -165,14 +176,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owner_params}
+        write_offsets = {id(p): 0 for p in owned_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:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 n = split_elems_for_src(p, state, src, num_ranks)
@@ -222,6 +233,10 @@ def _compute_u(p, state, steps, rank, compute_stream):
 
 @torch.no_grad()
 def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate scattered_u buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -234,18 +249,22 @@ def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
 
 
 def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
+    """
+    All2all scatters full gradients to all ranks
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
+        # Construct sending buffer
         per_dst = [[] for _ in range(num_ranks)]
         send_counts = [0] * num_ranks
 
-        if owner_params:
-            for p in owner_params:
+        if owned_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 if state.compute_event is None:
                     raise RuntimeError(
@@ -276,7 +295,9 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             # Even non-owner ranks must join the collective call
             send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
             for p in params:
@@ -290,6 +311,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
         assert recv_total > 0
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -298,6 +320,21 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             group=process_group,
         )
 
+        # Copy to pre-allocated scattered_u buffer from the received buffer
+        #
+        #                  recv_buf (num ranks = 3, local_rank = 0)
+        #
+        #      From rank 0        From rank 1       From rank 2
+        # | p1_0, p2_0, p3_0 |      p4_0       |    p5_0, p6_0    |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # src(0) :  p1_0 -> p2_0 -> p3_0
+        # src(1) :  p4_0
+        # src(2) :  p5_0 -> p6_0
+
         comm_stream.wait_event(alloc_event)
 
         off = 0
@@ -479,6 +516,11 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
+        overlap_step : How many all2all gather, compute operations are launched in advance
+                       before the corresponding all2all scatter steps begin.
+                       A higher overlap_step increases memory usage but can improve
+                       performance by overlapping communication.
+                       Parallel muon only. 
     """
 
     def __init__(self,
@@ -497,7 +539,8 @@ class Muon(torch.optim.Optimizer):
                      "k_indices": [],
                      "head_dim": 128,
                      "threshold": 100
-                 }):
+                 },
+                 overlap_step=5):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -521,15 +564,13 @@ class Muon(torch.optim.Optimizer):
 
         super().__init__(params, defaults)
 
-        if dist.is_initialized():
-            self.rank = dist.get_rank()
-        else:
-            self.rank = None
+        self.rank = None
 
         self.comm_stream = torch.cuda.Stream()
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
+        self.overlap_step = overlap_step
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -602,7 +643,11 @@ 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)
+                local_rank = dist.get_rank(group=process_group)
+                if self.rank is None:
+                    self.rank = dist.get_rank(group=process_group)
+                else:
+                    assert self.rank == local_rank
             elif mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
 
@@ -743,6 +788,8 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
+        assert self.rank is not None
+
         def enqueue_all2all_gather(start_idx, chunk_size):
             target_params = ordered_params[start_idx:start_idx + chunk_size]
             if target_params:
@@ -781,16 +828,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        PRE_STEP = 5
-        for i in range(0, PRE_STEP):
+        overlap_step = self.overlap_step
+        for i in range(0, overlap_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_computes(i + overlap_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_37be0d0_dirty
-ops = torch.ops._optimizer_37be0d0_dirty
+from . import _optimizer_6943c45_dirty
+ops = torch.ops._optimizer_6943c45_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_37be0d0_dirty::{op_name}"
+    return f"_optimizer_6943c45_dirty::{op_name}"

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

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

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

@@ -73,7 +73,7 @@ class _muon_state:
     qk_clip_state = None
 
 
-def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
+def split_elems_for_src(param, src_rank, num_ranks) -> int:
     rows = param.shape[0]
     cols = int(param.numel() // rows)
     base, rem = divmod(rows, num_ranks)
@@ -83,6 +83,10 @@ def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
 
 @torch.no_grad()
 def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate gathered_grad buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -102,13 +106,17 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
 @torch.no_grad()
 def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
                     alloc_event):
+    """
+    All2all gathers shards so each owner rank reconstructs its full gradient
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
 
-        # Calculate sending tensors
+        # Construct sending buffers
         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
@@ -125,15 +133,16 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             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 = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
-        # Calculate receiving tensors
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 total += split_elems_for_src(p, state, src, num_ranks)
@@ -141,6 +150,8 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -165,14 +176,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owner_params}
+        write_offsets = {id(p): 0 for p in owned_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:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 n = split_elems_for_src(p, state, src, num_ranks)
@@ -222,6 +233,10 @@ def _compute_u(p, state, steps, rank, compute_stream):
 
 @torch.no_grad()
 def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate scattered_u buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -234,18 +249,22 @@ def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
 
 
 def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
+    """
+    All2all scatters full gradients to all ranks
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
+        # Construct sending buffer
         per_dst = [[] for _ in range(num_ranks)]
         send_counts = [0] * num_ranks
 
-        if owner_params:
-            for p in owner_params:
+        if owned_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 if state.compute_event is None:
                     raise RuntimeError(
@@ -276,7 +295,9 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             # Even non-owner ranks must join the collective call
             send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
             for p in params:
@@ -290,6 +311,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
         assert recv_total > 0
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -298,6 +320,21 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             group=process_group,
         )
 
+        # Copy to pre-allocated scattered_u buffer from the received buffer
+        #
+        #                  recv_buf (num ranks = 3, local_rank = 0)
+        #
+        #      From rank 0        From rank 1       From rank 2
+        # | p1_0, p2_0, p3_0 |      p4_0       |    p5_0, p6_0    |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # src(0) :  p1_0 -> p2_0 -> p3_0
+        # src(1) :  p4_0
+        # src(2) :  p5_0 -> p6_0
+
         comm_stream.wait_event(alloc_event)
 
         off = 0
@@ -479,6 +516,11 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
+        overlap_step : How many all2all gather, compute operations are launched in advance
+                       before the corresponding all2all scatter steps begin.
+                       A higher overlap_step increases memory usage but can improve
+                       performance by overlapping communication.
+                       Parallel muon only. 
     """
 
     def __init__(self,
@@ -497,7 +539,8 @@ class Muon(torch.optim.Optimizer):
                      "k_indices": [],
                      "head_dim": 128,
                      "threshold": 100
-                 }):
+                 },
+                 overlap_step=5):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -521,15 +564,13 @@ class Muon(torch.optim.Optimizer):
 
         super().__init__(params, defaults)
 
-        if dist.is_initialized():
-            self.rank = dist.get_rank()
-        else:
-            self.rank = None
+        self.rank = None
 
         self.comm_stream = torch.cuda.Stream()
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
+        self.overlap_step = overlap_step
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -602,7 +643,11 @@ 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)
+                local_rank = dist.get_rank(group=process_group)
+                if self.rank is None:
+                    self.rank = dist.get_rank(group=process_group)
+                else:
+                    assert self.rank == local_rank
             elif mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
 
@@ -743,6 +788,8 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
+        assert self.rank is not None
+
         def enqueue_all2all_gather(start_idx, chunk_size):
             target_params = ordered_params[start_idx:start_idx + chunk_size]
             if target_params:
@@ -781,16 +828,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        PRE_STEP = 5
-        for i in range(0, PRE_STEP):
+        overlap_step = self.overlap_step
+        for i in range(0, overlap_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_computes(i + overlap_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_37be0d0_dirty
-ops = torch.ops._optimizer_37be0d0_dirty
+from . import _optimizer_6943c45_dirty
+ops = torch.ops._optimizer_6943c45_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_37be0d0_dirty::{op_name}"
+    return f"_optimizer_6943c45_dirty::{op_name}"

build/torch28-cxx11-rocm63-x86_64-linux/optimizer/{_optimizer_37be0d0_dirty.abi3.so → _optimizer_6943c45_dirty.abi3.so}

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

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

@@ -73,7 +73,7 @@ class _muon_state:
     qk_clip_state = None
 
 
-def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
+def split_elems_for_src(param, src_rank, num_ranks) -> int:
     rows = param.shape[0]
     cols = int(param.numel() // rows)
     base, rem = divmod(rows, num_ranks)
@@ -83,6 +83,10 @@ def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
 
 @torch.no_grad()
 def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate gathered_grad buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -102,13 +106,17 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
 @torch.no_grad()
 def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
                     alloc_event):
+    """
+    All2all gathers shards so each owner rank reconstructs its full gradient
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
 
-        # Calculate sending tensors
+        # Construct sending buffers
         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
@@ -125,15 +133,16 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             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 = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
-        # Calculate receiving tensors
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 total += split_elems_for_src(p, state, src, num_ranks)
@@ -141,6 +150,8 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -165,14 +176,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owner_params}
+        write_offsets = {id(p): 0 for p in owned_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:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 n = split_elems_for_src(p, state, src, num_ranks)
@@ -222,6 +233,10 @@ def _compute_u(p, state, steps, rank, compute_stream):
 
 @torch.no_grad()
 def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate scattered_u buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -234,18 +249,22 @@ def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
 
 
 def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
+    """
+    All2all scatters full gradients to all ranks
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
+        # Construct sending buffer
         per_dst = [[] for _ in range(num_ranks)]
         send_counts = [0] * num_ranks
 
-        if owner_params:
-            for p in owner_params:
+        if owned_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 if state.compute_event is None:
                     raise RuntimeError(
@@ -276,7 +295,9 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             # Even non-owner ranks must join the collective call
             send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
             for p in params:
@@ -290,6 +311,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
         assert recv_total > 0
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -298,6 +320,21 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             group=process_group,
         )
 
+        # Copy to pre-allocated scattered_u buffer from the received buffer
+        #
+        #                  recv_buf (num ranks = 3, local_rank = 0)
+        #
+        #      From rank 0        From rank 1       From rank 2
+        # | p1_0, p2_0, p3_0 |      p4_0       |    p5_0, p6_0    |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # src(0) :  p1_0 -> p2_0 -> p3_0
+        # src(1) :  p4_0
+        # src(2) :  p5_0 -> p6_0
+
         comm_stream.wait_event(alloc_event)
 
         off = 0
@@ -479,6 +516,11 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
+        overlap_step : How many all2all gather, compute operations are launched in advance
+                       before the corresponding all2all scatter steps begin.
+                       A higher overlap_step increases memory usage but can improve
+                       performance by overlapping communication.
+                       Parallel muon only. 
     """
 
     def __init__(self,
@@ -497,7 +539,8 @@ class Muon(torch.optim.Optimizer):
                      "k_indices": [],
                      "head_dim": 128,
                      "threshold": 100
-                 }):
+                 },
+                 overlap_step=5):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -521,15 +564,13 @@ class Muon(torch.optim.Optimizer):
 
         super().__init__(params, defaults)
 
-        if dist.is_initialized():
-            self.rank = dist.get_rank()
-        else:
-            self.rank = None
+        self.rank = None
 
         self.comm_stream = torch.cuda.Stream()
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
+        self.overlap_step = overlap_step
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -602,7 +643,11 @@ 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)
+                local_rank = dist.get_rank(group=process_group)
+                if self.rank is None:
+                    self.rank = dist.get_rank(group=process_group)
+                else:
+                    assert self.rank == local_rank
             elif mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
 
@@ -743,6 +788,8 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
+        assert self.rank is not None
+
         def enqueue_all2all_gather(start_idx, chunk_size):
             target_params = ordered_params[start_idx:start_idx + chunk_size]
             if target_params:
@@ -781,16 +828,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        PRE_STEP = 5
-        for i in range(0, PRE_STEP):
+        overlap_step = self.overlap_step
+        for i in range(0, overlap_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_computes(i + overlap_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_37be0d0_dirty
-ops = torch.ops._optimizer_37be0d0_dirty
+from . import _optimizer_6943c45_dirty
+ops = torch.ops._optimizer_6943c45_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_optimizer_37be0d0_dirty::{op_name}"
+    return f"_optimizer_6943c45_dirty::{op_name}"

build/torch28-cxx11-rocm64-x86_64-linux/optimizer/{_optimizer_37be0d0_dirty.abi3.so → _optimizer_6943c45_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:88f4b7caa9e13481a51ec93422477854ab460b2197c6b469ebe6d6cb8f81ddf5
+oid sha256:76ab2b9f9b47115d3ec6fbef5863bd80f3673a95ddf404ae03d8e516c3e3167a
 size 1750024

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

@@ -73,7 +73,7 @@ class _muon_state:
     qk_clip_state = None
 
 
-def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
+def split_elems_for_src(param, src_rank, num_ranks) -> int:
     rows = param.shape[0]
     cols = int(param.numel() // rows)
     base, rem = divmod(rows, num_ranks)
@@ -83,6 +83,10 @@ def split_elems_for_src(param, state, src_rank, num_ranks) -> int:
 
 @torch.no_grad()
 def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate gathered_grad buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -102,13 +106,17 @@ def _alloc_gathered_grad(params, param_to_state, rank, compute_stream):
 @torch.no_grad()
 def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
                     alloc_event):
+    """
+    All2all gathers shards so each owner rank reconstructs its full gradient
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
 
-        # Calculate sending tensors
+        # Construct sending buffers
         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
@@ -125,15 +133,16 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             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 = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
-        # Calculate receiving tensors
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
-            for p in owner_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 total += split_elems_for_src(p, state, src, num_ranks)
@@ -141,6 +150,8 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
 
         recv_total = sum(recv_counts)
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
+
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -165,14 +176,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         comm_stream.wait_event(alloc_event)
 
         off = 0
-        write_offsets = {id(p): 0 for p in owner_params}
+        write_offsets = {id(p): 0 for p in owned_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:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 assert state.worker_rank == rank
                 n = split_elems_for_src(p, state, src, num_ranks)
@@ -222,6 +233,10 @@ def _compute_u(p, state, steps, rank, compute_stream):
 
 @torch.no_grad()
 def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
+    """
+    Pre-allocate scattered_u buffer on compute_stream
+    before launching all2all gather
+    """
     with torch.cuda.stream(compute_stream):
         for p in params:
             state = param_to_state[id(p)]
@@ -234,18 +249,22 @@ def _alloc_scattered_u(params, param_to_state, rank, compute_stream):
 
 
 def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
+    """
+    All2all scatters full gradients to all ranks
+    """
     with torch.cuda.stream(comm_stream):
         process_group = param_to_state[id(params[0])].process_group
         num_ranks = dist.get_world_size(group=process_group)
-        owner_params = [
+        owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
 
+        # Construct sending buffer
         per_dst = [[] for _ in range(num_ranks)]
         send_counts = [0] * num_ranks
 
-        if owner_params:
-            for p in owner_params:
+        if owned_params:
+            for p in owned_params:
                 state = param_to_state[id(p)]
                 if state.compute_event is None:
                     raise RuntimeError(
@@ -276,7 +295,9 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             # Even non-owner ranks must join the collective call
             send_buf = torch.empty(0, dtype=COMM_DTYPE, device="cuda")
 
+        # Compute receive sizes and allocate receiving buffers
         recv_counts = [0] * num_ranks
+
         for src in range(num_ranks):
             total = 0
             for p in params:
@@ -290,6 +311,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
         assert recv_total > 0
         recv_buf = torch.empty(recv_total, dtype=COMM_DTYPE, device="cuda")
 
+        #All2All
         dist.all_to_all_single(
             recv_buf,
             send_buf,
@@ -298,6 +320,21 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             group=process_group,
         )
 
+        # Copy to pre-allocated scattered_u buffer from the received buffer
+        #
+        #                  recv_buf (num ranks = 3, local_rank = 0)
+        #
+        #      From rank 0        From rank 1       From rank 2
+        # | p1_0, p2_0, p3_0 |      p4_0       |    p5_0, p6_0    |
+        #
+        # Outer loop:
+        # rank 0 -> rank 1 -> rank2
+        #
+        # Inner loop:
+        # src(0) :  p1_0 -> p2_0 -> p3_0
+        # src(1) :  p4_0
+        # src(2) :  p5_0 -> p6_0
+
         comm_stream.wait_event(alloc_event)
 
         off = 0
@@ -479,6 +516,11 @@ class Muon(torch.optim.Optimizer):
                     "head_dim": 128,
                     "threshold": 100
                 }
+        overlap_step : How many all2all gather, compute operations are launched in advance
+                       before the corresponding all2all scatter steps begin.
+                       A higher overlap_step increases memory usage but can improve
+                       performance by overlapping communication.
+                       Parallel muon only. 
     """
 
     def __init__(self,
@@ -497,7 +539,8 @@ class Muon(torch.optim.Optimizer):
                      "k_indices": [],
                      "head_dim": 128,
                      "threshold": 100
-                 }):
+                 },
+                 overlap_step=5):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -521,15 +564,13 @@ class Muon(torch.optim.Optimizer):
 
         super().__init__(params, defaults)
 
-        if dist.is_initialized():
-            self.rank = dist.get_rank()
-        else:
-            self.rank = None
+        self.rank = None
 
         self.comm_stream = torch.cuda.Stream()
         self.compute_stream = torch.cuda.Stream()
         self.debug = debug
         self.clip_config = clip_config
+        self.overlap_step = overlap_step
 
     def _calc_flops(self, G, steps):
         assert len(G.shape) == 2
@@ -602,7 +643,11 @@ 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)
+                local_rank = dist.get_rank(group=process_group)
+                if self.rank is None:
+                    self.rank = dist.get_rank(group=process_group)
+                else:
+                    assert self.rank == local_rank
             elif mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
 
@@ -743,6 +788,8 @@ class Muon(torch.optim.Optimizer):
         param_to_state, ordered_params = self.init_state_and_assign_params(
             names, params, group, qk_logits)
 
+        assert self.rank is not None
+
         def enqueue_all2all_gather(start_idx, chunk_size):
             target_params = ordered_params[start_idx:start_idx + chunk_size]
             if target_params:
@@ -781,16 +828,16 @@ class Muon(torch.optim.Optimizer):
         # Wait grad update
         self.comm_stream.wait_stream(torch.cuda.current_stream())
 
-        PRE_STEP = 5
-        for i in range(0, PRE_STEP):
+        overlap_step = self.overlap_step
+        for i in range(0, overlap_step):
             enqueue_all2all_gather(i * chunk_size, chunk_size)
             enqueue_computes(i * chunk_size, chunk_size)
 
         for i in range(0, len(params) + chunk_size - 1, chunk_size):
             enqueue_all2all_scatter(i, chunk_size)
-            enqueue_all2all_gather(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_all2all_gather(i + overlap_step * chunk_size, chunk_size)
             enqueue_update_param(i, chunk_size)
-            enqueue_computes(i + PRE_STEP * chunk_size, chunk_size)
+            enqueue_computes(i + overlap_step * chunk_size, chunk_size)
 
         # Wait the last update_param to finish
         torch.cuda.current_stream().wait_stream(self.compute_stream)