torchmonarch-nightly 2025.6.27__cp311-cp311-manylinux2014_x86_64.whl

monarch/future.py

@@ -0,0 +1,76 @@
+# 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 asyncio
+from functools import partial
+from typing import Generator, Generic, Optional, TypeVar
+
+R = TypeVar("R")
+
+
+def _incomplete(impl, self):
+    try:
+        return self._set_result(impl())
+    except Exception as e:
+        self._set_exception(e)
+        raise
+
+
+async def _aincomplete(impl, self):
+    try:
+        return self._set_result(await impl())
+    except Exception as e:
+        self._set_exception(e)
+        raise
+
+
+# TODO: consolidate with monarch.common.future
+class ActorFuture(Generic[R]):
+    def __init__(self, impl, blocking_impl=None):
+        if blocking_impl is None:
+            blocking_impl = partial(asyncio.run, impl())
+        self._get = partial(_incomplete, blocking_impl)
+        self._aget = partial(_aincomplete, impl)
+
+    def get(self, timeout: Optional[float] = None) -> R:
+        if timeout is not None:
+            return asyncio.run(asyncio.wait_for(self._aget(self), timeout))
+        return self._get(self)
+
+    def __await__(self) -> Generator[R, None, R]:
+        return self._aget(self).__await__()
+
+    def _set_result(self, result):
+        def f(self):
+            return result
+
+        async def af(self):
+            return result
+
+        self._get, self._aget = f, af
+        return result
+
+    def _set_exception(self, e):
+        def f(self):
+            raise e
+
+        async def af(self):
+            raise e
+
+        self._get, self._aget = f, af
+
+    # compatibility with old tensor engine Future objects
+    # hopefully we do not need done(), add_callback because
+    # they are harder to implement right.
+    def result(self, timeout: Optional[float] = None) -> R:
+        return self.get(timeout)
+
+    def exception(self, timeout: Optional[float] = None):
+        try:
+            self.get(timeout)
+            return None
+        except Exception as e:
+            return e

monarch/gradient/__init__.py

@@ -0,0 +1,11 @@
+# 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
+
+from ._gradient_generator import GradientGenerator
+
+__all__ = ["GradientGenerator"]

monarch/gradient/_gradient_generator.pyi

@@ -0,0 +1,22 @@
+# 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 typing import Any, Optional
+
+import torch
+
+class GradientGenerator:
+    def __init__(
+        self,
+        roots_list: Any,
+        with_respect_to: Any,
+        grad_roots: Any,
+        context_restorer: Any,
+    ): ...
+    # pyre-ignore[11]: Annotation `torch.Tensor` is not defined as a type.
+    def __next__(self) -> Optional[torch.Tensor]: ...
+    def __iter__(self) -> "GradientGenerator": ...

monarch/gradient_generator.py

@@ -0,0 +1,185 @@
+# 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 math
+from contextlib import nullcontext
+from functools import partial
+from types import CellType, FunctionType
+from typing import Any, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.autograd.graph
+
+from monarch.common import device_mesh, stream
+from monarch.common.tensor import Tensor
+from monarch.common.tree import flatten
+from monarch.gradient import GradientGenerator as _GradientGenerator
+from torch._C._autograd import _get_sequence_nr  # @manual
+from torch.autograd.graph import get_gradient_edge, GradientEdge
+
+TensorOrEdge = Union[torch.Tensor, GradientEdge]
+
+
+class Context(NamedTuple):
+    device_mesh: "Optional[device_mesh.DeviceMesh]"
+    stream: "stream.Stream"
+
+    def enable(self):
+        if device_mesh is None:
+            activate_mesh = device_mesh.no_mesh.activate()
+        elif self.device_mesh is not device_mesh._active:
+            # XXX: something about activating device meshes from this object
+            # doesn't work correctly and somehow inactivates the device mesh
+            # if it is already enabled. This is a temporary workaround for
+            # the demo.
+            activate_mesh = self.device_mesh.activate()
+        else:
+            activate_mesh = nullcontext()
+        with activate_mesh, self.stream.activate(), torch.no_grad():
+            yield
+
+
+_sequence_nr_to_context: Dict[int, Context] = {}
+_sequence_nr_end = 0
+
+
+def restore_context(t: Optional[Tensor], sn: Optional[int], last: bool):
+    if sn is not None:
+        _update_context_map(Context(device_mesh._active, stream._active))
+        ctx = _sequence_nr_to_context.pop(sn) if last else _sequence_nr_to_context[sn]
+        return ctx.enable()
+    if t is not None:
+        return Context(t.mesh, t.stream).enable()
+    return Context(device_mesh._active, stream._active).enable()
+
+
+def _update_context_map(ctx: Context):
+    global _sequence_nr_end
+    next_sequence_nr = _get_sequence_nr()
+    for i in range(_sequence_nr_end, next_sequence_nr):
+        _sequence_nr_to_context[i] = ctx
+    _sequence_nr_end = _get_sequence_nr()
+
+
+device_mesh._on_change.append(
+    lambda old, mesh: _update_context_map(Context(old, stream._active))
+)
+stream._on_change.append(
+    lambda old, stream: _update_context_map(Context(device_mesh._active, old))
+)
+
+
+def grad_generator(
+    roots: Union[torch.Tensor, Sequence[TensorOrEdge]] = (),
+    with_respect_to: Sequence[TensorOrEdge] = (),
+    grad_roots: Sequence[Optional[torch.Tensor]] = (),
+):
+    if isinstance(roots, torch.Tensor):
+        roots = [roots]
+    return _GradientGenerator(
+        list(roots), list(with_respect_to), list(grad_roots), restore_context
+    )
+
+
+def _gradient_edge(a: TensorOrEdge) -> GradientEdge:
+    if isinstance(a, GradientEdge):
+        return a
+    return get_gradient_edge(a)
+
+
+class GradGenerator:
+    def __init__(self):
+        self.roots: List[torch.autograd.graph.GradientEdge] = []
+        self.with_respect_to: List[torch.autograd.graph.GradientEdge] = []
+        self.grad_roots: List[Optional[torch.Tensor]] = []
+        self.unflattens: List[Tuple[int, Any]] = []
+
+    def grad(self, tree: Any):
+        tensors, unflatten = flatten(tree, lambda x: isinstance(x, torch.Tensor))
+        self.unflattens.append((len(tensors), unflatten))
+        self.with_respect_to.extend(_gradient_edge(t) for t in tensors)
+
+    def root(self, r: TensorOrEdge, grad: Optional[torch.Tensor] = None):
+        self.roots.append(_gradient_edge(r))
+        self.grad_roots.append(grad)
+
+    def __iter__(self):
+        gi = _GradientGenerator(
+            self.roots,
+            list(reversed(self.with_respect_to)),
+            self.grad_roots,
+            restore_context,
+        )
+        for n, unflatten in reversed(self.unflattens):
+            yield unflatten(reversed([next(gi) for _ in range(n)]))
+
+
+class GradFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, fn, *args, **kwargs):
+        result, backward_continuation = fn(*args, **kwargs)
+        ctx.backward_continuation = backward_continuation
+        values = []
+        if backward_continuation.__closure__ is not None:
+            for cell in backward_continuation.__closure__:
+                values.append(cell.cell_contents)
+                cell.cell_contents = None
+        tensors, ctx.unflatten = flatten(values, lambda x: isinstance(x, torch.Tensor))
+        ctx.save_for_backward(*tensors)
+        return result
+
+    @staticmethod
+    def backward(ctx, *args, **kwargs):
+        closure = tuple(CellType(v) for v in ctx.unflatten(ctx.saved_tensors))
+        orig = ctx.backward_continuation
+        fn = FunctionType(
+            orig.__code__, orig.__globals__, orig.__name__, orig.__defaults__, closure
+        )
+        output = fn(*args, **kwargs)
+        if isinstance(output, tuple):
+            return None, *output
+        else:
+            return None, output
+
+
+def grad_function(fn):
+    return partial(GradFunction.apply, fn)
+
+
+def gradient_execution_order(
+    roots: Sequence[TensorOrEdge], with_respect_to: Sequence[TensorOrEdge]
+) -> List[int]:
+    """
+    Returns the order in which the gradients for `with_respect_to` would become available
+    if autograd were run on `roots`. This is the reverse order of each tensors
+    first use in the gradient computation.
+    """
+    with_respect_to = [_gradient_edge(g) for g in with_respect_to]
+    min_sequence_nr: Dict[Any, float] = {e: math.inf for e in with_respect_to}
+
+    to_scan = [_gradient_edge(r).node for r in roots]
+    scanned = set()
+    while to_scan:
+        node = to_scan.pop()
+        if node in scanned:
+            continue
+        scanned.add(node)
+        for key in node.next_functions:
+            nnode = key[0]
+            if nnode is None:
+                continue
+            to_scan.append(nnode)
+            value = min_sequence_nr.get(key)
+            if value is not None:
+                # pyre-ignore
+                min_sequence_nr[key] = min(node._sequence_nr(), value)
+
+    return sorted(
+        range(len(with_respect_to)),
+        key=lambda i: min_sequence_nr[with_respect_to[i]],
+        reverse=True,
+    )

monarch/memory.py

@@ -0,0 +1,43 @@
+# 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 pathlib import Path
+
+import torch
+from monarch.common.remote import remote
+
+
+PATH_KEY = "dir_snapshots"
+_counter = itertools.count()
+
+
+@remote(propagate="inspect")
+def record_memory_history() -> None:
+    torch.cuda.memory._record_memory_history()
+
+
+def dump_memory_snapshot(*args, **kwargs) -> None:
+    """
+    This function wraps torch.cuda.memory._dump_snapshot() to dump memory snapshot remotely.
+    """
+    assert isinstance(
+        kwargs.get(PATH_KEY, None), str
+    ), f"{PATH_KEY} must be passed and must be a string to represent the path to save the memory snapshots."
+    id = next(_counter)
+    _memory_controller_dump(id, *args, **kwargs)
+
+
+@remote(propagate="inspect")
+def _memory_controller_dump(ident, *args, **kwargs) -> None:
+    dir_path = Path(kwargs[PATH_KEY]).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
+    snapshot_path = f"{dir_path}/snapshot_{rank}.pickle"
+    torch.cuda.memory._dump_snapshot(filename=snapshot_path)

monarch/mesh_controller.py

@@ -0,0 +1,271 @@
+# 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 atexit
+import logging
+import os
+import time
+import traceback
+from collections import deque
+from logging import Logger
+from typing import List, NamedTuple, Optional, TYPE_CHECKING, Union
+
+import torch.utils._python_dispatch
+
+from monarch import NDSlice
+from monarch._rust_bindings.monarch_extension import client, debugger
+from monarch._rust_bindings.monarch_extension.client import (  # @manual=//monarch/monarch_extension:monarch_extension
+    WorldState,
+)
+from monarch._rust_bindings.monarch_extension.mesh_controller import _Controller
+from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ActorId,
+)
+
+if TYPE_CHECKING:
+    from monarch._rust_bindings.monarch_hyperactor.proc_mesh import (
+        ProcMesh as HyProcMesh,
+    )
+    from monarch.proc_mesh import ProcMesh
+
+from monarch._rust_bindings.monarch_hyperactor.shape import Point
+
+from monarch._rust_bindings.monarch_messages.debugger import DebuggerAction
+from monarch.common.client import Client
+from monarch.common.controller_api import LogMessage, MessageResult
+from monarch.common.device_mesh import DeviceMesh, no_mesh
+from monarch.common.invocation import DeviceException, RemoteException
+from monarch.controller.debugger import read as debugger_read, write as debugger_write
+from monarch.rust_local_mesh import _get_worker_exec_info
+from pyre_extensions import none_throws
+
+logger: Logger = logging.getLogger(__name__)
+
+
+class Controller(_Controller):
+    def __init__(self, workers: "HyProcMesh") -> None:
+        super().__init__()
+        # Buffer for messages unrelated to debugging that are received while a
+        # debugger session is active.
+        self._non_debugger_pending_messages: deque[
+            Optional[client.LogMessage | client.WorkerResponse]
+        ] = deque()
+        self._pending_debugger_sessions: deque[ActorId] = deque()
+
+    def next_message(
+        self, timeout: Optional[float]
+    ) -> Optional[LogMessage | MessageResult]:
+        if self._non_debugger_pending_messages:
+            msg = self._non_debugger_pending_messages.popleft()
+        else:
+            msg = self._get_next_message(timeout_msec=int((timeout or 0.0) * 1000.0))
+        if msg is None:
+            return None
+
+        if isinstance(msg, client.WorkerResponse):
+            return _worker_response_to_result(msg)
+        elif isinstance(msg, client.LogMessage):
+            return LogMessage(msg.level, msg.message)
+        elif isinstance(msg, client.DebuggerMessage):
+            self._run_debugger_loop(msg)
+
+    def send(
+        self,
+        ranks: Union[NDSlice, List[NDSlice]],
+        msg: NamedTuple,
+    ) -> None:
+        with torch.utils._python_dispatch._disable_current_modes():
+            return super().send(ranks, msg)
+
+    def drain_and_stop(
+        self,
+    ) -> List[LogMessage | MessageResult | client.DebuggerMessage]:
+        self._drain_and_stop()
+        return []
+
+    def _run_debugger_loop(self, message: client.DebuggerMessage) -> None:
+        if not isinstance(message.action, DebuggerAction.Paused):
+            raise RuntimeError(
+                f"Unexpected debugger message {message} when no debugger session is running"
+            )
+
+        self._pending_debugger_sessions.append(message.debugger_actor_id)
+        while self._pending_debugger_sessions:
+            debugger_actor_id = self._pending_debugger_sessions.popleft()
+            rank = debugger_actor_id.rank
+            proc_id = debugger_actor_id.proc_id
+            debugger_write(
+                f"pdb attached to proc {proc_id} with rank {rank}, debugger actor {debugger_actor_id} \n"
+            )
+
+            self._debugger_attach(debugger_actor_id)
+            while True:
+                # TODO: Add appropriate timeout.
+                msg = self._get_next_message(timeout_msec=None)
+
+                if not isinstance(msg, client.DebuggerMessage):
+                    self._non_debugger_pending_messages.append(msg)
+                    continue
+
+                if msg.debugger_actor_id != debugger_actor_id:
+                    if isinstance(msg.action, DebuggerAction.Paused):
+                        self._pending_debugger_sessions.append(msg.debugger_actor_id)
+                        continue
+                    else:
+                        raise RuntimeError(
+                            f"unexpected debugger message {msg} from rank {msg.debugger_actor_id.rank} "
+                            f"when debugging rank {debugger_actor_id.rank}"
+                        )
+
+                action = msg.action
+                if isinstance(action, DebuggerAction.Detach):
+                    break
+                elif isinstance(action, DebuggerAction.Read):
+                    self._debugger_write(
+                        debugger_actor_id, debugger_read(action.requested_size)
+                    )
+                elif isinstance(action, DebuggerAction.Write):
+                    debugger_write(
+                        debugger.get_bytes_from_write_action(action).decode()
+                    )
+                else:
+                    raise RuntimeError(
+                        f"unexpected debugger message {msg} when debugging rank {debugger_actor_id.rank}"
+                    )
+
+    def worker_world_state(self) -> WorldState:
+        raise NotImplementedError("worker world state")
+
+    def stop_mesh(self):
+        # I think this is a noop?
+
+        pass
+
+
+# TODO: Handling conversion of the response can move to a separate module over time
+# especially as we have structured error messages.
+def _worker_response_to_result(result: client.WorkerResponse) -> MessageResult:
+    if not result.is_exception():
+        # The result of the message needs to be unwrapped on a real device.
+        # Staying as a fake tensor will fail the tensor deserialization.
+        with no_mesh.activate():
+            return MessageResult(result.seq, result.result(), None)
+    exc = none_throws(result.exception())
+    if isinstance(exc, client.Error):
+        worker_frames = [
+            traceback.FrameSummary("<unknown>", None, frame)
+            for frame in exc.backtrace.split("\\n")
+        ]
+        return MessageResult(
+            seq=result.seq,
+            result=None,
+            error=RemoteException(
+                seq=exc.caused_by_seq,
+                exception=RuntimeError(exc.backtrace),
+                controller_frame_index=0,  # TODO: T225205291 fix this once we have recording support in rust
+                controller_frames=None,
+                worker_frames=worker_frames,
+                source_actor_id=exc.actor_id,
+                message=f"Remote function in {exc.actor_id} errored.",
+            ),
+        )
+    elif isinstance(exc, client.Failure):
+        frames = [
+            traceback.FrameSummary("<unknown>", None, frame)
+            for frame in exc.backtrace.split("\n")
+        ]
+        reason = f"Actor {exc.actor_id} crashed on {exc.address}, check the host log for details"
+        logger.error(reason)
+        return MessageResult(
+            seq=0,  # seq is not consumed for DeviceException; it will be directly thrown by the client
+            result=None,
+            error=DeviceException(
+                exception=RuntimeError(reason),
+                frames=frames,
+                source_actor_id=exc.actor_id,
+                message=reason,
+            ),
+        )
+    else:
+        raise RuntimeError(f"Unknown exception type: {type(exc)}")
+
+
+def _initialize_env(worker_point: Point, proc_id: str) -> None:
+    worker_rank = worker_point.rank
+    try:
+        _, worker_env = _get_worker_exec_info()
+        local_rank = worker_point["gpus"]
+        gpus_per_host = worker_point.size("gpus")
+        num_worker_procs = len(worker_point.shape)
+        process_env = {
+            **worker_env,
+            "HYPERACTOR_MANAGED_SUBPROCESS": "1",
+            "CUDA_VISIBLE_DEVICES": str(local_rank),
+            "NCCL_HOSTID": f"{proc_id}_host_{worker_rank // gpus_per_host}",
+            # 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",
+            "LOCAL_RANK": str(local_rank),
+            "RANK": str(worker_rank),
+            "WORLD_SIZE": str(num_worker_procs),
+            "LOCAL_WORLD_SIZE": str(gpus_per_host),
+        }
+        os.environ.update(process_env)
+    except Exception:
+        traceback.print_exc()
+        raise
+
+
+class MeshClient(Client):
+    def shutdown(
+        self,
+        destroy_pg: bool = True,
+        error_reason: Optional[RemoteException | DeviceException | Exception] = None,
+    ):
+        # return
+        if self.has_shutdown:
+            return
+        logger.info("shutting down the client gracefully")
+
+        atexit.unregister(self._atexit)
+        self._shutdown = True
+
+        # ensure all pending work is finished.
+        # all errors must be messaged back at this point
+        self.new_node_nocoalesce([], [], None, [])
+        self._request_status()
+
+        ttl = 60
+        start_time = time.time()
+        end_time = start_time + ttl
+        while ttl > 0 and self.last_assigned_seq > self.last_processed_seq:
+            ttl = end_time - time.time()
+            self.handle_next_message(ttl)
+            if self._pending_shutdown_error:
+                raise self._pending_shutdown_error
+
+        if ttl <= 0:
+            raise RuntimeError("shutdown timed out")
+
+        # we are not expecting anything more now, because we already
+        # waited for the responses
+        self.inner.drain_and_stop()
+
+
+def spawn_tensor_engine(proc_mesh: "ProcMesh") -> DeviceMesh:
+    # This argument to Controller
+    # is currently only used for debug printing. It should be fixed to
+    # report the proc ID instead of the rank it currently does.
+    gpus = proc_mesh.sizes.get("gpus", 1)
+    backend_ctrl = Controller(proc_mesh._proc_mesh)
+    client = MeshClient(backend_ctrl, proc_mesh.size(), gpus)
+    dm = DeviceMesh(
+        client,
+        NDSlice.new_row_major(list(proc_mesh.sizes.values())),
+        tuple(proc_mesh.sizes.keys()),
+    )
+    dm.exit = lambda: client.shutdown()
+    return dm