rgpu 0.1.0__py3-none-any.whl

rgpu/__init__.py

@@ -0,0 +1,18 @@
+"""A PyTorch device whose tensors live on a remote GPU.
+
+    import rgpu, torch
+    x = torch.randn(1024, 1024, device="rgpu")   # the data is on the server
+    y = (x @ x).relu().sum().item()               # one round trip, at .item()
+
+Set RGPU_OPSERVER=host:port to reach rgpu-opserver (default 127.0.0.1:9720),
+normally through an ssh tunnel: the server has no authentication.
+"""
+
+from .device import register_device as _register_device
+
+_register_device()
+
+from .compile import compile_backend  # noqa: E402
+from .session import RemoteError, SessionLost, stats  # noqa: E402
+
+__all__ = ["RemoteError", "SessionLost", "compile_backend", "stats"]

rgpu/autocast.py

@@ -0,0 +1,140 @@
+"""Autocast for rgpu: which ops run in half precision and which stay float32.
+
+PyTorch ships an autocast policy only for its own devices. A new device gets
+none, and under autocast every op fails until something is registered at its
+Autocast key. This registers a fallthrough for everything, then kernels for a
+documented subset of what PyTorch registers for AutocastCUDA: LOWER runs in
+the autocast dtype, FP32 widens half inputs back to float32, and BANNED
+raises the way CUDA does. Everything else falls through uncast.
+
+Falling through is right for most of what is missing, because much of CUDA's
+list is CompositeImplicitAutograd and decomposes into ops that are covered -
+einsum and layer_norm were checked and come out as they do on CUDA. These are
+the differences that remain, where rgpu is not CUDA:
+
+  - grid_sampler and the RNN cells (lstm_cell, gru_cell, rnn_tanh_cell,
+    rnn_relu_cell, and the fused _thnn_* forms), which CUDA runs in the
+    autocast dtype and rgpu leaves alone.
+  - dot, vdot, bilinear and conv_tbc, likewise on CUDA's lower-precision list
+    and not here.
+  - CUDA's whole "promote" category, which runs an op at the widest dtype
+    among its inputs: addcmul, addcdiv, atan2, cross, index_put, scatter_add
+    and the upsample_* family. On rgpu a mix of float16 and float32 inputs
+    reaches them unpromoted.
+"""
+
+import torch
+from torch.utils._pytree import tree_map
+
+# Ops that run in the autocast dtype (PyTorch's CUDA "lower precision" list).
+LOWER = [
+    "mm", "addmm", "bmm", "baddbmm", "addbmm", "addmv", "addr", "matmul", "mv",
+    "linear", "conv1d", "conv2d", "conv3d", "conv_transpose1d", "conv_transpose2d",
+    "conv_transpose3d", "convolution", "_convolution", "prelu", "chain_matmul",
+    "linalg_multi_dot", "scaled_dot_product_attention",
+]
+
+# Ops whose half-precision inputs are widened to float32 (the CUDA fp32 list).
+FP32 = [
+    "acos", "asin", "cosh", "erfinv", "exp", "expm1", "log", "log10", "log2",
+    "log1p", "reciprocal", "rsqrt", "sinh", "tan", "pow", "softplus", "layer_norm",
+    "group_norm", "norm", "nuclear_norm", "cosine_similarity", "poisson_nll_loss",
+    "cosine_embedding_loss", "nll_loss", "nll_loss2d", "hinge_embedding_loss",
+    "kl_div", "l1_loss", "smooth_l1_loss", "huber_loss", "mse_loss",
+    "margin_ranking_loss", "multilabel_margin_loss", "soft_margin_loss",
+    "triplet_margin_loss", "multi_margin_loss", "binary_cross_entropy_with_logits",
+    "dist", "pdist", "cdist", "renorm", "logsumexp", "softmax", "log_softmax",
+    "cross_entropy_loss", "cumprod", "cumsum", "prod", "sum",
+]
+
+# Ops CUDA refuses under autocast rather than casting, because they are
+# unsafe in half precision. Running one anyway would diverge silently from a
+# local GPU, so rgpu refuses them with the same advice.
+BANNED = {
+    "binary_cross_entropy": (
+        "torch.nn.functional.binary_cross_entropy and torch.nn.BCELoss are unsafe to "
+        "autocast. Many models use a sigmoid layer right before the binary cross "
+        "entropy layer. In this case, combine the two layers using "
+        "torch.nn.functional.binary_cross_entropy_with_logits or "
+        "torch.nn.BCEWithLogitsLoss. binary_cross_entropy_with_logits and "
+        "BCEWithLogits are safe to autocast."),
+}
+
+_KEY = torch._C.DispatchKey.AutocastPrivateUse1
+_libs = []
+
+
+def _kernel(op, target, eligible):
+    def kernel(*args, **kwargs):
+        dtype = target()
+
+        def cast(x):
+            if (isinstance(x, torch.Tensor) and x.device.type == "rgpu"
+                    and eligible(x.dtype) and x.dtype != dtype):
+                return x.to(dtype)
+            return x
+
+        with torch._C._ExcludeDispatchKeyGuard(torch._C.DispatchKeySet(_KEY)):
+            return op(*tree_map(cast, args), **tree_map(cast, kwargs))
+
+    return kernel
+
+
+def _refusing_kernel(message):
+    def kernel(*args, **kwargs):
+        raise RuntimeError(message)
+
+    return kernel
+
+
+def _has_out_argument(schema):
+    """True if any argument writes into a tensor the caller passed in.
+
+    An out= kernel's cast() would allocate a fresh cast tensor, compute into
+    that, and leave the caller's out tensor untouched - silently wrong
+    values, no exception. PyTorch's own CUDA autocast policy does not cover
+    out= variants either, so - like CUDA - they fall through to the uncast
+    path instead of getting a cast kernel here.
+    """
+    return any(a.alias_info is not None and a.alias_info.is_write for a in schema.arguments)
+
+
+def install():
+    if _libs:
+        return
+    everything = torch.library.Library("_", "IMPL")
+    everything.fallback(torch.library.fallthrough_kernel, "AutocastPrivateUse1")
+    lib = torch.library.Library("aten", "IMPL")
+
+    def lower_eligible(dt):
+        return dt.is_floating_point and dt != torch.float64
+
+    def fp32_eligible(dt):
+        return dt in (torch.float16, torch.bfloat16)
+
+    for names, target, eligible in [
+        (LOWER, lambda: torch.get_autocast_dtype("rgpu"), lower_eligible),
+        (FP32, lambda: torch.float32, fp32_eligible),
+    ]:
+        for name in names:
+            packet = getattr(torch.ops.aten, name, None)
+            if packet is None:
+                continue
+            for overload in packet.overloads():
+                op = getattr(packet, overload)
+                if _has_out_argument(op._schema):
+                    continue
+                qualified = name if overload == "default" else f"{name}.{overload}"
+                lib.impl(qualified, _kernel(op, target, eligible), "AutocastPrivateUse1")
+
+    for name, message in BANNED.items():
+        packet = getattr(torch.ops.aten, name, None)
+        if packet is None:
+            continue
+        # Every overload, out= included: this kernel computes nothing, so the
+        # reason out= variants are skipped above does not apply.
+        for overload in packet.overloads():
+            qualified = name if overload == "default" else f"{name}.{overload}"
+            lib.impl(qualified, _refusing_kernel(message), "AutocastPrivateUse1")
+
+    _libs.extend([everything, lib])

rgpu/compile.py

@@ -0,0 +1,169 @@
+"""torch.compile for rgpu: ship the graph once, then one message per call.
+
+The backend hooks aot_autograd, which hands over forward and backward graphs
+of aten ops. Each is serialized with the same op names and argument encoding
+as eager ops, sent once, and compiled on the server. After that a call is a
+single CALL message, and it streams: output shapes come from the traced
+graph, so nothing has to wait for them.
+
+Compiled code works on the unwrapped meta tensors, never on RemoteTensor
+itself, which is why each tensor's id lives on its meta tensor: inputs are
+looked up by it and outputs are given one.
+"""
+
+import itertools
+import logging
+import operator
+
+import torch
+from torch._dynamo.backends.common import aot_autograd
+from torch.utils._pytree import tree_map
+
+from . import session, wire
+from .tensor import RemoteTensor, id_of, meta_like, register
+
+log = logging.getLogger("rgpu")
+_graph_ids = itertools.count(1)
+
+
+class Unshippable(Exception):
+    """A graph containing something the wire cannot carry."""
+
+
+class UnsupportedOp(Unshippable):
+    """An op in the graph that rgpu-opserver would refuse to run.
+
+    Unlike the other refusals this one has no eager fallback. Running the
+    graph through the ordinary dispatch path posts the very same op as a RUN,
+    and the server resolves ops by name and runs only aten ops, so it fails
+    there instead - later, and with a worse message. Raising at compile time
+    says the true thing: this cannot run on rgpu at all.
+    """
+
+
+def serialize(gm):
+    nodes, index = [], {}
+
+    def ref(a):
+        if isinstance(a, torch.fx.Node):
+            return wire.NodeRef(index[a])
+        if isinstance(a, torch.Tensor):
+            raise Unshippable("a constant tensor inside the graph")
+        if isinstance(a, torch.device):
+            # An rgpu tensor is a meta tensor on this side, so a graph traced
+            # through one allocates on "meta". On the server that means rgpu.
+            return wire.Dev("rgpu" if a.type in ("rgpu", "meta") else str(a))
+        if isinstance(a, (torch.SymInt, torch.SymFloat)):
+            raise Unshippable("dynamic shapes; compile with dynamic=False")
+        try:
+            wire.encode(a)
+        except wire.EncodeError as e:
+            raise Unshippable(str(e)) from None
+        return a
+
+    for n in gm.graph.nodes:
+        if n.op == "placeholder":
+            nodes.append(["placeholder"])
+        elif n.op == "call_function":
+            t = n.target
+            if t is operator.getitem:
+                # Not an op: how a graph picks one output of an op that has
+                # several, which aot_autograd writes for max_pool2d_with_indices,
+                # native_batch_norm and the rest. It travels as its own kind of
+                # node, so the server still runs nothing but aten ops.
+                src, i = n.args
+                if not isinstance(src, torch.fx.Node) or not isinstance(i, int):
+                    raise Unshippable(f"getitem with a {type(i).__name__} index")
+                nodes.append(["getitem", wire.NodeRef(index[src]), i])
+                index[n] = len(nodes) - 1
+                continue
+            if not isinstance(t, torch._ops.OpOverload):
+                # A plain Python callable: not shippable as part of a graph,
+                # but the ops it calls are, so eager still works.
+                raise Unshippable(f"{t} is not an aten op")
+            if t.namespace != "aten":
+                raise UnsupportedOp(
+                    f"{t._schema.name}.{t._overloadname} cannot run on rgpu: "
+                    "rgpu-opserver runs aten ops only, so a custom op has nowhere "
+                    "to run, compiled or eager")
+            nodes.append(["call", t._schema.name, t._overloadname,
+                          tree_map(ref, list(n.args)),
+                          {k: tree_map(ref, v) for k, v in n.kwargs.items()}])
+        elif n.op == "output":
+            outs = n.args[0] if isinstance(n.args[0], (list, tuple)) else [n.args[0]]
+            nodes.append(["output", [tree_map(ref, o) for o in outs]])
+        else:
+            raise Unshippable(f"a {n.op} node")
+        index[n] = len(nodes) - 1
+    return nodes
+
+
+def _output_values(gm):
+    out = next(n for n in gm.graph.nodes if n.op == "output")
+    outs = out.args[0] if isinstance(out.args[0], (list, tuple)) else [out.args[0]]
+    return [o.meta.get("val") if isinstance(o, torch.fx.Node) else o for o in outs]
+
+
+def _boxed(fn):
+    def run(args):
+        return fn(*args)
+    run._boxed_call = True
+    return run
+
+
+def _ship(gid, nodes, compiler, mode):
+    session.get().post(wire.COMPILE, gid, nodes, compiler, mode)
+
+
+def _eager(gm):
+    """Run a graph that cannot be shipped op by op, through the normal path."""
+    def run(*inner):
+        wrapped = [RemoteTensor(x) if isinstance(x, torch.Tensor) else x for x in inner]
+        outs = gm(*wrapped)
+        return [o._rgpu_meta if isinstance(o, RemoteTensor) else o for o in outs]
+    return _boxed(run)
+
+
+def _compiler(mode, compiler):
+    def compile_graph(gm, example_inputs):
+        try:
+            nodes = serialize(gm)
+            values = _output_values(gm)
+            if any(isinstance(s, torch.SymInt) for v in values
+                   if isinstance(v, torch.Tensor) for s in v.shape):
+                raise Unshippable("dynamic shapes; compile with dynamic=False")
+        except UnsupportedOp:
+            raise   # no eager fallback exists for this one; see UnsupportedOp
+        except Unshippable as e:
+            log.warning("rgpu: running a graph eagerly because it has %s", e)
+            return _eager(gm)
+        gid = next(_graph_ids)
+        _ship(gid, nodes, compiler, mode)
+
+        def run(*inputs):
+            in_ids = [id_of(x) for x in inputs]
+            outs, out_ids = [], []
+            for v in values:
+                if isinstance(v, torch.Tensor):
+                    m = meta_like(v.dtype, list(v.shape), list(v.stride()), v.storage_offset(),
+                                  v.is_conj(), v.is_neg())
+                    out_ids.append(register(m))
+                    outs.append(m)
+                else:
+                    out_ids.append(None)
+                    outs.append(v)
+            session.get().post(wire.CALL, gid, in_ids, out_ids)
+            return outs
+
+        return _boxed(run)
+
+    return compile_graph
+
+
+def compile_backend(mode=None, compiler="inductor"):
+    """A torch.compile backend that compiles on the rgpu server."""
+    comp = _compiler(mode, compiler)
+    return aot_autograd(fw_compiler=comp, bw_compiler=comp)
+
+
+torch._dynamo.register_backend(name="rgpu")(compile_backend())

rgpu/device.py

@@ -0,0 +1,120 @@
+"""Makes "rgpu" a PyTorch device, from Python alone.
+
+PyTorch reserves one device slot for backends outside it (PrivateUse1). Its
+Python-backend hook, marked experimental, registers the device guard and
+hooks a backend would otherwise have to write in C++.
+
+Factory functions - torch.zeros(..., device="rgpu") and the like - have no
+tensor argument to dispatch on, so they need kernels of their own. Every
+factory op is registered, not just empty: ops such as arange build an empty
+tensor and resize it through out=, and a wrapper's shape is fixed when it is
+made, so going through empty would leave it describing a size-zero tensor.
+"""
+
+import types
+
+import torch
+
+from . import session, wire
+from .tensor import RemoteTensor, register
+
+_SKIP_FACTORIES = ("sparse", "quantized", "cudnn", "mkldnn", "nested", "_make_dep_token",
+                   "from_file")
+
+
+def synchronize(device=None):
+    session.get().request(wire.SYNC)
+
+
+def manual_seed(seed):
+    session.get().post(wire.SEED, int(seed))
+
+
+def _set_device(device):
+    index = device.index if isinstance(device, torch.device) else device
+    if index not in (0, None):
+        raise ValueError("rgpu has one device, rgpu:0")
+
+
+def _module():
+    m = types.ModuleType("torch.rgpu")
+    m.is_available = lambda: True
+    m.device_count = lambda: 1
+    m.current_device = lambda: 0
+    m.set_device = _set_device
+    m.is_initialized = lambda: True
+    m._is_in_bad_fork = lambda: False
+    m.synchronize = synchronize
+    m.manual_seed = manual_seed
+    m.manual_seed_all = manual_seed
+    m.get_amp_supported_dtype = lambda: [torch.float16, torch.bfloat16]
+    return m
+
+
+def _factory_kernel(op):
+    name, overload = op._schema.name, op._overloadname
+
+    def kernel(*args, **kwargs):
+        meta_kwargs = dict(kwargs)
+        meta_kwargs["device"] = torch.device("meta")
+        if "pin_memory" in meta_kwargs:
+            meta_kwargs["pin_memory"] = None
+        meta = op(*args, **meta_kwargs)
+        wire_kwargs = {k: v for k, v in kwargs.items() if k != "pin_memory"}
+        wire_kwargs["device"] = wire.Dev("rgpu")
+        session.get().post(wire.RUN, name, overload, list(args), wire_kwargs, [register(meta)])
+        return RemoteTensor(meta)
+
+    return kernel
+
+
+def _copy_from_kernel(self, dst, non_blocking=False):
+    """Fills a tensor this backend just made, in place of copy_.
+
+    torch.tensor(data, device="rgpu") builds the destination through a
+    factory kernel above, then reaches PrivateUse1's _copy_from directly
+    rather than going through __torch_dispatch__ - the same aten op that
+    the ordinary CPU->rgpu path already runs as copy_.default.
+    """
+    from . import dispatch
+    return dispatch.handle(torch.ops.aten.copy_.default, (dst, self, non_blocking), {})
+
+
+def _factories():
+    """Every aten overload with a device argument and no tensor arguments."""
+    for full in torch._C._dispatch_get_all_op_names():
+        if not full.startswith("aten::") or any(s in full for s in _SKIP_FACTORIES):
+            continue
+        name, _, overload = full[len("aten::"):].partition(".")
+        try:
+            op = getattr(getattr(torch.ops.aten, name), overload or "default")
+        except (AttributeError, RuntimeError):
+            continue
+        schema = op._schema
+        if len(schema.returns) != 1 or "Tensor" not in str(schema.returns[0].type):
+            continue
+        if not any(a.name == "device" for a in schema.arguments):
+            continue
+        if any("Tensor" in str(a.type) for a in schema.arguments):
+            continue
+        yield full[len("aten::"):], op
+
+
+_libs = []
+
+
+def register_device():
+    if _libs:
+        return
+    torch.utils.backend_registration._setup_privateuseone_for_python_backend(
+        "rgpu", backend_module=_module())
+    lib = torch.library.Library("aten", "IMPL")
+    for qualified, op in _factories():
+        lib.impl(qualified, _factory_kernel(op), "PrivateUse1")
+    # torch.tensor(data, device="rgpu") calls this directly, bypassing
+    # __torch_dispatch__: see _copy_from_kernel.
+    lib.impl("_copy_from", _copy_from_kernel, "PrivateUse1")
+    _libs.append(lib)   # registrations live only as long as the Library object
+
+    from . import autocast
+    autocast.install()