Add CPU test for use_syrk forwarding in newton_schulz_tp

Follow-up to #239. In test_distributed_muon_utils_cpu.py, using a real gloo
process group, verifies use_syrk is forwarded through newton_schulz_tp to the
underlying newton_schulz for both TP modes and both partition dims: a 3D input
with use_syrk=True raises the code-owned TypeError 'use_syrk does not support
N-d input' (mirroring test_newton_schulz_use_syrk_with_3d_raises_type_error),
while use_syrk=False stays on the dense batched path and succeeds.

Signed-off-by: pingtianl <pingtianl@nvidia.com>

Author: Wohox

tests/test_distributed_muon_utils_cpu.py

@@ -21,6 +21,7 @@
 from absl import flags, logging
 from absl.testing import absltest, parameterized
 
+from emerging_optimizers import utils
 from emerging_optimizers.orthogonalized_optimizers import muon_utils
 
 
@@ -227,6 +228,41 @@ def test_1step_close_to_non_distributed(self, shape, partition_dim, tp_mode):
 
         torch.testing.assert_close(ref_out.chunk(world_size, dim=partition_dim)[rank], dist_out, atol=1e-6, rtol=0)
 
+    @parameterized.product(
+        partition_dim=(0, 1),
+        tp_mode=("distributed", "duplicated"),
+    )
+    def test_use_syrk_forwarded_to_newton_schulz(self, partition_dim, tp_mode):
+        shape = (8, 16, 32)
+        world_size = torch.distributed.get_world_size()
+        if shape[partition_dim] % world_size != 0:
+            self.skipTest("Skipping because incompatible shape and world size")
+        rank = torch.distributed.get_rank()
+        x = torch.randint(-5, 5, shape, device="cpu", dtype=torch.float32)
+        torch.distributed.all_reduce(x, op=torch.distributed.ReduceOp.SUM)
+        local_x = x.chunk(world_size, dim=partition_dim)[rank]
+
+        with utils.fp32_matmul_precision("medium"):
+            muon_utils.newton_schulz_tp(
+                local_x,
+                steps=1,
+                coefficient_type="simple",
+                tp_group=torch.distributed.group.WORLD,
+                partition_dim=partition_dim,
+                tp_mode=tp_mode,
+                use_syrk=False,
+            )
+            with self.assertRaisesRegex(TypeError, "use_syrk does not support"):
+                muon_utils.newton_schulz_tp(
+                    local_x,
+                    steps=1,
+                    coefficient_type="simple",
+                    tp_group=torch.distributed.group.WORLD,
+                    partition_dim=partition_dim,
+                    tp_mode=tp_mode,
+                    use_syrk=True,
+                )
+
 
 if __name__ == "__main__":
     torch.distributed.init_process_group(backend="gloo")