PyTorch quickstart
Set up the PyTorch device and run a local program on a remote GPU.
Prerequisites
- A client with Python 3.10 or newer. macOS works without CUDA or Docker.
- A GPU host with an NVIDIA driver, Python 3, and SSH access.
- This repository checked out on the client. The deployment script copies the server package, so the GPU host does not need GitHub access or its own clone.
The supported setup uses matching PyTorch major.minor versions on the client
and server. The package currently requires torch>=2.14.
1. Install rGPU in the workload environment
Activate the environment that will run your program, then install rGPU from PyPI:
cd /path/to/workload
source venv/bin/activate
python -m pip install rgpu
python -c 'import torch, rgpu; print(torch.__version__)'To run unreleased changes instead, install your checkout in editable mode:
python -m pip install -e /path/to/rgpu/pythonImporting rgpu registers the device. Every program that uses
device="rgpu" must import it.
Save this as smoke.py in the workload directory:
import torch
import rgpu
x = torch.ones(4, device="rgpu")
print((x * 2).sum().item()) # 8.02. Deploy the server
Keep the workload environment active so the script can detect its PyTorch version:
/path/to/rgpu/scripts/deploy_opserver.sh \
user@gpu-host -p 2222 -i ~/.ssh/gpu_keyThe script works from any directory. It:
- copies the server package from the local checkout;
- creates a virtual environment on the GPU host;
- installs the matching CUDA-enabled PyTorch build;
- starts
rgpu-opserveron the host's loopback interface; - prints
LISTENINGafter verifying port 9720.
The first deployment downloads a multi-gigabyte PyTorch wheel and can take
several minutes. Installation progress is printed. Pass
--torch-version MAJOR.MINOR only when deliberately selecting the remote
version.
3. Run through SSH
Use the same SSH destination and options:
rgpu-run --host user@gpu-host --ssh-port 2222 -i ~/.ssh/gpu_key \
python smoke.pyrgpu-run opens a private tunnel, sets RGPU_OPSERVER for the child process,
runs the command, and closes the tunnel. It does not change the device used by
the program; the program must import rgpu and select device="rgpu".
The expected output is 8.0. You are ready to train a model.
If this fails, start with troubleshooting.
Jupyter notebooks
Start Jupyter through rgpu-run so the notebook server and its kernels inherit
the tunnel configuration:
rgpu-run --host user@gpu-host --ssh-port 2222 -i ~/.ssh/gpu_key \
python -m jupyter labSelect the kernel from the environment where rGPU is installed, then run the
same smoke.py code in a cell.
Remote tensors belong to that kernel's server session. Restart the kernel and
recreate them after redeploying rgpu-opserver or restarting the GPU host.
For manual tunnels or local CPU/MPS testing, see Operations.