LLM Distributed Training: DeepSpeed ZeRO, FSDP, and Multi-GPU Optimization

AI与大数据

Summary

  • Memory bottleneck in large model training: A 7B model in FP16 requires 28GB parameters + 14GB gradients + 14GB optimizer = 56GB; a single A100 80GB card cannot train it
  • DeepSpeed ZeRO-3 reduces memory usage by 8x through parameter sharding; 64 GPUs can train a 70B model
  • PyTorch FSDP is Meta's official solution, seamlessly integrated with the PyTorch ecosystem, and has become the mainstream choice in 2026
  • 3D parallelism (data + tensor + pipeline) is the standard for training ultra-large models; Megatron-LM is the de facto standard
  • This article provides a complete training solution from single-node multi-GPU to thousand-GPU clusters, including Slurm scheduling and fault recovery

Table of Contents

Memory Bottleneck in Large Model Training

Training Memory Breakdown

┌──────────────────────────────────────────────────────────────┐
│              Large Model Training Memory Breakdown            │
│                                                                │
│  Model Parameters(Weights):     2 × params × dtype_size      │
│  Gradients:                     2 × params × dtype_size      │
│  Optimizer States(Adam):        12 × params × dtype_size (FP32)│
│  Activations:                   batch × seq_len × hidden × layers│
│                                                                │
│  7B Model FP16 Training Memory:                               │
│  ┌──────────────────────────────────────────────────────┐    │
│  │ Parameters:     14GB                                  │    │
│  │ Gradients:      14GB                                  │    │
│  │ Optimizer:      56GB (FP32 m and v)                  │    │
│  │ Activations:    8-16GB (depends on batch and seq_len) │    │
│  │ ─────────────────────────                            │    │
│  │ Total:          92-100GB ❌ Exceeds single A100 80GB  │    │
│  └──────────────────────────────────────────────────────┘    │
└──────────────────────────────────────────────────────────────┘

Training Requirements by Model Scale

Model Parameters Training Memory (FP16+Adam) Minimum GPU Config Recommended Config
1.8B 1.8B 26GB 1×A100 40GB 1×A100 80GB
7B 7B 100GB 2×A100 80GB 4×A100 80GB
13B 13B 186GB 4×A100 80GB 8×A100 80GB
70B 70B 1TB 16×A100 80GB 32×H100 80GB
671B 671B 9.6TB 128×H100 80GB 256×H100 80GB

DeepSpeed ZeRO: 3-Stage Memory Optimization

ZeRO Three-Stage Principles

┌──────────────────────────────────────────────────────────────┐
│              ZeRO Three-Stage Memory Optimization             │
│                                                                │
│  ZeRO-1: Optimizer State Partitioning                         │
│  ┌──────────────────────────────────────────────────────┐    │
│  │ Each GPU stores only 1/N of optimizer states (m and v)│    │
│  │ Memory savings: 4× (optimizer is the largest portion) │    │
│  │ Communication: Same as DDP                            │    │
│  └──────────────────────────────────────────────────────┘    │
│                                                                │
│  ZeRO-2: Optimizer + Gradient Partitioning                    │
│  ┌──────────────────────────────────────────────────────┐    │
│  │ Each GPU stores only 1/N of optimizer states + gradients│   │
│  │ Memory savings: 8×                                    │    │
│  │ Communication: Same as DDP (gradient reduce-scatter)  │    │
│  └──────────────────────────────────────────────────────┘    │
│                                                                │
│  ZeRO-3: Optimizer + Gradient + Parameter Partitioning        │
│  ┌──────────────────────────────────────────────────────┐    │
│  │ Each GPU stores only 1/N of all states                │    │
│  │ Memory savings: N× (N=number of GPUs)                 │    │
│  │ Communication: Increases 1.5× (requires all-gather)   │    │
│  └──────────────────────────────────────────────────────┘    │
└──────────────────────────────────────────────────────────────┘

DeepSpeed ZeRO-3 Training Configuration

{
  "train_batch_size": 128,
  "train_micro_batch_size_per_gpu": 2,
  "gradient_accumulation_steps": 8,
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "bf16": {
    "enabled": true
  },
  "gradient_clipping": 1.0,
  "prescale_gradients": false,
  "wall_clock_breakdown": false
}

Memory Comparison Across ZeRO Stages (7B Model, 8×A100 80GB)

Stage Parameters Gradients Optimizer Activations Total Per GPU
DDP 14GB 14GB 56GB 12GB 96GB 96GB
ZeRO-1 14GB 14GB 7GB 12GB 47GB 47GB
ZeRO-2 14GB 1.75GB 7GB 12GB 35GB 35GB
ZeRO-3 1.75GB 0.22GB 0.88GB 12GB 15GB 15GB

DeepSpeed Launch Script

deepspeed --num_gpus=8 \
    --num_nodes=2 \
    --master_addr=$MASTER_ADDR \
    --master_port=29500 \
    train.py \
    --model_name_or_path Qwen/Qwen2.5-7B \
    --data_path ./data/train.jsonl \
    --bf16 \
    --deepspeed ds_config_z3.json \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --learning_rate 2e-5 \
    --num_train_epochs 3 \
    --warmup_ratio 0.03 \
    --lr_scheduler_type cosine \
    --weight_decay 0.01 \
    --save_strategy epoch \
    --output_dir ./output \
    --gradient_checkpointing true \
    --logging_steps 10 \
    --report_to tensorboard

PyTorch FSDP: Meta's Official Solution

FSDP vs DeepSpeed ZeRO

Dimension DeepSpeed ZeRO-3 PyTorch FSDP
Developer Microsoft Meta
Integration Standalone library PyTorch native
Parameter Sharding ✅ ✅
Gradient Sharding ✅ ✅
Optimizer Sharding ✅ ✅
CPU Offloading ✅ ✅
Mixed Precision ✅ ✅
Pipeline Parallelism ✅ (DeepSpeed PP) ⚠️ (Manual setup)
Debug Friendliness Medium High (PyTorch native)
Community Activity ⭐⭐⭐⭐ ⭐⭐⭐⭐⭐

FSDP Training Code

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy, MixedPrecision
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers import AutoModelForCausalLM, AutoTokenizer

torch.distributed.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B",
    torch_dtype=torch.bfloat16,
    device_map={"": local_rank},
)

mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

auto_wrap_policy = transformer_auto_wrap_policy(
    transformer_layer_names=["Qwen2DecoderLayer"],
)

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=mp_policy,
    auto_wrap_policy=auto_wrap_policy,
    device_id=local_rank,
    use_orig_params=True,
)

optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=1000)

for epoch in range(num_epochs):
    for batch in dataloader:
        batch = {k: v.to(local_rank) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

3D Parallelism: Data + Tensor + Pipeline

3D Parallelism Architecture

┌──────────────────────────────────────────────────────────────┐
│              3D Parallelism Architecture                      │
│                                                                │
│  Data Parallelism(DP): 4 groups × 4 GPUs per group            │
│  ┌──────────────────────────────────────────────────────┐    │
│  │  DP Group 0      DP Group 1      DP Group 2      DP Group 3│
│  │  ┌──────────┐  ┌──────────┐  ┌──────────┐  ┌──────┐│    │
│  │  │ TP=2     │  │ TP=2     │  │ TP=2     │  │TP=2  ││    │
│  │  │ PP=2     │  │ PP=2     │  │ PP=2     │  │PP=2  ││    │
│  │  │ ┌──┬──┐ │  │ ┌──┬──┐ │  │ ┌──┬──┐ │  │┌──┬─┐││    │
│  │  │ │G0│G1│ │  │ │G4│G5│ │  │ │G8│G9│ │  ││12│13│││    │
│  │  │ ├──┼──┤ │  │ ├──┼──┤ │  │ ├──┼──┤ │  │├──┼──┤││    │
│  │  │ │G2│G3│ │  │ │G6│G7│ │  │ │10│11│ │  ││14│15│││    │
│  │  │ └──┴──┘ │  │ └──┴──┘ │  │ └──┴──┘ │  │└──┴──┘││    │
│  │  └──────────┘  └──────────┘  └──────────┘  └──────┘│    │
│  └──────────────────────────────────────────────────────┘    │
│                                                                │
│  TP=Tensor Parallelism(intra-layer)  PP=Pipeline Parallelism(inter-layer)│
│  DP=Data Parallelism(data partitioning)  Total: 16 GPUs      │
└──────────────────────────────────────────────────────────────┘

3D Parallelism Decision Guide

Model Scale GPUs Recommended Strategy Notes
7B 8 DP=8 Pure data parallelism suffices
13B 16 DP=8, TP=2 Add tensor parallelism
70B 64 DP=8, TP=4, PP=2 3D parallelism
671B 256 DP=16, TP=8, PP=2 3D parallelism + ZeRO

Multi-Node Training and Fault Recovery

Slurm Launch Script

#!/bin/bash
#SBATCH --job-name=llm-train
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=8
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=8
#SBATCH --mem=512G
#SBATCH --time=72:00:00
#SBATCH --partition=gpu-a100

export MASTER_ADDR=$(scontrol show hostname $SLURM_JOB_NODELIST | head -n1)
export MASTER_PORT=29500
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0
export NCCL_IB_HCA=mlx5

srun torchrun \
    --nnodes=$SLURM_JOB_NUM_NODES \
    --nproc_per_node=8 \
    --rdzv_id=$SLURM_JOB_ID \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    train.py \
    --model_name_or_path Qwen/Qwen2.5-7B \
    --deepspeed ds_config_z3.json \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --learning_rate 2e-5 \
    --num_train_epochs 3 \
    --output_dir ./output \
    --save_strategy steps \
    --save_steps 500 \
    --save_total_limit 5

Training Fault Recovery

from transformers import Trainer

class FaultTolerantTrainer(Trainer):
    def _save_checkpoint(self, model, trial, metrics=None):
        super()._save_checkpoint(model, trial, metrics)
        self._save_training_state()

    def _save_training_state(self):
        state = {
            "global_step": self.state.global_step,
            "epoch": self.state.epoch,
            "rng_state": torch.cuda.get_rng_state().cpu().numpy().tolist(),
        }
        with open(f"{self.args.output_dir}/training_state.json", "w") as f:
            json.dump(state, f)

    def _load_training_state(self):
        state_path = f"{self.args.output_dir}/training_state.json"
        if os.path.exists(state_path):
            with open(state_path) as f:
                state = json.load(f)
            return state
        return None

Training Performance Optimization Checklist

Optimization Effect Configuration
Flash Attention 2 Training speed +30% attn_implementation="flash_attention_2"
Gradient Checkpointing Memory -40% gradient_checkpointing=True
BF16 Mixed Precision Memory -50% bf16=True
NCCL IB Communication Communication latency -60% NCCL_IB_DISABLE=0
Prefetch Sharded Parameters Communication-computation overlap stage3_prefetch_bucket_size
Gradient Accumulation Equivalent large batch gradient_accumulation_steps=8

The core of large model distributed training lies in memory optimization and communication optimization. DeepSpeed ZeRO-3 reduces memory by N times through parameter sharding, FSDP is the PyTorch native solution that is easier to debug, and 3D parallelism is the standard for ultra-large models.

Key Takeaways:

  1. ZeRO-3 is the ultimate memory optimization solution, but communication increases by 1.5x
  2. FSDP is the PyTorch native solution and has become the mainstream choice in 2026
  3. 3D parallelism: 7B uses DP, 13B adds TP, 70B+ uses DP+TP+PP
  4. Flash Attention 2 + Gradient Checkpointing are training essentials
  5. Fault recovery and checkpointing are must-haves for long-running training

Related Reading:

Authoritative References:

Try these browser-local tools — no sign-up required →

#大模型分布式训练#DeepSpeed ZeRO#FSDP训练#Megatron-LM#多GPU训练优化#2026