consider multi node

torch-ext/optimizer/muon.py

@@ -112,6 +112,7 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
         for p in params:
             state = param_to_state[id(p)]
             dst = state.worker_rank
+            assert dst < num_ranks
             shard_elems = split_elems_for_src(p, state, rank, num_ranks)
             g = p.grad
             g = g.to_local().to(COMM_DTYPE).contiguous().view(-1)
@@ -601,11 +602,13 @@ class Muon(torch.optim.Optimizer):
             if mesh is None:
                 mesh = p.device_mesh
                 shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
+                self.rank = dist.get_rank(group=process_group)
             elif mesh != p.device_mesh:
                 raise ValueError("All parameters must be on the same mesh.")
 
+            num_ranks = dist.get_world_size(group=process_group)
             param_to_state[id(p)] = _muon_state()
-            param_to_state[id(p)].worker_rank = shard_mesh[round_robin].item()
+            param_to_state[id(p)].worker_rank = shard_mesh[round_robin].item() % num_ranks
             param_to_state[id(p)].process_group = process_group
             qk_clip_state = self.get_qk_clip_info(n, qk_logits)
             param_to_state[id(p)].qk_clip_state = qk_clip_state