Hi all. I’ve been building rgpu, which lets you write device="rgpu" on a laptop while the tensors and kernels live on a GPU host you reach over SSH.
pip install rgpu
import torch, rgpu
x = torch.ones(4, device="rgpu")
print((x * 2).sum().item()) # 8.0, computed on the remote GPU
How it works
- rgpu is a PrivateUse1 backend with a
RemoteTensorwrapper subclass. Each op’s output shape is worked out locally on meta tensors, so the op returns at once with a handle, and it’s queued and sent to the server in batches. - Only
.item(),.cpu(), printing, ops whose output size depends on the data, and an explicit synchronize wait for the server. Training nanoGPT for 20 iterations sent about 140,000 messages but waited only 5 times. torch.compile(model, backend="rgpu")sends the whole graph once, instead of one message per op.- Autograd, AMP, optimizers and reconnecting after a dropped connection all work. A test suite checks results against CPU references.
A data point: training nanoGPT (10.8M parameters) from a Mac on an A40 about 116 ms away took 1.04 s per iteration, against 3.04 s per iteration on the Mac’s own GPU.
Limits:
- The protocol has no authentication, so it’s meant for an SSH tunnel to a machine you trust.
- The client and server need matching PyTorch major and minor versions.
- It’s not a replacement for running directly on the GPU host. On a single machine, eager mode costs real overhead, mostly CPU time on the client.
A question for PyTorch developers: once rename_privateuse1_backend has run, getAccelerator() returns PrivateUse1. After that, backward() on a local CUDA or MPS tensor in the same process fails an assertion in engine.cpp (pytorch/pytorch#161129, broken since 2.8). Is there a supported way for an out-of-tree backend to coexist with a local accelerator?
Docs: Overview | rGPU · Code: GitHub - ymcrcat/rgpu: Keep Python on your laptop. Run PyTorch operations and hold tensors on a remote GPU, including from a Mac with no CUDA installation. · GitHub
Feedback welcome, especially from anyone with a workload where the GPU is far from where they write code.