AI Model Wiki

    Skip to main content
    Back to Glossary

    Model Sharding

    Model sharding is a parallelism strategy that partitions a large neural network's layers or parameter tensors across multiple accelerators (GPUs or TPUs) so that no single device needs to hold the entire model in memory. It is distinct from data parallelism, where each device holds a full copy of the model and processes a different slice of the data; with model sharding, each device owns a slice of the model itself and must exchange activations or gradients with adjacent devices during forward and backward passes.

    Think of it like a relay race where each runner only carries the baton for one leg of the course. No single runner needs the stamina to complete the whole distance - they just hand off and pick up from where the previous runner stopped.

    A straightforward example is pipeline parallelism, one common form of model sharding:

    # Conceptual pipeline sharding with PyTorch RPC
    # Layer 0-11 on GPU 0, layers 12-23 on GPU 1
    model_shard_0 = GPT2Layers(0, 12).to("cuda:0")
    model_shard_1 = GPT2Layers(12, 24).to("cuda:1")
    
    def forward(x):
        mid = model_shard_0(x.to("cuda:0"))
        return model_shard_1(mid.to("cuda:1"))
    

    Tensor parallelism is a finer-grained variant where individual weight matrices are split column- or row-wise across devices, as used in Megatron-LM. Related concepts include the Transformer architecture that produces the oversized weight matrices requiring sharding, Parameter counts that drive memory requirements, Inference serving considerations where sharding reduces per-device VRAM, and Mixture of Experts which combines sparse routing with sharding to scale further.

    You will encounter model sharding whenever reading about training or deploying frontier-scale LLMs such as GPT-4, Llama 3, or Gemini, where single-GPU memory is orders of magnitude smaller than the full parameter footprint.

    AI Terminology
    Glossary