Week 13 — Tensor Parallelism

Learning Objectives

Key Concepts

Setup & Configuration

# Verify 2× H100 GPUs are visible
nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader

# Check NVLink connectivity
nvidia-smi topo -m

# Confirm vLLM installation supports multi-GPU
python -c "import vllm; print(vllm.__version__)"

# Download model weights (if not cached)
python -c "from huggingface_hub import snapshot_download; snapshot_download('meta-llama/Llama-3.1-8B-Instruct')"

Experiments

for tp in 1 2; do
  python benchmarks/benchmark_latency.py \
    --model meta-llama/Llama-3.1-8B-Instruct \
    --tensor-parallel-size $tp \
    --batch-size 1 \
    --input-len 512 \
    --output-len 128 \
    --num-iters 20 \
    2>&1 | tee results_tp${tp}_bs1_latency.txt
done

# Also test input-heavy (prefill-bound) workload
for tp in 1 2; do
  python benchmarks/benchmark_latency.py \
    --model meta-llama/Llama-3.1-8B-Instruct \
    --tensor-parallel-size $tp \
    --batch-size 1 \
    --input-len 2048 \
    --output-len 32 \
    --num-iters 10 \
    2>&1 | tee results_tp${tp}_prefill_latency.txt
done

# Start TP=1 server
vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --tensor-parallel-size 1 \
  --port 8000 \
  --disable-log-requests &
sleep 30

python benchmarks/benchmark_serving.py \
  --backend vllm \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --dataset-name sharegpt \
  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
  --num-prompts 200 --request-rate 4 \
  --port 8000 \
  2>&1 | tee results_serving_tp1.txt
kill %1; sleep 10

# Start TP=2 server
vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --tensor-parallel-size 2 \
  --port 8001 \
  --disable-log-requests &
sleep 30

python benchmarks/benchmark_serving.py \
  --backend vllm \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --dataset-name sharegpt \
  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
  --num-prompts 200 --request-rate 4 \
  --port 8001 \
  2>&1 | tee results_serving_tp2.txt
kill %1

for bs in 1 4 8 16 32 64; do
  for tp in 1 2; do
    python benchmarks/benchmark_throughput.py \
      --model meta-llama/Llama-3.1-8B-Instruct \
      --tensor-parallel-size $tp \
      --batch-size $bs \
      --input-len 512 \
      --output-len 128 \
      --num-iters 5 \
      2>&1 | tee results_tp${tp}_bs${bs}_throughput.txt
  done
done

nsys profile \
  -o tp2_trace \
  --trace=cuda,nvtx,nccl \
  --capture-range=cudaProfilerApi \
  python benchmarks/benchmark_latency.py \
    --model meta-llama/Llama-3.1-8B-Instruct \
    --tensor-parallel-size 2 \
    --batch-size 1 \
    --input-len 256 \
    --output-len 32 \
    --num-iters 3

# Open in nsys GUI: nsys-ui tp2_trace.nsys-rep
# Look for: ncclAllReduce kernel launches
# Count: number of all-reduces per decode step (should be 2× num_layers)
# Measure: duration of each all-reduce event

# Monitor NVLink bandwidth while benchmark runs
nvidia-smi nvlink --status -i 0 &
nvidia-smi dmon -s u -d 1 &

python benchmarks/benchmark_latency.py \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --tensor-parallel-size 2 \
  --batch-size 1 \
  --input-len 512 --output-len 64 --num-iters 10

Experiment Results

Hardware

Experiments run on four PACE Phoenix configurations: 2× NVIDIA H100 80GB HBM3 (NVLink), 2× NVIDIA H200 141GB HBM3e (NVLink), 4× NVIDIA H100 80GB HBM3 (full NVLink mesh), and 2× NVIDIA A100 80GB PCIe (NVLink bridge). Models: Llama-3.1-8B-Instruct (small-model TP control, all configs) and Llama-3.1-70B-Instruct (REQUIRES TP, H100/H200 configs). Both BF16.

TP=1 vs TP=2 on 8B — Throughput & Serving (H100 + A100)

Both H100 and A100 show the same pattern at 8B: offline throughput is unchanged (model fits on one GPU, no extra HBM bandwidth unlocked) but serving latency drops significantly because TP=2 halves the per-step GEMM tile size. A100 PCIe wins more (-39%) than H100 (-29%) because A100's slower memory bus makes the per-step GEMM a proportionally bigger share of step time, so halving it has a larger effect.

Batch Size Scaling Under TP=2 — H100 vs A100

A100 latency is roughly 1.94× H100 across all batch sizes — close to the H100/A100 HBM bandwidth ratio \((3.35 / 1.935 \approx 1.73)\), plus a small extra penalty for slower NVLink-bridge all-reduce vs full NVLink. The shape (latency growing only modestly with batch size) is identical on both hardware classes — confirming the memory-bandwidth-bound decode regime is GPU-architecture-independent at 8B scale.

Charts

Figure 1: Left — 8B serving TTFT and ITL on 1× vs 2× H100 and 1× vs 2× A100. TP=2 cuts TTFT and ITL on BOTH platforms (H100 −29%, A100 −39%) even though offline throughput is unchanged. Right — TP=2 batch size scaling on H100 and A100: A100 throughput is consistently ~52% of H100, matching the HBM bandwidth ratio.

Llama-3.1-70B Tensor Parallelism (the real use case)

To see why TP exists, we ran the same TP=2 setup on the much larger Llama-3.1-70B (140 GB BF16). This model cannot fit on a single H100 (80 GB), so TP is mandatory, not optional.

70B TP=2 vs TP=4 batch scaling on H100×4

Doubling TP from 2 to 4 on the same 4-GPU node yields 1.31–1.33× speedup — far below the 2× theoretical ceiling because (a) all-reduce traffic doubles (4-way ring vs 2-way) and (b) per-GPU GEMM tile size halves, dropping arithmetic intensity. Still, for workloads that need lowest per-token latency and fit in 4 GPUs, TP=4 is the right call.

70B batch size scaling on H200×2

Latency only grows ~5% from bs=1 to bs=8 — even on a 70B model with TP=2, batching is essentially free in the memory-bound decode regime. Throughput grows nearly 8× linearly.

Figure 2: 8B vs 70B at increasing TP. 8B (TP=1 or TP=2 on H100) is bandwidth-limited at ~7 ms/tok and ~11K tok/s. 70B is bandwidth-limited at TP=2 (29.9 ms on H100, 24.4 ms on H200), and TP=4 on H100×4 brings it down to 22.2 ms/tok and 1136 tok/s — a 1.35× speedup over TP=2 on the same hardware (well below the 2× ceiling because of the heavier 4-way ring all-reduce).

Expected vs Actual

Metrics to Collect

Source Code Reading

# Per decode step (simplified) for one layer:
# 1. Input token embedding (shape: [batch, hidden]) — replicated on both GPUs
# 2. Q/K/V projection: ColumnParallelLinear splits output heads
#    GPU0: W_q[:, :heads/2], GPU1: W_q[:, heads/2:]
# 3. Attention: each GPU computes on its half of heads
# 4. Output projection: RowParallelLinear, requires ALL-REDUCE
# 5. MLP gate/up: ColumnParallelLinear (no comm needed)
# 6. MLP down: RowParallelLinear, requires ALL-REDUCE
# Result: 2 all-reduces per layer × 32 layers = 64 all-reduces per step

Core Source Code Walkthrough

Below are real excerpts from vllm-continuum that implement the concepts you measured. Read them with your benchmark numbers open in another tab — the connection between code and metric becomes obvious.

vllm/model_executor/layers/linear.py:371 — ColumnParallelLinear

class ColumnParallelLinear(LinearBase):
    """Linear layer with column parallelism.
    The linear layer is defined as Y = XA + b. A is parallelized along
    its second dimension as A = [A_1, ..., A_p].
    """
    def __init__(self, ...):
        # Divide the weight matrix along the last dimension.
        self.tp_rank = get_tensor_model_parallel_rank()
        self.tp_size = get_tensor_model_parallel_world_size()
        self.input_size_per_partition = input_size
        self.output_size_per_partition = divide(output_size, self.tp_size)

This is how Tensor Parallelism is actually implemented. For each Linear layer \(Y = XA\), the matrix \(A\) is split column-wise across tp_size GPUs. Each GPU computes its slice \(Y_i = X \cdot A_i\) in parallel. After the operation, an all_reduce (in MLP) or all_gather (when output needs gathering) combines partial results. This is why our 70B TP=2 needs 32 layers × 2 all-reduces per decode step.

Written Analysis — Reference Answers

Answers below cite the measured numbers from the Results section: 8B TP=1 vs TP=2 on 1–2× H100, plus 70B TP=2 on 2× H100 (29.90 ms/tok, 818 tok/s) and 2× H200 (24.43 ms/tok, 995 tok/s).

Q1: Does TP=2 reduce decode latency by exactly 2×?

Q2: What % of decode time is NCCL all-reduce?

Q3: At what batch size does TP=2 beat TP=1?

Q4: NVLink bandwidth utilization during all-reduce