Jesse A. Green / notebook

DeepSpeed User Guide

Generated June 2026 from deepspeedai/DeepSpeed at 28a196f

AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it.

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:

  1. Replace your optimizer/scheduler setup with deepspeed.initialize(...).
  2. Call model_engine.backward(loss) and model_engine.step() instead of the usual loss.backward() / optimizer.step() / optimizer.zero_grad().
  3. Launch the script with the deepspeed command instead of python.

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, and LRRangeTest.
  • 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_config CLI flag is only needed if your script reads it from args.deepspeed_config.

What to try next


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 torch is 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=1 together with DS_BUILD_<OP>=0 to pre-build everything except an incompatible op. Set DS_ENABLE_NINJA=1 to force ninja for the build, or DS_NVCC_THREADS to 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 raise Duplicate 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": true in params to force the stock PyTorch optimizer, or "adam_w_mode": false to 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

fp16 and bf16 cannot both be enabled — initialization asserts bfloat16 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(), or optimizer.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 distributed training_dataloader (sharded across data-parallel ranks). Otherwise build your own DistributedSampler-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_dict by default. To save a consolidated 16-bit model, enable stage3_gather_16bit_weights_on_model_save and call engine.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 GatheredParameters context, modifying ZeRO-3 parameters with modifier_rank=None raises Detected in-place modification of ZeRO-3 parameters ... Use modifier_rank=<rank> to broadcast updates across ranks. Always pass an explicit modifier_rank when 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/build cpu_adam for good performance. NVMe offload requires the async_io op and libaio. Partial optimizer offload (ratio < 1.0) is only supported with Stage 3.

Tip

Tune your NVMe settings with ds_nvme_tune before 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.
  • Muon is supported with ZeRO-3 (with constraints — see Troubleshooting); set per-parameter use_muon flags 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": true to 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=None to deepspeed.initialize (the PipelineModule provides its own model-parallel unit). partition_method accepts parameters, uniform, or type:<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 mpu and a recorded tp_group, or mismatched autotp_size/dtype between tp_model_init and the config, raises Conflicting 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=True to 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_checkpoint ZeRO option is no longer supported, and ZeRO-3 raises ZeRO-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 2 or by setting CUDA_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.

  1. 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 } }
    }
    
  2. Launch:
    deepspeed train.py --deepspeed --deepspeed_config ds_config.json
    
  3. Verify: loss prints every steps_per_print and GPU memory is well below capacity.
  4. 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)

  1. Use Stage 3 with optimizer + param CPU offload (see 4.3).
  2. Build the model inside deepspeed.zero.Init(...) if construction itself OOMs.
  3. Verify: training starts; host RAM usage rises (offload working).
  4. If it fails: reduce stage3_max_live_parameters and stage3_prefetch_bucket_size; consider NVMe offload.

Task: Multi-node training over SSH (PDSH)

  1. Create a hostfile:
    worker-0 slots=8
    worker-1 slots=8
    
  2. Ensure passwordless SSH between nodes.
  3. Launch from worker-0:
    deepspeed --hostfile=hostfile train.py --deepspeed --deepspeed_config ds_config.json
    
  4. Verify: WORLD INFO DICT lists both nodes; dist_world_size=16.
  5. 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_env file in your home or current directory; the launcher exports them to all workers. Change the filename with DS_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/--exclude or --num_nodes/--num_gpus filtering 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_inference accepts config-and-kwargs together; kwargs take precedence, but a key set to two different values in both raises Conflicting 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:

  1. Provide only two of the three and let DeepSpeed compute the third.
  2. 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:

  1. Set up SSH keys so ssh worker-1 hostname works without a password.
  2. Or pass --ssh_port <port>, or --no_ssh_check to skip the probe, or --no_ssh to 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.