GutenOCR-3B RMSNorm 的 PyPTO 融合算子集成指南从模型打点到整网注入【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym导读本文围绕 PyPTO-Gym 仓库中 GutenOCR-3B 模型 RMSNorm 的 PyPTO 融合算子集成实践展开完整梳理了从模型打点采集真实 shape/dtype、编写pypto.tensor()pypto.rms_norm融合 kernel到单算子精度验证与整网注入的完整流程。读完本文你将掌握在 Ascend NPU 上用 PyPTO 替换 HuggingFace Transformers 模型中RMSNorm模块的落地方法、tile 动态设置技巧以及当前实现的精度状态与性能边界。一、集成背景为什么替换 RMSNormGutenOCR-3B 是基于 Qwen2.5-VL 架构的多模态 OCR 模型其文本分支在多个位置使用了 RMSNorm在仓库的 modeling_qwen2_5_vl.py 中以GutenOcr3BVLRMSNorm类实现包括视觉分支的ln_q第 167 行与norm1/norm2第 323-324 行文本分支的input_layernorm/post_attention_layernorm第 777-778 行与最终norm第 851 行。这些模块默认路径均为 PyTorch 基础算子组合先hidden_states.pow(2).mean(-1)求方差再rsqrt归一化并与权重相乘见 modeling_qwen2_5_vl.py 的 fallback 分支。PyPTO-Gym 的目标是验证将这类“仅使用 torch 基础算子”的原始实现即 README 技术说明中的场景 A替换为 PyPTO 融合算子是否能在 NPU 上获得收益。RMSNorm 集成即其中一个最小化验证样例D2048、BF16单算子场景。二、产品支持情况与适用前提关联文档 rms_norm/README.md 明确了该实现的产品支持矩阵产品支持情况Ascend 950PR不支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持需要强调的是这是仓库文档声明的产品适配结论实际运行前应以当前 CANN 与 PyPTO 版本在目标硬件上的验证结果为准。从modeling/transformers/gutenocr_3b/README.md记录的环境信息看该项目此前在Ascend 910B torch 2.9.0 torch_npu 2.9.0.post2 transformers 5.8.1环境下完成过整网集成与性能对比。三、算子实现无 shape 声明的pypto.rms_norm融合 kernel3.1 核心实现RMSNorm 的 PyPTO 实现位于 rms_norm_impl.py完整代码如下import logging import torch import pypto _logger logging.getLogger(__name__) PYPTO_AVAILABLE True PYPTO_KERNEL_AVAILABLE True pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU}) def rms_norm_kernel(hidden_states: pypto.tensor(), weight: pypto.tensor(), output: pypto.tensor(), eps): rank hidden_states.dim tile_shapes [128 for _ in range(rank)] tile_shapes[-1] 2048 pypto.set_vec_tile_shapes(*tile_shapes) y pypto.rms_norm(hidden_states, weight, eps) output[:] y def rms_norm_pto_native(hidden_states, weight, eps1e-6): out torch.empty_like(hidden_states) rms_norm_kernel(hidden_states, weight, out, eps) return out3.2 逐行解读pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU})将 Python 函数编译为 NPU 上执行的 kernel。run_mode指定运行模式为 NPU编译产物可直接在昇腾设备上运行。参数全部声明为pypto.tensor()且不带 shape这正是关联文档技术说明中“无 shape 声明”的含义——kernel 通过运行时张量的实际维度确定形状使得同一个 kernel 可以同时服务 prefill[1, seq_len, 2048]与 decode[1, 1, 2048]两种 shape无需按 seq_len 分别编译。tile 根据 dim 动态设置rank hidden_states.dim tile_shapes [128 for _ in range(rank)] tile_shapes[-1] 2048 pypto.set_vec_tile_shapes(*tile_shapes)先以hidden_states.dim获取输入张量的维度数例如 3 维时得到rank3随后将除最后一维外的所有维度 tile 设为 128最后一维固定为 2048——这与模型隐藏维度 D2048 完全对齐确保向量运算沿隐藏维度整行切分。pypto.rms_norm(hidden_states, weight, eps)调用 PyPTO 的 RMSNorm 融合 API内部完成mean(x^2) rsqrt(meaneps) x * rsqrt * weight的融合计算。output[:] y将计算结果写回预先分配的输出张量避免 kernel 内部隐式分配便于上层在整网中复用输出缓冲区。3.3 与 torch 参考实现的语义对应Golden 参考实现位于 rms_norm_golden_gutenocr_3b.py其数学公式为RMS sqrt(mean(x^2) eps) output (x / RMS) * weight对应代码为variance hidden_states.pow(2).mean(-1, keepdimTrue) hidden_states hidden_states * torch.rsqrt(variance eps) return weight * hidden_statesPyPTO kernel 输出的正是与上述公式等价的融合结果这也是后续精度对比的基准。四、测试用例从模型打点采集的真实 shape4.1 shape/dtype 来源关联文档说明测试用例的 shape 全部来自模型打点采集的真实数据而非人工编造prefill 主 norm[1, seq_len, 2048]decode 主 norm[1, 1, 2048]q_norm[1, seq_len, 16, 128]k_norm[1, seq_len, 8, 128]其中 prefill/decode 主 norm 对应文本分支的input_layernorm/post_attention_layernorm等q_norm/k_norm 则对应注意力中 Q/K 的逐头归一化。4 维的 q_norm/k_norm 用例也解释了 kernel 中rank动态获取、逐维设置 tile 的设计动机。4.2 测试用例 JSON测试用例以 JSON 驱动当前包含两个用例用例 ID描述hidden_statesweightepsrtol/atolrms_norm_001minimal prefill(1, 1, 2048)fp16[1, 1, 2048]fp16[2048]fp161e-61e-2rms_norm_002prefillseq_len2(1, 2, 2048)fp16[1, 2, 2048]fp16[2048]fp161e-61e-2每个用例包含seed随机数种子保证可复现、input输入 shape/dtype/eps、output期望输出 shape/dtype以及rtol/atol精度阈值。注意 JSON 中 dtype 映射支持float16/float32/bfloat16三种测试脚本通过dtype_map完成字符串到 torch dtype 的转换。4.3 运行测试关联文档给出的运行方式export TILE_FWK_DEVICE_ID0 python3 tests/ops/gutenocr_3b/test_rms_norm.py注意README 中记录的命令路径对应仓库早期归档位置在当前仓库结构中该测试实际位于 tests/ops/gutenocr_3b/rms_norm/test_rms_norm_gutenocr_3b.py完整运行方式为export TILE_FWK_DEVICE_ID0 python3 tests/ops/gutenocr_3b/rms_norm/test_rms_norm_gutenocr_3b.py # 或 pytest tests/ops/gutenocr_3b/rms_norm/test_rms_norm_gutenocr_3b.py -v -sTILE_FWK_DEVICE_ID指定使用的 NPU 设备号测试脚本会通过torch.npu.set_device(device_id)完成设备绑定见 test_rms_norm_gutenocr_3b.py。4.4 测试执行逻辑测试脚本 test_rms_norm_gutenocr_3b.py 的执行流程环境准备读取TILE_FWK_DEVICE_ID默认 0并设置 NPU 设备加载用例解析同目录下的rms_norm_test_cases.json构造输入按用例 seed 在 CPU 上用torch.randn生成 hidden_states 与 weight再.npu()搬移到设备Golden 对比分别计算rms_norm_goldentorch 参考与rms_norm_pto_nativePyPTO 实现的输出精度断言计算max diff对 bfloat16 输出先转 float32 再使用assert_allclose按rtol1e-2, atol1e-2校验同时断言输出 shape 与 dtype 与用例声明一致结果输出以[PRECISION_PASS]标记精度通过。该测试同时验证了 BF16 路径README 声称的 D2048、BF16 场景——这正是__init__.py中“BF16 优化实现”注释所指向的优化目标。五、整网集成如何注入到 Transformers 模型5.1 注入机制RMSNorm 的整网注入采用模块替换 环境开关的方式不修改 HF 原始建模文件在 src/pypto_gym/ops/pypto_tensor/gutenocr_3b/init.py 中定义开关USE_PTO_RMS_NORM False当前默认关闭需显式开启通过rms_norm_wrapper(hidden_states, weight, eps)桥接rms_norm_pto_native直接调用 PyPTO kernel在 modeling_qwen2_5_vl.py 的GutenOcr3BVLRMSNorm.forward中运行时探测sys.modules.get(pto_kernels)若存在且USE_PTO_RMS_NORM为真则走pto_kernels.rms_norm_wrapper(...)否则回退到 torch 基础算子路径。def forward(self, hidden_states: torch.Tensor) - torch.Tensor: # PyPTO RMS Norm injection pto_kernels sys.modules.get(pto_kernels) if pto_kernels is not None and getattr(pto_kernels, USE_PTO_RMS_NORM, False): return pto_kernels.rms_norm_wrapper(hidden_states, self.weight, self.variance_epsilon) input_dtype hidden_states.dtype hidden_states hidden_states.to(torch.float32) variance hidden_states.pow(2).mean(-1, keepdimTrue) hidden_states hidden_states * torch.rsqrt(variance self.variance_epsilon) return self.weight * hidden_states.to(input_dtype)这种设计保证PyPTO 环境不可用时自动走 torch fallback整网功能不受影响。5.2 ACLGraph 状态关联文档技术说明中“ACLGraph 未注册 torch.library”——即当前实现尚未通过torch.library注册为可供图模式ACLGraph直接调度的算子仅以 Python 层包装函数方式注入。这对应 README 状态栏中的“⏳ ACLGraph”进行中也意味着当前集成方式更接近 eager 模式调用。5.3 整网运行与已知边界整网验证可参考 modeling/transformers/gutenocr_3b/README.md 的流程模型权重下载、restore_model_patch.sh入网适配、--use_pto开关运行。该文档同时记录了 RMSNorm 单算子替换的真实性能边界单算子替换时 PTO 比基线慢属正常kernel launch 开销 单算子收益。收益来自多算子融合。这与 rms_norm/init.py 中“⚠️ 不推荐启用固化开销抵消优化”的注释相互印证RMSNorm 单算子集成验证了正确性单算子精度 ✅、整网集成 ✅但尚未带来端到端收益性能调优 ⏳。集成目标更偏重于作为 PyPTO 融合算子能力的最小验证样例而非当前即用的性能优化手段。六、技术说明与状态总结6.1 技术说明表关联文档给出的技术说明汇总如下项目说明场景A — 原始实现只使用 torch 基础算子实现pypto.tensor()无 shape 声明 pypto.rms_norm融合 APItile 根据 dim 动态设置ACLGraph未注册 torch.library6.2 集成状态状态项结论单算子精度✅ 通过fp16/BF16 用例rtol/atol1e-2整网集成✅ 已完成模块替换 开关注入ACLGraph⏳ 进行中未注册 torch.library性能调优⏳ 进行中单算子收益被 kernel launch 开销抵消七、仓库内延伸阅读若希望继续深入可对照阅读以下路径rms_norm_impl.pyRMSNorm PyPTO kernel 实现rms_norm_test_cases.json驱动测试的真实 shape 用例rms_norm_golden_gutenocr_3b.pytorch 参考实现modeling_qwen2_5_vl.py整网注入点GutenOcr3BVLRMSNorm及use_kernel_forward_from_hub(RMSNorm)gutenocr_3b/init.pyUSE_PTO_RMS_NORM开关与 wrapper 桥接层modeling/transformers/gutenocr_3b/README.md整网运行、入网适配与环境信息。本文所介绍的“场景 A纯 torch 基础算子→ PyPTO 融合算子”替换模式同样适用于该目录下其他已验证算子如 SwiGLU MLP、MRoPE可作为在 PyPTO 框架下进行算子级模型适配的通用参考。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 GutenOCR-3B RMSNorm 的 PyPTO 融合算子集成指南：从模型打点到整网注入
GutenOCR-3B RMSNorm 的 PyPTO 融合算子集成指南从模型打点到整网注入【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym导读本文围绕 PyPTO-Gym 仓库中 GutenOCR-3B 模型 RMSNorm 的 PyPTO 融合算子集成实践展开完整梳理了从模型打点采集真实 shape/dtype、编写pypto.tensor()pypto.rms_norm融合 kernel到单算子精度验证与整网注入的完整流程。读完本文你将掌握在 Ascend NPU 上用 PyPTO 替换 HuggingFace Transformers 模型中RMSNorm模块的落地方法、tile 动态设置技巧以及当前实现的精度状态与性能边界。一、集成背景为什么替换 RMSNormGutenOCR-3B 是基于 Qwen2.5-VL 架构的多模态 OCR 模型其文本分支在多个位置使用了 RMSNorm在仓库的 modeling_qwen2_5_vl.py 中以GutenOcr3BVLRMSNorm类实现包括视觉分支的ln_q第 167 行与norm1/norm2第 323-324 行文本分支的input_layernorm/post_attention_layernorm第 777-778 行与最终norm第 851 行。这些模块默认路径均为 PyTorch 基础算子组合先hidden_states.pow(2).mean(-1)求方差再rsqrt归一化并与权重相乘见 modeling_qwen2_5_vl.py 的 fallback 分支。PyPTO-Gym 的目标是验证将这类“仅使用 torch 基础算子”的原始实现即 README 技术说明中的场景 A替换为 PyPTO 融合算子是否能在 NPU 上获得收益。RMSNorm 集成即其中一个最小化验证样例D2048、BF16单算子场景。二、产品支持情况与适用前提关联文档 rms_norm/README.md 明确了该实现的产品支持矩阵产品支持情况Ascend 950PR不支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持需要强调的是这是仓库文档声明的产品适配结论实际运行前应以当前 CANN 与 PyPTO 版本在目标硬件上的验证结果为准。从modeling/transformers/gutenocr_3b/README.md记录的环境信息看该项目此前在Ascend 910B torch 2.9.0 torch_npu 2.9.0.post2 transformers 5.8.1环境下完成过整网集成与性能对比。三、算子实现无 shape 声明的pypto.rms_norm融合 kernel3.1 核心实现RMSNorm 的 PyPTO 实现位于 rms_norm_impl.py完整代码如下import logging import torch import pypto _logger logging.getLogger(__name__) PYPTO_AVAILABLE True PYPTO_KERNEL_AVAILABLE True pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU}) def rms_norm_kernel(hidden_states: pypto.tensor(), weight: pypto.tensor(), output: pypto.tensor(), eps): rank hidden_states.dim tile_shapes [128 for _ in range(rank)] tile_shapes[-1] 2048 pypto.set_vec_tile_shapes(*tile_shapes) y pypto.rms_norm(hidden_states, weight, eps) output[:] y def rms_norm_pto_native(hidden_states, weight, eps1e-6): out torch.empty_like(hidden_states) rms_norm_kernel(hidden_states, weight, out, eps) return out3.2 逐行解读pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU})将 Python 函数编译为 NPU 上执行的 kernel。run_mode指定运行模式为 NPU编译产物可直接在昇腾设备上运行。参数全部声明为pypto.tensor()且不带 shape这正是关联文档技术说明中“无 shape 声明”的含义——kernel 通过运行时张量的实际维度确定形状使得同一个 kernel 可以同时服务 prefill[1, seq_len, 2048]与 decode[1, 1, 2048]两种 shape无需按 seq_len 分别编译。tile 根据 dim 动态设置rank hidden_states.dim tile_shapes [128 for _ in range(rank)] tile_shapes[-1] 2048 pypto.set_vec_tile_shapes(*tile_shapes)先以hidden_states.dim获取输入张量的维度数例如 3 维时得到rank3随后将除最后一维外的所有维度 tile 设为 128最后一维固定为 2048——这与模型隐藏维度 D2048 完全对齐确保向量运算沿隐藏维度整行切分。pypto.rms_norm(hidden_states, weight, eps)调用 PyPTO 的 RMSNorm 融合 API内部完成mean(x^2) rsqrt(meaneps) x * rsqrt * weight的融合计算。output[:] y将计算结果写回预先分配的输出张量避免 kernel 内部隐式分配便于上层在整网中复用输出缓冲区。3.3 与 torch 参考实现的语义对应Golden 参考实现位于 rms_norm_golden_gutenocr_3b.py其数学公式为RMS sqrt(mean(x^2) eps) output (x / RMS) * weight对应代码为variance hidden_states.pow(2).mean(-1, keepdimTrue) hidden_states hidden_states * torch.rsqrt(variance eps) return weight * hidden_statesPyPTO kernel 输出的正是与上述公式等价的融合结果这也是后续精度对比的基准。四、测试用例从模型打点采集的真实 shape4.1 shape/dtype 来源关联文档说明测试用例的 shape 全部来自模型打点采集的真实数据而非人工编造prefill 主 norm[1, seq_len, 2048]decode 主 norm[1, 1, 2048]q_norm[1, seq_len, 16, 128]k_norm[1, seq_len, 8, 128]其中 prefill/decode 主 norm 对应文本分支的input_layernorm/post_attention_layernorm等q_norm/k_norm 则对应注意力中 Q/K 的逐头归一化。4 维的 q_norm/k_norm 用例也解释了 kernel 中rank动态获取、逐维设置 tile 的设计动机。4.2 测试用例 JSON测试用例以 JSON 驱动当前包含两个用例用例 ID描述hidden_statesweightepsrtol/atolrms_norm_001minimal prefill(1, 1, 2048)fp16[1, 1, 2048]fp16[2048]fp161e-61e-2rms_norm_002prefillseq_len2(1, 2, 2048)fp16[1, 2, 2048]fp16[2048]fp161e-61e-2每个用例包含seed随机数种子保证可复现、input输入 shape/dtype/eps、output期望输出 shape/dtype以及rtol/atol精度阈值。注意 JSON 中 dtype 映射支持float16/float32/bfloat16三种测试脚本通过dtype_map完成字符串到 torch dtype 的转换。4.3 运行测试关联文档给出的运行方式export TILE_FWK_DEVICE_ID0 python3 tests/ops/gutenocr_3b/test_rms_norm.py注意README 中记录的命令路径对应仓库早期归档位置在当前仓库结构中该测试实际位于 tests/ops/gutenocr_3b/rms_norm/test_rms_norm_gutenocr_3b.py完整运行方式为export TILE_FWK_DEVICE_ID0 python3 tests/ops/gutenocr_3b/rms_norm/test_rms_norm_gutenocr_3b.py # 或 pytest tests/ops/gutenocr_3b/rms_norm/test_rms_norm_gutenocr_3b.py -v -sTILE_FWK_DEVICE_ID指定使用的 NPU 设备号测试脚本会通过torch.npu.set_device(device_id)完成设备绑定见 test_rms_norm_gutenocr_3b.py。4.4 测试执行逻辑测试脚本 test_rms_norm_gutenocr_3b.py 的执行流程环境准备读取TILE_FWK_DEVICE_ID默认 0并设置 NPU 设备加载用例解析同目录下的rms_norm_test_cases.json构造输入按用例 seed 在 CPU 上用torch.randn生成 hidden_states 与 weight再.npu()搬移到设备Golden 对比分别计算rms_norm_goldentorch 参考与rms_norm_pto_nativePyPTO 实现的输出精度断言计算max diff对 bfloat16 输出先转 float32 再使用assert_allclose按rtol1e-2, atol1e-2校验同时断言输出 shape 与 dtype 与用例声明一致结果输出以[PRECISION_PASS]标记精度通过。该测试同时验证了 BF16 路径README 声称的 D2048、BF16 场景——这正是__init__.py中“BF16 优化实现”注释所指向的优化目标。五、整网集成如何注入到 Transformers 模型5.1 注入机制RMSNorm 的整网注入采用模块替换 环境开关的方式不修改 HF 原始建模文件在 src/pypto_gym/ops/pypto_tensor/gutenocr_3b/init.py 中定义开关USE_PTO_RMS_NORM False当前默认关闭需显式开启通过rms_norm_wrapper(hidden_states, weight, eps)桥接rms_norm_pto_native直接调用 PyPTO kernel在 modeling_qwen2_5_vl.py 的GutenOcr3BVLRMSNorm.forward中运行时探测sys.modules.get(pto_kernels)若存在且USE_PTO_RMS_NORM为真则走pto_kernels.rms_norm_wrapper(...)否则回退到 torch 基础算子路径。def forward(self, hidden_states: torch.Tensor) - torch.Tensor: # PyPTO RMS Norm injection pto_kernels sys.modules.get(pto_kernels) if pto_kernels is not None and getattr(pto_kernels, USE_PTO_RMS_NORM, False): return pto_kernels.rms_norm_wrapper(hidden_states, self.weight, self.variance_epsilon) input_dtype hidden_states.dtype hidden_states hidden_states.to(torch.float32) variance hidden_states.pow(2).mean(-1, keepdimTrue) hidden_states hidden_states * torch.rsqrt(variance self.variance_epsilon) return self.weight * hidden_states.to(input_dtype)这种设计保证PyPTO 环境不可用时自动走 torch fallback整网功能不受影响。5.2 ACLGraph 状态关联文档技术说明中“ACLGraph 未注册 torch.library”——即当前实现尚未通过torch.library注册为可供图模式ACLGraph直接调度的算子仅以 Python 层包装函数方式注入。这对应 README 状态栏中的“⏳ ACLGraph”进行中也意味着当前集成方式更接近 eager 模式调用。5.3 整网运行与已知边界整网验证可参考 modeling/transformers/gutenocr_3b/README.md 的流程模型权重下载、restore_model_patch.sh入网适配、--use_pto开关运行。该文档同时记录了 RMSNorm 单算子替换的真实性能边界单算子替换时 PTO 比基线慢属正常kernel launch 开销 单算子收益。收益来自多算子融合。这与 rms_norm/init.py 中“⚠️ 不推荐启用固化开销抵消优化”的注释相互印证RMSNorm 单算子集成验证了正确性单算子精度 ✅、整网集成 ✅但尚未带来端到端收益性能调优 ⏳。集成目标更偏重于作为 PyPTO 融合算子能力的最小验证样例而非当前即用的性能优化手段。六、技术说明与状态总结6.1 技术说明表关联文档给出的技术说明汇总如下项目说明场景A — 原始实现只使用 torch 基础算子实现pypto.tensor()无 shape 声明 pypto.rms_norm融合 APItile 根据 dim 动态设置ACLGraph未注册 torch.library6.2 集成状态状态项结论单算子精度✅ 通过fp16/BF16 用例rtol/atol1e-2整网集成✅ 已完成模块替换 开关注入ACLGraph⏳ 进行中未注册 torch.library性能调优⏳ 进行中单算子收益被 kernel launch 开销抵消七、仓库内延伸阅读若希望继续深入可对照阅读以下路径rms_norm_impl.pyRMSNorm PyPTO kernel 实现rms_norm_test_cases.json驱动测试的真实 shape 用例rms_norm_golden_gutenocr_3b.pytorch 参考实现modeling_qwen2_5_vl.py整网注入点GutenOcr3BVLRMSNorm及use_kernel_forward_from_hub(RMSNorm)gutenocr_3b/init.pyUSE_PTO_RMS_NORM开关与 wrapper 桥接层modeling/transformers/gutenocr_3b/README.md整网运行、入网适配与环境信息。本文所介绍的“场景 A纯 torch 基础算子→ PyPTO 融合算子”替换模式同样适用于该目录下其他已验证算子如 SwiGLU MLP、MRoPE可作为在 PyPTO 框架下进行算子级模型适配的通用参考。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考