Compiling updates in Pytorch

When we train models using PyTorch we inadvertently store every gradient in memory until the backward pass is complete. Keeping every gradient in memory is inefficient because we could free memory as the gradient is used. First i’ll show you the memory consumption of a typical update, then we can see how a hook into autograd can fix this, and then i’ll show you that the compiler gives you this for free.

TLDR: Here is the eager-mode hook backprop.py, and here is the compiled example compiled_backprop.py, use the compiled version.

The Problem and fixing it with an eagar hook

Gradients are stored for the full backward pass which is the issue I’d like to address. If we could reduce the peak memory usage, we could train larger models, or run more experiments in parallel. Below is an interactive chart showing a typical backward pass.

Using an autograd hook to fix it

The desired behavior is that unused memory is freed as soon as possible. A quick way to do this is to initalise the optimiser separately for each layer, then pytorch will automatically return unused tensor shapes to CUDA’s reserved pool of memory to be re-used in the next layer. Below I compare to using my eager hook and we can see the hook cuts peak memory down to slightly more than the activations.

Usage

This solution is similar to what the Pytorch Lightning team did in a blog post: Faster PyTorch Training by Reducing Peak Memory. They observed slow-downs caused by memory management from not having enough storage to keep all tensor shapes cached in memory, and cache invalidations cost time. My implementation is different in the semantics on what the .backward() and .step() calls do; their approach makes .backward() do everything, including the optimiser step, and .step() becomes a no-op.

My hook makes .backward() compute the gradient of the final layer and halt, and .step() continues the backward pass, freeing memory as it computes gradients, steps the optimiser and continues the backward pass, or in code:

loss = loss_fn(model(inputs), targets) # as usual: loss is a scalar tensor
loss.backward()        # different: computes the gradient of the final layer and halts
optimizer.step()       # different: propagates the backward pass freeing memory as it goes
optimizer.zero_grad()  # as usual: frees gradients into the memory pool

Here is the code backprop.py.

Explanation

It might be helpful to see why the paradigm of .backward() and .step() existed in a way that shows you why the hook could work in the first place.

Below I have an illustration with three columns, we have our usual forward pass, backward pass, and optimiser step. The columns correspond to code:

1. loss = loss_fn(model(inputs), targets) is the “forward pass” column, then;
2. loss.backward() is the “backward pass” column, and;
3. optimizer.step() is the “optimiser step” column. The arrow going from the “backward pass” column to the “optimiser step” column is the reason the gradients are kept until the optimiser step.

Why we should be compiling PyTorch

Compilers should be able to statically analyse the workload to see that memory can be freed without the hook to make an optimiser for each layer. Normally we can’t fully compile PyTorch’s backward pass as there are a number of dynamic control flow decisions that are made at runtime under the hood and torch.inductor pulls an absolute fit about it.

What we can do is re-implement the optimiser and find much less memory is used and the backward pass is much faster due to operation fusion and/or fewer kernel launches. I have a similar graph below, but without the memory segmentation seen in the eager version.

The compiled step runs faster at 194 ms compared to 248 ms for the eager wrapper, and both using the same peak memory pool of ~12 GB. The code for this is compiled_backprop.py.

Teaser

Next I will showcase compiling experiments end-to-end in PyTorch. I do this to run as many experiments as possible, where without fully compiling my experiment I would need to spawn a process for each experiment which is much slower. After 4-6 processes we’re thrashing the GPU with work and the GPU’s SMs are still stalled and under-utilised. Instead, we can coalesce the experiments into a single shared graph to run our seeds/hypers/games in parallel.