Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

FSDP › Practical Considerations

Michael BrenndoerferJanuary 22, 20262 min read

Language AI HandbookAll chapters

Part 31 of 67Training Infrastructure

Chapter 226FSDP › Practical Considerations

Reading section 10 of 13 · Practical Considerations · 2 minSection 10 of 13 · 2 min

Section 10 of 13 · 2 min read

FSDP › Practical ConsiderationsOne section · Sections in this chapter10 / 13 ↓

Sections in this chapter13

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Initializing the Process GroupLink Copied

In a real distributed FSDP job, you need to initialize the process group before wrapping the model. The standard pattern using torchrun (PyTorch's distributed launcher) looks like this:

In[2]:

Code

import os

import torch.distributed as dist


def setup_distributed():
    """
    Initialize the distributed process group.
    Called once at the start of each worker process.
    """
    dist.init_process_group(backend="nccl")  # NCCL for GPU-GPU communication
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank


def cleanup_distributed():
    dist.destroy_process_group()

The torchrun launcher sets environment variables (RANK, LOCAL_RANK, WORLD_SIZE, MASTER_ADDR, MASTER_PORT) that dist.init_process_group() reads to coordinate across processes. You launch a training script on 8 GPUs with:

torchrun --nproc_per_node=8 train.py

For multi-node training across 4 nodes of 8 GPUs each (32 total):

torchrun --nproc_per_node=8 --nnodes=4 --node_rank=0 \
         --master_addr=10.0.0.1 --master_port=29500 train.py

Gradient ClippingLink Copied

Gradient clipping with FSDP requires a special API. The standard torch.nn.utils.clip_grad_norm_() function does not work correctly with FSDP because it does not account for the sharded nature of the parameters. FSDP provides FSDP.clip_grad_norm_() as a replacement:

In[4]:

Code

# In a real FSDP training loop, after loss.backward():
fsdp_model.clip_grad_norm_(max_norm=1.0)
optimizer.step()

The reason the standard function fails is that it computes the gradient norm locally on each GPU (seeing only shards), which underestimates the true global gradient norm. The FSDP version coordinates across GPUs to compute the correct global norm before clipping.

Activation CheckpointingLink Copied

Activation checkpointing (gradient checkpointing) works with FSDP and is essential for training with long sequences or large batch sizes. The combination of FSDP and activation checkpointing is particularly powerful: FSDP reduces the persistent memory for parameters and optimizer states, while activation checkpointing reduces the peak activation memory during the backward pass.

To use activation checkpointing with FSDP, you wrap the checkpointing around the individual layers before or after the FSDP wrapping:

In[6]:

Code

from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing,
    checkpoint_wrapper,
)

# Apply activation checkpointing to all TransformerLayer modules
check_fn = lambda submodule: isinstance(submodule, TransformerLayer)
apply_activation_checkpointing(
    fsdp_model,
    checkpoint_wrapper_fn=checkpoint_wrapper,
    check_fn=check_fn,
)

Activation checkpointing trades compute for memory: activations are recomputed during the backward pass rather than stored, roughly doubling the compute cost for the checkpointed layers but reducing activation memory to $O(\sqrt{L})$ where $L$ is the number of layers, instead of $O(L)$ .