融合优化器与状态精度：把参数更新的读写和状态分开优化

01 · backward结束，训练step还没结束

反向产生梯度后，优化器读取参数、梯度和历史状态，算出下一版参数。模型有很多Tensor时，逐Tensor、逐操作启动小kernel可能产生大量固定成本；大Tensor则常需要搬运许多参数与状态。

融合优化器主要改变这些更新怎样执行。减少moment存储精度则改变状态表示和数值行为。二者可组合，但不能把一种省容量的方法都归为“算子融合”。

依据：PyTorch · 优化器for-loop / foreach / fused执行 ↗ · PyTorch · AdamW定义与执行选项 ↗

02 · AdamW的一次更新需要哪些数

参数θ、梯度g，以及一阶moment m和二阶moment v。m跟踪梯度方向的历史平均，v跟踪梯度平方的历史平均；偏置修正处理零初始化带来的起步偏差。AdamW的weight decay与这些moment更新分开。

教学示意 · 不代表硬件实测

t = t + 1
m = beta1*m + (1-beta1)*g
v = beta2*v + (1-beta2)*g*g
m_hat = m / (1-beta1**t)
v_hat = v / (1-beta2**t)
theta = theta*(1-lr*weight_decay) - lr*m_hat/(sqrt(v_hat)+eps)
# 本篇采用无AMSGrad、最小化目标、每参数有效更新的基本AdamW。

自设θ=[1,−2]、g=[0.1,−0.2]，m/v初始0，β1=0.9、β2=0.999、lr=0.01、decay=0.1。第一步m=[0.01,−0.02]、v=[0.00001,0.00004]，修正后m_hat=g、v_hat=g²。eps很小时，更新约为[0.989,−1.988]。

把weight decay加进g再更新moment，是另一种耦合L2路线，不能作为同定义AdamW的融合替代。epsilon位置、step计数、grad=None、梯度裁剪及混合精度溢出跳步都需要保留原合同。

依据：PyTorch · AdamW定义与执行选项 ↗

03 · 横向合批与纵向融合，分别省什么

执行方式	主要变化	代价或边界
逐Tensor循环	每个Tensor分别执行m/v/参数更新	小操作多，启动成本可能明显。
foreach / multi-tensor	一次组织许多Tensor的同类操作	减少逐Tensor调用，可能增加Tensor列表中间量。
融合更新	同一元素或块内完成多步更新，再写回	减少中间读写及启动，但受dtype、设备、参数组与实现限制。

同一优化器可以有很多参数组：不同lr、decay、step状态要按各自合同使用。没有梯度的参数通常不执行这一轮更新；不能为了整批融合，把它当零梯度继续衰减或推进moment。

PyTorch文档区分for-loop、foreach和fused，并提示foreach额外峰值内存。理论上融合有机会减少更多启动，不代表每个shape或设备都必然最快。

依据：PyTorch · 优化器for-loop / foreach / fused执行 ↗ · PyTorch · AdamW定义与执行选项 ↗

04 · 一个简化读写账本

假设θ/g/m/v全为FP32，并把每次数组访问都计4B。三个分离循环：更新m读m/g写m，共12B；更新v读v/g写v，共12B；更新θ读θ/m/v写θ，共16B；合计40B/元素。

融合循环读θ/g/m/v一次，共16B，写θ/m/v，共12B，总28B/元素。偏置修正作为标量处理，不另物化m_hat/v_hat。这里只比较自设三循环与一循环，真实实现还可能有其他缓冲或缓存命中。

更新公式不变，逻辑读写少了多少？

元素数参数Tensor数

θ/g/m/v均为FP32；三个分离循环40B/元素，一融合循环28B/元素。调用数按每Tensor三次与整组融合一次的自设模型计；真实kernel数取决于实现、设备和参数组，非实测。

05 · 少存moment，不等于少存整套训练状态

若有一亿个参数，两个FP32 moment原始载荷为800 MB；换成两个16bit moment为400 MB，两个8bit码为200 MB，再加低位元数据。参数、梯度、master副本、激活和分片布局都要另外统计。

一种教学量化模型把一块m除以该块scale，四舍五入成有符号8bit码，恢复时乘回scale。多个元素共用尺度，离群值会让小moment难以表示；每轮解码、更新、重新编码会让误差沿训练累积。真实8bit优化器可能使用非线性码本和分块统计，不等于这段均匀量化。

教学示意 · 不代表硬件实测

教学signed moment块：
scale = max(abs(moment)) / 127
q = round(moment / scale)
moment_approx = q * scale
# 全零块单独处理；只解释表示误差，不是完整8bit Adam配方。

PyTorch main文档提供特定fused CUDA路径的BF16 moment示例；bitsandbytes则有8bit状态和选择性高精度设置。它们的支持范围要按安装版本核对，不能手动把任意optimizer状态cast后就认为正确。

依据：PyTorch main · 特定融合路径的BF16优化器状态 ↗ · bitsandbytes · 8bit优化器与选择性精度 ↗

06 · 正确更新要和整个step配套

混合精度训练常先做梯度unscale、检查非有限值和裁剪，再执行更新。若这一步按合同跳过，参数、moment及相关有效step计数都不能被融合kernel提前改掉。梯度累积则在规定更新时刻才调用优化器，不是每个microbatch都推进Adam状态。

ZeRO/FSDP分片的是谁拥有状态，融合优化器组织的是本地状态怎样更新；两者可组合。状态卸载又引入搬运，不能只测设备上更新kernel忽略来回传输。图捕获还需要支持对应的step计数及可捕获模式。

教学示意 · 不代表硬件实测

本轮规定梯度已累积并同步
unscale / finite检查 / 裁剪（按配方）
若允许更新：按每组参数合同执行优化器
若跳过：保持相应更新状态不变
清理梯度，推进调度器等（按训练合同）
# 调度器、AMP scaler和optimizer的step含义需要一致。

依据：PyTorch · AMP 示例（main） ↗ · PyTorch · AdamW定义与执行选项 ↗

07 · CPU验证多步状态，不只看第一步参数

下载脚本比较分离循环与逐元素融合的多轮AdamW参数、m/v和有效step，包含零梯度、grad=None以及整步溢出跳过。另展示耦合衰减的错误替代，并检查均匀8bit moment的尺度、全零块和误差界。没有调用真实fused或8bit优化器。

教学示意 · 不代表硬件实测

记录算法与超参数、参数组、dtype、step/跳步合同
用多轮固定梯度比较参数与全部moment
若压缩状态：核对还原、误差、尺度与恢复
测更新kernel、峰值状态和完整训练step
最后比较Loss/收敛与达到目标质量的总时间

依据：PyTorch · AdamW定义与执行选项 ↗ · bitsandbytes · 8bit优化器与选择性精度 ↗

08 · 验收和跨平台读法

同精度融合：参数和状态多步误差，缺失梯度、参数组、epsilon、衰减及跳步。
低位状态：状态容量、选择性高精度对象、数值误差和收敛。
性能：启动数、实际读写量、临时内存、更新与完整step。
恢复：检查点应保存状态dtype/scale、step和配方，不能只恢复参数。

fused CUDA、bitsandbytes和昇腾训练算子是不同实现入口。平台之间应匹配相同更新定义和训练配方，不把一个框架开关当所有GPU/NPU共有能力。

依据：vLLM Ascend · 支持功能矩阵（main） ↗

grad=None为什么不一定等同于一张全零梯度？
融合更新能否自动把moment从FP32变成8bit？
只验证第一步θ，为什么可能漏掉错误的m/v和step？

流程位置与关联方法

这个方法可能出现在多个位置；计算线包含模型计算与算子实现，执行线包含运行时、内存和通信。

T-A4 · 优化器更新参数 T-B6 · 调度参数更新与预取

资料与核查范围

原理依据以原始论文和官方文档为主；算例与时间线是教学推导。文档核查日期：2026-10-10。引用 main / latest 文档时，实际接口以安装版本为准。

代码为接入片段或教学伪代码；本页没有声称完成 NVIDIA / 昇腾硬件实测。