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 epochfor data in trainloader双层循环做前向/反向/optimizer.zero_grad()/optimizer.step()。二、让 CIFAR-10 支持 DeepSpeed四个改造步骤把原生模型接入 DeepSpeed 只需要四处改动① 命令行参数解析、② 引擎初始化、③ 训练 API、④ JSON 配置文件。下面逐步展开。步骤 1接入命令行参数解析在add_argument()中通过deepspeed.add_config_arguments(parser)把 DeepSpeed 需要的命令行参数注册进原有的argparse.ArgumentParserimport argparse import deepspeed def add_argument(): parserargparse.ArgumentParser(descriptionCIFAR) # Data. # Cuda. parser.add_argument(--with_cuda, defaultFalse, actionstore_true, helpuse CPU in case there\s no GPU support) parser.add_argument(--use_ema, defaultFalse, actionstore_true, helpwhether use exponential moving average) # Train. parser.add_argument(-b, --batch_size, default32, typeint, helpmini-batch size (default: 32)) parser.add_argument(-e, --epochs, default30, typeint, helpnumber of total epochs (default: 30)) parser.add_argument(--local_rank, typeint, default-1, helplocal rank passed from distributed launcher) # Include DeepSpeed configuration arguments. parser deepspeed.add_config_arguments(parser) argsparser.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, optimizerNone, model_paramsNone, training_dataNone, lr_schedulerNone, mpuNone, dist_init_requiredTrue, collate_fnNone):初始化 DeepSpeed传入 CIFAR-10 网络net、args、需要优化的参数以及训练集trainsetparameters filter(lambda p: p.requires_grad, net.parameters()) argsadd_argument() # Initialize DeepSpeed to use the following features # 1) Distributed model. # 2) Distributed data loader. # 3) DeepSpeed optimizer. model_engine, optimizer, trainloader, _ deepspeed.initialize(argsargs, modelnet, model_parametersparameters, training_datatrainset)初始化完成之后原脚本中手动搬移设备、手动创建优化器的代码就可以删除/注释掉了#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(), lr0.001, momentum0.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 的训练 APIdeepspeed.initialize返回的model是DeepSpeed 模型引擎Model Engine训练循环改用它的forward / backward / stepAPIfor 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_size4因此在单卡运行时引擎会推断出train_micro_batch_size_per_gpu4可从引擎启动日志中印证。steps_per_print每隔 N 个训练 step 打印一次进度报告包含步数、跳过的优化器更新次数、当前学习率与动量等默认值为 10。教程设置为 2000 以降低日志频率。optimizer优化器声明。type: Adam表示使用 DeepSpeed 内置 Adam 优化器params中的lr/betas/eps/weight_decay直接对应优化器构造参数。GPU 且 CUDA 扩展可用时Adam 通常对应到融合实现的 FusedAdam。scheduler学习率调度器声明。type: WarmupLRparams包含warmup_min_lr起始最小学习率、warmup_max_lrwarmup 结束后固定的最大学习率、warmup_num_steps从 min 升到 max 所需步数默认 1000。注意使用 DeepSpeed 内置调度器时引擎会在每次model_engine.step()时自动调用调度器的step()。wall_clock_breakdown布尔值是否统计前向/反向/更新各训练阶段耗时默认false。三、启动 DeepSpeed 训练改造完成后使用 DeepSpeed 自带的deepspeed启动器运行默认使用宿主机上检测到的全部 GPUdeepspeed cifar10_deepspeed.py --deepspeed_config ds_config.jsondeepspeed启动器的实现在 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: version2.1, git-hashfa937e7, git-branchmaster [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, SamplesPerSec1303.6726433398537 [INFO] 0/150, SamplesPerSec1304.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 开关等逐项打印出来方便核对配置是否如预期生效。吞吐统计SamplesPerSecxxx表示每秒处理的样本数来自引擎内部的计时与计数逻辑。学习率/动量快照rank:0 step0, skipped0, 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-L267engine包装了用户模型的 DeepSpeed 运行时引擎负责分布式训练optimizer若 JSON 配置中声明了优化器则返回被引擎包装的优化器对象training_dataloader只有当传入training_datatorch.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此后固定在最大学习率不再变化。同文件中还提供了在此基础上继续线性衰减的WarmupDecayLRdeepspeed/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 系列显存优化能力的统一入口。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 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 epochfor data in trainloader双层循环做前向/反向/optimizer.zero_grad()/optimizer.step()。二、让 CIFAR-10 支持 DeepSpeed四个改造步骤把原生模型接入 DeepSpeed 只需要四处改动① 命令行参数解析、② 引擎初始化、③ 训练 API、④ JSON 配置文件。下面逐步展开。步骤 1接入命令行参数解析在add_argument()中通过deepspeed.add_config_arguments(parser)把 DeepSpeed 需要的命令行参数注册进原有的argparse.ArgumentParserimport argparse import deepspeed def add_argument(): parserargparse.ArgumentParser(descriptionCIFAR) # Data. # Cuda. parser.add_argument(--with_cuda, defaultFalse, actionstore_true, helpuse CPU in case there\s no GPU support) parser.add_argument(--use_ema, defaultFalse, actionstore_true, helpwhether use exponential moving average) # Train. parser.add_argument(-b, --batch_size, default32, typeint, helpmini-batch size (default: 32)) parser.add_argument(-e, --epochs, default30, typeint, helpnumber of total epochs (default: 30)) parser.add_argument(--local_rank, typeint, default-1, helplocal rank passed from distributed launcher) # Include DeepSpeed configuration arguments. parser deepspeed.add_config_arguments(parser) argsparser.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, optimizerNone, model_paramsNone, training_dataNone, lr_schedulerNone, mpuNone, dist_init_requiredTrue, collate_fnNone):初始化 DeepSpeed传入 CIFAR-10 网络net、args、需要优化的参数以及训练集trainsetparameters filter(lambda p: p.requires_grad, net.parameters()) argsadd_argument() # Initialize DeepSpeed to use the following features # 1) Distributed model. # 2) Distributed data loader. # 3) DeepSpeed optimizer. model_engine, optimizer, trainloader, _ deepspeed.initialize(argsargs, modelnet, model_parametersparameters, training_datatrainset)初始化完成之后原脚本中手动搬移设备、手动创建优化器的代码就可以删除/注释掉了#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(), lr0.001, momentum0.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 的训练 APIdeepspeed.initialize返回的model是DeepSpeed 模型引擎Model Engine训练循环改用它的forward / backward / stepAPIfor 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_size4因此在单卡运行时引擎会推断出train_micro_batch_size_per_gpu4可从引擎启动日志中印证。steps_per_print每隔 N 个训练 step 打印一次进度报告包含步数、跳过的优化器更新次数、当前学习率与动量等默认值为 10。教程设置为 2000 以降低日志频率。optimizer优化器声明。type: Adam表示使用 DeepSpeed 内置 Adam 优化器params中的lr/betas/eps/weight_decay直接对应优化器构造参数。GPU 且 CUDA 扩展可用时Adam 通常对应到融合实现的 FusedAdam。scheduler学习率调度器声明。type: WarmupLRparams包含warmup_min_lr起始最小学习率、warmup_max_lrwarmup 结束后固定的最大学习率、warmup_num_steps从 min 升到 max 所需步数默认 1000。注意使用 DeepSpeed 内置调度器时引擎会在每次model_engine.step()时自动调用调度器的step()。wall_clock_breakdown布尔值是否统计前向/反向/更新各训练阶段耗时默认false。三、启动 DeepSpeed 训练改造完成后使用 DeepSpeed 自带的deepspeed启动器运行默认使用宿主机上检测到的全部 GPUdeepspeed cifar10_deepspeed.py --deepspeed_config ds_config.jsondeepspeed启动器的实现在 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: version2.1, git-hashfa937e7, git-branchmaster [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, SamplesPerSec1303.6726433398537 [INFO] 0/150, SamplesPerSec1304.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 开关等逐项打印出来方便核对配置是否如预期生效。吞吐统计SamplesPerSecxxx表示每秒处理的样本数来自引擎内部的计时与计数逻辑。学习率/动量快照rank:0 step0, skipped0, 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-L267engine包装了用户模型的 DeepSpeed 运行时引擎负责分布式训练optimizer若 JSON 配置中声明了优化器则返回被引擎包装的优化器对象training_dataloader只有当传入training_datatorch.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此后固定在最大学习率不再变化。同文件中还提供了在此基础上继续线性衰减的WarmupDecayLRdeepspeed/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 系列显存优化能力的统一入口。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 RELATED READING
延伸阅读
更多一线实战笔记与深度复盘，助您持续精进
