DeepSpeed is a deep learning optimization library that makes distributed training and inference of large models easy, efficient, and effective. You wrap your existing PyTorch model and optimizer with a single deepspeed.initialize(...) call and a JSON configuration file, then launch with the deepspeed command — and DeepSpeed handles memory partitioning (ZeRO), CPU/NVMe offloading, mixed precision, gradient accumulation, learning-rate scheduling, checkpointing, and multi-GPU/multi-node orchestration for you.
This guide is for users — people training or fine-tuning models, or running inference, who want to get DeepSpeed working and use its features correctly. It covers installation, the configuration file, the engine API, every command-line tool, and the day-to-day workflows. For architecture, the complete configuration reference, deployment, security, and internals, see the companion DeepSpeed Administrator & Developer Guide.
1. Quick Start
DeepSpeed slots into an existing PyTorch training loop. You change three things:
- Replace your optimizer/scheduler setup with
deepspeed.initialize(...). - Call
model_engine.backward(loss)andmodel_engine.step()instead of the usualloss.backward()/optimizer.step()/optimizer.zero_grad(). - Launch the script with the
deepspeedcommand instead ofpython.
Key features
- ZeRO (Zero Redundancy Optimizer) — partitions optimizer states (Stage 1), gradients (Stage 2), and parameters (Stage 3) across data-parallel GPUs to train models far larger than a single GPU can hold.
- ZeRO-Offload / ZeRO-Infinity — offload optimizer state and parameters to CPU RAM or NVMe disk, enabling multi-billion-parameter training on modest hardware.
- Mixed precision — FP16, BF16, and PyTorch autocast with automatic loss scaling.
- A rich optimizer & scheduler set — fused Adam/Lamb/Lion, CPU Adam/Adagrad/Lion, 1-bit Adam/Lamb, ZeRO-One Adam, Muon, plus
WarmupLR,WarmupDecayLR,WarmupCosineLR,OneCycle, andLRRangeTest. - Pipeline parallelism and tensor parallelism (AutoTP) for 3D parallelism.
- Mixture-of-Experts (MoE) layers.
- DeepSpeed-Inference —
deepspeed.init_inference(...)with kernel injection, tensor parallelism, and quantization. - A multi-node launcher supporting PDSH, OpenMPI, MPICH, Intel MPI, MVAPICH, and SLURM.
- Universal checkpointing and a checkpoint-to-FP32 converter.
Prerequisites
You need Python 3.x and PyTorch (torch>=2.0.0) with the right accelerator support (CUDA, ROCm, XPU, HPU, NPU, MLU, SDAA, or CPU). A C++ compiler is needed on Linux only if you pre-build C++/CUDA ops (most ops are compiled just-in-time at runtime).
Install
pip install deepspeed
Successfully installed deepspeed-x.y.z
Verify
ds_report
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
runtime if needed. Op compatibility means that your system
meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
...
--------------------------------------------------
Minimal example
Create a config file ds_config.json:
{
"train_batch_size": 16,
"gradient_accumulation_steps": 1,
"optimizer": {
"type": "Adam",
"params": { "lr": 0.0001 }
},
"fp16": { "enabled": true },
"zero_optimization": { "stage": 1 }
}
Adapt your training script (train.py):
import deepspeed
import torch
model = MyModel() # your torch.nn.Module
model_engine, optimizer, _, _ = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
config="ds_config.json",
)
for batch in dataloader:
inputs, labels = batch
inputs = inputs.to(model_engine.device)
labels = labels.to(model_engine.device)
outputs = model_engine(inputs) # forward
loss = loss_fn(outputs, labels)
model_engine.backward(loss) # backward (handles grad accumulation + scaling)
model_engine.step() # optimizer + LR scheduler step + zero_grad
Launch on all local GPUs:
deepspeed train.py --deepspeed --deepspeed_config ds_config.json
[INFO] DeepSpeed info: version=x.y.z, git-hash=..., git-branch=...
[INFO] Using /home/user/.cache/torch_extensions ...
Tip
You can pass the config as a path or a Python dict directly:
deepspeed.initialize(model=model, config=ds_config_dict). The--deepspeed_configCLI flag is only needed if your script reads it fromargs.deepspeed_config.
What to try next
- Hit a memory wall? Move from
"stage": 1to"stage": 2, then"stage": 3— see Core Usage → ZeRO. - Still out of memory? Add
offload_optimizer/offload_param— see ZeRO-Offload and ZeRO-Infinity. - Going multi-node? See Common Tasks → Multi-node training.
2. Installation
System requirements
| Requirement | Detail |
|---|---|
| Python | 3.8 – 3.12 (classifiers in setup.py) |
| PyTorch | torch>=2.0.0, installed before DeepSpeed |
| Accelerator | NVIDIA CUDA, AMD ROCm, Intel XPU, Habana HPU, Ascend NPU, Cambricon MLU, SDAA, or CPU |
| Compiler (Linux, source/pre-build only) | A C++ compiler and, for CUDA ops, a matching CUDA toolkit |
ninja |
Required for JIT op compilation |
libaio |
Required for the async_io op (NVMe offloading) |
Core Python dependencies (installed automatically): einops, hjson, msgpack, ninja, numpy, packaging>=20.0, psutil, py-cpuinfo, pydantic>=2.0.0, torch>=2.0.0, tqdm. On NVIDIA GPUs nvidia-ml-py is also installed.
Important
Install PyTorch first. If
torchis not importable when DeepSpeed is installed or built, pre-compiling C++/CUDA ops is disabled and you will see:[WARNING] Unable to import torch, pre-compiling ops will be disabled.
Install methods
pip
pip install deepspeed
Source
git clone https://github.com/deepspeedai/DeepSpeed
cd DeepSpeed
pip install .
install.sh (multi-node)
# builds a wheel and distributes it across hosts in the hostfile via pdsh/pdcp
bash install.sh
# local machine only:
bash install.sh --local_only
conda
# environment.yml in the repo pins pytorch, cuda, ninja, libaio, pydantic, etc.
conda env create -f environment.yml
JIT vs. pre-compiled ops
By default on Linux, DeepSpeed relies on just-in-time (JIT) compilation: ops are built the first time they are needed at runtime (this requires ninja). On Windows, ops are pre-built at install time.
To pre-compile all compatible ops at install time on Linux:
DS_BUILD_OPS=1 pip install deepspeed
To pre-compile a single op, set its build variable:
DS_BUILD_FUSED_ADAM=1 pip install deepspeed
| Op | Build variable | Purpose |
|---|---|---|
cpu_adam |
DS_BUILD_CPU_ADAM |
CPU Adam (used by optimizer offload) |
cpu_adagrad |
DS_BUILD_CPU_ADAGRAD |
CPU Adagrad |
cpu_lion |
DS_BUILD_CPU_LION |
CPU Lion |
fused_adam |
DS_BUILD_FUSED_ADAM |
Fused Adam kernel |
fused_lamb |
DS_BUILD_FUSED_LAMB |
Fused Lamb kernel |
fused_lion |
DS_BUILD_FUSED_LION |
Fused Lion kernel |
async_io |
DS_BUILD_AIO |
NVMe async I/O (needs libaio) |
gds |
DS_BUILD_GDS |
NVIDIA GPUDirect Storage |
transformer |
DS_BUILD_TRANSFORMER |
Training transformer kernel |
transformer_inference |
DS_BUILD_TRANSFORMER_INFERENCE |
Inference transformer kernel |
quantizer / fp_quantizer |
DS_BUILD_QUANTIZER / DS_BUILD_FP_QUANTIZER |
Quantization kernels |
sparse_attn |
DS_BUILD_SPARSE_ATTN |
Sparse attention |
evoformer_attn |
DS_BUILD_EVOFORMER_ATTN |
EvoFormer attention (needs CUTLASS) |
dc |
DS_BUILD_DEEP_COMPILE |
DeepCompile |
Tip
Use
DS_BUILD_OPS=1together withDS_BUILD_<OP>=0to pre-build everything except an incompatible op. SetDS_ENABLE_NINJA=1to force ninja for the build, orDS_NVCC_THREADSto parallelize CUDA compilation.
Windows
DeepSpeed installs and runs natively on Windows for single-GPU training, fine-tuning, and inference (version 0.14.5 and above).
pip (Windows)
pip install deepspeed
On Windows the wheel ships with all operators pre-built, so no CUDA SDK or C++ compiler is required.
Source (Windows)
build_win.bat
The wheel is produced under dist\*.whl. Building from source needs PyTorch with CUDA, the Visual C++ build tools, and the CUDA toolkit; launch the build console with Administrator privilege so the required symlink folders can be created.
Note
The prerequisites in both Windows cases are Python 3.x and PyTorch with CUDA support. Multi-GPU on Windows was on the roadmap at the time of writing — for production multi-GPU/multi-node, use Linux.
Post-install verification
ds_report
Then a 10-second smoke test:
import deepspeed, torch
print(deepspeed.__version__)
model = torch.nn.Linear(8, 8)
engine, opt, _, _ = deepspeed.initialize(
model=model, model_parameters=model.parameters(),
config={"train_batch_size": 1, "optimizer": {"type": "Adam", "params": {"lr": 1e-4}}})
print("device:", engine.device)
Install troubleshooting
| Symptom | Fix |
|---|---|
Unable to import torch, pre-compiling ops will be disabled |
Install a matching PyTorch build first, then reinstall DeepSpeed. |
ninja: command not found during JIT |
pip install ninja (JIT op compilation requires it). |
async_io shows [NO]/incompatible in ds_report |
Install libaio (apt-get install libaio-dev / yum install libaio-devel); rerun. NVMe offload needs it. |
RuntimeError: launcher 'pdsh' not installed. |
Install the launcher backend (pdsh) or pick another with --launcher. |
pip is unable to install the latest version ... remains stuck at an older version |
Clear pip cache and ensure your Python/torch versions match a published wheel; build from source if no wheel matches (see issue #8037). |
Wrong CUDA arch / nvcc mismatch |
Set TORCH_CUDA_ARCH_LIST (e.g. "8.0;9.0") before installing; DS_SKIP_CUDA_CHECK=1 skips the version check. |
/dev/shm too small warning in Docker |
Run the container with --shm-size=1gb or larger. |
3. Configuration
DeepSpeed is driven by a single JSON (or HJSON) configuration object — either a path you pass to --deepspeed_config / deepspeed.initialize(config=...), or a Python dict. This section covers the settings you will use most often. For the complete reference of every key, default, and validation rule, see the Admin Guide → Configuration Reference.
Important
Unknown keys are rejected. DeepSpeed config sections are strictly validated (
extra="forbid"), so a typo like"zero_optimisation"will raise a validation error rather than being silently ignored. Duplicate keys in the JSON also raiseDuplicate keys in DeepSpeed config.
Batch size: the three linked numbers
DeepSpeed ties three values together by the identity train_batch_size = train_micro_batch_size_per_gpu × gradient_accumulation_steps × world_size (world_size = number of data-parallel GPUs). You must provide enough of them for DeepSpeed to derive the rest:
| Key | Meaning |
|---|---|
train_batch_size |
Effective global batch per optimizer step. |
train_micro_batch_size_per_gpu |
Samples processed per GPU per forward/backward. |
gradient_accumulation_steps |
Micro-steps accumulated before an optimizer step. |
{
"train_micro_batch_size_per_gpu": 4,
"gradient_accumulation_steps": 8
}
With 8 GPUs this yields train_batch_size = 4 × 8 × 8 = 256 automatically.
Warning
If you provide all three and they don't satisfy the identity, initialization fails with
Check batch related parameters. train_batch_size is not equal to micro_batch_per_gpu * gradient_acc_step * world_size. Provide any two and let DeepSpeed compute the third.
Optimizer
{
"optimizer": {
"type": "AdamW",
"params": {
"lr": 0.0003,
"betas": [0.9, 0.999],
"eps": 1e-8,
"weight_decay": 0.01
}
}
}
Built-in optimizer type values: Adam, AdamW, Adagrad, Lamb, Lion, OneBitAdam, ZeroOneAdam, OneBitLamb, MuAdam, MuAdamW, MuSGD, Muon. You can also pass a constructed optimizer object to deepspeed.initialize(optimizer=...), which overrides the config.
Note
For
Adam/AdamW, DeepSpeed uses its fused/CPU implementations by default. Set"torch_adam": trueinparamsto force the stock PyTorch optimizer, or"adam_w_mode": falseto use Adam (L2) instead of AdamW (decoupled weight decay) semantics.
Learning-rate scheduler
{
"scheduler": {
"type": "WarmupDecayLR",
"params": {
"warmup_min_lr": 0,
"warmup_max_lr": 0.0003,
"warmup_num_steps": 1000,
"total_num_steps": 100000
}
}
}
Scheduler type |
Use for |
|---|---|
WarmupLR |
Linear/log warmup to a fixed max LR, then hold. |
WarmupDecayLR |
Warmup then linear decay to zero over total_num_steps. |
WarmupCosineLR |
Warmup then cosine decay (warmup_min_ratio, cos_min_ratio). |
OneCycle |
1cycle LR/momentum policy (cycle_min_lr, cycle_max_lr, ...). |
LRRangeTest |
LR range test to find a good max LR. |
Mixed precision
FP16
{
"fp16": {
"enabled": true,
"loss_scale": 0,
"initial_scale_power": 16,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
}
}
loss_scale: 0 means dynamic loss scaling (recommended).
BF16
{
"bf16": { "enabled": true }
}
BF16 needs no loss scaling. Requires hardware/PyTorch BF16 support.
Torch autocast
{
"torch_autocast": {
"enabled": true,
"dtype": "bfloat16",
"lower_precision_safe_modules": [
"torch.nn.modules.linear.Linear",
"torch.nn.modules.conv.Conv2d"
]
}
}
Caution
fp16andbf16cannot both be enabled — initialization assertsbfloat16 and fp16 modes cannot be simultaneously enabled.
Gradient clipping and accumulation dtype
{
"gradient_clipping": 1.0,
"data_types": { "grad_accum_dtype": "fp32" }
}
A complete starter config (ZeRO-2 + BF16)
{
"train_micro_batch_size_per_gpu": 4,
"gradient_accumulation_steps": 4,
"gradient_clipping": 1.0,
"optimizer": { "type": "AdamW", "params": { "lr": 0.0002, "weight_decay": 0.01 } },
"scheduler": { "type": "WarmupDecayLR",
"params": { "warmup_min_lr": 0, "warmup_max_lr": 0.0002,
"warmup_num_steps": 500, "total_num_steps": 50000 } },
"bf16": { "enabled": true },
"zero_optimization": {
"stage": 2,
"contiguous_gradients": true,
"overlap_comm": true,
"reduce_scatter": true,
"reduce_bucket_size": 5e8,
"allgather_bucket_size": 5e8
},
"steps_per_print": 100,
"wall_clock_breakdown": false
}
Tip
See the full set of config keys (eigenvalue, progressive layer drop, sparse attention, monitoring, autotuning, NVMe AIO, compile, etc.) in the Admin Guide. This guide focuses on the ones you reach for daily.
4. Core Usage
This section covers what you actually do with DeepSpeed, grouped by goal.
4.1 The training loop and engine API
After deepspeed.initialize(...) you get back an engine that replaces both your model and your optimizer step. The essential methods:
| Method | Replaces | Notes |
|---|---|---|
engine(*inputs) / engine.forward(...) |
model(...) |
Forward pass. |
engine.backward(loss) |
loss.backward() |
Scales loss for grad accumulation, applies loss scaling for FP16. |
engine.step() |
optimizer.step() + scheduler.step() + optimizer.zero_grad() |
Runs only on gradient-accumulation boundaries. |
engine.device |
— | The accelerator device to move tensors to. |
engine.local_rank / engine.global_rank |
— | Rank info. |
engine.train() / engine.eval() |
model.train() / model.eval() |
Mode switch. |
engine.is_gradient_accumulation_boundary() |
— | True on steps where the optimizer actually updates. |
model_engine, optimizer, training_dataloader, lr_scheduler = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
training_data=train_dataset, # optional: DeepSpeed builds a distributed dataloader
config="ds_config.json",
)
for step, batch in enumerate(training_dataloader):
batch = {k: v.to(model_engine.device) for k, v in batch.items()}
loss = model_engine(**batch).loss
model_engine.backward(loss)
model_engine.step()
Important
Do not call
loss.backward(),optimizer.step(), oroptimizer.zero_grad()yourself when using the DeepSpeed engine.engine.step()already performs the optimizer update, LR scheduler step, and gradient reset, and only fires on accumulation boundaries.
Note
If you pass
training_data, DeepSpeed returns a ready-to-use distributedtraining_dataloader(sharded across data-parallel ranks). Otherwise build your ownDistributedSampler-based loader.
4.2 ZeRO — fitting big models on limited memory
ZeRO eliminates memory redundancy across data-parallel GPUs. Choose a stage with zero_optimization.stage:
| Stage | Partitions | Typical use |
|---|---|---|
0 |
Nothing (ZeRO disabled, plain data parallel) | Small models that fit comfortably. |
1 |
Optimizer states | First step up when optimizer memory dominates (e.g. Adam). |
2 |
Optimizer states + gradients | Most common training choice; big savings, low overhead. |
3 |
Optimizer states + gradients + parameters | Largest models; parameters live sharded and are gathered on demand. |
{ "zero_optimization": { "stage": 3 } }
Key Stage-2 tuning keys:
{
"zero_optimization": {
"stage": 2,
"contiguous_gradients": true,
"overlap_comm": true,
"reduce_scatter": true,
"reduce_bucket_size": 5e8,
"allgather_bucket_size": 5e8,
"round_robin_gradients": false
}
}
Key Stage-3 tuning keys:
{
"zero_optimization": {
"stage": 3,
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_prefetch_bucket_size": 5e7,
"stage3_param_persistence_threshold": 1e5,
"sub_group_size": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
}
}
| Key | Effect |
|---|---|
stage3_max_live_parameters |
Max parameter elements resident per GPU before releasing. Lower = less memory, more communication. |
stage3_max_reuse_distance |
Keep a parameter if it will be reused within this many parameters. |
stage3_prefetch_bucket_size |
How many parameter elements to fetch ahead of use. |
stage3_param_persistence_threshold |
Parameters smaller than this stay unpartitioned (saves latency-bound comm). |
sub_group_size |
Tile size for processing massive models. |
stage3_gather_16bit_weights_on_model_save |
Gather sharded weights into a 16-bit checkpoint on save. |
Warning
Under Stage 3, parameters are sharded and only materialized when needed. They are not in
state_dictby default. To save a consolidated 16-bit model, enablestage3_gather_16bit_weights_on_model_saveand callengine.save_16bit_model(...), or convert a ZeRO checkpoint later (see 4.7).
Building huge Stage-3 models with deepspeed.zero.Init
For models too large to even construct on one GPU, build them inside the ZeRO-3 init context so parameters are partitioned as they are created:
import deepspeed
with deepspeed.zero.Init(config_dict_or_path="ds_config.json"):
model = MyLargeModel(hidden_size=8192, num_layers=96)
model_engine, *_ = deepspeed.initialize(model=model, config="ds_config.json")
Caution
Inside a
GatheredParameterscontext, modifying ZeRO-3 parameters withmodifier_rank=NoneraisesDetected in-place modification of ZeRO-3 parameters ... Use modifier_rank=<rank> to broadcast updates across ranks.Always pass an explicitmodifier_rankwhen you intend to update gathered parameters.
4.3 ZeRO-Offload and ZeRO-Infinity (CPU/NVMe)
When even Stage 3 won't fit, offload optimizer state and/or parameters off the GPU.
Optimizer → CPU
{
"zero_optimization": {
"stage": 2,
"offload_optimizer": { "device": "cpu", "pin_memory": true }
}
}
Optimizer + params → CPU (Stage 3)
{
"zero_optimization": {
"stage": 3,
"offload_optimizer": { "device": "cpu", "pin_memory": true },
"offload_param": { "device": "cpu", "pin_memory": true }
}
}
NVMe (ZeRO-Infinity)
{
"zero_optimization": {
"stage": 3,
"offload_optimizer": { "device": "nvme", "nvme_path": "/local_nvme", "pin_memory": true },
"offload_param": { "device": "nvme", "nvme_path": "/local_nvme", "buffer_count": 5, "buffer_size": 1e8 }
}
}
| Offload key | Applies to | Meaning |
|---|---|---|
device |
both | cpu, nvme, or none. |
nvme_path |
both | Filesystem path on the NVMe device. |
pin_memory |
both | Use page-locked CPU memory (faster transfers, more host RAM). |
buffer_count / buffer_size |
both | NVMe buffer pool size. Optimizer needs ≥ states-per-param (Adam = 4). |
ratio |
optimizer | Fraction of optimizer states offloaded to CPU Adam (Stage 3 only). |
max_in_cpu |
param | Parameter elements kept in CPU when offloading to NVMe. |
Important
Optimizer computation is performed on CPU whenever optimizer offload is enabled (regardless of
device), so install/buildcpu_adamfor good performance. NVMe offload requires theasync_ioop andlibaio. Partial optimizer offload (ratio < 1.0) is only supported with Stage 3.
Tip
Tune your NVMe settings with
ds_nvme_tunebefore committing to an offload config (see 6.x CLI tools).
4.4 Optimizers and schedulers in practice
- Prefer DeepSpeed's fused optimizers (
Adam/AdamW→ fused kernel) on GPU; they are faster than stock PyTorch. - When offloading optimizer state to CPU, DeepSpeed uses CPU Adam automatically.
Muonis supported with ZeRO-3 (with constraints — see Troubleshooting); set per-parameteruse_muonflags are applied automatically by DeepSpeed based on parameter shape and name.- 1-bit optimizers (
OneBitAdam,OneBitLamb,ZeroOneAdam) reduce communication volume; they require the matching extras and an MPI/compressed backend.
Note
If you supply a custom optimizer that DeepSpeed hasn't validated with ZeRO, set
"zero_allow_untested_optimizer": trueto bypass the safety check.
4.5 Pipeline parallelism
Split a model's layers across GPUs as pipeline stages using PipelineModule and train with engine.train_batch(...):
import torch, deepspeed
from deepspeed.pipe import PipelineModule, LayerSpec
layers = [LayerSpec(EmbeddingLayer, vocab, hidden)]
layers += [LayerSpec(TransformerLayer, hidden) for _ in range(num_layers)]
layers += [LayerSpec(OutputLayer, hidden, vocab)]
net = PipelineModule(layers=layers, num_stages=4,
loss_fn=torch.nn.CrossEntropyLoss(),
partition_method="parameters")
engine, *_ = deepspeed.initialize(model=net, config="ds_config.json")
for _ in range(num_steps):
loss = engine.train_batch(data_iter=data_iter)
pipeline config block:
{ "pipeline": { "pipe_partitioned": true, "grad_partitioned": true,
"activation_checkpoint_interval": 0 } }
Note
With pipeline parallelism, pass
mpu=Nonetodeepspeed.initialize(thePipelineModuleprovides its own model-parallel unit).partition_methodacceptsparameters,uniform, ortype:<regex>.
4.6 Tensor parallelism (AutoTP)
Enable automatic tensor parallelism through the tensor_parallel config block:
{ "tensor_parallel": { "autotp_size": 2 } }
AutoTP sharding is applied during deepspeed.initialize(...). The older deepspeed.tp_model_init(...) API is record-only and kept for backward compatibility.
Warning
Don't combine conflicting tensor-parallel sources. Passing both an
mpuand a recordedtp_group, or mismatchedautotp_size/dtypebetweentp_model_initand the config, raisesConflicting tp_model_init arguments/Conflicting tensor_parallel.autotp_size in config ... and tp_model_init.
4.7 Saving and loading checkpoints
# Save a full DeepSpeed checkpoint (sharded across ranks)
engine.save_checkpoint("checkpoints/", tag="step_1000",
client_state={"step": 1000})
# Resume
_, client_state = engine.load_checkpoint("checkpoints/", tag="step_1000")
Methods:
| Method | Saves / loads |
|---|---|
save_checkpoint(save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False) |
Full sharded checkpoint (model + optimizer + scheduler + your client_state). |
load_checkpoint(load_dir, tag=None, ...) |
Restores engine state; returns (load_path, client_state). |
save_16bit_model(save_dir, save_filename="pytorch_model.bin", exclude_frozen_parameters=False) |
A single 16-bit consolidated weight file. |
save_fp16_model(save_dir, save_filename="pytorch_model.bin") |
FP16 consolidated weights. |
Tip
Set
exclude_frozen_parameters=Trueto skip frozen (e.g. non-LoRA) weights and shrink checkpoints.
Converting a ZeRO checkpoint to a plain FP32 model
After training with ZeRO, convert the sharded checkpoint into a single FP32 state_dict for downstream use:
python -m deepspeed.utils.zero_to_fp32 checkpoints/ pytorch_model.fp32.bin
| Argument | Default | Purpose |
|---|---|---|
checkpoint_dir |
— | Path to the checkpoint folder (e.g. path/checkpoint-12). |
output_dir |
— | Where to write the FP32 output. |
--max_shard_size |
5GB |
Shard the output above this size. |
--safe_serialization |
off | Save with safetensors instead of pickle. |
-t / --tag |
latest | Specific checkpoint tag (e.g. global_step1). |
--exclude_frozen_parameters |
off | Skip frozen params. |
-d / --debug |
off | Verbose debugging. |
Universal checkpointing
To resume a job on a different GPU count / parallelism layout, convert to a universal checkpoint and enable loading it:
python -m deepspeed.checkpoint.ds_to_universal \
--input_folder checkpoints/global_step1000 \
--output_folder checkpoints/global_step1000_universal \
--num_extract_workers 4 --num_merge_workers 2
{ "checkpoint": { "load_universal": true } }
Warning
The legacy
elastic_checkpointZeRO option is no longer supported, and ZeRO-3 raisesZeRO-3 does not yet support elastic checkpointing, please disable for now.Use universal checkpointing instead for cross-topology resume.
4.8 Command-line tools
DeepSpeed installs several CLI entry points (Linux):
| Command | Purpose |
|---|---|
deepspeed |
The main launcher (alias: ds, deepspeed.pt). Runs your script across GPUs/nodes. |
ds_report |
Report op install/compatibility status and environment info. |
ds_ssh |
Run a command across all hosts in the hostfile over SSH. |
ds_bench |
Communication/throughput benchmarks. |
ds_elastic |
Helpers for elastic training. |
ds_nvme_tune |
Sweep and tune NVMe offload I/O parameters. |
ds_io |
NVMe async-I/O micro-benchmark. |
dsr |
Runner helper. |
On Windows the scripts are deepspeed.bat, ds, ds_report.bat, ds_report.
The deepspeed launcher
deepspeed [launcher options] your_script.py [your script args]
Most-used options:
| Option | Default | Purpose |
|---|---|---|
-H, --hostfile |
/job/hostfile |
MPI-style hostfile (worker-0 slots=8). |
-i, --include |
— | Resources to use: NODE[:SLOT,SLOT][@NODE...], e.g. -i "worker-0@worker-1:0,2". |
-e, --exclude |
— | Resources to skip (mutually exclusive with --include). |
--num_nodes |
-1 (all) | Use the top N hosts from the hostfile. |
--num_gpus / --num_accelerators |
-1 (all) | Max GPUs per node. |
--master_addr / --master_port |
inferred / 29500 |
Rendezvous address/port. |
--launcher |
pdsh |
Multi-node backend: pdsh, openmpi, mpich, impi, mvapich, slurm. |
--launcher_args |
— | Extra args passed to the backend (single quoted string). |
--no_ssh |
off | Launch independently per node without SSH (set --node_rank). |
--bind_cores_to_rank |
off | Pin each rank to distinct CPU cores (helps CPU backend). |
--autotuning |
— | tune or run — run the autotuner first. |
--module / --no_python |
off | Treat the script as python -m / run it directly. |
--save_pid |
off | Write launcher PID to /tmp/<pid>.ds. |
-q, --quiet |
off | Quiet logging (aliases to --log_level error). |
Tip
On a single node you don't need a hostfile —
deepspeed train.py ...uses all visible GPUs. Limit GPUs with--num_gpus 2or by settingCUDA_VISIBLE_DEVICES(DeepSpeed respects it and maps it to--include localhost:<ids>).
5. Common Tasks
Task: Fine-tune a Hugging Face model with ZeRO-2
Goal: fine-tune a Transformers model that nearly fits, with memory headroom.
- Config
ds_config.json:{ "train_micro_batch_size_per_gpu": 8, "gradient_accumulation_steps": 2, "bf16": { "enabled": true }, "zero_optimization": { "stage": 2, "overlap_comm": true, "contiguous_gradients": true }, "optimizer": { "type": "AdamW", "params": { "lr": 0.00002 } } } - Launch:
deepspeed train.py --deepspeed --deepspeed_config ds_config.json - Verify: loss prints every
steps_per_printand GPU memory is well below capacity. - If it fails with OOM: raise
gradient_accumulation_steps, lower micro batch, or move to Stage 3.
Task: Train a model that doesn't fit on one GPU (ZeRO-3 + CPU offload)
- Use Stage 3 with optimizer + param CPU offload (see 4.3).
- Build the model inside
deepspeed.zero.Init(...)if construction itself OOMs. - Verify: training starts; host RAM usage rises (offload working).
- If it fails: reduce
stage3_max_live_parametersandstage3_prefetch_bucket_size; consider NVMe offload.
Task: Multi-node training over SSH (PDSH)
- Create a hostfile:
worker-0 slots=8 worker-1 slots=8 - Ensure passwordless SSH between nodes.
- Launch from
worker-0:deepspeed --hostfile=hostfile train.py --deepspeed --deepspeed_config ds_config.json - Verify:
WORLD INFO DICTlists both nodes;dist_world_size=16. - If
host ... was not reachable via ssh: fix passwordless SSH, or use--no_ssh_check, or pass--ssh_port.
Tip
Put per-node environment variables (e.g.
NCCL_*) in a.deepspeed_envfile in your home or current directory; the launcher exports them to all workers. Change the filename withDS_ENV_FILE.
Task: Multi-node training without SSH
# on each node, with its own --node_rank
deepspeed --no_ssh --node_rank=0 --master_addr=10.0.0.1 --master_port=29500 \
--num_nodes=2 train.py --deepspeed --deepspeed_config ds_config.json
Task: Launch with SLURM / MPI
deepspeed --launcher slurm train.py --deepspeed_config ds_config.json
# or: --launcher openmpi / mpich / impi / mvapich
Warning
The OpenMPI, MPICH, IMPI, and MVAPICH backends do not support
--include/--excludeor--num_nodes/--num_gpusfiltering and raise an error if you pass them. Use a hostfile or the launcher's own resource controls.
Task: Select specific GPUs
deepspeed --include localhost:0,1 train.py --deepspeed_config ds_config.json
# or
CUDA_VISIBLE_DEVICES=0,1 deepspeed train.py --deepspeed_config ds_config.json
Task: Run inference with DeepSpeed-Inference
import deepspeed, torch
model = deepspeed.init_inference(
model,
tensor_parallel={"tp_size": 2},
dtype=torch.half,
replace_with_kernel_inject=True,
)
output = model(input_ids)
Or via a config dict / JSON:
model = deepspeed.init_inference(model, config={"dtype": torch.half}, replace_with_kernel_inject=True)
| Inference option | Meaning |
|---|---|
tensor_parallel={"tp_size": N} |
Shard the model across N GPUs. |
dtype |
torch.half, torch.bfloat16, torch.int8, or torch.float. |
replace_with_kernel_inject (alias kernel_inject) |
Use DeepSpeed's optimized transformer kernels. |
max_out_tokens (alias max_tokens) |
Max output tokens (default 1024). |
enable_cuda_graph |
Capture a CUDA graph for repeated shapes. |
injection_policy |
Manual layer-replacement policy for custom models. |
Note
init_inferenceaccepts config-and-kwargs together; kwargs take precedence, but a key set to two different values in both raisesConflicting argument '<key>'.
Task: ZeRO-Inference (offload weights for inference)
For inference on memory-limited hardware, offload model weights to CPU/NVMe with Stage 3 offload and generate normally — useful for running large models on small GPUs.
Task: Profile performance
Enable the FLOPS profiler in your config:
{
"flops_profiler": {
"enabled": true,
"profile_step": 1,
"module_depth": -1,
"top_modules": 1,
"detailed": true
}
}
Set "wall_clock_breakdown": true to log per-stage timing.
Task: Gradient checkpointing (activation checkpointing)
import deepspeed
deepspeed.checkpointing.configure(mpu_=None, partition_activations=True)
# wrap forward of expensive blocks with deepspeed.checkpointing.checkpoint(...)
Config block:
{ "activation_checkpointing": { "partition_activations": true, "contiguous_memory_optimization": true, "cpu_checkpointing": false } }
Task: Autotune the config
deepspeed --autotuning tune train.py --deepspeed_config ds_config.json
The autotuner explores ZeRO/batch settings and writes an optimal config; use --autotuning run to tune then run.
6. Troubleshooting
A quick decision tree:
- Crash at startup → usually a config or batch-size mismatch (read the assertion message).
- CUDA OOM → increase ZeRO stage / add offload / lower micro-batch.
- NaN/Inf loss → loss-scaling or precision issue; check FP16 settings.
- Hangs → communication/launcher/SSH or dynamic-shape issue.
Error: train_batch_size is not equal to micro_batch_per_gpu * gradient_acc_step * world_size
Cause: You supplied all three batch parameters and they violate the identity, or world_size differs from what you assumed.
Solution:
- Provide only two of the three and let DeepSpeed compute the third.
- Remember
world_size= number of data-parallel GPUs (affected by TP/PP/SP).
Prevention: Set train_micro_batch_size_per_gpu + gradient_accumulation_steps and omit train_batch_size.
Error: bfloat16 and fp16 modes cannot be simultaneously enabled
Cause: Both "fp16": {"enabled": true} and "bf16": {"enabled": true} are set.
Solution: Enable exactly one precision mode.
Prevention: Pick BF16 on Ampere/Hopper-class hardware; FP16 otherwise.
Error: Unable to proceed, no GPU resources available
Cause: The launcher found no accelerators (no hostfile and device_count()==0).
Solution: Verify drivers (nvidia-smi), CUDA_VISIBLE_DEVICES, and that PyTorch sees the GPU (torch.cuda.is_available()).
Prevention: Run ds_report first.
Error: host=<name> was not reachable via ssh
Cause: Multi-node launch but passwordless SSH to a worker failed.
Solution:
- Set up SSH keys so
ssh worker-1 hostnameworks without a password. - Or pass
--ssh_port <port>, or--no_ssh_checkto skip the probe, or--no_sshto launch per-node manually.
Prevention: Test SSH to every host in the hostfile before launching.
Error: launcher '<name>' not installed.
Cause: The chosen multi-node backend binary isn't on PATH.
Solution: Install the backend (pdsh, mpirun, srun, ...) or pick another with --launcher.
Error: Hostfile contains a bad entry / Hostfile is empty or not formatted correctly
Cause: Hostfile lines must be hostname slots=N.
Solution: Fix the format; one host per line, e.g. worker-0 slots=8. Comments (#) and blank lines are allowed.
Error: Environment variable KMP_AFFINITY conflicts with numactl
Cause: --bind_cores_to_rank with KMP_AFFINITY set interferes with core binding.
Solution: unset KMP_AFFINITY before running deepspeed.
Error: DeepSpeed requires --deepspeed_config to specify configuration file
Cause: Neither a config path nor a config dict reached deepspeed.initialize.
Solution: Pass config="ds_config.json" (or a dict) to initialize, or --deepspeed_config if your code reads args.deepspeed_config.
Error: Expected ... an existing deepspeed config, or a dictionary or a valid base64
Cause: The config argument isn't a real path, dict, or base64 string (e.g. wrong working directory).
Solution: Use an absolute path or pass the dict directly.
Error: Detected in-place modification of ZeRO-3 parameters inside GatheredParameters with modifier_rank=None
Cause: You modified sharded parameters while gathering them without telling DeepSpeed which rank's values to broadcast.
Solution: Use with deepspeed.zero.GatheredParameters(params, modifier_rank=0): ....
Error: ZeRO-3 does not yet support elastic checkpointing, please disable for now.
Cause: elastic_checkpoint enabled under Stage 3.
Solution: Disable it; use universal checkpointing for cross-topology resume.
Error: Cannot use fp16 without accelerator. / Accelerator is not detected, cannot perform low precision training
Cause: FP16/BF16 requested on a system with no usable accelerator.
Solution: Run on a GPU/accelerator, or disable low-precision modes for CPU runs.
Error: Muon and reduce scatter cannot be used together / Muon and all2all process group cannot be used together
Cause: The Muon optimizer has restrictions under ZeRO-3.
Solution: Disable reduce_scatter / avoid the all2all process group when using Muon, or use a different optimizer.
Error: narrow() cannot be applied to a 0-dim tensor (ZenFlow Stage 3)
Cause: Known edge case with ZenFlow + Stage 3 (issue #7686).
Solution: Update DeepSpeed; if it persists, file/track the issue and adjust ZenFlow settings.
Issue: DeepCompile training hang on random-sized inputs
Cause: Dynamic input shapes can hang DeepCompile (issue #7611).
Solution: Pad to fixed shapes/buckets, or disable DeepCompile for variable-length data.
Issue: ZeRO++ IndexError sharding a small parameter
Cause: Very small parameters under ZeRO++ hierarchical partitioning (issue #6659).
Solution: Raise stage3_param_persistence_threshold so tiny params stay unpartitioned; update DeepSpeed.
Issue: /dev/shm size might be too small (Docker)
Cause: Default container shared-memory is too small for NCCL.
Solution: Add --shm-size=1gb (or more) to docker run.
7. FAQ
Q: Do I still call loss.backward() and optimizer.step()?
No. Use model_engine.backward(loss) then model_engine.step(). The engine handles loss scaling, gradient accumulation, the optimizer update, the LR scheduler step, and zeroing gradients.
Q: Path or dict for the config?
Either. deepspeed.initialize(config="ds_config.json") or config=ds_config_dict. A base64-encoded JSON string also works (used internally by the launcher).
Q: Which ZeRO stage should I use? Start at the lowest that fits: Stage 1 if optimizer memory dominates, Stage 2 for most training, Stage 3 for the largest models. Add CPU/NVMe offload only when Stage 3 alone won't fit.
Q: FP16 or BF16?
BF16 on hardware that supports it (no loss scaling, more stable). FP16 with dynamic loss scaling (loss_scale: 0) elsewhere. Never enable both.
Q: How do I get a normal model file out of a ZeRO checkpoint?
Run python -m deepspeed.utils.zero_to_fp32 <checkpoint_dir> <output> for an FP32 state_dict, or call engine.save_16bit_model(...) (with stage3_gather_16bit_weights_on_model_save for Stage 3).
Q: My optimizer isn't in DeepSpeed's list — can I still use ZeRO?
Yes, pass it to deepspeed.initialize(optimizer=...). If DeepSpeed warns it's untested with ZeRO, set "zero_allow_untested_optimizer": true.
Q: How many GPUs am I actually using?
On one node, all visible GPUs unless limited by --num_gpus, --include, or CUDA_VISIBLE_DEVICES. Multi-node, it's the sum of slots= across active hosts.
Q: How do I limit training to certain GPUs?
--include localhost:0,1, --num_gpus 2, or CUDA_VISIBLE_DEVICES=0,1.
Q: Does DeepSpeed work on AMD / Intel / Habana / Ascend hardware?
Yes — ROCm (AMD), XPU (Intel), HPU (Habana), NPU (Ascend), MLU (Cambricon), and SDAA accelerators are supported, in addition to NVIDIA CUDA and CPU. The accelerator is auto-detected; DS_ACCELERATOR can override it.
Q: Does it run on Windows? Yes, for single-GPU training, fine-tuning, and inference (0.14.5+). The Windows wheel ships with ops pre-built.
Q: How do I pass environment variables to all nodes?
Put them in a .deepspeed_env file (home or CWD); the launcher exports matching variables (e.g. NCCL_*, UCX_*) to every worker. Override the filename with DS_ENV_FILE.
Q: What's the difference between deepspeed and ds?
They're the same launcher — ds, deepspeed, and deepspeed.pt are installed as aliases.
Q: How do I check whether an op is available without installing it?
ds_report lists each op's installed/compatible status. "Compatible but not installed" means it will be JIT-compiled on first use (needs ninja).
Q: How do I resume on a different number of GPUs?
Convert the checkpoint with deepspeed.checkpoint.ds_to_universal and set "checkpoint": {"load_universal": true}. The legacy ZeRO elastic_checkpoint path is no longer supported.
Q: Where is the complete list of every config option? In the DeepSpeed Administrator & Developer Guide → Configuration Reference. This guide covers the day-to-day subset.
Q: How do I set the gradient accumulation dtype?
"data_types": { "grad_accum_dtype": "fp32" } (or "bf16" / "fp16").
Q: Can I overlap communication with computation?
Yes — "overlap_comm": true in zero_optimization (default for Stage 3). It overlaps gradient reduction with the backward pass.