OpenAI Techniques for Training Large Neural Networks

Training large neural networks requires orchestrating clusters of GPUs to perform single synchronized calculations. Because model and cluster sizes have grown, practitioners use various parallelism techniques to distribute the computational load and memory requirements across multiple hardware accelerators.

Data Parallelism

Data parallelism enables the use of multiple GPUs by copying the same model parameters to each worker and assigning different subsets of the training batch to be processed simultaneously. While this allows the compute power of many GPUs to be utilized, the model must still fit within a single GPU's memory.

To maintain consistency across workers, the system must coordinate parameter updates. The standard synchronous approach involves three steps:

  1. Independently computing the gradient on each worker.
  2. Averaging the gradients across all workers (a blocking communication step).
  3. Independently computing the new parameters on each worker.

While asynchronous synchronization schemes exist to reduce the overhead of the blocking average, they often decrease learning efficiency, leading most practitioners to prefer synchronous methods.

Pipeline Parallelism

Pipeline parallelism partitions sequential chunks of the model (consecutive layers) across different GPUs, reducing the memory footprint per device.

Addressing the "Bubble" Problem

Naive pipeline implementation creates "bubbles"—periods of idle time where a worker waits for the outputs of the previous machine. To minimize these bubbles, batches are split into smaller microbatches. Each worker begins processing the next microbatch as soon as it becomes available, overlapping computation with wait time. Gradients are averaged across microbatches, and parameter updates occur only after all microbatches are completed.

Scheduling Strategies

Two primary scheduling schemes are used to manage forward and backward passes:

  • GPipe: Workers process forward and backward passes consecutively, aggregating gradients from multiple microbatches synchronously at the end.
  • PipeDream: Workers alternatively process forward and backward passes, which can increase efficiency but may result in some computations using stale parameters.

Tensor Parallelism

Tensor parallelism splits individual operations within a layer "horizontally" across GPUs. In modern architectures like the Transformer, the primary bottleneck is multiplying an activation batch matrix with a large weight matrix. Tensor parallelism shards the weight matrix into even-sized pieces, hosting each shard on a different GPU to compute parts of the overall matrix product before communicating to combine the results.

Specific implementations include:

  • Megatron-LM: Parallelizes matrix multiplications within the Transformer's MLP and self-attention layers.
  • PTD-P: Combines tensor, data, and pipeline parallelism, assigning non-consecutive layers to each device to reduce bubble overhead.
  • Sequence Parallelism: Splits an input sequence across time into multiple sub-examples to decrease peak memory consumption.

Mixture-of-Experts (MoE)

Mixture-of-Experts (MoE) allows a model to scale its parameter count without a proportional increase in computation cost. This is achieved by using a gating mechanism that selects only a fraction of the network's weights (the "experts") to compute the output for any given input. Because different experts can be hosted on different GPUs, MoE provides a scalable way to increase the number of GPUs used for a model.

Memory Saving Designs

Beyond parallelism, several strategies are employed to reduce the device RAM required for training:

  • Checkpointing (Activation Recomputation): Instead of saving all original activations for the gradient computation, checkpointing stores a subset and recomputes intermediate ones just-in-time during the backward pass. Selective activation recomputation further optimizes this by checkpointing only activations that are expensive to store but cheap to compute.
  • Mixed Precision Training: Training with lower-precision numbers (typically FP16) increases FLOP counts and saves device RAM with minimal loss in accuracy.
  • Offloading: Temporarily moving unused data to the CPU or other devices. The ZeRO implementation splits parameters, gradients, and optimizer states across hardware and materializes them as needed.
  • Memory Efficient Optimizers: Using optimizers like Adafactor to reduce the memory footprint of the running state.
  • Compression: Compressing intermediate results, such as Gist compressing activations for the backward pass or DALL·E compressing gradients before synchronization.

Sources