DeepSpeed 入门实战指南安装、模型接入、配置与分布式训练启动全解析【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeedDeepSpeed 是一个用于深度学习的优化库它让分布式训练与推理变得简单、高效且有效。本文以官方入门教程为主体结合仓库内源码实现完整介绍从安装验证、把任意torch.nn.Module接入 DeepSpeed 引擎、编写训练与断点checkpoint逻辑、编写ds_config.json配置文件再到使用deepspeed命令在单机/多机上启动分布式训练的全过程。读完本文你将掌握一套可直接落地的最小可运行训练闭环并能理解其底层引擎 API 的实际行为。一、安装与安装验证安装 DeepSpeed 最简单的方式是通过 pip 安装最新 release 版本该版本不绑定特定的 PyTorch 或 CUDA 版本pip install deepspeedDeepSpeed 包含若干 C/CUDA 扩展官方称为 ops。默认情况下这些扩展会通过 PyTorch 的 JIT C 扩展加载器底层依赖 ninja在运行时被**即时编译just-in-timeJIT**并动态链接因此在首次运行用到对应算子时才编译安装本身很快。更详细的安装方式源码编译、预编译算子、构建 wheel、conda 环境等参见高级安装教程。安装完成后可用 DeepSpeed 环境报告工具验证安装并检查当前机器支持哪些算子。ds_report在排查安装与兼容性问题时非常有用ds_report # 等价形式 python -m deepspeed.env_report该命令会输出 PyTorch/CUDA 版本、GPU 架构、各DS_BUILD_*算子是否可编译等信息对应实现位于 deepspeed/env_report.py。如果你希望安装预编译全部与机器兼容的算子可在安装时设置DS_BUILD_OPS1DS_BUILD_OPS1 pip install deepspeed若只想安装某一个算子例如FusedLamb可用对应的DS_BUILD_*环境变量精确控制DS_BUILD_FUSED_LAMB1 pip install deepspeed常用DS_BUILD选项包括均用于预编译JIT 模式下无需设置环境变量作用DS_BUILD_OPS编译全部算子DS_BUILD_AIO编译异步NVMeI/O 算子DS_BUILD_CPU_ADAM/DS_BUILD_CPU_ADAGRAD/DS_BUILD_CPU_LION编译 CPU 版 Adam/Adagrad/LionDS_BUILD_FUSED_ADAM/DS_BUILD_FUSED_LAMB/DS_BUILD_FUSED_LION编译融合版 Adam/LAMB/LionDS_BUILD_CCL_COMM编译通信集合库DS_BUILD_EVOFORMER_ATTN编译 EvoformerAttnAlphaFold 相关算子DS_BUILD_QUANTIZER/DS_BUILD_RANDOM_LTD/DS_BUILD_SPARSE_ATTN编译量化、Random-LTD、稀疏注意力算子DS_BUILD_TRANSFORMER/DS_BUILD_TRANSFORMER_INFERENCE/DS_BUILD_STOCHASTIC_TRANSFORMER编译 Transformer 训练/推理/随机 Transformer 算子两个常见安装/运行排障点CUDA kernel image 错误若运行时报RuntimeError: CUDA error: no kernel image is available for execution on the device说明 CUDA 扩展不是针对你所用 GPU 架构编译的。可设置TORCH_CUDA_ARCH_LIST按需指定架构重新构建例如TORCH_CUDA_ARCH_LIST6.1;7.5;8.6 pip install ...。CUDA 版本不匹配若报Installed CUDA version ... does not match the version torch was compiled with ...说明本机 CUDA 与编译 torch 所用的 CUDA 主版本不一致。可用nvcc --version与python3 -c import torch; print(torch.__version__)核对确需跳过检查可设DS_SKIP_CUDA_CHECK1不推荐可能引发未预期行为。二、用 DeepSpeed 编写训练脚本deepspeed.initializeDeepSpeed 模型训练由DeepSpeed 引擎engine完成。引擎可以包装任意torch.nn.Module模型并提供一套极简的训练与断点 API。初始化引擎的入口是deepspeed.initializeimport deepspeed model_engine, optimizer, _, _ deepspeed.initialize(argscmd_args, modelmodel, model_parametersparams)deepspeed.initialize会在底层替你完成分布式数据并行DDP或混合精度训练所需的全部初始化。除了包装模型DeepSpeed 还能根据传入deepspeed.initialize的参数以及 DeepSpeed 配置文件自动创建并管理优化器、数据加载器与学习率调度器。从源码看deepspeed/init.py 中的initialize返回一个四元组其完整签名为返回项含义满足何种条件时非Noneengine包装了用户模型的 DeepSpeed 运行时引擎分布式训练用恒有optimizer被包装的优化器用户在initialize传入优化器或 JSON 配置中声明了优化器training_dataloaderDeepSpeed 数据加载器传入了training_datalr_scheduler被包装的学习率调度器用户在initialize传入或 JSON 配置中声明了调度器也就是说上文示例中第 3、4 个返回值用_忽略即可。引擎会根据实际训练路径自动选择包装器普通模型对应DeepSpeedEnginePipeline 模型对应PipelineEngine启用 Hybrid EngineRLHF 场景时对应DeepSpeedHybridEngine。另外注意initialize要求必须提供模型与配置--deepspeed_config指定的 JSON 文件路径或以字典形式传入config参数二者缺一会直接断言报错。分布式环境初始化deepspeed.init_distributed()如果你已经自行搭建过分布式环境需要把原来这行torch.distributed.init_process_group(...)替换为deepspeed.init_distributed()默认使用 NCCL 后端——这是 DeepSpeed 经过充分测试的后端——你也可覆盖默认后端。如果你不需要在deepspeed.initialize()之前建立分布式环境则不必显式调用该函数因为 DeepSpeed 会在其initialize过程中自动完成分布式环境初始化。但无论如何如果你原本有torch.distributed.init_process_group调用都需要删除否则会发生重复初始化。分布式初始化在底层由 deepspeed/comm/comm.py 的init_distributed完成它会按加速器类型GPU/XPU/HPU/NPU 等选择对应的通信后端。如果你使用模型并行或流水线并行需要在deepspeed.initialize()之前完成某些torch.distributed调用同样用上面的deepspeed.init_distributed()替换即可。训练循环forward / backward / step引擎初始化后只需三个简单的 API 即可训练前向传播引擎对象可调用、反向传播backward与权重更新stepfor step, batch in enumerate(data_loader): # forward() 方法 loss model_engine(batch) # 反向传播 model_engine.backward(loss) # 权重更新 model_engine.step()注意这里传入model_engine(batch)的batch会作为位置参数直接传给被包装的模型即等价于model(*batch)因此多数教程会先batch tuple(batch)或用batch {k: v for ...}处理后调用。上文的枚举步长也可写成对数据集长度取整后的num_training_steps。在引擎内部deepspeed/runtime/engine.py这套调用会自动完成分布式数据并行、混合精度下所必需的操作1. 梯度平均Gradient Averaging在分布式数据并行训练中backward会确保每个train_batch_size训练步之后各数据并行进程间的梯度完成平均。其背后是梯度累加步数gradient accumulation steps与梯度的 AllReduce/Reduce-Scatter 通信。引擎维护一个微步micro-step计数器is_gradient_accumulation_boundary通过(micro_steps 1) % gradient_accumulation_steps 0判断当前是否为梯度累加边界只有到达边界时才会触发梯度归约与真正的优化器更新step。2. 损失缩放Loss Scaling在 FP16/混合精度训练中引擎自动处理损失缩放以避免梯度精度损失。backward内部对 ZeRO 优化器调用optimizer.scale_if_loss(loss)完成缩放使用 Apex AMP 时则通过amp.scale_loss(...)在梯度累加边界内延迟 unscaleengine.py 中 AMP 分支。引擎还会在backward中对 loss 除以gradient_accumulation_steps使累加后的梯度等价于对一个完整train_batch_size求平均。3. 学习率调度器Learning Rate Scheduler当使用 DeepSpeed 的学习率调度器在ds_config.json中指定时每次执行model_engine.step()DeepSpeed 都会自动调用调度器的step()方法当不使用 DeepSpeed 调度器时若调度器应在每个训练步更新可将调度器对象传给deepspeed.initialize由 DeepSpeed 代为更新与保存/恢复若调度器应在其他周期如每个 epoch更新则不应在初始化时传给 DeepSpeed而必须在训练循环中自行显式管理。断点保存与恢复Model Checkpointing训练状态的保存与加载由save_checkpoint与load_checkpointAPI 完成两者通过两个参数唯一定位一个断点ckpt_dir源码中为save_dir/load_dir断点保存目录ckpt_id源码中为tag目录内唯一标识该断点的标签。下述示例用 loss 值作为断点标识。完整示例包含断点恢复、推进 dataloader 到对应步、周期性保存并携带客户端状态# 加载断点 _, client_sd model_engine.load_checkpoint(args.load_dir, args.ckpt_id) step client_sd[step] # 将 data loader 推进到断点对应的 step dataloader_to_step(data_loader, step 1) for step, batch in enumerate(data_loader): # forward() 方法 loss model_engine(batch) # 反向传播 model_engine.backward(loss) # 权重更新 model_engine.step() # 保存断点 if step % args.save_interval 0: client_sd[step] step ckpt_id loss.item() model_engine.save_checkpoint(args.save_dir, ckpt_id, client_sdclient_sd)DeepSpeed 会自动保存并恢复模型、优化器与学习率调度器的状态将这些细节对用户隐藏。但对于某个训练任务特有的额外数据例如上面示例中的step你可以通过save_checkpoint的client state 字典client_sd一并保存并通过load_checkpoint的返回值取回。从源码签名看load_checkpoint返回(load_path, client_state)二元组若未指定tag会先读取load_dir下的latest文件获取最新断点标签。此外 ZeRO-3 支持矩阵 提示在 ZeRO Stage 3 下不能在engine.save_checkpoint()之后立刻engine.load_checkpoint()因为此时engine.module已被切分加载需要一个未被切分的“原始”模型确有需求应先重新初始化引擎。重要提醒save_checkpoint/load_checkpoint必须由所有进程调用而不是只在 rank 0 进程调用。因为每个进程都需要保存自己的 master weights 与调度器优化器状态如果只在一个进程中调用该调用会一直挂起等待与其他进程同步。五、DeepSpeed 配置文件ds_config.jsonDeepSpeed 的各类特性可通过一个 JSON 配置文件启用、禁用或调整该文件通过args.deepspeed_config命令行--deepspeed_config指定。入门示例配置如下{ train_batch_size: 8, gradient_accumulation_steps: 1, optimizer: { type: Adam, params: { lr: 0.00015 } }, fp16: { enabled: true }, zero_optimization: true }下面结合完整配置参数文档对核心字段作进一步说明。5.1 批次大小相关参数参数说明默认值train_batch_size有效训练批次大小即触发一次模型更新所需的数据样本量可省略train_micro_batch_size_per_gpu单张 GPU 单次前向/反向处理的批大小不做梯度累加取train_batch_sizegradient_accumulation_steps在平均并应用梯度前累积的训练步数1managed_gradient_accumulation是否由 DeepSpeed 内部管理梯度累加边界默认true设为false时由客户端在每个累加边界自行调用step()true三者必须满足恒等关系train_batch_size train_micro_batch_size_per_gpu × gradient_accumulation_steps × GPU 数量为简化起见可只显式指定三者中的任意两个最后一个由 DeepSpeed 自动推算。增加gradient_accumulation_steps可以减少梯度通信频率有利于扩展性也支持在单卡上使用更大有效批次。5.2 优化器与调度器optimizer是一个字典包含type与params两个字段。DeepSpeed 原生支持Adam、AdamW、OneBitAdam、Lamb、OneBitLamb、Muon等优化器也可以导入 torch 自带优化器params中的参数名必须与对应优化器构造函数签名一致如{lr: 0.001, eps: 1e-8}。例如带 betas/eps/weight_decay 的 Adam 配置optimizer: { type: Adam, params: { lr: 0.001, betas: [0.8, 0.999], eps: 1e-8, weight_decay: 3e-7 } }optimizer.params还支持两个扩展键torch_adam使用 torch 的 Adam 实现而非 DeepSpeed 融合 Adam默认false与adam_w_mode使用 L2 正则即 AdamW 语义默认true。lr_scheduler/scheduler同理type指定调度器名称如WarmupLRparams对应构造参数scheduler: { type: WarmupLR, params: { warmup_min_lr: 0, warmup_max_lr: 0.001, warmup_num_steps: 1000 } }5.3 混合精度fp16/bf16/amp入门示例中fp16: { enabled: true }只是开启 FP16 训练的最小配置。完整可配字段包括fp16: { enabled: true, auto_cast: false, loss_scale: 0, initial_scale_power: 16, loss_scale_window: 1000, hysteresis: 2, consecutive_hysteresis: false, min_loss_scale: 1 }字段说明默认值enabled是否启用 FP16 训练falseloss_scale损失缩放值0表示动态损失缩放否则为静态固定缩放0initial_scale_power动态损失缩放的初始值取2^initial_scale_power16loss_scale_window动态缩放升/降的窗口长度1000hysteresis动态缩放中的延迟位移2min_loss_scale动态损失缩放的最小值1使用 ZeRO 时必须配合 FP16或 BF16模式FP16 模式不能与 Apex 的amp模式混用。BF16如 A100 等硬件支持可配置为bf16: { enabled: true, bf16_master_weights_and_grads: true, bf16_optimizer_states: true }amp模式Apex AMP如opt_level: O1则不能与fp16模式或 ZeRO 同时使用。5.4 ZeRO 优化zero_optimization入门示例的zero_optimization: true是启用 ZeRO 记忆优化的历史写法完整形态是字典并显式声明 stagezero_optimization: { stage: 2, allgather_partitions: true, allgather_bucket_size: 5e8, overlap_comm: false, reduce_scatter: true, reduce_bucket_size: 5e8, contiguous_gradients: true }字段说明默认值stageZeRO 阶段0/1/2/3 分别表示关闭、优化器状态切分、优化器梯度切分、优化器梯度参数切分0allgather_partitions每步结束后用 allgather 集合通信而非一串 broadcast聚合更新后的参数trueallgather_bucket_size每次 allgather 的元素数量限制大模型下的 allgather 显存占用5e8reduce_scatter用 reduce-scatter而非 allreduce做梯度平均truereduce_bucket_size每次归约的元素数量限制显存占用5e8contiguous_gradients反向时把梯度拷贝进连续缓冲避免显存碎片trueoverlap_comm让梯度归约与反向计算重叠falseoffload_param参数卸载到 CPU/NVMe仅 Stage 3 有效falseoffload_optimizer优化器状态卸载到 CPU/NVMeStage 1/2/3 有效falsestage3_max_live_parameters每卡驻留的最大参数数量调小省显存但增加通信1e9stage3_max_reuse_distance距下次重用距离小于该阈值时不释放参数1e9stage3_prefetch_bucket_size参数预取固定缓冲大小5e8stage3_param_persistence_threshold小于该阈值的参数不做切分1e5stage3_gather_16bit_weights_on_model_save保存save_16bit_model()前自动聚合被切分的权重falseround_robin_gradientsStage 1/2 的 CPU 卸载优化按 rank 细粒度切分并行拷贝梯度到 CPUfalse完整字段与取值说明含参数/优化器卸载、NVMe 异步 I/O、AutoTP、AutoEP 等请查阅配置 JSON 参考文档与 ZeRO 教程。六、启动 DeepSpeed 训练DeepSpeed 安装后会注册可执行入口deepspeed用于启动分布式训练。假设你已经满足① 模型已集成 DeepSpeed②client_entry.py是你的模型入口脚本③client args是 argparse 命令行参数④ds_config.json是 DeepSpeed 配置文件则启动命令形如deepspeed client_entry.py client args \ --deepspeed --deepspeed_config ds_config.json其中--deepspeed与--deepspeed_config这两个核心参数在 deepspeed/init.py 中注册另有两个已废弃的--deepscale/--deepscale_config别名。6.1 多节点资源配置hostfileDeepSpeed 使用与 OpenMPI、Horovod 兼容的hostfile配置多节点计算资源。hostfile 是若干hostname或 SSH 别名需可通过免密 SSH 访问及slot 数该机器上可用的 GPU 数组成的列表例如worker-1 slots4 worker-2 slots4表示名为worker-1与worker-2的两台机器各有 4 张 GPU 用于训练。hostfile 通过命令行选项--hostfile指定。若未指定DeepSpeed 会查找/job/hostfile该默认路径定义于 deepspeed/launcher/runner.py若既未指定也未找到DeepSpeed 则查询本机 GPU 数量来推算本地 slot 数。在myhostfile中列出的全部节点和 GPU 上启动训练deepspeed --hostfilemyhostfile client_entry.py client args \ --deepspeed --deepspeed_config ds_config.json6.2 限定节点与 GPU 资源DeepSpeed 允许把分布式训练限制在部分节点/GPU 上通过--num_nodes与--num_gpus两个参数实现。例如只使用其中两个节点deepspeed --num_nodes2 \ client_entry.py client args \ --deepspeed --deepspeed_config ds_config.json也可通过--include/--exclude精确包含或排除指定资源。例如排除worker-2上的 GPU 0以及worker-3上的 GPU 0 和 1deepspeed --excludeworker-2:0worker-3:0,1 \ client_entry.py client args \ --deepspeed --deepspeed_config ds_config.json类似地仅使用worker-2上的 GPU 0 和 1deepspeed --includeworker-2:0,1 \ client_entry.py client args \ --deepspeed --deepspeed_config ds_config.json这些参数都由启动器在 deepspeed/launcher/runner.py 中解析-H/--hostfile、-i/--include、-e/--exclude、--num_nodes、--num_gpus、--node_rank、--master_addr、--master_port、--no_ssh等均在add_argument列表中。6.3 无需免密 SSH 的启动模式DeepSpeed 支持无需免密 SSH 地启动训练任务该模式在 Kubernetes 等云环境中尤为实用——容器化调度下再搭建 leader-worker 免密 SSH 架构会徒增复杂度。用法是在所有节点上分别运行 DeepSpeed 命令deepspeed --hostfilemyhostfile --no_ssh --node_rankn \ --master_addraddr --master_portport \ client_entry.py client args \ --deepspeed --deepspeed_config ds_config.json--hostfilemyhostfile指定包含节点与 GPU 信息的 hostfile--no_ssh开启无 SSH 模式--node_rankn节点 rank应为 0 到 n-1 之间的唯一整数--master_addraddrleader 节点rank 0的地址--master_portportleader 节点的端口。在该模式下hostfile 中的主机名无需可通过免密 SSH 访问但 hostfile 仍然必须提供启动器要据此收集环境信息节点数与每节点 GPU 数。每个节点必须以唯一的node_rank启动且所有节点都要获知 leader 节点的地址与端口。此模式下启动器的行为与 PyTorch 的torchrun启动器类似。6.4 多节点环境变量.deepspeed_env跨节点训练时DeepSpeed 支持传播用户自定义的环境变量。默认情况下 DeepSpeed 会传播所有已设置的 NCCL 与 PYTHON 相关环境变量如需传播额外变量可在名为.deepspeed_env的点文件中写若干VARVAL条目每行一条。该文件会从当前执行目录与主目录两处查找若想覆盖默认文件名/路径可通过环境变量DS_ENV_FILE指定该默认名即DS_ENV_FILE的值见 deepspeed/launcher/runner.py多任务需要不同变量时尤为有用。例如某些集群需要在训练前设置特殊的 NCCL 变量在主目录的.deepspeed_env中写入NCCL_IB_DISABLE1 NCCL_SOCKET_IFNAMEeth0DeepSpeed 会确保每个节点上的每个进程启动时都已设置这些环境变量。6.5 MPI 与 AzureML 兼容DeepSpeed 提供了自己的并行启动器来启动多节点/多卡训练。如果你更倾向用 MPI如 mpirun启动DeepSpeed 也提供支持——需要说明的是DeepSpeed 底层仍使用 torch 分布式 NCCL 后端而非 MPI 后端。若用 mpirun DeepSpeed 或借助 AzureML其以 mpirun 作为启动后端启动任务只需安装mpi4pyPython 包DeepSpeed 会用它来探测 MPI 环境并把 world size、rank 等必要状态传给 torch 分布式后端。若使用模型并行、流水线并行或需要在deepspeed.initialize()之前调用torch.distributed同样把最初的torch.distributed.init_process_group(..)换成deepspeed.init_distributed()即可获得等价支持。6.6 单节点资源配置如果只在单节点单卡或多卡上运行DeepSpeed不需要hostfile——未检测到或未传入 hostfile 时DeepSpeed 会查询本机 GPU 数来发现可用 slot 数。--include/--exclude参数照常工作但主机名应写localhostdeepspeed --include localhost:0,1 ...此外CUDA_VISIBLE_DEVICES也可与deepspeed配合用于控制单节点上使用的设备。以下两种方式等价都只在当前节点的设备 0 和 1 上启动deepspeed --include localhost:0,1 ...CUDA_VISIBLE_DEVICES0,1 deepspeed ...七、小结至此你已掌握 DeepSpeed 从零启动训练的最小闭环pip install deepspeed完成安装并用ds_report验证通过deepspeed.initialize包装任意torch.nn.Module必要时用deepspeed.init_distributed替换分布式初始化用forward/backward/step三段式 API 编写训练循环并用save_checkpoint/load_checkpoint做断点管理通过ds_config.json声明批次参数、优化器、FP16/BF16 与 ZeRO 配置最后用deepspeed启动器配合 hostfile、--include/--exclude、--no_ssh或.deepspeed_env在单机或多机上拉起任务。想进一步深入可继续阅读仓库内相关教程ZeRO 优化详解、安装细节、多加速器适配指南以及完整配置参数参考对配置项在引擎中的实际解析行为可直接阅读 DeepSpeedConfig 与引擎核心实现以加深理解。【免费下载链接】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），仅供参考