Update muon.py

torch-ext/optimizer/muon.py

@@ -597,7 +597,7 @@ class Muon(torch.optim.Optimizer):
         adjusted_lr = lr * adjusted_ratio
         return adjusted_lr
 
-    def get_shard_mesh(self, p, rank):
+    def get_shard_mesh(self, p):
         """
         Get the shard mesh for a parameter p on the given rank.
         """
@@ -609,8 +609,13 @@ class Muon(torch.optim.Optimizer):
             return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
         elif p.placements == (Replicate(), Shard(dim=0)):
             # Case for HSDP
+            process_group = p.device_mesh.get_group(mesh_dim=1)
+            if self.rank is None:           
+                self.rank = dist.get_rank(group=process_group)
+            else:
+                assert self.rank == dist.get_rank(group=process_group)
             for i, shard_mesh in enumerate(p.device_mesh.mesh):
-                if rank in shard_mesh:
+                if self.rank in shard_mesh:
                     return shard_mesh, p.device_mesh.get_group(mesh_dim=1)
         else:
             raise ValueError(f"Unsupported placements ({p.placements}).")
@@ -651,15 +656,9 @@ class Muon(torch.optim.Optimizer):
         for n, p in zip(ordered_names, ordered_params):
             if mesh is None:
                 mesh = p.device_mesh
-                shard_mesh, process_group = self.get_shard_mesh(p, self.rank)
-                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
+                shard_mesh, process_group = self.get_shard_mesh(p)                
             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(