Model parallelism

What model parallelism is

Model parallelism trains a single neural network across several accelerators by dividing the model itself — its parameters and the computation that applies them — instead of dividing the training data. In the more familiar data-parallel setup, every device holds a complete copy of the model and works on a different batch of examples; in model parallelism, each device holds a different slice of the model, and the work of a single layer is shared out among them. The point is capacity: some models are too large to fit on one device, and partitioning is what makes training them possible at all.^[1]

Why memory, not arithmetic, is the constraint

The Megatron-LM paper frames memory rather than raw arithmetic throughput as the binding constraint on very large models.^[1] This is a claim about where the wall sits. Adding more floating-point operations per second to one accelerator does not help if the parameters, their gradients, the optimizer state and the intermediate activations all have to be resident at the same time. Once the weights no longer fit, training cannot proceed however fast the hardware is, and the way forward is to keep different parts of the model — and the activations produced inside a layer — on different devices.^[1]

How the split is made

The 2019 Megatron-LM paper from NVIDIA set out an intra-layer implementation in which the matrix multiplications inside a transformer layer are partitioned across GPUs.^[1] The matrix multiplication is the natural unit to cut. Picture a weight matrix divided into blocks: each device multiplies its block by the incoming vector, and the partial results are then combined. What is being split is the arithmetic of one layer itself, not the stack of layers.

What made this implementation notable was its engineering footprint. It needs no new compiler and no new library; it can be built by inserting a few communication operations into native PyTorch code.^[1] The tensor program stays recognisable, and the parallel structure is expressed as a handful of sends and receives rather than as a rewritten framework.

The paper also records a numerical detail that grows in importance with size: the placement of layer normalization matters for BERT-like models as they scale, because it changes how the partitioned layers behave numerically.^[1]

Relation to pipeline parallelism

Model parallelism comes in more than one shape. Pipeline model parallelism places different layers on different devices and streams activations from one to the next. The paper describes its own intra-layer method as orthogonal and complementary to pipeline model parallelism, so the two can be combined.^[1] Splitting within a layer and splitting across layers answer different questions — how to fit one layer's arithmetic, and how to fit a whole stack of layers — and a training system can do both at once.

Reported results

The implementation was used with 512 GPUs to converge transformer models of up to 8.3 billion parameters.^[1] The training run sustained 15.1 PetaFLOPs across the whole application, which worked out to a scaling efficiency of 76 percent against a strong single-GPU baseline.^[1] Scaling efficiency measures the share of the ideal speedup that a multi-device run actually delivers; because devices must exchange results, measured throughput falls short of simply multiplying single-device performance by the number of devices.

Limits

The cost of model parallelism is written into its definition. If one layer is computed by several devices together, those devices must keep exchanging intermediate results, and that communication is an unavoidable overhead. It addresses whether a model fits, not whether the arithmetic becomes cheaper: partitioning does not reduce the number of operations, it only spreads them. The paper's remark about layer normalization points to a second, quieter cost — once a layer is cut apart, its numerical behaviour has to be watched separately.^[1]