AI Model Wiki

    Skip to main content
    Back to Glossary

    Pipeline Parallelism

    Pipeline parallelism is a strategy for training or running neural networks that are too large to fit on a single accelerator by splitting the model's layers into sequential stages, each assigned to a different device. Device 1 holds the first set of layers, device 2 holds the next, and so on. During the forward pass, activations flow from one device to the next like items on an assembly line; during the backward pass, gradients flow in reverse. Naive pipeline parallelism leaves most devices idle while one runs its stage (the "bubble" problem), so modern implementations use micro-batching — splitting a mini-batch into smaller chunks that keep all pipeline stages busy in an overlapping schedule.

    Think of it like a car assembly line where each station performs one step: welding at station 1, painting at station 2, interior fitting at station 3. No single station builds the whole car, but all stations work simultaneously on different cars at different stages, so the factory produces vehicles continuously rather than waiting for each car to fully complete before starting the next.

    Google's GPipe paper introduced the foundational micro-batch schedule, and Megatron-LM combines pipeline parallelism with tensor parallelism for models with hundreds of billions of parameters:

    # Conceptual pipeline stage assignment in PyTorch
    import torch.nn as nn
    stage_0 = nn.Sequential(layer0, layer1, layer2).to("cuda:0")
    stage_1 = nn.Sequential(layer3, layer4, layer5).to("cuda:1")
    
    x = stage_0(inputs)         # runs on GPU 0
    x = x.to("cuda:1")
    output = stage_1(x)         # runs on GPU 1
    

    The key engineering challenge is balancing compute across stages so that no single device becomes the bottleneck, and minimizing the GPU memory holding intermediate activations at stage boundaries.

    Pipeline parallelism complements tensor parallelism (splitting individual layer computations across devices) and data parallelism (replicating the full model and splitting the data). Together they enable training the largest Foundation Model and LLM architectures. The approach is inherently tied to Transformer layer structure, which partitions cleanly into stages.

    You will encounter pipeline parallelism when reading about how large models like GPT-4 or Gemini are trained across hundreds or thousands of GPUs, or when setting up multi-node training for your own large-scale models using frameworks like DeepSpeed or Megatron.

    AI Terminology
    Glossary