DeepSpeed: Training Huge Models on Limited GPUs with ZeRO and Offload
Hey everyone, today we go into slightly heavier territory that becomes very useful once you start working with large models: DeepSpeed, Microsoft's library for making giant model training possible on a limited amount of hardware. If you have ever hit CUDA out of memory even after shrinking the batch size to one, or looked at a 7 billion parameter model and thought "no way this fits on my GPU", DeepSpeed is what you need to learn.
The core magic is a technology called ZeRO, short for Zero Redundancy Optimizer. The idea is clever and actually quite simple: in ordinary multi-GPU training, every GPU keeps a full copy of the model, the gradients, and the optimizer state. That is enormously wasteful. ZeRO shards all of that and spreads the pieces across GPUs so nothing is stored redundantly.
In this tutorial we cover where training memory actually goes, the three ZeRO stages and when to use each, installation, the JSON configuration file, integration with the HuggingFace Trainer and with Accelerate, CPU offload, and debugging tips that will save you hours.
Introduction
Before the code, you need to understand what your GPU memory is really spent on.
Where GPU Memory Goes
Many people assume training memory is dominated by model weights. It is not. Let us do the math for a 1 billion parameter model trained with Adam in mixed precision.
Model weights in fp16 take about 2 bytes per parameter, so 2 GB.
Gradients in fp16 also take 2 bytes per parameter, another 2 GB.
Adam's optimizer state keeps an fp32 copy of the weights, fp32 momentum, and fp32 variance. That is 4 plus 4 plus 4, so 12 bytes per parameter, or 12 GB.
That already totals 16 GB of state before counting activations. Activations, the intermediate outputs kept for the backward pass, can add several more GB depending on batch size and sequence length.
Now imagine you have 4 GPUs. In plain data-parallel training every number above is multiplied by four, because each GPU keeps a full copy. You burn 64 GB to hold information that is genuinely only 16 GB. That redundancy is exactly what ZeRO removes.
The Three ZeRO Stages
ZeRO splits its work into three stages you can choose from.
Stage 1 shards the optimizer state across GPUs. This removes the single largest memory block (12 GB in the example above) at very low communication cost. It is a safe choice for almost everyone.
Stage 2 shards the optimizer state and the gradients. Bigger savings, still modest communication overhead. This is the sweet spot most people run.
Stage 3 shards the optimizer state, the gradients, and the model parameters themselves. No single GPU ever holds the whole model. Parameters are gathered momentarily as each layer executes, then released. The memory savings are dramatic, but inter-GPU traffic increases substantially, so you want fast interconnect such as NVLink or InfiniBand to avoid a slowdown.
On top of that sits offload, which moves optimizer state or even parameters to CPU RAM or NVMe. That lets you train models that truly do not fit on the GPU, at a significant speed cost.
My rule of thumb: start at stage 2. If you still run out of memory, go to stage 3. If it is still not enough, enable CPU offload. Do not jump straight to stage 3 with offload, because you will pay in throughput for no reason.
Installation
DeepSpeed needs Linux, a CUDA-enabled PyTorch, and a C++ compiler. On Windows, use WSL2.
python -m venv venv
source venv/bin/activate
pip install torch
pip install deepspeed
Verify the installation and see which operators are available:
dsreport
That prints a table of extension statuses. Not everything needs to be green. What matters is that torch and deepspeed are detected, and if you plan to use CPU offload, that cpuadam can compile.
When extension compilation fails, it is usually because CUDAHOME is unset or the compiler version mismatches. Setting these environment variables often fixes it:
export CUDAHOME=/usr/local/cuda
export DSBUILDCPUADAM=1
The Configuration File
All DeepSpeed behavior is driven by one JSON file. Here is a ZeRO stage 2 configuration you can use directly.
{
"trainmicrobatchsizepergpu": 4,
"gradientaccumulationsteps": 8,
"gradientclipping": 1.0,
"bf16": {
"enabled": true
},
"zerooptimization": {
"stage": 2,
"overlapcomm": true,
"contiguousgradients": true,
"reducebucketsize": 5e7,
"allgatherbucketsize": 5e7
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": 2e-5,
"betas": [0.9, 0.999],
"eps": 1e-8,
"weightdecay": 0.01
}
},
"scheduler": {
"type": "WarmupDecayLR",
"params": {
"warmupminlr": 0,
"warmupmaxlr": 2e-5,
"warmupnumsteps": 100,
"totalnumsteps": 5000
}
},
"stepsperprint": 50,
"wallclockbreakdown": false
}
Save it as dsconfig.json. A few notes.
trainmicrobatchsizepergpu is what actually enters one GPU in a single forward pass. Your effective batch is that value times gradientaccumulationsteps times the number of GPUs.
overlapcomm overlaps inter-GPU communication with computation and usually buys free speed.
bf16 is safer than fp16 for large models because it needs no loss scaling and rarely produces NaNs. If your GPU does not support bf16 (a V100, for example), replace that block with "fp16": {"enabled": true}.
For stage 3 with offload, the zero block becomes:
"zerooptimization": {
"stage": 3,
"offload
optimizer": { "device": "cpu", "pinmemory": true },
"offload
param": { "device": "cpu", "pinmemory": true },
"overlap
comm": true,
"contiguousgradients": true,
"stage3gather16bitweightsonmodelsave": true
}
Note stage3gather16bitweightsonmodelsave. Without it, your saved checkpoint contains only shards rather than a whole model, and you will be very confused when loading it back.
Path 1: Through the HuggingFace Trainer
This is the easiest and most common route. You write no DeepSpeed code at all, you just point at the config file.
from transformers import (
AutoModelForCausalLM, AutoTokenizer,
TrainingArguments, Trainer, DataCollatorForLanguageModeling,
)
from datasets import loaddataset
modelname = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.frompretrained(modelname)
tokenizer.padtoken = tokenizer.eostoken
model = AutoModelForCausalLM.frompretrained(modelname)
ds = loaddataset("wikitext", "wikitext-2-raw-v1", split="train[:2000]")
def tokenize(batch):
return tokenizer(batch["text"], truncation=True, maxlength=512)
ds = ds.map(tokenize, batched=True, removecolumns=ds.columnnames)
args = TrainingArguments(
outputdir="out",
perdevicetrainbatchsize=4,
gradientaccumulationsteps=8,
numtrainepochs=1,
loggingsteps=50,
savestrategy="epoch",
bf16=True,
deepspeed="dsconfig.json", # the only DeepSpeed line
)
trainer = Trainer(
model=model,
args=args,
traindataset=ds,
datacollator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()
Launch it with:
deepspeed --numgpus=4 train.py
One thing to know: if you specify optimizer and scheduler in the JSON, those override your TrainingArguments. To avoid confusion you can use the special "auto" value in the JSON and let the Trainer decide.
"optimizer": {
"type": "AdamW",
"params": { "lr": "auto", "weightdecay": "auto" }
}
With "auto", DeepSpeed pulls the values from TrainingArguments, giving you a single source of truth.
Path 2: Through Accelerate
If you already use Accelerate, enabling DeepSpeed is purely a configuration matter.
accelerate config
Answer yes when asked about DeepSpeed, then point at your dsconfig.json or answer the ZeRO stage prompts. After that your existing Accelerator() based script runs unchanged:
accelerate launch train.py
This is my favorite combination, because the training code stays pure PyTorch and can still run without DeepSpeed while you debug on a laptop.
Path 3: The DeepSpeed API Directly
When you need full control, here is the native form.
import deepspeed
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
model = nn.Sequential(
nn.Linear(512, 2048), nn.ReLU(),
nn.Linear(2048, 2048), nn.ReLU(),
nn.Linear(2048, 10),
)
x = torch.randn(4096, 512)
y = torch.randint(0, 10, (4096,))
loader = DataLoader(TensorDataset(x, y), batchsize=4)
modelengine, optimizer, , scheduler = deepspeed.initialize(
model=model,
modelparameters=model.parameters(),
config="dsconfig.json",
)
criterion = nn.CrossEntropyLoss()
for epoch in range(3):
for xb, yb in loader:
xb = xb.to(modelengine.device)
yb = yb.to(modelengine.device)
loss = criterion(modelengine(xb), yb)
modelengine.backward(loss) # not loss.backward()
modelengine.step() # includes zerograd and scheduler
if modelengine.globalrank == 0:
print(f"epoch {epoch} loss {loss.item():.4f}")
modelengine.savecheckpoint("checkpoints", tag="epochfinal")
Three differences from plain PyTorch matter here. You call modelengine.backward(loss) rather than loss.backward(). You call modelengine.step(), which already handles the optimizer step, scheduler step, and gradient zeroing. And you must not call optimizer.zerograd() yourself, because it breaks DeepSpeed's internal gradient accumulation.
Saving and Converting Checkpoints
DeepSpeed checkpoints are sharded. They are not a single pytorchmodel.bin but many per-rank files. To convert them into standard weights anyone can load, DeepSpeed ships a helper.
python -m deepspeed.utils.zerotofp32 checkpoints/epochfinal modelfp32.bin
Or from Python:
from deepspeed.utils.zerotofp32 import loadstatedictfromzerocheckpoint
model = load
statedictfromzerocheckpoint(model, "checkpoints/epochfinal")
If you use the HuggingFace Trainer with stage3gather16bitweightsonmodelsave enabled, trainer.savemodel() already writes a consolidated model, so this conversion step is unnecessary.
Estimating Memory Before You Train
This feature is not widely known but saves a lot of money. DeepSpeed can estimate memory requirements before you run anything.
from transformers import AutoModel
from deepspeed.runtime.zero.stage3 import estimatezero3modelstatesmemneedsalllive
model = AutoModel.frompretrained("meta-llama/Llama-3.2-1B")
estimatezero3modelstatesmemneedsalllive(model, numgpuspernode=4, numnodes=1)
The output shows estimated per-GPU memory for several offload combinations. Run this before renting expensive GPUs by the hour.
Tips and Best Practices
Start at stage 2 without offload. Escalate only when you genuinely run out of memory, because each level costs throughput.
Use bf16 if your GPU supports it. NaN problems in large model training almost always come from fp16 with unstable loss scaling.
Enable gradient checkpointing when activations are the bottleneck rather than optimizer state. In HuggingFace it is just gradientcheckpointing=True in TrainingArguments. It trades roughly 30 percent throughput for a large activation memory saving.
Remember that CPU offload requires a lot of host RAM and a decent CPU. If host memory is tight, offload can get your process killed by the OOM killer instead of helping.
Do not mix the deepspeed launcher with torchrun. Pick one. If you use Accelerate, just use accelerate launch.
Rerun dsreport after upgrading PyTorch or CUDA. DeepSpeed extensions are compiled against specific versions and can quietly break.
For debugging, shrink the model and dataset until a single training step succeeds, then scale up. DeepSpeed errors are long and intimidating, yet the cause is often something simple like batch size values that disagree between the JSON and your code.
Conclusion
DeepSpeed opens the door to training large models without owning a giant cluster. The key takeaways:
Training memory is dominated by optimizer state, not model weights, and that is the first thing ZeRO shards.
ZeRO stage 1 shards optimizer state, stage 2 adds gradients, stage 3 adds model parameters. Escalate only as needed.
All behavior lives in one JSON file, and the "auto" value avoids conflicts with the HuggingFace Trainer.
There are three usage paths: through the Trainer (easiest), through Accelerate (most flexible), or the native API (most control).
Stage 3 checkpoints are sharded and need zeroto_fp32 conversion before use elsewhere.
Use the built-in memory estimator before renting GPUs, and combine it with gradient checkpointing when activations are the problem.
Try fine-tuning a 1 billion parameter model with ZeRO stage 2 on two GPUs. The moment a model that used to go out of memory trains smoothly, you will understand why this library is a backbone of so many open source LLM projects. Happy training.