Transformers 中 DeepSpeed ZeRO 集成实战从配置到 ZeRO-3 分片训练【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers本篇基于 Transformers 官方文档 deepspeed.md 展开系统讲解如何通过TrainingArguments的deepspeed参数或 Accelerate 配置文件接入 DeepSpeed ZeRO覆盖 ZeRO-1/2/3 三阶段分片配置、auto占位符机制、offload 与优化器/调度器同步、以及 DeepSpeed 分片 checkpoint 的正确保存方式。读完后你能够在多卡环境下为Trainer配置一套可直接运行的 ZeRO 训练方案并理解其背后的源码级调用链。DeepSpeed ZeRO 是什么ZeROZero Redundancy Optimizer通过把优化器状态、梯度和参数分片到多张 GPU消除分布式训练中每个 GPU 各自持有完整副本造成的显存冗余。ZeRO 分为三个阶段阶段越高分片越多params grads opt states ┌──────────┐ ┌──────────┐ ┌──────────┐ ZeRO-1 │██████████│ │██████████│ │███░░░░░░░│ GPU 0 │██████████│ │██████████│ │░░░███░░░░│ GPU 1 │██████████│ │██████████│ │░░░░░░████│ GPU 2 └──────────┘ └──────────┘ └──────────┘ ┌──────────┐ ┌──────────┐ ┌──────────┐ ZeRO-2 │██████████│ │███░░░░░░░│ │███░░░░░░░│ GPU 0 │██████████│ │░░░███░░░│ │░░░███░░░░│ GPU 1 │██████████│ │░░░░░░████│ │░░░░░░████│ GPU 2 └──────────┘ └──────────┘ └──────────┘ ┌──────────┐ ┌──────────┐ ┌──────────┐ ZeRO-3 │███░░░░░░░│ │███░░░░░░░│ │███░░░░░░░│ GPU 0 │░░░███░░░░│ │░░░███░░░░│ │░░░███░░░░│ GPU 1 │░░░░░░████│ │░░░░░░████│ │░░░░░░████│ GPU 2 └──────────┘ └──────────┘ └──────────┘ █ resident ░ held on another GPUZeRO-1参数与梯度每卡保留完整副本只分片优化器状态ZeRO-2在 ZeRO-1 基础上再分片梯度通信开销仍低于 ZeRO-3ZeRO-3参数、梯度、优化器状态全部分片是模型在 ZeRO-2 下无法放下时的选择。官方文档给出的经验法则是优先用 ZeRO-2只有当模型在 ZeRO-2 下仍放不进各卡显存时才上 ZeRO-3因为每提升一个阶段都会增加 GPU 间通信。此外 DeepSpeed 还支持把优化器状态甚至参数 offload 到 CPU 或 NVMe进一步压缩显存代价是更长的传输时间。安装从 PyPI 安装 DeepSpeed或者直接安装带deepspeedextra 的 Transformerspip install deepspeed pip install transformers[deepspeed]在 setup.py 中可以确认deepspeedextra 的构成extras[deepspeed] deps_list(deepspeed, accelerate)且依赖表中要求deepspeed0.9.3见 setup.py。也就是说启用 DeepSpeed 集成的前提是同时具备deepspeed与accelerate两个依赖Transformers 的 DeepSpeed 胶水层正是构建在 Accelerate 的DeepSpeedPlugin之上的。如果安装时遇到 CUDA 相关报错可参考 调试文档中的 DeepSpeed CUDA 章节从源码编译安装而非用 PyPI 预编译包通常更稳妥因为它会匹配你实际的 GPU 架构。配置方式一TrainingArguments 的 deepspeed 参数TrainingArguments定义了deepspeed字段# --- DeepSpeed --- deepspeed: dict | str | None field( defaultNone, metadata{help: Enable DeepSpeed integration. Value is a path to a JSON config file or a dict.}, )它接受一个 DeepSpeed JSON 配置文件路径字符串或已加载的字典。在TrainingArguments的 docstringtraining_args.py中也有明确提示若使用 ZeRO 初始化必须先初始化TrainingArguments再实例化模型否则 ZeRO 不会被应用。用 auto 让 DeepSpeed 从 TrainingArguments 取默认值DeepSpeed JSON 中凡是希望由 Transformers 代填的字段都可以写成字符串autotrain_micro_batch_size_per_gpu: auto, // ← per_device_train_batch_size in TrainingArguments gradient_accumulation_steps: auto, // ← gradient_accumulation_steps in TrainingArguments optimizer.params.lr: auto, // ← learning_rate in TrainingArguments fp16.enabled: auto, // ← fp16 flag in TrainingArguments其底层实现是 HfTrainerDeepSpeedConfig 类中的fill_match方法auto会被直接替换为TrainingArguments中对应的值如果某个字段显式写了具体数值则会与TrainingArguments的值做一致性校验不一致时记录到mismatches最终在校验阶段报错。源码中的映射关系deepspeed.py 的trainer_config_process包括DeepSpeed JSON 字段对应 TrainingArguments 字段train_micro_batch_size_per_gpuper_device_train_batch_sizetrain_batch_size计算值world_size × per_device_train_batch_size × gradient_accumulation_stepsgradient_accumulation_stepsgradient_accumulation_stepsgradient_clippingmax_grad_normoptimizer.params.lr / betas / eps / weight_decaylearning_rate/adam_beta1adam_beta2/adam_epsilon/weight_decayscheduler.params.warmup_max_lrlearning_ratefp16.enabled/bf16.enabledfp16/bf16或对应full_eval标志关键结论如果你在 JSON 中手动写死了 batch size、梯度累积步数等值且与TrainingArguments不一致训练会静默地以错误值继续文档原文如此警告。因此这些字段一律建议写auto。拿到配置后通过deepspeed参数传入并选择启动方式from transformers import TrainingArguments args TrainingArguments( deepspeedpath/to/deepspeed_config.json, ... )# DeepSpeed launcher deepspeed --num_gpus 4 train.py # torchrun torchrun --nproc_per_node 4 train.py # Accelerate accelerate launch --num_processes 4 train.py从源码结构看TrainingArguments初始化末尾training_args.py会按如下链路装配 DeepSpeed先创建HfTrainerDeepSpeedConfig(self.deepspeed)并调用trainer_config_process(self)完成auto填充再用它构造 Accelerate 的DeepSpeedPlugin存入self.deepspeed_plugin。此外还有一个兜底路径即使没传deepspeed参数只要环境变量ACCELERATE_USE_DEEPSPEEDtrue例如用 Accelerate 启动器启动同样会创建DeepSpeedPlugin并同步混合精度设置。Trainer初始化时trainer.py则通过accelerator.state.deepspeed_plugin是否存在判定is_deepspeed_enabled并在未传deepspeed参数时调用propagate_args_to_deepspeed把训练参数补推进插件。与 FSDP 互斥Trainer中有明确校验--fsdp与--deepspeed不能同时使用trainer.py 会抛出Using --fsdp xxx together with --deepspeed is not possible。配置方式二Accelerate 配置文件也可以绕过TrainingArguments.deepspeed改用 Accelerate 的配置文件。先运行accelerate config交互式生成default_config.yaml然后配置distributed_type: DEEPSPEED deepspeed_config: deepspeed_config_file: path/to/ds_config.json machine_rank: 0 num_machines: 1 num_processes: 4再用accelerate launch启动基于Trainer的脚本accelerate launch --config_file deepspeed_config.yaml train.py文档特别强调使用 Accelerate 启动时TrainingArguments中的deepspeed参数会被忽略DeepSpeed 配置完全来自 Accelerate 配置文件。更多用法可参考 Accelerate 集成文档。ZeRO 阶段配置示例官方文档为三个阶段各提供了一份可直接作为起点的配置。ZeRO-1{ bf16: { enabled: auto }, zero_optimization: { stage: 1 }, gradient_clipping: auto, train_micro_batch_size_per_gpu: auto, train_batch_size: auto, gradient_accumulation_steps: auto }ZeRO-2{ bf16: { enabled: auto }, zero_optimization: { stage: 2, overlap_comm: true, allgather_bucket_size: 2e8, reduce_bucket_size: 2e8, contiguous_gradients: true }, gradient_clipping: auto, train_micro_batch_size_per_gpu: auto, train_batch_size: auto, gradient_accumulation_steps: auto }ZeRO-3注意ZeRO-3 在初始化阶段就对参数分片。你必须先实例化TrainingArguments再加载模型——如果模型在 DeepSpeed 配置完成前已经加载到每张 GPU 上就不会节省任何显存。{ bf16: { enabled: auto }, zero_optimization: { stage: 3, overlap_comm: true, contiguous_gradients: true, reduce_bucket_size: auto, stage3_prefetch_bucket_size: auto, stage3_param_persistence_threshold: auto, stage3_gather_16bit_weights_on_model_save: true, offload_optimizer: { device: cpu, pin_memory: true }, // optional offloading offload_param: { device: cpu, pin_memory: true } // optional offloading }, gradient_clipping: auto, train_micro_batch_size_per_gpu: auto, train_batch_size: auto, gradient_accumulation_steps: auto }关键字段详解以下字段对自定义训练最为重要均出自 deepspeed.md 原文并逐项展开zero_optimization设置 ZeRO 阶段。{ zero_optimization: { stage: 3 } }batch size 与梯度累积train_micro_batch_size_per_gpu、train_batch_size、gradient_accumulation_steps一律设为auto。如前所述手动设置且与TrainingArguments不一致时训练会以错误值静默继续。{ train_micro_batch_size_per_gpu: auto, train_batch_size: auto, gradient_accumulation_steps: auto, gradient_clipping: auto }bf16控制训练精度。设为auto后与TrainingArguments的bf16标志联动。从源码看trainer_config_process会依据bf16.enabled/fp16.enabled的最终取值把训练 dtype 定为torch.bfloat16、torch.float16或torch.float32deepspeed.py后续Trainer._prepare_input移动输入张量到设备时会以该 dtype 为准trainer.py。{ bf16: { enabled: auto } }stage3_gather_16bit_weights_on_model_saveZeRO-3 专属。保存前在所有 GPU 间做 all-gather把分片还原成完整张量。这是保存完整 16 位模型权重含 16 位 state dict的必要开关。{ zero_optimization: { stage: 3, stage3_gather_16bit_weights_on_model_save: true, } }overlap_comm/allgather_bucket_size/reduce_bucket_sizeoverlap_comm: true让 all-reduce 通信与反向传播重叠隐藏通信延迟两个 bucket 参数在通信速度与显存之间做权衡——取值越小显存越省但通信越慢。{ zero_optimization: { stage: 2, overlap_comm: true, allgather_bucket_size: 2e8, reduce_bucket_size: 2e8, contiguous_gradients: true } }offload_optimizer/offload_paramoffload_optimizer把优化器状态挪到 CPU 内存offload_param进一步把模型参数也 offload仅 ZeRO-3。pin_memory: true能加速 CPU-GPU 传输但代价是锁定的 RAM 无法被其他进程使用。{ zero_optimization: { stage: 3, offload_optimizer: { device: cpu, pin_memory: true }, offload_param: { device: cpu, pin_memory: true } } }optimizer与scheduler默认继承TrainingArguments中配置的优化器与调度器。参数写auto即可让 DeepSpeed 从TrainingArguments读取除非你需要 DeepSpeed 原生优化器如 LAMB。{ optimizer: { type: AdamW, params: { lr: auto, betas: auto, eps: auto, weight_decay: auto } }, scheduler: { type: WarmupDecayLR, params: { total_num_steps: auto, warmup_min_lr: auto, warmup_max_lr: auto, warmup_num_steps: auto } } }若启用了优化器 offload可把zero_force_ds_cpu_optimizer设为false以启用 DeepSpeed 的 CPU Adam 优化器{ zero_force_ds_cpu_optimizer: false }从源码结构看优化器与调度器真正被 DeepSpeed 接管发生在Trainer的_prepare_for_training阶段is_deepspeed_enabled时调用 deepspeed_init 生成optimizer, lr_scheduler之后再经accelerator.prepare包裹模型纯推理场景如evaluate前未训练也会调用deepspeed_init(..., inferenceTrue)且此时只有 ZeRO-3 有意义源码会直接报错提示调整配置。Checkpoint分片保存与权重还原DeepSpeed 的 checkpoint 采用分片格式无法直接用from_pretrained加载。推荐的完整工作流是设置load_best_model_at_endTrue让Trainer在训练结束时追踪并重新加载最佳 checkpointTrainer内部通过 deepspeed_load_checkpoint 从global_step*目录恢复模型、优化器与调度器状态训练结束后调用trainer.save_model()导出一个标准 Transformers checkpointfrom transformers import TrainingArguments, Trainer args TrainingArguments( deepspeedds_config_zero3.json, load_best_model_at_endTrue, ... ) # after training, save a normal transformers checkpoint trainer.save_model(./best-model)需要注意两个陷阱save_only_modelTrue与load_best_model_at_endTrue不兼容跳过完整优化器状态保存意味着训练结束时无法重新加载最佳模型Trainer会直接抛出ValueErrortrainer.py 中的校验逻辑DeepSpeed cant be used with save_only_model along with load_best_model_at_end。ZeRO-3 下必须设置stage3_gather_16bit_weights_on_model_save: true才能从分片重建完整权重。Trainer._save在 DeepSpeed 分支中trainer.py会先用accelerator.get_state_dict(self.deepspeed)拿到完整 state dict 再走常规保存路径PEFT ZeRO-3 场景下则优先调用model_wrapped._zero3_consolidated_16bit_state_dict(exclude_frozen_parametersTrue)避免拉取冻结参数。如果你需要跨不同并行配置恢复训练例如换 GPU 数继续训需要使用 DeepSpeed 的 Universal Checkpointing 功能可在 DeepSpeed 官方教程中查阅。小结Trainer通过TrainingArguments.deepspeed参数接入 DeepSpeed本质是HfTrainerDeepSpeedConfig AccelerateDeepSpeedPlugin的组合integrations/deepspeed.pyauto占位符机制保证了 DeepSpeed 配置与TrainingArguments的单一数据源优先 ZeRO-2模型放不下再上 ZeRO-3offload 是最后的显存手段所有可同步的字段batch size、梯度累积、学习率、精度、梯度裁剪写auto避免两处配置不一致导致的静默错误ZeRO-3 必须先生成TrainingArguments再加载模型保存权重时开启stage3_gather_16bit_weights_on_model_save配合load_best_model_at_end与训练结束后的save_model得到标准可加载产物。想进一步理解 ZeRO 算法本身可以阅读 ZeRO 系列论文Memory Optimizations Toward Training Trillion Parameter Models、Democratizing Billion-Scale Model Training、Breaking the GPU Memory Wall for Extreme Scale Deep Learning以及 HuggingFace 的 The Ultra-Scale Playbook 中关于 Zero Redundancy Optimizer 的章节。【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考