Every .item() call makes the GPU wait for Python

Training loops collect numbers: the loss, gradient norms, per-layer statistics for a dashboard. It feels free, because each number is tiny. I wanted to measure it, because of how a GPU is driven. When Python calls a CUDA operation, PyTorch queues it and returns immediately, so the Python thread can race ahead and queue the next operations while the GPU works through the earlier ones (Paszke et al., 2019¹). That overlap is what keeps a GPU busy. Asking for a value on the GPU as a Python number, with .item(), float() or .tolist(), ends the overlap: Python has to wait until every queued operation before it has finished, and the GPU then waits for Python to queue more.

The measurement

A small convolutional network, five convolution layers with batch normalisation and 22 parameter tensors, trained on 64×64 images on one GPU. I timed 300 steps after a warm-up, three repetitions, at batch sizes 16 and 128, with four kinds of logging:

1. none;
2. the loss read as a Python number every step;
3. the gradient norm of every parameter tensor read as a Python number every step, 22 reads;
4. the same statistics stacked into one tensor on the GPU and copied to the CPU every 50 steps.

Medians of three runs.

Reading the loss every step cost 6% at batch 16 and under 1% at batch 128: one wait per step, and at the larger batch the GPU is busy long enough that the wait barely shows. Reading 22 norms one at a time was a different matter: it made the step 5.6 times slower at batch 16, and even at batch 128 it was 1.6 times slower. Each read is a separate wait. While Python waits, nothing new is queued, so each small norm computation runs alone, followed by an idle gap. The same 22 numbers collected on the GPU and copied in one transfer every 50 steps cost 3% at batch 16 and almost nothing at 128.

Why the batch size changes the cost

At the small batch the GPU finishes a step's work in a few milliseconds, so a stall per read is a large fraction of the step. At batch 128 each step is long enough to hide some of it. That doesn't mean small batches are inherently wasteful. The GPU was already close to full: eight times the batch delivered only 13% more samples per second. A loop that stalls on reads can look like a case for larger batches, when the real fix is to stop stalling.

Finding the reads

The stalls are easy to add without noticing, because each looks like ordinary Python:

• .item(), float(), int() or .tolist() on a GPU tensor;
• print or an f-string of a GPU tensor;
• Python conditions on GPU values, if loss > threshold, which need the value on the CPU;
• .cpu() or .numpy() inside the loop;
• loops over parameters that compute and read one statistic each.

Two practices remove most of them. Keep statistics as tensors on the GPU, accumulate them, and transfer them in one copy at the logging interval. And when a check must happen every step, such as a guard against a non-finite loss, read one value for the whole step, not one per tensor.

The framework's other hidden choice, which kernel runs a given operation, can change the outputs themselves, in the same weights, two execution paths, two answers.

What I still don't know

My network is small and runs one stream on one GPU. In larger models each step is long and the relative cost of a few reads shrinks, but loops there also tend to log far more statistics, and multi-GPU training adds its own synchronisation points that interact with these. I also measured wall-clock time without a profiler trace. A trace would show the idle gaps directly, and the next step is to check them against the timings here.

References

1. PyTorch: An Imperative Style, High-Performance Deep Learning Library — Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai and Soumith Chintala, NeurIPS 2019