Training Phase

During training, tensor parallelism splits each linear layer’s weight matrix across the tensor‑parallel ranks. The forward pass first all‑gather the input activation shards, then each rank performs a local matrix multiplication, and the backward pass uses a reduce‑scatter to combine gradient shards. The overall per‑layer communication pattern is all‑gather → compute → reduce‑scatter. This pattern appears for every transformer block and adds a bandwidth cost that can be expressed as $$T_{\text{train}} = T_{\text{comp}} + 2(\alpha + \beta \cdot \frac{M}{P})$$ where $M$ is the matrix dimension, $P$ the tensor‑parallel size, $\alpha$ the latency term and $\beta$ the inverse bandwidth.

In addition, after the backward pass an all‑reduce across the data‑parallel group synchronises the weight gradients, further increasing communication.

Inference Phase

Inference only needs the forward pass. Many engines fuse the all‑gather with the matmul or pre‑shard the weights so that each rank can compute its output shard directly. The final logits are then assembled with a single all‑reduce (or an all‑gather of the output) instead of a reduce‑scatter. The communication cost drops to $$T_{\text{infer}} = T_{\text{comp}} + (\alpha + \beta \cdot \frac{M}{P})$$ and latency becomes the dominant metric. Because gradients are not exchanged, the bandwidth overhead is roughly halved.

Implementation details differ as well: training frameworks (Megatron‑LM, DeepSpeed‑ZeRO) use torch.distributed.all_gather and torch.distributed.reduce_scatter kernels for each layer, while inference engines (DeepSpeed‑Inference, TensorRT‑LLM) replace them with fused all‑gather/all‑reduce kernels and often keep the sharded weights in FP16/BF16 to reduce memory traffic. Memory usage is also lower at inference because activation checkpoints are not needed.

Performance trade‑offs: training is bandwidth‑bound and scales with the number of tensor‑parallel ranks until the all‑gather/reduce‑scatter overhead dominates; inference is latency‑bound, so reducing the number of collective calls (e.g., by fusing or by using pipeline parallelism) yields higher throughput. Choosing the tensor‑parallel size therefore balances compute‑to‑communication ratio differently for training (favoring larger batches) and inference (favoring minimal per‑token latency).

6. Tensor Parallelism Comparison

Similar Questions