A model’s weights, activations, and data are stored as tensors: containers ranging from single numbers to multidimensional arrays. Sharding splits tensor contents across devices, so one GPU need not store everything. AlexNet used hand-built sharding, an early response to limits in GPU memory and interconnect bandwidth. section 6.3 section 6.3
Tensor parallelism partitions tensor axes across devices. Mesh-TensorFlow made the choice explicit: splitting the vocabulary or feed-forward dimension is tensor parallelism, while splitting the batch is data parallelism. Pipeline parallelism instead divides computation into stages; it can be combined with tensor parallelism, with each pipeline stage distributed across devices. section 6.3.3 section 6.3.3
From hand-built slices to large-scale training
Megatron-LM’s tensor slicing was designed for a particular transformer layer. Mesh-TensorFlow generalized the approach by mapping tensor axes onto device-mesh axes. Megatron-LM later trained an 8.3-billion-parameter GPT-2-style transformer on 512 GPUs, demonstrating the scale these parallelism strategies could support. section 6.3.3 section 6.3.3
Distributing computation does not by itself eliminate memory redundancy: data-parallel GPUs still hold full model weights. Tensor parallelism addresses how tensor computations are distributed, while ZeRO targets duplicate storage of parameters, gradients, and optimizer states. section 6.3.3 section 6.3.3
