add doc strings + init self rank on init_assign_params

torch-ext/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
@@ -129,8 +137,9 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             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:
@@ -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,
@@ -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,6 +249,9 @@ 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)
@@ -241,6 +259,7 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
             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
 
@@ -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
@@ -521,10 +558,7 @@ 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()
@@ -602,7 +636,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 +781,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: