Tensor parallelism and FSDP

You have eight GPUs with 80 GB each. The model has 7 billion parameters.

The weights fit. In bfloat16, they occupy about 14 GB. You start training with AdamW, and the job dies before processing its first batch:

CUDA out of memory. Tried to allocate ...

The missing memory is not a mystery. Training stores gradients, optimizer state, and usually a substantial pile of activations. With ordinary data parallelism, every GPU stores a complete copy of most of that state.

Two ideas solve different parts of this problem:

• FSDP, Fully Sharded Data Parallel, divides the model’s state between data-parallel ranks.
• Tensor parallelism divides the arithmetic inside a layer between devices.

They are often used together, but they are not interchangeable.

The training memory budget

A parameter is a learned number. A gradient is the derivative computed for it during backpropagation. Optimizer state is information AdamW uses to update parameters. Activations are intermediate forward-pass values saved for backpropagation.

For P = 7,000,000,000 parameters, assume:

• bfloat16 parameters: 2 bytes each
• bfloat16 gradients: 2 bytes each
• two fp32 Adam moment buffers: 8 bytes total per parameter
• an fp32 master copy of the weights: 4 bytes each

Thing	Bytes per parameter	Memory for 7B parameters
bfloat16 parameters	2	14 GB
bfloat16 gradients	2	14 GB
fp32 Adam moments, two of them	8	56 GB
fp32 master parameters	4	28 GB
Total, before activations	16	112 GB

These are decimal gigabytes; displayed GPU capacities use GiB, so exact numbers differ slightly. The optimizer and master weights alone consume 84 GB. AdamW dominates because this common mixed-precision setup stores three fp32-sized values per parameter.

Implementations vary: some optimizers do not keep a separate master copy, and some use less state. The planning rule is that training state can be several times larger than inference weights.

Activations are another bill. A layer output with batch size 2, sequence length 2,048, and hidden size 4,096 uses:

2 × 2,048 × 4,096 × 2 bytes = 33,554,432 bytes

That is about 33.6 MB for one bfloat16 tensor. Thirty-two saved tensors would already be about 1.07 GB, before residuals, attention intermediates, temporary workspaces, and allocator overhead.

Activation memory depends on batch size, sequence length, layer count, attention implementation, and checkpointing. Activation checkpointing trades extra recomputation for less resident memory.

With eight ordinary data-parallel replicas, each GPU tries to hold the whole 112 GB of model, gradient, and optimizer state. The aggregate 896 GB does not help a single GPU. Replication is the problem.

Sharding the state: ZeRO and FSDP

Sharding divides an object so each rank stores only one piece. A rank is one participating process, usually one GPU process.

ZeRO, or Zero Redundancy Optimizer, removes data-parallel replicas in three stages. With N data-parallel ranks, the idealized model-state memory per rank is:

Scheme	Sharded state	Approximate memory per rank
Ordinary DDP, sometimes called Stage 0	Nothing	16P
ZeRO Stage 1	Optimizer state	4P + 12P/N
ZeRO Stage 2	Optimizer state and gradients	2P + 14P/N
ZeRO Stage 3	Optimizer state, gradients, and parameters	16P/N

The formulas exclude activations, temporary buffers, communication storage, and framework overhead. For the 7B example on eight ranks:

• Stage 0: 112 GB per rank
• Stage 1: 28 GB + 84 GB / 8 = 38.5 GB
• Stage 2: 14 GB + 98 GB / 8 = 26.25 GB
• Stage 3: 112 GB / 8 = 14 GB

Stage 1 can be a useful first move because optimizer state is the largest replicated object. Each stage exchanges memory savings for communication:

• Stage 1 keeps replicated parameters and gradients. DDP reduces gradients, while optimizer states remain private shards.
• Stage 2 partitions gradients during reduction, commonly with reduce-scatter. Updated parameter pieces must still be made available to ranks holding a full model.
• Stage 3 gathers parameter shards when a layer needs them, computes with the temporary full layer, and releases them. Gradients are reduce-scattered back to their owners; optimizer state stays sharded.

An all-gather gives every rank a complete assembled tensor. A reduce-scatter sums tensors across ranks and gives each rank its assigned slice. These collectives make full sharding possible.

FSDP is PyTorch’s Fully Sharded Data Parallel implementation. Its FULL_SHARD strategy is broadly the same kind of state partitioning as ZeRO Stage 3. FSDP commonly wraps a transformer block.

Before the block runs, parameter shards are all-gathered. Afterward, the full parameters can be discarded. During backward, gradients are reduced and scattered to their owning ranks.

The wrapping boundary matters. Wrapping the entire model can make FSDP gather the entire model at once. Wrapping individual transformer blocks lets it gather one block, compute it, release it, and then gather the next.

Checkpoint saving must also use sharded model and optimizer state; gathering everything on rank zero can recreate the original memory problem.

FSDP is a strong choice when replicated state is the obstacle. It can be slower than DDP for a small model because each step adds parameter gathers and gradient reductions. Memory was not abolished; it was exchanged for network traffic.

Splitting the computation: tensor parallelism

FSDP still makes one rank perform a layer’s matrix multiplication. A very wide matrix can itself be too large or expensive for one GPU. Tensor parallelism splits that matrix across a group of GPUs, called a tensor-parallel group.

Consider the feed-forward pair:

H = activation(X W1)

Y = H W2

Let:

• X be 1024 × 4096, representing 1,024 tokens with hidden size 4,096.
• W1 be 4096 × 11008.
• Four tensor-parallel ranks be available.

A column-wise split divides W1 into four matrices of shape 4096 × 2752:

W1 = [W1_0, W1_1, W1_2, W1_3]

Every rank has X, but each computes only its local product, such as X W1_0. Each result is 1024 × 2752, so the intermediate hidden dimension is distributed.

W2 has shape 11008 × 4096. Split it row-wise into four matrices of shape 2752 × 4096. Each rank multiplies its local H_i by its local W2_i, producing a partial 1024 × 4096 result:

Y = Y_0 + Y_1 + Y_2 + Y_3

An all-reduce sums those partial results, leaving every rank with Y.

The pairing is:

1. Column-split the first linear layer.
2. Keep the intermediate activation sharded.
3. Row-split the second linear layer.
4. All-reduce the partial output.

The local matrix dimensions are smaller, and the intermediate activation does not need to be gathered between the two projections.

In this example, the forward all-reduce carries the output tensor:

1024 × 4096 × 2 bytes = 8,388,608 bytes

That is about 8.4 MB. A four-rank ring all-reduce sends roughly 2 × (4 - 1) / 4 = 1.5 times that payload per rank, or about 12.6 MB, ignoring protocol details.

The important detail is the frequency: tensor parallelism communicates inside nearly every layer. A slow cross-node link can make these repeated collectives more expensive than the matrix multiplication. Fast GPU-to-GPU links inside one server are therefore the natural home for tensor parallelism.

Backward propagation also communicates. “One all-reduce per block” describes the paired forward path for this two-linear pattern, not an entire training step or every transformer implementation.

3D parallelism: choosing the axes

Large jobs commonly use three axes:

• Data parallelism, degree D, gives different batches to replicas and synchronizes gradients.
• Tensor parallelism, degree T, splits individual layers.
• Pipeline parallelism, degree P, places consecutive layer groups on different devices and sends activations between them.

The process count is:

world size = D × T × P

For 64 GPUs, consider D = 2, T = 8, and P = 4. Each model replica uses 32 GPUs. Each pipeline stage owns one quarter of the layers and uses eight GPUs for tensor parallelism. Two replicas process different data.

Put the most frequent, latency-sensitive communication on the fastest links:

• Tensor groups should usually stay inside one server.
• Pipeline traffic crosses only stage boundaries, so it is a better candidate for crossing servers.
• Data parallelism communicates gradients once per step and can also cross servers, depending on bandwidth.

This is a rule of thumb, not a law. Huge boundary activations or unusually fast cross-node links can change the answer.

Pipeline parallelism also introduces bubbles, when a stage is idle while work moves through the pipeline. More microbatches reduce the relative bubble but increase scheduling and activation-memory complexity. See pipeline parallelism.

Long sequences: sequence and context parallelism

Sequence length is a major source of activation memory. Doubling a sequence from 4,096 to 8,192 doubles many token-shaped activations, while some attention intermediates can grow quadratically.

Sequence parallelism

Sequence parallelism shards the sequence dimension across tensor-parallel ranks for operations that do not need the whole sequence on every rank. Layer normalization, dropout, and some residual paths are common targets.

Each rank holds a slice of token positions rather than a replicated tokens × hidden activation. Communication is needed when the next operation requires another layout.

Context parallelism

Context parallelism partitions a long context across devices, including attention. A rank can own a range of query positions, but each query still needs keys and values from other positions.

Ranks therefore exchange K and V blocks through a ring or all-gather-like schedule. Per-rank memory falls for partitioned tensors, while communication rises because attention has global dependencies.

Sequence parallelism usually extends tensor parallelism; context parallelism specifically targets very long contexts. Neither is a free “divide the sequence” switch: masks, positional encodings, KV-cache layout, and attention kernels must agree on ownership.

Combine them with mixed precision and activation checkpointing when activation memory is the limit.

Which technique should carry the problem?

Technique	Main bottleneck	Communication pattern	Use it when
DDP	Too little aggregate compute	Gradient all-reduces each step	Full training state fits on each GPU
FSDP or ZeRO-3	Replicated model state	Parameter all-gathers and gradient reduce-scatters	State does not fit per GPU
Tensor parallelism	A layer is too wide or expensive	Collectives inside each layer	Fast intra-node links are available
Pipeline parallelism	Too many layers for one device group	Activation sends between stages	Layers can be split into balanced stages
Activation checkpointing	Saved activations	Extra forward computation	Compute is available but memory is not
Sequence or context parallelism	Long-token activations and attention	Sequence-layout or K/V exchange	Sequence length is the limit

Production systems assign each bottleneck to an axis: FSDP handles replicated state, tensor parallelism handles width, pipeline parallelism handles depth, and checkpointing or sequence/context parallelism handles activations.

Failure modes and the trade-off

Common failure modes have different fixes:

• Tensor parallelism across a slow interconnect makes collectives dominate the step; keep tensor groups inside nodes.
• An FSDP memory spike at a block boundary usually means the wrapping unit is too large; wrap individual transformer blocks and leave room for activations and temporary buffers.
• A rank-zero checkpoint failure means the save path is gathering sharded state; use sharded model and optimizer checkpoints.

The honest limitation is complexity. You must reason about:

• process groups
• topology
• wrapping boundaries
• microbatches
• checkpoint formats
• failure recovery

FSDP can make a model fit while reducing throughput. Tensor parallelism can reduce per-GPU compute while making the network the new bottleneck. A smaller model with ordinary DDP may finish sooner and be easier to operate.

What to remember

• Mixed-precision AdamW can need about 16 bytes per parameter before activations.
• ZeRO Stage 1 shards optimizer state, Stage 2 adds gradients, and Stage 3 shards parameters too.
• FSDP gathers a module when it runs, then keeps only each rank’s owned state.
• Column-parallel first projections plus row-parallel second projections keep intermediate activations sharded and all-reduce the output.
• Put tensor parallelism on fast intra-node links; use data and pipeline parallelism across slower boundaries when topology permits.

Quick check

0/3

Q1

14 GB28 GB84 GB112 GB

Check answer

Q2

The first layer creates sharded intermediate activations that match the rows of the second layer, so only the partial final outputs need to be summed.The first layer replicates all activations, so the second layer needs no communication.The second layer can be removed because its weights are stored in the optimizer.Both layers must be placed on the same GPU.

Check answer

Q3

Use tensor parallelism across all 32 GPUs, because the largest possible group always minimizes communication.Use an eight-way tensor-parallel group inside each server, then use data or pipeline parallelism across servers.Use no sharding and replicate the full model on every GPU.Use context parallelism for the wide matrix even though the sequence is short.

Check answer