DeepSpeed ZeRO

Overview

DeepSpeed ZeRO is a memory optimization technique, and its implementation inside the DeepSpeed library, that lets you train models far larger than a single GPU can hold. ZeRO stands for Zero Redundancy Optimizer. The core idea is simple: in ordinary data parallel training, every GPU keeps a full copy of the model state, which is wasteful. ZeRO splits that state across the GPUs so each one stores only a share, and brings pieces together only when they are needed.

ZeRO comes in stages. Stage 1 splits optimizer states, stage 2 also splits gradients, and stage 3 also splits the parameters themselves. Extra options can push state out to CPU memory or fast disk. The result is that memory per GPU falls as you add GPUs, and large transformers become trainable on clusters of ordinary accelerators. This is the same family of ideas that appears in PyTorch as FSDP, and it is used by many training stacks, including LLaMA-Factory, Hugging Face training libraries and Megatron-LM based systems. For tabular work, it matters whenever a table transformer, or the language model around it, grows too big for one card.

What it is

DeepSpeed is an open-source library for large-scale deep learning training and inference. ZeRO is its best-known training feature. You use it by wrapping your model and optimizer with a DeepSpeed engine, or by enabling it through a framework integration, and by writing a JSON configuration that selects the stage and related options.

The library offers more than ZeRO. It includes mixed precision support, gradient accumulation, activation checkpointing helpers, mixture-of-experts support, communication optimizations and inference tools. This page focuses on ZeRO.

The three kinds of model state

To see why ZeRO helps, look at what training must keep in memory, per parameter, when using mixed precision and an Adam-style optimizer.

Under plain data parallel training, every GPU holds all 16 bytes for every parameter. A model with 7 billion parameters would need about 112 GB for these states alone, before counting activations. That does not fit on a 40 GB or 80 GB card. Notice that the optimizer states dominate, which is why stage 1 already helps a lot.

The stages

Think of a library. Stage 0: every reader buys all the books. Stage 1: the heavy reference volumes are shared, one reader holds each. Stage 2: the notes are shared too. Stage 3: even the main books are split, and a reader borrows a book from a neighbor only for the minutes it is needed.

Offload options

• ZeRO-Offload. Moves optimizer states, and the optimizer step, to CPU memory. It trades speed for the ability to fit bigger models.
• ZeRO-Infinity. Extends the idea to fast NVMe storage as well, so the usable memory pool becomes CPU memory plus disk.
• ZeRO++. A family of communication reductions that lowers the traffic of stage 3, for example by compressing some transfers. Availability and gains depend on version and hardware.

Where it came from

Training large models hit a memory wall in the late 2010s. Data parallel training scaled compute but not model size. Model parallel methods such as tensor and pipeline parallelism scaled model size but needed changes to the model code, see tensor parallelism and pipeline parallelism.

The ZeRO paper, first posted in 2019 and published in 2020, proposed removing the redundancy in data parallel training by partitioning the model states. It showed that the memory per device could shrink in proportion to the number of devices while keeping the simple data parallel programming model. This idea was implemented in DeepSpeed, an open-source library from a large technology company's research group.

Follow-up work extended it: ZeRO-Offload in 2021 for CPU offload, ZeRO-Infinity in 2021 for using NVMe storage, and later ZeRO++ for reducing communication. The idea influenced other systems. FairScale offered a sharded wrapper, and PyTorch later added native FSDP, which follows the same principle as stage 3. Today ZeRO is a standard part of the toolbox for training and tuning large models.

How it works

Imagine a team of builders, each with their own toolbox. In plain data parallel mode every builder carries every tool, a huge duplication. ZeRO says: let each builder carry only a part of the tools. When a task needs a tool, whoever owns it passes it around for a moment, then everyone puts it back. The team can now build a much bigger structure because nobody is weighed down.

Stage 1 step by step

1. Every GPU runs forward and backward on its own micro-batch, with full parameters and full gradients.
2. Gradients are averaged across GPUs with an all-reduce.
3. Each GPU holds only a slice of the optimizer states. It updates only the parameter slice that matches its states.
4. The updated slices are shared with an all-gather so every GPU again has full parameters.

Stage 2 step by step

1. Forward and backward run as before.
2. Instead of an all-reduce, gradients go through a reduce-scatter. Each GPU receives only the averaged gradient slice that it owns.
3. Each GPU updates its slice using its optimizer-state slice.
4. An all-gather restores full parameters on all GPUs.

Gradients are freed as soon as they are reduced, which saves memory.

Stage 3 step by step

1. Parameters live as slices across GPUs. No GPU holds a complete copy.
2. Before a layer runs forward, the GPUs all-gather that layer's parameters. After the layer, non-owned pieces are released.
3. In the backward pass, the same all-gather happens again for each layer.
4. Gradients are reduce-scattered, so each GPU holds only its slice.
5. Each GPU updates its parameter slice with its optimizer-state slice. No final all-gather of the whole model is needed.

Stage 3 uses prefetching so that the gather for the next layer overlaps with the compute of the current one. With a fast interconnect such as that managed by NCCL, much of the cost is hidden.

Memory arithmetic

For N parameters on W GPUs, with the 16 bytes per parameter breakdown above, the model-state memory per GPU is roughly:

For a 7 billion parameter model on 8 GPUs: stage 0 is about 112 GB, stage 1 about 38.5 GB, stage 2 about 26.25 GB, and stage 3 about 14 GB per GPU, counting model states only. Activations, temporary buffers and fragmentation come on top. These are textbook estimates and real numbers vary.

Where activations fit

ZeRO does not partition activations. Large batches and long sequences can still run out of memory. The usual remedy is gradient checkpointing, which recomputes activations in the backward pass. Some ZeRO versions also support partitioning or offloading checkpointed activations.

Communication cost

Stage 1 and 2 keep communication about the same as plain data parallel training. Stage 3 adds a parameter all-gather in both forward and backward, so total communication is higher, often described as about one and a half times. The benefit is memory. If the network is slow, stage 3 may run slower, so use the lowest stage that fits your model.

Mixed precision

DeepSpeed integrates with bf16 and fp16 training. With fp16, it handles loss scaling to avoid underflow. With bf16, loss scaling is generally not needed. The master copy of weights stays in full precision inside the optimizer states. See mixed precision for the background.

Key concepts and terms

A worked example

This example shows a minimal DeepSpeed config and a training loop for a small model, then how the same config is reused through a higher level library.

The configuration file

{
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 8,
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 200000000,
    "allgather_bucket_size": 200000000
  },
  "gradient_clipping": 1.0,
  "optimizer": {
    "type": "AdamW",
    "params": { "lr": 0.0001, "betas": [0.9, 0.999], "weight_decay": 0.01 }
  },
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 0.0001,
      "warmup_num_steps": 100,
      "total_num_steps": 10000
    }
  }
}

A stage 3 variation with CPU offload

{
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 16,
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "offload_param": { "device": "none" },
    "overlap_comm": true,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "gradient_clipping": 1.0
}

The training loop

import torch
import deepspeed
from torch import nn

class TinyModel(nn.Module):
    def __init__(self, d=1024, layers=8):
        super().__init__()
        self.blocks = nn.Sequential(*[
            nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
            for _ in range(layers)
        ])
        self.head = nn.Linear(d, 1)

    def forward(self, x):
        return self.head(self.blocks(x)).squeeze(-1)

model = TinyModel()

engine, optimizer, _, scheduler = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json",
)

for step in range(1000):
    x = torch.randn(4, 1024, device=engine.device, dtype=torch.bfloat16)
    y = torch.randn(4, device=engine.device, dtype=torch.bfloat16)
    pred = engine(x)
    loss = nn.functional.mse_loss(pred, y)
    engine.backward(loss)   # handles scaling, accumulation and partitioned grads
    engine.step()           # optimizer step on the local partition
    if engine.global_rank == 0 and step % 100 == 0:
        print(step, loss.item())

Launching

deepspeed --num_gpus 8 train.py

For multiple nodes, a host file lists the machines and the launcher starts one process per GPU. Use placeholder host names such as node1.example.internal in documentation and keep real addresses in your own secure configuration.

node1.example.internal slots=8
node2.example.internal slots=8

Using it through a higher level library

Many libraries accept a DeepSpeed config path. In a Hugging Face style trainer, you pass the config in the training arguments, and use values set to "auto" so that the library fills them from its own settings.

{
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "bf16": { "enabled": "auto" },
  "zero_optimization": { "stage": 3, "overlap_comm": true }
}

This keeps one source of truth for batch size and precision, which avoids a common source of mismatch errors.

Saving and loading

With stage 3, the model parameters are sharded, so a simple state dictionary save would only hold local pieces. DeepSpeed writes sharded checkpoints by default, and offers a utility to consolidate them into a single full-precision file for inference or sharing. Make sure to save from all ranks, since each rank holds its own part. You can enable gathering of 16-bit weights on save when you need a single file directly, at the cost of a large temporary memory use on one rank.

Choosing the stage

A planning walk-through

Suppose you want to train a table transformer with 3 billion parameters on 8 GPUs with 40 GB each, using bf16 and an Adam-style optimizer.

Stage 1 already fits the state, leaving about 23 GB for activations and buffers. If activations are heavy because the context holds many rows, move to stage 2, and turn on checkpointing before moving to stage 3. This is a way to pick the cheapest option in communication terms that still fits. The figures are estimates, so confirm with a short real run and look at peak memory.

Reading the memory report

At startup the engine can print an estimate of memory per GPU. Compare it with what a short run actually reaches. If the real peak is much higher than the estimate, look for large temporary buffers, big activations or fragmentation. If the estimate is near the real peak, you can trust it for planning the next size up.

Troubleshooting table

Extra ideas for table workloads

• Variable-size tasks. Pre-training tasks vary in rows and columns. Group tasks of similar size into the same batch to avoid padding waste.
• Context length control. Cap the number of context rows during early training, then raise it, so memory stays predictable.
• Mixed objectives. If you train a shared encoder with several heads, balance the losses and watch gradient norms per head.
• Evaluation hooks. Run a small in-context evaluation on held-out tables every few thousand steps to see whether the prior is improving, not only the loss.

Strengths and limitations

Strengths

• Large memory savings. Per-GPU model state shrinks roughly in proportion to GPU count at stage 3.
• Simple programming model. Your model code stays that of ordinary data parallel training.
• Flexible. Stages and offload options can be tuned to hardware.
• Wide integration. Works with many training frameworks.
• Extends beyond GPU memory. Offload to CPU and NVMe allows very large models on modest hardware.
• Mature. Used widely in open training stacks.

Limitations

• Communication overhead. Stage 3 needs more traffic and may be slow on weak networks.
• Offload slows training. CPU and disk are far slower than GPU memory.
• Configuration complexity. Many options, some of which interact in subtle ways.
• Activations are not partitioned. You still need checkpointing for long sequences.
• Checkpoint handling. Sharded checkpoints need care when moving between different numbers of GPUs.
• Debugging. Partitioned parameters make inspection harder.
• Overlap with other tools. Mixing with other parallelism methods needs testing and version care.

How it compares with alternatives

Frequently the winning design is a combination, such as ZeRO or FSDP for data parallel sharding plus tensor parallelism within a node. Tuning jobs with LoRA often need no sharding at all for models up to a certain size.

Role in tabular foundation models

Tabular foundation models learn directly from raw tables. They must cope with row and column structure, a mix of numeric, categorical, date and text values, missing cells, and column names that may be absent or unhelpful. Models such as TabPFN are pre-trained on huge numbers of synthetic tables drawn from a prior over data-generating processes. At inference they use in-context learning: the labeled rows go into the context and the model predicts for new rows without weight updates. Related designs like TabICL, TabSTAR and FT-Transformer use other encoders and training schemes.

ZeRO is a training infrastructure tool, so its role is to make those models trainable at larger scale. It does not change the model. It changes how state is stored and moved.

Scale pressure in table models

Table transformers can be memory hungry in an unusual way. Attention may run across rows, across columns, or both, as in axial attention designs. The context contains many rows from the training split, plus the query rows. Each cell becomes a token or a part of one. A table with thousands of rows and many columns produces long sequences. Model size also grows as builders aim for richer priors over tasks. When parameters, gradients and Adam states exceed one card, ZeRO stage 2 or 3 is a direct fix, and it is often paired with gradient checkpointing and FlashAttention style kernels for the attention part.

Pre-training on synthetic and real tables

Pre-training draws a continuous stream of fresh synthetic tasks, so data loading and generation can be a bottleneck. ZeRO helps with the model side only. A balanced design generates tasks on CPU workers, keeps GPUs fed, and uses ZeRO to fit and scale the transformer. Many steps of masked-cell prediction, see masked cell modeling, are cheap per step but numerous, which suits data parallel scaling. See also self-supervised learning and synthetic data generation.

Continual learning

Continual learning for table models means adding new domains without erasing old skills. Strategies include replaying old synthetic tasks, regularizing weights as in EWC, distilling from an earlier model as in LwF, or adding adapters. Replay and distillation increase compute and sometimes require two models in memory, which makes sharding more valuable. ZeRO offload can hold the frozen teacher's state cheaply while the student trains.

Imputation, anomaly detection and semantic classification

Heads on top of a shared table encoder can serve several jobs: filling missing cells, scoring rows for anomaly detection, and assigning semantic or sector labels to columns or tables, see semantic sector heads and imputation. Training one encoder with many heads, each with its own loss, grows optimizer state. ZeRO stage 1 or 2 is often enough to absorb that growth, because the optimizer states are the largest component.

Language models in the table stack

The language models used for text-to-SQL, schema reading and column description are tuned with the same stack. Full fine-tuning of a mid-size model on a few GPUs is a classic ZeRO use. For lighter tuning, adapter methods need less sharding. Compare options on the pages for Unsloth, Torchtune and LLaMA-Factory.

Where it is less needed

Many tabular models used in practice are small. Gradient boosted trees such as XGBoost, LightGBM and CatBoost do not use this kind of sharding at all. Small neural table models fit on one GPU. Use ZeRO only when memory or scale demands it. For background on how large table models are designed and evaluated, see the architecture page, the research page and the benchmarks page.

Practical guidance and best practices

Pick the lowest stage that fits

1. Try plain data parallel with bf16 and checkpointing.
2. Move to stage 1 if optimizer states are the issue.
3. Use stage 2 for more savings with similar traffic.
4. Use stage 3 only when parameters themselves do not fit.
5. Add offload last, since it costs speed.

Tune the knobs

Hardware and network

• A fast interconnect inside a node and a good network between nodes matter most for stage 3.
• Use the same GPU model throughout.
• Check that the NCCL version and driver stack are consistent.
• For hardware reference, see NVIDIA H100, H200 and B200.

Reproducibility and checkpoints

• Save sharded checkpoints often. Consolidate only for release.
• Record world size, stage, library versions and config with each checkpoint.
• Test resume on a small run before a long one.
• When changing the number of GPUs, use the conversion utilities and verify results on a fixed batch.

Monitoring

• Track step time, tokens per second, memory use and loss.
• Watch for a rising step time, which can signal network problems.
• Log gradient norms to catch instability early.
• Use a tracker such as MLflow or Weights & Biases.

Config hygiene

• Keep one config per experiment under version control.
• Use "auto" values in framework integrations to avoid conflicts.
• Validate the config with a short run of five to ten steps.
• Keep a baseline config that you know works for your hardware.

Rollout plan for a new training job

1. Run a single-GPU smoke test with a tiny model and the same data code.
2. Run on two GPUs at stage 0 and confirm the loss matches the single-GPU curve closely.
3. Switch to the target stage and confirm the loss still matches for the first hundred steps.
4. Scale to the full GPU count and record tokens per second.
5. Kill the job on purpose and resume from a checkpoint; confirm the loss continues smoothly.
6. Only then start the long run.

Security and data notes

Training jobs often read raw tables that contain sensitive fields. Keep data access controlled, avoid putting credentials in launch scripts or host files, and mask personal data before it reaches the training pipeline. Use placeholder hostnames and secrets managers in shared documentation.

Common mistakes

1. Using stage 3 by default. It adds traffic. Use it only when needed.
2. Forgetting activations. ZeRO does not partition them. Enable checkpointing.
3. Conflicting settings. A batch size in the framework and a different one in the DeepSpeed config cause errors or silent surprises.
4. Saving from one rank only with stage 3. You end up with partial weights.
5. Offloading when it is not needed. Training becomes slow for no gain.
6. Ignoring the network. A slow link can erase the benefit of more GPUs.
7. Mismatched library versions across nodes. Hangs and cryptic errors follow.
8. Not testing resume. A crash after a day of training is costly.
9. Comparing speed without fixing batch and sequence length. Results are not comparable.
10. Assuming ZeRO makes training faster. It makes larger models possible. Speed depends on communication.

Frequently asked questions

What does ZeRO stand for?

ZeRO stands for Zero Redundancy Optimizer. The name reflects the goal of removing duplicated copies of model state across data parallel workers. By partitioning optimizer states, gradients and parameters, each GPU stores only its share instead of a full copy.

What is the difference between ZeRO stage 1, 2 and 3?

Stage 1 partitions the optimizer states, stage 2 also partitions gradients, and stage 3 also partitions the model parameters. Higher stages save more memory per GPU but add more communication. Stage 3 is the one that lets a model larger than one GPU's memory be trained with simple data parallel code.

How is ZeRO related to FSDP?

FSDP in PyTorch follows the same principle as ZeRO stage 3: shard parameters, gradients and optimizer states, gather weights just in time, and release them afterward. ZeRO came first and lives in the DeepSpeed library, while FSDP is native to PyTorch. The choice often depends on which ecosystem your training stack uses.

Does ZeRO speed up training?

Not by itself. It reduces memory so that larger models or larger batches become possible. Stage 1 and 2 add little overhead, while stage 3 and offload can slow down each step. The gain is in what you can train, and sometimes in better GPU use through larger batches.

When should I use CPU or NVMe offload?

Use offload when the model still does not fit after sharding and checkpointing, or when you have few GPUs. It lets you train larger models, at the cost of speed because CPU and disk are slower than GPU memory. Try it as the last step, and measure the slowdown.

Can ZeRO be combined with LoRA?

Yes. LoRA reduces the number of trainable parameters, which shrinks gradients and optimizer states. ZeRO can still shard the frozen base weights at stage 3 for very large models. For small and mid-size models, LoRA alone is often enough and the sharding is not needed.

Does ZeRO help with long sequences?

Only indirectly. ZeRO partitions model states, not activations. For long sequences you need activation checkpointing and memory-efficient attention kernels. ZeRO frees memory from the model state, which leaves more room for activations.

Can ZeRO be used to train tabular foundation models?

Yes, when the model is large enough to need it. A table transformer with attention across rows and columns, trained on streams of synthetic tasks, can use ZeRO to fit its parameters and optimizer states. Smaller tabular models and tree-based methods do not need it.

How do I save a usable model after stage 3 training?

Stage 3 saves sharded checkpoints. For a single full file, use the consolidation utility offered by the library, or enable gathering of 16-bit weights on save. Do the save from all ranks and test loading the result in a clean environment.

What should I watch for in multi-node runs?

Check that all nodes have identical software versions, that the network is fast and stable, and that the launcher and host file list the right slots. Start with a small test of a few steps, confirm memory and speed, and only then launch the long run.

Glossary

• All-gather: A collective that gives every rank all slices.
• All-reduce: A collective that sums values and returns the sum to every rank.
• Bucket: A group of tensors sent in one communication call.
• Data parallelism: Training where each worker handles a different batch slice.
• Engine: The DeepSpeed object that runs the training step.
• Offload: Storing state in CPU memory or on disk.
• Optimizer states: Extra tensors such as moment estimates kept by the optimizer.
• Partition: The slice of a tensor owned by one rank.
• Prefetch: Starting communication early to overlap with compute.
• Rank: A process index in a distributed job.
• Reduce-scatter: A collective that sums values and gives each rank one slice.
• Sharding: Splitting a tensor across devices.
• World size: The number of ranks in the job.

Summary

DeepSpeed ZeRO removes the memory duplication of data parallel training by partitioning optimizer states, gradients and parameters across GPUs, in three stages. With optional offload to CPU and disk, it allows very large models to be trained on modest clusters while keeping the simple programming model of data parallelism. The price is extra communication at higher stages and extra configuration.

For tabular foundation models, ZeRO is infrastructure. Table transformers that attend across rows and columns, learn from streams of synthetic and real tables, and carry several heads for imputation, anomaly scoring and semantic classification can outgrow one GPU. ZeRO, like its PyTorch cousin FSDP, makes that scale reachable. Use the lowest stage that fits, pair it with checkpointing, test checkpoint resume, and keep the network and software versions consistent.

Related

It belongs to Training tooling in my stack, under LLMs, retrieval, and evaluation. Others in the same group: