torchmonarch-nightly 2025.6.27__cp312-cp312-manylinux2014_x86_64.whl

tests/test_grad_generator.py

@@ -0,0 +1,121 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+from unittest import main, TestCase
+
+import torch
+from monarch.gradient._gradient_generator import GradientGenerator
+from monarch.gradient_generator import gradient_execution_order
+
+
+class TestGradIter(TestCase):
+    def checkEqual(self, r, r2):
+        self.assertEqual(len(r), len(r2))
+        for i, i2 in zip(r, r2):
+            self.assertTrue((i is None and i2 is None) or torch.allclose(i, i2))
+
+    def test_simple(self):
+        t = torch.rand(2, requires_grad=True)
+        t2 = torch.rand(2, requires_grad=True)
+
+        _ = t + t2
+        a, b = torch.std_mean(t + t2)
+
+        r2 = torch.autograd.grad([a, b], [t2, t], retain_graph=True)
+        r = list(GradientGenerator([a, b], [t2, t]))
+        print(a, b)
+        print(a.grad_fn, b.grad_fn)
+
+        print(r)
+        self.checkEqual(r, r2)
+
+    def test_pipeline_like(self):
+        t = torch.rand(3, 3, requires_grad=True)
+
+        w1 = torch.rand(3, 2, requires_grad=True)
+        w2 = torch.rand(3, 2, requires_grad=True)
+        w3 = torch.rand(3, 2, requires_grad=True)
+
+        u = torch.rand(3, 2, requires_grad=True)
+
+        _ = u * u
+
+        w4 = torch.rand(2, 3, requires_grad=True)
+        w5 = torch.rand(2, 3, requires_grad=True)
+        w6 = torch.rand(2, 3, requires_grad=True)
+
+        from torch.nn.functional import relu
+
+        a = relu(t @ (w1 @ w4))
+        a = relu(a @ (w2 @ w5))
+        a = relu(a @ (w3 @ w6))
+
+        std, mean = torch.std_mean(a)
+        loss = std + std
+
+        cgrads = torch.autograd.grad(
+            [loss], [t, w3, w6, u, w2, w5], allow_unused=True, retain_graph=True
+        )
+        gi = GradientGenerator([loss], [t, w3, w6, u, w2, w5])
+        grads = [*gi]
+        self.checkEqual(grads, cgrads)
+
+    def test_tree(self):
+        t = torch.rand(3, 3, requires_grad=True)
+
+        t2 = t + t
+        t3 = t * t
+        t4 = t / t
+        t5 = t - t
+
+        t6 = t2 * t3
+        t7 = t4 * t5
+        t8 = t2 * t4
+        t9 = t3 * t5
+        t10 = t6 + t7 + t8 + t9
+
+        t11 = t10.sum()
+
+        cgrads = torch.autograd.grad([t11], [t2, t], retain_graph=True)
+        gi = GradientGenerator([t11], [t2, t])
+        grads = [*gi]
+        self.checkEqual(grads, cgrads)
+
+    def test_broadcast(self):
+        t = torch.rand(3, 3, requires_grad=True)
+        t2 = torch.rand(3, requires_grad=True)
+        t3 = t2 / t2
+
+        r = (t * t3).sum()
+        cgrads = torch.autograd.grad([r], [t, t2], retain_graph=True)
+        gi = GradientGenerator([r], [t, t2])
+        grads = [*gi]
+        self.checkEqual(grads, cgrads)
+
+    def test_grad_order(self):
+        t = torch.rand(3, 3, requires_grad=True)
+        w1 = torch.rand(3, 3, requires_grad=True)
+        w2 = torch.rand(3, 3, requires_grad=True)
+        w3 = torch.rand(3, 3, requires_grad=True)
+
+        u = torch.rand(3, 2, requires_grad=True)
+        _ = u * u
+        from torch.nn.functional import relu
+
+        a = relu(t @ w1)
+        a = relu(a @ w2)
+        a = relu(a @ w3)
+
+        std, mean = torch.std_mean(a)
+        loss = std + std
+
+        order = gradient_execution_order([loss], [w2, w3, w1, a])
+        self.assertEqual(order, [3, 1, 0, 2])
+
+
+if __name__ == "__main__":
+    main()

tests/test_mock_cuda.py

@@ -0,0 +1,74 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+from unittest import main, TestCase
+
+import pytest
+import torch
+import monarch.common.mock_cuda  # usort: skip
+
+
+def simple_forward_backward(device: str) -> None:
+    torch.manual_seed(123)
+    m = torch.nn.Sequential(torch.nn.Linear(3, 3), torch.nn.ReLU()).to(device)
+    x = torch.rand(10, 3).to(device)
+    y = m(x)
+    loss_fn = torch.nn.CrossEntropyLoss()
+    loss = loss_fn(y, torch.randint(3, (10,)).to(device))
+    # Under the hood, enabling/disabling CUDA mocking is done with a thread-local
+    # flag. By default, backward() executes ops on a different thread than the one
+    # we enabled mocking on, which would lead to an invalid memory access. So we need
+    # to disable multithreading for backward.
+    with torch.autograd.set_multithreading_enabled(False):
+        loss.backward()
+    # pyre-ignore: Incompatible return type [7]: Expected `None` but got `Tuple[typing.Any, Union[None, Tensor, Module], Union[None, Tensor, Module]]`.
+    return y, m[0].weight.grad, m[0].bias.grad
+
+
+# Mock cuda depends on initialization load order
+# For OSS, run this test separately until it can be run in a subprocess.
+@pytest.mark.oss_skip
+class TestMockCuda(TestCase):
+    def setUp(self) -> None:
+        return super().setUp()
+
+    def test_output_is_garbage(self):
+        with monarch.common.mock_cuda.mock_cuda_guard():
+            x = torch.arange(9, device="cuda", dtype=torch.float32).reshape(3, 3)
+            y = 2 * torch.eye(3, device="cuda")
+            true_output = torch.tensor(
+                [[0, 2, 4], [6, 8, 10], [12, 14, 16]], dtype=torch.float32
+            )
+            self.assertFalse(torch.equal((x @ y).cpu(), true_output))
+
+    def test_simple_forward_backward(self):
+        # This test just makes sure that the forward and backward pass work
+        # and don't crash.
+        simple_forward_backward("cuda")
+
+    def test_turn_mock_on_and_off(self):
+        cpu_y, cpu_dw, cpu_db = simple_forward_backward("cpu")
+
+        real_y, real_dw, real_db = simple_forward_backward("cuda")
+        self.assertTrue(torch.allclose(cpu_y, real_y.cpu()))
+        self.assertTrue(torch.allclose(cpu_dw, real_dw.cpu()))
+        self.assertTrue(torch.allclose(cpu_db, real_db.cpu()))
+
+        with monarch.common.mock_cuda.mock_cuda_guard():
+            mocked_y, mocked_dw, mocked_db = simple_forward_backward("cuda")
+            self.assertFalse(torch.allclose(cpu_y, mocked_y.cpu()))
+            self.assertFalse(torch.allclose(cpu_dw, mocked_dw.cpu()))
+            self.assertFalse(torch.allclose(cpu_db, mocked_db.cpu()))
+
+        real_y, real_dw, real_db = simple_forward_backward("cuda")
+        self.assertTrue(torch.allclose(cpu_y, real_y.cpu()))
+        self.assertTrue(torch.allclose(cpu_dw, real_dw.cpu()))
+        self.assertTrue(torch.allclose(cpu_db, real_db.cpu()))
+
+
+if __name__ == "__main__":
+    main()

tests/test_pdb_actor.py

@@ -0,0 +1,110 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+import sys
+import traceback
+from contextlib import contextmanager
+from typing import Generator
+
+import pytest
+
+import torch
+
+from monarch import DeviceMesh, fetch_shard, remote, rust_local_mesh
+from monarch._rust_bindings.monarch_extension.client import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ClientActor,
+    DebuggerMessage as ClientDebuggerMessage,
+)
+
+from monarch._rust_bindings.monarch_extension.debugger import (
+    DebuggerMessage as PdbDebuggerMessage,
+    get_bytes_from_write_action,
+)
+from monarch._rust_bindings.monarch_messages.debugger import DebuggerAction
+from monarch.rust_local_mesh import LoggingLocation, SocketType
+from monarch_supervisor.logging import fix_exception_lines
+
+
+def custom_excepthook(exc_type, exc_value, exc_traceback):
+    tb_lines = fix_exception_lines(
+        traceback.format_exception(exc_type, exc_value, exc_traceback)
+    )
+    print("\n".join(tb_lines), file=sys.stderr)
+
+
+sys.excepthook = custom_excepthook
+
+
+@contextmanager
+def local_mesh(
+    hosts: int = 1, gpu_per_host: int = 2, activate: bool = True
+) -> Generator[DeviceMesh, None, None]:
+    with rust_local_mesh.local_mesh(
+        hosts=hosts,
+        gpus_per_host=gpu_per_host,
+        socket_type=SocketType.UNIX,
+        logging_location=LoggingLocation.DEFAULT,
+    ) as dm:
+        try:
+            if activate:
+                with dm.activate():
+                    yield dm
+            else:
+                yield dm
+            dm.exit()
+        except Exception:
+            dm.client._shutdown = True
+            raise
+
+
+remote_test_pdb_actor = remote(
+    "monarch.worker._testing_function.test_pdb_actor",
+    propagate=lambda: torch.zeros(1),
+)
+
+
+@pytest.mark.skipif(
+    torch.cuda.device_count() < 2,
+    reason="Not enough GPUs, this test requires at least 2 GPUs",
+)
+# Set global timeout--sandcastle's timeout is 600s. A test that sandcastle times
+# out is not counted as a failure, so we set a more restrictive timeout to
+# ensure we see a hard failure in CI.
+@pytest.mark.timeout(120)
+class TestPdbActor:
+    def test_pdb_actor(self):
+        with local_mesh(1, 1) as dm:
+            with dm.activate():
+                client = dm.client.inner._actor
+                assert isinstance(client, ClientActor)
+                fut = fetch_shard(remote_test_pdb_actor())
+                msg = client.get_next_message(timeout_msec=None)
+                assert isinstance(msg, ClientDebuggerMessage)
+                assert isinstance(msg.action, DebuggerAction.Paused)
+                client.send(
+                    msg.debugger_actor_id,
+                    PdbDebuggerMessage(action=DebuggerAction.Attach()).serialize(),
+                )
+                msg = client.get_next_message(timeout_msec=None)
+                assert isinstance(msg, ClientDebuggerMessage)
+                assert isinstance(msg.action, DebuggerAction.Read)
+                assert msg.action.requested_size == 4
+                client.send(
+                    msg.debugger_actor_id,
+                    PdbDebuggerMessage(
+                        action=DebuggerAction.Write(b"1234")
+                    ).serialize(),
+                )
+                msg = client.get_next_message(timeout_msec=None)
+                assert isinstance(msg, ClientDebuggerMessage)
+                assert isinstance(msg.action, DebuggerAction.Write)
+                assert get_bytes_from_write_action(msg.action) == b"5678"
+                client.send(
+                    msg.debugger_actor_id,
+                    PdbDebuggerMessage(action=DebuggerAction.Detach()).serialize(),
+                )
+                fut.result()