Optimizing AI with Pipeline Parallelism and Micro-Batching

As artificial intelligence models grow to billions of parameters, fitting an entire neural network onto a single Graphics Processing Unit (GPU) has become impossible. Pipeline parallelism addresses this memory bottleneck by splitting the layers of a neural network sequentially across multiple GPUs. However, basic layer splitting leads to severe inefficiency, as GPUs spend most of their time waiting for preceding devices to finish processing data. To solve this, engineers use a technique called micro-batching, which breaks incoming training data into smaller chunks to keep all GPUs active simultaneously, significantly reducing idle time—known as pipeline bubbles—and drastically accelerating the AI training process.

Splitting Neural Network Layers Across GPUs

Pipeline parallelism is a form of model parallelism that divides a deep neural network vertically. Because neural networks are structured as a sequence of layers, pipeline parallelism assigns consecutive groups of layers to different GPUs.

For example, if a network has 24 layers and there are four available GPUs, the model is partitioned into four stages:

  • GPU 1: Layers 1 to 6
  • GPU 2: Layers 7 to 12
  • GPU 3: Layers 13 to 18
  • GPU 4: Layers 19 to 24

During the forward pass of training, GPU 1 takes the raw input data, processes it through layers 1 to 6, and sends the resulting intermediate outputs (activations) over an interconnect to GPU 2. GPU 2 processes those activations through layers 7 to 12 and passes the results to GPU 3, continuing until GPU 4 calculates the final output and loss. During the backward pass, error gradients flow in the exact opposite direction, from GPU 4 back to GPU 1, to update the model's weights.

The Problem of Pipeline Bubbles

While splitting layers solves the memory limit of a single GPU, a naive implementation creates a massive computational bottleneck.

In a standard setup, GPU 2 cannot begin its work until GPU 1 has completely finished processing the entire batch of data. Similarly, while GPU 4 is processing the final layers of the forward pass and the beginning of the backward pass, GPUs 1, 2, and 3 sit completely idle, waiting for the gradients to return. This idle time across the GPU cluster is referred to as a "pipeline bubble." Without optimization, up to 75% of a four-GPU cluster's compute capacity is wasted waiting for data transfers.

How Micro-Batching Minimizes Idle Time

To shrink the pipeline bubble, AI frameworks use micro-batching. Instead of feeding a massive mini-batch of data into the pipeline all at once, the mini-batch is split into smaller segments called micro-batches.

Here is how micro-batching changes the workflow across four GPUs using four micro-batches (M1, M2, M3, M4):

  1. Step 1: GPU 1 processes the forward pass for M1.
  2. Step 2: GPU 1 passes the M1 activations to GPU 2. Immediately, GPU 1 begins processing M2 instead of waiting.
  3. Step 3: GPU 2 passes M1 to GPU 3, GPU 1 passes M2 to GPU 2, and GPU 1 starts on M3.
  4. Step 4: All four GPUs are now computing different micro-batches simultaneously.

By streaming these smaller data chunks like an assembly line, the GPUs are kept busy. The idle time is strictly limited to the brief startup phase (filling the pipeline) and the cooldown phase (draining the pipeline). Once all micro-batches have completed their forward and backward passes, the gradients are accumulated across all micro-batches, and the model weights are updated synchronously.

Advanced Scheduling: 1F1B

Modern pipeline parallelism goes a step further by optimizing the order of operations using the One-Forward-One-Backward (1F1B) schedule.

In a standard micro-batching setup, a GPU processes the forward pass for all micro-batches before doing any backward passes. This requires the GPU to store the intermediate activations for every single micro-batch in its memory at the same time, which can quickly cause out-of-memory errors.

With the 1F1B schedule, once the pipeline is full, each GPU alternates between executing one forward pass for a new micro-batch and one backward pass for an older micro-batch. Because backward passes are completed earlier, the memory used to store activations can be freed immediately. This keeps the pipeline bubble equally small while significantly reducing the peak memory required on each GPU, allowing for the training of even larger AI models.