fix bug in fsdp

torch-ext/optimizer/muon.py

@@ -606,6 +606,11 @@ class Muon(torch.optim.Optimizer):
 
         if p.placements == (Shard(dim=0), ):
             # Case for FSDP
+            process_group = p.device_mesh.get_group(mesh_dim=0)
+            if self.rank is None:
+                self.rank = dist.get_rank(group=process_group)
+            else:
+                assert self.rank == dist.get_rank(group=process_group)
             return p.device_mesh.mesh, p.device_mesh.get_group(mesh_dim=0)
         elif p.placements == (Replicate(), Shard(dim=0)):
             # Case for HSDP