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

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.