Distributed Training

Things that we need to keep in memory while training a model include:

• Model parameters: weights and biases in every layer
• Optimizer states: for SGD, just the parameters, for Adam, the momentum estimates (if using AdamW, capture the first/second momentum estimates)
• Model activations: intermediate outputs of each layer during the forward pass; stored to be used during backprop to compute gradients
• Gradients: the partial derivatives of the loss with respect to each model parameter, computed ruing the backward pass (storing $\partial L/\partial \theta$ for each parameter $\theta$ before applying the optimizer step)
• Input data: current batch of training data being processed

Ways to reduce memory:

• Gradient Accumulation:
  • Solves the problem of the batch size being too big to fit in memory
  • Instead of doing:

loss = model(big_batch) loss.backward() optimized.step()

	- You split the batch into smaller chunks and do:

	```python
for micro_batch in chunks:
	loss = model(micro_batch)
	loss.backward()
optimizer.step()
optimizer.zero_grad()

- This allows you to store activations for one micro batch at a time, but does cause more forward/backward passes -> slower training per epoch

• Activation Checkpointing:
  • Solve the problem of forward activations taking too much memory during training
  • Instead of storing every activation for backpropagation, you “forget” some and recompute them during the backward pass

from torch.utils.checkpoint import checkpoint
def forward(x):
	x = checkpoint(block1, x)
	x = checkpoint(block2, x)
	return x

• CPU offloading:
  • Since GPU memory is limited, we can move things like optimizer states, gradients, or activations to the CPU (much slower though)

Scaling beyond a single GPU

• Data Parallelism
  • Clone the model on each GPU, each GPU processes a different slice of the batch.
    1. Each GPU runs forward +backward pass independently
    2. Gradients are averaged across GPUs
    3. All GPUs update their copy of the model the same way
  • Using the all-reduce communication primitive
    • Ensuring all methods stay synchronized
    • Fast implementations use ring all-reduce (each GPU shares chunks with neighbors in a bandwith efficient ring, accumulating as they go)
• Model Parallelism:
  • When the model itself is too big for one GPU
  • Can split layer-wise (assigning layer 1 to GPU 0, layer 2 to GPU 1, etc)
  • Tensor-wise: splitting individual tensors across GPUs - If a single layer (e.g., a 16K-dim transformer MLP) is too wide to fit: - Column-wise split: each GPU handles a slice of the output dimensions. - Forward: input is broadcasted; each GPU computes partial outputs. - All-gather to combine final output.
    - Row-wise split: each GPU handles a slice of the input dimensions - Forward: input is sharded; GPUs compute and all-reduce outputs.

References: Training extremely large neural networks across thousands of GPUs. 💥 Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups | by Thomas Wolf | HuggingFace | Medium