示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载导读本文以 DeepSpeedExamples 仓库中的 ZenFlow Llama-2 微调示例 为骨架系统讲解如何利用ZenFlow——一种面向大规模模型训练的无阻塞stall-free卸载引擎——在 CPU 卸载offloading场景下微调 Llama-2 模型。读完本文你将掌握 ZenFlow 的核心机制异步更新 Top-k 选择性参数更新、zf_config.json 中每一个关键配置项的含义与调优思路、Alpaca-52K 微调脚本的完整运行流程以及如何用基准测试脚本定量评估不同配置组合的加速效果。一、ZenFlow 是什么为什么 CPU 卸载会成为瓶颈在单卡或多卡显存受限的环境中训练大模型通常依赖 ZeRO-Offload 将优化器状态optimizer state卸载到 CPU 内存用 CPU 承担部分 Adam 更新计算。这种做法大幅降低了 GPU 显存压力却引入了新的问题CPU 上的优化器步骤optimizer step与 GPU 上的前向/反向计算天然是串行的GPU 不得不长时间空闲等待 CPU 完成参数更新训练出现明显的停顿stall。ZenFlow论文《ZenFlow: Enabling Stall-Free Offloading Training via Asynchronous Updates》的核心思路是异步重叠不再等待完整的优化器步骤结束后才进入下一轮训练而是把一次完整的参数更新拆解将其中一部分工作异步地偷偷安排在多个梯度累积步骤的间隙里执行Top-k 选择性更新在更新周期内只把最重要的一部分梯度按 topk_ratio 比例筛选立即更新到 GPU 参数上让训练进度持续向前推进而不是等所有参数一起更新完毕。本仓库的微调示例正是这一思路的落地实现配置了 ZenFlow 的 DeepSpeed 会在日志中打印ZenFlowCPUAdam initialized with overlap step.表明 CPU Adam 已进入重叠overlap模式。二、环境准备与依赖安装示例的依赖清单见 requirements.txt安装方式如下pip install -r requirements.txt各依赖的最低版本要求为依赖包版本要求用途torch 2.5.1深度学习框架deepspeed 0.16.0分布式训练引擎内置 ZenFlow 支持datasets 2.14.1加载 Alpaca-52K 数据集transformers 4.37.2加载 Llama-2 模型与分词器numpy 1.21.0数值计算需要注意的是ZenFlow 的zenflow配置块是 DeepSpeed 0.16.0 及以上版本引入的能力如果本地 DeepSpeed 版本过低zf_config.json中的zenflow字段将无法被解析示例会在启动时直接报错。三、配置文件精讲zf_config.json 的每一个字段ZenFlow 微调的核心配置集中在 zf_config.json该文件同时承担了 DeepSpeed 训练超参与 ZenFlow 卸载策略的双重职责。完整内容如下{ train_batch_size: 8, bf16: { enabled: true }, zero_optimization: { stage: 2, offload_optimizer: { device: cpu, pin_memory: true }, zenflow: { topk_ratio: 0.1, update_interval: 4, full_warm_up_rounds: 0, overlap_step: true } }, optimizer: { type: AdamW, params: { lr: 2e-5, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0.01 } }, gradient_accumulation_steps: 1, gradient_clipping: 1.0, zero_allow_untested_optimizer: true }下面逐项拆解其含义3.1 基础训练配置train_batch_size: 8全局训练 batch size。训练脚本 finetune_llama.py 中 DataLoader 的batch_size1只是占位实际批大小由 DeepSpeed 依据该配置自动改写。bf16.enabled: true启用 bfloat16 混合精度训练。脚本中模型加载也显式指定了torch_dtypetorch.bfloat16见 finetune_llama.py与配置保持一致。BF16 相比 FP16 拥有与 FP32 相同的指数位宽在大模型微调中更稳定。gradient_clipping: 1.0梯度裁剪阈值。zero_allow_untested_optimizer: true允许 DeepSpeed 使用未列入官方已验证清单的优化器这里为 AdamW否则初始化会因校验失败而中止。3.2 优化器配置optimizer: { type: AdamW, params: { lr: 2e-5, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0.01 } }lr: 2e-5学习率也是微调 Llama-2 的典型取值。betas: [0.9, 0.999]Adam 的一阶/二阶动量衰减系数。weight_decay: 0.01权重衰减。需要说明的是Shell 脚本 finetune_llama.sh 中同样定义了LR、BATCH_SIZE、WARMUP、WEIGHT_DECAY等变量但脚本注释明确指出这些参数仅作回退fallback之用实际生效的是配置文件中的值。3.3 卸载与 ZenFlow 配置本文核心zero_optimization: { stage: 2, offload_optimizer: { device: cpu, pin_memory: true }, zenflow: { topk_ratio: 0.1, update_interval: 4, full_warm_up_rounds: 0, overlap_step: true } }stage: 2ZeRO 第二阶段将优化器状态分片并卸载。当前示例在 2 卡GPUS_PER_NODE2上运行每张卡只保存部分优化器状态。offload_optimizer.device: cpu优化器状态卸载到 CPU 内存这是 ZenFlow 发挥作用的前提场景。offload_optimizer.pin_memory: true使用锁页pinned内存可显著降低 CPU↔GPU 数据传输延迟。基准测试 README 的分析也印证了这一点pin_memoryTrue通常能减少传输延迟。zenflow.topk_ratio: 0.1每个更新周期内仅选择占比 10% 的最重要参数进行立即更新。这个比例直接影响部分更新的粒度比例越小单次立即更新的成本越低但剩余参数的完整更新周期越长。zenflow.update_interval: 4每 4 个训练步执行一次完整的剩余参数更新。update_interval4意味着在 4 个步中前 3 步只做后台重叠的部分更新步时约 700ms第 4 步完成剩余参数的完整优化器步骤并更新回 GPU步时约 2 秒。zenflow.full_warm_up_rounds: 0预热轮数。设为 0 表示不进行预热立即进入重叠模式若设为 N则前 N 轮执行完整更新以积累稳定的梯度统计之后再切换为 Top-k 选择性更新。zenflow.overlap_step: true开启优化器步骤与前向/反向计算的重叠。日志中ZenFlowCPUAdam initialized with overlap step.正是该开关生效的标志。调优建议从基准测试的结论看update_interval越大卸载频率越低、累积开销越小overlap_steptrue在模型更新阶段较长时能更充分地隐藏卸载延迟。三者topk_ratio、update_interval、overlap_step与pin_memory共同构成可搜索的配置空间建议用第四节介绍的命令行基准工具做系统性扫描。四、微调脚本剖析从数据预处理到 checkpoint 保存4.1 训练主流程 finetune_llama.py脚本整体流程如下设置随机种子set_seed(args.seed)默认 42保证可复现性见 finetune_llama.py。加载模型与分词器AutoTokenizer.from_pretrained/AutoModelForCausalLM.from_pretrained模型以 bfloat16 加载。若分词器无 pad token则自动用 eos token 补齐见 finetune_llama.py。加载并预处理 Alpaca-52K使用tatsu-lab/alpaca数据集preprocess_alpaca函数将 instruction/input/output 组装为### Instruction:...### Input:...### Response:...的提示词模板并做max_length512的截断与填充labels直接复用input_ids见 finetune_llama.py。DeepSpeed 初始化deepspeed.initialize(argsargs, modelmodel, model_parameters..., training_data..., collate_fn...)自动解析通过--deepspeed_config传入的配置文件返回model_engine / optimizer / train_dataloader / lr_scheduler见 finetune_llama.py。训练循环每个 step 执行model_engine(**batch)得到 loss随后model_engine.backward(loss)与model_engine.step()并打印Step N, Loss: xxxx, Time: xxxms见 finetune_llama.py。注意脚本的打印注释为 Print every 10 steps但实际实现是每个 step 都打印rank 0。保存仅 rank 0 通过model_engine.save_checkpoint(args.output_dir)保存 DeepSpeed checkpoint并调用tokenizer.save_pretrained保存分词器见 finetune_llama.py。4.2 启动脚本 finetune_llama.sh脚本的默认运行参数为参数默认值说明模型meta-llama/Llama-2-7b-hf可通过修改MODEL_NAME更换GPU 数GPUS_PER_NODE2NNODES1单机双卡输出目录./alpaca_outputmkdir -p自动创建Epochs3EPOCHS种子42SEED实际启动命令等价于deepspeed --num_gpus2 finetune_llama.py \ --deepspeed_config./zf_config.json \ --model_name meta-llama/Llama-2-7b-hf \ --num_train_epochs 3 \ --lr 2e-5 \ --batch_size 8 \ --weight_decay 0.01 \ --output_dir ./alpaca_output \ --seed 42需要特别强调两点脚本先检查zf_config.json是否存在不存在则直接报错退出见 finetune_llama.sh避免配置缺失导致静默错误换模型、改 batch size 或 epochs直接修改脚本头部变量即可所有 DeepSpeed / ZenFlow 选项仍统一收敛在zf_config.json。五、运行结果解读如何从日志看出无阻塞效果原 README 给出的示例日志是理解 ZenFlow 行为的最佳教材ZenFlowCPUAdam initialized with overlap step. Step 5, Loss: 1.2599, Time: 719.58ms Step 6, Loss: 0.9847, Time: 702.81ms -- gradient accumulation with overlapped update Step 7, Loss: 0.6220, Time: 705.50ms Step 8, Loss: 0.5173, Time: 1912.92ms -- full optimizer step of remaining part and update parameters Step 9, Loss: 0.4557, Time: 890.60ms Step 10, Loss: 0.3882, Time: 740.11ms Step 11, Loss: 0.3627, Time: 731.95ms Step 12, Loss: 0.3341, Time: 2221.18ms Step 13, Loss: 0.2453, Time: 1061.80ms解读要点Step 5/6/7 是累积步ZenFlow 在后台异步重叠执行了部分优化器步骤因此这些步依然很快约 700msloss 也在持续下降——这正是训练进度不被阻塞的直接证据。Step 8 是完整更新步执行剩余部分的优化器步骤并将参数同步回 GPU耗时约 1.9~2.2 秒每update_interval4出现一次。对比收益若无 ZenFlow一次完整更新约需近 4 秒ZenFlow 将其中约一半的代价通过异步重叠分摊到前面的累积步中从而隐藏了绝大部分 CPU 卸载成本。关键洞察ZenFlow 不仅仅是把 CPU 优化器步骤藏起来更重要的是通过立即更新最重要Top-k的梯度让 GPU 上的训练进度持续前进这也是其论文标题中异步更新Asynchronous Updates的含义所在。从源码结构看zenflow_report.log的生成逻辑zf_benchmark.py会区分accumulation_step与update_step并分别统计平均耗时正是为了量化这两类步长的差异。六、用基准脚本定量评估配置组合除了微调示例仓库还提供了独立的基准测试模块benchmark 目录用于在不依赖真实模型的前提下扫描 ZenFlow 配置空间。6.1 单次运行deepspeed --num_gpus4 zf_benchmark.py --hidden_dim 4096 --nlayers 4 --iteration 5 \ --pin_memory_opts 1 --topk_ratios 0.1 --update_intervals 2 --overlap_stepszf_benchmark.py 使用一个由若干nn.Linear组成的SimpleModel其内部config_dict见 zf_benchmark.py与微调配置保持同构ZeRO Stage 2 CPU 卸载 zenflow四参数并开启wall_clock_breakdown输出细粒度耗时。输出示例Step 1 time: 242.16ms time (ms) | fwd: 55.15 | bwd: 180.82 | bwd_inner: 13.28 | bwd_allreduce: 160.15 | step: 183.61 [Summary] pin_memoryTrue topk_ratio0.1 update_interval2 overlap_stepTrue avg_accumulation_step16.77ms avg_update_step171.38ms可以看到累积步平均仅 16.77ms而完整更新步平均 171.38ms两者差距印证了 ZenFlow 把完整更新的成本集中分摊到了少量步中。6.2 全量扫描与结果汇总run_benchmark.sh 会对pin_memory × topk_ratio × update_interval × overlap_step的笛卡尔积组合默认TRIALS1、NGPUS2逐一运行并为每次运行随机选择 20000~65000 之间的MASTER_PORT避免端口冲突所有输出追加写入zf_benchmark.log。最后调用 output_table.py 用正则解析日志、按配置聚合求平均生成如下形式的对照表-------------------------------------------------------------------------------------------------------------------------------------------------------------- | trial | topk_ratio | update_interval | overlap_step | pin_memory | avg_step (ms) | avg_bwd (ms) | avg_fwd (ms) | avg_selective_optimizer_step (ms) |各指标含义依据 benchmark/README.md指标含义分析价值avg_step (ms)每步端到端耗时训练性能的首要指标越低越好avg_bwd (ms)反向传播含梯度计算与 allreduce耗时对比重叠是否挤压了计算资源avg_fwd (ms)前向传播耗时同上avg_selective_optimizer_step (ms)选择性优化器步骤耗时衡量 ZenFlow 卸载逻辑自身的开销分析要点对比不同配置即可定位权衡关系——pin_memoryTrue通常降低传输延迟更高的update_interval降低卸载频率与开销overlap_stepTrue在模型更新阶段较长时收益更明显。你可以结合自己的硬件情况GPU 型号、CPU 核数、内存带宽从中选择最优组合。七、常见注意事项与 FAQ换模型/改超参编辑 finetune_llama.sh 头部变量MODEL_NAME、EPOCHS、SEED、OUTPUT_DIRlr、batch_size、weight_decay等以 zf_config.json 为准脚本中的同名变量仅作回退。checkpoint 位置训练完成后保存在./alpaca_output由--output_dir指定包含 DeepSpeed checkpoint 与分词器文件。多机扩展脚本当前为单机多卡NNODES1若要多机训练需参照 DeepSpeed 常规做法补充 hostfile 与 node rank 环境变量。显存受限怎么办ZenFlow 的价值正是在显存不足、必须依赖 CPU 卸载的场景中体现若显存充足ZeRO Stage 2 不卸载本身就是更快的选择无需开启。版本前提ZenFlow 依赖deepspeed0.16.0请先确认本地版本再运行示例。八、引用与进一步阅读ZenFlow 的学术依据来自 arXiv 论文《ZenFlow: Enabling Stall-Free Offloading Training via Asynchronous Updates》仓库 README 中给出了完整 BibTeX 条目作者包括 Tingfeng Lan、Yusen Wu、Bin Ma 等。若在你的研究中使用了本示例可引用该论文。想继续深入建议按以下顺序阅读仓库内文件微调入口finetune_llama.py 与 finetune_llama.sh配置全貌zf_config.json基准测试zf_benchmark.py、run_benchmark.sh、output_table.py 及 benchmark/README.md。总结ZenFlow 通过异步重叠 Top-k 选择性更新两条路径将 CPU 卸载场景下原本串行、昂贵的优化器步骤拆解并隐藏到梯度累积步的间隙中实现了接近无阻塞stall-free的训练体验。本文以 DeepSpeedExamples 中的 Llama-2 微调示例为主线完整覆盖了从依赖安装、配置解读、脚本运行到日志分析与基准扫描的全流程。你可以直接复制 zf_config.json 中的配置骨架结合基准测试工具为自己的模型与硬件定制一套最优的卸载训练方案。赞分享示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载相关推荐基于 Ray Train DeepSpeed ZeRO-3 微调 Llama-2 系列大模型全参微调与 LoRA 实战指南基于 Ray Train DeepSpeed ZeRO 3 微调 Llama 2 系列大模型全参微调与 LoRA 实战指南 导读 本指南基于开源仓库 Ra人工智能分布式训练强化学习任务调度模型推理服务后端MiniCPM-SALA 微调实战基于 Transformers Trainer / DeepSpeed / Accelerate 与 LLaMA-Factory 的全参数与 LoRA 微调指南MiniCPM SALA 微调实战基于 Transformers Trainer / DeepSpeed / Accelerate 与 LLaMA Facto大模型本地部署模型量化微调LoRA工具调用openBMBAscendGLM-OCR 微调实战基于 LLaMA-Factory 的全参微调与 LoRA 微调完整指南GLM OCR 微调实战基于 LLaMA Factory 的全参微调与 LoRA 微调完整指南 本文是 GLM OCR 官方微调教程 examples/fi人工智能大模型计算机视觉OCR本地部署AI 技能上一篇三步跑通 Wren AI自然语言转 SQL 取数快速上手指南下一篇YooAsset资源管理框架解决Unity游戏开发中资源加载痛点的完整解决方案创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考