Scaling LLM Training: Multi-Node Distributed Training with PyTorch FSDP2 on H100 & A100 Clusters

1. Introduction

The Challenge of LLM Training

Training a large language model on a single GPU or even a single 8-GPU node runs into a hard wall well before you reach useful model sizes. A 70B-parameter model in FP32 needs roughly 280GB just to hold the weights, before accounting for gradients, optimizer states (Adam alone typically adds two extra copies of the parameters), and activations. Even at BF16, the combined footprint of weights, gradients, and optimizer state for a 70B model can exceed 1TB.

No single GPU, and no single 8x80GB node, has that much memory. Once a model's total training memory footprint exceeds what a node can hold, multi-node scaling stops being an optimization and becomes a requirement.

What Is FSDP2?

Fully Sharded Data Parallel (FSDP) is a data-parallel training strategy that, instead of replicating the full model on every GPU, shards the model's parameters, gradients, and optimizer states across all participating devices. Each GPU holds only a slice of the model at rest, and full parameters are reconstructed on the fly layer by layer only when needed for computation, then discarded again immediately after.

FSDP2 is the second-generation implementation of this idea in PyTorch, built around a different internal parameter representation than the original FSDP.

FSDP1 vs. FSDP2: The Paradigm Shift

The original FSDP (now referred to as FSDP1) was a massive leap forward for distributed training, but it had architectural limitations. It represented sharded parameters as FlatParameter objects — essentially concatenating many individual tensors into one massive 1D buffer and splitting that buffer across ranks.

While this worked for basic data parallelism, it made partial parameter freezing (like LoRA), memory accounting, and exporting checkpoints incredibly complex.

FSDP2 replaces the flat buffer approach with DTensor (Distributed Tensor) based per-parameter sharding. Each original parameter keeps its own identity and metadata, but is simply sharded along dimension 0.

Key Architectural Differences

Why This Matters: When you run a 70B parameter model on an 8x H100 cluster, FSDP2 ensures that memory usage scales exactly linearly across the mesh, preventing unexpected Out-Of-Memory (OOM) crashes during gradient syncing and checkpoint saving.

2. Hardware & Infrastructure Overview

When scaling LLMs, your underlying hardware determines whether your cluster actually scales or simply bottlenecks at the network layer.

Network Topology: Why Interconnect Matters

FSDP2 relies on continuous All-Gather (reconstructing weights) and Reduce-Scatter (syncing gradients) operations. Within a single node, NVLink provides the massive bandwidth required. Across nodes, this traffic crosses the network fabric. InfiniBand or RoCE (RDMA over Converged Ethernet) is mandatory here; standard Ethernet will bottleneck a multi-node FSDP2 job badly enough that adding GPUs yields zero added throughput.

3. Setting Up the Distributed Environment

Essential NCCL Environment Variables

NCCL (NVIDIA Collective Communications Library) handles the GPU-to-GPU communication. Set these environment variables before running your script:

export NCCL_IB_DISABLE=0          # Ensure InfiniBand is used when available
export NCCL_NET_GDR_LEVEL=2       # Enable GPUDirect RDMA for lower-latency transfers
export NCCL_DEBUG=INFO            # Verbose logging to validate InfiniBand/RDMA links
export NCCL_SOCKET_IFNAME=eth0    # Set explicitly if the node has multiple network interfaces

Launch Configuration with torchrun

To launch your training script across two 8-GPU nodes, execute this on your master node:

torchrun \
  --nnodes=2 \
  --nproc_per_node=8 \
  --rdzv_id=fsdp2_run \
  --rdzv_backend=c10d \
  --rdzv_endpoint=<master_node_ip>:29500 \
  train.py

4. The Complete train.py Implementation

Here is the complete, cohesive implementation for setting up the Device Mesh, wrapping a dummy LLM architecture with fully_shard, handling mixed precision, and saving checkpoints via DCP.

import os
import torch
import torch.nn as nn
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_model_state_dict, set_model_state_dict

# 1. Initialize Distributed Process Group
torch.distributed.init_process_group(backend="nccl")
local_rank = int(os.environ.get("LOCAL_RANK", 0))
torch.cuda.set_device(local_rank)

# 2. Define the Device Mesh (e.g., 2 Nodes x 8 GPUs)
# Ensure these dimensions match your torchrun --nnodes and --nproc_per_node
world_size = int(os.environ.get("WORLD_SIZE", 16))
num_nodes = world_size // 8
mesh = init_device_mesh("cuda", (num_nodes, 8))

# 3. Define a Dummy Model (Replace with your actual LLM architecture)
class SimpleTransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.ffn = nn.Linear(4096, 4096)
        self.activation = nn.GELU()
        
    def forward(self, x):
        return self.activation(self.ffn(x))

class DummyLLM(nn.Module):
    def __init__(self, num_layers=8):
        super().__init__()
        self.layers = nn.ModuleList([SimpleTransformerBlock() for _ in range(num_layers)])
        self.head = nn.Linear(4096, 1000)
        
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return self.head(x)

# Instantiate model and move to CUDA
model = DummyLLM().cuda()

# 4. Configure Mixed Precision Policy
# BF16 is ideal for A100. On H100, native FP8 optimizations can be utilized via Transformer Engine.
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
)

# 5. Apply FSDP2 Bottom-Up Sharding
# Wrap individual layers first to allow granular all-gather/reduce-scatter operations
for layer in model.layers:
    fully_shard(layer, mesh=mesh, mp_policy=mp_policy)

# Finally, wrap the root model
fully_shard(model, mesh=mesh, mp_policy=mp_policy)

if local_rank == 0:
    print(f"Model successfully sharded across the {num_nodes}x8 device mesh.")

# --- DUMMY TRAINING LOOP ---
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
dummy_input = torch.randn(32, 4096, device="cuda", dtype=torch.bfloat16)

# Forward & Backward Pass
output = model(dummy_input)
loss = output.sum()
loss.backward()
optimizer.step()

# 6. Distributed Checkpointing (DCP)
# DCP safely saves DTensor states without OOM errors by avoiding a full-gather onto a single rank.
checkpoint_dir = "fsdp2_checkpoints/step_1"

# Save State
model_state = get_model_state_dict(model)
dcp.save(model_state, checkpoint_id=checkpoint_dir)
if local_rank == 0:
    print(f"Checkpoint safely saved at {checkpoint_dir}")

# Cleanup
torch.distributed.destroy_process_group()

5. Conclusion & Next Steps

FSDP2's move to DTensor-based, per-parameter sharding makes multi-node LLM training significantly easier to reason about than it was under FSDP1. Memory usage is predictable, checkpointing avoids OOM disasters natively via DCP, and the abstraction scales effortlessly as your cluster grows.

If you are ready to put this into practice, deploying it directly on dedicated bare-metal infrastructure is the most efficient path forward.