DeepSpeed 发布了面向模型训练开发者的 0.20.0 版本。用户现在可以选择用于在相同条件下复现结果的执行策略,记录并重放计算流程的功能也扩展至更多加速器。此外,训练期间保留计算值和恢复已保存训练状态的过程中的错误也得到修复。
1. 确定性执行作为可选功能引入
此版本新增了由用户选择开启的确定性执行策略。确定性执行旨在相同条件下重复计算时复现结果。对于需要比较训练结果的开发者,这提供了一种新的执行方式选择。
此次发布还包含与训练计算准确性相关的修复。在将多个子神经网络作为专家使用的模型中,限制梯度大小的计算已修正为每个专家的值只计入一次。此时,用作大小基准的值以 32 位浮点数计算。
在结合专家划分与张量划分的 AutoEP 和 AutoTP 配置中,专家梯度调整错误也得到修复。这意味着除了新增执行策略,实际的训练计算路径也进行了修正。
2. 图捕获扩展至非 CUDA 加速器
Hybrid Engine 现在也支持在非 CUDA 加速器上进行图捕获。图捕获是预先记录待执行的计算流程并反复使用的方式。此次变更扩大了能够使用该功能的硬件范围。
图捕获需要针对各加速器显式启用。此次也新增了昆仑芯 XPU 的图捕获支持。因此,评估是否采用时,应先同时检查各加速器的支持情况和启用要求。
内存管理也根据执行环境作出调整。CPU 加速器识别的内存上限将受应用于进程的资源管理上限约束。这一变更区分了系统总内存容量和单个进程获准使用的容量。
3. 修正 ZeRO 的梯度保留和内存处理
ZeRO 是对训练状态进行分片管理的功能,其第 1、2 阶段(ZeRO-1/2)现在会跨请求边界保留梯度,直至更新权重。梯度是训练中用于计算权重应向哪个方向调整的值。此次修正了汇集多个请求的计算结果并用于一次更新的路径。
在将训练状态卸载至 CPU 的配置中,每组设置一个累加器,并将未使用的梯度置为 0。ZeRO-2 在此配置下改为在分片区域中汇集并累加跨请求的梯度。这里的累加器是汇集组内梯度的 CPU 存储空间,分片区域则用于存放分片管理的梯度中对应的那一份。
ZeRO-3 通过页锁定内存缓冲区,将存放在设备外的更新用主权重复制到设备上。如果权重分片已经与主权重具有相同的数据类型,则直接复用,而不另建副本。此次同时修正了向设备传输数据的路径和副本管理方式。
4. 训练状态恢复和存储设备错误也得到修复
通用检查点是可供多种训练配置共同读取的存储格式。加载该格式时,融合式优化器和支持图捕获的优化器的训练步数值现在会恢复到正确的位置。优化器是利用梯度更新权重的组件。
对于存放在 CPU 上的优化器状态,原本为 32 位浮点数的状态在读取已保存状态后,仍会保持相同精度。
在将训练状态存放于存储设备的 ZeRO-3 配置中,预取参数时若交换缓冲区已分配,原先会发生错误,此次已修复。交换缓冲区是设备间传输所用的临时存储空间。DeepNVMe 也修复了未对齐缓冲区引发的输入输出失败。
引擎和 ZeRO 优化器中的内存泄漏也得到修复。此次修复不仅涵盖训练计算本身,还涉及状态恢复、数据传输和内存管理。
5. 可选配置的灵活性与验证负担
在汇集多个请求的梯度用于一次更新的 ZeRO-1/2 训练中,跨请求边界保留梯度是判断是否采用的依据。在 ZeRO-3 中,权重分片与主权重数据类型相同的配置,需要评估复用现有分片、而非另建主权重副本这一变更。
在进程受到资源上限约束的 CPU 训练环境中,应根据该进程获准使用的内存,而非系统总容量来评估配置。将这一限制反映到 CPU 加速器所识别的内存上限中的变更,与此类环境直接相关。
发行说明未解释确定性执行策略的具体启用方法及其适用的运算范围。
此次发布的核心在于扩展可复现性选项和加速器范围,同时提高 ZeRO 训练与内存路径的稳定性。
