torchmonarch-nightly 2025.6.27__cp312-cp312-manylinux2014_x86_64.whl

monarch/profiler.py

@@ -0,0 +1,160 @@
+# 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 itertools
+import os
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from typing import Any, Dict, NamedTuple, Optional, Tuple
+
+import torch
+from monarch.common.remote import remote
+from monarch.remote_class import ControllerRemoteClass, WorkerRemoteClass
+
+
+class Schedule(NamedTuple):
+    wait: int
+    warmup: int
+    active: int
+    repeat: int = 0
+    skip_first: int = 0
+
+
+class profile:
+    """
+    The class wraps `torch.profiler.profile()` to allow invoking the profiler remotely.
+    There are two main differences:
+    1) `on_trace_ready` can only be a string, indicating the folder where the traces
+        will be saved.
+    2) `schedule` must be of type `monarch.profiler.Schedule`.
+    """
+
+    PATH_KEY = "on_trace_ready"
+    _counter = itertools.count()
+
+    def __init__(self, *args, **kwargs) -> None:
+        assert isinstance(kwargs.get(self.PATH_KEY, None), str), (
+            f"{self.PATH_KEY} must be passed and must be a string to represent the "
+            "path to save the profiler."
+        )
+        schedule = kwargs.get("schedule", None)
+        assert (
+            isinstance(schedule, Schedule) or schedule is None
+        ), "schedule can only be monarch.profiler.Schedule or None."
+        self.id = next(self._counter)
+        _profiler_controller_init(self.id, *args, **kwargs)
+
+    def __enter__(self) -> "profile":
+        _profiler_controller_enter(self.id)
+        return self
+
+    def __exit__(self, *args, **kwargs) -> None:
+        _profiler_controller_exit(self.id)
+
+    def step(self) -> None:
+        _profiler_controller_step(self.id)
+
+
+@dataclass
+class _Profiler:
+    args: Tuple[Any, ...]
+    kwargs: Dict[str, Any]
+    profiler: Optional[torch.profiler.profile] = None
+
+
+_profilers: Dict[int, _Profiler] = {}
+
+
+def _profiler_init(ident, *args, **kwargs) -> None:
+    global _profilers
+    assert (
+        ident not in _profilers
+    ), f"Initializing an already existing profiler, {ident=}"
+    _profilers[ident] = _Profiler(args, kwargs)
+    # It's unclear why we cannot create the profiler here. Even though
+    # the thread is the same, profiler complains thread id mismatch.
+
+
+def _profiler_enter(ident, *args, **kwargs) -> None:
+    def on_trace_ready(prof, dir_path):
+        dir_path = Path(dir_path).absolute()
+        os.makedirs(dir_path, exist_ok=True)
+        # This is not a synchronized call, so it is okay to call without
+        # device mesh.
+        rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+        prof.export_chrome_trace(f"{dir_path}/trace_{rank}.json")
+
+    profiler = _profilers[ident]
+    profiler.kwargs[profile.PATH_KEY] = partial(
+        on_trace_ready, dir_path=profiler.kwargs[profile.PATH_KEY]
+    )
+    schedule = profiler.kwargs.get("schedule", None)
+    if schedule is not None:
+        profiler.kwargs["schedule"] = torch.profiler.schedule(**schedule._asdict())
+    profiler.profiler = torch.profiler.profile(*profiler.args, **profiler.kwargs)
+
+    profiler.profiler.__enter__()
+
+
+def _profiler_exit(ident, *args, **kwargs) -> None:
+    profiler = _profilers[ident].profiler
+    assert profiler is not None
+    profiler.__exit__(None, None, None)
+    _profilers.pop(ident)
+
+
+def _profiler_step(ident, *args, **kwargs) -> None:
+    profiler = _profilers[ident].profiler
+    assert profiler is not None
+    profiler.step()
+
+
+_profiler_controller_init = remote(
+    "monarch.profiler._profiler_init", propagate="inspect"
+)
+
+_profiler_controller_enter = remote(
+    "monarch.profiler._profiler_enter", propagate="inspect"
+)
+
+_profiler_controller_exit = remote(
+    "monarch.profiler._profiler_exit", propagate="inspect"
+)
+
+_profiler_controller_step = remote(
+    "monarch.profiler._profiler_step", propagate="inspect"
+)
+
+
+class record_function(ControllerRemoteClass):
+    """
+    The class wraps `torch.profiler.record_function()` to allow invoking the
+    record_function remotely.
+    """
+
+    def __init__(self, name: str, args: Optional[str] = None) -> None:
+        super().__init__("monarch.profiler.WorkerRecordFunction", name, args)
+
+    @ControllerRemoteClass.remote_method
+    def __enter__(self) -> "record_function":
+        return self
+
+    @ControllerRemoteClass.remote_method
+    def __exit__(self, *args, **kwargs) -> None:
+        return
+
+
+class WorkerRecordFunction(WorkerRemoteClass):
+    def __init__(self, *args, **kwargs) -> None:
+        self._record_function = torch.profiler.record_function(*args, **kwargs)
+
+    def __enter__(self) -> None:
+        self._record_function.__enter__()
+
+    def __exit__(self, *args, **kwargs) -> None:
+        self._record_function.__exit__(*args, **kwargs)

monarch/python_local_mesh.py

@@ -0,0 +1,107 @@
+# 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 os
+import subprocess
+from time import sleep
+from typing import Optional, TYPE_CHECKING
+
+import monarch_supervisor
+from monarch.common._device_utils import _local_device_count
+from monarch.common.fake import fake_call
+from monarch.common.invocation import DeviceException, RemoteException
+from monarch.world_mesh import world_mesh
+from monarch_supervisor import Context, HostConnected
+from monarch_supervisor.python_executable import PYTHON_EXECUTABLE
+
+if TYPE_CHECKING:
+    from monarch.common.device_mesh import DeviceMesh
+
+
+class PythonLocalContext:
+    def __init__(self, N: int):
+        # do a fake call to instantiate ThreadPoolExecutor so we don't block GIL later
+        fake_call(lambda: 0)
+
+        self.ctx = ctx = Context()
+        ctx.request_hosts(N)
+
+        # we want ctx to start its listener threads
+        # before creating the hosts because
+        # initialization will happen faster in this case
+        sleep(0)
+        supervisor_addr = f"tcp://127.0.0.1:{ctx.port}"
+
+        env = {
+            **os.environ,
+            "TORCH_SUPERVISOR_HEARTBEAT_INTERVAL": str(
+                monarch_supervisor.HEARTBEAT_INTERVAL
+            ),
+            # This is needed to avoid a hard failure in ncclx when we do not
+            # have backend topology info (eg. on RE).
+            "NCCL_IGNORE_TOPO_LOAD_FAILURE": "true",
+        }
+
+        # start_new_session=True, because we want the host managers to be able to kill
+        # any worker processes before they exit, even if the supervisor crashes, or we ctrl-c
+        # it in testing.
+        self.host_managers = [
+            subprocess.Popen(
+                [
+                    PYTHON_EXECUTABLE,
+                    "-m",
+                    "monarch_supervisor.host",
+                    supervisor_addr,
+                ],
+                env=env,
+                start_new_session=True,
+            )
+            for _ in range(N)
+        ]
+        connections = ctx.messagefilter(HostConnected)
+        self.hosts = [connections.recv(timeout=30).sender for _ in range(N)]
+
+    def shutdown(self):
+        self.ctx.shutdown()
+        for host_manager in self.host_managers:
+            host_manager.wait(timeout=10)
+
+
+def python_local_mesh(*, gpus: Optional[int] = None, hosts: int = 1) -> "DeviceMesh":
+    """
+    Creates a local device mesh with the given number of hosts and gpus per host.
+    Easy way to use PythonLocalContext.
+
+    Args:
+        gpus (Optional[int]): number of gpus per host.
+                              Default: the number of GPUs this machine has.
+
+        hosts (int): number of hosts, primarily used for simulating multiple machines locally.
+                     Default: 1
+
+    Example::
+        local_mesh = python_local_mesh(gpus=2)
+        with local_mesh.activate():
+            x = torch.rand(3, 4)
+            local_tensor = fetch_shard(x).result()
+
+        # Cleanly shut down the local mesh and exit.
+        local_mesh.exit()
+    """
+    ctx = PythonLocalContext(hosts)
+    if gpus is None:
+        gpus = _local_device_count()
+    dm = world_mesh(ctx.ctx, ctx.hosts, gpus)
+
+    def exit(
+        error: Optional[RemoteException | DeviceException | Exception] = None,
+    ) -> None:
+        dm.client.shutdown(True, error)
+        ctx.shutdown()
+
+    dm.exit = exit
+    return dm

monarch/random.py

@@ -0,0 +1,61 @@
+# 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.
+
+import os
+from typing import NamedTuple, Tuple
+
+import torch
+from monarch.common.remote import remote
+from monarch.common.tensor import Tensor
+
+
+class State(NamedTuple):
+    cpu: Tensor
+    cuda: Tensor
+
+
+@remote(
+    propagate=lambda: (
+        torch.empty(5056, dtype=torch.uint8),
+        torch.empty(16, dtype=torch.uint8),
+    )
+)
+def _get_state() -> Tuple[torch.Tensor, torch.Tensor]:
+    return (torch.get_rng_state(), torch.cuda.get_rng_state())
+
+
+@remote(propagate=lambda state: None)
+def set_state(state: Tuple[Tensor, Tensor]):
+    cpu, device = state
+    torch.set_rng_state(cpu)
+    torch.cuda.set_rng_state(device)
+
+
+@remote(propagate=lambda _: None)
+def _manual_seed(seed: torch.Tensor):
+    torch.manual_seed(seed.item())
+
+
+@remote(propagate=lambda: None)
+def make_deterministic():
+    torch.use_deterministic_algorithms(True)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    # env var for deterministic CuBLAS
+    # https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+
+def get_state() -> State:
+    return State(*_get_state())
+
+
+def new_state(seed: Tensor) -> State:
+    orig = get_state()
+    _manual_seed(seed)
+    mine = get_state()
+    set_state(orig)
+    return mine

monarch/rdma.py

@@ -0,0 +1,162 @@
+# 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.
+
+import ctypes
+
+from dataclasses import dataclass
+from typing import cast, Dict, Optional, Tuple
+
+import torch
+
+from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
+
+from monarch.actor_mesh import (
+    _ActorMeshRefImpl,
+    Actor,
+    ActorMeshRef,
+    endpoint,
+    MonarchContext,
+)
+
+
+@dataclass
+class LocalRDMARecord:
+    data: torch.Tensor
+
+
+_local_buffers: Dict[int, "LocalRDMARecord"] = {}
+
+
+def _get_bytes(storage: torch.Tensor, offset: int, size: int) -> bytearray:
+    """Extracts a bytearray from a 1D, 1byte per item tensor."""
+    if offset + size > storage.numel():
+        raise ValueError(f"Read out of range: {offset + size} > {storage.size()}")
+    addr = storage.data_ptr()
+    if storage.device.type != "cpu":
+        result = bytearray(size)
+        result_tensor = torch.frombuffer(
+            result,
+            dtype=torch.uint8,
+        )
+        source_tensor = storage[offset:]
+        result_tensor.copy_(source_tensor)
+    else:
+        ctypes_array = (ctypes.c_byte * size).from_address(addr)
+        result = bytearray(ctypes_array)
+    return result
+
+
+class RDMAManager(Actor):
+    @staticmethod
+    def on_proc(proc_id: str) -> "RDMAManager":
+        ctx = MonarchContext.get()
+        return cast(
+            RDMAManager,
+            ActorMeshRef(
+                RDMAManager,
+                _ActorMeshRefImpl.from_actor_id(
+                    ctx.mailbox,
+                    ActorId.from_string(f"{proc_id}.rdma_manager[0]"),
+                ),
+                ctx.mailbox,
+            ),
+        )
+
+    @endpoint
+    async def drop(self, addr: int) -> None:
+        if addr in _local_buffers:
+            del _local_buffers[addr]
+
+    @endpoint
+    async def fetch(self, addr: int, offset: int, nbytes: int) -> bytearray:
+        if addr not in _local_buffers:
+            raise ValueError(f"Unknown buffer {addr}")
+        storage = _local_buffers[addr].data
+        return _get_bytes(storage, offset, nbytes)
+
+    @endpoint
+    async def put(self, addr: int, offset: int, bytes: bytearray) -> None:
+        if addr not in _local_buffers:
+            raise ValueError(f"Unknown buffer {addr}")
+        storage = _local_buffers[addr].data
+        storage[offset : offset + len(bytes)] = torch.frombuffer(
+            bytes, dtype=storage.dtype
+        )
+
+
+def _assert_tensor_is_1d_contiguous_uint8(t: torch.Tensor) -> None:
+    if t.ndim != 1:
+        raise ValueError(f"Tensor must be 1D, got {t.ndim}D")
+    if t.dtype != torch.uint8:
+        raise ValueError(f"Tensor must be uint8, got {t.dtype}")
+    if not t.is_contiguous():
+        raise ValueError("Tensor must be contiguous")
+
+
+class RDMABuffer:
+    def __init__(self, data: torch.Tensor) -> None:
+        """
+        RDMABuffer only supports 1D contiguous tensors that are 1 byte per item.
+
+        To create a 1 byte, 1D view, use t.view(torch.uint8).flatten()
+
+        TODO: Create TensorBuffer, which will be main user API supporting non-contiguous , multi-byte-per-elment tensors
+        """
+        _assert_tensor_is_1d_contiguous_uint8(data)
+        assert data.storage_offset() == 0
+        storage = data.untyped_storage()
+        self.addr: int = storage.data_ptr()
+        self.begin = 0
+        self.end: int = storage.size()
+        self.proc_id: str = MonarchContext.get().proc_id
+        self.local_data: object = None
+        _local_buffers[self.addr] = LocalRDMARecord(data)
+
+    def drop(self) -> None:
+        if self.proc_id is None:
+            del _local_buffers[self.addr]
+            return
+        rmda_actor = RDMAManager.on_proc(self.proc_id)
+        # pyre-ignore[16]: Undefined attribute [16]: `Endpoint` has no attribute `cast`.
+        rmda_actor.drop.cast(self.addr)
+
+    def __getstate__(self) -> Tuple[int, int, int, Optional[str]]:
+        proc_id = self.proc_id
+        # locally created RDMABuffer being set remotely,
+        # record its proc_id so we know how to establish connections to it
+        if proc_id is None:
+            proc_id = MonarchContext.get().proc_id
+        return (self.addr, self.begin, self.end, proc_id)
+
+    def __setstate__(self, state: Tuple[int, int, int, str]) -> None:
+        self.local_data = None
+        self.addr, self.begin, self.end, self.proc_id = state
+
+    async def read_into(self, dst: torch.Tensor, offset: int = 0) -> None:
+        """
+        Read data from the RDMABuffer into a destination tensor.
+
+        The destination tensor must be contiguous and 1 byte per item.
+        """
+        _assert_tensor_is_1d_contiguous_uint8(dst)
+        bytes = await RDMAManager.on_proc(self.proc_id).fetch.call_one(
+            self.addr, offset, dst.numel()
+        )
+        dst.copy_(torch.frombuffer(bytes, dtype=torch.uint8))
+
+    async def write(self, src: torch.Tensor, offset: int = 0) -> None:
+        """
+        Write data from a source tensor into the RDMABuffer.
+
+        The source tensor must be contiguous and 1 byte per item.
+        """
+        _assert_tensor_is_1d_contiguous_uint8(src)
+        bytes = _get_bytes(
+            src,
+            cast(int, src.storage_offset()),
+            src.numel(),
+        )
+        await RDMAManager.on_proc(self.proc_id).put.call_one(self.addr, offset, bytes)

monarch/remote_class.py

@@ -0,0 +1,114 @@
+# 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 importlib
+import itertools
+from typing import Any, Dict
+
+from monarch.common import device_mesh
+from monarch.common.remote import remote
+
+
+class ControllerRemoteClass:
+    """
+    This class simplifies the creation and management of remote classes.  It serves as
+    the controller side of a remote class architecture.  Classes that are intended to be
+    controlled remotely should inherit from this class. The constructor of the inheriting
+    class must invoke `super().__init__()` with the path to the remote class that will be
+    used on the worker nodes.  Methods that are intended for remote execution must be
+    decorated with `ControllerRemoteClass.remote_method`.
+
+    Note: This class is designed for use by the controller developer only and should
+    not be directly used in model code.
+
+    Example usage:
+
+    class ControllerMyClass(ControllerRemoteClass):
+        def __init__(self, *args, **kwargs) -> None:
+            super().__init__("my_package.my_class", *args, **kwargs)
+
+        @ControllerRemoteClass.remote_method
+        def some_method(self, *args, **kwargs) -> None:
+            # This method is intended for remote execution and does nothing locally.
+            pass
+    """
+
+    _counter = itertools.count()
+
+    def __init__(self, cls_path: str, *args, **kwargs) -> None:
+        self.ident = next(ControllerRemoteClass._counter)
+        self.cls_path = cls_path
+        self.mesh = device_mesh._active
+        _controller_remote_class_method(
+            cls_path, "remote_init", self.ident, *args, **kwargs
+        )
+
+    def __del__(self) -> None:
+        mesh = getattr(self, "mesh", None)
+        if mesh is not None and not mesh.client._shutdown:
+            with self.mesh.activate():
+                _controller_remote_class_method(
+                    self.cls_path,
+                    "remote_del",
+                    self.ident,
+                )
+
+    @staticmethod
+    def remote_method(fn):
+        def wrapper(self, *args, **kwargs) -> None:
+            _controller_remote_class_method(
+                self.cls_path, "remote_method", self.ident, fn.__name__, *args, **kwargs
+            )
+
+        return wrapper
+
+
+# Add the logic as a separate private function instead of adding ita to
+# ResolvableFunctionFromPath. This avoids users to using this directly.
+_controller_remote_class_method = remote(
+    "monarch.remote_class._remote_class_method", propagate="inspect"
+)
+
+
+def _remote_class_method(cls_path: str, method_name: str, *args, **kwargs) -> None:
+    modulename, classname = cls_path.rsplit(".", 1)
+    module = importlib.import_module(modulename)
+    cls = getattr(module, classname)
+    method = getattr(cls, method_name)
+    method(*args, **kwargs)
+
+
+class WorkerRemoteClass:
+    """
+    This class is designed to be used alongside ``ControllerRemoteClass`` and represents
+    the worker-side of a remote class architecture. Instances of this class should just
+    mimic standard Python classes, with the notable exception that all methods must
+    return None -- the current RemoteClass architecture does not support methods that
+    return values.
+
+    The `ident` attribute is used for tracking object instances created via `remote_init`.
+    This tracking is necessary because the remote function would otherwise lose the
+    reference to the object.
+    """
+
+    _objects: Dict[int, Any] = {}
+
+    @classmethod
+    def remote_init(cls, ident: int, *args, **kwargs) -> None:
+        WorkerRemoteClass._objects[ident] = cls(*args, **kwargs)
+
+    @classmethod
+    def remote_del(cls, ident) -> None:
+        WorkerRemoteClass._objects.pop(ident)
+
+    @classmethod
+    def remote_method(cls, ident: int, method_name, *args, **kwargs) -> None:
+        instance = WorkerRemoteClass._objects[ident]
+        assert (
+            cls == instance.__class__
+        ), "Mismatched class type {cls} {instance.__class__}"
+        getattr(instance, method_name)(*args, **kwargs)