GPU Memory Sharding with FSDP

DISTRIBUTED DEEP LEARNING — PART 1/5

Day 1: Demystifying GPU Memory & The ZeRO Revolution — DeepSpeed, FSDP & State Sharding

25 min read

Series: The Dharma of Development

Distributed DL (Day 1 / 5)

Level: Principal / Systems AI Engineer

💥 Context: You attempt to fine-tune or pretrain a 13-billion parameter dense transformer on an 80GB NVIDIA H100 or A100 GPU using standard PyTorch. You set your batch size to 1. You press enter. Within three seconds, your terminal explodes with the most dreaded error in artificial intelligence: torch.cuda.OutOfMemoryError: CUDA out of memory. How is this possible? In 16-bit precision, 13 billion parameters occupy only ~26 GB of disk space. Why can't an 80 GB state-of-the-art GPU train a 26 GB model? Because the naive mental model of machine learning memory is deeply flawed. Today, we demystify the true $16\Phi$ memory footprint of neural networks and master the architecture that made modern large language models possible: ZeRO (Zero Redundancy Optimizer) and PyTorch Fully Sharded Data Parallel (FSDP).

🕉️ The Dharma of Cosmic Unity in Diversity: Avibhaktaṁ Vibhaktam

"अविभक्तं च भूतेषु विभक्तमिव च स्थितम्।
भूतभर्तृ च तज्ज्ञेयं ग्रसिष्णु प्रभविष्णु च॥"

(Avibhaktaṁ cha bhūteṣu vibhaktam iva sthitam, bhūta-bhartṛi cha taj jñeyaṁ grasiṣhṇu prabhaviṣhṇu cha — Bhagavad Gita 13.17)

The Translation: Indivisible, yet existing as if divided among all living beings; know That to be the sustainer of all entities, devouring them and bringing them forth anew.

The Engineering Parallel: The human mind assumes that to understand or operate upon an entity, the entire entity must reside within a single localized container. In classical deep learning (DDP), every GPU clings greedily to an identical, monolithic copy of the entire network. But the sheer scale of intelligence surpasses the bounds of any individual silicon chip. True mastery lies in recognizing that a network can be undivided in its mathematical essence, yet divided across an entire cluster of silicon (Vibhaktam iva sthitam). Through Fully Sharded Data Parallelism, no single GPU bears the burden of the whole; each holds only a fragment, yet in collective rhythm, they materialize, compute, and dissolve the parameters seamlessly.

1. The True Anatomy of GPU VRAM: The $16\Phi$ Tax

When training with mixed precision (FP16 or BF16) using the standard AdamW optimizer, memory is divided into two broad categories: Model States and Residual Memory.

For a model with $\Phi$ parameters, static model states consume a staggering 16 bytes per parameter:

Memory Component	Precision / Representation	Bytes per Parameter	Memory for 13B Model
Model Parameters ($\Phi$)	BF16 / FP16	$2$ bytes	26 GB
Gradients ($g$)	BF16 / FP16	$2$ bytes	26 GB
Adam Master Weights	FP32 (Numerical stability)	$4$ bytes	52 GB
Adam Momentum ($m$)	FP32 First Moment	$4$ bytes	52 GB
Adam Variance ($v$)	FP32 Second Moment	$4$ bytes	52 GB
Total Static States	Parameters + Grads + Adam	$16\Phi$ bytes	208 GB (OOM!)

The 208 GB Baseline Reality

Before allocating a single byte for training activations, CUDA workspace buffers, or batch data, a 13B model requires 208 GB of VRAM. Attempting to run standard PyTorch Distributed Data Parallel (DDP)—which replicates this exact 208 GB footprint on every single GPU—guarantees immediate failure.

2. The ZeRO Paradigm: Sharding Without Communication Penalty

In 2019, Microsoft Research introduced the Zero Redundancy Optimizer (ZeRO), which eliminates memory redundancy in DDP by sharding model states across $N$ GPUs:

• ZeRO-Stage 1 (Optimizer State Sharding): The $12\Phi$ Adam optimizer states are partitioned across $N$ devices. Each GPU updates only $1/N$ of the parameters. Static memory reduces from $16\Phi \to 4\Phi + \frac{12\Phi}{N}$. For $N=8$, memory drops from 208 GB to 71.5 GB with zero additional communication overhead!
• ZeRO-Stage 2 (Gradient + Optimizer Sharding): Gradients are also sharded ($2\Phi/N$). As gradients are computed during the backward pass, they are immediately reduced into their assigned shard via Reduce-Scatter and discarded locally. Static memory drops to $2\Phi + \frac{14\Phi}{N}$. For $N=8$, memory drops to 43.2 GB.
• ZeRO-Stage 3 / PyTorch FSDP (Full Parameter Sharding): Every state—parameters, gradients, and optimizer states—is sharded across all GPUs. Each GPU holds only $\frac{16\Phi}{N}$ bytes! For $N=8$, the 13B model static footprint drops to 26 GB, comfortably fitting into standard GPUs.

3. How PyTorch FSDP Executes Under the Hood

If parameters are sharded across 8 GPUs, how does a GPU execute a matrix multiplication that requires the full layer weight matrix? FSDP implements an on-demand transient gather pipeline:

1. Forward Pass: When execution reaches Layer $L$, FSDP issues an asynchronous All-Gather across NVLink to reconstruct the full weights of Layer $L$ in temporary memory. The forward computation runs. Immediately upon completion, the full weights are discarded from memory, keeping only the local shard.
2. Backward Pass: When backpropagation reaches Layer $L$, FSDP again performs an All-Gather to reconstruct the full weights to compute input gradients. Once gradients are calculated, FSDP executes a Reduce-Scatter on gradients to send the partial gradient sums to their corresponding owner GPU, and immediately discards both the full weights and full gradients.
3. Communication Cost: FSDP increases communication volume by exactly $1.5\times$ compared to standard DDP, but unlocks the ability to train models $N\times$ larger without complex model parallelism frameworks.

4. Production PyTorch FSDP Implementation

Here is an enterprise-grade PyTorch implementation using native FullyShardedDataParallel with custom transformer block auto-wrapping and mixed precision:

PyTorch FSDP Distributed Training Script (Python)

import os
import functools
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    ShardingStrategy,
    MixedPrecision,
    BackwardPrefetch,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

# 1. Define Standard Transformer Block
class TransformerBlock(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.Linear(dim, dim, bias=False)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, 4 * dim, bias=False),
            nn.GELU(),
            nn.Linear(4 * dim, dim, bias=False),
        )

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

# 2. Deep Transformer Model
class DeepTransformerModel(nn.Module):
    def __init__(self, num_layers: int = 24, dim: int = 4096):
        super().__init__()
        self.layers = nn.ModuleList([TransformerBlock(dim) for _ in range(num_layers)])
        self.head = nn.Linear(dim, 32000, bias=False)

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return self.head(x)

# 3. Distributed Training Setup
def train_fsdp():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    # Precision Policy: Compute in BF16, gradients in FP32
    mixed_precision_policy = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
        buffer_dtype=torch.bfloat16,
    )

    # Auto-wrap each TransformerBlock independently
    auto_wrap_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={TransformerBlock},
    )

    raw_model = DeepTransformerModel().cuda(local_rank)

    # Wrap with Fully Sharded Data Parallelism
    model = FSDP(
        raw_model,
        auto_wrap_policy=auto_wrap_policy,
        mixed_precision=mixed_precision_policy,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
        device_id=local_rank,
    )

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.1)

    # Synthetic Training Step
    inputs = torch.randn(2, 512, 4096, device=local_rank, dtype=torch.bfloat16)
    outputs = model(inputs)
    loss = outputs.sum()
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

    if local_rank == 0:
        max_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
        print(f"[+] FSDP Step Complete! Peak GPU VRAM: {max_mem_gb:.2f} GB")

    dist.destroy_process_group()

The Shareable Quote: "Do not restrict intelligence to the boundaries of a single silicon chip; shard state across the cluster and let parameters materialize only in the moment of action."

🛠️ Day 1 Actionable Project: FSDP vs. DDP Memory Profiling

Construct a reproducible multi-GPU memory benchmark comparing standard DDP against FSDP:

• Write a script initializing a 3-billion parameter transformer model.
• Run a forward and backward training step using standard PyTorch DDP and log peak VRAM using torch.cuda.max_memory_allocated().
• Re-wrap the identical model using FSDP(sharding_strategy=ShardingStrategy.FULL_SHARD) and measure the exact percentage of memory reduction across a 2-GPU or 4-GPU setup.

🔥 TOMORROW: PART 2 / 5

Tomorrow in Part 2, we tackle models too large to fit a single layer on a single GPU: Day 2: Tensor Parallelism (TP) & Sequence Parallelism — Megatron-LM Matrix Sharding (Vibhūti Yoga).

Architectural & AI Engineering Consulting

If you are architecting distributed pretraining clusters, multi-node fine-tuning infrastructure, or memory-efficient LLM training pipelines, I am available for direct engineering engagements.

Explore Enterprise Engagements →

← Previous Series

ML Systems Day 5: End-to-End MLOps & Canary

Next →

Distributed DL Day 2: Tensor Parallelism