Mesh LLM Core Concepts: Parallelism and Sharding

LLM Scale: Why Single-Device Inference Fails

Modern Large Language Models (LLMs) frequently exceed the memory capacity of single GPU devices. For instance, the Llama 2 70B model contains 70 billion parameters. When stored in 16-bit floating-point format (bfloat16 or float16), each parameter requires 2 bytes.

Calculating the memory footprint for model weights alone:

# Calculate model weight memory footprint
parameters = 70_000_000_000  # Llama 2 70B parameters
bytes_per_parameter = 2      # bfloat16 or float16

total_bytes = parameters * bytes_per_parameter
total_gb = total_bytes / (1024**3)

print(f"Llama 2 70B model weights: {total_gb:.2f} GB")

Llama 2 70B model weights: 130.37 GB

A high-end GPU like the NVIDIA A100 typically offers 80 GB of HBM2e VRAM, while the H100 provides 80 GB of HBM3. The calculated 130.37 GB for Llama 2 70B weights significantly exceeds the 80 GB available on these single devices. This memory constraint prevents loading the full model onto one GPU.

Beyond model weights, LLM inference demands additional memory for activations and the KV (Key-Value) cache. Activations store intermediate tensor computations during the forward pass, scaling with batch size and sequence length. The KV cache stores previously computed attention keys and values, growing proportionally with the generated output sequence length. For long sequences or large batch sizes, these components can consume tens of gigabytes.

A single GPU cannot simultaneously hold the full model weights, activations, and the KV cache required for state-of-the-art LLM inference. This fundamental limitation necessitates distributing the model and its computations across multiple devices.

Model Sharding: Distributing LLM Components

Training large language models (LLMs) often exceeds the memory capacity of a single GPU, even high-end ones. A 70-billion parameter LLM, using bfloat16 precision for weights and an AdamW optimizer, can require over 500 GB of GPU memory. This far exceeds the 80 GB available on a single NVIDIA A100 GPU, making distributed training a necessity.

Model sharding partitions the LLM’s components across multiple devices. Each device stores and processes a subset of the model, enabling the full model to operate within a distributed system and potentially accelerate training by parallelizing computations. This is distinct from data parallelism, where the entire model resides on each device and data batches are split.

Typically, sharding applies to model weights (parameters), intermediate activations, and optimizer states. For example, a single transformer layer’s weight matrices might be split. One GPU holds the first half of a matrix, another GPU holds the second half.

During a forward pass, each GPU computes its portion of the layer’s output using its local weights. These partial results are then exchanged and aggregated across devices to reconstruct the complete layer output before processing continues. The reverse occurs during the backward pass for gradient computation.

This distribution reduces per-device memory load, allowing larger models to fit. However, it introduces significant communication overhead. Data must be transferred between GPUs to synchronize computations and aggregate results, which can become a bottleneck. The design of efficient sharding strategies focuses on minimizing this inter-device communication.

How Tensor Parallelism Divides Layers

A core operation within any Large Language Model layer is the linear transformation, represented as Y = X @ W + B. Here, X is the input activation tensor, W is the layer’s weight matrix, B is the bias vector, and Y is the output activation. For models with billions of parameters, the W matrix can become too large to store on a single GPU’s memory.

Tensor parallelism addresses this by sharding the W matrix across multiple devices. The input X or intermediate activations are also distributed or replicated to allow parallel computation. This approach splits computations within a single layer, unlike pipeline parallelism which splits layers themselves.

Consider a weight matrix W with dimensions (input_features, output_features). There are two primary ways to shard W:

Column-wise Sharding In column-wise sharding, the W matrix is split horizontally into N parts, where N is the number of parallel devices. Each part W_i has dimensions (input_features, output_features / N). The input tensor X with dimensions (batch_size, input_features) is replicated across all N devices.

Each device i then computes a partial output Y_i = X @ W_i. Since X is (batch_size, input_features) and W_i is (input_features, output_features / N), Y_i will be (batch_size, output_features / N). The final output Y is formed by concatenating these partial outputs across devices: Y = [Y_0, Y_1, ..., Y_{N-1}]. This concatenation happens logically, with each device holding its segment of the output.

# Assuming N devices, device 0 holds W_0, device 1 holds W_1, etc.
# X is replicated on all devices.
# Each device computes:
Y_i = X @ W_i
# Logical concatenation across devices forms the full Y

Row-wise Sharding Row-wise sharding splits the W matrix vertically into N parts. Each W_i has dimensions (input_features / N, output_features). For this split, the input tensor X must also be sharded along its feature dimension, so X_i has dimensions (batch_size, input_features / N).

Each device i computes Y_i = X_i @ W_i. The resulting Y_i will have dimensions (batch_size, output_features). To form the complete output Y, these partial results must be summed across all devices: Y = Y_0 + Y_1 + ... + Y_{N-1}. This summation requires an all-reduce collective communication operation, which adds communication overhead.

# Assuming N devices, device 0 holds X_0 and W_0, device 1 holds X_1 and W_1, etc.
# Each device computes:
Y_i = X_i @ W_i
# All-reduce operation sums Y_i across all devices to get the full Y on each device
Y = all_reduce(Y_i)

Often, a full linear layer within an LLM combines these two sharding types. A ColumnParallelLinear operation is typically followed by an all-reduce to gather the full intermediate activation, which is then fed into a RowParallelLinear operation. This pattern minimizes communication by ensuring that only one all-reduce is performed per block of operations.

Pipeline Parallelism: Staging Sequential Layers

Large language models often exceed the memory capacity of a single GPU. Pipeline parallelism addresses this by distributing sequential layers of a model across multiple devices, forming a processing pipeline. Each device is responsible for a distinct set of consecutive layers.

Consider a model with L layers. Device 0 might handle layers 0 through L/3 - 1, Device 1 takes L/3 through 2L/3 - 1, and Device 2 processes 2L/3 through L - 1. An input batch first computes on Device 0. Its intermediate activations are then sent to Device 1 for further processing, and finally to Device 2.

This forms an assembly line. As Device 0 finishes processing Batch A and sends its output to Device 1, Device 0 can immediately begin processing Batch B. Simultaneously, Device 1 works on Batch A. This overlap of computation and communication reduces idle time compared to processing the entire model sequentially on each device.

To maximize throughput, pipeline parallelism often uses micro-batching. A large logical batch is split into smaller micro-batches. Device 0 processes Micro-batch 1, then Micro-batch 2, and so on. As Micro-batch 1 moves to Device 1, Device 0 starts Micro-batch 2. This keeps the pipeline full.

# Conceptual layer distribution for a 12-layer model across 3 devices
total_layers = 12
num_devices = 3
layers_per_device = total_layers // num_devices

device_layer_map = {}
for device_id in range(num_devices):
    start_layer = device_id * layers_per_device
    end_layer = start_layer + layers_per_device - 1
    device_layer_map[f"Device {device_id}"] = f"Layers {start_layer}-{end_layer}"

print(device_layer_map)

Output:

{'Device 0': 'Layers 0-3', 'Device 1': 'Layers 4-7', 'Device 2': 'Layers 8-11'}

The primary tradeoff with pipeline parallelism is the “pipeline bubble” or “stall”. At the beginning of processing a new batch, and at the end, not all devices are active. For example, Device 1 must wait for Device 0 to complete its first micro-batch before it can start. Similarly, when the last micro-batch leaves Device 0, Device 0 becomes idle while Device 1 and Device 2 finish their work. These periods of inactivity reduce overall hardware use. Communication overhead also increases as intermediate activations must be transferred between devices.

This method is effective for models with many sequential layers that exceed single-device memory. It allows for larger model sizes by distributing the memory footprint of activations and parameters.

Mesh LLM: Orchestrating Parallelism & Sharding

Distributing large language models across many devices requires coordinated strategies for data and computation. Mesh LLM provides a framework to manage tensor parallelism, pipeline parallelism, and overall model sharding, abstracting the complexities of multi-device execution.

Tensor parallelism splits individual tensor operations, such as matrix multiplications, across multiple devices. Each device computes a fragment of the output tensor, then communicates necessary intermediate results. This method reduces the memory footprint for large tensors and accelerates operations that are otherwise memory-bound on a single device. The primary cost is the increased inter-device communication bandwidth.

Pipeline parallelism distributes model layers sequentially across different devices. Device 0 computes layer 1, passes its output to Device 1 for layer 2, and so on. This keeps GPUs busy by overlapping computation and communication, improving overall throughput. However, it introduces latency due to the sequential nature of layer execution and requires careful management of micro-batches to fill the pipeline efficiently.

Mesh LLM combines these techniques to achieve comprehensive model sharding. It defines a logical mesh of devices and maps model components—weights, optimizer states, and gradients—onto this mesh. This mapping dictates how tensor and pipeline parallelism are applied across the device topology. For instance, one dimension of the mesh might handle tensor parallelism, while another handles pipeline parallelism.

Consider a model distributed across a 2D mesh of devices:

# Conceptual Mesh LLM configuration using mesh_tensorflow
import mesh_tensorflow as mtf

# Define a logical mesh shape: 2 pipeline stages, 4 devices per stage for tensor parallelism
mesh_shape = mtf.Shape([("pipeline", 2), ("tensor", 4)])
mesh = mtf.Mesh(mesh_shape, mtf.DeviceMesh(devices=["gpu:0", "gpu:1", "gpu:2", "gpu:3",
                                                    "gpu:4", "gpu:5", "gpu:6", "gpu:7"]))

# Example: A tensor parallel dimension for weights, a pipeline dimension for layers.
# Model layers would be placed along the "pipeline" dimension.
# Individual weight matrices within a layer might be sharded along the "tensor" dimension.

This configuration allows the system to automatically manage the placement and communication patterns. The user defines the desired parallelism dimensions, and Mesh LLM handles the underlying data movement and synchronization. This abstraction simplifies the deployment of models that exceed the capacity of a single GPU or even a single node. The cost of this flexibility is the initial setup complexity in defining the mesh and mapping operations.

Parallelism Pitfalls: Common Sharding Misconfigurations

Distributed LLM training often encounters performance bottlenecks due to sharding misconfigurations. Effective sharding distributes both data and computation evenly across devices. Failure to achieve this balance leads to underutilized resources and increased training times.

One common issue is data skew, where certain shards process significantly more data or computations than others. If a dataset of text is sharded by language, and one language comprises 80% of the corpus, that shard becomes a bottleneck. Other shards remain idle, negating the benefits of parallelism. Monitoring shard-level resource use, such as GPU memory or compute cycles, identifies these imbalances.

Choosing an ineffective sharding key also creates problems. A sharding key should distribute data uniformly and minimize cross-shard dependencies. Sharding a model’s weights by the first layer, for instance, forces all subsequent layers to communicate with the first, creating a sequential bottleneck. This approach limits true parallel execution, as dependent computations must wait for preceding ones.

Excessive communication overhead can negate the gains from parallel computation. While sharding aims to distribute work, frequent data exchange between shards, such as for gradient synchronization or activation passing, consumes significant network bandwidth and latency. Over-sharding, where the number of shards is too high relative to the model size or data volume, exacerbates this. The coordination cost for many small shards can outweigh the compute benefit.

Consider a 100 million parameter model sharded across 1000 GPUs. Each GPU holds only 100,000 parameters. The time spent communicating updates between these tiny fragments will dominate the actual computation time. Conversely, under-sharding a 10 billion parameter model onto two GPUs limits parallelism, hitting a scaling ceiling prematurely. The goal is to balance local computation with necessary global communication.

Tools like torch.distributed.monitor or NVIDIA’s nvprof can help diagnose communication patterns. For example, profiling a distributed run might reveal high NCCL AllReduce times:

nvprof --print-gpu-trace python train_llm.py

This output helps identify if collective operations are consuming an unexpectedly large portion of the training step. Addressing these misconfigurations is crucial for efficient distributed LLM development.