DeepSpeed: Training Foundation Models Beyond a Single GPU

Ever wonder how Foundation Models — too large to fit into a single GPU — are trained?

Enter DeepSpeed, an open-source framework from Microsoft designed to scale deep learning training to billions of parameters. It integrates seamlessly with the PyTorch ecosystem and uses sharding strategies to distribute model weights, gradients, and optimizer states across GPUs, CPUs, and even NVMe storage.

Why Optimizer States Matter

Before diving into DeepSpeed, let’s recap what makes training large models so memory-intensive.

During backpropagation, optimizers like Adam help update weights using gradients. But Adam (and similar algorithms) maintains optimizer states such as running averages of gradients(1st moment) and squared gradients (2nd moment).

Here’s the problem:

• Optimizer states can take 3-4x more memory than the weights themselves.
• These are usually stored in FP32 for numerical stability.
• That means less GPU memory, especially when you have a large model that barely fits into a single GPU.

DeepSpeed addresses this problem by splitting the optimizer states among GPUs (Zero-1,2,3) and also offloading them to CPU RAM or even NVMe disks (ZeRO-3 Offload) — freeing GPU memory for actual training.

ZeRO Optimization Stages: The Heart of DeepSpeed

The ZeRO (Zero Redundancy Optimizer) in DeepSpeed is implemented in stages, each addressing a specific kind of memory redundancy:

With ZeRO-3 Offload, optimizer states can even be pushed to CPU or NVMe storage. This allows training massive models on relatively modest GPUs — but introduces communication overhead.

Zero vs. FSDP

• ZeRO Stage 3: Shards tensors (parameters, gradients, optimizer states) across GPUs. It doesn’t care about layers; sharding happens automatically at the tensor level.
• FSDP: Shards full layers (or submodules) across GPUs. You wrap modules in FSDP, and it shards the parameters of that module.

FSDP shards module by module, each GPU only keeps its shard of the parameters for any module. During forward and backward which happens sequentially module after module, the whole module parameters will be gathered at each GPU. In contrast, ZeRO reconstructs full tensors on the fly during forward and backward passes.

In other words, FSDP, through its auto_wrap_policy, gives users finer control over how a model is sharded across GPUs, directly impacting memory efficiency.

Bottom line: ZeRO-3 is very powerful for parameter sharding, but without module-level control (like FSDP wrapping), forward/backward memory is less predictable and can be higher than expected.

How do model parallel and data parallel work together?

Suppose you have 8 GPUs:

• 4 GPUs: form a ZeRO-3 shard group (Data Parallel replica)
• Each GPU stores only a part of each weight
• Another 4 GPUs: another ZeRO-3 shard group (a data parallel replica)
• Each replica sees a different batch of data

Putting it together:

Data Parallel Split: the full batch is divided across data parallel replicas:

Within each replica:

• ZeRO-3 shards the model across GPUs
• Each GPU in the group computes its local forward using full X_i for the shard

After backward:

• Each ZeRO-3 group computes gradients (sharded)
• Data-parallel all-reduce happens across groups to sync optimizer updates

Diagrammatic intuition

Full batch X
│
├─ Data Parallel Split ─> X1 (replica 1) → ZeRO-3 shard GPUs 0–3
│                         X2 (replica 2) → ZeRO-3 shard GPUs 4–7
│
Within each ZeRO-3 group j:
GPU0  GPU1  GPU2  GPU3
 W0    W1    W2    W3
 Xj    Xj    Xj    Xj   → compute Y^j
 → concat Y_i → Y^j

• In column-wise sharding, the partial outputs from each GPU (Y_i=X*W_i) are combined by concatenation along the output dimension
• Each data-parallel replica computes its own Y1, Y2, … independently
• Gradients are all-reduced across data-parallel replicas

Now it’s time to get your hand dirty

Installing DeepSpeed

DeepSpeed comes with many optimized CUDA/C++ extensions (called ops):

• AdamCPU
• FusedAdam
• AIO …and more.

There are two installation modes:

1. Default mode — installs ops dynamically using Ninja (assumes CUDA/GPU compatibility).
2. Pre-install mode — explicitly builds ops based on environment flags.

For example, to install DeepSpeed with CPU Adam optimizer support:

DS_BUILD_CPU_ADAM=1 pip install deepspeed

Since DeepSpeed requires CUDA for many ops, it is recommended to install torch version with cuda from torch-cuda wheel before installing DeepSpeed with CUDA ops:

uv pip install --no-build-isolation --no-cache torch==2.7.1+cu128 torchvision==0.22.1+cu128 \
    --extra-index-url https://download.pytorch.org/whl/cu128
DS_BUILD_AIO=1 DS_BUILD_CPU_ADAM=1 DS_BUILD_FUSED_ADAM=1 uv pip install --no-build-isolation --index-strategy unsafe-best-match deepspeed==0.17.6 --no-cache

💡 Tip:

1. You can specify a torch wheel with Cuda inside your pyproject.toml:

[tool.uv.sources]
torch = { index = "pytorch-cu128" }
torchvision = { index = "pytorch-cu128" }

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

2. Run the environment report to verify DeepSpeed installation:

ds_report
# or
python -m deepspeed.env_report

Using DeepSpeed in PyTorch Lightning

The great news: DeepSpeed integrates gracefully into PyTorch Lightning.

You can use it with a simple trainer’s strategy:

from lightning.pytorch import Trainer
trainer = Trainer(strategy="deepspeed_stage_3", accelerator="gpu", devices=4)

Or pass in a custom config with DeepSpeedStrategy:

from lightning.pytorch.strategies import DeepSpeedStrategy
deepspeed_config = {
    "zero_allow_untested_optimizer": True,
    "optimizer": {
        "type": "OneBitAdam",
        "params": {"lr": 3e-5, "betas": [0.998, 0.999], "eps": 1e-5,
                   "weight_decay": 1e-9, "cuda_aware": True},
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {"warmup_max_lr": 3e-5, "warmup_num_steps": 100},
    },
    "zero_optimization": {
        "stage": 2,  # ZeRO Stage 2
        "offload_optimizer": {"device": "cpu"},
        "contiguous_gradients": True,
        "overlap_comm": True,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
    },
}

trainer = Trainer(
    accelerator="gpu", devices=4,
    strategy=DeepSpeedStrategy(config=deepspeed_config),
    precision=16
)

When using offload, make sure you also specify a DeepSpeed optimizer that is compatible, for example for the above ‘zero_optimization’ with offload_optimizer to CPU, use DeepSpeedCPUAdam (https://www.deepspeed.ai/tutorials/zero-offload/):

  optimizer:
    class_path: deepspeed.ops.adam.DeepSpeedCPUAdam
    init_args:
      lr: 0.0001

Checkpointing with ZeRO-3

One caveat: ZeRO-3 saves sharded checkpoints by default.

Checkpoint with Zero-3

To merge them back into a single .pt file:

from lightning.pytorch.utilities.deepspeed import convert_zero_checkpoint_to_fp32_state_dict

save_path = "lightning_logs/version_0/checkpoints/epoch=0-step=0.ckpt/"
output_path = "lightning_model.pt"

convert_zero_checkpoint_to_fp32_state_dict(save_path, output_path)

Data Parallelism in DeepSpeed

In DeepSpeed, **train_batch_size is a crucial configuration parameter, but it’s important to understand what it actually represents** because it behaves a bit differently than in vanilla DP (refer to the explanation of Model sharding above):

• train_batch_size is the global batch size for training.
• It represents the total number of samples processed per training step across all GPUs and all gradient accumulation steps.

Formally:

train_batch_size = train_micro_batch_size_per_gpu * gradient_accumulation_steps * num_gpus

You can specify two of these parameters, and DeepSpeed infers the third automatically.

Conclusion

DeepSpeed shines when training models that don’t fit on a single GPU. By sharding weights, gradients, and optimizer states, it enables training on a cluster of smaller GPUs.

But there’s a tradeoff:

• Pros: Train bigger models, fit larger batch sizes, leverage CPU/NVMe offloading.
• Cons: Extra communication overhead.

In short:

• If your model doesn’t fit → DeepSpeed is a lifesaver.
• If it does fit → Stick to Data Parallel.