Glossary · Computing
CUDA graph
Also called CUDA graphs, graph capture, graphed, torch.cuda.graph
A recording of a sequence of GPU kernels that can be replayed as one unit, so the CPU doesn't have to launch each kernel again.
Normally the CPU tells the GPU what to do one kernel at a time. When the same long sequence of small kernels runs again and again, as in a training loop, that is wasted effort: the CPU repeats the same instructions every time, and the GPU waits for each one.
A CUDA graph (CUDA is NVIDIA’s platform for GPU programming) records the sequence once, with every kernel and the memory it reads and writes. After that one call replays the whole thing. The GPU knows everything that is coming and runs it back to back.
The catch: a replay uses the same memory addresses every time. So the data has to be copied into the same fixed blocks of memory before each replay, and the code inside can’t make decisions on the CPU that depend on the GPU’s results.
In this project
Training on a rented GPU recorded two parts of the trainer as CUDA graphs. One PPO update on a minibatch (run the network, work out the gradients, clip them, update the weights) went from 0.6–0.7 s an iteration to 0.2 s on an RTX 3090 in a desktop. On servers with slower processors it had been 0.7–1.4 s. The whole environment step, the MuJoCo Warp physics included, is a graph too, which saved about 5 %.