torchmonarch-nightly 2025.6.27__cp313-cp313-manylinux2014_x86_64.whl

monarch/tensorboard.py

@@ -0,0 +1,84 @@
+# 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 logging
+from typing import Any
+
+from monarch.common.device_mesh import DeviceMesh
+from monarch.remote_class import ControllerRemoteClass, WorkerRemoteClass
+from torch.utils.tensorboard import SummaryWriter
+
+logger = logging.getLogger(__name__)
+
+
+class Tensorboard(ControllerRemoteClass):
+    def __init__(self, coordinator: DeviceMesh, path: str, *args, **kwargs) -> None:
+        from monarch import IN_PAR
+
+        self.path = path
+        self.url: str = ""
+        self.coordinator = coordinator
+        if path.startswith("manifold://"):
+            if not IN_PAR:
+                raise RuntimeError(
+                    "Cannot save tensorboard to manifold with conda environment. "
+                    "Save to the local filesystem or oilfs instead"
+                )
+
+            manifold_url = f"https://internalfb.com/intern/tensorboard/?dir={path}"
+            self.url = manifold_url
+        else:
+            self.url = path
+
+        # Only create tensorboard for the coordinator rank.
+        with self.coordinator.activate():
+            super().__init__(
+                "monarch.tensorboard._WorkerSummaryWriter",
+                path,
+                *args,
+                **kwargs,
+            )
+
+        logger.info("Run `tensorboard --logdir %s` to launch the tensorboard.", path)
+
+    @ControllerRemoteClass.remote_method
+    def _log(self, name: str, data: Any, step: int) -> None:
+        pass
+
+    @ControllerRemoteClass.remote_method
+    def _flush(self) -> None:
+        pass
+
+    @ControllerRemoteClass.remote_method
+    def _close(self) -> None:
+        pass
+
+    def log(self, name: str, data: Any, step: int) -> None:
+        with self.coordinator.activate():
+            self._log(name, data, step)
+
+    def flush(self) -> None:
+        with self.coordinator.activate():
+            self._flush()
+
+    def close(self) -> None:
+        with self.coordinator.activate():
+            self._close()
+
+
+class _WorkerSummaryWriter(WorkerRemoteClass):
+    def __init__(self, path: str, *args, **kwargs) -> None:
+        self._writer = SummaryWriter(path, *args, **kwargs)
+
+    def _log(self, name: str, data: Any, step: int) -> None:
+        self._writer.add_scalar(name, data, global_step=step, new_style=True)
+
+    def _flush(self) -> None:
+        self._writer.flush()
+
+    def _close(self) -> None:
+        self._writer.close()

monarch/timer/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+
+from .execution_timer import (
+    execution_timer_start,
+    execution_timer_stop,
+    ExecutionTimer,
+    get_execution_timer_average_ms,
+    get_latest_timer_measurement,
+)
+
+__all__ = [
+    "ExecutionTimer",
+    "execution_timer_start",
+    "execution_timer_stop",
+    "get_latest_timer_measurement",
+    "get_execution_timer_average_ms",
+]

monarch/timer/example_monarch.py

@@ -0,0 +1,78 @@
+# 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.
+
+"""An example that demonstrates how to use ExecutionTimer with a Monarch program.
+
+Run this with
+buck run //monarch/python/monarch/timer:example_monarch
+
+"""
+# pyre-unsafe
+
+import logging
+
+import torch
+
+from monarch import inspect, remote
+from monarch.rust_local_mesh import local_mesh
+
+logger = logging.getLogger(__name__)
+
+
+execution_timer_start = remote(
+    "monarch.timer.remote_execution_timer.execution_timer_start", propagate="inspect"
+)
+
+execution_timer_stop = remote(
+    "monarch.timer.remote_execution_timer.execution_timer_stop", propagate="inspect"
+)
+
+get_execution_timer_average_ms = remote(
+    "monarch.timer.remote_execution_timer.get_execution_timer_average_ms",
+    propagate=lambda: torch.tensor(0.0, dtype=torch.float64),
+)
+
+get_time_perfcounter = remote(
+    "monarch.timer.remote_execution_timer.get_time_perfcounter",
+    propagate=lambda: torch.tensor(0.0, dtype=torch.float64),
+)
+
+
+def main() -> None:
+    with local_mesh(hosts=1, gpus_per_host=1) as mesh:
+        with mesh.activate():
+            num_iterations = 5
+
+            a = torch.randn(1000, 1000, device="cuda")
+            b = torch.randn(1000, 1000, device="cuda")
+            torch.matmul(a, b)
+
+            total_dt = torch.zeros(1, dtype=torch.float64)
+
+            for _ in range(num_iterations):
+                t0 = get_time_perfcounter()
+                torch.matmul(a, b)
+                total_dt += get_time_perfcounter() - t0
+
+            for _ in range(num_iterations):
+                execution_timer_start()
+                torch.matmul(a, b)
+                execution_timer_stop()
+
+            cuda_average_ms = get_execution_timer_average_ms()
+            local_total_dt = inspect(total_dt)
+            local_cuda_avg_ms = inspect(cuda_average_ms)
+
+        local_total_dt = local_total_dt.item()
+        local_cuda_avg_ms = local_cuda_avg_ms.item()
+        mesh.exit()
+    avg_perfcounter_ms = local_total_dt / num_iterations * 1000
+    print(f"average time w/ perfcounter: {avg_perfcounter_ms:.4f} (ms)")
+    print(f"average time w/ ExecutionTimer:   {local_cuda_avg_ms:.4f} (ms)")
+
+
+if __name__ == "__main__":
+    main()

monarch/timer/example_spmd.py

@@ -0,0 +1,55 @@
+# 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.
+
+"""An example that demonstrates how to use ExecutionTimer in a SPMD style program.
+
+Run this with:
+buck run //monarch/python/monarch/timer:example_spmd
+"""
+
+import time
+
+# pyre-strict
+
+import torch
+from monarch.timer import ExecutionTimer
+
+
+def main() -> None:
+    # Check if CUDA is available
+    if not torch.cuda.is_available():
+        print("CUDA is not available. Exiting.")
+        return
+
+    device = torch.device("cuda")
+
+    num_iterations = 5
+
+    a = torch.randn(1000, 1000, device=device)
+    b = torch.randn(1000, 1000, device=device)
+
+    # Warmup
+    torch.matmul(a, b)
+    torch.cuda.synchronize()
+
+    cpu_timings = []
+    for _ in range(num_iterations):
+        t0 = time.perf_counter()
+        torch.matmul(a, b)
+        cpu_timings.append(time.perf_counter() - t0)
+
+    for _ in range(num_iterations):
+        with ExecutionTimer.time("matrix_multiply"):
+            torch.matmul(a, b)
+
+    mean_cuda_ms = ExecutionTimer.summary()["matrix_multiply"]["mean_ms"]
+    mean_perfcounter_ms = sum(cpu_timings) / len(cpu_timings) * 1000
+    print("mean perf counter times: ", mean_perfcounter_ms)
+    print("mean cuda times: ", mean_cuda_ms)
+
+
+if __name__ == "__main__":
+    main()

monarch/timer/execution_timer.py

@@ -0,0 +1,199 @@
+# 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.
+
+"""A simple timer that utilizes CUDA events to measure time spent in GPU kernels."""
+
+# pyre-strict
+import logging
+import threading
+import time
+from contextlib import contextmanager
+from typing import Any, Dict, Generator, List, Optional, Tuple
+
+import torch
+
+
+class ExecutionTimer:
+    """
+    A lightweight timer for measuring CPU or GPU execution time.
+    """
+
+    _enable_cuda: bool = torch.cuda.is_available()
+    _times: Dict[str, List[float]] = {}
+    _lock = threading.Lock()
+    _threads: Dict[str, List[threading.Thread]] = {}
+    _events: Dict[
+        str, List[Tuple[torch.cuda.Event, torch.cuda.Event, torch.cuda.Stream]]
+    ] = {}
+    _cuda_warning_shown: bool = False
+    _cpu_start_times: Dict[str, List[float]] = {}
+
+    @classmethod
+    @contextmanager
+    # pyre-fixme[3]: Return type must be specified as type that does not contain `Any`.
+    def time(
+        cls, name: Optional[str] = None, use_cpu: bool = False
+    ) -> Generator[None, Any, Any]:
+        """
+        Context manager for timing an operation.
+        Args:
+            name (str): Name of the timing section
+            use_cpu (bool): Whether to use CPU time instead of CUDA time. Defaults to false.
+        Example:
+            with ExecutionTimer.time("matrix_multiply"):
+                result = torch.matmul(a, b)
+
+            with ExecutionTimer.time("sleep", use_cpu=True):
+                time.sleep(1)
+        """
+        cls.start(name, use_cpu)
+        try:
+            yield
+        finally:
+            cls.stop(name, use_cpu)
+
+    @classmethod
+    def start(cls, name: Optional[str] = None, use_cpu: bool = False) -> None:
+        if not cls._enable_cuda and cls._cuda_warning_shown:
+            logging.warning("CUDA not available, falling back to CPU timing")
+            cls._cuda_warning_shown = True
+
+        if not name:
+            name = "default"
+        if name not in cls._times:
+            cls._times[name] = []
+
+        if not cls._enable_cuda or use_cpu:
+            if name not in cls._cpu_start_times:
+                cls._cpu_start_times[name] = []
+            cls._cpu_start_times[name].append(time.perf_counter())
+        else:
+            stream = torch.cuda.current_stream()
+            start_event = torch.cuda.Event(enable_timing=True)
+            end_event = torch.cuda.Event(enable_timing=True)
+            start_event.record(stream)
+            if name not in cls._events:
+                cls._events[name] = []
+            cls._events[name].append((start_event, end_event, stream))
+
+    @classmethod
+    def stop(cls, name: Optional[str] = None, use_cpu: bool = False) -> None:
+        if not name:
+            name = "default"
+
+        if not cls._enable_cuda or use_cpu:
+            assert (
+                name in cls._cpu_start_times
+            ), f"No CPU start time found for {name}, did you run start()?"
+            start_time = cls._cpu_start_times[name].pop()
+            elapsed_time_ms = (time.perf_counter() - start_time) * 1000
+            with cls._lock:
+                cls._times[name].append(elapsed_time_ms)
+
+        if name in cls._events and cls._events[name]:
+            start_event, end_event, stream = cls._events[name].pop()
+            end_event.record(stream)
+
+            # We create a separate thread to poll on the event status
+            # to avoid blocking the main thread.
+            thread = threading.Thread(
+                target=cls._check_event_completion, args=(name, start_event, end_event)
+            )
+            thread.start()
+            if name not in cls._threads:
+                cls._threads[name] = []
+            cls._threads[name].append(thread)
+
+    @classmethod
+    def _check_event_completion(
+        cls, name: str, start_event: torch.cuda.Event, end_event: torch.cuda.Event
+    ) -> None:
+        while True:
+            if end_event.query():
+                with cls._lock:
+                    cuda_time = start_event.elapsed_time(end_event)
+                    cls._times[name].append(cuda_time)
+                break
+            time.sleep(0.01)
+
+    @classmethod
+    def reset(cls) -> None:
+        """Clear all timing data."""
+        with cls._lock:
+            cls._times = {}
+            cls._threads = {}
+
+    @classmethod
+    def summary(cls) -> Dict[str, Dict[str, float]]:
+        """
+        Get summary of all timing data.
+        Returns:
+            Dict containing timing statistics for each section
+        """
+        # Wait for all in-flight measurements to complete
+        for _, threads in cls._threads.items():
+            for thread in threads:
+                thread.join()
+        with cls._lock:
+            result = {}
+            for name, times in cls._times.items():
+                if not times:
+                    continue
+                result[name] = {
+                    "count": len(times),
+                    "mean_ms": sum(times) / len(times),
+                    "total_ms": sum(times),
+                    "min_ms": min(times),
+                    "max_ms": max(times),
+                }
+            return result
+
+    @classmethod
+    def get_latest_measurement(cls, name: Optional[str] = None) -> float:
+        """Get the latest measurement (in ms) for a given section."""
+        if not name:
+            name = "default"
+        if name in cls._threads:
+            for thread in cls._threads[name]:
+                thread.join()
+            cls._threads[name] = []
+        with cls._lock:
+            if name not in cls._times or not cls._times[name]:
+                logging.warning(f"Section {name} not found in timing data.")
+                return 0.0
+            return cls._times[name][-1]
+
+
+def execution_timer_start(name: Optional[str] = None, use_cpu: bool = False) -> None:
+    """Start the ExecutionTimer."""
+    ExecutionTimer.start(name=name, use_cpu=use_cpu)
+
+
+def execution_timer_stop(name: Optional[str] = None, use_cpu: bool = False) -> None:
+    """Stop the ExecutionTimer."""
+    ExecutionTimer.stop(name=name, use_cpu=use_cpu)
+
+
+def get_execution_timer_average_ms(name: str = "default") -> torch.Tensor:
+    """Get the ExecutionTimer results."""
+    return torch.tensor(ExecutionTimer.summary()[name]["mean_ms"], dtype=torch.float64)
+
+
+def get_latest_timer_measurement(name: Optional[str] = None) -> torch.Tensor:
+    """Get the latest ExecutionTimer results."""
+    return torch.tensor(
+        ExecutionTimer.get_latest_measurement(name), dtype=torch.float64
+    )
+
+
+def get_execution_timer_summary() -> Dict[str, Dict[str, float]]:
+    """Get the ExecutionTimer summary."""
+    return ExecutionTimer.summary()
+
+
+def get_time_perfcounter() -> torch.Tensor:
+    """Get the time performance counter. Should be used only for debugging."""
+    return torch.tensor(time.perf_counter(), dtype=torch.float64)

monarch/timer/execution_timer_test.py

@@ -0,0 +1,131 @@
+# 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.
+
+"""Test suite for ExecutionTimer class."""
+
+# pyre-strict
+
+import time
+import unittest
+from unittest.mock import MagicMock, patch
+
+from monarch.timer.execution_timer import ExecutionTimer
+
+
+class TestExecutionTimer(unittest.TestCase):
+    """Test suite for the ExecutionTimer class."""
+
+    def setUp(self) -> None:
+        """Reset the profiler state before each test."""
+        ExecutionTimer.reset()
+
+    def test_basic_timing(self) -> None:
+        """Test basic CUDA timing functionality."""
+        with ExecutionTimer.time("test_section"):
+            time.sleep(0.01)  # Sleep for 10ms
+
+        # Get the stats
+        stats = ExecutionTimer.summary()
+
+        # Check that our section exists
+        self.assertIn("test_section", stats)
+
+        # Check timing (should be at least 10ms, allow some overhead)
+        section_stats = stats["test_section"]
+        self.assertEqual(section_stats["count"], 1)
+        self.assertGreaterEqual(section_stats["mean_ms"], 10)  # At least 10ms
+        self.assertLess(section_stats["mean_ms"], 50)  # Reasonable upper bound
+
+    def test_multiple_timing_same_section(self) -> None:
+        """Test timing the same section multiple times."""
+        for _ in range(5):
+            with ExecutionTimer.time("repeated_section"):
+                time.sleep(0.01)
+
+        stats = ExecutionTimer.summary()
+        self.assertIn("repeated_section", stats)
+
+        section_stats = stats["repeated_section"]
+        self.assertEqual(section_stats["count"], 5)
+        self.assertGreaterEqual(section_stats["mean_ms"], 10)
+        self.assertGreaterEqual(section_stats["total_ms"], 50)
+
+    def test_reset(self) -> None:
+        """Test that reset clears all timing data."""
+        with ExecutionTimer.time("before_reset"):
+            time.sleep(0.01)
+
+        # Verify the section exists
+        stats_before = ExecutionTimer.summary()
+        self.assertIn("before_reset", stats_before)
+
+        # Reset the profiler
+        ExecutionTimer.reset()
+
+        # Timing section should be gone
+        stats_after = ExecutionTimer.summary()
+        self.assertNotIn("before_reset", stats_after)
+
+    @patch("torch.cuda.is_available")
+    @patch("torch.cuda.Event")
+    def test_cuda_mocked(
+        self, mock_event: MagicMock, mock_is_available: MagicMock
+    ) -> None:
+        """Test CUDA timing with mocked CUDA functions."""
+
+        mock_is_available.return_value = True
+
+        mock_event_instance = MagicMock()
+        mock_event_instance.elapsed_time.return_value = 15.0  # 15ms
+        mock_event.return_value = mock_event_instance
+
+        with ExecutionTimer.time("mocked_cuda"):
+            time.sleep(0.01)
+
+        stats = ExecutionTimer.summary()
+
+        # Should have CUDA timings
+        self.assertIn("mocked_cuda", stats)
+        self.assertEqual(stats["mocked_cuda"]["mean_ms"], 15.0)
+
+    def test_get_latest_measurement(self) -> None:
+        """Test get_latest_measurement."""
+        with ExecutionTimer.time("latest_measurement_test"):
+            time.sleep(0.01)  # Sleep for 10ms
+
+        # Get the latest measurement
+        latest_measurement = ExecutionTimer.get_latest_measurement(
+            "latest_measurement_test"
+        )
+
+        self.assertGreaterEqual(latest_measurement, 5)
+        self.assertLess(latest_measurement, 50)  # Reasonable upper bound
+
+        # Test for a non-existent section
+        non_existent_measurement = ExecutionTimer.get_latest_measurement(
+            "non_existent_section"
+        )
+        self.assertEqual(non_existent_measurement, 0.0)
+
+    def test_cpu_timing(self) -> None:
+        """Test CPU timing functionality."""
+        with ExecutionTimer.time("cpu_section", use_cpu=True):
+            time.sleep(0.01)  # Sleep for 10ms
+        # Get the stats
+        stats = ExecutionTimer.summary()
+
+        # Check that our section exists
+        self.assertIn("cpu_section", stats)
+
+        # Check timing (should be at least 10ms, allow some overhead)
+        section_stats = stats["cpu_section"]
+        self.assertEqual(section_stats["count"], 1)
+        self.assertGreaterEqual(section_stats["mean_ms"], 10)  # At least 10ms
+        self.assertLess(section_stats["mean_ms"], 50)  # Reasonable upper bound
+
+
+if __name__ == "__main__":
+    unittest.main()

monarch/tools/__init__.py

@@ -0,0 +1,7 @@
+# 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-strict