← Back to reference

Deepspeed Integration Guide

Install

pip install --upgrade denpex deepspeed

Wrap a DeepSpeed Engine

import denpex
import deepspeed

ds_engine = deepspeed.initialize(
    model=model,
    optimizer=optimizer,
    config="ds_config.json"
)[0]

from denpex.deepspeed import attach_deepspeed_instrument
instrument = attach_deepspeed_instrument(
    ds_engine,
    sample_every_n_steps=10
)

for step in range(steps):
    loss = ds_engine(batch)
    ds_engine.backward(loss)
    ds_engine.step()
    instrument.flush_if_due(step)

What We Capture

  • model_engine.communication_data_type
  • model_engine.optimizer.curriculum
  • model_engine.optimizer.partition_count
  • ZeRO stage + offload configuration
  • Partition parameter shape

SDC Canary

denpex.sdc.SdcRunner operates unchanged on top of the DeepSpeed microbatch.

On Crash

DeepSpeed-specific rule IDs:

  • DS_ZERO3_PARTITION_OOM
  • DS_HOLD_TIMEOUT
  • DS_GRAD_SKEW
  • DS_BACKWARD_FAIL

Compatibility Tested

  • DeepSpeed 0.13, 0.14, 0.15.
  • ZeRO stage 1, 2, 3.
  • Offload to CPU + NVMe.

Debug Recipes

docs/user-guides/deepspeed-recipes.md covers:

  • Symptom → cause matrix.
  • How to enable per-rank logging without overhead.
  • Common ZeRO-3 partition blow-up patterns.