Tensor & Pipeline Parallelism

Overview

When data parallelism (DDP/FSDP) stops scaling— because the model is too big for one GPU, or the interconnect is too slow— you split the model itself:

• Tensor Parallelism (TP): split each layer's weight matrices across ranks; each rank computes a slice of the output. High comm per layer (all-to-all style) → needs fast NVLink.
• Pipeline Parallelism (PP): split the model by layers into stages; ranks run different stages with micro-batch pipelining. Comm is small (stage boundaries) but bubbles appear.
• PyTorch implements TP via torch.distributed.tensor.parallel (DTensor), PP via torch.distributed.pipelining (new) or torch.distributed.pipeline.sync (legacy, Pipe).

Golden rule: TP trades bandwidth for memory; PP trades latency/bubble for memory. Combine with DDP/FSDP (hybrid) for big models.

---

Tensor Parallelism— How a Linear Splits

For y = x @ W^T with 2 ranks:

• Split W: rank 0 gets W[:,:K], rank 1 gets W[:, K:] (column split).
• Each rank computes x @ W_part^T → partial y.
• All-reduce the partials → full y.

Row-split variant: for the preceding layer, split x vertically so the matmul of the next layer starts from reduced data— PyTorch's DTensor handles this transparently.

import torch, torch.nn as nn, torch.distributed as dist
from torch.distributed.tensor.parallel import parallelize_module
from torch.distributed.tensor import DeviceMesh

def main(rank, world_size):
 dist.init_process_group("nccl", rank=rank, world_size=world_size)
 torch.cuda.set_device(rank)

 mesh = DeviceMesh("cuda", list(range(world_size))) # 1 mesh dim
 model = nn.Linear(8, 8).cuda()

 # parallelize the Linear's weight by "colwise" (output dim)
 tp_model = parallelize_module(model, mesh, {"": {"linear": "ColwiseParallel"}})

 x = torch.randn(4, 8).cuda()
 out = tp_model(x)
 # out now requires all-reduce internally; result same as serial model
 torch.cuda.synchronize()
 if rank == 0:
 print("TP output works, shape:", out.shape)
 dist.destroy_process_group()

Attention-specific TP (classic LLM pattern)

OP	Split scheme
QKV projection	colwise → each rank has full Q/K/V rows
attention softmax + matmul	all-reduce partial outputs
out projection	rowwise → gather+reduce

torch.distributed.tensor.parallel provides ColwiseParallel, RowwiseParallel, SequenceParallel; always verify equivalence to the un-parallelized model on a tiny seed.

-

Pipeline Parallelism— The Bubble

Micro-batches pipelined across stages fill the pipeline; throughput gap = bubble time. Standard schedule (GPipe in PyTorch):

stream: 1F1B (one-forward-one-backward) — classic:
 f0 f1 f2 f3
 stage0: [1][2][3][4]...
Schedule example (4 stages, 4 micro-batches):
step: 1 2 3 4 5 6 7 8
s0: f1 f2 f3 f4 b4 b3 b2 b1
s1: f1 f2 f3 f4 b4 b3 b2 b1
s2: f1 f2 f3 f4 b4 b3 b2 b1
s3: f1 f2 f3 f4 b4 b3 b2 b1

The bubble = idle slots at ramp-up/down ≈ (P-1)/P of pipeline depth.

More micro-batches = smaller bubble ratio, but if you don't have enough memory to hold them, you can't win.

-

Manual PP (the honest version)

import torch, torch.nn as nn, torch.distributed as dist

def make_stage(rank):
 return nn.Sequential(nn.Linear(8, 8), nn.ReLU())

def main(rank, world_size):
 dist.init_process_group("nccl", rank=rank, world_size=world_size)
 torch.cuda.set_device(rank)
 stage = make_stage(rank).cuda()

 x = torch.randn(4, 8).cuda()
 if rank == 0:
 # forward stage 0 -> send to stage 1
 out = stage(x)
 dist.send(out.contiguous(), dst=1)
 elif rank == world_size - 1:
 # receive from previous, compute last stage
 recv = torch.empty(4, 8).cuda()
 dist.recv(recv, src=rank - 1)
 out = stage(recv)
 if rank == world_size - 1:
 print("PP forward complete, shape:", out.shape)
 else:
 recv = torch.empty(4, 8).cuda()
 dist.recv(recv, src=rank - 1)
 out = stage(recv)
 dist.send(out.contiguous(), dst=rank + 1)
 dist.destroy_process_group()

if __name__ == "__main__":
 import torch.multiprocessing as mp
 mp.spawn(main, args=(3,), nprocs=3, join=True)

Real libraries (torch.distributed.pipelining.Pipeline, torchpipe, DeepSpeed PP) implement 1F1B + scheduling + backward for you— manual send/recv above is for understanding, not production.

-

Hybrid Parallelism (TP × PP × DP)— the LLM recipe

mesh: [DP=2, PP=2, TP=2]
rank layout: dp-rank * (pp * tp)... organized as a 3D mesh
- FSDP/DDP over DP dim, pipeline over PP, TensorParallel over TP dim
- Communication: TP on NVLink; PP on node-links; DP across nodes

from torch.distributed.tensor import DeviceMesh
mesh = DeviceMesh("cuda", list(range(8)), mesh_dim_names=["dp", "pp", "tp"])
# PP and TP use DeviceMesh from torch.distributed.tensor.parallel / pipelining

Practical sizing: TP ≤ 8 (NVLink width), PP = #stages/GPU count, DP fills the rest. Keep TP fastest interconnect, PP to the next, DP slowest.

---

Which to Use When (Decision Table)

Situation	Strategy	Why
Model fits one GPU	DDP/FSDP	simplest
Model > 1 GPU memory (params)	FSDP or TP	shard params
Model huge, nodes slow	TP + PP + FSDP	hybrid
Latency-sensitive inference	TP (not PP)	no bubble
Training on 2 GPUs, small model	DDP	TP overhead not worth it

---

Verification & Debugging

1. Numerical check: TP must equal the serial model (same seed)— any mismatch = split bug.
2. Comm %: torch.profiler— if all-reduce >30% of step, interconnect-bound: reduce TP size or use FSDP instead.
3. Bubble measuring: PP idle kernel slots in the timeline— increase micro-batches.
4. Use TORCH_DISTRIBUTED_DEBUG=DETAIL to spot deadlocks and shape mismatches.

-

Key Takeaways

• TP splits weight matrices; comm per layer; needs fast interconnect; verify numerics.
• PP splits by layers; small comm, bubble cost; micro-batches fill the pipe.
• Hybrid (DP × PP × TP) is the production LLM pattern— arrange mesh dims by speed.
• Never reimplement send/recv in prod— use torch.distributed.pipelining / tensor.parallel APIs.
• Measure comm% vs compute% before and after every parallelizaton change.

-

Related Topics

• Ddp
• Fsdp
• [Checkpointing](/06-pytorch/05-distributed-training/(04-checkpointing-fault-tolerance/)