torchmonarch-nightly 2025.6.4__cp310-cp310-manylinux2014_x86_64.whl

monarch/controller/rust_backend/__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

monarch/controller/rust_backend/controller.py

@@ -0,0 +1,245 @@
+# 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
+
+import logging
+import traceback
+from collections import deque
+from logging import Logger
+from typing import List, NamedTuple, Optional, Sequence, Union
+
+from monarch._rust_bindings.monarch_extension import (
+    client,
+    controller,
+    debugger,
+    tensor_worker,
+)
+from monarch._rust_bindings.monarch_extension.client import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ClientActor,
+    SystemSnapshotFilter,
+    WorldState,
+)
+from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ActorId,
+    Proc,
+)
+
+from monarch._rust_bindings.monarch_messages.debugger import DebuggerAction
+from monarch.common.controller_api import LogMessage, MessageResult
+from monarch.common.device_mesh import no_mesh
+from monarch.common.invocation import DeviceException, RemoteException
+from monarch.common.messages import SupportsToRustMessage
+from monarch.common.shape import NDSlice
+from monarch.common.tensor import Tensor
+from monarch.controller.debugger import read as debugger_read, write as debugger_write
+from pyre_extensions import none_throws
+
+logger: Logger = logging.getLogger(__name__)
+
+
+class RustController:
+    def __init__(
+        self,
+        proc: Proc,
+        client_actor: ClientActor,
+        controller_id: ActorId,
+        worker_world_name: str,
+    ) -> None:
+        self._controller_actor = controller_id
+        self._proc = proc
+        self._actor = client_actor
+        # Attach the client to the controller
+        # Errors will be raised if someone else has attached it already.
+        self._actor.attach(self._controller_actor)
+        self._worker_world_name = worker_world_name
+
+        # 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 send(
+        self,
+        ranks: Union[NDSlice, List[NDSlice]],
+        msg: NamedTuple,
+    ) -> None:
+        self._actor.send_obj(self._controller_actor, ranks, msg)
+
+    def drop_refs(self, refs: Sequence[tensor_worker.Ref]) -> None:
+        self._actor.drop_refs(self._controller_actor, list(refs))
+
+    def node(
+        self,
+        seq: int,
+        defs: Sequence["Tensor"],
+        uses: Sequence["Tensor"],
+    ) -> None:
+        node = controller.Node(
+            seq=seq,
+            defs=[tensor_worker.Ref(id=t.ref) for t in defs if t.ref is not None],
+            uses=[tensor_worker.Ref(id=t.ref) for t in uses if t.ref is not None],
+        )
+
+        self._actor.send(self._controller_actor, node.serialize())
+
+    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._actor.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 stop_mesh(self) -> None:
+        logger.info("rust controller stopping the system")
+        self._actor.stop_worlds(
+            [self._controller_actor.world_name, self._worker_world_name]
+        )
+
+    def drain_and_stop(
+        self,
+    ) -> List[LogMessage | MessageResult | client.DebuggerMessage]:
+        logger.info("rust controller shutting down")
+        results = []
+        for msg in self._actor.drain_and_stop():
+            if isinstance(msg, client.WorkerResponse):
+                results.append(_worker_response_to_result(msg))
+            elif isinstance(msg, client.LogMessage):
+                results.append(LogMessage(msg.level, msg.message))
+            elif isinstance(msg, client.DebuggerMessage):
+                results.append(msg)
+            else:
+                raise RuntimeError(f"Unexpected message type {type(msg)}")
+        return results
+
+    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._actor.send(
+                debugger_actor_id,
+                debugger.DebuggerMessage(action=DebuggerAction.Attach()).serialize(),
+            )
+
+            while True:
+                # TODO: Add appropriate timeout.
+                msg = self._actor.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._actor.send(
+                        debugger_actor_id,
+                        debugger.DebuggerMessage(
+                            action=DebuggerAction.Write(
+                                bytes=debugger_read(action.requested_size)
+                            )
+                        ).serialize(),
+                    )
+                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:
+        worlds_state = self._actor.world_state(
+            SystemSnapshotFilter(worlds=[self._worker_world_name])
+        )
+
+        return worlds_state[self._worker_world_name]
+
+
+# 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")
+        ]
+        logger.error(f"Worker {exc.actor_id} failed")
+        return MessageResult(
+            seq=result.seq,
+            result=None,
+            error=RemoteException(
+                seq=exc.caused_by_seq,
+                exception=RuntimeError(exc.backtrace),
+                controller_frame_index=0,  # TODO: fix this once we have recording support in rust
+                controller_frames=None,
+                worker_frames=worker_frames,
+                source_actor_id=exc.actor_id,
+                message=f"Worker {exc.actor_id} failed",
+            ),
+        )
+    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)}")

monarch/fetch.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.
+
+# pyre-unsafe
+"""
+This is a utility file for fetching a shard of a tensor from remote.
+"""
+
+from typing import TypeVar
+
+from monarch.common.device_mesh import no_mesh
+
+from monarch.common.future import Future
+
+from monarch.common.remote import _call_on_shard_and_fetch
+
+T = TypeVar("T")
+
+
+def fetch_shard(
+    obj: T, shard: dict[str, int] | None = None, **kwargs: int
+) -> Future[T]:
+    """
+    Retrieve the shard at `coordinates` of the current device mesh of each
+    tensor in obj. All tensors in `obj` will be fetched to the CPU device.
+        obj - a pytree containing the tensors the fetch
+        shard - a dictionary from mesh dimension name to coordinate of the shard
+                If None, this will fetch from coordinate 0 for all dimensions (useful after all_reduce/all_gather)
+        preprocess - a
+        **kwargs - additional keyword arguments are added as entries to the shard dictionary
+    """
+    if kwargs:
+        if shard is None:
+            shard = {}
+        shard.update(kwargs)
+
+    return _call_on_shard_and_fetch(
+        None, lambda *args, **kwargs: None, obj, shard=shard
+    )
+
+
+def show(obj: T, shard: dict[str, int] | None = None, **kwargs: int) -> object:
+    v = inspect(obj, shard=shard, **kwargs)
+    # pyre-ignore
+    from torchshow import show  # @manual
+
+    with no_mesh.activate():
+        return show(v)
+
+
+def inspect(obj: T, shard: dict[str, int] | None = None, **kwargs: int) -> T:
+    return fetch_shard(obj, shard=shard, **kwargs).result()

monarch/future.py

@@ -0,0 +1,25 @@
+# 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 typing import Generator, Generic, TypeVar
+
+R = TypeVar("R")
+
+
+# TODO: consolidate with monarch.common.future
+class ActorFuture(Generic[R]):
+    def __init__(self, impl, blocking_impl=None):
+        self._impl = impl
+        self._blocking_impl = blocking_impl
+
+    def get(self) -> R:
+        if self._blocking_impl is not None:
+            return self._blocking_impl()
+        return asyncio.run(self._impl())
+
+    def __await__(self) -> Generator[R, None, R]:
+        return self._impl().__await__()

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)