尧图网站开发 
PyPTO-Gym 优化器更新算子模式 AT-19 深度解析AdamW/RMSProp 在昇腾 NPU 上的状态更新内核设计【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym导读本篇文章围绕 PyPTO-Gym 仓库中 AT-19 优化器更新模式 展开系统讲解动量累积 参数衰减 参数写回这一优化器算子族在 PyPTO 框架下的内核设计方法。文中以 AdamW 与 RMSProp 为实例结合仓库内 ApplyAdamWV2 实现 与 ApplyRMSProp 实现 的源码与测试覆盖计算流分解、泛化变体、PyPTO 编程约束、tile 策略与精度验证。读者学完后可以照此模式独立设计并实现任意优化器更新算子AdamW、RMSProp、Lion、LAMB 等的昇腾 NPU kernel并理解其与 FlashAttention、RMSNorm 等 V 型局部模式在设计方法上的共通之处。一、AT-19 模式定位V 型局部计算模式中的优化器步AT-19Optimizer Update是 PyPTO-Gym 算子设计工作流pypto-op-design/SKILL.md中 patterns/atoms 索引下的局部计算模式卡片tag 为optimizer-stepflow_pattern 为纯VVector。在 AT 索引 中它与 Online SoftmaxAT-01、RMSNormAT-03、SwiGLUAT-07等同属局部计算模式即描述单个计算段的排布方式而不像 SKSkeleton卡片那样描述整个算子的宏观结构。所谓纯 V 排布是指优化器更新整条计算链——一阶动量累积、二阶动量累积、偏置修正、开方取倒数、衰减项、参数写回——全部由 Vector 单元完成不涉及 Cube矩阵乘单元。这是因为更新公式中每个元素的计算只依赖同位置的参数、梯度和状态张量属于典型的 element-wise 运算天然适配 V 单元。这一特征使 AT-19 与 AT-03-rmsnorm 等 V 型模式共享同一套 tile、缓存与循环设计约束。设计流程提示按 pypto-op-design/SKILL.md 的规则设计算子时应先读 SK 索引与 AT 索引再读取候选卡片正文。命中 AT-19 时还需遵守AT 模式与 API 功能重叠时必须机制对比择优的纪律若某个更新段既可用本模式实现、也可由单一 PyPTO API 实现必须在 DESIGN.md 中对比搬运开销、指令开销与硬件适配后择优。二、计算流分解从 AdamW 实例到泛化六步骨架2.1 AdamW 实例五段计算链AT-19 卡片给出的 AdamW 计算流是理解整个模式的最佳入口# 一阶动量 m_new add(mul(m, beta1), mul(grad, 1-beta1)) m_hat div(m_new, bias_correction1) # 二阶动量 v_new add(mul(v, beta2), mul(mul(grad, grad), 1-beta2)) v_hat div(v_new, bias_correction2) # 参数更新 denom add(sqrt(v_hat), eps) update add(div(m_hat, denom), mul(w, weight_decay)) w_new sub(w, mul(update, lr))这条链在仓库的 golden 参考实现 apply_adam_w_v2_golden.py 中被逐行印证为标准的单步 AdamW 更新Loshchilov Hutter, 2019m_t beta1 * m (1 - beta1) * g v_t beta2 * v (1 - beta2) * g * g m_hat m_t / (1 - beta1**t) v_hat v_t / (1 - beta2**t) update m_hat / (sqrt(v_hat) eps) weight_decay * weight w_new weight - lr * update需要注意三个易错点偏置修正项1 - beta**step在 host 端预计算。仓库 kernel 中step不进入设备侧不产生任何 device-sidepow运算而是由 wrapper 在 host 上算好bc1、bc2后以标量浮点传入见 apply_adam_w_v2_impl.py。这对 kernel 性能与数值稳定性都很关键。分母加eps的位置在sqrt之外sqrt(v_hat) eps而不是sqrt(v_hat eps)两种写法在v_hat接近 0 时数值行为不同golden 与 kernel 必须保持一致。weight decay 是解耦decoupled形式update weight_decay * w衰减直接作用于原参数而非梯度这是 AdamW 区别于传统 Adam 的本质。2.2 泛化六步骨架所有优化器的统一抽象将 AdamW 实例中与算法无关的结构抽取出来AT-19 给出覆盖整个优化器算子族的泛化骨架# 1. 动量/状态更新按变体配置 state_new state_update_fn(state_old, grad, hyperparams) # 2. 偏置修正可选 state_hat bias_correction_fn(state_new, step) # 3. 自适应学习率可选 adaptive_lr adaptive_lr_fn(state_hat) # 4. 权重衰减可选 decay_term mul(w, weight_decay) # 5. 参数更新 w_new sub(w, mul(adaptive_lr decay_term, lr)) # 6. 状态写回in-place state[:] state_new六个阶段中阶段 1/5/6 是强制性的任何优化器都要更新状态、更新参数、写回阶段 2/3/4 按变体取舍。这一骨架的价值在于设计新优化器算子时只需要替换state_update_fn与adaptive_lr_fn两个插件点其余结构与 PyPTO 编程约束全部复用。例如 RMSProp 只需删掉一阶状态与偏置修正、把自适应 LR 换成g / (sqrt(v) eps)就能在骨架上快速完成设计。三、泛化变体整个优化器算子族的覆盖矩阵AT-19 卡片用一张表系统归纳了该模式覆盖的算子族这是设计时选择变体的直接依据完整继承如下变体一阶状态二阶状态自适应 LR衰减形式典型实现AdamW当前样本m β₁m (1-β₁)gv β₂v (1-β₂)g²m_hat / (√v_hat eps)decoupledw*wdApplyAdamWV2RMSProp当前样本无v β v (1-β)g²g / (√v eps)可选w*wdApplyRMSPropSGD with Momentumm μm g无无可选w*wd训练框架常见AdaGrad无v v g²(累加)g / (√v eps)可选稀疏特征Lionm β m (1-β) g无sign(β₂m (1-β₂)g)decoupledGoogle LionLAMB / LARSm β₁m (1-β₁)gv β₂v (1-β₂)g²m_hat/(√v_hateps) layer-wise scale可选大 batch 训练AdaFactor无行/列各一行/列分离v_r, v_crank-1 重构decoupled显存优化Muon / Sophiam β m (1-β)g二阶信息Hessian 近似矩阵预处理牛顿步可选新兴优化器解读这张表的关键维度状态数量决定资源占用无状态SGD、单状态RMSProp、AdaGrad、Lion、双状态AdamW、LAMB、Sophia依次递增。AT-19 编程约束明确指出多状态m, v, w, grad同时存在时UB 容量约束更紧TILE 需相应缩小这是设计 LAMB 等双状态优化器时的首要资源预算依据。二阶状态的更新方式区分实现AdamW/RMSProp 是滑动平均指数衰减AdaGrad 是累加单调不减AdaFactor 是行/列分离的 rank-1 近似——后两者对数值范围管理防溢出有额外要求。自适应 LR 的形状决定是否引入额外机制LAMB/LARS 的自适应项含 layer-wise scale需要额外的归约与标量广播运算已超出纯 element-wise 范围设计时需评估是否仍走纯 V排布或引入跨核归约。四、仓库落地实现ApplyAdamWV2 源码级拆解AT-19 卡片标注使用算子: ApplyAdamWV2, ApplyRMSProp两者在仓库中都有完整的实现与测试是理解该模式的最佳参考代码。4.1 Host wrapper参数解析与 kernel 分派apply_adam_w_v2_impl.py 的入口是apply_adam_w_v2_wrapper其职责链清晰展示了 AT-19 六步骨架的 host 侧配合参数解析与校验_parse_adam_args严格校验beta1, beta2, lr, weight_decay, eps, step六个参数wrapper 断言四张输入张量 shape 一致、weight/graddtype 匹配、m/v必须为 fp32、step 1L195-L228。host 端预计算标量bc1 1.0 - beta1**step、bc2 1.0 - beta2**step、one_m_b1、one_m_b2L232-L235。按 dtype 分派 kernelweight为 bf16 走apply_adam_w_v2_kernel_bf16fp32 走apply_adam_w_v2_kernel_fp32其余 dtype 直接抛TypeErrorL245-L250。按形状选择 tile 配置_select_adam_tile_config针对三类形状分别返回 DEFAULT / LARGE_M / TARGET_K_AXIS 三套配置见 4.3 节。调用方式来自 READMEfrom apply_adam_w_v2_impl import apply_adam_w_v2_wrapper w_new, m_new, v_new apply_adam_w_v2_wrapper( weight, grad, m, v, beta10.9, beta20.999, lr1e-3, weight_decay0.01, eps1e-8, step1, )4.2 Kernel 主体双循环 valid_shape 尾块 三次 assemble以 fp32 kernel 为例其结构完整映射 AT-19 的 AdamW 计算链L100-L127M/K 双循环外层m_loop、内层n_loop对应参数矩阵的 M 轴与 K 轴均以pypto.loop构建unroll_list[1]保持结构循环不展开尾块处理每层循环计算valid_m/valid_n (dim - offset).min(tile)通过pypto.view(..., valid_shapevalid_shape)传入K 为 N_TILE 整数倍时尾块自动退化为 no-op计算链逐行写出m_new、grad_sq、v_new、m_hat、v_hat、denom、update、w_new与 AT-19 伪代码逐行对应其中div使用precision_typepypto.PrecisionType.INTRINSIC利用硬件 intrinsic 除法的吞吐优势状态写回w_new、m_new、v_new三个结果在同一 kernel 迭代内通过三次独立的pypto.assemble写回对应骨架第 6 步状态写回in-place。bf16 kernel 在此基础上增加了混合精度路径L177-L188入口处w_tile、g_tile先pypto.cast到 FP32全部中间运算保持 FP32仅参数结果在assemble前 cast 回 BF16m_new/v_new始终以 FP32 输出。这正是 AT-19 全程 FP32 计算参数最终 cast 回 BF16约束的实现证据。4.3 Tile 策略形状驱动的三套配置AT-19 卡片给出的经验值是沿参数维度K 轴分块N_TILE2048而仓库实现将其细化为按输入形状自动选择的三档配置L31-L70配置触发条件M_TILEVecTileN_TILE(K 向)设计意图DEFAULT_TILE_CONFIG通用 2D 小形状32(32, 512)1024常规路径LARGE_M_TILE_CONFIGM 7168 且 K 40967168(16, 512)1024M 大 K 小整 M 块减少循环TARGET_K_AXIS_TILE_CONFIGbf16 且M7168 且 K 40967168(16, 512)4096网络级[7168, K]K 动态 2048-24576K 向大块降 prologue 占比三个配置共同印证了 AT-19 的 tiling 原则N_TILE 存在一个平衡点——过小则 prologue循环头开销占比高、带宽利用不足过大则单个 tile 装不进 UBUnified Buffer。仓库采用 1024/4096 两档 K 向 tile 配合pypto.set_vec_tile_shapes设置 V 单元 tile并在TARGET_K_AXIS路径额外开启pypto.experimental.set_operation_options(combine_axisTrue)与vec_nbuffer_setting{DEFAULT: 4}L160-L162来压深流水。4.4 产品支持与输入规格根据 ApplyAdamWV2 README产品支持Ascend 950PR、Atlas A3 训练/推理系列、Atlas A2 训练/推理系列均支持Shape 与 dtypeweight/grad为[7168, K]的bfloat16或float32K 动态范围 2048-24576m/v恒为[7168, K]的float32输出 dtype 与输入一致m/v输出恒为 fp322026-06-11 更新kernel 已支持动态 M 与动态 K 的 2D 张量并新增level5~level7小形状/尾块验证层级[7168, K]网络形状保留 Large-M 策略以维持既有性能覆盖。五、第二个样本ApplyRMSProp 实现对照ApplyRMSProp 是 AT-19 覆盖矩阵中无一阶状态的代表其公式README与骨架第 2/3/4 步的可选项取舍形成直接对照grad_sq grad * grad ms_new ms (grad_sq - ms) * (1 - rho) mom_new mom * momentum (grad * lr) / sqrt(ms_new epsilon) var_new var - mom_newPyPTO 算子映射同样来自 README §3grad * grad计算平方梯度ms (grad_sq - ms) * (1.0 - rho)更新ms注意这是以增量形式书写的滑动平均与 AdamW 的beta*v (1-beta)*g²形式等价但写法不同pypto.sqrt(ms_new epsilon)计算分母mom * momentum (grad * lr) / sqrt(...)更新动量var - mom_new更新参数通过move(...)将结果原地写回var/ms/mom。该实现同样设置了pypto.experimental.set_operation_options(combine_axisTrue)与pypto.set_vec_tile_shapes(32, 512)与 ApplyAdamWV2 的常规路径保持一致。输入输出规格为var/ms/mom/grad均为[rows, cols]的 fp32 张量标量lr默认 0.001、rho默认 0.9、momentum默认 0.9、epsilon默认 1e-7。已知限制当前仅覆盖 fp32、输入必须是二维张量、主验证路径依赖 NPU 环境与torch_npuREADME §9。两个样本的对照价值在于同一个 AT-19 骨架通过增删偏置修正自适应学习率两个可选阶段即可派生出 AdamW 与 RMSProp 两种截然不同的优化器实现这正体现了模式抽象的复用能力。六、PyPTO 特化编程约束带宽、缓存与精度的三条红线AT-19 卡片的编程约束PyPTO 特化部分是设计落地时的硬性要求逐条展开如下6.1 全程 FP32 计算参数最终 cast 回 BF16优化器更新对数值精度高度敏感动量项在长期训练中累积二阶矩的舍入误差会通过自适应学习率被放大。因此 kernel 内部无论输入是 bf16 还是 fp32全部中间计算保持在 FP32仅在写回参数前 cast 回 bf16。仓库的 golden 参考实现 明确将这一行为定义为混合精度契约Mixed precision contractweight/grad in bf16 are cast up to fp32 internally, all math is done in fp32, weight is cast back to its original dtype on return。kernel 与 golden 必须遵守同一契约否则精度验证无法通过。6.2 必须使用 submit_before_loopTrue 实现 pipeline overlap优化器更新是典型的带宽密集型memory-bound算子每个参数元素要读入 grad、w、m、v 四个张量并写回三个张量计算量相对搬运量很小。若不使用submit_before_loopTrue循环迭代之间无法与搬运流水重叠带宽利用率不满。仓库设计规范 pypto-kernel-design-format.md 对此给出精确表述只要迭代之间存在跨块依赖例如 block loop 携带状态就必须在pypto.loop上打开submit_before_loopTrue同时强调循环必须用pypto.loop而非 Pythonfor否则会在构图期被展开、破坏框架的 tiling 与并行调度。6.3 必须配置 set_cache_policy(NONE_CACHEABLE)参数张量在大模型训练中体积很大例如[7168, 24576]的 fp32 张量若允许其驻留 L2 cache会污染缓存、挤占其他数据因此必须显式设置NONE_CACHEABLE缓存策略。这也是该约束在仓库 执行约束文档 与其他 V 型模式如缓存写回类算子中反复出现的原因。6.4 K 轴分块与 N_TILE 的平衡经验值与实现取值沿参数维度K 轴分块N_TILE2048 是经验值过小 prologue 占比高过大 UB 装不下——这条约束在仓库实现中体现为多档 N_TILE1024 常规 / 4096 网络级 K 向并给出量化的边界条件过小循环次数增多每次循环头prologue的固定开销占比升高带宽利用不足过大单个 tile 连同 m、v、w、grad 四个驻留张量超出 UB 容量。AT-19 特别强调多状态同时存在时 UB 容量约束更紧TILE 需相应缩小。AdamW 双状态m、v加 w、grad 共四张驻留而 RMSProp 也是 var/ms/mom/grad 四张因此两者的 vec tile 都收敛到 512 量级。设计 LAMB同样双状态时可复用这一经验值作为起点再按目标形状实测调整。七、精度验证golden 对照与测试入口7.1 ApplyAdamWV2 的验证体系test_apply_adam_w_v2.py 以纯 PyTorch 的 apply_adam_w_v2_golden.py 为参照覆盖level0fp32 精度与level1bf16 精度两级并附带多步偏置修正step10的数值 sanity 检查。运行方式README §Running the precision testcd repo-root source env_setup.sh python3 tests/ops/experimental/vector/ApplyAdamWV2/test_apply_adam_w_v2.py通过标准atol1e-4、rtol0.0078125测试打印[PRECISION_PASS]exit 0或[PRECISION_FAIL]exit 1。此外 golden 自带 smoke test在mv0、step1的解析可验条件下做解析核对此时update g/(sqrt(g²)eps) wd*w并将 bf16 回程误差上限放宽到4e-3以容纳 bf16 舍入apply_adam_w_v2_golden.py#L139-L148。测试用例配置见 test_cases.json。7.2 ApplyRMSProp 的验证体系ApplyRMSProp 的测试以 test_apply_rms_prop.py 为入口支持level08x8、level11024x1024、level216x16三个层级阈值atol1e-5、rtol1e-5。每个 level 构造随机输入 → 调用 kernel → 与 NumPy golden 对比 → 分别输出var_diff、ms_diff、mom_diff摘要格式如level0: shape(8, 8), var_diff0.00000000, ms_diff0.00000000, mom_diff0.00000000。README §10 还给出了典型排障路径TILE_FWK_DEVICE_ID未设置时先export TILE_FWK_DEVICE_ID0Invalid Device时用npu-smi info核对设备号精度断言失败时优先检查输入 dtype、超参与 golden 的一致性以及apply_rms_prop_impl.py中的运算顺序是否被改动。7.3 环境准备运行上述测试需要 CANN、torch_npu与pto-isa环境典型配置ApplyRMSProp README §6source /path/to/ascend-toolkit/set_env.sh export TILE_FWK_DEVICE_ID0 export PTO_TILE_LIB_CODE_PATHpto-isa 路径如需重编译安装 Python 包python3 build_ci.py -f python3 --disable_auto_execute。八、如何在 DESIGN.md 中接入 AT-19 模式按 pypto-op-design/SKILL.md 的工作流命中 AT-19 后应在 DESIGN.md 中完成以下对接范式选择在范式与设计决策中记录采用 AT-19V 型局部模式并对照 AT 索引 确认无更匹配的候选若同一计算段可由单一 API 实现按机制对比规则说明择优理由。变体确认按第三节矩阵确定目标优化器在六步骨架中的可选阶段取舍明确一阶/二阶状态、自适应 LR、衰减形式并同步确认状态张量的 dtype恒为 fp32与驻留数量决定 tile 缩放。Tiling 与配置以 N_TILE 经验值2048为起点结合目标形状与 UB 容量确定 M/K 向 tile 与 vec tile在参数表中把执行切块参数标注为tunable把硬件容量UB 上限、32B 对齐标注为constant。循环与状态写回使用pypto.loop双循环 submit_before_loopTrueunroll_list只选一个展开因子三个输出在同一迭代内独立assemble。精度契约与 golden 对齐全 FP32 中间计算 参数 cast 回原 dtype的混合精度契约并把atol/rtol阈值写入设计文档供后续验证阶段引用。至此从模式卡片到落地 kernel 再到精度验证的完整闭环已经打通——这正是 AT-19 作为局部计算模式在 PyPTO-Gym 仓库中的真实价值。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
PyPTO-Gym 优化器更新算子模式 AT-19 深度解析AdamW/RMSProp 在昇腾 NPU 上的状态更新内核设计【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym导读本篇文章围绕 PyPTO-Gym 仓库中 AT-19 优化器更新模式 展开系统讲解动量累积 参数衰减 参数写回这一优化器算子族在 PyPTO 框架下的内核设计方法。文中以 AdamW 与 RMSProp 为实例结合仓库内 ApplyAdamWV2 实现 与 ApplyRMSProp 实现 的源码与测试覆盖计算流分解、泛化变体、PyPTO 编程约束、tile 策略与精度验证。读者学完后可以照此模式独立设计并实现任意优化器更新算子AdamW、RMSProp、Lion、LAMB 等的昇腾 NPU kernel并理解其与 FlashAttention、RMSNorm 等 V 型局部模式在设计方法上的共通之处。一、AT-19 模式定位V 型局部计算模式中的优化器步AT-19Optimizer Update是 PyPTO-Gym 算子设计工作流pypto-op-design/SKILL.md中 patterns/atoms 索引下的局部计算模式卡片tag 为optimizer-stepflow_pattern 为纯VVector。在 AT 索引 中它与 Online SoftmaxAT-01、RMSNormAT-03、SwiGLUAT-07等同属局部计算模式即描述单个计算段的排布方式而不像 SKSkeleton卡片那样描述整个算子的宏观结构。所谓纯 V 排布是指优化器更新整条计算链——一阶动量累积、二阶动量累积、偏置修正、开方取倒数、衰减项、参数写回——全部由 Vector 单元完成不涉及 Cube矩阵乘单元。这是因为更新公式中每个元素的计算只依赖同位置的参数、梯度和状态张量属于典型的 element-wise 运算天然适配 V 单元。这一特征使 AT-19 与 AT-03-rmsnorm 等 V 型模式共享同一套 tile、缓存与循环设计约束。设计流程提示按 pypto-op-design/SKILL.md 的规则设计算子时应先读 SK 索引与 AT 索引再读取候选卡片正文。命中 AT-19 时还需遵守AT 模式与 API 功能重叠时必须机制对比择优的纪律若某个更新段既可用本模式实现、也可由单一 PyPTO API 实现必须在 DESIGN.md 中对比搬运开销、指令开销与硬件适配后择优。二、计算流分解从 AdamW 实例到泛化六步骨架2.1 AdamW 实例五段计算链AT-19 卡片给出的 AdamW 计算流是理解整个模式的最佳入口# 一阶动量 m_new add(mul(m, beta1), mul(grad, 1-beta1)) m_hat div(m_new, bias_correction1) # 二阶动量 v_new add(mul(v, beta2), mul(mul(grad, grad), 1-beta2)) v_hat div(v_new, bias_correction2) # 参数更新 denom add(sqrt(v_hat), eps) update add(div(m_hat, denom), mul(w, weight_decay)) w_new sub(w, mul(update, lr))这条链在仓库的 golden 参考实现 apply_adam_w_v2_golden.py 中被逐行印证为标准的单步 AdamW 更新Loshchilov Hutter, 2019m_t beta1 * m (1 - beta1) * g v_t beta2 * v (1 - beta2) * g * g m_hat m_t / (1 - beta1**t) v_hat v_t / (1 - beta2**t) update m_hat / (sqrt(v_hat) eps) weight_decay * weight w_new weight - lr * update需要注意三个易错点偏置修正项1 - beta**step在 host 端预计算。仓库 kernel 中step不进入设备侧不产生任何 device-sidepow运算而是由 wrapper 在 host 上算好bc1、bc2后以标量浮点传入见 apply_adam_w_v2_impl.py。这对 kernel 性能与数值稳定性都很关键。分母加eps的位置在sqrt之外sqrt(v_hat) eps而不是sqrt(v_hat eps)两种写法在v_hat接近 0 时数值行为不同golden 与 kernel 必须保持一致。weight decay 是解耦decoupled形式update weight_decay * w衰减直接作用于原参数而非梯度这是 AdamW 区别于传统 Adam 的本质。2.2 泛化六步骨架所有优化器的统一抽象将 AdamW 实例中与算法无关的结构抽取出来AT-19 给出覆盖整个优化器算子族的泛化骨架# 1. 动量/状态更新按变体配置 state_new state_update_fn(state_old, grad, hyperparams) # 2. 偏置修正可选 state_hat bias_correction_fn(state_new, step) # 3. 自适应学习率可选 adaptive_lr adaptive_lr_fn(state_hat) # 4. 权重衰减可选 decay_term mul(w, weight_decay) # 5. 参数更新 w_new sub(w, mul(adaptive_lr decay_term, lr)) # 6. 状态写回in-place state[:] state_new六个阶段中阶段 1/5/6 是强制性的任何优化器都要更新状态、更新参数、写回阶段 2/3/4 按变体取舍。这一骨架的价值在于设计新优化器算子时只需要替换state_update_fn与adaptive_lr_fn两个插件点其余结构与 PyPTO 编程约束全部复用。例如 RMSProp 只需删掉一阶状态与偏置修正、把自适应 LR 换成g / (sqrt(v) eps)就能在骨架上快速完成设计。三、泛化变体整个优化器算子族的覆盖矩阵AT-19 卡片用一张表系统归纳了该模式覆盖的算子族这是设计时选择变体的直接依据完整继承如下变体一阶状态二阶状态自适应 LR衰减形式典型实现AdamW当前样本m β₁m (1-β₁)gv β₂v (1-β₂)g²m_hat / (√v_hat eps)decoupledw*wdApplyAdamWV2RMSProp当前样本无v β v (1-β)g²g / (√v eps)可选w*wdApplyRMSPropSGD with Momentumm μm g无无可选w*wd训练框架常见AdaGrad无v v g²(累加)g / (√v eps)可选稀疏特征Lionm β m (1-β) g无sign(β₂m (1-β₂)g)decoupledGoogle LionLAMB / LARSm β₁m (1-β₁)gv β₂v (1-β₂)g²m_hat/(√v_hateps) layer-wise scale可选大 batch 训练AdaFactor无行/列各一行/列分离v_r, v_crank-1 重构decoupled显存优化Muon / Sophiam β m (1-β)g二阶信息Hessian 近似矩阵预处理牛顿步可选新兴优化器解读这张表的关键维度状态数量决定资源占用无状态SGD、单状态RMSProp、AdaGrad、Lion、双状态AdamW、LAMB、Sophia依次递增。AT-19 编程约束明确指出多状态m, v, w, grad同时存在时UB 容量约束更紧TILE 需相应缩小这是设计 LAMB 等双状态优化器时的首要资源预算依据。二阶状态的更新方式区分实现AdamW/RMSProp 是滑动平均指数衰减AdaGrad 是累加单调不减AdaFactor 是行/列分离的 rank-1 近似——后两者对数值范围管理防溢出有额外要求。自适应 LR 的形状决定是否引入额外机制LAMB/LARS 的自适应项含 layer-wise scale需要额外的归约与标量广播运算已超出纯 element-wise 范围设计时需评估是否仍走纯 V排布或引入跨核归约。四、仓库落地实现ApplyAdamWV2 源码级拆解AT-19 卡片标注使用算子: ApplyAdamWV2, ApplyRMSProp两者在仓库中都有完整的实现与测试是理解该模式的最佳参考代码。4.1 Host wrapper参数解析与 kernel 分派apply_adam_w_v2_impl.py 的入口是apply_adam_w_v2_wrapper其职责链清晰展示了 AT-19 六步骨架的 host 侧配合参数解析与校验_parse_adam_args严格校验beta1, beta2, lr, weight_decay, eps, step六个参数wrapper 断言四张输入张量 shape 一致、weight/graddtype 匹配、m/v必须为 fp32、step 1L195-L228。host 端预计算标量bc1 1.0 - beta1**step、bc2 1.0 - beta2**step、one_m_b1、one_m_b2L232-L235。按 dtype 分派 kernelweight为 bf16 走apply_adam_w_v2_kernel_bf16fp32 走apply_adam_w_v2_kernel_fp32其余 dtype 直接抛TypeErrorL245-L250。按形状选择 tile 配置_select_adam_tile_config针对三类形状分别返回 DEFAULT / LARGE_M / TARGET_K_AXIS 三套配置见 4.3 节。调用方式来自 READMEfrom apply_adam_w_v2_impl import apply_adam_w_v2_wrapper w_new, m_new, v_new apply_adam_w_v2_wrapper( weight, grad, m, v, beta10.9, beta20.999, lr1e-3, weight_decay0.01, eps1e-8, step1, )4.2 Kernel 主体双循环 valid_shape 尾块 三次 assemble以 fp32 kernel 为例其结构完整映射 AT-19 的 AdamW 计算链L100-L127M/K 双循环外层m_loop、内层n_loop对应参数矩阵的 M 轴与 K 轴均以pypto.loop构建unroll_list[1]保持结构循环不展开尾块处理每层循环计算valid_m/valid_n (dim - offset).min(tile)通过pypto.view(..., valid_shapevalid_shape)传入K 为 N_TILE 整数倍时尾块自动退化为 no-op计算链逐行写出m_new、grad_sq、v_new、m_hat、v_hat、denom、update、w_new与 AT-19 伪代码逐行对应其中div使用precision_typepypto.PrecisionType.INTRINSIC利用硬件 intrinsic 除法的吞吐优势状态写回w_new、m_new、v_new三个结果在同一 kernel 迭代内通过三次独立的pypto.assemble写回对应骨架第 6 步状态写回in-place。bf16 kernel 在此基础上增加了混合精度路径L177-L188入口处w_tile、g_tile先pypto.cast到 FP32全部中间运算保持 FP32仅参数结果在assemble前 cast 回 BF16m_new/v_new始终以 FP32 输出。这正是 AT-19 全程 FP32 计算参数最终 cast 回 BF16约束的实现证据。4.3 Tile 策略形状驱动的三套配置AT-19 卡片给出的经验值是沿参数维度K 轴分块N_TILE2048而仓库实现将其细化为按输入形状自动选择的三档配置L31-L70配置触发条件M_TILEVecTileN_TILE(K 向)设计意图DEFAULT_TILE_CONFIG通用 2D 小形状32(32, 512)1024常规路径LARGE_M_TILE_CONFIGM 7168 且 K 40967168(16, 512)1024M 大 K 小整 M 块减少循环TARGET_K_AXIS_TILE_CONFIGbf16 且M7168 且 K 40967168(16, 512)4096网络级[7168, K]K 动态 2048-24576K 向大块降 prologue 占比三个配置共同印证了 AT-19 的 tiling 原则N_TILE 存在一个平衡点——过小则 prologue循环头开销占比高、带宽利用不足过大则单个 tile 装不进 UBUnified Buffer。仓库采用 1024/4096 两档 K 向 tile 配合pypto.set_vec_tile_shapes设置 V 单元 tile并在TARGET_K_AXIS路径额外开启pypto.experimental.set_operation_options(combine_axisTrue)与vec_nbuffer_setting{DEFAULT: 4}L160-L162来压深流水。4.4 产品支持与输入规格根据 ApplyAdamWV2 README产品支持Ascend 950PR、Atlas A3 训练/推理系列、Atlas A2 训练/推理系列均支持Shape 与 dtypeweight/grad为[7168, K]的bfloat16或float32K 动态范围 2048-24576m/v恒为[7168, K]的float32输出 dtype 与输入一致m/v输出恒为 fp322026-06-11 更新kernel 已支持动态 M 与动态 K 的 2D 张量并新增level5~level7小形状/尾块验证层级[7168, K]网络形状保留 Large-M 策略以维持既有性能覆盖。五、第二个样本ApplyRMSProp 实现对照ApplyRMSProp 是 AT-19 覆盖矩阵中无一阶状态的代表其公式README与骨架第 2/3/4 步的可选项取舍形成直接对照grad_sq grad * grad ms_new ms (grad_sq - ms) * (1 - rho) mom_new mom * momentum (grad * lr) / sqrt(ms_new epsilon) var_new var - mom_newPyPTO 算子映射同样来自 README §3grad * grad计算平方梯度ms (grad_sq - ms) * (1.0 - rho)更新ms注意这是以增量形式书写的滑动平均与 AdamW 的beta*v (1-beta)*g²形式等价但写法不同pypto.sqrt(ms_new epsilon)计算分母mom * momentum (grad * lr) / sqrt(...)更新动量var - mom_new更新参数通过move(...)将结果原地写回var/ms/mom。该实现同样设置了pypto.experimental.set_operation_options(combine_axisTrue)与pypto.set_vec_tile_shapes(32, 512)与 ApplyAdamWV2 的常规路径保持一致。输入输出规格为var/ms/mom/grad均为[rows, cols]的 fp32 张量标量lr默认 0.001、rho默认 0.9、momentum默认 0.9、epsilon默认 1e-7。已知限制当前仅覆盖 fp32、输入必须是二维张量、主验证路径依赖 NPU 环境与torch_npuREADME §9。两个样本的对照价值在于同一个 AT-19 骨架通过增删偏置修正自适应学习率两个可选阶段即可派生出 AdamW 与 RMSProp 两种截然不同的优化器实现这正体现了模式抽象的复用能力。六、PyPTO 特化编程约束带宽、缓存与精度的三条红线AT-19 卡片的编程约束PyPTO 特化部分是设计落地时的硬性要求逐条展开如下6.1 全程 FP32 计算参数最终 cast 回 BF16优化器更新对数值精度高度敏感动量项在长期训练中累积二阶矩的舍入误差会通过自适应学习率被放大。因此 kernel 内部无论输入是 bf16 还是 fp32全部中间计算保持在 FP32仅在写回参数前 cast 回 bf16。仓库的 golden 参考实现 明确将这一行为定义为混合精度契约Mixed precision contractweight/grad in bf16 are cast up to fp32 internally, all math is done in fp32, weight is cast back to its original dtype on return。kernel 与 golden 必须遵守同一契约否则精度验证无法通过。6.2 必须使用 submit_before_loopTrue 实现 pipeline overlap优化器更新是典型的带宽密集型memory-bound算子每个参数元素要读入 grad、w、m、v 四个张量并写回三个张量计算量相对搬运量很小。若不使用submit_before_loopTrue循环迭代之间无法与搬运流水重叠带宽利用率不满。仓库设计规范 pypto-kernel-design-format.md 对此给出精确表述只要迭代之间存在跨块依赖例如 block loop 携带状态就必须在pypto.loop上打开submit_before_loopTrue同时强调循环必须用pypto.loop而非 Pythonfor否则会在构图期被展开、破坏框架的 tiling 与并行调度。6.3 必须配置 set_cache_policy(NONE_CACHEABLE)参数张量在大模型训练中体积很大例如[7168, 24576]的 fp32 张量若允许其驻留 L2 cache会污染缓存、挤占其他数据因此必须显式设置NONE_CACHEABLE缓存策略。这也是该约束在仓库 执行约束文档 与其他 V 型模式如缓存写回类算子中反复出现的原因。6.4 K 轴分块与 N_TILE 的平衡经验值与实现取值沿参数维度K 轴分块N_TILE2048 是经验值过小 prologue 占比高过大 UB 装不下——这条约束在仓库实现中体现为多档 N_TILE1024 常规 / 4096 网络级 K 向并给出量化的边界条件过小循环次数增多每次循环头prologue的固定开销占比升高带宽利用不足过大单个 tile 连同 m、v、w、grad 四个驻留张量超出 UB 容量。AT-19 特别强调多状态同时存在时 UB 容量约束更紧TILE 需相应缩小。AdamW 双状态m、v加 w、grad 共四张驻留而 RMSProp 也是 var/ms/mom/grad 四张因此两者的 vec tile 都收敛到 512 量级。设计 LAMB同样双状态时可复用这一经验值作为起点再按目标形状实测调整。七、精度验证golden 对照与测试入口7.1 ApplyAdamWV2 的验证体系test_apply_adam_w_v2.py 以纯 PyTorch 的 apply_adam_w_v2_golden.py 为参照覆盖level0fp32 精度与level1bf16 精度两级并附带多步偏置修正step10的数值 sanity 检查。运行方式README §Running the precision testcd repo-root source env_setup.sh python3 tests/ops/experimental/vector/ApplyAdamWV2/test_apply_adam_w_v2.py通过标准atol1e-4、rtol0.0078125测试打印[PRECISION_PASS]exit 0或[PRECISION_FAIL]exit 1。此外 golden 自带 smoke test在mv0、step1的解析可验条件下做解析核对此时update g/(sqrt(g²)eps) wd*w并将 bf16 回程误差上限放宽到4e-3以容纳 bf16 舍入apply_adam_w_v2_golden.py#L139-L148。测试用例配置见 test_cases.json。7.2 ApplyRMSProp 的验证体系ApplyRMSProp 的测试以 test_apply_rms_prop.py 为入口支持level08x8、level11024x1024、level216x16三个层级阈值atol1e-5、rtol1e-5。每个 level 构造随机输入 → 调用 kernel → 与 NumPy golden 对比 → 分别输出var_diff、ms_diff、mom_diff摘要格式如level0: shape(8, 8), var_diff0.00000000, ms_diff0.00000000, mom_diff0.00000000。README §10 还给出了典型排障路径TILE_FWK_DEVICE_ID未设置时先export TILE_FWK_DEVICE_ID0Invalid Device时用npu-smi info核对设备号精度断言失败时优先检查输入 dtype、超参与 golden 的一致性以及apply_rms_prop_impl.py中的运算顺序是否被改动。7.3 环境准备运行上述测试需要 CANN、torch_npu与pto-isa环境典型配置ApplyRMSProp README §6source /path/to/ascend-toolkit/set_env.sh export TILE_FWK_DEVICE_ID0 export PTO_TILE_LIB_CODE_PATHpto-isa 路径如需重编译安装 Python 包python3 build_ci.py -f python3 --disable_auto_execute。八、如何在 DESIGN.md 中接入 AT-19 模式按 pypto-op-design/SKILL.md 的工作流命中 AT-19 后应在 DESIGN.md 中完成以下对接范式选择在范式与设计决策中记录采用 AT-19V 型局部模式并对照 AT 索引 确认无更匹配的候选若同一计算段可由单一 API 实现按机制对比规则说明择优理由。变体确认按第三节矩阵确定目标优化器在六步骨架中的可选阶段取舍明确一阶/二阶状态、自适应 LR、衰减形式并同步确认状态张量的 dtype恒为 fp32与驻留数量决定 tile 缩放。Tiling 与配置以 N_TILE 经验值2048为起点结合目标形状与 UB 容量确定 M/K 向 tile 与 vec tile在参数表中把执行切块参数标注为tunable把硬件容量UB 上限、32B 对齐标注为constant。循环与状态写回使用pypto.loop双循环 submit_before_loopTrueunroll_list只选一个展开因子三个输出在同一迭代内独立assemble。精度契约与 golden 对齐全 FP32 中间计算 参数 cast 回原 dtype的混合精度契约并把atol/rtol阈值写入设计文档供后续验证阶段引用。至此从模式卡片到落地 kernel 再到精度验证的完整闭环已经打通——这正是 AT-19 作为局部计算模式在 PyPTO-Gym 仓库中的真实价值。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考