In 2012, the paper Large Scale Distributed Deep Networks addressed the training of neural networks with billions of parameters on tens of thousands of CPU cores, by running copies (replicas) of the model on different subsets of the training data. The replicas sent their updates to a parameter server split across many machines, while modern frameworks such as PyTorch exchange the gradients directly between processes, with a collective communication at every training step.

Distributed data parallelism (DDP) refers to the parallel execution of different input samples across processors. Consider any data input of shape \(B \times S \times E\) (batch size, sequence/temporal size and embedding/feature size): data parallelism refers to the use case where the data is split across \(P\) processes, leading to a local input of shape \(B/P \times S \times E\):

An illustration of the DDP data layout, split across 4 processes color-coded as blue, yellow, red and green.

This post applies distributed data parallelism, sharding and CPU offloading to the training of the GPTlite model built in the previous post, on a network of 8 GPUs. The implementations use PyTorch DDP (Distributed Data Parallel) and DeepSpeed ZeRO (Zero Redundancy Optimizer, a lightweight wrapper on PyTorch).

There are two main data parallelism approaches:

Distributed Data Parallelism keeps a full copy of the model (weights, optimizer states and gradients) on each processor. All models are initialized equally. Each processor takes a different mini-batch as input and performs a forward pass to compute the loss of that batch.

On the backward pass, at every layer of the model, all processes compute the gradients of their batch and perform an all-reduce to get the mean gradients across all processors. These are then used to update the parameters and the optimizer states. Because parameters are initialized equally, gradients are mean-reduced and all processes perform the same updates, the models across all processors are kept in an identical state throughout the execution. The following diagram illustrates this workflow:

Distributed data parallelism (DDP) on 4 processes. Each process holds a full copy of a 4-layer model (the parameters, optimizer states and gradients of every layer), initialized identically, and takes its own input. The forward pass (left arrow) computes the activations and the loss of each process. In the backward pass (right arrow), each process computes the local gradients of every layer, and an all-reduce averages them across the processes. The optimizer then applies the same mean gradients on every process, so all copies stay identical.

The backward pass overlaps computation and communication: it computes the gradients of some layers while the gradients of other layers are sent across the network, which hides the communication latency behind the computation.

The diagram above shows that each processor holds a copy of the model, so the model states are stored \(P\) times, which leads to superfluous memory usage. That is where sharding comes into play.

(Fully-)Sharded Data Parallelism (FSDP), a.k.a. sharding, is a distributed setup where processors do not hold a full copy of the model, but only the parameters, optimizer states and gradients of a subset of layers. As before, each processor takes a different mini-batch as input.

In DeepSpeed terminology, sharding goes by the name of ZeRO (Zero Redundancy Optimizer). ZeRO has several alternative execution modes (stages). Each stage represents a different level of memory redundancy, corresponding to different variable types being distributed or replicated across processes:

  • ZeRO stage 0: no sharding of any variables, equivalent to Distributed Data Parallelism.
  • ZeRO stage 1 (ZeRO-1): the optimizer states (e.g., for the Adam optimizer in mixed precision, an fp32 copy of the weights and the first and second moment estimates) are partitioned across the processes. Affects only the backward pass.
  • ZeRO stage 2 (ZeRO-2): the gradients used to update the model weights are also partitioned, so that each process retains only the gradients corresponding to its portion of the optimizer states. Also relevant only for the backward pass.
  • ZeRO stage 3 (ZeRO-3): the model parameters are also partitioned across the processes. Adds a communication overhead to both the forward and backward passes.

An important remark is that activations are not sharded, i.e. they are kept in full shape on each processor. And in modern models, a huge chunk of memory is allocated to residual memory (activations, temporary buffers, fragmented memory, buffers of normalization layers, etc.), which is not sharded by FSDP. With that in mind, the following diagram illustrates the workflow of stage 3 sharding (of parameters, gradients and optimizer states):

Stage 3 sharding (ZeRO-3) on 4 processes. Process k holds the parameters, optimizer states and gradients of layer k only (solid boxes; the dashed boxes are layers held by other processes), takes its own input and keeps its own activations. Forward pass (left): before layer k is computed, process k broadcasts its parameters, and the other processes discard them after use. Backward pass (right): process k broadcasts the parameters of layer k again, every process computes its local gradients of layer k, and these gradients are summed (reduced) on process k, which updates the parameters of layer k. Each pass overlaps the computation of a layer with the communication of its neighbours: the all-gather of the parameters of the next layer to compute and, in the backward pass, the reduce-scatter of the gradients of the layer computed before. Inspired by: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Instead of a single all-reduce of the gradients, sharding uses two collective operations. The all-gather temporarily brings together the shards of the parameters, so that each process rebuilds the full weights of the layer it computes, in the forward and in the backward pass. The reduce-scatter, in the backward pass, sums the gradients across all processes and scatters the result, so that each process only stores and updates the gradient shard of its optimizer states. Note that DeepSpeed ZeRO-3 and PyTorch FSDP split the parameters of every layer into \(P\) shards, one per process, while the diagram assigns whole layers to processes: there, the all-gather of a layer is a broadcast from the process that holds it, and the reduce-scatter is a reduction on that process.

The higher the stage, the less memory is consumed. Stages 1 and 2 communicate as much as regular data parallelism, and stage 3 communicates 1.5x as much, as it gathers the parameters in both the forward and the backward passes. These memory improvements are summarized in the Microsoft Research blog as:

Per-GPU memory of the model states (parameters, gradients and optimizer states) with plain data parallelism (baseline) and with the three ZeRO stages: \(P_{os}\) partitions the optimizer states across the GPUs, \(P_{os+g}\) also the gradients, and \(P_{os+g+p}\) also the parameters. \(\Psi\) is the number of parameters, \(K\) the memory multiplier of the optimizer states and \(N_d\) the data-parallel degree; the example uses \(K=12\) (mixed-precision Adam), \(\Psi=7.5B\) and \(N_d=64\). The last column is the communication volume relative to the baseline. Adapted from: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models and the Microsoft Research blog

CPU Offloading

Sometimes the model is so big that, even with sharding, it will not fit in a single process. A common technique to handle such memory limitations is CPU offloading, also referred to as virtualized Deep Neural Networks by vDNN (Rhu et al.) and vDNN+ (Shriram et al.).

The main goal of this method is to iteratively move to the GPU only the portions of activations and model that are required for the current and next computation steps. Previously processed layers are moved from GPU to CPU, while upcoming layers are moved from CPU to GPU. vDNN offloads only the activations (the feature maps), as they take most of the GPU memory of convolutional neural networks.

This is possible because the loss (and its derivative) can be written as a composition of the functions of the layers, e.g. for the mean absolute error (MAE):

\[\mathcal{L} = \frac{1}{N} \sum_{n=1}^N | y_n - f^{(L+1)} \circ ... \circ f^{(2)} \circ f^{(1)} (x_n^{(0)}) |\]

where \(f^{(l)}\) is the function implemented by layer \(l\), \(L\) is the number of hidden layers, \(y\) is the ground truth, \(x_n^{(0)}\) is the input of sample \(n\) and \(N\) is the number of samples.

The important concept here is the composition of the functions \(f\) across layers. The computation of a layer needs only its weights, its input and its output (and, in the backward pass, the gradient coming from the next layer), so the data of all other layers can be kept in CPU memory. To reduce the time spent waiting to push (pull) a layer to (from) the GPU, the computation of the current layers can be overlapped with the communication of the upcoming ones. The forward pass can be illustrated as:

The forward pass of vDNN on a convolutional neural network, at layer N. Blue arrows are the input and output feature maps \(X\) and \(Y\) of the current layer, the purple arrow its weights \(W\) and the yellow arrow its cuDNN workspace \(WS\). White arrows are data that the current layer does not use. The input feature maps \(X\) of the preceding layers, marked with a red cross, are not used by the current layer and can be released from GPU memory. Adapted from: vDNN: Virtualized Deep Neural Networks for Scalable, Memory-Efficient Neural Network Design

The backward propagation phase is trickier, as it also includes the gradients being propagated down the model:

The backward pass of vDNN at layer 2. Blue arrows are the feature maps \(X\) and \(Y\) of the current layer, green arrows its gradients \(dX\) and \(dY\), the purple arrow its weights \(W\) and the yellow arrow its cuDNN workspace \(WS\). White arrows are data that the current layer does not use. The data used only by the layers whose backward pass is complete (the layers after layer 2 and the loss function), marked with a red cross, is not reused for this batch and can be released from GPU memory. Adapted from: vDNN: Virtualized Deep Neural Networks for Scalable, Memory-Efficient Neural Network Design

Model and dataset setup

The implementation starts by taking the previous GPTlite model, with specs matching the GPT-3 Small model in Language Models are Few-Shot Learners (Table 2.1):

n_layer = 12   # depth of the network as number of decoder blocks.
n_embd = 768   # size of the embeddings (d_model)
n_head = 12   # number of attention heads in the Multi-Attention mechanism
block_size = 2048   # block size, i.e. the maximum sequence length, the n_ctx in the paper
dropout = 0.1   # dropout rate (variable p) for dropout units

It then defines the methods get_model() and get_dataset() that return the model and the tiny shakespeare dataset:

## gptlite.py

def get_dataset():
  
  class GPTliteDataset(torch.utils.data.Dataset):

      def __init__(self, train_data, block_size):
        self.train_data = train_data
        self.block_size = block_size

      def __len__(self):
        return len(self.train_data)

      def __getitem__(self, idx):
        # generate 1 random offset on the data
        ix = torch.randint(len(self.train_data)-self.block_size , size=())
        # input is a random subset of tokens
        x = self.train_data[ix   : ix+self.block_size]
        # target is just x shifted right (ie the next predicted word)
        y = self.train_data[ix+1 : ix+1+self.block_size]
        return x, y

  train_data, valid_data, vocab_size = load_tiny_shakespeare_data()
  train_dataset = GPTliteDataset(train_data, block_size)
  valid_dataset = GPTliteDataset(valid_data, block_size)
  return train_dataset, valid_dataset, vocab_size


def get_model(vocab_size):
  return GPTlite(vocab_size)

Detour: using a pre-existing model

Note that this code is applicable to any model of type torch.nn.Module and any dataset of type torch.utils.data.Dataset. As an example, to perform multi-class classification using the ResNet network on the CIFAR10 dataset available in torchvision, the previous 2 methods could be defined as:

import torchvision

def get_dataset():
  import torchvision.transforms as transforms
  transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
  ])
  dataset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform)
  return dataset

def get_model(num_classes):
  return torchvision.models.resnet18(num_classes=num_classes)

PyTorch implementation

This section implements data parallelism in PyTorch. First, the code collects the environment variables that are set by the torchrun launcher (detailed below), as they uniquely identify processes in the network and GPU devices within a node:

import os
rank = int(os.environ['RANK'])   # the unique id across all processes in all nodes
local_rank = int(os.environ['LOCAL_RANK'])   # the unique id across this node
world_size = int(os.environ['WORLD_SIZE'])   # the number of processes across all nodes

Now the code defines the DataLoader that tells each process how to iterate through the data:

dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, sampler=sampler)

Note the argument sampler, a torch.utils.data.distributed.DistributedSampler that assigns different samples of the dataset to different processes. Without it, all processes would load exactly the same datapoints in every iteration. The sampler is created as:

sampler = torch.utils.data.distributed.DistributedSampler(
  dataset, num_replicas=world_size, rank=rank)

Each model is then placed on a different GPU, with the correct data type (here bfloat16):

device = f"cuda:{local_rank}"
dtype = torch.bfloat16
model = model.to(device=device, dtype=dtype)

and finally wrapped with DistributedDataParallel for the DDP implementation:

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

or with the FullyShardedDataParallel wrapper for the sharded implementation:

model = torch.distributed.fsdp.FullyShardedDataParallel(
    model,
    device_id=local_rank,
    sharding_strategy=torch.distributed.fsdp.api.ShardingStrategy.SHARD_GRAD_OP,  # sharded grads + optimizer states (ZeRO-2-like)
)

The argument sharding_strategy plays the role of the ZeRO stage: NO_SHARD replicates all variables as DDP does, SHARD_GRAD_OP shards the gradients and optimizer states similarly to ZeRO-2, and FULL_SHARD (the default) also shards the parameters, as ZeRO-3.

And that is all. The training loop is then written as usual, and PyTorch does all the communication and synchronization under the hood:

for step, data in enumerate(dataloader):
  inputs = data[0].to(device)
  labels = data[1].to(device)
  outputs = model(inputs)  # fwd pass
  loss = torch.nn.functional.cross_entropy(outputs, labels) # loss
  loss.backward() # compute gradients
  optimizer.step() # update parameters
  optimizer.zero_grad(set_to_none=True)

During the backward pass, gradients are computed and mean-reduced from the top to the bottom layer of the model. Therefore, the computation of the gradients of the lower layers can be overlapped with the communication of the gradients of the upper layers.

According to the PyTorch DDP documentation and the PyTorch DDP paper, the backward pass iteratively computes gradients (from last to first layer) and collects blocks of gradients (buckets) to be communicated. These blocks are mean-reduced asynchronously during the backward pass, while the computation for the backward pass proceeds. At the end of the backward pass, all GPUs wait for all gradient all-reduces to finish, and then the optimizer updates the parameters.

For extra memory savings, tensors can also be offloaded in PyTorch by using custom hooks for autograd saved tensors.

Finally, the application is launched by calling torchrun. Torchrun is a network bootstrapper that, started on each compute node in a network, spawns one process of a Python script per GPU and sets the environment variables mentioned above, which uniquely identify processes across the network. The usage is simple:

$ torchrun --standalone --nproc_per_node=4 train.py

where nproc_per_node is the number of GPUs on each node.

DeepSpeed implementation

Porting existing code to DeepSpeed is pretty simple. To start, DeepSpeed features can be activated via the DeepSpeed API or its Configuration JSON. The number of possible optimizations is large, as the config can define parallelism, floating-point precision, logging, communication parameters, etc.

This implementation starts with a simple config file. The file configures the DeepSpeed logger to output memory and throughput info every 10 steps (steps_per_print), sets the batch size to 256 and (optionally) defines the settings of the optimizer (optimizer) and learning rate scheduler (scheduler):

{
  "train_batch_size": 256,
  "steps_per_print": 10,
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 0.001, "betas": [ 0.8, 0.999  ], "eps": 1e-8, "weight_decay": 3e-7 }
  },
  "scheduler": {
    "type": "WarmupLR",
    "params": { "warmup_min_lr": 0, "warmup_max_lr": 0.001, "warmup_num_steps": 1000 }
  }
}

Note that, in DeepSpeed terminology, train_micro_batch_size_per_gpu refers to the batch size loaded per dataloader (i.e. per GPU, per gradient accumulation step), while train_batch_size refers to the batch size across all gradient accumulation steps and processes in the network, i.e.:

train_batch_size = train_micro_batch_size_per_gpu * num_gpus * num_nodes * gradient_accumulation_steps

Therefore, as train_batch_size is already set, gradient accumulation (a technique that simulates a large mini-batch as an iteration across several micro-batches) can be enabled by setting train_micro_batch_size_per_gpu.

ZeRO Fully-Sharded Data Parallelism can be activated by specifying the relevant stage in the config file. If omitted, or when passing stage 0, ZeRO is disabled and the execution follows a regular distributed data parallel workflow:

{
  "zero_optimization": { "stage": 3 }
}

CPU offloading in DeepSpeed is called ZeRO-Offload, a system for offloading optimizer and gradient states to CPU memory, introduced in the ZeRO-Offload paper. On top of stage 3, it is possible to enable ZeRO-Infinity, also an offloading engine, which extends ZeRO-Offload with support for NVMe memory. According to the ZeRO-3 documentation, “ZeRO-Infinity has all of the savings of ZeRO-Offload, plus is able to offload more the model weights and has more effective bandwidth utilization and overlapping of computation and communication”. They can be enabled via:

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "offload_param": { "device": "cpu", "pin_memory": true }
  }
}

Almost everything is ready now. Once the config file is properly calibrated, the implementation is straightforward. All boilerplate that PyTorch requires for parallelism and data loaders is managed internally by DeepSpeed, so the code only needs to set up DeepSpeed as:

def main_deepspeed(n_epochs=100, random_seed=42):

  torch.manual_seed(random_seed)  #set random seed (used by DataLoader)
  deepspeed.runtime.utils.set_random_seed(random_seed) #set DeepSpeed seed
  deepspeed.init_distributed()  # initialize distributed DeepSpeed
  config = 'ds_config.json'  # load the DeepSpeed config
  criterion = torch.nn.CrossEntropyLoss()  # initialize loss function
  train_dataset, _, vocab_size = gptlite.get_dataset()  # initialize GPTlite dataset
  model = gptlite.get_model(vocab_size)  # initialize GPTlite model

  engine, optimizer, train_dataloader , _ = deepspeed.initialize(
    config=config, model=model, training_data=train_dataset,) # initialize deepspeed

The training loop is then written with a structure very similar to the PyTorch implementation. The only exception is that the gradients are not zeroed, as this is managed internally by DeepSpeed. Also, initialize() already returns a train_dataloader that assigns disjoint subsets of data to each process, so there is no need to fiddle with distributed dataloaders and samplers.

  for epoch in range(n_epochs):
    for step, data in enumerate(train_dataloader):
      inputs = data[0].to(engine.device)
      labels = data[1].to(engine.device)
              
      outputs = engine(inputs)  # fwd pass
      loss = criterion(outputs, labels)
      engine.backward(loss)  # backprop
      engine.step()  # update weights, no need for zero-ing

Finally, the run is launched with the torchrun launcher as before, or with the launcher included in DeepSpeed as:

$ deepspeed --num_gpus=8 train.py --deepspeed --deepspeed_config ds_config.json

Detour: measuring memory allocated to parameters

The DeepSpeed API to estimate the memory requirements of model parameters for different ZeRO implementations can be used by calling the following method at the onset of execution:

def measure_parameters_memory(model):
  param_size_GB = sum([p.nelement() * p.element_size() for p in model.parameters()])/1024**3
  print(f"Native model parameters size: {round(param_size_GB, 2)}GB.")

  from deepspeed.runtime.zero.stage_1_and_2 import estimate_zero2_model_states_mem_needs_all_live
  estimate_zero2_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)

  from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live
  estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=8, num_nodes=1)

The output shows that:

  • the base model requires about 0.323GB per GPU to store its parameters;
  • DeepSpeed ZeRO-2 requires 0.161GB and 0.484GB with and without offloading, respectively;
  • DeepSpeed ZeRO-3 requires 0.009GB and 0.190GB with and without offloading, respectively.

This metric is very useful as it gives a quick overview of scaling and is very fast to compute. However, it has many limitations: it only measures the memory of the parameters, gradients and optimizer states, and it does not take activations or other residual buffers (e.g. normalization variables) into account. The batch size and numerical precision (or any field in the config file) and the temporary (e.g. communication) buffers are also ignored.

Results

To measure performance, the DeepSpeed logger was used to extract the following metrics from different runs every 10 steps: model throughput as the average number of samples per second, the average allocated memory, and the maximum allocated memory. The runs used pytorch==2.0.1, CUDA 11.7 and deepspeed==0.10.3.

All implementations use the same numerical precision (fp32), communication bucket sizes, and other config parameters. The following implementations (and configs) were benchmarked:

  1. The distributed data parallel (DDP) implementation, i.e. no DeepSpeed (ds_config.json with 'stage': 0);
  2. The fully-sharded data parallel implementation with ZeRO 1, 2 and 3 (ds_config.json with 'stage' :1, 2 or 3);
  3. The ZeRO-3 implementation with ZeRO-Infinity for CPU offloading (ds_config_offload.json).

The GPTlite model was tested with a micro-batch size of 1 sample per GPU, and the results are:

Maximum and average memory allocated, and throughput, of GPTlite (GPT-3 small) with distributed data parallelism, DeepSpeed ZeRO-1 to ZeRO-3, and ZeRO-3 with CPU offloading, at a micro-batch size of 1 sample per GPU.

Memory overhead from communication buffers. Looking at the max vs average memory, note that, in theory, the max memory should be much higher at high ZeRO stages than at low ZeRO stages and DDP. This is because more parameters are communicated, which requires more communication buffers. However, setting the communication bucket sizes (the number of elements communicated at a time) to a low value in the config file overcomes this effect. In fact, several runs were also benchmarked with the default communication bucket sizes (5e8), which led to higher memory usage, as expected (approximately double in stages 2 and 3), and became prohibitive for some runs.

Communication vs computation trade-off from different stages in ZeRO. In ideal scenarios, as the execution moves from DDP to ZeRO-1, ZeRO-2, ZeRO-3 and ZeRO-Infinity, the memory consumption and throughput are reduced. As expected, data locality is traded for communication of parameters, and a price in performance is paid for the communication/offloading of parameters. In practice, the fastest choice is therefore the lowest stage that fits the model in memory.

Offloaded vs in-memory parameters. Offloading proved to be a consistent way to reduce memory usage, with the drawback of a small reduction in throughput, as expected.

Lack of superlinear speedup. A small improvement in memory efficiency was observed, but still far from the claims of the original DeepSpeed team. One explanation is that a small network of 8 GPUs was used, compared to the 64 to 400 GPUs used by the authors in their benchmarks, so the memory reduction is much smaller. Moreover, most of the memory of a model as small as GPTlite is taken by its activations, which are not sharded. A large network of GPUs also underlies their claim of superlinear speedup (with more GPUs, each GPU holds fewer model states and fits a larger batch, at a higher throughput), which was not observed here.

Finally, communication quantization, which reduces the floating-point representation of the communicated tensors, was not used, as it did not result in any improvement. This is in line with the ZeRO++ paper, which claims this method is applicable to small batch sizes or low-bandwidth networks.

Resources and code

This post explored only the dimension of data parallelism. More details about DeepSpeed can be found in the DeepSpeed API documentation, the training features page, the tutorials page, the HuggingFace page for DeepSpeed, and the examples at DeepSpeedExamples.

As an important remark, finding the best parallelism strategy is hard, due to the high number of hyper-parameters (ZeRO stages, offloading, activation checkpointing intervals, pipeline parallelism stages, data parallelism, tensor parallelism, etc.). The best strategy also depends on the ML model, the data and the hardware. In practice, the DeepSpeed config files of this post are a manually-optimized ballpark figure of the default config file with some parameter grid search. On this topic, there is still plenty of work to be done to make it optimal, possibly by exploring the autotuning tools in DeepSpeed.

This post only scratched the surface of the capabilities of DeepSpeed. There are plenty of features that should also be explored. To name a few:

  • autotuning (README.md) for parallelism hyper-parameter discovery;
  • flops profiler to measure the time, flops and parameters of individual layers;
  • sparse attention kernels (API) to support long sequences of model inputs, such as text, image, or sound;
  • communication optimizers that offer the same convergence as Adam/LAMB but incur up to 26x less communication on BERT-large pretraining;
  • monitor to log live training metrics to TensorBoard, csv file or other backend;
  • model compression (API) via layer reduction, weight quantization, activation quantization, sparse pruning, row pruning, head pruning and channel pruning, to deliver faster speed and smaller model size.

There are a lot of results and food for thought here. The code of all experiments is available in the distributed-training repo.