N-dimensional parallelism combines several ways of distributing training work instead of relying on one strategy alone. Data parallelism replicates models and splits data; model parallelism divides model computation and parameters; pipeline parallelism assigns sequential stages; tensor parallelism partitions layers or tensors. section 16.12 section 16.6
The dimensions can be matched to the hardware. Tensor parallelism usually stays within one host, where fast GPU interconnects reduce its communication cost, while pipeline parallelism can distribute stages across hosts. Combining dimensions also reduces bottlenecks such as increasingly inefficient allreduce operations at very large scale. section 16.6
A practical n-dimensional design may combine:
- Data parallelism for multiple model replicas processing different data batches. section 16.12
- Model or tensor parallelism for distributing parameters, computation, or individual layers across devices. section 16.12
- Pipeline parallelism for processing different microbatches concurrently through sequential model stages. section 16.12
- Context or other specialized dimensions when scaling large language models. section 16.12
A two-dimensional device mesh
Represent two model replicas, each split across four model partitions, with a 2 × 4 mesh. One mesh dimension groups replica devices for parameter synchronization, while the other selects devices for model-parallel computation. PyTorch’s DeviceMesh creates the corresponding process groups and communication patterns. LLaMA 3 demonstrates the larger idea with 4D parallelism combining data, tensor, pipeline, and context parallelism. section 16.6 section 16.6
Figure 16.9 Visualization of a 2 × 4 device mesh, where the rows represent model replicas and the columns represent model partitions. The dashed arrows represent how each device is communicating.
Visualize how rows can represent model replicas and columns can represent model partitions, with communication occurring along each mesh dimension. section 16.6
