[RFC]: TENPU API 一致性接入 ATK-TransformerEngineNPU-AtomGit

已关闭

[RFC]: TENPU API 一致性接入 ATK #34

clc2025创建于 8月24日关闭于 5 天前

clc2025成员

8月24日创建

状态（Status）： Draft
作者（Authors）： @clc2025
创建日期（Created）： 2026-08-24
更新日期（Updated）： 2026-08-24
相关 Issue/PR： #123（关联 Issue/PR 以便追踪背景）

1. 概述

1.1 简介

TENPU（TransformerEngineNPU）是对标 NVIDIA Transformer Engine（NVTE）的昇腾实现。为保证 TENPU 与 NVTE 的 API 一致性（签名、语义、数值），本提案设计一套基于 ATK 算子测试工具 的一致性看护体系：以 Megatron-Core 对 TE 的静态调用分析（110 个唯一 API / 507 处调用）为需求源，对需要数值看护的接口建立 ATK golden 闭环（NPU 执行 TENPU ↔ CUDA 采集 NVTE golden ↔ 数值 + API 契约对比），对纯接口兼容类用轻量 pytest 契约检查。

1.2 动机

一致性风险：TENPU 与 NVTE 是两个独立实现，API 签名、导入路径、数值行为可能不一致（如 apply_rotary_pos_emb 导入路径、Float8Tensor legacy shim、LayerNorm dummy、ops.GroupedLinear 未导出、npu_layernorm_grad 与 aclnnLayerNormBackward 数值差）。
看护需求：Megatron 依赖 TENPU 全链路（训练/推理/checkpoint），需持续防回归。已有 tests/pytorch 单元测试，但缺 与 NVTE 的跨实现数值一致性 看护。
现状：人工逐个比对成本高、易漏；需要一个可重复、可 CI 化的看护框架。ATK 提供用例生成、多节点执行、精度/性能对比、报告输出能力，是本看护体系的载体。
不做此提案的影响：TENPU 与 NVTE 的 API/数值差异只能靠人工发现，回归无法自动化拦截。

1.3 目标

目标：

对需要数值看护的 TE 接口建立 ATK golden 闭环（NPU vs CUDA-NVTE）
对 [仅API兼容] 接口建立 pytest 契约检查
建立可复用的模块模板（7 件套）与踩坑约束清单
支撑后续"新增 API 的 ATK 适配自动化生成"

非目标：

性能/内存看护（P3，数值闭环稳定后补）
多卡分布式集合通信的 ATK 看护（框架单进程限制，走 pytest 分布式）
[语义+性能+内存] 链路级行为（checkpoint/offload/fp8 上下文，表格注明由上层 case 覆盖）

2. 用例分析

用例	功能点	关键要求
算子数值一致性	前向/反向/梯度/状态 vs NVTE golden	容差内一致（mare/mixed_tolerance），确定性可复现
API 签名/导入一致性	符号路径、构造参数、枚举、序列化	与 NVTE 对齐；缺失导出（如 legacy shim）可被拦截
FP8 量化精度	量化/反量化 round-trip、raw 字节、scale 元数据	raw 逐位一致（跨实现），dequant 容差
优化器/融合 kernel	FusedAdam/SGD 多步后参数/状态；multi_tensor	多步演化 vs NVTE；容量含近零病态
fuser 基础 op	ops.BasicLinear/Bias/LayerNorm/RMSNorm	与 module 走同 kernel → 免独立 guard
CI 回归	预设提交/合入前跑 small/real 两层	快（small）/ 全（real），golden 可重采

约束：NPU 需 910B 系列；CUDA 主机需 pinned NVTE + GPU；golden 与运行端必须同版适配器代码（契约 exact-compare）。

3. 方案设计

3.1 总体方案

以 Megatron-Core r0.18.0 TE API 分析表（Megatron_Core_r0.18.0_TE_API_Analysis_一致性看护细化.xlsx）为需求源，按表格分类分层看护：

看护类型	数量	手段
[仅API兼容]	33	pytest 契约检查（符号/签名/枚举/序列化，CPU 跑，不用 ATK）
[精度+性能] / [三项看护]	58	ATK golden 闭环（本体系核心）
[语义+性能+内存]	19	链路级 case 覆盖（checkpoint/offload/fp8 上下文）

核心模式：no-CPU-reference golden 闭环

atk case（生成用例 JSON）
  → collect_nvte_golden.py（CUDA 主机，跑 NVTE，产出 input.bin + output_*.pt + manifest）
  → atk task --input_data（NPU 主机，跑 TENPU，accuracy_load 节点对比 golden）
  → 数值（mare/mixed_tolerance）+ 契约（exact）对比 → 报告

适配器不实现 CPU 参考（新模块一律如此），数值一致性完全由 golden 承担
每模块 7 件套：<name>.yaml / generate_<name>.py / function_<name>.py / collect_nvte_golden.py / nodes_npu_vs_golden.yaml / result_schema.json / README.md

3.2 技术选型

决策点	方案	理由
对比基准	CUDA 采集 NVTE golden（accuracy_load 节点）	ATK `BenchmarkDevice` 无 GPU，不用其内置 single_bm
CPU 参考	不实现（no-CPU-reference 模板）	避免双份实现漂移；golden 即真值
精度标准	mare（mean/max 双门）默认，mixed_tolerance（ULP backstop）对近零病态	近零元素相对误差爆炸（rope/grouped_linear/fused_optimizer 均踩过），mixed_tolerance 最稳健
FP8 raw 对比	位重解释 uint8 exact	TENPU `_data` 是 float8、NVTE 是 uint8，需统一位视图
契约	结构化字段 exact-compare（数据依赖浮点不入契约）	amax/scale 等浮点 repr 跨端可能不同，只留结构/格式字段
优化器/多步	固定步数（small=7/real=3）	状态充分演化；golden 与运行同步数

符号级判定（精度看护范围收窄）：

cast_master_weights_to_fp8 / QuantizedTensor / 通信算子（AllReduce/ReduceScatter/gather/reduce_scatter_along_first_dim）/ 抽象基类：无数值分叉目标 → 免 ATK 精度 guard，走 pytest 语义/API 兼容
BasicLinear/Bias/LayerNorm：与 module 同 kernel → 免独立 ops guard（module guard 覆盖）

3.3 功能与性能设计

模块清单（已建 10 个）：linear、rmsnorm、layernorm_linear、grouped_linear、dot_product_attention、rope、fused_activation、quantization、fused_optimizer、multi_tensor。

核心流程（以量化为例）：

atk case 生成用例（small/real 两层，128 对齐预留）
CUDA collect_nvte_golden.py 产 golden（含 _sanitize_args 去 numpy 标量）
NPU atk task --input_data 加载同一 input.bin，TENPU 执行，与 golden 对比
报告输出通过率/是否达标

契约设计：只含跨端必然一致的字段（fp8_format/shapes/dtypes 等），数据依赖浮点（amax/scale）排除——因契约 exact-compare，任何端差异字段都会误报。

性能：small 层 presubmit 快跑，real 层合入前跑；用例数 8~32/模块，单模块数秒~分钟。

3.4 安全隐私与 DFX 设计

兼容性：ATK 版本适配（共享比较器 _comparators.py 兼容老 ATK）；导入路径按报告 API 清单字段取、不假设顶层导出（FusedAdam 在 optimizers、Float8Tensor 需 legacy shim）。
可测试性：确定性可复现（--deterministic-mode/seed）；golden 可重采。
可维护性：7 件套模板统一；踩坑约束沉淀为清单（见 §3.5 约束）。
可靠性：近零病态用 mixed_tolerance 吸收；kernel 非确定用容差兜底。

已沉淀的硬约束（自动化生成前置）：

yaml inputs 块数必须与 generator 产出一致（否则 atk case 静态统计 KeyError）
契约字段跨端字节一致（数据依赖浮点不入契约；golden 需同版重采）
比较器按算子性质选（近零病态→mixed_tolerance）
api_type 不能含 "tensor"/"method" 子串（ATK 子串匹配误判需额外 input 文件）
NVTE pybind 绑定函数只接受位置参数（multi_tensor_l2norm 已踩）
用例矩阵满足 API 超参约束（nesterov 需 momentum>0 等）
导入路径按报告 API 清单取、不假设顶层导出

3.5 编程与调用设计

3.5.1 编程模型基本设计

开发环境：Ascend 910B3（NPU）+ CANN + torch_npu；CUDA 主机（NVTE golden）；transformer_engine（TENPU）editable 安装。
开发约束：适配器 function_<name>.py 只走设备执行路径，cpu backend 显式 raise。
可验收设计：模块在确定性下 NPU vs golden 通过率达标（mare/mixed_tolerance）。

3.5.2 接口定义与设计（模块文件）

文件	职责	关键点
`<name>.yaml`	设计 + 容差	`standard.acc` 选比较器（`<name>_mare`/`mixed_tolerance`）
`generate_<name>.py`	两层 P0 矩阵	固定 seed；yaml inputs 块数一致
`function_<name>.py`	设备适配器 + 契约 + 比较器	`register_module_comparators`；BaseApi 拒 cpu；符号按报告路径导入
`collect_nvte_golden.py`	CUDA golden 采集	`_sanitize_args` 去 numpy 标量；manifest 记录 SHA
`nodes_npu_vs_golden.yaml`	npu + cpu_nvte 节点	accuracy + accuracy_load
`result_schema.json`	输出语义映射	与 RESULT_SCHEMA 一致
`README.md`	范围 + 验证清单 + golden 重采铁律

调用示例（fused_optimizer 流程）：

atk case -f tests/atk/module/fused_optimizer/fused_optimizer.yaml -p tests/atk/module/fused_optimizer/generate_fused_optimizer.py
# CUDA 主机
python tests/atk/module/fused_optimizer/collect_nvte_golden.py --cases result/.../all_fused_optimizer.json --input-dir ./fused_optimizer_inputs --golden-dir ./fused_optimizer_goldens
# NPU 主机
atk task -c result/.../all_fused_optimizer.json -n tests/atk/module/fused_optimizer/nodes_npu_vs_golden.yaml -p tests/atk/module/fused_optimizer/function_fused_optimizer.py --input_data ./fused_optimizer_inputs -tk accuracy

3.5.3 编程手册设计

各模块 README 即使用手册（范围/工作流/设备验证清单/golden 重采规则）；总体流程见 tests/atk/REPORT.md §5。

4. 测试设计

单元：比较器（mare/mixed_tolerance/allclose）判定逻辑；_sanitize_args。
集成（NPU standalone）：适配器前向/反向与参考逐位一致（如 LayerNorm vs functional.layer_norm）。
端到端（golden 闭环）：每模块 atk task NPU vs CUDA-NVTE golden，确定性下通过率。
- 已闭环：linear/rmsnorm/layernorm_linear/grouped_linear/dot_product_attention/rope/fused_activation/multi_tensor（8/8）
- 待补：quantization（9/10，case9 stale golden 已重采，10/10 复测待低精设备）、fused_optimizer（8/8 待复测）
回归：改适配器契约/输出 → 重采 golden → 重跑（铁律写入各 README）。

5. 缺点和风险

golden 与代码版本耦合：契约 exact-compare，golden 过期表现为"个别 case contract size 不等"（已踩两次）→ README 铁律 + 自动化生成前置。
近零元素病态：mare 对近零相对误差爆炸（rope/grouped_linear/fused_optimizer/multi_tensor 均踩）→ 统一 mixed_tolerance。
框架限制：单进程无法做多卡集合通信；性能看护未开。
实现成本：每模块 7 件套 + 设备验证（需 CUDA + NPU 双环境）；人力投入集中在"确认能走到修改点 + 数值对比"。
跨端符号差异：FP8 dtype（tex vs torch.dtype）、raw 表示（float8 vs uint8）、pybind 位置参数——均需逐模块适配（已沉淀为约束清单）。

6. 现有技术

ATK（/home/c00893505/atk-tenpu/ATK）：昇腾算子测试工具，提供用例生成/多节点执行/精度对比/报告。
NVTE：CUDA 参考实现（golden 真值来源）；架构与 TENPU 对齐（ops/module 分层）。
torch_npu / op-plugin：CANN 算子层；aclnnLayerNorm(Backward) 等；npu_layernorm_grad 等独立自定义 op。
借鉴：no-CPU-reference + golden 闭环思路来自 ATK 的 accuracy_load 通路；符号级判定（无数值分叉→免精度 guard）是对报告分类的修正。

7. 未解决问题

新增 API 的 ATK 适配自动化生成（P2）：基于 7 条硬约束 + 模板，自动生成 7 件套（减少手写）。
性能/内存看护（P3）：数值闭环稳定后补耗时/峰值内存对比。
分布式集合通信看护：待 ATK 支持分布式节点后按需建模块（当前 pytest 分布式覆盖）。

附录

参考资料：Megatron_Core_r0.18.0_TE_API_Analysis_一致性看护细化.xlsx；tests/atk/REPORT.md；TENPU_LayerNorm_设计文档.md；ATK atk/case_generator、atk/tasks/executors。
术语表：golden 闭环 = CUDA 采 NVTE + NPU 对比；no-CPU-reference = 适配器不实现 CPU 参考；mare/mixed_tolerance = ATK 单标杆/混合容差精度标准。
文档更新计划：随模块增加更新模块清单与符号级判定；自动化生成落地后更新 §7。