DeepSpeed 入门实战：为 CIFAR-10 图像分类模型接入 DeepSpeed 引擎的完整改造教程-平芜编程栈

DeepSpeed 入门实战：为 CIFAR-10 图像分类模型接入 DeepSpeed 引擎的完整改造教程

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

导读

本篇教程面向第一次接触 DeepSpeed 的开发者，以 CIFAR-10 这一经典小型图像分类模型为载体，完整演示如何将一个普通 PyTorch 训练脚本一步步改造成由DeepSpeed 引擎（Engine）驱动的分布式训练程序。你将掌握deepspeed.add_config_arguments参数接入、deepspeed.initialize初始化、Engine 的 forward/backward/step 训练 API，以及ds_config.json配置文件的编写与deepspeed启动器的用法——这套改造方法论可平移到任何torch.nn.Module模型上。

本教程对应仓库中的官方教程文档 cifar-10.md，属于"Getting Started"（快速上手）系列，建议先阅读 getting-started.md 了解安装与基础概念后再继续。

一、运行原始 CIFAR-10 模型（改造前基线）

CIFAR-10 模型是来自 PyTorch 官方示例的一个小型图像分类卷积网络。教程将其以子模块（submodule）的形式组织在 DeepSpeedExamples 配套仓库的training/cifar/目录下。首先拉取子模块代码：

git submodule update --init --recursive

进入示例目录并安装模型依赖：

cd DeepSpeedExamples/cifar pip install -r requirements.txt

安装完成后直接运行原生训练脚本python cifar10_tutorial.py，脚本会在首次运行时自动下载 CIFAR-10 数据集。首次运行的控制台输出大致如下：

Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./data/cifar-10-python.tar.gz 170500096it [00:02, 61124868.24it/s] Extracting ./data/cifar-10-python.tar.gz to ./data Files already downloaded and verified cat frog frog frog [1, 2000] loss: 2.170 [1, 4000] loss: 1.879 [1, 6000] loss: 1.690 [1, 8000] loss: 1.591 [1, 10000] loss: 1.545 ... [2, 12000] loss: 1.287 Finished Training GroundTruth: cat ship ship plane Predicted: cat ship plane plane Accuracy of the network on the 10000 test images: 53 %

说明：以上输出与教程文档中记录的某次运行结果一致，具体 loss 与精度会随随机种子、数据顺序等因素波动，此处仅作为"改造前基线"的直观参照，并非可复现的承诺值。

原生脚本中训练相关的主要逻辑（后续改造会逐一替换）包括：把网络移动到 GPU、创建optim.SGD优化器、用for epoch+for data in trainloader双层循环做前向/反向/optimizer.zero_grad()/optimizer.step()。

二、让 CIFAR-10 支持 DeepSpeed：四个改造步骤

把原生模型接入 DeepSpeed 只需要四处改动：① 命令行参数解析、② 引擎初始化、③ 训练 API、④ JSON 配置文件。下面逐步展开。

步骤 1：接入命令行参数解析

在add_argument()中，通过deepspeed.add_config_arguments(parser)把 DeepSpeed 需要的命令行参数注册进原有的argparse.ArgumentParser：

import argparse import deepspeed def add_argument(): parser=argparse.ArgumentParser(description='CIFAR') # Data. # Cuda. parser.add_argument('--with_cuda', default=False, action='store_true', help='use CPU in case there\'s no GPU support') parser.add_argument('--use_ema', default=False, action='store_true', help='whether use exponential moving average') # Train. parser.add_argument('-b', '--batch_size', default=32, type=int, help='mini-batch size (default: 32)') parser.add_argument('-e', '--epochs', default=30, type=int, help='number of total epochs (default: 30)') parser.add_argument('--local_rank', type=int, default=-1, help='local rank passed from distributed launcher') # Include DeepSpeed configuration arguments. parser = deepspeed.add_config_arguments(parser) args=parser.parse_args() return args

从源码 deepspeed/init.py#L305-L318 可以看到，add_config_arguments实际调用内部辅助函数_add_core_arguments，为命令行注册的核心参数有：

参数	含义
`--deepspeed`	布尔开关，启用 DeepSpeed（辅助 flag，供用户代码分支判断，不影响 DeepSpeed 后端行为）
`--deepspeed_config <json>`	指定 DeepSpeed 运行时配置 JSON 文件的路径
`--deepscale`/`--deepscale_config`	已废弃的历史名称（早期 DeepSpeed 曾叫 DeepScale），用于向后兼容

其中--deepspeed_config是真正驱动训练的核心参数，deepspeed.initialize会读取args.deepspeed_config指向的配置文件（见后文）。

步骤 2：通过 deepspeed.initialize 初始化引擎

改造的核心是把原来的模型包装、设备搬运、优化器创建等工作统一交给deepspeed.initialize。教程给出的接口形态（与仓库中 deepspeed/init.py#L93-L107 的实际签名一致，当前版本还额外支持config/config_params、mesh_param等）：

def initialize(args, model, optimizer=None, model_params=None, training_data=None, lr_scheduler=None, mpu=None, dist_init_required=True, collate_fn=None):

初始化 DeepSpeed，传入 CIFAR-10 网络net、args、需要优化的参数以及训练集trainset：

parameters = filter(lambda p: p.requires_grad, net.parameters()) args=add_argument() # Initialize DeepSpeed to use the following features # 1) Distributed model. # 2) Distributed data loader. # 3) DeepSpeed optimizer. model_engine, optimizer, trainloader, _ = deepspeed.initialize(args=args, model=net, model_parameters=parameters, training_data=trainset)

初始化完成之后，原脚本中手动搬移设备、手动创建优化器的代码就可以删除/注释掉了：

#from deepspeed.accelerator import get_accelerator #device = torch.device(get_accelerator().device_name(0) if get_accelerator().is_available() else "cpu") #net.to(device) #optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

deepspeed.initialize内部做了什么？结合源码实现可以梳理出完整调用链：

校验输入：model不能为空，且必须提供配置（--deepspeed_config文件路径或config字典），否则抛出断言（deepspeed/init.py#L165、L192）。
自动初始化分布式环境：根据get_accelerator().communication_backend_name()选择的通信后端（GPU 环境默认 NCCL）调用dist.init_distributed(...)。这意味着你不必再手动执行torch.distributed.init_process_group(...)，DeepSpeed 会在内部替你完成。
读取并解析配置：把 JSON 文件加载成配置字典，交由DeepSpeedConfig统一管理。
创建 Engine：根据是否启用流水线并行（PipelineModule）、Hybrid Engine 等选择构造DeepSpeedEngine、DeepSpeedHybridEngine或PipelineEngine（见 deepspeed/init.py#L213-L256）。
返回四元组(engine, optimizer, training_dataloader, lr_scheduler)。其中optimizer是由 JSON 配置创建并被引擎包装的优化器；training_dataloader是传入training_data时自动构造的分布式数据加载器；lr_scheduler来自 JSON 配置中的 scheduler 字段。

步骤 3：切换到 Engine 的训练 API

deepspeed.initialize返回的model是DeepSpeed 模型引擎（Model Engine），训练循环改用它的forward / backward / stepAPI：

for i, data in enumerate(trainloader): # Get the inputs; data is a list of [inputs, labels]. inputs = data[0].to(model_engine.device) labels = data[1].to(model_engine.device) outputs = model_engine(inputs) loss = criterion(outputs, labels) model_engine.backward(loss) model_engine.step()

与原生写法相比有三个关键差异：

设备句柄由引擎提供：输入统一搬到model_engine.device，引擎已在内部把模型放置到正确的加速设备上，无需自己维护device变量。
model_engine.backward(loss)取代loss.backward()：Engine 的backward在分布式数据并行下会自动完成梯度平均（gradient averaging）等操作，实现位于 deepspeed/runtime/engine.py#L3288 附近。
model_engine.step()取代optimizer.step()：Engine 在参数更新后会自动清零梯度（zeroing the gradients is handled automatically after the weights have been updated），因此无需再显式调用optimizer.zero_grad()。step()实现位于 deepspeed/runtime/engine.py#L3500 附近，且每当一个有效 batch（见下文train_batch_size语义）完成时，引擎还会自动推进学习率调度器并输出训练进度报告。

步骤 4：编写 ds_config.json 配置文件

要真正启用 DeepSpeed，需要创建一个 JSON 配置文件（惯例命名为ds_config.json）。该文件集中声明批大小、优化器、调度器等 DeepSpeed 专属参数：

{ "train_batch_size": 4, "steps_per_print": 2000, "optimizer": { "type": "Adam", "params": { "lr": 0.001, "betas": [ 0.8, 0.999 ], "eps": 1e-8, "weight_decay": 3e-7 } }, "scheduler": { "type": "WarmupLR", "params": { "warmup_min_lr": 0, "warmup_max_lr": 0.001, "warmup_num_steps": 1000 } }, "wall_clock_breakdown": false }

各配置项的含义如下（详细字段定义可查阅 config-json.md）：

配置项	含义与要点
`train_batch_size`	有效训练批大小：指"累积到一次模型参数更新"所需的样本总量。它由单卡单步批大小`train_micro_batch_size_per_gpu`、梯度累积步数`gradient_accumulation_steps`与 GPU 数量共同决定，即`train_batch_size = train_micro_batch_size_per_gpu × gradient_accumulation_steps × GPU数`。三个量只需指定其中两个，第三个会被 DeepSpeed 自动推算。本配置只给出`train_batch_size=4`，因此在单卡运行时引擎会推断出`train_micro_batch_size_per_gpu=4`（可从引擎启动日志中印证）。
`steps_per_print`	每隔 N 个训练 step 打印一次进度报告（包含步数、跳过的优化器更新次数、当前学习率与动量等），默认值为 10。教程设置为 2000 以降低日志频率。
`optimizer`	优化器声明。`type: "Adam"`表示使用 DeepSpeed 内置 Adam 优化器；`params`中的`lr`/`betas`/`eps`/`weight_decay`直接对应优化器构造参数。GPU 且 CUDA 扩展可用时，"Adam" 通常对应到融合实现的 FusedAdam。
`scheduler`	学习率调度器声明。`type: "WarmupLR"`，`params`包含`warmup_min_lr`（起始最小学习率）、`warmup_max_lr`（warmup 结束后固定的最大学习率）、`warmup_num_steps`（从 min 升到 max 所需步数，默认 1000）。注意：使用 DeepSpeed 内置调度器时，引擎会在每次`model_engine.step()`时自动调用调度器的`step()`。
`wall_clock_breakdown`	布尔值，是否统计前向/反向/更新各训练阶段耗时，默认`false`。

三、启动 DeepSpeed 训练

改造完成后，使用 DeepSpeed 自带的deepspeed启动器运行（默认使用宿主机上检测到的全部 GPU）：

deepspeed cifar10_deepspeed.py --deepspeed_config ds_config.json

deepspeed启动器的实现在 deepspeed/launcher/launch.py 中，会解析节点/GPU 资源信息（hostfile、--num_gpus、--num_nodes、--include/--exclude等），为每个进程设置正确的环境变量并拉起训练。单机场景下无需 hostfile，启动器会自动探测本机 GPU 数量；也可通过CUDA_VISIBLE_DEVICES或--include localhost:0,1限定使用的设备。更完整的启动参数说明见 getting-started.md 的 "Launching DeepSpeed Training" 章节。

与原生脚本相比，DeepSpeed 会在控制台输出更丰富的训练信息，包括训练配置、性能统计与 loss 走势。教程文档记录的启动输出节选如下（该日志录制于较早版本，当前版本输出格式可能略有差异，但关键信息一致）：

DeepSpeed info: version=2.1, git-hash=fa937e7, git-branch=master [INFO] Set device to local rank 0 within node. [INFO] Using DeepSpeed Optimizer param name adam as basic optimizer DeepSpeed Basic Optimizer = FusedAdam ( Parameter Group 0 betas: [0.8, 0.999] bias_correction: True eps: 1e-08 lr: 0.001 weight_decay: 3e-07 ) [INFO] DeepSpeed using configured LR scheduler = WarmupLR DeepSpeedLight configuration: train_batch_size ............. 4 train_micro_batch_size_per_gpu 4 optimizer_name ............... adam scheduler_name ............... WarmupLR steps_per_print .............. 2000 wall_clock_breakdown ......... False ... [INFO] 0/100, SamplesPerSec=1303.6726433398537 [INFO] 0/150, SamplesPerSec=1304.4251022567403 ...... [2, 12000] loss: 1.247 Finished Training GroundTruth: cat ship ship plane Predicted: cat car car plane Accuracy of the network on the 10000 test images: 57 %

日志中值得关注的信息点：

引擎配置回显：启动时引擎会把解析、推算后的完整配置（包括推断出的train_micro_batch_size_per_gpu、world size、FP16/Zeo 开关等）逐项打印出来，方便核对配置是否如预期生效。
吞吐统计：SamplesPerSec=xxx表示每秒处理的样本数，来自引擎内部的计时与计数逻辑。
学习率/动量快照：rank:0 step=0, skipped=0, lr=[0.001], mom=[[0.8, 0.999]]一类输出用于监控优化器状态，skipped字段在混合精度训练发生溢出跳步时会大于 0。

说明：日志中 53%（原生）与 57%（DeepSpeed）均为教程文档记录的单次示例运行结果，二者差异来自随机性，不应解读为"接入 DeepSpeed 必然提升精度"。接入 DeepSpeed 的价值在于分布式扩展、混合精度、ZeRO 显存优化等工程能力，而非改变模型收敛精度。

四、源码级机制解读：initialize、Engine 与调度器

1.`deepspeed.initialize`的返回语义

initialize总是返回(engine, optimizer, training_dataloader, lr_scheduler)四元组（deepspeed/init.py#L261-L267）：

engine：包装了用户模型的 DeepSpeed 运行时引擎，负责分布式训练；
optimizer：若 JSON 配置中声明了优化器，则返回被引擎包装的优化器对象；
training_dataloader：只有当传入training_data（torch.utils.data.Dataset）时才非空。Engine 会借助 DeepSpeedDataLoader 依据分布式 world size 切分数据（内置 distributed sampler 处理数据分片），这也是 CIFAR-10 示例能直接拿到trainloader的原因；
lr_scheduler：由 JSON 的 scheduler 字段创建，供用户侧手动管理或引擎托管。

2.`model_engine.step()`的自动梯度管理

DeepSpeedEngine 的定义位于 deepspeed/runtime/engine.py#L334。backward()内部负责梯度累积与（分布式）梯度归并；当累积达到一个train_batch_size后，step()会执行优化器更新、按steps_per_print决定是否打印进度，并推进学习率。原脚本中手工书写的optimizer.zero_grad()全部被引擎接管，这正是"把循环体重写为forward → backward → step三段式"能够成立的根本原因。

3. WarmupLR 等内置调度器

scheduler.type = "WarmupLR"对应的类定义在 deepspeed/runtime/lr_schedules.py#L703。其语义是：学习率在warmup_num_steps（校验为正整数）步内从warmup_min_lr线性爬升到warmup_max_lr，此后固定在最大学习率不再变化。同文件中还提供了在此基础上继续线性衰减的WarmupDecayLR（deepspeed/runtime/lr_schedules.py#L791）等变体，可按训练总步数选择。由于调度器由引擎托管，训练循环里不需要任何手动scheduler.step()调用。

4. 进一步扩展

CIFAR-10 示例只用了 DeepSpeed 最基础的能力。同一份配置骨架可以通过追加字段开启更多特性，例如：

混合精度：在配置中加入"fp16": {"enabled": true}（或 BF16）；
显存优化：加入"zero_optimization": {"stage": 2}（或 Stage 3）；
梯度累积：显式声明train_micro_batch_size_per_gpu与gradient_accumulation_steps。

这些字段的行为定义、约束与完整参数表都收录在 config-json.md 中；各类进阶能力（ZeRO、Offload、MoE、自动调参等）对应教程可参考 docs/_tutorials 目录下的相关文档。对本仓库只读使用而言，读者可以直接基于本教程的配置骨架，在本地逐步叠加上述字段来观察引擎启动日志与显存/吞吐变化。

小结

至此，你已完成一次完整的"原生 PyTorch → DeepSpeed"改造闭环：

用deepspeed.add_config_arguments(parser)让脚本识别--deepspeed_config；
用deepspeed.initialize(...)一次性拿到分布式模型引擎、优化器与分布式数据加载器；
把训练循环改为model_engine(inputs) → backward(loss) → step()三段式，梯度清零交给引擎；
编写声明train_batch_size、优化器与学习率调度器的ds_config.json；
用deepspeed cifar10_deepspeed.py --deepspeed_config ds_config.json启动。

这套模式不局限于图像分类：任何torch.nn.Module模型都可以沿用相同思路接入 DeepSpeed 引擎，从而获得分布式训练、混合精度与后续 ZeRO 系列显存优化能力的统一入口。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

DeepSpeed 入门实战：为 CIFAR-10 图像分类模型接入 DeepSpeed 引擎的完整改造教程