fix assert in a2a gather scatter

torch-ext/optimizer/muon.py

@@ -128,11 +128,14 @@ def _all2all_gather(params, param_to_state, rank, comm_stream, none_grad,
             per_dst[dst].append(g)
             send_counts[dst] += shard_elems
 
-        assert all(
-            len(v) > 0
-            for v in per_dst), "all params should be sharded to all devices"
+        assert any(
+            len(v) > 0 for v in per_dst
+        ), "At least one destination rank must receive a sharded tensor"
+        # list[list[Tensor]] -> list[Tensor]
+        per_dst = [t for dst in per_dst for t in dst]
+
+        send_buf = torch.cat(per_dst, dim=0)
 
-        send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
         owned_params = [
             p for p in params if param_to_state[id(p)].worker_rank == rank
         ]
@@ -288,8 +291,14 @@ def _all2all_scatter(params, param_to_state, rank, comm_stream, alloc_event):
 
                 assert offset == u_full.numel()
 
-        if any(len(v) > 0 for v in per_dst):
-            send_buf = torch.cat([torch.cat(v, dim=0) for v in per_dst], dim=0)
+        lengths = [len(v) for v in per_dst]
+        if all(l > 0 for l in lengths):
+            assert all(
+                l == lengths[0] for l in lengths
+            ), "All destination ranks must have the same number of sharded tensor"
+            # list[list[Tensor]] -> list[Tensor]
+            per_dst = [t for dst in per_dst for t in dst]
+            send_buf = torch.cat(per_dst, dim=0)
         else:
             # all_to_all requires participation from all ranks
             # Even non-owner ranks must join the collective call