# DeepSpeed User Guide

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

```bash
pip install deepspeed
```

```text
Successfully installed deepspeed-x.y.z
```

### Verify

```bash
ds_report
```

```text
--------------------------------------------------
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`:

```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`):

```python
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:

```bash
deepspeed train.py --deepspeed --deepspeed_config ds_config.json
```

```text
[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

- Hit a memory wall? Move from `"stage": 1` to `"stage": 2`, then `"stage": 3` — see [Core Usage → ZeRO](#4-core-usage).
- Still out of memory? Add `offload_optimizer` / `offload_param` — see [ZeRO-Offload and ZeRO-Infinity](#43-zero-offload-and-zero-infinity-cpunvme).
- Going multi-node? See [Common Tasks → Multi-node training](#5-common-tasks).

---

## 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"
    ```bash
    pip install deepspeed
    ```

=== "Source"
    ```bash
    git clone https://github.com/deepspeedai/DeepSpeed
    cd DeepSpeed
    pip install .
    ```

=== "install.sh (multi-node)"
    ```bash
    # 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"
    ```bash
    # 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:

```bash
DS_BUILD_OPS=1 pip install deepspeed
```

To pre-compile a **single** op, set its build variable:

```bash
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)"
    ```bash
    pip install deepspeed
    ```
    On Windows the wheel ships with all operators pre-built, so no CUDA SDK or C++ compiler is required.

=== "Source (Windows)"
    ```bash
    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

```bash
ds_report
```

Then a 10-second smoke test:

```python
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](https://github.com/deepspeedai/DeepSpeed/issues/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. |

```json
{
  "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

```json
{
  "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

```json
{
  "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"
    ```json
    {
      "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"
    ```json
    {
      "bf16": { "enabled": true }
    }
    ```
    BF16 needs no loss scaling. Requires hardware/PyTorch BF16 support.

=== "Torch autocast"
    ```json
    {
      "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

```json
{
  "gradient_clipping": 1.0,
  "data_types": { "grad_accum_dtype": "fp32" }
}
```

### A complete starter config (ZeRO-2 + BF16)

```json
{
  "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. |

```python
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. |

```json
{ "zero_optimization": { "stage": 3 } }
```

Key Stage-2 tuning keys:

```json
{
  "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:

```json
{
  "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](#47-saving-and-loading-checkpoints)).

#### 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:

```python
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"
    ```json
    {
      "zero_optimization": {
        "stage": 2,
        "offload_optimizer": { "device": "cpu", "pin_memory": true }
      }
    }
    ```

=== "Optimizer + params → CPU (Stage 3)"
    ```json
    {
      "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param": { "device": "cpu", "pin_memory": true }
      }
    }
    ```

=== "NVMe (ZeRO-Infinity)"
    ```json
    {
      "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](#48-command-line-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(...)`:

```python
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:

```json
{ "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:

```json
{ "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

```python
# 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:

```bash
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:

```bash
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
```

```json
{ "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

```bash
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`:
   ```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:
   ```bash
   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](#43-zero-offload-and-zero-infinity-cpunvme)).
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:
   ```text
   worker-0 slots=8
   worker-1 slots=8
   ```
2. Ensure passwordless SSH between nodes.
3. Launch from `worker-0`:
   ```bash
   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

```bash
# 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

```bash
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

```bash
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

```python
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:

```python
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:

```json
{
  "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)

```python
import deepspeed
deepspeed.checkpointing.configure(mpu_=None, partition_activations=True)
# wrap forward of expensive blocks with deepspeed.checkpointing.checkpoint(...)
```

Config block:

```json
{ "activation_checkpointing": { "partition_activations": true, "contiguous_memory_optimization": true, "cpu_checkpointing": false } }
```

### Task: Autotune the config

```bash
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](https://github.com/deepspeedai/DeepSpeed/issues/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](https://github.com/deepspeedai/DeepSpeed/issues/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](https://github.com/deepspeedai/DeepSpeed/issues/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.

---

_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. Unofficial; not affiliated with or reviewed by the project._