GPU overhead: why our MNIST model trains faster on CPU

GPUs can speed up neural-network training by processing many calculations in parallel. Small models, however, may not provide enough work to offset the cost of dispatching operations, moving data, and synchronizing results.

We compared our MNIST model on a Quadro RTX 5000 GPU and an Intel Core i9-10885H CPU. In one Keras run, the CPU was faster: 4.4s vs 6.5s per epoch. A separate sweep recorded as little as 0.14s for a JAX training loop on CPU, but its timed workload differed from Keras’s. We’ll examine both the sources of overhead and what these measurements can tell us.

Let’s start by looking at what happens during the training.

Breaking down one training step

A training step runs a forward pass, computes the loss gradient, backpropagates it, and updates the weights. This excerpt computes the combined softmax–cross-entropy gradient directly, without calculating the scalar loss:

# Forward pass
z1 = xb @ w1 + b1                                # matrix multiply + bias (hidden layer)
a1 = np.maximum(0, z1)                            # ReLU activation
z2 = a1 @ w2 + b2                                # matrix multiply + bias (output layer)
exp_z = np.exp(z2 - z2.max(axis=1, keepdims=True))
probs = exp_z / exp_z.sum(axis=1, keepdims=True)  # softmax → probabilities

# Loss gradient
dz2 = probs.copy()
dz2[np.arange(bs), yb] -= 1                       # how far off were the predictions
dz2 /= bs

# Backward pass — compute gradients
dw2 = a1.T @ dz2                                  # gradient for W2
da1 = dz2 @ w2.T                                  # gradient flowing backward
dz1 = da1 * (z1 > 0)                              # ReLU gradient
dw1 = xb.T @ dz1                                  # gradient for W1

# Update weights
w1 -= lr * dw1
b1 -= lr * dz1.sum(axis=0)
w2 -= lr * dw2
b2 -= lr * dz2.sum(axis=0)

Where does the time go? To find out, we wrapped each operation with time.perf_counter():

t = time.perf_counter()
z1 = xb @ w1 + b1
timings["fwd: X @ W1 + b1"] += time.perf_counter() - t

t = time.perf_counter()
a1 = np.maximum(0, z1)
timings["fwd: ReLU"] += time.perf_counter() - t
# ... and so on for each operation

GPU operations are asynchronous, so we added tf.test.experimental.sync_devices() to wait for completion. Synchronizing after every operation also adds overhead and prevents overlap, so these timings should not be treated as the cost of an optimized training loop.

Here are the results for a single training step (batch_size=32, averaged over 1,000 steps, time in microseconds μs — millionths of a second, lower is better):

The CPU is faster in this per-operation test. NumPy uses optimized BLAS routines for matrix multiplication, while the TensorFlow GPU measurements also include dispatch and synchronization. SIMD instructions let the CPU process several values per instruction; fused multiply-add combines multiplication and addition. This comparison measures the implementations and timing method together, rather than isolating the arithmetic speed of the two devices.

Let’s look at xb @ w1 + b1 and a1 @ w2 + b2 — the matrix multiplies from each layer — and why GPUs are designed to make them fast. Here’s the forward pass from our NumPy implementation:

class HiddenLayer:
    def forward(self, x):
        self.z = self.W @ x + self.b        # matrix multiply + bias
        self.out = np.maximum(0, self.z)     # ReLU activation
        return self.out

class OutputLayer:
    def forward(self, x):
        self.z = self.W @ x + self.b        # matrix multiply + bias
        exp = np.exp(self.z - np.max(self.z))
        self.probs = exp / np.sum(exp)       # softmax → probabilities
        return self.probs

For one image, self.W @ x multiplies a 128×784 weight matrix by a vector of 784 pixel values. Each neuron computes a dot product: 784 multiplications and 783 additions. Adding the bias gives 128 pre-activation values, after which ReLU produces the hidden activations:

W (128 × 784) @ x (784,) + b (128,) → z (128,)

For a batch of 32 images, X has shape (32, 784). Using the transpose of the weights stored in our per-image implementation gives W1 with shape (784, 128). The product has shape (32, 128); bias addition and ReLU follow it:

      X (32 × 784)                W1 (784 × 128)            result (32 × 128)
 ┌─────────────────┐         ┌──────────────────┐       ┌──────────────────┐
 │ img1:  p1 … p784│         │  n1   n2  … n128 │       │ img1: z1  … z128 │
 │ img2:  p1 … p784│    @    │  w    w   …  w   │   =   │ img2: z1  … z128 │
 │ ...             │         │  ...  ...    ... │       │ ...              │
 │ img32: p1 … p784│         │  w    w   …  w   │       │ img32:z1  … z128 │
 └─────────────────┘         └──────────────────┘       └──────────────────┘
  32 images, 784 pixels each   784 weights per neuron     32 × 128 values

Batching reuses the same weights across 32 images and expresses the work as one matrix multiplication. This reduces the number of separate dispatches and gives the implementation more opportunities to reuse data and parallelize computation.

The 4,096 output values can be computed independently, but optimized GPU matrix multiplication usually distributes tiles of the matrices across groups of threads. We cannot infer the number of active cores or SMs simply by counting output values.

For an illustrative estimate, assume 3,072 FP32 cores at 1.8 GHz and two floating-point operations per multiply-add. That gives a theoretical peak of about 11 TFLOPS. Actual throughput depends on the GPU variant, clock speed, and workload.

The matrix multiplication requires approximately $2 \times 32 \times 784 \times 128 = 6{,}422{,}528$ floating-point operations, counting a multiply and an add separately. At 11 TFLOPS, the arithmetic-only lower bound is about 0.58 μs. This assumes peak throughput and excludes memory access and dispatch. It is not a measured kernel time, and dividing the operation count by peak FLOPS gives a time, not a percentage of GPU utilization.

Where the time actually goes

A GPU training step includes several costs beyond arithmetic. Their relative importance depends on the workload and execution mode:

1. Kernel dispatch. The CPU and runtime must submit work to the GPU. For small operations, dispatch can take a substantial part of the total time. A framework operation may use several kernels, while a compiler may fuse several operations into one kernel; the number of Python expressions does not determine the launch count.

2. Data movement and synchronization. Input data may need to move from CPU memory to GPU memory. A batch of 32 float32 images with 784 pixels occupies about 100 KB. Transfer time depends on the connection, memory allocation, and synchronization. During ordinary GPU training, weights and gradients can stay on the GPU; gradients do not need to return to the CPU after every step.

3. Python and framework work. Dispatch, memory allocation, metrics, and data handling add costs. In eager execution, Python participates in individual operations. Graph execution and compilation can remove much of that repeated work. XLA compilation occurs when a compiled function is first needed or retraced, not for every operation on every step.

Profiling helps distinguish these costs, provided we keep host API duration, GPU execution, and total elapsed time separate.

Profiling the CUDA overhead

The per-operation measurements suggest that overhead matters for this workload. To examine it, we also need a profile of the full training pipeline.

To find out, we used NVIDIA’s nsys (Nsight Systems) profiler, which intercepts every CUDA API call:

nsys profile -o keras-gpu-profile python mnist-keras.py
nsys stats --force-export=true keras-gpu-profile.nsys-rep

CUDA is NVIDIA’s software layer between your code and the GPU hardware. When TensorFlow wants to do a matrix multiply, it doesn’t talk to the GPU directly — it calls CUDA functions like “allocate memory,” “copy this data,” “launch this kernel.” Each of these calls goes through the driver and has its own overhead. The nsys profiler records every one of them, so we can see exactly where the time goes.

The recorded GPU epoch took 6.5s. Below is the accompanying CUDA API summary. An epoch with 48,000 training images and batch size 32 has 1,500 training steps, but this summary records 1,875 graph launches. The counts alone do not establish which calls belong to training, validation, or setup.

These are durations of CUDA API calls on the host, not GPU kernel durations. Calls can overlap GPU work or activity on other CPU threads. Their sum cannot be subtracted from 6.5s to identify the remaining time as Python overhead. A timeline with a clearly defined measurement range is needed for that breakdown. Nsight Systems distinguishes API, queue, and kernel time.

The cuGraphLaunch entries show that this run used CUDA Graphs to submit recorded work. Graph replay can reduce dispatch overhead, but the call count does not establish one launch per training batch or prove that GPU computation is negligible.

In the earlier Keras comparison, the CPU epoch took 4.4s, versus 6.5s on GPU. CPU execution avoids CUDA calls and transfers to a discrete GPU, though it still has framework, scheduling, and memory-access costs.

A separate TensorFlow microbenchmark timed (32, 784) @ (784, 128) without synchronizing after every operation. It reported these average times:

The GPU measurement is about 2.5 times faster in this test. Both numbers include software and memory-access costs; neither is a pure arithmetic measurement. They also use a different timing method from the first table, so the apparent reversal is not a contradiction.

The accompanying compiled-step script does different work on the two devices: its GPU path updates weights, while its CPU path only computes gradients. Those step timings therefore cannot support a direct CPU–GPU comparison. A fair comparison must run the same computation and wait for both devices to finish.

During training, nvidia-smi may show 90–100% GPU utilization even for a workload that runs faster on CPU.

GPU utilization in nvidia-smi measures the share of the sampling interval during which at least one kernel is executing. It does not measure the fraction of cores occupied or peak FLOPS achieved. Small kernels running continuously can therefore produce a high utilization reading without using the GPU’s full throughput.

Making a single training run faster

We tried five ways to change the training workload: larger batches, Keras on CPU, a NumPy loop, JAX compilation, and a direct CuPy port. They affect different parts of the execution cost.

1. Larger batch size

Larger batches reduce the number of updates per epoch. With 48,000 training images, batch size 32 gives 1,500 steps. Batch size 4,096 gives 12 steps if the final partial batch is included. The NumPy and JAX loops in our sweep instead drop that remainder, processing 11 full batches, or 45,056 images.

A larger batch also gives each matrix multiplication more work, which can help use GPU parallelism. The number of active SMs depends on the selected kernel and must be measured; it cannot be read directly from the batch size.

However, you can’t increase the batch size indiscriminately — it has a direct impact on model accuracy. Larger batches produce smoother but less frequent gradient updates, which can lead to worse generalization. The right batch size is something you need to experiment with for your specific model.

2. Keras on CPU — skip CUDA overhead

Disabling the GPU with tf.config.set_visible_devices([], 'GPU') before device initialization runs TensorFlow on CPU. This removes CUDA work from the execution path, but it does not imply an elapsed-time saving equal to the sum of CUDA API durations.

3. Pure NumPy — skip framework overhead too

A NumPy training loop avoids TensorFlow’s runtime and uses BLAS for matrix multiplication. It still pays for Python calls, array allocation, and memory access.

4. JAX JIT — compile the entire step

JAX’s jit compiles a function so repeated calls avoid dispatching every operation through Python. Compilation can fuse operations and reduce overhead, but one compiled function is not necessarily one GPU kernel. Compilation time should be separated from steady-state timing.

5. CuPy — what if we just move NumPy to GPU?

We also ported the per-image NumPy loop to CuPy. Small vector operations then required repeated GPU dispatches. This version took 443 seconds for five epochs, over six times the time reported for the corresponding NumPy loop. That result applies to this implementation; a batched CuPy implementation would be a different comparison.

Putting it all together

The following table records a separate sweep across batch sizes. The workloads are not identical: the Keras calls include validation and metrics, while the NumPy and JAX loops time training only and drop incomplete batches. These are timings of the scripts as written, not a controlled ranking of framework speed.

The batch-size-32 CPU result here is 11.31s, rather than the earlier 4.4s. These separate measurements should not be combined into a single speedup claim. A controlled comparison would use the same data, updates, validation work, warmup, and timing boundaries.

Keras GPU approaches 2.3s in this sweep. The table alone does not show which component limits further improvement.

Keras CPU falls from 11.31s to 1.39s as batch size increases. This measures the full fit() call used in the sweep.

NumPy reaches 0.84s at batch size 2,048 and takes 1.03s at 4,096. The timings alone do not identify the cause of that slowdown.

JAX records 0.14s at batch size 4,096 for its training-only loop. Dividing 6.26s by 0.14s gives about 45, but that ratio combines different batch sizes and workloads. It does not demonstrate a 45× speedup for equivalent training or equal validation accuracy.

For this small model, reducing repeated dispatch and changing batch size are worth testing before changing hardware. Compare time to the same validation quality, as well as time per epoch.

Using parallelization for hyperparameter sweeps

The approaches above make a single training run faster. But when you’re doing hyperparameter sweeps — trying 5 learning rates, or 4 architectures — each run is completely independent. They don’t share weights, gradients, or state.

JAX’s vmap adds a model dimension to our update function, letting one call update five independent parameter sets on the same images. Combined with jit, this avoids a Python loop over models. The compiler chooses how to execute the resulting operations; it does not guarantee one kernel or simultaneous execution of every model.

In this example, all models have the same parameter shapes. The following code shows the vectorized update, assuming x_batch and y_batch have already been prepared:

import jax
import jax.numpy as jnp
from jax import vmap, jit, grad, random

def init_params(key):
    # Same model as before: 784→128→10, random weights
    k1, k2 = random.split(key)
    w1 = random.normal(k1, (784, 128)) * jnp.sqrt(2.0 / 784)
    b1 = jnp.zeros(128)
    w2 = random.normal(k2, (128, 10)) * jnp.sqrt(2.0 / 128)
    b2 = jnp.zeros(10)
    return (w1, b1, w2, b2)

def loss_fn(params, x, y):
    # Forward pass + cross-entropy loss — same math as our NumPy version
    w1, b1, w2, b2 = params
    h = jnp.maximum(0, x @ w1 + b1)       # hidden layer + ReLU
    logits = h @ w2 + b2                    # output layer
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    return -jnp.mean(log_probs[jnp.arange(y.shape[0]), y])

def sgd_step(params, x, y, lr):
    # One training step: compute gradients, update weights
    grads = grad(loss_fn)(params, x, y)     # JAX auto-differentiates loss_fn
    return tuple(p - lr * g for p, g in zip(params, grads))

# 5 learning rates, 5 sets of weights, updated in one vectorized call
lr_array = jnp.array([0.001, 0.01, 0.1, 1.0, 10.0])
keys = random.split(random.PRNGKey(42), len(lr_array))
batched_params = vmap(init_params)(keys)
batched_step = jit(vmap(sgd_step, in_axes=(0, None, None, 0)))
batched_params = batched_step(batched_params, x_batch, y_batch, lr_array)

Another option is to train each configuration in a separate process. This sketch assumes Keras and the training data are available. Run it as a script with a main guard, and account for each worker’s memory use and CPU threads when choosing the process count:

import multiprocessing as mp

def train_one_config(args):
    # Train one model with one learning rate — runs in its own process
    lr, X_train, y_train = args
    model = keras.Sequential([
        keras.layers.Dense(128, activation="relu", input_shape=(784,)),
        keras.layers.Dense(10, activation="softmax"),
    ])
    model.compile(optimizer=keras.optimizers.SGD(learning_rate=lr),
                  loss="sparse_categorical_crossentropy")
    history = model.fit(X_train, y_train, epochs=10, batch_size=32,
                        validation_split=0.2, verbose=0)
    return lr, history.history

if __name__ == "__main__":
    with mp.get_context("spawn").Pool(5) as pool:
        results = pool.map(train_one_config,
                           [(lr, X_train, y_train) for lr in [0.001, 0.01, 0.1, 1.0, 10.0]])

What to take away

1. A small model can run faster on CPU. Our initial Keras comparison recorded 4.4s per epoch on CPU and 6.5s on GPU; the separate sweep produced different timings.
2. Peak TFLOPS is not a utilization measurement. An arithmetic-only lower bound cannot tell us the actual kernel duration or number of active cores.
3. Measure the full workload. Synchronization, validation, dropped batches, and compilation all affect what a timing means.
4. Compare equivalent training. Larger batches and compiled loops can reduce elapsed time, but compare validation quality before claiming a speedup.
5. Benchmark the devices you have. Parameter count alone does not determine whether CPU or GPU will be faster.

Word2vec offers another example of why the access pattern matters. With negative sampling, an update touches selected rows of embedding matrices rather than multiplying large dense matrices. The relative performance of CPU and GPU implementations depends on how those lookups and updates are batched; it is not a universal CPU advantage.