rGPU

nanoGPT example

Adapt the GPT lecture script to train through the PyTorch device.

This example uses Andrej Karpathy's GPT lecture script from ng-video-lecture, revision 5220142. It is the character-level model from Let's build GPT, rather than the separate nanoGPT repository's training program.

Get the example

In a new workload directory, download the script and its dataset at the same pinned revision:

curl -fL -o gpt.py https://raw.githubusercontent.com/karpathy/ng-video-lecture/52201428ed7b46804849dea0b3ccf0de9df1a5c3/gpt.py
curl -fL -o input.txt https://raw.githubusercontent.com/karpathy/ng-video-lecture/52201428ed7b46804849dea0b3ccf0de9df1a5c3/input.txt

Complete the PyTorch quickstart in your workload's Python environment to install rGPU, deploy the server, and verify the connection.

Change the device

Make these two edits to gpt.py:

+import rgpu
 import torch
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = 'rgpu'

The script already uses device for the model, input batches, position indices, and sampling context. It moves the model before the first forward pass and before constructing the optimizer, as required by the training guide.

Run it

From the directory containing gpt.py and input.txt, use the SSH destination and options from your quickstart:

rgpu-run --host user@gpu-host --ssh-port 2222 -i ~/.ssh/gpu_key \
  python gpt.py

The script prints the parameter count, periodic training and validation losses, and generated text. Its defaults run 5,000 training iterations and 200 evaluation batches per split. For a short connection check, reduce max_iters and eval_iters, and shorten the final max_new_tokens=500 sample.

Evaluation calls .item() for each batch, which waits for the server. Sampling also introduces frequent host waits. Both can be slow over a high-latency link; use a short sample while checking the setup.

Compilation and measurements

For optional graph compilation, follow the training guide after eager training works. Measure completed work before deciding whether compilation helps your run.

The recorded performance results come from a local yangpt variant with modified settings, not this pinned script. They document short runs and their limitations; they are not expected timings for this example.

On this page