ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

DeepSpeed ZeRO-3 参数分区实战:让单卡 OOM 的 7B 模型装进显存

DeepSpeed ZeRO-3 参数分区实战:让单卡 OOM 的 7B 模型装进显存 DeepSpeed ZeRO-3 参数分区实战让单卡 OOM 的 7B 模型装进显存【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed用一张 24G 显卡微调 7B 模型第一步前向就报 GPU 显存不足OOM。问题不在模型太大而在每张卡都冗余存了一份完整参数、梯度和优化器状态。DeepSpeed 的 ZeRO-3Zero Redundancy Optimizer第三阶段把模型状态切分存到多卡上计算到谁才聚集谁的参数配合 CPU 卸载Offload7B 模型在单卡也能训起来。本文讲清参数分区的原理、三步配置和关键调优参数。一个 7B 模型切到 8 张卡ZeRO-3 参数分区原理 打个比方。把 7B 模型想象成一整套 70 本图书传统数据并行Data Parallelism等于每个分馆都买齐整套8 个分馆就买 8 套。ZeRO-3 则把书拆成 8 份分库存放每个分馆只留 1/8哪个分馆要查某本书就临时从所有分馆把相关页码借齐聚集Gather用完立刻归还释放Free。任意时刻每卡上只躺着 1/N 的参数显存自然降下来。官方教程里有个真实数据1.5B 参数的 GPT-2 在 8 张 V100 上裸跑数据并行直接 OOM仅 Adam 优化器状态就占 18GB启用 ZeRO 阶段 1 后这部分降到每卡 2.25GB。官方文档docs/_tutorials/zero.md三个阶段的切分对象逐步扩大阶段切分对象每卡显存收益典型规模普通数据并行不切分无基线约 1B 以下Stage 1优化器状态省 4-8 倍1B-10BStage 2优化器状态 梯度省 8-16 倍10B 级Stage 3参数 梯度 优化器状态省 32-64 倍随卡数线性扩展百B 以上收益倍数引自 ZeRO 论文结论实测截图见官方文档docs/_tutorials/zero.md。实现层面每个参数内部挂了一个三态状态机可用、不在本卡、传输中前向走到哪一层才把哪一层的分片聚齐用完即释放。核心代码在 deepspeed/runtime/zero/partition_parameters.py。ZeRO-3 三步配置安装、写配置、启动第一步安装 DeepSpeed。以下命令克隆仓库并本地安装需要编译 CUDA 算子装完会用git clone https://gitcode.com/GitHub_Trending/de/DeepSpeed cd DeepSpeed pip install .第二步写一份最简 JSON 配置。这里只打开一个开关其余参数全部走默认值先跑通再调优{ zero_optimization: { stage: 3 } }第三步用 DeepSpeed 启动器拉起训练。--num_gpus指定本机卡数--deepspeed_config指定上一步的配置文件python -m deepspeed.launcher.launch --num_gpus8 --deepspeed_config ds_zero3.json train.py三步跑通后nvidia-smi里每卡的显存占用会明显低于裸跑数据并行。stage3 显存参数怎么调一张表看懂关键配置参数名作用建议值调错了会怎样stage切分阶段开关3设成 1/2 时参数不分片大模型照样 OOMoffload_param把分片参数卸载到 CPU{device: cpu}不开则单卡放不下超大模型开了吞吐下降offload_optimizer优化器状态卸载到 CPU{device: cpu}同上显存和速度的直接权衡stage3_param_persistence_threshold小于该阈值的小参数常驻显存不释放1e5调太大显存悄悄涨调小反复聚集拖慢速度stage3_max_live_parameters显存中同时驻留的参数上限1e9太小则聚集/释放频繁通信开销暴涨stage3_max_reuse_distance参数预取Prefetch的提前窗口1e9太小预取来不及白等通信太大占用显存stage3_prefetch_bucket_size异步预取的通信批大小1e7太小通信碎片化太大显存瞬时冲高stage3_gather_16bit_weights_on_model_save保存模型时自动聚齐全量权重true保持 false 时保存下来的是碎片分片无法直接加载以上建议值与默认值见 deepspeed/runtime/zero/config.py 及官方教程docs/_tutorials/zero.md。调参的主线只有一条先用默认值跑通OOM 就开 offload慢了再往上抬max_live_parameters和max_reuse_distance。这两个值是一对矛盾——抬得越高一次前向里能借住的参数越多预取命中率越高但常驻显存也越多。建议小步上调每次看nvidia-smi峰值留 20% 余量。开启卸载或大规模卡数时官方教程中的 ZeRO-InfinityZeRO-3 CPU/NVMe 卸载引擎方案可直接参考docs/_posts/2021-03-08-zero3-offload.md。ZeRO-3 避坑清单现象、原因与解法现象报错提示参数为空或不可用访问到空张量原因在前向/反向之外的代码里直接读了别的卡才持有的参数分片解法用GatheredParameters上下文包住访问逻辑跨模块共享的参数在模块初始化时调用register_external_parameter注册源码见 deepspeed/runtime/zero/init.py现象还没开始训练Model(...)构建阶段就 OOM原因每张卡先完整建了一份模型再切分峰值显存等于整模解法把模型构建包进deepspeed.zero.Init上下文参数在构造时就分片落盘用法见 docs/_tutorials/zero.md现象训练结束存的权重文件别的程序加载报缺 key 或形状不对原因stage 3 下默认每张卡只保存自己那份分片解法配置里加stage3_gather_16bit_weights_on_model_save: true或调用引擎的save_16bit_model接口保存全量权重现象切到 ZeRO-3 后 step 时间比 stage 2 明显变慢原因参数每次前向/反向都要现场聚集预取窗口没吃满解法按上一张表的顺序调大stage3_max_reuse_distance和stage3_prefetch_bucket_size并确认contiguous_gradients已开启完整字段说明可查官方配置手册docs/_pages/config-json.md。下一步建议先按 docs/_tutorials/zero.md 里的 1.5B GPT-2 示例把 stage 1/2/3 各跑一遍对比显存曲线再把自己模型的配置迁过来。【免费下载链接】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),仅供参考
返回列表