ZeRO and FSDP: Revolutionizing AI Training on Multiple GPUs

Introduction to the Challenges of AI Training on GPU

In the field of artificial intelligence, training massive models like GPT-3, which has 175 billion parameters, requires efficient resource management. Distributed Data Parallelism (DDP) is a commonly used method to accelerate training by distributing data across multiple GPUs. However, this technique introduces a major issue: memory redundancy. Each GPU retains a complete copy of the model parameters, gradients, and optimizer states, leading to a significant waste of valuable VRAM.

Memory Redundancy in DDP

To understand the impact of memory redundancy, let's examine what actually consumes memory during the training of a model. A model with N parameters must manage several elements:

• Model parameters, which are the weights of the neurons.
• Gradients, which are calculated for each parameter.
• Optimizer states, such as those used by Adam, which require storage of the first and second moments for each parameter.
• Activations, which are the intermediate outputs stored during the forward pass to be used in the backward pass.

The first three elements increase proportionally to the size of the model and are duplicated across each GPU in DDP. Activations, on the other hand, depend on the batch size, sequence length, and number of neurons, and are unique to each GPU since each GPU processes different data. ZeRO does not alter the management of activations.

For a model with 7 billion parameters using the Adam optimizer and FP32 precision format, the memory used breaks down as follows:

• Parameters: 7 billion multiplied by 4 bytes, totaling 28 GB.
• Gradients: 7 billion multiplied by 4 bytes, totaling 28 GB.
• Optimizer states: 7 billion multiplied by 2 times 4 bytes, totaling 56 GB.

Thus, each GPU in a DDP system uses 112 GB of memory, not counting activations which add an additional load. Techniques like activation checkpointing can be used to reduce this load by discarding certain activations and recomputing them as needed, but this is beyond the scope of this article.

ZeRO: An Innovative Solution

ZeRO, or Zero Redundancy Optimizer, offers an approach to reduce this memory redundancy. It comes in several versions, from ZeRO-1 to ZeRO-3, each providing incremental improvements.

ZeRO-1: Reducing Optimizer States

The first version, ZeRO-1, focuses on partitioning the optimizer states. In this model, each GPU retains a complete copy of the model parameters and gradients but only stores a fraction of the optimizer states, specifically 1/N where N is the number of GPUs. This means that each GPU only updates a corresponding portion of the parameters.

The training process with ZeRO-1 follows these steps:

• Forward pass: each GPU processes its own micro-batch of data.
• Backward pass: gradients are calculated.
• All-reduce of gradients: each GPU receives all gradients.
• Optimizer step: each GPU updates its partition of parameters.
• All-gather of parameters: synchronization of updated parameters between GPUs.

Implementation of ZeRO-1

To illustrate how ZeRO-1 works, here is a simplified implementation in Python:

import torch.distributed as dist

def __init__(self, model, optimizer_cls):
    self.model = model
    self.rank = dist.get_rank()
    self.world_size = dist.get_world_size()
    self.param_shards = list()  # each rank only holds its share of the optimizer states
    self.param_metadata = list()  # metadata to reconstruct the shards
    
    for param in self.model.parameters():
        original_shape = param.data.shape
        flat = param.data.view(-1)
        numel = flat.numel()
        remainder = numel % self.world_size
        pad_size = (self.world_size - remainder) % self.world_size
        padded_numel = numel + pad_size
        shard_size = padded_numel // self.world_size
        shard_start = self.rank * shard_size
        shard_end = shard_start + shard_size
        
        self.param_metadata.append({
            "original_shape": original_shape,
            "padded_numel": padded_numel,
            "shard_size": shard_size,
            "shard_start": shard_start,
            "shard_end": shard_end,
        })
        
        if pad_size > 0:
            flat_padded = torch.cat([flat, flat.new_zeros(pad_size)])
        else:
            flat_padded = flat
            
        shard = flat_padded[shard_start:shard_end].clone()
        shard.requires_grad_(True)
        self.param_shards.append(shard)
        
    self.optimizer = optimizer_cls(self.param_shards)

def training_step(self, inputs, targets, loss_fn):
    output = self.model(inputs)  # forward pass
    loss = loss_fn(output, targets)  # loss calculation
    loss.backward()  # backward pass
    self._sync_gradients()  # all-reduce of gradients between GPUs
    self.optimizer.step()  # update local parameter shard
    self._sync_params()  # synchronize model parameters
    # reset gradients for the next step
    for param in self.model.parameters():
        param.grad = None

def _sync_gradients(self):
    for idx, param in enumerate(self.model.parameters()):
        # Code to synchronize gradients

This implementation demonstrates how ZeRO-1 manages to partition the optimizer states while retaining complete parameters and gradients on each GPU. This allows for more efficient memory usage, thereby facilitating the training of large-scale AI models.