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

Technology stack

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

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.

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

3.2 Optimizer & scheduler

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

3.3 Mixed precision

fp16

bf16

torch_autocast

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

offload_param (Stage 3)

offload_optimizer (Stage 1/2/3)

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:

3.6 Inference config (init_inference)

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

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:

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

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

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

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

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

NCCL_IB_DISABLE=0
NCCL_SOCKET_IFNAME=eth0

Elastic training

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

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

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

Checkpointing operations

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:

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

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

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

Autotuning

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

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

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)

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), DeepCompile hang on dynamic shapes (#7611), missing zero-guards causing ZeroDivisionError/non-finite values (#7838), no public API to reassemble Adam moments from sharded checkpoints (#7961), and EvoformerAttention requiring CUTLASS_PATH instead of auto-detecting (#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

Key environment variables (reference)

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.