A05 · 第 5 章 · 第三编 训练工程与综合复习

混合精度、分布式训练与检查点

比较 FP32、FP16 与 bfloat16 的范围和精度,以主权重、损失缩放和非有限检测保护更新,再推导数据并行的全局批量、梯度同步、累积、扩展效率及可恢复检查点。

报告页面错误
预备知识正则化、数据增强与早停随机梯度、动量与自适应方法浮点数、条件数与误差传播并行计算、可复现性与性能测量

本章目标

  1. 比较 FP32、FP16 与 bfloat16 的指数位、有效数字和动态范围。
  2. 解释低精度计算副本、FP32 主权重及优化器状态为何需要分工。
  3. 按正确顺序执行损失缩放、反缩放、非有限检测、裁剪和参数更新。
  4. 推导数据并行与梯度累积下的全局批量和平均梯度。
  5. 用通信重叠、吞吐、步延迟、强扩展与弱扩展诊断并行收益。
  6. 设计包含模型、优化器、缩放器、随机状态、数据游标和清单的可恢复检查点。
页面阅读位置0% · 仅保存在此浏览器
章节未开始
本册完成进度0/6 章 · 0%
本页目录

低精度先改变范围和舍入间隔

常见二进制浮点格式由符号、指数和尾数组成。FP32 有八位指数和二十三位显式尾数,正规数约有七位十进制有效数字,最大有限值约 3.4×10383.4\times10^{38}。IEEE FP16 有五位指数和十位尾数,最大值为 6550465504,最小正规正数约 6.10×1056.10\times10^{-5},加上次正规数可到约 5.96×1085.96\times10^{-8}。bfloat16 保留 FP32 的八位指数,却只有七位尾数,因此动态范围接近 FP32、精度比 FP16 更粗。

FP16 更容易让小梯度下溢或大激活溢出;bfloat16 缓解范围问题,却仍会因尾数短丢掉小增量。格式名字也不完全决定算子行为:矩阵输入可以是低精度而累加为 FP32,归一化和归约可能提升精度,具体由实现、算子和硬件决定。报告应记录每类关键算子的输入、累加和输出精度,不能只写“开启半精度”。

softmax、对数、方差、范数和长归约对范围与抵消敏感,常保留较高精度。将所有张量强制转成低精度不叫可靠混合精度;真正目标是在误差可控处减少存储或提高算术吞吐,在敏感状态和归约处保留范围。

例 1:尾数间隔与主权重

在数值一附近,FP16 相邻正规数间隔约为 2109.77×1042^{-10}\approx9.77\times10^{-4},bfloat16 间隔约为 27=7.8125×1032^{-7}=7.8125\times10^{-3}。值 1+291+2^{-9} 能由 FP16 表示,却小于 bfloat16 的一个间隔,转换后可能舍入回一。

若权重约为一、每步更新 10410^{-4},直接在 FP16 权重上相减可能连续舍入为零。把 FP32 主权重用于优化器更新,十步累计约 10310^{-3} 后,再转换出的 FP16 计算副本才出现可见变化。主权重保护小更新;它不会修复前向激活溢出,也不会免除验证数值误差。

主权重和状态按职责分层

一种常见方案保留 FP32 主参数、梯度归约结果和优化器矩状态,前向与反向使用 FP16 或 bfloat16 参数副本与激活,部分算子以 FP32 累加。每次成功更新后把主参数转换回计算副本。这样参数小增量和 Adam 的一阶、二阶矩不会持续在粗尾数中丢失。

并非每种 bfloat16 训练都必须复制一套主权重,也并非所有框架梯度都以同一格式存储。选择应基于模型状态大小、算子支持、数值诊断和恢复要求。混合精度常显著减少激活内存,但带 FP32 主权重和矩状态时,模型状态内存未必减少,甚至可能多出一份低精度参数副本。

边界判断应基于实际峰值内存分解:参数、梯度、优化器状态、激活、临时工作区、通信桶和内存碎片。只用参数量乘字节数无法预测能否放下一个小批量。

损失缩放把小梯度移入可表示范围

设损失为 LL,缩放因子为 S>0S>0。反向传播 SLSL 得到缩放梯度 SgSg,随后在较高精度中除以 SS 恢复 gg。这不会改变理想实数梯度,却能让原本低于 FP16 次正规范围的中间梯度先变大。SS 太大又会使激活梯度或归约溢出为无穷。

动态缩放器在每个候选更新检查所有梯度是否有限。若任一参数出现 NaN 或无穷,所有数据并行进程必须一致跳过该优化器更新,并降低 SS;连续若干成功步后可谨慎增大。反缩放后再做梯度范数裁剪,因为裁剪阈值定义在真实梯度尺度。调度器、全局步和权重衰减是否前进也要与“成功优化器更新”语义一致。

不同实现可能先归约缩放梯度再统一反缩放,也可能局部反缩放后归约;两条路径要保证缩放约定、非有限标志和平均因子一致。一个进程溢出而其他进程继续更新,会立即让模型副本分叉。下溢成零不一定触发非有限检测,因此还需监控梯度零比例、最小非零量级和长时间无更新参数。

例 2:缩放同时面对下溢与溢出

某个理想梯度为 2302^{-30},小于 FP16 最小次正规数 2242^{-24},直接存储会变为零。取 S=210S=2^{10},反向中对应值变为 2202^{-20},可以作为次正规数保留;在 FP32 中除以 2102^{10} 后恢复近似 2302^{-30}

但若另一梯度本来接近 100100,同一缩放会得到 102400102400,超过 FP16 最大有限值 6550465504。动态检测发现无穷后,整组进程跳过本步并减小缩放。它不能只保留小梯度,也不能只避免大梯度,实际 SS 是当前梯度范围的折中。

确定性比固定随机种子更严格

浮点加法不满足结合律。all-reduce 的分块、网络拓扑、线程调度和原子更新会改变求和顺序,低精度使末位差异更明显;后续非线性与随机采样可放大差异。某些高性能算子使用非确定性并行算法,即使输入和种子相同也不保证逐位一致。

可复算协议应保存模型初始化、数据洗牌、增强、dropout、工作器和每个进程的随机状态,并记录确定性开关、库、驱动、算子选择、进程数和归约配置。确定性模式可能降低吞吐或缺少某些算子,应实测并披露。跨硬件无法保证位级相同时,预先定义损失、参数摘要或最终指标的数值容差,并用多种子报告统计波动。

精确恢复还要求在相同数据边界保存。若检查点位于梯度累积一半,必须连同已累积梯度和微步索引保存;更简单的策略是在完整优化器步后同步保存。随机种子本身不能重建已消费到中途的随机流,必须保存生成器状态或确定数据位置。

数据并行同步同一份模型

同步数据并行在 WW 个进程上复制模型,每个进程处理不同本地小批量,反向后用 all-reduce 汇总梯度,再执行相同优化器更新。若每个进程本地损失是 bb 个样本的均值、各进程样本数相等,先求和再除以 WW 得到全局 WbWb 个样本的均值梯度。

若最后一批大小不等,简单平均各进程均值会给小批量过高权重。应按本地样本数加权,或保证共同截断和填充策略不会重复评价样本。数据采样器要让训练单位不重叠,并把 epoch、分片和洗牌种子纳入检查点。验证阶段要去除为对齐进程而复制的样本,否则指标分母错误。

数据并行复制全部模型状态,适合模型能放进单设备、数据批量可分的场景。模型本身放不下时需要参数、张量或流水分片;这会引入新的通信、分区和检查点格式,不应与简单数据并行混称。

梯度累积改变全局批量和更新频率

每进程微批量为 bb,累计 AA 个微步后同步并更新一次,则每个优化器步的名义全局批量为

Bglobal=WbA.B_{\mathrm{global}}=W\,b\,A.

每个微步损失若已取均值,可除以 AA 后反向,或累积梯度后统一除以 AA;再对进程取平均。重复除法或漏除会让有效学习率改变。通信可以只在最后一个微步进行,但任何进程提前失败都需让整个累积窗口作废或一致恢复。

例 3:四进程八微步的全局梯度

四个进程各用微批量十六,累计八个微步,故每次更新覆盖 4×16×8=5124\times16\times8=512 个样本。若数据集有 5120051200 个互异训练样本,每个 epoch 约一百次优化器更新,而不是每个进程的八百次微步。

各微步均值损失先除以八再反向,最后 all-reduce 对四个进程梯度取平均,可得到这五百一十二个样本的均值梯度。若归约库已经取平均,代码又除以四,更新会缩小四倍。运行记录应同时写微批量、累计步数、进程数、全局批量和归约语义。

梯度累积与一次真实大批量并非所有方面等价。固定参数且无状态随机层时,样本梯度求和可一致;但 BatchNorm 每个微步只看 bb 个样本,dropout 掩码与随机增强时序不同,混合精度舍入路径也不同。显存不足时累积是有效工具,不应据此声称完全模拟了同样统计的全局批。

全局批量变化必须重新解释学习率

增加进程或累积步会减少每个 epoch 的优化器更新数并改变梯度噪声。学习率线性缩放、平方根缩放和预热是经验起点,不是无条件定理。模型、优化器、归一化、数据相关和训练预算都会改变适用范围。固定原学习率可能过小也可能正好,机械放大又可能早期发散。

比较配置时用已处理样本数、token 数或 epoch 标记调度位置,而不只看本地 step。权重衰减和动量按更新次数作用,更新数改变时也要审查。先在较小规模建立收敛基线,再扩展并比较达到同一验证目标的墙钟时间;吞吐提高但需要更多样本才收敛,不一定节省总训练成本。

通信与计算重叠需要看关键路径

反向传播从后层向前层产生梯度。把相邻梯度装入桶,桶就绪后立即启动异步 all-reduce,可让后续反向计算遮蔽一部分通信。桶太小会产生大量启动延迟,太大又要等很久才开始;最佳大小取决于模型层次、网络带宽和拓扑。

例 4:吞吐增加不等于线性扩展

单设备每步计算一百二十毫秒,处理六十四个样本。四设备在相同本地批量下,计算仍约一百二十毫秒,all-reduce 原需六十毫秒;分桶重叠后只剩二十毫秒暴露在关键路径,总步延迟约一百四十毫秒。全局吞吐约为 256/0.141829256/0.14\approx1829 样本每秒,单设备约 64/0.1253364/0.12\approx533,速度比约 3.433.43,不是四倍。

这里同时增大全局批量,属于弱扩展式测量,不能直接说明固定问题的强扩展。还需比较达到同一验证质量的时间,并记录数据读取是否成为瓶颈。换模型、桶、网络或精度后数字都会变化,不能把该比值当作特定硬件的固定保证。

强扩展固定总工作或全局批量,增加设备以缩短完成时间,效率可写为 Es=T1/(WTW)E_s=T_1/(W T_W)。弱扩展固定每设备工作量,观察规模随 WW 增长时步延迟能否保持,效率常与单设备和多设备步时间之比比较。训练还应报告样本吞吐、步延迟分位数、通信暴露时间、数据等待、峰值内存和达到目标的时间。

检查点是分布式状态的一致快照

可精确继续训练的检查点至少包含:FP32 主权重和计算副本所需状态、优化器矩与步数、学习率调度器、损失缩放器、归一化运行统计、梯度累积位置、全局更新与已处理样本数、每进程随机生成器、数据采样器和游标、最佳指标状态、代码配置及数据版本。只保存模型权重只能做推理或暖启动,不能称为原训练续跑。

大状态可按进程分片,每个分片写内容摘要,顶层清单记录世界大小、张量布局和完整性。先写临时路径,所有进程完成并校验后再原子发布清单;任一分片缺失都不能把快照标为可用。改变进程数恢复时,需要明确的重分片器;简单让每个新进程读取旧编号分片可能重复或漏掉状态。

例 5:从故障点恢复而不重复数据

八进程训练在完整更新 12000 后同步保存。清单记录下一数据游标、epoch 内分片位置、八个随机状态、缩放因子、Adam 矩、调度位置和各分片摘要。进程随后在更新 12340 故障,恢复器验证快照后从 12001 重新开始。

若只恢复权重而把 Adam 矩、缩放器和调度器归零,第一步就与原轨迹不同;若采样器从 epoch 开头开始,还会重复已处理数据。位级重放另需相同拓扑和确定性算子。若只要求统计复现,可允许轨迹差异,但必须用预先容差和最终指标复跑验证。

检查点频率在写入成本与预期重算之间权衡。保存墙钟间隔、失败率、快照耗时和存储预算可决定策略;最新检查点之外还应保留里程碑与最佳验证模型,防止静默损坏被新快照覆盖。

练习

练习 1:比较 FP16 与 bfloat16
为何 bfloat16 较少溢出却不一定更精确?
查看提示
比较指数位和尾数位。
查看解答
bfloat16 用八位指数,范围接近 FP32,但七位尾数使精度较粗;FP16 只有五位指数、范围窄,却有十位尾数,在相同量级上常比 bfloat16 精细。
练习 2:排列更新顺序
动态损失缩放、检测和裁剪应怎样排序?
查看提示
裁剪阈值定义在反缩放后的真实梯度。
查看解答
缩放损失并反向,完成同步约定后反缩放,跨进程汇总非有限标志;若有限则按真实尺度裁剪并更新,否则所有进程跳步并降低缩放。
练习 3:计算全局批量
八进程每卡十二样本、累计四步时批量是多少?
查看提示
进程数乘微批量再乘累计微步。
查看解答
W=8、b=12、A=4 时全局批量为 8×12×4=3848\times 12\times 4=384;优化器和调度器每四个微步前进一步。
练习 4:解释累积差异
为何梯度累积不总等于大批量?
查看提示
考察 BatchNorm 和随机层在每个微批上的行为。
查看解答
固定参数的样本梯度和可等价,但 BatchNorm 只看微批统计,dropout 和增强随机序列不同,舍入顺序也不同,因此训练轨迹不保证等同一次真实大批量。
练习 5:评价扩展
怎样判断增加设备真正缩短了训练?
查看提示
同时报告吞吐、步延迟和达到同一质量的时间。
查看解答
区分固定总工作或批量的强扩展与固定每设备工作量的弱扩展;测量通信暴露、数据等待、峰值内存,并比较达到同一验证目标的墙钟时间,不能只看样本每秒。
练习 6:设计可恢复快照
为什么只保存模型权重不能精确续跑?
查看提示
权重之外还有优化、随机和数据位置。
查看解答
保存主权重、优化器、调度器、缩放器、归一化状态、累计位置、全局步、每进程随机状态、采样器游标、代码配置和数据版本;分片以清单与摘要原子发布。

关系与资源

书籍 · 2016

Deep Learning

Ian Goodfellow, Yoshua Bengio, Aaron Courville

适合作为反向传播和优化章节的完整参考。

打开官方来源

《Deep Learning》用于核对深度网络数值计算、优化和训练工程的基础口径。具体低精度算子、通信库和硬件能力会变化,本文不承诺固定加速倍数;任何收益都应连同数值误差、收敛质量和恢复能力实测。