add muon tests

PrimeIntellect-ai · Dec 5, 2024 · 02e1caa · 02e1caa
1 parent 5cd347c
commit 02e1caa
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/tests/test_torchrun/test_train.py b/tests/test_torchrun/test_train.py
@@ -114,8 +114,12 @@ def test_packing(packing: bool):
     _test_multi_gpu(num_gpus, "debug/normal.toml", extra_args=[packing_arg])
 
 
-@pytest.mark.parametrize("diloco", [False])
+@pytest.mark.parametrize("diloco", [False, True])
 def test_muon(diloco: bool):
-    num_gpus = [1, 2]
-    config_file = "debug/diloco.toml" if diloco else "debug/normal.toml"
-    _test_multi_gpu(num_gpus, config_file, extra_args=["--optim.optim.pseudo_order_steps", "6"])
+    num_gpus = [1, 2] if diloco else [2, 1]
+    _test_multi_gpu(
+        num_gpus,
+        "debug/diloco.toml" if diloco else "debug/normal.toml",
+        extra_args=["--optim.optim.pseudo_order_steps", "6"],
+        diloco=diloco,
+    )