Model Sharding

Model sharding is the practice of partitioning a neural network's parameters, and the activations that flow through them, across multiple devices so that the full model never resides on any single machine. The technique is invoked when an individual model is larger than the memory of any one accelerator, or when an operator wants to spread a model's footprint across machines that are individually too small to load it. Sharding is the unifying term for what gets called by different names in different contexts: pipeline parallelism when layers are split sequentially across stages, tensor parallelism when the computation inside individual layers is sliced across devices, and expert parallelism when the parameters being routed to are themselves scattered. The generic point in every case is the same: the model is divided along some axis, and the pieces coordinate to produce a single forward pass.

The mechanism varies with which axis of the model gets partitioned. In pipeline parallelism, layers are grouped into stages, the stages run on different devices, and an input mini-batch is sliced into microbatches that flow from one stage to the next, with each stage computing on whatever activations it receives from its predecessors. The technique is simple in principle but introduces a bubble of idle time at the start and end of each pipeline which becomes the throughput bottleneck for deep models. In tensor parallelism, individual matrix multiplications are sliced so that different devices each compute a portion of the result and exchange partial sums over the network fabric; this requires a fast interconnect and is most often used within a single server where GPUs are tightly coupled. Expert parallelism handles the specific case of mixture-of-experts models by holding each expert on a different device and routing tokens to whichever expert is selected, with the heavy lifting done by an all-to-all communication primitive rather than the point-to-point streaming that pipeline parallelism needs. In production systems the techniques are typically combined: a model sharded across eight nodes within a server, with eight servers connected through pipeline parallelism, and experts scattered across the wider cluster.

The tradeoffs with running a single device, or with running a smaller model that does not need sharding, are well-documented. Sharding unlocks larger models than would otherwise fit, but it adds significant engineering complexity: every sharded system needs a way to coordinate the participating devices, a way to handle the failure of one shard without losing the rest, and a way to manage the heavy communication overhead that comes from streaming activations or partial sums across the network. Latency for a single inference often goes up because the work has to be split across multiple devices that have to wait on each other. Throughput per device goes down because some of the device's compute is spent on communication rather than math. The economics that justify sharding are usually one of two: either the model simply cannot fit elsewhere, or the user is willing to accept slower inference per request in exchange for serving a larger and more capable model. Distributed systems for serving frontier models increasingly rely on sharding; smaller models on commodity hardware rarely need it.

Open questions in the field are mostly about reducing the friction. The communication cost of sharded inference is dominated by the interconnect, so techniques that minimize cross-device traffic — better scheduling of computation, compression of activations, overlapping communication with computation, specialized networking hardware — are an active area of research. The question of how to shard a model across devices that are not in the same data center, with longer and more variable network latencies between them, is what several peer-to-peer inference systems are currently trying to solve, and the answer is not yet clear. Memory-efficient serving techniques like paged attention, quantization during inference, and speculative decoding interact with sharding in non-trivial ways, and the optimal policy depends heavily on the model, the workload, and the available hardware.