# DeepSpeed Administrator & Developer Guide

This guide is the operational and engineering reference for DeepSpeed: architecture, the **complete** configuration reference, deployment across clusters, operations, performance tuning, security, and the codebase internals. It is written for ML platform engineers, infrastructure admins, and contributors. For installation and day-to-day usage, see the companion **DeepSpeed User Guide**.

DeepSpeed is a pure-Python library (with C++/CUDA extension ops) that wraps a user's PyTorch model and optimizer to provide memory-efficient, large-scale distributed training and inference. It introduces no standalone server or daemon — it is invoked in-process via `deepspeed.initialize(...)` and orchestrated across hosts by the `deepspeed` launcher.

---

## 1. Architecture Overview

DeepSpeed sits between the user's PyTorch training script and the accelerator/communication layer. The user calls `deepspeed.initialize(...)`, which parses the JSON config into a `DeepSpeedConfig`, selects an engine implementation, and returns a wrapped engine that intercepts forward/backward/step.

```
                 user training script (train.py)
                              │
                 deepspeed.initialize(model, optimizer, config)
                              │
        ┌─────────────────────────────────────────────┐
        │              DeepSpeedConfig                   │
        │  (parses JSON/HJSON/dict/base64; pydantic v2  │
        │   sub-configs: ZeRO, offload, fp16/bf16,      │
        │   optimizer, scheduler, comms, monitor, ...)  │
        └─────────────────────────────────────────────┘
                              │ selects engine
        ┌──────────────┬──────────────┬────────────────┐
        ▼              ▼              ▼                ▼
  DeepSpeedEngine  PipelineEngine  HybridEngine   InferenceEngine
  (data-parallel   (pipeline       (RLHF train+   (init_inference;
   + ZeRO 0-3)      parallel)       inference)     kernel inject, TP)
        │
        ▼
  ZeRO optimizer (stage_1_and_2 / stage3) ── offload (CPU/NVMe via swap_tensor + async_io)
        │
        ▼
  deepspeed.comm  ──────────►  torch.distributed (NCCL / Gloo / MPI / CCL / ...)
        │
        ▼
  accelerator abstraction (get_accelerator) ──► CUDA | ROCm | XPU | HPU | NPU | MLU | SDAA | CPU
        │
        ▼
  op_builder + csrc (C++/CUDA/HIP kernels, JIT or pre-built)
```

> [!NOTE]
> The engine choice is made inside `deepspeed.initialize`: a `PipelineModule` yields a `PipelineEngine`; `hybrid_engine.enabled` yields a `DeepSpeedHybridEngine`; otherwise a `DeepSpeedEngine`. `init_inference` is a separate path producing an `InferenceEngine`. This is why the same `initialize` call transparently supports very different parallelism strategies.

### Component inventory

| Component | Path | Role |
|---|---|---|
| Public API | `deepspeed/__init__.py` | `initialize`, `init_inference`, `zero`, `add_config_arguments`, `tp_model_init`. |
| Runtime engine | `deepspeed/runtime/engine.py` | `DeepSpeedEngine`: forward/backward/step, checkpointing, optimizer/scheduler wiring. |
| Config system | `deepspeed/runtime/config.py`, `config_utils.py`, `constants.py` | `DeepSpeedConfig` + pydantic config models + key/default constants. |
| ZeRO | `deepspeed/runtime/zero/` | `stage_1_and_2.py`, `stage3.py`, `partition_parameters.py`, `config.py`, `offload_config.py`, MiCS, tiling. |
| Pipeline parallel | `deepspeed/runtime/pipe/`, `deepspeed/pipe/` | `PipelineModule`, `LayerSpec`, `PipelineEngine`. |
| Tensor parallel | `deepspeed/runtime/tensor_parallel/` | AutoTP init utilities. |
| Sequence parallel | `deepspeed/sequence/` | Ulysses sequence parallelism. |
| Communication | `deepspeed/comm/` | `init_distributed`, collective wrappers over torch.distributed, SDMA allgather. |
| Launcher | `deepspeed/launcher/` | `runner.py` (front-end), `launch.py` (per-node), multi-node runners. |
| Accelerator abstraction | `accelerator/` | `get_accelerator()` → CUDA/ROCm/XPU/HPU/NPU/MLU/SDAA/CPU/MPS backends. |
| Op builders | `op_builder/` | Build/JIT-compile C++/CUDA ops; per-accelerator subfolders. |
| Native kernels | `csrc/` | C++/CUDA/HIP sources (adam, transformer, quantization, aio, gds, ...). |
| Inference | `deepspeed/inference/`, `deepspeed/module_inject/` | `InferenceEngine`, kernel injection, layer replacement. |
| MoE | `deepspeed/moe/` | Mixture-of-Experts layers and sharded gating. |
| Checkpointing | `deepspeed/checkpoint/`, `deepspeed/runtime/model_checkpointing/` | Universal checkpoint conversion, parallel write. |
| Autotuning | `deepspeed/autotuning/` | `Autotuner` config search. |
| Elasticity | `deepspeed/elasticity/` | Elastic training agent + config. |
| Monitoring | `deepspeed/monitor/` | TensorBoard/WandB/CSV/Comet monitor backends. |
| Profiling | `deepspeed/profiling/flops_profiler/` | FLOPS/latency profiler. |
| Compression | `deepspeed/compression/` | Quantization, pruning, layer reduction. |
| Compile | `deepspeed/compile/`, `csrc/compile/` | DeepCompile graph optimization. |
| NVMe / I/O | `deepspeed/nvme/`, `deepspeed/io/`, `csrc/aio/`, `csrc/gds/` | Async NVMe I/O for offload; DeepNVMe tools. |
| Utilities | `deepspeed/utils/` | Logging, groups, NUMA binding, `zero_to_fp32`, tensor fragments. |

### Technology stack

| Layer | Technology |
|---|---|
| Language | Python 3.8–3.12; C++/CUDA/HIP for kernels |
| DL framework | PyTorch ≥ 2.0 |
| Config validation | pydantic v2 (`DeepSpeedConfigModel`) |
| Config format | JSON / HJSON / dict / base64 |
| Communication | torch.distributed (NCCL, Gloo, MPI, oneCCL), custom SDMA allgather |
| Build | setuptools + ninja; per-op `DS_BUILD_*` env vars |
| Multi-node launch | PDSH, OpenMPI, MPICH, Intel MPI, MVAPICH, SLURM |
| Accelerators | CUDA, ROCm, XPU, HPU, NPU, MLU, SDAA, MPS, CPU |

### Data flow for a training step

1. **Forward** — `engine.forward()` runs the (possibly ZeRO-3-gathered) module; parameters are fetched/prefetched and released around submodule execution.
2. **Backward** — `engine.backward(loss)` scales the loss (grad accumulation + FP16 loss scaling), runs autograd, and reduces/partitions gradients (reduce-scatter under ZeRO-2/3, overlapped with compute when `overlap_comm`).
3. **Step** — on a gradient-accumulation boundary, `engine.step()` runs the (possibly CPU-offloaded) optimizer, the LR scheduler, gradient clipping, and zeroes gradients.

> [!NOTE]
> ZeRO's invariant is that no two data-parallel ranks hold redundant copies of the partitioned state (optimizer states at Stage 1, +gradients at Stage 2, +parameters at Stage 3). Parameters are transiently materialized only while a submodule executes — the basis for ZeRO-Infinity's ability to train trillion-parameter models with NVMe offload.

---

## 2. System Requirements

| Requirement | Detail |
|---|---|
| Python | 3.8 – 3.12 |
| PyTorch | ≥ 2.0.0, installed before DeepSpeed |
| Accelerator | NVIDIA CUDA, AMD ROCm, Intel XPU, Habana HPU, Ascend NPU, Cambricon MLU, SDAA, Apple MPS, or CPU |
| Build toolchain (Linux, for pre-built/JIT ops) | C++ compiler; CUDA toolkit matching torch for CUDA ops; `ninja` for JIT |
| NVMe offload | `libaio` (for `async_io`); optionally NVIDIA GDS (`gds`) |
| Shared memory | `/dev/shm` ≥ 512 MB (NCCL); in Docker use `--shm-size=1gb`+ |
| Network (multi-node) | Passwordless SSH (PDSH) or an MPI/SLURM scheduler; high-bandwidth interconnect (InfiniBand recommended) |

Auto-installed Python deps: `einops`, `hjson`, `msgpack`, `ninja`, `numpy`, `packaging>=20.0`, `psutil`, `py-cpuinfo`, `pydantic>=2.0.0`, `torch>=2.0.0`, `tqdm` (+ `nvidia-ml-py` on NVIDIA).

Optional extras (from `setup.py`): `1bit`, `1bit_mpi`, `readthedocs`, `dev`, `autotuning`, `autotuning_ml`, `sparse_attn`, `sparse`, `inf`, `sd`, `triton`, `deepcompile`, and `all`.

```bash
pip install deepspeed[autotuning,inf]
```

---

## 3. Configuration Reference

The configuration object is parsed by `DeepSpeedConfig` (`deepspeed/runtime/config.py`). It accepts a dict, a path to a JSON/HJSON file, or a base64-encoded JSON string. Sub-sections are pydantic v2 models with `extra="forbid"`, `validate_default=True`, and `validate_assignment=True` — **unknown keys raise validation errors**, and duplicate JSON keys are rejected.

> [!IMPORTANT]
> Because validation is strict, "auto" sentinel values are stripped for non-strict construction (Hugging Face compatibility), and deprecated keys emit warnings and are re-mapped to their replacements. Treat the keys below as the source-of-truth surface as of commit `28a196f`.

### 3.1 Top-level training keys

| Key | Type | Default | Description |
|---|---|---|---|
| `train_batch_size` | int | `None` | Effective global batch per step. Derivable from the other two + world_size. |
| `train_micro_batch_size_per_gpu` | int | `None` | Samples per GPU per micro-step. |
| `gradient_accumulation_steps` | int | `None` | Micro-steps per optimizer step. |
| `steps_per_print` | int | `None` | Log cadence for training stats. |
| `gradient_clipping` | float | `0.0` | Global gradient-norm clip (0 = off). |
| `prescale_gradients` | bool | `false` | Scale gradients before all-reduce. |
| `gradient_predivide_factor` | float | `1.0` | Pre-divide factor for gradients. |
| `sparse_gradients` | bool | `false` | Enable CSR sparse-gradient handling. |
| `disable_allgather` | bool | `false` | Use broadcasts instead of all-gather. |
| `communication_data_type` | str | `None` | `fp16`/`bf16`/`fp32` for gradient comm. |
| `seq_parallel_communication_data_type` | str | `fp32` | Comm dtype for sequence parallelism. |
| `dump_state` | bool | `false` | Dump engine state for debugging. |
| `wall_clock_breakdown` | bool | `false` | Per-stage timing breakdown. |
| `memory_breakdown` | bool | `false` | Memory usage breakdown. |
| `dataloader_drop_last` | bool | `false` | Drop the last incomplete batch. |
| `vocabulary_size` | int | `None` | For tensor-core alignment warnings. |
| `graph_harvesting` | bool | `false` | Capture graphs for short kernel sequences. |
| `data_types.grad_accum_dtype` | str | `None` | `bf16`/`fp16`/`fp32` accumulation dtype. |

### 3.2 Optimizer & scheduler

```json
{
  "optimizer": { "type": "AdamW", "params": { "lr": 3e-4, "weight_decay": 0.01 },
                 "legacy_fusion": false },
  "scheduler": { "type": "WarmupDecayLR",
                 "params": { "warmup_min_lr": 0, "warmup_max_lr": 3e-4,
                             "warmup_num_steps": 1000, "total_num_steps": 100000 } }
}
```

Optimizer `type`: `Adam`, `AdamW`, `Adagrad`, `Lamb`, `Lion`, `OneBitAdam`, `ZeroOneAdam`, `OneBitLamb`, `MuAdam`, `MuAdamW`, `MuSGD`, `Muon`. Extra param toggles: `torch_adam` (force stock PyTorch Adam), `adam_w_mode` (default `true`), `max_grad_norm`.

Scheduler `type`: `LRRangeTest`, `OneCycle`, `WarmupLR`, `WarmupDecayLR`, `WarmupCosineLR`. Selected scheduler params include `warmup_min_lr`, `warmup_max_lr`, `warmup_num_steps`, `warmup_type` (`log`/`linear`), `total_num_steps`, `cycle_min_lr`, `cycle_max_lr`, `cycle_first_step_size`, `decay_lr_rate`, `warmup_min_ratio`, `cos_min_ratio`.

| ZeRO-related optimizer key | Default | Meaning |
|---|---|---|
| `zero_allow_untested_optimizer` | `false` | Permit a non-validated optimizer with ZeRO. |
| `zero_force_ds_cpu_optimizer` | `true` | Force DeepSpeed CPU optimizer when offloading. |

### 3.3 Mixed precision

#### `fp16`

| Key | Default | Meaning |
|---|---|---|
| `enabled` | `false` | Enable FP16. |
| `loss_scale` | `0` | 0 = dynamic loss scaling. |
| `initial_scale_power` | `16` | Initial dynamic scale = 2^power. |
| `loss_scale_window` | `1000` | Window for scale adjustment. |
| `hysteresis` | `2` | Overflow tolerance. |
| `consecutive_hysteresis` | `false` | Require consecutive good steps. |
| `min_loss_scale` | `1` | Lower bound on scale. |
| `auto_cast` | `false` | Auto-cast inputs to FP16. |
| `fp16_master_weights_and_grads` | `false` | FP16 master copy (ZeRO 1/2/3 only). |

#### `bf16`

| Key | Default | Meaning |
|---|---|---|
| `enabled` | `false` | Enable BF16 (no loss scaling needed). |
| `immediate_grad_update` | `true` | Update grads immediately. |
| `check_overflow` | `false` | Check for overflow. |
| `bf16_master_weights_and_grads` | `false` | BF16 master weights (ZeRO 1/2/3). |
| `bf16_optimizer_states` | `false` | BF16 optimizer states (requires master weights). |

#### `torch_autocast`

| Key | Default | Meaning |
|---|---|---|
| `enabled` | `false` | Use `torch.autocast`. |
| `dtype` | — | `float16`/`bfloat16`. |
| `lower_precision_safe_modules` | — | List of module class paths kept in low precision. |

> [!CAUTION]
> `fp16.enabled` and `bf16.enabled` are mutually exclusive (asserted at init). `*_master_weights_and_grads` and `bf16_optimizer_states` are only valid with ZeRO Stages 1–3; misuse raises a config assertion.

### 3.4 `zero_optimization` — complete

| Key | Type | Default | Meaning |
|---|---|---|---|
| `stage` | 0–3 | `0` | Disabled / optimizer / +gradients / +parameters. |
| `contiguous_gradients` | bool | `true` | Copy grads to a contiguous buffer (less fragmentation). |
| `reduce_scatter` | bool | `true` | Use reduce-scatter instead of all-reduce. |
| `reduce_bucket_size` | int | `5e8` | Elements reduced at a time. |
| `use_multi_rank_bucket_allreduce` | bool | `true` | Combine rank buckets into one all-reduce. |
| `allgather_partitions` | bool | `true` | All-gather vs. broadcast to gather params. |
| `allgather_bucket_size` | int | `5e8` | Elements all-gathered at a time. |
| `overlap_comm` | bool | `None`→`stage==3` | Overlap gradient reduction with backward. |
| `load_from_fp32_weights` | bool | `true` | Init FP32 master from FP32 checkpoint copies. |
| `elastic_checkpoint` | bool | `false` | (No longer supported.) |
| `offload_param` | dict | `None` | Param offload config (Stage 3). |
| `offload_optimizer` | dict | `None` | Optimizer offload config (Stage 1/2/3). |
| `zenflow` | dict | `None` | ZenFlow config. |
| `sub_group_size` | int | `1e9` | Tile size for massive models. |
| `stage3_prefetch_bucket_size` | int | `5e7` | Param elements to prefetch (alias `prefetch_bucket_size`). |
| `stage3_param_persistence_threshold` | int | `1e5` | Don't partition params smaller than this. |
| `stage3_model_persistence_threshold` | int | `sys.maxsize` | Max unpartitioned param elements. |
| `stage3_max_live_parameters` | int | `1e9` | Max resident params per GPU before release. |
| `stage3_max_reuse_distance` | int | `1e9` | Keep params reused within this distance. |
| `stage3_gather_16bit_weights_on_model_save` | bool | `false` | Gather sharded weights into a 16-bit save. |
| `stage3_module_granularity_threshold` | int | `0` | Treat fine-grained modules as units when fetching. |
| `stage3_use_all_reduce_for_fetch_params` | bool | `false` | Use all-reduce when fetching params. |
| `stage3_allgather_sequential` | bool | `false` | Per-param sequential all-gather (lower peak memory). |
| `ignore_unused_parameters` | bool | `true` | Don't error on unused params (Stage 2). |
| `legacy_stage1` | bool | `false` | Old Stage-1 implementation (deprecated). |
| `round_robin_gradients` | bool | `false` | Parallelize grad copy to CPU across ranks. |
| `zero_hpz_partition_size` | int | `1` | Secondary (hierarchical) partition group size (ZeRO++). |
| `zero_quantized_weights` | bool | `false` | Quantize weights for all-gather comm. |
| `zero_quantized_nontrainable_weights` | bool | `false` | Store non-trainable weights quantized (LoRA). |
| `zero_quantized_gradients` | bool | `false` | Quantize gradients for comm. |
| `zeropp_loco_param` | dict | `None` | LoCo-Zero++ params (`err_beta`, `reset_T`). |
| `mics_shard_size` | int | `-1` | MiCS shard size. |
| `mics_hierarchical_params_gather` | bool | `false` | MiCS hierarchical gather. |
| `memory_efficient_linear` | bool | `true` | Memory-efficient linear (Stage 3). |
| `override_module_apply` | bool | `true` | Override `nn.Module.apply` (Stage 3). |
| `pipeline_loading_checkpoint` | bool | `false` | Force checkpoint load in pipeline mode (Stage 3). |
| `log_trace_cache_warnings` | bool | `false` | Log trace-cache invalidations. |
| `enable_sanity_checks` | bool | `false` | Internal sanity checks (debugging). |
| `save_muon_momentum_buffer_in_memory` | bool | `false` | Keep Muon momentum in memory vs. NVMe (Stage 3). |
| `leaf_module` | dict | — | Modules treated as ZeRO-3 leaf units. |

#### `offload_param` (Stage 3)

| Key | Default | Meaning |
|---|---|---|
| `device` | `none` | `cpu`/`nvme`/`none`. |
| `nvme_path` | `None` | NVMe filesystem path. |
| `buffer_count` | `5` | NVMe buffer-pool size. |
| `buffer_size` | `1e8` | NVMe buffer size. |
| `max_in_cpu` | `1e9` | Param elements kept in CPU when NVMe offloading. |
| `pin_memory` | `false` | Page-locked CPU memory. |

#### `offload_optimizer` (Stage 1/2/3)

| Key | Default | Meaning |
|---|---|---|
| `device` | `none` | `cpu`/`nvme`/`none`. Optimizer compute runs on CPU regardless. |
| `nvme_path` | `None` | NVMe path. |
| `buffer_count` | `4` | NVMe buffers (≥ optimizer states per param; Adam=4). |
| `pin_memory` | `false` | Page-locked CPU memory. |
| `pipeline_read` / `pipeline_write` | `false` | Overlap tile read/write with compute (ZeRO-Infinity). |
| `fast_init` | `false` | Fast init when offloading to NVMe. |
| `ratio` | `1.0` | Fraction of optimizer states offloaded to CPU Adam (Stage 3). |
| `super_offload` | `false` | High-performance CPU offload for Superchips (Stage 3). |
| `cpuadam_cores_perc` | `0.8` | CPU cores for CPU Adam (with `super_offload`). |

> [!WARNING]
> Deprecated ZeRO keys still parse but warn and re-map: `cpu_offload` → `offload_optimizer`, `cpu_offload_param` → `offload_param`, `cpu_offload_use_pin_memory`, and `stage3_gather_fp16_weights_on_model_save` → `gather_16bit_weights_on_model_save`. Migrate configs off them.

### 3.5 Other config blocks

These blocks exist and are parsed by `DeepSpeedConfig`; see source for full sub-schemas:

| Block | Purpose |
|---|---|
| `activation_checkpointing` | Partition/CPU-offload activations (`partition_activations`, `contiguous_memory_optimization`, `cpu_checkpointing`, `number_checkpoints`, `synchronize_checkpoint_boundary`, `profile`). |
| `flops_profiler` | FLOPS/latency profiler (`enabled`, `profile_step`, `module_depth`, `top_modules`, `detailed`, `output_file`). |
| `autotuning` | Autotuner search settings. |
| `comms_logger` | Communication logging. |
| `monitor_config` (`tensorboard`, `wandb`, `csv_monitor`, `comet`) | Metric monitoring backends. |
| `compression_training` | Quantization/pruning/layer-reduction schedule. |
| `nebula` | Nebula fast checkpointing. |
| `data_states` | DataStates checkpointing. |
| `checkpoint` | `tag_validation` (`Warn`/`Ignore`/`Fail`), `load_universal`, `use_node_local_storage`, `parallel_write.pipeline_stage`. |
| `aio` | NVMe async-I/O tuning (`block_size`, `queue_depth`, `thread_count`, `single_submit`, `overlap_events`). |
| `tensor_parallel` | AutoTP (`autotp_size`, `tp.tp_size`, `tp.tp_grain_size`). |
| `data_efficiency` | Curriculum learning & data sampling/routing. |
| `compile` | DeepCompile (`CompileConfig`). |
| `weight_quantization` | Post-training weight quantization. |
| `hybrid_engine` | RLHF hybrid engine (`enabled`, `max_out_tokens`, `inference_tp_size`, `release_inference_cache`, `pin_parameters`, `tp_gather_partition_size`). |
| `eigenvalue` | Eigenvalue computation (`enabled`, `max_iter`, `tol`, `stability`, ...). |
| `progressive_layer_drop` | PLD (`theta`, `gamma`). |
| `sparse_attention` | Sparse attention (`mode`: dense/fixed/variable/bigbird/bslongformer). |
| `elasticity` | Elastic training (`max_train_batch_size`, `micro_batch_sizes`, `min_gpus`, `max_gpus`, `model_parallel_size`, `num_gpus_per_node`). |

### 3.6 Inference config (`init_inference`)

| Key (alias) | Default | Meaning |
|---|---|---|
| `replace_with_kernel_inject` (`kernel_inject`) | `false` | Inject optimized transformer kernels. |
| `tensor_parallel.tp_size` | `1` | Inference tensor-parallel degree. |
| `dtype` | — | `fp16`/`bf16`/`int8`/`fp32`. |
| `enable_cuda_graph` | `false` | Capture CUDA graph. |
| `use_triton` / `triton_autotune` | `false` | Triton kernels / autotune. |
| `max_out_tokens` (`max_tokens`) | `1024` | Max output tokens. |
| `min_out_tokens` (`min_tokens`) | `1` | Min output tokens. |
| `injection_policy` (`injection_dict`) | `None` | Manual layer-replacement policy. |
| `quant` | — | Quantization (`num_bits`, `q_groups`). |
| `checkpoint` | `None` | Checkpoint dir/dict to load. |
| `keep_module_on_host` | `false` | Keep module on host during replace. |
| `mp_size` (deprecated → `tensor_parallel.tp_size`) | `1` | Legacy MP size. |
| `moe.ep_size` (deprecated `ep_size`) | `1` | Expert-parallel size. |

---

## 4. Deployment & Scaling

DeepSpeed has no service to deploy — you deploy your **training job**. The `deepspeed` launcher (`deepspeed/launcher/runner.py`) is the front-end; it discovers resources, encodes a base64 "world info" map, and spawns per-node `deepspeed.launcher.launch` processes (one process per rank).

### Single node

```bash
deepspeed train.py --deepspeed_config ds_config.json          # all visible GPUs
deepspeed --num_gpus 4 train.py --deepspeed_config ds_config.json
CUDA_VISIBLE_DEVICES=0,1 deepspeed train.py --deepspeed_config ds_config.json
```

On a single node with no hostfile, the runner sets `localhost` to `device_count()` and `master_addr=127.0.0.1`.

### Hostfile (multi-node)

MPI-style, one host per line:

```text
worker-0 slots=8
worker-1 slots=8
worker-2 slots=8
```

```bash
deepspeed --hostfile=/path/hostfile train.py --deepspeed_config ds_config.json
```

The default hostfile path is `/job/hostfile` (`DLTS_HOSTFILE`). Resource selection:

```bash
deepspeed -i "worker-0@worker-1:0,2" train.py ...   # include
deepspeed -e "worker-1:0" train.py ...              # exclude (mutually exclusive with -i)
deepspeed --num_nodes 2 --num_gpus 4 train.py ...   # top-N hosts, N GPUs each
```

### Multi-node launcher backends

| `--launcher` | Backend binary | Notes |
|---|---|---|
| `pdsh` (default) | `pdsh` | Parallel SSH; supports `--include`/`--exclude`; sets `PDSH_RCMD_TYPE=ssh`. |
| `openmpi` | `ompi_info`/`mpirun` | No include/exclude/num filtering; default `--mca btl ^openib`. |
| `mpich` | `mpirun` | Same devices per node required. |
| `impi` | `mpirun` (Intel MPI) | Per-rank env mapping; disables IMPI core binding. |
| `mvapich` | `mpiname`/`mpirun` | Requires MVAPICH2-GDR; sets CUDA-aware flags. |
| `slurm` | `srun` | Supports `--include`/`--exclude`/`--num_nodes`/`--num_gpus`, `--comment`. |

> [!IMPORTANT]
> PDSH requires passwordless SSH between nodes. The runner probes the first host (`ssh -o PasswordAuthentication=no <host> hostname`) and aborts with `host=<name> was not reachable via ssh` if it fails. Use `--no_ssh_check` to skip the probe, `--ssh_port` for non-standard ports, or `--no_ssh` (with `--node_rank` per node) to launch without SSH.

### Environment propagation

The runner exports variables matching `MLFLOW*`, `PYTHON*`, `MV2*`, `UCX*`, Nebula vars, plus accelerator-specific exports, and reads a `.deepspeed_env` file from `~` or CWD (filename via `DS_ENV_FILE`, default `.deepspeed_env`):

```text
NCCL_IB_DISABLE=0
NCCL_SOCKET_IFNAME=eth0
```

### Elastic training

```bash
deepspeed --elastic_training --master_addr <addr> --num_nodes MIN:MAX train.py ...
```

Elastic mode uses the `elasticity` config block to compute a valid batch size for the current GPU count and overrides batch params at runtime (warns if you also set them manually unless `ignore_non_elastic_batch_info: true`).

> [!WARNING]
> Under elasticity, batch-related keys (`train_batch_size`, `train_micro_batch_size_per_gpu`, `gradient_accumulation_steps`) are taken over by the elastic agent and will be overridden. Set `ignore_non_elastic_batch_info: true` to silence the resulting errors.

### Containers

- Mount or build with PyTorch + CUDA matching your driver.
- Use `--shm-size=1gb` (or larger) so NCCL has enough `/dev/shm`.
- Pre-build ops in the image (`DS_BUILD_OPS=1`) to avoid first-run JIT latency and to make images reproducible.

### CPU core binding

```bash
deepspeed --bind_cores_to_rank --bind_core_list 0-31 train.py ...
```

Improves host efficiency (especially CPU backend / offload). Sets `OMP_NUM_THREADS` per rank via `numactl`.

> [!CAUTION]
> Core binding conflicts with `KMP_AFFINITY` and raises an error instructing you to `unset KMP_AFFINITY` first.

---

## 5. Operations

### Health & inventory

```bash
ds_report          # op install/compatibility + torch/cuda/nccl versions, /dev/shm size
```

`ds_report` enumerates every op (`async_io`, `cpu_adam`, `fused_adam`, `transformer`, `transformer_inference`, `quantizer`, `gds`, ...) with `[YES]/[NO]` installed and `[OKAY]/[NO]` compatible, and reports torch/CUDA/HIP/NCCL versions and `/dev/shm` size with warnings.

### Logging

| Control | Effect |
|---|---|
| `--log_level {debug,info,warning,error,critical}` | Launcher/runner verbosity (default `info`). |
| `-q` / `--quiet` | Aliases to `--log_level error`. |
| `--enable_each_rank_log <dir>` | Per-rank stdout/stderr to separate files. |
| `DISABLE_LOGS_WHILE_COMPILING` | Suppress logs during compilation. |
| `steps_per_print` (config) | Training-stat cadence. |

### Checkpointing operations

```python
engine.save_checkpoint(save_dir, tag, client_state={...}, exclude_frozen_parameters=False)
load_path, client_state = engine.load_checkpoint(load_dir, tag)
engine.save_16bit_model(save_dir, "pytorch_model.bin")
```

Offline conversions:

```bash
# ZeRO checkpoint -> single FP32 state_dict
python -m deepspeed.utils.zero_to_fp32 <checkpoint_dir> <output> [--safe_serialization] [--max_shard_size 5GB]

# -> universal checkpoint for cross-topology resume
python -m deepspeed.checkpoint.ds_to_universal --input_folder <in> --output_folder <out> \
    --num_extract_workers 4 --num_merge_workers 2 [--inject_missing_state] [--keep_temp_folder]
```

`checkpoint.tag_validation` controls whether a tag mismatch across ranks is `Warn` (default), `Ignore`, or `Fail`. `checkpoint.parallel_write.pipeline_stage` enables pipelined parallel checkpoint writes; `use_node_local_storage` writes to node-local disk.

> [!NOTE]
> For very large jobs, enable `nebula` or `data_states` checkpoint engines (config blocks) for faster, asynchronous checkpoint writes. The checkpoint engine is pluggable (`deepspeed/runtime/checkpoint_engine/`).

### Monitoring

Configure a monitor backend to stream metrics:

```json
{ "tensorboard": { "enabled": true, "output_path": "logs/", "job_name": "run1" } }
```

Also supported: `wandb`, `csv_monitor`, `comet`. Enable `wall_clock_breakdown` and `flops_profiler` for performance telemetry.

### Cross-node command execution

```bash
ds_ssh -f hostfile "nvidia-smi --query-gpu=memory.used --format=csv"
```

> [!TIP]
> Standardize a `.deepspeed_env` and a versioned `ds_config.json` per job, and pre-build ops in your image. This makes runs reproducible and removes first-step JIT stalls from your timing data.

---

## 6. Performance

### Choosing and tuning ZeRO

1. **Stage selection:** lowest stage that fits. Stage 2 is the usual sweet spot; Stage 3 for the largest models.
2. **Bucket sizes:** `reduce_bucket_size` / `allgather_bucket_size` trade memory for comm efficiency (default `5e8`). Lower under memory pressure.
3. **Overlap:** `overlap_comm: true` overlaps gradient reduction with backward (default on for Stage 3).
4. **Stage 3 residency:** lower `stage3_max_live_parameters` and `stage3_prefetch_bucket_size` to cut memory; raise `stage3_param_persistence_threshold` to keep small params unpartitioned (cuts latency-bound comm).
5. **Sequential allgather:** `stage3_allgather_sequential: true` reduces peak memory at some throughput cost (high-pressure scenarios).

### Offload performance

- CPU optimizer offload runs the optimizer on CPU — install/build `cpu_adam`.
- `pin_memory: true` speeds host↔device transfers (uses more host RAM).
- NVMe offload (`async_io`) needs `libaio`; tune with `ds_nvme_tune` and the `aio` block (`queue_depth`, `block_size`, `thread_count`, `overlap_events`).
- Partial offload (`offload_optimizer.ratio < 1.0`) keeps some optimizer states on GPU (Stage 3 only).
- `super_offload` + `cpuadam_cores_perc` for Superchip CPU offload.

### Communication optimizations

| Feature | Config | Benefit |
|---|---|---|
| ZeRO++ hierarchical partitioning | `zero_hpz_partition_size` | Reduce cross-node all-gather. |
| Quantized weights/grads | `zero_quantized_weights`, `zero_quantized_gradients` | Smaller comm volume. |
| LoCo-Zero++ | `zeropp_loco_param` | Error-compensated low-comm. |
| Multi-rank bucket all-reduce | `use_multi_rank_bucket_allreduce` | Fewer, larger messages on small models. |
| 1-bit optimizers | `OneBitAdam`/`OneBitLamb`/`ZeroOneAdam` | Drastically reduced comm. |
| Communication dtype | `communication_data_type` | Reduce comm bytes. |

### Autotuning

```bash
deepspeed --autotuning tune train.py --deepspeed_config ds_config.json
deepspeed --autotuning run  train.py --deepspeed_config ds_config.json
```

The `Autotuner` (`deepspeed/autotuning/`) sweeps ZeRO/batch/bucket settings against your model and writes an optimal config; `run` tunes then trains.

### Profiling

```json
{ "flops_profiler": { "enabled": true, "profile_step": 1, "module_depth": -1, "top_modules": 1, "detailed": true },
  "wall_clock_breakdown": true }
```

> [!TIP]
> Activation checkpointing (`activation_checkpointing.partition_activations`, `cpu_checkpointing`) trades compute for activation memory — combine with ZeRO for the largest models. Pair with `sub_group_size` tuning under Stage 3 to balance memory and step time.

---

## 7. Security

DeepSpeed is an in-process library, not a network service, so its attack surface is the code/configs it executes and the cluster it orchestrates — not exposed endpoints.

> [!WARNING]
> **Configs and checkpoints are executable trust boundaries.** A DeepSpeed config can be a path, a dict, or a **base64-encoded** string that is decoded and parsed (`base64.urlsafe_b64decode` in `DeepSpeedConfig`). Checkpoints are loaded with PyTorch (pickle) by default. Treat both as code: only load configs and checkpoints from trusted sources. Prefer `--safe_serialization` (safetensors) when converting models for distribution.

Operational security guidance:

- **SSH for multi-node** — the PDSH launcher relies on passwordless SSH. Restrict keys to the training subnet, use dedicated job accounts, and avoid running as root (the `install.sh` script explicitly refuses to run as root without `--allow_sudo`).
- **Environment propagation** — `.deepspeed_env` and exported variables are pushed to all nodes; don't put secrets in them. The runner deliberately excludes some host-specific vars (e.g. under `AISC_JOB_NAME`).
- **Process control** — `--save_pid` writes a PID file under `/tmp/<pid>.ds`; the launcher installs SIGINT/SIGTERM handlers that terminate the full process tree and (PDSH) issue a remote `pkill`. Ensure `/tmp` is not world-writable on shared hosts.
- **Resource isolation** — on shared clusters, scope GPUs with `--include`/`CUDA_VISIBLE_DEVICES` and CPU cores with `--bind_core_list` to avoid noisy-neighbor interference.
- **Supply chain** — pin DeepSpeed, PyTorch, and accelerator runtime versions; build ops in a controlled image. `DS_SKIP_CUDA_CHECK`/`DS_IGNORE_CUDA_DETECTION` bypass safety checks — use sparingly and only in known-good environments.

> [!CAUTION]
> Never load a checkpoint, universal-checkpoint folder, or base64 config received from an untrusted party into a privileged training environment — deserialization can execute arbitrary code.

---

## 8. Developer Internals

### Repository structure

| Path | Contents |
|---|---|
| `deepspeed/` | The Python package (runtime, zero, comm, launcher, inference, moe, etc.). |
| `accelerator/` | Accelerator abstraction (`get_accelerator`, per-device backends). |
| `op_builder/` | Op build/JIT logic; per-accelerator subfolders (`cpu`, `xpu`, `npu`, `hpu`, `mlu`, `sdaa`). |
| `csrc/` | C++/CUDA/HIP kernel sources. |
| `tests/` | `unit/`, `model/`, `benchmarks/`, accelerator/onebit/lightning suites (pytest). |
| `docs/` | Jekyll site + Sphinx `code-docs` (readthedocs). |
| `blogs/` | Feature announcement write-ups (ZeRO++, DeepNVMe, Ulysses, FastGen, Windows, ...). |
| `benchmarks/`, `examples/` | Benchmark and example scripts. |
| `bin/` | Installed CLI scripts (`deepspeed`, `ds_report`, `ds_ssh`, ...). |
| `setup.py`, `version.txt`, `requirements/` | Build + dependency definitions. |

### Architecture deep dive

**Initialization path** (`deepspeed/__init__.py::initialize`):
1. Shut down any active `zero.Init` context, assert a model was passed.
2. `init_distributed(...)` via `deepspeed.comm` (auto-detects backend from the accelerator).
3. Resolve the config (dict / path / base64), optionally build a mesh device for sequence+data parallelism.
4. Merge tensor-parallel init args; if `tensor_parallel.autotp_size > 0`, enable AutoTP training mode.
5. Construct `DeepSpeedConfig`, set optimizer flags (e.g. Muon per-param tagging), and instantiate the engine (`DeepSpeedEngine` / `DeepSpeedHybridEngine` / `PipelineEngine`).
6. Return `(engine, optimizer, dataloader, lr_scheduler)`.

**Config system** (`config.py` + `config_utils.py`): `DeepSpeedConfig._initialize_params` wires dozens of sub-configs (zero, fp16/bf16, optimizer, scheduler, comms, monitor, flops_profiler, autotuning, nebula, datastates, aio, tensor_parallel, data_efficiency, compile, ...). `_configure_train_batch_size` derives the missing batch parameter and `_batch_assertion` enforces the batch identity. Sub-configs extend `DeepSpeedConfigModel` (pydantic v2) with a deprecation mechanism (`json_schema_extra={"deprecated": True, "new_param": ...}`).

**ZeRO** (`runtime/zero/`): `stage_1_and_2.py` and `stage3.py` implement the partitioned optimizers; `partition_parameters.py` provides the `zero.Init` context and `GatheredParameters`; `partitioned_param_coordinator.py` manages prefetch/release with a trace cache; `offload_config.py` defines CPU/NVMe offload; `mics.py` implements MiCS hierarchical sharding; `tiling.py` provides tiled linear layers.

**Communication** (`comm/`): wraps `torch.distributed` collectives behind `deepspeed.comm`, with a config (`DeepSpeedCommsConfig`), MPI/CCL utilities, and an SDMA-based all-gather path (`mori.py`, gated by `DS_SDMA_ALLGATHER`).

**Accelerator abstraction** (`accelerator/real_accelerator.py`): `get_accelerator()` returns the active backend; `DS_ACCELERATOR` overrides auto-detection. Each backend implements device, memory, and communication-backend queries. This is what lets the same code target CUDA/ROCm/XPU/HPU/NPU/MLU/SDAA/CPU.

**Op builders** (`op_builder/`): each op declares `NAME` and `BUILD_VAR` (e.g. `fused_adam`/`DS_BUILD_FUSED_ADAM`); `is_compatible()` gates pre-build/JIT. `all_ops.py` enumerates them and `installed_ops`/`compatible_ops` are recorded at build time into `git_version_info_installed.py`.

### Module reference (HIGH-priority modules)

| Module | Files / size | Highlights |
|---|---|---|
| `deepspeed/runtime` | 20 files, ~523 KB | `engine.py` (the engine), `config.py`, `constants.py`, `lr_schedules.py`, `bf16_optimizer.py`, `utils.py`. |
| `deepspeed/runtime/zero` | 18 files, ~597 KB | `stage3.py`, `stage_1_and_2.py`, `partition_parameters.py`, `config.py`, `offload_config.py`, `mics.py`. |
| `deepspeed/module_inject` | 15 files, ~204 KB | Inference layer replacement / kernel injection containers. |
| `deepspeed/utils` | 18 files, ~156 KB | `logging.py`, `groups.py`, `numa.py`, `zero_to_fp32.py`, `tensor_fragment.py`, `z3_leaf_module.py`. |
| `deepspeed/sequence` | 9 files, ~148 KB | Ulysses sequence parallelism. |
| `deepspeed/compression` | 8 files, ~109 KB | Quantization/pruning/layer reduction. |
| `deepspeed/compile` | 16 files, ~104 KB | DeepCompile passes and custom ops. |
| `accelerator` | 11 files, ~99 KB | Per-device accelerator backends. |
| `op_builder` | 26 files, ~96 KB | Op build/JIT logic. |
| `deepspeed/autotuning` | 6 files, ~96 KB | `Autotuner` and config templates. |
| `deepspeed/comm` | 10 files, ~82 KB | Distributed comm wrappers. |
| `deepspeed/checkpoint` | 10 files, ~74 KB | Universal checkpoint conversion. |
| `deepspeed/launcher` | 6 files, ~65 KB | Runner + per-node launch + multi-node runners. |
| `csrc/includes` | 30 files, ~231 KB | Shared CUDA/C++ headers. |
| `csrc/transformer` | ~48 KB | Transformer training/inference kernels. |

### Code patterns & conventions

- **Pydantic-based config** with strict validation and a custom deprecation path; prefer adding new keys as typed fields on the relevant `DeepSpeedConfigModel` with a docstring (the docstrings are the config docs).
- **Constants module** — JSON key names and defaults live as `NAME`/`NAME_DEFAULT` pairs in `constants.py`; `get_<thing>(param_dict)` accessors centralize parsing.
- **Accelerator-agnostic** — never call `torch.cuda.*` directly in portable code; use `get_accelerator()`. A pre-commit hook (`check-torchcuda.py`) enforces this.
- **Comm wrapping** — use `deepspeed.comm`, not `torch.distributed` directly (pre-commit `check-torchdist.py` enforces this, with documented exceptions).
- **Formatting/linting** — pre-commit runs `yapf`, `flake8`, `clang-format`, `codespell`, license/header checks. Run `pre-commit install` before contributing.

### Testing strategy

- pytest suites under `tests/unit/` (the bulk), plus `tests/model/`, `tests/benchmarks/`, `tests/onebit/`, `tests/torch_compile/`, `tests/accelerator/`, `tests/lightning/`, `tests/hybrid_engine/`.
- Distributed tests use a fixed timeout (`DS_UNITTEST_TIMEOUT`, default 600s) and pytest-xdist workers; `DS_DISABLE_REUSE_DIST_ENV` controls dist-env reuse.
- CI definitions live under `ci/` and `.github/workflows/` (e.g. `ci/torch_latest.py` for latest-torch runs via Modal).

### Potential improvements (from source signals)

- **Deprecation cleanup** — several ZeRO keys (`cpu_offload*`, `legacy_stage1`, `stage3_gather_fp16_weights_on_model_save`, `elastic_checkpoint`) remain for back-compat; a future major version could remove them.
- **Comm wrapper completeness** — `.pre-commit-config.yaml` excludes specific files from the `check-torchdist` hook "until we wrap ProcessGroup in deepspeed.comm," indicating an in-progress migration.
- **Known open issues** to track: ZeRO++ small-parameter `IndexError` ([#6659](https://github.com/deepspeedai/DeepSpeed/issues/6659)), DeepCompile hang on dynamic shapes ([#7611](https://github.com/deepspeedai/DeepSpeed/issues/7611)), missing zero-guards causing `ZeroDivisionError`/non-finite values ([#7838](https://github.com/deepspeedai/DeepSpeed/issues/7838)), no public API to reassemble Adam moments from sharded checkpoints ([#7961](https://github.com/deepspeedai/DeepSpeed/issues/7961)), and EvoformerAttention requiring `CUTLASS_PATH` instead of auto-detecting ([#8001](https://github.com/deepspeedai/DeepSpeed/issues/8001)).

### Contributing

1. Fork, clone, `pip install .[dev]`, and `pre-commit install`.
2. Follow the accelerator-agnostic and comm-wrapping conventions (hooks enforce them).
3. Add config fields as typed pydantic fields with docstrings; add key/default constants in `constants.py`.
4. Add tests under the appropriate `tests/` suite; keep them accelerator-agnostic where possible.
5. License headers (`Apache-2.0`, "DeepSpeed Team") are checked by `check-license.py`.

> [!NOTE]
> The internal docs build (Sphinx `code-docs`, readthedocs) is configured in `.readthedocs.yml` and `docs/code-docs/source/conf.py`, and the marketing site is a Jekyll project under `docs/`. Config field docstrings are the canonical source for the published config reference.

---

## 9. Integration Points

| Integration | How it connects |
|---|---|
| **PyTorch** | Wraps `torch.nn.Module` + `torch.optim.Optimizer`; uses `torch.distributed` for collectives; honors `torch.dtype`. |
| **Hugging Face Transformers / Accelerate** | HF Trainer and Accelerate drive DeepSpeed via a config; `DeepSpeedConfigModel` non-strict mode strips `"auto"` values so HF can fill them. ZeRO-3 integrates with HF model init via `zero.Init`. |
| **torch.distributed backends** | NCCL (CUDA), Gloo, MPI, oneCCL (`CCL_ROOT`); the backend is chosen from the accelerator's `communication_backend_name()`. |
| **Schedulers** | PDSH, OpenMPI, MPICH, Intel MPI, MVAPICH, SLURM via `--launcher`. |
| **Experiment tracking** | TensorBoard, Weights & Biases, CSV, Comet via `monitor_config`. |
| **Storage / checkpointing** | Local FS, node-local storage, Nebula, DataStates; NVMe via `async_io`/GDS for offload. |
| **Accelerators** | CUDA, ROCm, Intel XPU, Habana HPU, Ascend NPU, Cambricon MLU, SDAA, Apple MPS, CPU — selected by `get_accelerator()` / `DS_ACCELERATOR`. |
| **DeepSpeedExamples** | Companion example repo (CIFAR, BERT, DeepSpeed-Chat) referenced by tutorials. |

### Key environment variables (reference)

| Variable | Purpose |
|---|---|
| `DS_ACCELERATOR` | Force a specific accelerator backend. |
| `DS_BUILD_OPS` | Pre-build all compatible ops at install. |
| `DS_BUILD_<OP>` | Pre-build a single op (e.g. `DS_BUILD_FUSED_ADAM`). |
| `DS_ENABLE_NINJA` | Force ninja for builds. |
| `DS_NVCC_THREADS` | Parallelism for CUDA compilation. |
| `DS_SKIP_CUDA_CHECK` / `DS_IGNORE_CUDA_DETECTION` | Bypass CUDA version/detection checks. |
| `TORCH_CUDA_ARCH_LIST` | Target GPU architectures for builds. |
| `CUTLASS_PATH` / `DS_EVOFORMER_GPU_ARCH` | EvoFormer attention build. |
| `DEEPSPEED_TIMEOUT` | Default process-group timeout (minutes; default 30). |
| `DS_ENV_FILE` | Name of the per-node env file (default `.deepspeed_env`). |
| `DS_SDMA_ALLGATHER` / `DS_SDMA_ALLGATHER_MAX_NUMEL` | Enable/limit SDMA all-gather. |
| `LOCAL_RANK` / `RANK` / `WORLD_SIZE` / `CROSS_RANK` / `CROSS_SIZE` / `LOCAL_SIZE` | Distributed rank/size (set by the launcher). |
| `MASTER_ADDR` / `MASTER_PORT` | Rendezvous (default port `29500`). |
| `CUDA_VISIBLE_DEVICES` | GPU selection (respected by the launcher). |
| `DISABLE_LOGS_WHILE_COMPILING` | Quiet logs during compilation. |

> [!IMPORTANT]
> The launcher sets `LOCAL_RANK`, `RANK`, `WORLD_SIZE`, `CROSS_RANK`, `CROSS_SIZE`, `LOCAL_SIZE`, `MASTER_ADDR`, and `MASTER_PORT` for each spawned process. Your training script should read `LOCAL_RANK` (or accept `--local_rank`) to place tensors on the correct device; suppress the injected `--local_rank` argument with `--no_local_rank` if your script doesn't expect it.

---

_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._