Gradients for Backpropagation #1 - MatMul Forward + Backward

let the gradients flow freely!

Whenever we train an LLM or any neural network we’re essentially performing a giant series of matrix multiplications and rely on the underlying framework’s (JAX / PyTorch) autograd engine to derive gradients for our model’s individual layers without us ever touching a single calculus formula. We’ve become so detached from this lower level detail that most of us have forgotten how to derive gradients manually should there ever be a need. In this post we will learn how to manually implement the backward pass of nn.Linear. We’ll also explore how this knowledge can help us decide how to efficiently shard our layers for faster training.

Let’s consider a simple example where X is the 3x3 activation matrix , W is 2x3 matrix representing weights for the layer and O is the output of size 3x2. In PyTorch we usually store the weights with output dimension as first dim because underlying implementation uses W_t (transposed) for the multiplication, this makes gradient calculation simpler as we see later

Thanks for reading! Subscribe for free to receive new posts and support my work.

Now let’s reason about the gradients of X and W assuming we already have access to O_grad - the gradients of loss w.r.t O.

If we focus on gradient for x_00, we need to locate elements of O that have contributions from x_00 i.e o_00, o_01 and to see why this is the case we can expand the computation for these elements as :

o_00 = x_00 * w_00 + x_01 * w_10 + x_02 * w_20

o_01 = x_00 * w_01 + x_01 * w_11 + x_02 * w_21

So if we do have access to dL/do_00 and dL/do_01 we can write the gradient for x_00 as :

dL/dx_00 = dL/do_00 * w_00 + dL/do_01 * w_01

One observation we can make here is that the gradient of x_ij depends on dot product of i-th row of O_grad and j-th row of W or in matrix form we could note that

dL/dX = dL/dO @ W

Visually this is easier to reason about as seen below

Similarly for elements of W if we focus on w_00 and see it has contributions to 0th column of O with weights x_00 , x_10 and x_20 so we could write

dL/dW_t = X_t @ dL/dO

If we take transpose on both sides and noting that transpose(AB) = transpose(B) * transpose(A) we get

Now let’s implement our own custom forward + backward pass and verify that it matches with PyTorch default nn.Linear without bias.

#@title Custom Linear layer backward pass

from torch.autograd.function import once_differentiable
import torch.nn as nn
import torch

# ctx is the context helper that saves tensors from forward pass for later use
# in backward pass and without advanced techniques like activation checkpointing 
# this would mean most activations occupy memory between forward pass for layer_i 
# till gradient is computed for layer_i’s weights and activations
class CustomLinearLayerNoBias(torch.autograd.Function):

  @staticmethod
  def forward(ctx, X, W):
    ctx.save_for_backward(X, W)
    return X @ W.t()

  @staticmethod
  @once_differentiable
  def backward(ctx, grad_output):
    X, W = ctx.saved_tensors

    grad_X = grad_output @ W
    grad_W = grad_output.t() @ X

    return grad_X, grad_W

In order to compare with torch’s default nn.Linear we turn off bias and also make sure to set weights equal between the 2 paths.

# Setup inputs and torch nn.Linear ground-truth
X = torch.rand(3, 3)
W = torch.rand(2, 3)

X.requires_grad = True
W.requires_grad = True

torch_layer = nn.Linear(3, 2, bias=False)
torch_layer.weight.data.copy_(W)
torch_layer.weight.requires_grad = True

custom_out = CustomLinearLayerNoBias.apply(X, W)

X_copy = X.detach().clone()
X_copy.requires_grad = True
ground_truth_out = torch_layer(X_copy)

torch.testing.assert_close(custom_out, ground_truth_out)

custom_loss = custom_out.mean()
ground_truth_loss = ground_truth_out.mean()

# Since torch backward can be called on scalar output we mean the tensors to get a scalar.
custom_loss.backward()
ground_truth_loss.backward()

torch.testing.assert_close(W.grad, torch_layer.weight.grad)
torch.testing.assert_close(X.grad, X_copy.grad)

This code can also be seen in this gist and run on public google colab if you want to tinker around with more optimizations. This post is also inspired by Josh Levy’s awesome gradients for backpropagation series and here I wanted to explain it without tensor calculus as much as possible so it’s simpler to understand.

Bonus - RowParallel Sharding for W

Now let’s assume we have sharded W row-wise on 2 GPU devices due to memory constraints is it possible to do forward pass and backward pass without gathering the shards entirely into each device ?? Yes!

Let’s make the scenario concrete, we have O = X @ W_t and since we have sharded W row-wise, we have W_t sharded column-wise as shown below and X is replicated meaning if we have access to full X on both devices.

As seen above for forward pass since columns of W are on same device we can compute parts of the output without any communication needed between devices i.e O becomes sharded column-wise on it’s own because W was row-sharded and we compute O = X @ W_t … so that’s good news! but let’s look at gradient computation for W — dL/dW_t = X_t @ dL/dO … this also doesn’t require an all-reduce as we have everything needed to compute gradient of the weight shard locally … primarily enabled by the fact that X is replicated on each device.

Does this still hold true even for computing dL / dX ? unfortunately since dL /dX = dL/dO @ W and both operands on the right hand have shardings (column parallel for dL/dO and row parallel for W) we cannot compute this using just information locally on device so what’s the best we can do ? compute parts of correct output locally and then do all-reduce to get overall correct result

dL/dX_00 = dL/do_00 * w_00 + dL/do_01 * w_10

where dL/do_00 * w_00 is computed on device_0 and dL/do_01 * w_10 can be computed on device_1 then all_reduced with sum operation to get correct output on each device

How cool is that!!! We can correctly compute forward pass and backward pass without transferring weight shards entirely onto other device and simply doing one all-reduce which on modern GPUs is pretty fast inter-node thanks to NVLink! We’ll explore this topic deeper in a future post on FSDP and how this enables training LLMs that don’t fit on single GPU!

Thanks for reading! Subscribe for free to receive new posts and support my work.