Tensor Parallelism: Splitting a Model Across GPUs

A 70B model in FP16 (16-bit floating point) requires roughly 130 GB of memory. No single A100–80GB card can hold it. The intro and memory table use GiB (params × 2 / 102⁴³); the recommender uses decimal GB (params × 2 / 1⁰⁹). The two conventions differ by ~7% — 70B is 130.4 GiB ≈ 140 GB, and 405B is 754.4 GiB ≈ 810 GB. The naive solution — data parallelism — doesn’t help because data parallelism replicates the model, not shards it. Tensor parallelism (introduced at scale by Megatron-LM, Shoeybi et al. 2019) shards individual weight matrices across GPUs: each GPU holds 1/N of every layer, computes its partial result, and synchronizes via AllReduce. The efficiency depends almost entirely on interconnect bandwidth. At NVLink 600 GB/s, 8-GPU tensor parallelism achieves 93.5% scaling efficiency on a 7B model at seq_len=2048; at PCIe 4.0, it drops to 74.9%.

Disclaimer: The opinions expressed in this article are my own and do not represent the views of Google. This content is based solely on publicly available information.

How Tensor Parallelism Shards Weight Matrices

Every transformer layer is dominated by matrix multiplications of the form Y = XW, where X is the activation matrix [batch × seq × d_in] and W is the weight matrix [d_in × d_out]. Tensor parallelism exploits a basic property of matrix multiplication: you can split W along either axis and recombine the results exactly.

Column-parallel split: Partition W by output columns. Each GPU i holds W[:, i*k:(i+1)*k] and computes x @ W_shard, producing a partial output slice. The slices are disjoint, so no communication is needed — only an AllGather to concatenate the output columns. This is used for the QKV (query/key/value projections) (split the output head dimension) and the first FFN (feed-forward network) layer (split the hidden dimension).

Row-parallel split: Partition W by input rows. Each GPU i holds W[i*k:(i+1)*k, :] and receives the matching rows of X. Each GPU computes a partial product; the full output is the sum of all partial products. An AllReduce (sum) is required to combine them. This is used for the output projection and the second FFN layer. Together, a column-parallel layer followed by a row-parallel layer creates one complete tensor-parallel block with a single AllReduce at the boundary.

The code below implements both patterns using NumPy to simulate what each GPU computes, then verifies numerical equivalence with a single-GPU reference and tabulates how much memory each model size requires per GPU at different TP degrees. (TFLOPS — tera floating-point operations per second — appears in the compute calculations later.)

import numpy as np

def column_parallel_linear(x: np.ndarray, w: np.ndarray, n_gpus: int,
                             bias: np.ndarray | None = None) -> np.ndarray:
    """
    Column-parallel linear: split W columns across GPUs.
    Each GPU computes x @ W_shard, then AllGather (concatenate) to get full output.
    """
    out_features = w.shape[1]
    shard_size = out_features // n_gpus

    full_output = np.zeros((x.shape[0], out_features))

    for gpu_id in range(n_gpus):
        col_start = gpu_id * shard_size
        col_end = col_start + shard_size
        w_shard = w[:, col_start:col_end]  # each GPU holds these columns
        partial = x @ w_shard              # local matmul
        full_output[:, col_start:col_end] = partial  # AllGather (not AllReduce for column-parallel)

    if bias is not None:
        full_output += bias

    return full_output

def row_parallel_linear(x: np.ndarray, w: np.ndarray, n_gpus: int) -> np.ndarray:
    """
    Row-parallel linear: split W rows (and x columns) across GPUs.
    Each GPU computes x_shard @ W_shard, then AllReduce (sum) partial outputs.
    """
    in_features = w.shape[0]
    shard_size = in_features // n_gpus

    partial_sum = np.zeros((x.shape[0], w.shape[1]))

    for gpu_id in range(n_gpus):
        row_start = gpu_id * shard_size
        row_end = row_start + shard_size
        x_shard = x[:, row_start:row_end]  # each GPU holds these input features
        w_shard = w[row_start:row_end, :]  # each GPU holds these weight rows
        partial_sum += x_shard @ w_shard   # AllReduce sum across GPUs

    return partial_sum

if __name__ == "__main__":
    rng = np.random.default_rng(42)
    batch, d_in, d_out = 4, 4096, 4096

    x = rng.normal(0, 0.01, (batch, d_in)).astype(np.float32)
    w = rng.normal(0, 0.01, (d_in, d_out)).astype(np.float32)

    # Reference: full matmul on one GPU
    y_ref = x @ w

    print("=== Tensor Parallel Correctness Check ===")
    for n_gpus in [1, 2, 4, 8]:
        y_col = column_parallel_linear(x, w, n_gpus)
        y_row = row_parallel_linear(x, w, n_gpus)

        err_col = np.max(np.abs(y_col - y_ref))
        err_row = np.max(np.abs(y_row - y_ref))

        print(f"  n_gpus={n_gpus}:  column-parallel max_err={err_col:.2e}  "
              f"row-parallel max_err={err_row:.2e}")

    print()
    print("=== Memory Per GPU at Different TP Degrees ===")
    model_sizes = {"7B": 7e9, "13B": 13e9, "70B": 70e9, "405B": 405e9}
    for model_name, n_params in model_sizes.items():
        total_gb = n_params * 2 / 1024**3  # FP16
        print(f"\n  {model_name} ({total_gb:.1f} GB FP16):")
        for n_gpus in [1, 2, 4, 8, 16]:
            per_gpu = total_gb / n_gpus
            fits_a100 = "✓" if per_gpu <= 80 else "✗"
            print(f"    {n_gpus} GPUs: {per_gpu:6.1f} GB/GPU  A100-80GB: {fits_a100}")

Output:

=== Tensor Parallel Correctness Check ===
  n_gpus=1:  column-parallel max_err=0.00e+00  row-parallel max_err=0.00e+00
  n_gpus=2:  column-parallel max_err=0.00e+00  row-parallel max_err=6.15e-08
  n_gpus=4:  column-parallel max_err=0.00e+00  row-parallel max_err=4.77e-08
  n_gpus=8:  column-parallel max_err=0.00e+00  row-parallel max_err=5.44e-08

=== Memory Per GPU at Different TP Degrees ===

7B (13.0 GB FP16):
    1 GPUs:   13.0 GB/GPU  A100-80GB: ✓
    2 GPUs:    6.5 GB/GPU  A100-80GB: ✓
    4 GPUs:    3.3 GB/GPU  A100-80GB: ✓
    8 GPUs:    1.6 GB/GPU  A100-80GB: ✓
    16 GPUs:    0.8 GB/GPU  A100-80GB: ✓

13B (24.2 GB FP16):
    1 GPUs:   24.2 GB/GPU  A100-80GB: ✓
    2 GPUs:   12.1 GB/GPU  A100-80GB: ✓
    4 GPUs:    6.1 GB/GPU  A100-80GB: ✓
    8 GPUs:    3.0 GB/GPU  A100-80GB: ✓
    16 GPUs:    1.5 GB/GPU  A100-80GB: ✓

70B (130.4 GB FP16):
    1 GPUs:  130.4 GB/GPU  A100-80GB: ✗
    2 GPUs:   65.2 GB/GPU  A100-80GB: ✓
    4 GPUs:   32.6 GB/GPU  A100-80GB: ✓
    8 GPUs:   16.3 GB/GPU  A100-80GB: ✓
    16 GPUs:    8.1 GB/GPU  A100-80GB: ✓

405B (754.4 GB FP16):
    1 GPUs:  754.4 GB/GPU  A100-80GB: ✗
    2 GPUs:  377.2 GB/GPU  A100-80GB: ✗
    4 GPUs:  188.6 GB/GPU  A100-80GB: ✗
    8 GPUs:   94.3 GB/GPU  A100-80GB: ✗
    16 GPUs:   47.1 GB/GPU  A100-80GB: ✓

Column-parallel reproduces the full matmul exactly (each GPU produces a disjoint output slice that is then concatenated). Row-parallel sums partial products in a different order, so floating-point non-associativity yields a small residual error around 5e-8 — well below FP16 precision and harmless for inference. The memory table makes the cost of single-GPU serving concrete: a 70B model overflows an 80 GB card by 50 GB, and a 405B model needs at least 10 cards just for weights at FP16; the recommender below adds an 85% headroom factor for activations and KV-cache (cached key/value tensors from previous decode steps), raising the minimum to 12. Sharding is not optional at frontier scale.

Figure 1: Per-GPU memory by TP degree (left) and TP vs pipeline-parallel overhead by batch size (right).

The left panel of Figure 1 shows the same dynamic: the 70B line crosses the 80 GB A100 limit between 1 and 2 GPUs, while a 175B model needs four cards minimum. The right panel previews a result derived later — tensor parallelism’s overhead is roughly batch-independent, whereas pipeline-parallel bubbles shrink as batch size grows. Both panels frame the rest of the article: memory sets the floor on GPU count, and bandwidth sets the ceiling on efficiency.

AllReduce Communication Cost

After each row-parallel layer, every GPU holds a partial sum of the output activation. An AllReduce combines these partial sums into the full activation on every GPU. The dominant algorithm is ring-AllReduce, which proceeds in two phases:

1. Reduce-scatter: the ring of N GPUs exchanges chunks in a round-robin; after N-1 steps each GPU holds 1/N of the fully reduced output, having sent (N-1)/N × D bytes.
2. AllGather: another N-1 steps broadcast each chunk around the ring, sending another (N-1)/N × D bytes.

Total bytes sent per GPU = 2 × (N-1)/N × D, which approaches 2D as N grows. The wall-clock time is 2(N-1)/N × D / BW where BW is the unidirectional link bandwidth. For NVLink 4.0 at 600 GB/s and a 16 MB activation tensor: 2 × 7/8 × 0.016 / 600 ≈ 0.047 ms — negligible. For PCIe 4.0 at 32 GB/s the same tensor takes ≈ 0.875 ms, comparable to the compute time of a single transformer layer at batch=1.

The following code applies these formulas directly, sweeping across all five common interconnect classes and six GPU counts to quantify exactly how much overhead each configuration incurs.

import numpy as np

def allreduce_cost_ms(tensor_gb: float, bandwidth_gb_s: float, n_gpus: int) -> float:
    """
    AllReduce cost using ring-allreduce algorithm.
    Ring AllReduce sends 2(N-1)/N × data per GPU, running at peak bandwidth.
    """
    # Ring AllReduce: each GPU sends (N-1)/N × data in reduce-scatter,
    # then (N-1)/N × data in allgather. Total = 2(N-1)/N × data.
    data_sent_gb = 2 * (n_gpus - 1) / n_gpus * tensor_gb
    time_s = data_sent_gb / bandwidth_gb_s
    return time_s * 1000  # ms

def compute_cost_ms(batch: int, seq_len: int, d_model: int,
                     tflops: float) -> float:
    """
    Forward pass compute cost for one transformer layer on one GPU.
    Each layer: ~6 × B × S × D² FLOPs (QKV, O proj, 2 FFN)
    *Excludes the O(S²·D) attention matmuls — at S=2048, D=4096 they're a small fraction of the matmul total.*
    """
    flops = 6 * batch * seq_len * d_model ** 2
    time_s = flops / (tflops * 1e12)
    return time_s * 1000  # ms

def scaling_efficiency(n_gpus: int, allreduce_ms: float, compute_ms: float) -> float:
    """
    Scaling efficiency = ideal speedup / actual speedup.
    Actual speedup is limited by communication overhead.
    """
    total_ms = compute_ms + allreduce_ms  # communication runs after compute
    ideal_ms = compute_ms                  # perfect scaling, no comm cost
    actual_speedup = compute_ms / total_ms * n_gpus  # comp/total × n_gpus
    return actual_speedup / n_gpus        # fraction of ideal

if __name__ == "__main__":
    # A100: 312 TFLOPS FP16, 80GB
    GPU_TFLOPS = 312.0
    BATCH, SEQ, D_MODEL = 1, 2048, 4096

    compute_per_layer_ms = compute_cost_ms(BATCH, SEQ, D_MODEL, GPU_TFLOPS)
    n_layers = 32
    compute_total_ms = compute_per_layer_ms * n_layers

    # AllReduce tensor: activations per layer (batch × seq × d_model)
    activation_gb = BATCH * SEQ * D_MODEL * 2 / 1024**3  # BF16

    print("=== AllReduce Cost vs Interconnect (8 GPUs) ===")
    print(f"Compute per 32-layer model (batch=1): {compute_total_ms:.1f} ms")
    print(f"Activation size per AllReduce: {activation_gb*1024:.1f} MB")
    print()

    interconnects = [
        ("PCIe 4.0", 32),
        ("PCIe 5.0", 64),
        ("NVLink 3.0", 300),
        ("NVLink 4.0 (A100)", 600),
        ("NVLink 5.0 (H100)", 1800),
    ]

    n_gpus = 8
    print(f"{'Interconnect':<22} {'BW (GB/s)':>10} {'Comm/layer (ms)':>17} "
          f"{'Comm total (ms)':>17} {'Comm %':>8} {'Efficiency':>12}")
    print("-" * 92)
    for name, bw in interconnects:
        comm_per_layer = allreduce_cost_ms(activation_gb, bw, n_gpus)
        comm_total = comm_per_layer * n_layers
        total_ms = compute_total_ms + comm_total
        comm_pct = comm_total / total_ms * 100
        eff = scaling_efficiency(n_gpus, comm_total, compute_total_ms)
        speedup = eff * n_gpus
        print(f"{name:<22} {bw:>10} {comm_per_layer:>16.2f}  {comm_total:>15.1f}  "
              f"{comm_pct:>6.1f}%  {eff*100:>9.1f}% ({speedup:.2f}×)")

    print()
    print("=== Efficiency vs GPU Count (NVLink 4.0) ===")
    bw_nvlink = 600
    print(f"{'N GPUs':>8} {'Speedup':>10} {'Efficiency':>12} {'Comm %':>8}")
    print("-" * 42)
    for n in [1, 2, 4, 8, 16, 32]:
        if n == 1:
            print(f"{n:>8} {'1.00×':>10} {'100.0%':>12} {'0.0%':>8}")
            continue
        comm = allreduce_cost_ms(activation_gb, bw_nvlink, n) * n_layers
        eff = scaling_efficiency(n, comm, compute_total_ms)
        speedup = eff * n
        comm_pct = comm / (compute_total_ms + comm) * 100
        print(f"{n:>8} {speedup:>9.2f}×  {eff*100:>11.1f}%  {comm_pct:>7.1f}%")

Output:

=== AllReduce Cost vs Interconnect (8 GPUs) ===
Compute per 32-layer model (batch=1): 21.1 ms
Activation size per AllReduce: 16.0 MB

Interconnect            BW (GB/s)   Comm/layer (ms)   Comm total (ms)   Comm %   Efficiency
--------------------------------------------------------------------------------------------
PCIe 4.0                       32             0.85             27.3    56.4%       43.6% (3.49×)
PCIe 5.0                       64             0.43             13.7    39.3%       60.7% (4.86×)
NVLink 3.0                    300             0.09              2.9    12.1%       87.9% (7.03×)
NVLink 4.0 (A100)             600             0.05              1.5     6.5%       93.5% (7.48×)
NVLink 5.0 (H100)            1800             0.02              0.5     2.2%       97.8% (7.82×)

=== Efficiency vs GPU Count (NVLink 4.0) ===
  N GPUs    Speedup   Efficiency   Comm %
------------------------------------------
       1      1.00×       100.0%     0.0%
       2      1.92×         96.2%      3.8%
       4      3.78×         94.4%      5.6%
       8      7.48×         93.5%      6.5%
      16     14.90×         93.1%      6.9%
      32     29.73×         92.9%      7.1%

At seq=2048 with the real FLOPs estimate for compute, the NVLink 4.0 efficiency lands around 93.5%. The lesson is the same as the shorter-sequence regime: PCIe 4.0 is a non-starter for tensor parallelism — communication consumes more than half the wall-clock time, dropping the 8-GPU speedup to 3.49×, well below the value of running 8 GPUs in the first place. NVLink 4.0 keeps comm under 7% even at this longer sequence, and NVLink 5.0 stays near 98%. (Note: BF16 (bfloat16) and FP16 are both 2 bytes/element, so the byte counts in this article are the same in either dtype.)

Figure 2: Speedup vs GPU count on NVLink (left) and 8-GPU speedup by interconnect class (right).

The left panel of Figure 2 traces the speedup curve from 1 to 16 GPUs at seq_len=2048 — efficiency stays above 93% throughout, hugging the linear-scaling line. The right panel quantifies what changes when only the interconnect is swapped: PCIe 4.0 loses 25% of the ideal 8× speedup to communication, PCIe 5.0 recovers half of that, and NVLink generations cluster near the ideal. The cliff between PCIe and NVLink — roughly an order of magnitude in bandwidth — is what forces tensor-parallel groups inside a single node. For cross-node tensor parallelism (InfiniBand at 25–400 GB/s), the overhead is even worse, which is why production stacks pair intra-node TP with inter-node pipeline or data parallelism.

Tensor Parallelism in Transformer Layers

The Megatron-LM approach (Shoeybi et al. 2019) applies the column/row split pattern to both attention and FFN sub-layers in a way that minimizes AllReduce count to exactly two per transformer layer: one after the attention output projection and one after the second FFN linear.

Attention: the multi-head attention weight matrices W_Q, W_K, W_V (each [d × d]) are split column-wise so each GPU owns n_heads/N complete heads. The head dimension is the natural split axis because each head is independent — no cross-head communication is needed during attention computation. The output projection W_O is split row-wise (by head), and its partial results are summed via AllReduce to produce the full post-attention activation.

FFN: the first linear W_1 ([d × 4d]) is split column-wise across the 4d hidden dimension; the second linear W_2 ([4d × d]) is split row-wise. One AllReduce after W_2 collapses the partial outputs. This pattern means a 32-layer model requires exactly 64 AllReduces per forward pass — two per layer, each operating on a [batch × seq × d_model] activation tensor.

The code below computes the exact AllReduce volume and parameter split for the 7B configuration, confirming both the parameter-per-GPU numbers and the fixed 16 MB AllReduce size regardless of TP degree.

import numpy as np

def megatron_attention_tp(batch: int, seq: int, d: int, n_heads: int, n_gpus: int) -> dict:
    """
    Megatron-LM attention tensor parallelism.
    Shards attention heads across GPUs: each GPU handles n_heads/n_gpus heads.
    Column-parallel QKV projection, row-parallel output projection.
    """
    heads_per_gpu = n_heads // n_gpus
    d_head = d // n_heads

    # Each GPU: Q, K, V projections for its heads (column-parallel)
    qkv_params_per_gpu = 3 * d * (heads_per_gpu * d_head)

    # Each GPU: attention computation for its heads
    attn_flops_per_gpu = 4 * batch * seq * seq * heads_per_gpu * d_head

    # Row-parallel output projection (results summed via AllReduce)
    out_params_per_gpu = heads_per_gpu * d_head * d

    # One AllReduce after output projection (size: batch × seq × d)
    allreduce_gb = batch * seq * d * 2 / 1024**3

    return {
        "params_per_gpu": qkv_params_per_gpu + out_params_per_gpu,
        "total_params": (qkv_params_per_gpu + out_params_per_gpu) * n_gpus,
        "allreduce_gb": allreduce_gb,
        "n_allreduces": 1,  # one per attention layer
    }

def megatron_ffn_tp(batch: int, seq: int, d: int, d_ffn: int, n_gpus: int) -> dict:
    """
    Megatron-LM FFN tensor parallelism.
    Column-parallel first linear (splits d_ffn across GPUs).
    Row-parallel second linear (results summed via AllReduce).
    """
    ffn1_params_per_gpu = d * (d_ffn // n_gpus)  # column-parallel
    ffn2_params_per_gpu = (d_ffn // n_gpus) * d  # row-parallel

    allreduce_gb = batch * seq * d * 2 / 1024**3

    return {
        "params_per_gpu": ffn1_params_per_gpu + ffn2_params_per_gpu,
        "total_params": (ffn1_params_per_gpu + ffn2_params_per_gpu) * n_gpus,
        "allreduce_gb": allreduce_gb,
        "n_allreduces": 1,  # one per FFN layer
    }

if __name__ == "__main__":
    # 7B model configuration
    config = {"batch": 1, "seq": 2048, "d": 4096, "n_heads": 32}
    d_ffn = 11008
    n_layers = 32

    print("=== Megatron-LM Tensor Parallel Communication Pattern (7B, 32 layers) ===")
    print()
    print(f"{'GPUs':>6} {'Attn AllReduce/layer':>22} {'FFN AllReduce/layer':>22} "
          f"{'Total/layer':>14} {'Total model':>14}")
    print("-" * 82)

    for n_gpus in [1, 2, 4, 8]:
        attn = megatron_attention_tp(n_gpus=n_gpus, **config)
        ffn = megatron_ffn_tp(n_gpus=n_gpus, d_ffn=d_ffn, **{k: v for k, v in config.items()
                                                                if k != "n_heads"})
        total_ar_gb = (attn["allreduce_gb"] + ffn["allreduce_gb"]) * n_layers
        total_gb_per_layer = attn["allreduce_gb"] + ffn["allreduce_gb"]

        attn_mb = attn["allreduce_gb"] * 1024
        ffn_mb = ffn["allreduce_gb"] * 1024
        total_mb = total_gb_per_layer * 1024

        if n_gpus == 1:
            print(f"{n_gpus:>6} {'no AllReduce needed':>22} {'no AllReduce needed':>22} "
                  f"{'0':>14} {'0 GB':>14}")
        else:
            print(f"{n_gpus:>6} {attn_mb:>18.0f} MB  {ffn_mb:>19.0f} MB  "
                  f"{total_mb:>10.0f} MB  {total_ar_gb*1024:.0f} MB")

    print()
    print("Each layer generates 2 AllReduces (1 attn + 1 FFN).")
    print("32-layer model: 64 total AllReduces per forward pass.")
    print()

    print("=== Params Per GPU vs Total ===")
    for n_gpus in [1, 2, 4, 8]:
        attn = megatron_attention_tp(n_gpus=n_gpus, **config)
        ffn = megatron_ffn_tp(n_gpus=n_gpus, d_ffn=d_ffn, **{k: v for k, v in config.items()
                                                                if k != "n_heads"})
        total_params_per_gpu = (attn["params_per_gpu"] + ffn["params_per_gpu"]) * n_layers
        total_params = (attn["total_params"] + ffn["total_params"]) * n_layers
        mem_per_gpu = total_params_per_gpu * 2 / 1024**3  # FP16
        print(f"  {n_gpus} GPUs: {total_params_per_gpu/1e6:.0f}M params/GPU  "
              f"({mem_per_gpu:.1f} GB/GPU in FP16)")

Output:

=== Megatron-LM Tensor Parallel Communication Pattern (7B, 32 layers) ===

  GPUs   Attn AllReduce/layer    FFN AllReduce/layer    Total/layer    Total model
----------------------------------------------------------------------------------
     1    no AllReduce needed    no AllReduce needed              0           0 GB
     2                 16 MB                   16 MB          32 MB  1024 MB
     4                 16 MB                   16 MB          32 MB  1024 MB
     8                 16 MB                   16 MB          32 MB  1024 MB

Each layer generates 2 AllReduces (1 attn + 1 FFN).
32-layer model: 64 total AllReduces per forward pass.

=== Params Per GPU vs Total ===
  1 GPUs: 5033M params/GPU  (9.4 GB/GPU in FP16)
  2 GPUs: 2517M params/GPU  (4.7 GB/GPU in FP16)
  4 GPUs: 1258M params/GPU  (2.3 GB/GPU in FP16)
  8 GPUs: 629M params/GPU  (1.2 GB/GPU in FP16)

The AllReduce tensor size is fixed regardless of GPU count — it’s always batch × seq × d_model. The total volume across the model (1024 MB for a 32-layer 7B at batch=1, seq=2048) is what NVLink has to move per forward pass; the per-AllReduce size (16 MB) is what determines whether ring-AllReduce can hide latency behind bandwidth. The bandwidth-to-compute ratio is fixed by the architecture and hardware — which is why doubling GPU count on the same interconnect does not double overhead, and why the same model gets better TP efficiency on H100 (NVLink 5.0) than on A100. Note that this count covers attention and FFN AllReduces only; sequence-parallel layer-norm and sharded embeddings add a small additional AllReduce/AllGather when those optimizations are enabled; the baseline Megatron with replicated layer-norm has zero extra traffic from them.

Tensor Parallelism vs Pipeline Parallelism

Tensor parallelism is not the only way to distribute a model across GPUs. Pipeline parallelism assigns whole layers to different devices, which avoids AllReduce communication but introduces a “bubble” — idle GPU time while earlier stages finish. The following simulation quantifies that bubble against the AllReduce overhead from tensor parallelism.

import numpy as np

def pipeline_parallel_overhead(n_stages: int, batch_size: int) -> float:
    """
    Pipeline parallel bubble overhead.
    Bubble = (n_stages - 1) / (n_stages - 1 + batch_size) × 100%
    With micro-batching: bubble fraction = (n_stages - 1) / (batch_size + n_stages - 1)
    """
    return (n_stages - 1) / (batch_size + n_stages - 1) * 100

def allreduce_cost_ms(tensor_gb: float, bandwidth_gb_s: float, n_gpus: int) -> float:
    """Ring AllReduce: 2(N-1)/N × data per GPU at peak bandwidth."""
    data_sent_gb = 2 * (n_gpus - 1) / n_gpus * tensor_gb
    return data_sent_gb / bandwidth_gb_s * 1000

def tensor_parallel_overhead(n_gpus: int, activation_gb: float,
                               bw_gb_s: float, compute_ms: float) -> float:
    """
    Tensor parallel overhead from AllReduce.
    Returns overhead as percentage of compute time.
    """
    comm_ms = allreduce_cost_ms(activation_gb, bw_gb_s, n_gpus)
    return comm_ms / compute_ms * 100

if __name__ == "__main__":
    activation_gb = 1 * 2048 * 4096 * 2 / 1024**3  # batch=1, seq=2048, d=4096, BF16
    compute_ms_per_layer = 1.37  # from compute_cost_ms(1, 2048, 4096, 312 TFLOPS)
    nvlink_bw = 600  # GB/s

    print("=== Pipeline Parallelism: Bubble Overhead ===")
    print(f"{'Stages':>8} {'Batch=1':>10} {'Batch=4':>10} {'Batch=16':>11} {'Batch=64':>11}")
    print("-" * 52)
    for stages in [2, 4, 8, 16, 32]:
        row = f"{stages:>8}"
        for batch in [1, 4, 16, 64]:
            ov = pipeline_parallel_overhead(stages, batch)
            row += f" {ov:>9.1f}%"
        print(row)

    print()
    print("=== Tensor Parallelism: AllReduce Overhead (NVLink 4.0) ===")
    print(f"{'GPUs':>8} {'Overhead %':>12}")
    print("-" * 22)
    for n in [2, 4, 8, 16, 32]:
        ov = tensor_parallel_overhead(n, activation_gb, nvlink_bw, compute_ms_per_layer)
        print(f"{n:>8} {ov:>11.2f}%")

    print()
    print("=== Side-by-Side: 8-way Parallelism Overhead ===")
    print(f"{'Method':>25} {'Batch=1':>10} {'Batch=4':>10} {'Batch=16':>11}")
    print("-" * 58)

    # PP with 8 stages
    pp_row = "Pipeline (8 stages)"
    for batch in [1, 4, 16]:
        ov = pipeline_parallel_overhead(8, batch)
        pp_row += f" {ov:>9.1f}%"
    print(f"{pp_row:>55}")

    # TP with 8 GPUs at NVLink
    tp_nvlink = tensor_parallel_overhead(8, activation_gb, nvlink_bw, compute_ms_per_layer)
    tp_pcie = tensor_parallel_overhead(8, activation_gb, 32, compute_ms_per_layer)  # PCIe 4.0

    print(f"{'TP 8 GPUs (NVLink 4.0)':>25} {tp_nvlink:>9.2f}%  (same overhead at all batch sizes)")
    print(f"{'TP 8 GPUs (PCIe 4.0)':>25} {tp_pcie:>9.2f}%  (same overhead at all batch sizes)")

    print()
    print("Key: PP bubble scales with 1/batch_size; TP overhead is batch-independent.")
    print("PP wins at large batch sizes where bubble shrinks; TP wins for latency (batch=1).")

Output:

=== Pipeline Parallelism: Bubble Overhead ===
  Stages    Batch=1    Batch=4    Batch=16    Batch=64
----------------------------------------------------
       2      50.0%      20.0%       5.9%       1.5%
       4      75.0%      42.9%      15.8%       4.5%
       8      87.5%      63.6%      30.4%       9.9%
      16      93.8%      78.9%      48.4%      19.0%
      32      96.9%      88.6%      66.0%      32.6%

=== Tensor Parallelism: AllReduce Overhead (NVLink 4.0) ===
    GPUs   Overhead %
----------------------
       2        1.90%
       4        2.85%
       8        3.33%
      16        3.56%
      32        3.68%

=== Side-by-Side: 8-way Parallelism Overhead ===
                   Method    Batch=1    Batch=4    Batch=16
----------------------------------------------------------
   Pipeline (8 stages)      87.5%      63.6%      30.4%
   TP 8 GPUs (NVLink 4.0)      3.33%  (same overhead at all batch sizes)
     TP 8 GPUs (PCIe 4.0)     62.37%  (same overhead at all batch sizes)

Key: PP bubble scales with 1/batch_size; TP overhead is batch-independent.
PP wins at large batch sizes where bubble shrinks; TP wins for latency (batch=1).

Pipeline parallelism at batch=1, 8 stages has 87.5% overhead — nearly useless for latency-sensitive inference. Tensor parallelism on NVLink stays around 3.3% overhead independent of batch size. The two methods cross over at large training batch sizes (64+), where pipeline parallelism falls to ~10% and becomes competitive on bandwidth-poor interconnects. This is why real distributed training systems (Megatron-DeepSpeed, NVIDIA NeMo, DeepSeek’s training stack) combine the two — TP inside each node and PP across nodes — rather than picking one.

Figure 3: Compute vs AllReduce time as GPU count grows (left) and 8-GPU efficiency by interconnect (right).

The left panel of Figure 3 makes the dominance of compute over communication visible at intra-node scale: even at 16 GPUs on NVLink, the red AllReduce sliver is barely perceptible against the green compute bar. The right panel collapses the whole story to a single ranking: PCIe 4.0 sits at 74.9% efficiency, PCIe 5.0 at 85.6%, NVLink 3.0–5.0 between 96.5% and 99.4%. For a frontier-scale serving system the right answer is rarely “more PCIe GPUs” — it is “fewer GPUs on a better link, or pipeline parallel across nodes.” This is the same principle that motivated NVIDIA’s NVSwitch fabric and AMD’s xGMI (AMD’s inter-GPU interconnect): tensor parallelism only pays off when the link can keep up.

Practical Parallelism Decisions

With the overhead numbers in hand, the decision reduces to a lookup: does the model fit on one GPU, what interconnect is available, and what is the serving batch size? The function below encodes these rules for the most common hardware configurations.

import numpy as np

def recommend_parallelism(model_params_b: float, gpu_vram_gb: float,
                            n_gpus: int, interconnect: str,
                            serving_batch: int) -> dict:
    """
    Recommend parallelism strategy based on model size and hardware.
    """
    model_gb_fp16 = model_params_b * 2  # FP16

    # Check if model fits on one GPU
    fits_one = model_gb_fp16 <= gpu_vram_gb * 0.85

    # Interconnect bandwidth
    bw_map = {"nvlink": 600, "pcie5": 64, "pcie4": 32, "infiniband": 25}
    # InfiniBand range: NDR-class fabrics reach 50-400 GB/s; the 25 GB/s used here is the conservative HDR baseline.
    bw = bw_map.get(interconnect.lower(), 32)

    # Tensor parallelism overhead at this interconnect
    activation_gb = serving_batch * 2048 * 4096 * 2 / 1024**3  # rough estimate
    comm_overhead_pct = 2 * (n_gpus - 1) / n_gpus * activation_gb / bw * 1000 / 1.37 * 100

    # Minimum GPUs needed
    # 85% headroom factor accounts for activations + workspace; this is why the recommender
    # is more conservative than the raw param-memory table above.
    min_gpus = int(np.ceil(model_gb_fp16 / (gpu_vram_gb * 0.85)))

    recommendations = []
    if fits_one and n_gpus > 1:
        recommendations.append("Data parallel (model fits on one GPU; replicate for throughput)")
    if not fits_one and bw >= 300:
        recommendations.append(f"Tensor parallel: {min_gpus}+ GPUs (NVLink ≥300 GB/s required)")
    if not fits_one and bw < 300:
        recommendations.append(f"Pipeline parallel or quantize first ({interconnect} bandwidth too low for TP)")
    if model_params_b > 100:
        recommendations.append("Consider 3D parallelism: TP within node + PP across nodes")

    return {
        "model_gb_fp16": model_gb_fp16,
        "fits_one_gpu": fits_one,
        "min_gpus_needed": min_gpus,
        "tp_comm_overhead_pct": comm_overhead_pct,
        "recommendations": recommendations,
    }

if __name__ == "__main__":
    scenarios = [
        (7, 80, 8, "nvlink", 1, "7B on DGX-A100 (8×80GB NVLink)"),  # DGX (NVIDIA's 8-GPU server)
        (70, 80, 8, "nvlink", 1, "70B on DGX-A100 (8×80GB NVLink)"),
        (7, 24, 4, "pcie4", 1, "7B on workstation (4×RTX3090 PCIe)"),
        (405, 80, 16, "nvlink", 1, "405B on 2× DGX-A100"),
        (70, 80, 8, "infiniband", 32, "70B on cluster (cross-node InfiniBand)"),
    ]

    print("=== Parallelism Recommendations ===")
    print()
    for params, vram, n_gpus, interconnect, batch, desc in scenarios:
        rec = recommend_parallelism(params, vram, n_gpus, interconnect, batch)
        print(f"Scenario: {desc}")
        print(f"  Model: {params}B ({rec['model_gb_fp16']:.0f} GB FP16)  "
              f"Min GPUs: {rec['min_gpus_needed']}")
        for r in rec["recommendations"]:
            print(f"  → {r}")
        print()

Output:

=== Parallelism Recommendations ===
Scenario: 7B on DGX-A100 (8×80GB NVLink)
  Model: 7B (14 GB FP16)  Min GPUs: 1
  → Data parallel (model fits on one GPU; replicate for throughput)

Scenario: 70B on DGX-A100 (8×80GB NVLink)
  Model: 70B (140 GB FP16)  Min GPUs: 3
  → Tensor parallel: 3+ GPUs (NVLink ≥300 GB/s required)

Scenario: 7B on workstation (4×RTX3090 PCIe)
  Model: 7B (14 GB FP16)  Min GPUs: 1
  → Data parallel (model fits on one GPU; replicate for throughput)

Scenario: 405B on 2× DGX-A100
  Model: 405B (810 GB FP16)  Min GPUs: 12
  → Tensor parallel: 12+ GPUs (NVLink ≥300 GB/s required)
  → Consider 3D parallelism: TP within node + PP across nodes

Scenario: 70B on cluster (cross-node InfiniBand)
  Model: 70B (140 GB FP16)  Min GPUs: 3
  → Pipeline parallel or quantize first (infiniband bandwidth too low for TP)

The recommendations collapse into three patterns. If the model fits on one GPU with headroom, data parallelism is the right answer — replicate to scale throughput, never shard to add overhead. If the model overflows one GPU and the interconnect is NVLink-class, tensor parallelism is the answer (the recommender includes a 15% activation+workspace headroom on top of weights; that’s why it suggests 3 GPUs for the 70B where the raw param table shows 2 fits). If the model overflows and the interconnect is PCIe or InfiniBand, the only viable options are pipeline parallelism, quantization (which can shrink a 70B INT4/INT8 (4-bit / 8-bit integer quantization) model below 40 GB), or ZeRO (Zero Redundancy Optimizer) data parallelism. The 405B case is the only one that needs all three together: tensor parallelism within a DGX node, pipeline parallelism across nodes, and often data parallelism to scale total throughput — the “3D parallelism” that Megatron-DeepSpeed introduced for trillion-parameter training.

Thank you for reading my post, and I hope it was useful for you. If you enjoyed the article and would like to show your support, please consider taking the following actions:

👏 Give the story a round of applause (clap) to help it gain visibility.

📖 Follow me on Medium to access more of the content on my profile. Follow Now

🔔 Subscribe to the newsletter to not miss my latest posts: Subscribe Now

🛎 Connect with me on LinkedIn for updates.