deepspeed

name	deepspeed
description	Comprehensive reference documentation and skill for DeepSpeed - the distributed deep learning training and inference optimization library. Covers ZeRO optimization (Stages 0-3), ZeRO-Offload, ZeRO-Infinity, SuperOffload, ZenFlow, pipeline parallelism, tensor parallelism (AutoTP), sequence parallelism (Ulysses/ALST), MoE (Mixture of Experts), inference engines (v1/v2), quantization and compression, custom optimizers (Adam, LAMB, LION, Muon, 1-bit Adam), mixed precision training (FP16/BF16/AMP), activation checkpointing, model checkpointing, communication primitives, launcher, autotuning, elasticity, monitoring, profiling, DeepCompile, DeepNVMe, accelerator abstraction, module injection, and supported model implementations. Based on DeepSpeed source code analysis.
version	0.16.x

DeepSpeed - Distributed Deep Learning Training & Inference

Overview

DeepSpeed is a deep learning optimization library that provides distributed training and inference capabilities for extreme-scale models. It powers some of the world's largest language models including MT-530B, BLOOM-176B, and GLM-130B.

Core Capabilities:

1. Memory Optimization - ZeRO (Zero Redundancy Optimizer) stages 0-3, ZeRO-Offload, ZeRO-Infinity, SuperOffload
2. Parallelism - 3D parallelism (data, pipeline, tensor), sequence parallelism (Ulysses/ALST)
3. Inference - High-performance inference with kernel injection, ragged batching, quantization
4. Communication - Compressed communication (1-bit Adam, ZeRO++), efficient collectives
5. System - DeepCompile, DeepNVMe, accelerator abstraction across 8+ hardware platforms

Supported Hardware: NVIDIA (CUDA), AMD (ROCm), Intel Gaudi (HPU), Intel XPU, Huawei Ascend (NPU), Cambricon (MLU), Tecorigin (SDAA), CPU

Key Publications: ZeRO (SC'20), ZeRO-Offload (ATC'21), ZeRO-Infinity (SC'21), DeepSpeed-MoE (ICML'22), DeepSpeed Inference (SC'22), ZenFlow (2025), SuperOffload (ASPLOS'26)

Architecture

┌─────────────────────────────────────────────────────────────────┐
│                      User Application                            │
│         model_engine.forward() / backward() / step()            │
├─────────────────────────────────────────────────────────────────┤
│                    DeepSpeed Python API                          │
│  initialize() │ init_inference() │ PipelineEngine │ HybridEngine │
├────────────────────────┬────────────────────────────────────────┤
│    Training Runtime    │          Inference Runtime              │
│  ┌──────────────────┐  │  ┌──────────────────────────────────┐  │
│  │ DeepSpeedEngine  │  │  │ InferenceEngine (v1)             │  │
│  │  - ZeRO 0/1/2/3  │  │  │  - Kernel Injection             │  │
│  │  - Pipeline      │  │  │  - Tensor Parallel               │  │
│  │  - Tensor Para   │  │  │ InferenceEngineV2                │  │
│  │  - Sequence Para │  │  │  - Ragged Batching               │  │
│  │  - Mixed Prec    │  │  │  - Blocked KV Cache              │  │
│  │  - Checkpointing │  │  │  - MoE Support                   │  │
│  └──────────────────┘  │  └──────────────────────────────────┘  │
├─────────────────────────────────────────────────────────────────┤
│                      Communication Layer                         │
│           NCCL │ CCL │ MPI │ HCCL │ Custom Backends             │
├─────────────────────────────────────────────────────────────────┤
│                     Accelerator Abstraction                      │
│    CUDA │ CPU │ HPU │ NPU │ XPU │ MLU │ SDAA │ MPS             │
├─────────────────────────────────────────────────────────────────┤
│                   Custom Ops & CUDA Kernels                      │
│  FusedAdam │ FusedLamb │ Transformer │ SparseAttn │ Quantizer   │
│  AsyncIO │ GDS │ CPU Adam │ Evoformer │ Spatial                   │
└─────────────────────────────────────────────────────────────────┘

Quick Reference

Basic Training

import deepspeed

# Initialize DeepSpeed engine
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    optimizer=optimizer,
    args=args,
    config_params=ds_config,  # or path to ds_config.json
)

# Training loop
for step, batch in enumerate(dataloader):
    inputs, labels = batch
    outputs = model_engine(inputs)
    loss = criterion(outputs, labels)
    model_engine.backward(loss)
    model_engine.step()

Basic Inference

import deepspeed

# Initialize inference engine
model = deepspeed.init_inference(
    model=model,
    mp_size=2,              # tensor parallel size
    dtype=torch.float16,
    replace_with_kernel_inject=True,
)

# Run inference
outputs = model(inputs)

Launch with DeepSpeed

# Single-node
deepspeed --num_gpus=4 train.py --deepspeed ds_config.json

# Multi-node
deepspeed --num_nodes=2 --hostfile=myhostfile train.py --deepspeed ds_config.json

Minimal Configuration (ZeRO-2)

{
    "train_batch_size": 32,
    "gradient_accumulation_steps": 1,
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu"
        }
    },
    "bf16": {
        "enabled": true
    }
}

Reference Chapters

#	Chapter	Description
01	Overview & Architecture	DeepSpeed architecture, design philosophy, component overview
02	Installation & Setup	Installation methods, requirements, environment setup
03	Configuration Reference	Complete ds_config.json schema with all fields
04	Getting Started	Quick-start guide, basic training and inference
05	DeepSpeed Engine	DeepSpeedEngine, PipelineEngine, HybridEngine classes
06	ZeRO Optimization	ZeRO stages 0-3, parameter/gradient partitioning
07	ZeRO Offload & Infinity	ZeRO-Offload, ZeRO-Infinity, NVMe offloading
08	SuperOffload & ZenFlow	SuperOffload, ZenFlow stall-free offloading
09	Mixed Precision Training	FP16, BF16, AMP, loss scaling
10	Pipeline Parallelism	Pipeline stages, topology, scheduling
11	Tensor Parallelism	AutoTP, manual TP, partition configs
12	Sequence Parallelism	Ulysses, ALST, long sequence training
13	MoE (Mixture of Experts)	Expert parallelism, sharded MoE, routing
14	Inference Engine V1	Kernel injection, model replacement, quantization
15	Inference Engine V2	Ragged batching, blocked KV cache, model implementations
16	Quantization & Compression	ZeroQuant, MoQ, weight/activation quantization, pruning
17	Optimizers	FusedAdam, FusedLamb, FusedLion, 1-bit Adam, Muon
18	Schedulers	WarmupLR, OneCycle, LRRangeTest, custom schedulers
19	Activation Checkpointing	Memory optimization through activation recomputation
20	Model Checkpointing	Save/load, universal checkpoint, ZeRO checkpoint
21	Communication Primitives	Comm module, NCCL, CCL, coalesced collectives
22	Launcher	Multi-node launch, elastic training, resource management
23	Autotuning	Automatic hyperparameter optimization
24	Elasticity	Dynamic resource scaling, fault tolerance
25	Monitoring & Profiling	TensorBoard, WandB, FLOPS profiler
26	Accelerator Abstraction	Multi-hardware support, custom accelerators
27	Custom Ops & Kernels	CUDA kernels, op builder, JIT compilation
28	Module Inject & AutoTP	Model replacement, automatic tensor parallelism
29	Model Implementations	Supported models: LLaMA, Mistral, Falcon, Qwen, etc.
30	DeepCompile	Compiler optimizations, FX passes, Inductor integration
31	DeepSpeed4Science	Scientific computing, Evoformer attention
32	Data Pipeline	Data efficiency, curriculum learning, data routing
33	IO & DeepNVMe	Async IO, GPU Direct Storage, NVMe offloading
34	Communication Compression	1-bit Adam, 0/1 Adam, ZeRO++ quantization
35	Nebula & DataStates	Asynchronous checkpointing, data state management
36	Muon Optimizer	MomentUm Orthogonalized by Newton-Schulz
37	Constants & Utilities	All constants, utility functions, helpers
38	Debugging & Troubleshooting	Common issues, debugging tools, environment report
39	Integrations	HuggingFace, Accelerate, Lightning, MosaicML
40	API Reference	Complete public API reference

Key Configuration Sections

ZeRO Stages

{"zero_optimization": {"stage": 0}}  // Disabled (standard DDP)
{"zero_optimization": {"stage": 1}}  // Optimizer state partitioning
{"zero_optimization": {"stage": 2}}  // + Gradient partitioning
{"zero_optimization": {"stage": 3}}  // + Parameter partitioning

Offloading Options

{"offload_optimizer": {"device": "cpu"}}     // CPU optimizer offload
{"offload_optimizer": {"device": "nvme"}}    // NVMe optimizer offload
{"offload_param": {"device": "cpu"}}         // CPU param offload (stage 3)
{"offload_param": {"device": "nvme"}}        // NVMe param offload (stage 3)

Supported Optimizers

Adam, AdamW, Lamb, OneBitAdam, OneBitLamb, ZeroOneAdam, Lion, Muon, MuAdam, MuAdamW, MuSGD, Adagrad

Supported Models for Inference

LLaMA 2, Mistral, Mixtral, Qwen, Qwen v2, Falcon, Phi, Phi-3, OPT, ExaOne4, and custom models via injection policies.