GutenOCR-3B 整网集成实战基于 PyPTO 的 RMSNorm / MRoPE / SwiGLU 算子注入指南【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本篇技术指南面向在 CANN Ascend 平台上使用 PyPTO 框架对多模态 OCR 模型做算子级整网集成的开发者以 GutenOCR-3B基于 Qwen2.5-VL 架构的 3B 级光学字符识别模型为例完整讲解其权重获取、模型下载、PyPTO 补丁注入、推理运行与精度验证的端到端流程并深入解析ask_gutenocr_3b.py脚本、modeling_qwen2_5_vl.py注入点与三个 PyPTO 融合算子RMSNorm、MRoPE、SwiGLU MLP的源码级实现原理。读完本文你将掌握在 pypto-gym 仓库中复现 Baseline 与 PTO 两种模式推理、按需启停单算子、以及读懂单算子替换阶段性能回落现象背后的原因。一、GutenOCR-3B 整网集成概览GutenOCR-3B 是一个约 30 亿参数的多模态 OCR 模型其基座为 Qwen2.5-VL 架构。pypto-gym 仓库在 modeling/transformers/gutenocr_3b/ 目录下完整承载了该模型的整网集成样例核心思路可以概括为一条公式trust_remote_codeHuggingFace 自带 Qwen2.5-VL 代码 pypto-gym 补丁PyPTO RMSNorm / MRoPE / SwiGLU 注入具体来说模型权重目录中的 Python 建模代码仍来自 HuggingFace 的rootsautomation/GutenOCR-3B而 pypto-gym 以补丁方式restore_model_patch.sh向其中注入三类 PyPTO 融合算子RMSNormBF16 优化的pypto.rms_norm融合实现替换Qwen2RMSNormMRoPEMultimodal Rotary Position Embedding替换多模态三维位置编码temporal height widthSwiGLU MLPgate_proj SiLU up_proj element-wise mul down_proj五合一融合 kernel。这三类算子的独立说明分别记录在 rms_norm/README.md、mrope/README.md 与 swiglu_mlp/README.md 中建议作为本文的配套阅读材料。二、模型权重与代码来源2.1 HuggingFace 模型模型原产地为 HuggingFace 上的rootsautomation/GutenOCR-3B。下载到本地后权重目录的定位方式有两种二选一即可环境变量export GUTENOCR_MODEL_PATH/path/to/gutenocr_3b命令行参数--model-path /path/to/gutenocr_3b在 ask_gutenocr_3b.py 的参数解析中可以看到两者的优先级关系parser.add_argument(--model-path, defaultos.environ.get(GUTENOCR_MODEL_PATH, .), help模型路径)即--model-path未显式给出时回退读取环境变量GUTENOCR_MODEL_PATH再回退到当前目录.。2.2 代码组成权重目录内的模型代码由两部分拼装而成HuggingFace 远程代码通过trust_remote_codeTrue加载其auto_map指向仓库内置的configuration_qwen2_5_vl.py与modeling_qwen2_5_vl.pypypto-gym 归档代码pypto-gym 将适配后的模型代码归档到了src/pypto_gym/transformers/gutenocr_3b/其中 modeling_qwen2_5_vl.py 头部注释明确标注了该文件的三处改造点绝对导入兼容 transformers 5.6.0PTO RMSNorm 注入通过sys.modules.get(pto_kernels)探测自定义GutenOcr3BVL*类名避免与官方类名冲突。从 config.json 可以看到加载时的auto_map定义auto_map: { AutoConfig: configuration_qwen2_5_vl.GutenOcr3BVLConfig, AutoModelForCausalLM: modeling_qwen2_5_vl.GutenOcr3BVLForConditionalGeneration, AutoModelForImageTextToText: modeling_qwen2_5_vl.GutenOcr3BVLForConditionalGeneration }这意味着推理脚本中使用的AutoModelForImageTextToText.from_pretrained(..., trust_remote_codeTrue)会实际加载 pypto-gym 的这份改造代码。2.3 关键模型配置从 config.json 中可以提取出与三个 PyPTO 算子直接相关的配置项配置项值关联算子hidden_size2048RMSNorm / SwiGLU MLPintermediate_size11008SwiGLU MLPnum_hidden_layers36全模型num_attention_heads/num_key_value_heads16 / 2MRoPErms_norm_eps1e-06RMSNormrope_parameters.mrope_section[16, 24, 24]MRoPErope_theta1000000.0MRoPEdtypebfloat16全部其中mrope_section[16, 24, 24]表示多维旋转位置编码对 temporal / height / width 三个维度的切分方式这一数值在推理脚本中还会被用于修复权重加载后的rope_scaling字段见下文第四节。三、模型下载与 PyPTO 入网适配3.1 下载模型pypto-gym 复用pypto-fused-op-integration技能中的下载脚本拉取 HuggingFace 权重python3 cannbot-skills/model/pypto-fused-op-integration/scripts/download_hf_model.py \ --model-id rootsautomation/GutenOCR-3B \ --output-dir ${GUTENOCR_MODEL_PATH}--output-dir应与GUTENOCR_MODEL_PATH保持一致保证后续加载路径统一。3.2 PyPTO 入网适配补丁注入下载完成后执行补丁脚本完成 PyPTO 注入bash cannbot-skills/model/pypto-fused-op-integration/scripts/restore_model_patch.sh \ ${GUTENOCR_MODEL_PATH} gutenocr_3b该脚本以gutenocr_3b为模型标识把 pypto-gym 归档的 modeling_qwen2_5_vl.py、configuration_qwen2_5_vl.py 与config.json还原到本地权重目录使trust_remote_code加载到的就是注入点已就位的模型代码。另外还需要准备cann_pow_patch并通过环境变量指定export CANN_POW_PATCH_PATH/path/to/cann_pow_patchask_gutenocr_3b.py 启动时会优先将该路径插入sys.path并尝试导入cann_pow_patch导入失败仅静默跳过except ImportError: pass不影响 Baseline 模式运行。四、推理运行Baseline 与 PTO 两种模式4.1 最小运行命令准备就绪后先设置两个环境变量export GUTENOCR_MODEL_PATH/path/to/gutenocr_3b export CANN_POW_PATCH_PATH/path/to/cann_pow_patch然后分别运行两种模式# Baseline默认不注入 PyPTO python3 scripts/ask_gutenocr_3b.py --prompt 你好 --device 0 # PTO 模式启用 PyPTO 融合算子 python3 scripts/ask_gutenocr_3b.py --prompt 你好 --device 0 --use_pto说明上文的scripts/即仓库内的 modeling/transformers/gutenocr_3b/完整命令为python3 modeling/transformers/gutenocr_3b/ask_gutenocr_3b.py ...。4.2 完整命令行参数ask_gutenocr_3b.py 提供了一组覆盖性能测试需求的参数参数默认值说明--prompt你好提问文本--device1NPU 卡号--model-pathGUTENOCR_MODEL_PATH环境变量模型路径--warmup5Warmup 迭代次数--use_ptoFalse启用 PyPTO 融合算子--use_compileFalse启用torch.compile--use_acl_graphFalse启用 aclgraph 模式需 torchair--use_dynamic_configFalse启用动态配置按 batch 自动选择最优算子--simple_promptFalse简化 prompt跳过apply_chat_template高性能--backendinductortorch.compilebackendinductor/npugraphs/npu--modeNonetorch.compilemodereduce-overhead/max-autotune/default--batch1Batch size--output_length50生成 token 数--report-fileNone性能报告 JSON 输出路径推理流程上脚本采用「Warmup默认 5 次每次max_new_tokens10→ 正式推理5 次取平均max_new_tokens--output_length→ 统计」的结构并通过torch.npu.synchronize()保证计时准确torch.npu.max_memory_allocated()采集峰值显存MB最终结果可落盘为 JSON 报告。4.3 PyPTO 注入的核心机制sys.modules 注册--use_pto模式的关键动作发生在脚本的 Step 22 阶段ask_gutenocr_3b.py其原理是模块别名注册将--model-path插入sys.path保证能导入权重目录下的pto_kernels模块import pto_kernels后将其以别名gutenocr_3b_pto_kernels注册进sys.modulessys.modules[gutenocr_3b_pto_kernels] pto_kernels手动模式仅--use_pto下将三个开关全部置 Truepto_kernels.USE_PTO_RMS_NORM True pto_kernels.USE_PTO_MROPE True pto_kernels.USE_PTO_SWIGLU_MLP True模型代码侧的注入点则通过sys.modules.get(pto_kernels)反向探测注意模型代码探测的是pto_kernels这个原始名字因此脚本必须确保该名字可被 import。以 modeling_qwen2_5_vl.py 中的三个注入点为例RMSNorm 注入点GutenOcr3BVLRMSNorm.forward约 L91-L101pto_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) # fallback原始 FP32 RMSNorm 实现SwiGLU MLP 注入点Qwen2MLP.forward约 L648-L653pto_kernels sys.modules.get(pto_kernels) if pto_kernels is not None and getattr(pto_kernels, USE_PTO_SWIGLU_MLP, False): return pto_kernels.swiglu_mlp_wrapper(self, x) down_proj self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_projMRoPE 注入点apply_multimodal_rotary_pos_emb约 L656-L671pto_kernels sys.modules.get(pto_kernels) if pto_kernels is not None and getattr(pto_kernels, USE_PTO_MROPE, False): return pto_kernels.mrope_pto_correct(q, k, cos, sin, mrope_section, unsqueeze_dim) # fallbackmrope_section * 2 cat/split rotate_half 的 torch 路径这种「开关变量 wrapper 桥接」的设计使得算子级 A/B 对照非常方便不改模型代码只需翻转开关即可在 PTO 与 Baseline 之间切换。4.4 torch.compile 与 aclgraph 扩展模式脚本还支持在 PyPTO 之上叠加图编译模式Step 23--use_compile调用torch.compile(model, backendargs.backend, modeargs.mode)支持reduce-overhead、max-autotune等 mode--use_acl_graph优先走 torchair 路径使用CompilerConfig配置frozen_parameterTrue与tiling_schedule_optimizeTrue然后torch.compile(model, dynamicTrue, fullgraphTrue, backendnpu_backend)若 torchair 不可用则回退到普通torch.compile。若同时启用--use_dynamic_config脚本会从dynamic_pto_config.py导入DynamicPTOConfig根据 batch 大小自动选择最优算子组合如自动开启 aclgraph 并给出预期吞吐与选择理由并同步打印各开关的最终状态。该机制对应 pto_kernels 包__init__.py中的推荐结论SwiGLU MLP所有 batch 推荐启用端到端 3%~16%MRoPE仅 Batch ≤ 4 推荐启用存在固化开销问题RMSNorm不推荐启用固化开销抵消优化收益。这也是为什么手动模式默认三算子全开而动态配置会按 batch 做差异化取舍。4.5 模型加载细节与 rope_scaling 修复模型加载使用 BF16 eager attentionmodel AutoModelForImageTextToText.from_pretrained( args.model_path, torch_dtypetorch.bfloat16, device_map{: device}, local_files_onlyTrue, trust_remote_codeTrue, attn_implementationeager )加载后脚本会遍历每个 decoder layer修复rope_scaling字段若为 None 则补上{mrope_section: [16, 24, 24], type: mrope}确保 MRoPE 位置编码行为与配置一致ask_gutenocr_3b.py。五、三个 PyPTO 融合算子的实现剖析PyPTO 算子库归档在 src/pypto_gym/ops/pypto_tensor/gutenocr_3b/包含rms_norm/、mrope/、swiglu_mlp/三个子包对外统一通过init.py 暴露三个开关与两个 wrapperUSE_PTO_SWIGLU_MLP False USE_PTO_MROPE False USE_PTO_RMS_NORM False5.1 RMSNormBF16 优化版替换对象Qwen2RMSNorm融合范围为mean(x^2) rsqrt(meaneps) x * rsqrt * weight实现方式pypto.tensor()无 shape 声明 pypto.rms_norm融合 APItile 根据 dim 动态设置见 rms_norm/README.md精度测试export TILE_FWK_DEVICE_ID0 python3 tests/ops/gutenocr_3b/rms_norm/test_rms_norm_gutenocr_3b.py。测试脚本 test_rms_norm_gutenocr_3b.py 的结构很典型从rms_norm_test_cases.json加载用例含 seed、shape、dtype、eps在 CPU 上生成 golden 参考将输入搬移到 NPU 后分别跑rms_norm_golden与rms_norm_pto_native最后用assert_allclose按 rtol/atol默认 1e-2断言同时校验输出 shape 与 dtype。测试用例来源是从模型打点采集的真实 shape/dtypeprefill 主 norm[1, seq_len, 2048]decode 主 norm[1, 1, 2048]q_norm[1, seq_len, 16, 128]k_norm[1, seq_len, 8, 128]产品支持情况Atlas A2 / A3 训练与推理系列支持Ascend 950PR 不支持。5.2 MRoPEMultimodal Rotary Position Embedding替换对象多模态三维旋转位置编码temporal height widthmrope_section[16, 24, 24]BF16实现方式目前为等价 torch 路径融合范围是concat(temporal, height, width) cos/sin rotation concat output核心调用为mrope_pto_correct与 Baseline 行为一致见 mrope/README.md精度测试export TILE_FWK_DEVICE_ID0 python3 tests/ops/gutenocr_3b/mrope/test_mrope.py。测试用例来源同样为模型打点采集q/k 形如[1, 16, seq_len, 128]/[1, 2, seq_len, 128]cos/sin 形如[3, 1, seq_len, 128]三个通道对应 temporal/height/width 三维。5.3 SwiGLU MLP融合版替换对象Qwen2MLP融合范围为gate_proj(x) SiLU up_proj(x) element-wise mul down_proj(hidden)D2048I11008BF16实现方式swiglu_mlp_fused/swiglu_mlp_fused_statickernel__init__.py中的swiglu_mlp_wrapper(mlp_module, hidden_states)完成nn.Module → kernel的桥接读取gate_proj/up_proj/down_proj权重并转置为连续内存、按需补零 bias、将输入展平为[batch*seq, hidden_size]后调用融合 kernel异常时回退到 torch 路径见init.py精度测试export TILE_FWK_DEVICE_ID0 python3 tests/ops/gutenocr_3b/swiglu_mlp/test_swiglu_mlp.py。从其 README 的状态栏可以看到当前进度为「内核适配中 / 整网集成fallback torch 路径」即整网已接入 wrapper 桥接、但融合 kernel 仍在适配中精度与性能调优尚未完成——这正是阅读源码时需要注意的边界开关置位后实际走的是 wrapper 内部 torch fallback 还是融合 kernel取决于swiglu_mlp_fused的可用性。六、环境信息与性能对比6.1 已验证环境模型集成在以下组件版本组合下完成验证来自 README 的环境信息表组件版本torch2.9.0cputorch_npu2.9.0.post2transformers5.8.1CANNAscend 910B6.2 单算子替换的性能观测README 记录了同机同输入的实测对比--prompt 你好 --device 0 --output_length 50模式命令推理耗时吞吐峰值显存baseline--prompt 你好 --device 0 --output_length 504.51s11.1 tok/s7182 MBpto (RMSNormMRoPE)--prompt 你好 --device 0 --output_length 50 --use_pto5.85s8.6 tok/s7182 MB需要特别强调的是 README 给出的结论性注记单算子替换时 PTO 比基线慢属正常现象根因是 kernel launch 开销大于单算子替换带来的收益PyPTO 的真正收益来自多算子融合如 SwiGLU MLP 将三个 Linear 融合为一个 kernel这解释了为何 pto_kernels__init__.py中只有 SwiGLU MLP 被标注为所有 batch 均推荐启用端到端 3%~16%而 MRoPE、RMSNorm 因固化开销问题需要按 batch 取舍。此外仓库还提供了批量基准脚本 bench_gutenocr_3b.sh它会依次运行 Baseline 与 PyPTO 两个阶段--output_length 100分别输出bench_baseline.json与bench_pypto.json最后用一段内嵌 Python 汇总模型加载、推理耗时、生成 token 数、吞吐与峰值显存的差值百分比适合做多轮迭代回归对比。七、仓库归档映射源码到哪里找为了便于在仓库内快速定位README 给出了「归档前 → pypto-gym」的映射关系结合仓库实际路径整理如下来源models/gutenocr_3b/目标pypto-gym 仓库scripts/modeling/transformers/gutenocr_3b/含ask_gutenocr_3b.py、bench_gutenocr_3b.sh、README.mdconfig.jsonsrc/pypto_gym/transformers/gutenocr_3b/config.jsonmodeling_qwen2_5_vl.py、configuration_qwen2_5_vl.pysrc/pypto_gym/transformers/gutenocr_3b/pto_kernels/src/pypto_gym/ops/pypto_tensor/gutenocr_3b/rms_norm/、mrope/、swiglu_mlp/配套的单算子精度测试则位于 tests/ops/gutenocr_3b/ 下按算子分子目录组织rms_norm/、mrope/、swiglu_mlp/每个目录包含 golden 参考实现、测试用例 JSON 与 pytest 测试脚本是理解算子行为与复现精度验证的最佳入口。八、小结GutenOCR-3B 整网集成样例完整展示了 PyPTO 在真实多模态模型上的落地路径通过trust_remote_code加载改造后的 Qwen2.5-VL 代码、以sys.modules别名注册机制将 PyPTO 算子库注入模型、用开关变量实现算子级 A/B 切换再以「打点采集真实 shape golden 对比」的方式逐算子验证精度。其中最有工程价值的经验是单算子替换阶段的性能回落不代表 PyPTO 无效收益验证必须放到多算子融合的端到端场景中而具体的算子取舍如 SwiGLU MLP 全开、MRoPE 限 batch≤4、RMSNorm 暂不建议应依据实测数据而非直觉判断。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 GutenOCR-3B 整网集成实战：基于 PyPTO 的 RMSNorm / MRoPE / SwiGLU 算子注入指南
GutenOCR-3B 整网集成实战基于 PyPTO 的 RMSNorm / MRoPE / SwiGLU 算子注入指南【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本篇技术指南面向在 CANN Ascend 平台上使用 PyPTO 框架对多模态 OCR 模型做算子级整网集成的开发者以 GutenOCR-3B基于 Qwen2.5-VL 架构的 3B 级光学字符识别模型为例完整讲解其权重获取、模型下载、PyPTO 补丁注入、推理运行与精度验证的端到端流程并深入解析ask_gutenocr_3b.py脚本、modeling_qwen2_5_vl.py注入点与三个 PyPTO 融合算子RMSNorm、MRoPE、SwiGLU MLP的源码级实现原理。读完本文你将掌握在 pypto-gym 仓库中复现 Baseline 与 PTO 两种模式推理、按需启停单算子、以及读懂单算子替换阶段性能回落现象背后的原因。一、GutenOCR-3B 整网集成概览GutenOCR-3B 是一个约 30 亿参数的多模态 OCR 模型其基座为 Qwen2.5-VL 架构。pypto-gym 仓库在 modeling/transformers/gutenocr_3b/ 目录下完整承载了该模型的整网集成样例核心思路可以概括为一条公式trust_remote_codeHuggingFace 自带 Qwen2.5-VL 代码 pypto-gym 补丁PyPTO RMSNorm / MRoPE / SwiGLU 注入具体来说模型权重目录中的 Python 建模代码仍来自 HuggingFace 的rootsautomation/GutenOCR-3B而 pypto-gym 以补丁方式restore_model_patch.sh向其中注入三类 PyPTO 融合算子RMSNormBF16 优化的pypto.rms_norm融合实现替换Qwen2RMSNormMRoPEMultimodal Rotary Position Embedding替换多模态三维位置编码temporal height widthSwiGLU MLPgate_proj SiLU up_proj element-wise mul down_proj五合一融合 kernel。这三类算子的独立说明分别记录在 rms_norm/README.md、mrope/README.md 与 swiglu_mlp/README.md 中建议作为本文的配套阅读材料。二、模型权重与代码来源2.1 HuggingFace 模型模型原产地为 HuggingFace 上的rootsautomation/GutenOCR-3B。下载到本地后权重目录的定位方式有两种二选一即可环境变量export GUTENOCR_MODEL_PATH/path/to/gutenocr_3b命令行参数--model-path /path/to/gutenocr_3b在 ask_gutenocr_3b.py 的参数解析中可以看到两者的优先级关系parser.add_argument(--model-path, defaultos.environ.get(GUTENOCR_MODEL_PATH, .), help模型路径)即--model-path未显式给出时回退读取环境变量GUTENOCR_MODEL_PATH再回退到当前目录.。2.2 代码组成权重目录内的模型代码由两部分拼装而成HuggingFace 远程代码通过trust_remote_codeTrue加载其auto_map指向仓库内置的configuration_qwen2_5_vl.py与modeling_qwen2_5_vl.pypypto-gym 归档代码pypto-gym 将适配后的模型代码归档到了src/pypto_gym/transformers/gutenocr_3b/其中 modeling_qwen2_5_vl.py 头部注释明确标注了该文件的三处改造点绝对导入兼容 transformers 5.6.0PTO RMSNorm 注入通过sys.modules.get(pto_kernels)探测自定义GutenOcr3BVL*类名避免与官方类名冲突。从 config.json 可以看到加载时的auto_map定义auto_map: { AutoConfig: configuration_qwen2_5_vl.GutenOcr3BVLConfig, AutoModelForCausalLM: modeling_qwen2_5_vl.GutenOcr3BVLForConditionalGeneration, AutoModelForImageTextToText: modeling_qwen2_5_vl.GutenOcr3BVLForConditionalGeneration }这意味着推理脚本中使用的AutoModelForImageTextToText.from_pretrained(..., trust_remote_codeTrue)会实际加载 pypto-gym 的这份改造代码。2.3 关键模型配置从 config.json 中可以提取出与三个 PyPTO 算子直接相关的配置项配置项值关联算子hidden_size2048RMSNorm / SwiGLU MLPintermediate_size11008SwiGLU MLPnum_hidden_layers36全模型num_attention_heads/num_key_value_heads16 / 2MRoPErms_norm_eps1e-06RMSNormrope_parameters.mrope_section[16, 24, 24]MRoPErope_theta1000000.0MRoPEdtypebfloat16全部其中mrope_section[16, 24, 24]表示多维旋转位置编码对 temporal / height / width 三个维度的切分方式这一数值在推理脚本中还会被用于修复权重加载后的rope_scaling字段见下文第四节。三、模型下载与 PyPTO 入网适配3.1 下载模型pypto-gym 复用pypto-fused-op-integration技能中的下载脚本拉取 HuggingFace 权重python3 cannbot-skills/model/pypto-fused-op-integration/scripts/download_hf_model.py \ --model-id rootsautomation/GutenOCR-3B \ --output-dir ${GUTENOCR_MODEL_PATH}--output-dir应与GUTENOCR_MODEL_PATH保持一致保证后续加载路径统一。3.2 PyPTO 入网适配补丁注入下载完成后执行补丁脚本完成 PyPTO 注入bash cannbot-skills/model/pypto-fused-op-integration/scripts/restore_model_patch.sh \ ${GUTENOCR_MODEL_PATH} gutenocr_3b该脚本以gutenocr_3b为模型标识把 pypto-gym 归档的 modeling_qwen2_5_vl.py、configuration_qwen2_5_vl.py 与config.json还原到本地权重目录使trust_remote_code加载到的就是注入点已就位的模型代码。另外还需要准备cann_pow_patch并通过环境变量指定export CANN_POW_PATCH_PATH/path/to/cann_pow_patchask_gutenocr_3b.py 启动时会优先将该路径插入sys.path并尝试导入cann_pow_patch导入失败仅静默跳过except ImportError: pass不影响 Baseline 模式运行。四、推理运行Baseline 与 PTO 两种模式4.1 最小运行命令准备就绪后先设置两个环境变量export GUTENOCR_MODEL_PATH/path/to/gutenocr_3b export CANN_POW_PATCH_PATH/path/to/cann_pow_patch然后分别运行两种模式# Baseline默认不注入 PyPTO python3 scripts/ask_gutenocr_3b.py --prompt 你好 --device 0 # PTO 模式启用 PyPTO 融合算子 python3 scripts/ask_gutenocr_3b.py --prompt 你好 --device 0 --use_pto说明上文的scripts/即仓库内的 modeling/transformers/gutenocr_3b/完整命令为python3 modeling/transformers/gutenocr_3b/ask_gutenocr_3b.py ...。4.2 完整命令行参数ask_gutenocr_3b.py 提供了一组覆盖性能测试需求的参数参数默认值说明--prompt你好提问文本--device1NPU 卡号--model-pathGUTENOCR_MODEL_PATH环境变量模型路径--warmup5Warmup 迭代次数--use_ptoFalse启用 PyPTO 融合算子--use_compileFalse启用torch.compile--use_acl_graphFalse启用 aclgraph 模式需 torchair--use_dynamic_configFalse启用动态配置按 batch 自动选择最优算子--simple_promptFalse简化 prompt跳过apply_chat_template高性能--backendinductortorch.compilebackendinductor/npugraphs/npu--modeNonetorch.compilemodereduce-overhead/max-autotune/default--batch1Batch size--output_length50生成 token 数--report-fileNone性能报告 JSON 输出路径推理流程上脚本采用「Warmup默认 5 次每次max_new_tokens10→ 正式推理5 次取平均max_new_tokens--output_length→ 统计」的结构并通过torch.npu.synchronize()保证计时准确torch.npu.max_memory_allocated()采集峰值显存MB最终结果可落盘为 JSON 报告。4.3 PyPTO 注入的核心机制sys.modules 注册--use_pto模式的关键动作发生在脚本的 Step 22 阶段ask_gutenocr_3b.py其原理是模块别名注册将--model-path插入sys.path保证能导入权重目录下的pto_kernels模块import pto_kernels后将其以别名gutenocr_3b_pto_kernels注册进sys.modulessys.modules[gutenocr_3b_pto_kernels] pto_kernels手动模式仅--use_pto下将三个开关全部置 Truepto_kernels.USE_PTO_RMS_NORM True pto_kernels.USE_PTO_MROPE True pto_kernels.USE_PTO_SWIGLU_MLP True模型代码侧的注入点则通过sys.modules.get(pto_kernels)反向探测注意模型代码探测的是pto_kernels这个原始名字因此脚本必须确保该名字可被 import。以 modeling_qwen2_5_vl.py 中的三个注入点为例RMSNorm 注入点GutenOcr3BVLRMSNorm.forward约 L91-L101pto_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) # fallback原始 FP32 RMSNorm 实现SwiGLU MLP 注入点Qwen2MLP.forward约 L648-L653pto_kernels sys.modules.get(pto_kernels) if pto_kernels is not None and getattr(pto_kernels, USE_PTO_SWIGLU_MLP, False): return pto_kernels.swiglu_mlp_wrapper(self, x) down_proj self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_projMRoPE 注入点apply_multimodal_rotary_pos_emb约 L656-L671pto_kernels sys.modules.get(pto_kernels) if pto_kernels is not None and getattr(pto_kernels, USE_PTO_MROPE, False): return pto_kernels.mrope_pto_correct(q, k, cos, sin, mrope_section, unsqueeze_dim) # fallbackmrope_section * 2 cat/split rotate_half 的 torch 路径这种「开关变量 wrapper 桥接」的设计使得算子级 A/B 对照非常方便不改模型代码只需翻转开关即可在 PTO 与 Baseline 之间切换。4.4 torch.compile 与 aclgraph 扩展模式脚本还支持在 PyPTO 之上叠加图编译模式Step 23--use_compile调用torch.compile(model, backendargs.backend, modeargs.mode)支持reduce-overhead、max-autotune等 mode--use_acl_graph优先走 torchair 路径使用CompilerConfig配置frozen_parameterTrue与tiling_schedule_optimizeTrue然后torch.compile(model, dynamicTrue, fullgraphTrue, backendnpu_backend)若 torchair 不可用则回退到普通torch.compile。若同时启用--use_dynamic_config脚本会从dynamic_pto_config.py导入DynamicPTOConfig根据 batch 大小自动选择最优算子组合如自动开启 aclgraph 并给出预期吞吐与选择理由并同步打印各开关的最终状态。该机制对应 pto_kernels 包__init__.py中的推荐结论SwiGLU MLP所有 batch 推荐启用端到端 3%~16%MRoPE仅 Batch ≤ 4 推荐启用存在固化开销问题RMSNorm不推荐启用固化开销抵消优化收益。这也是为什么手动模式默认三算子全开而动态配置会按 batch 做差异化取舍。4.5 模型加载细节与 rope_scaling 修复模型加载使用 BF16 eager attentionmodel AutoModelForImageTextToText.from_pretrained( args.model_path, torch_dtypetorch.bfloat16, device_map{: device}, local_files_onlyTrue, trust_remote_codeTrue, attn_implementationeager )加载后脚本会遍历每个 decoder layer修复rope_scaling字段若为 None 则补上{mrope_section: [16, 24, 24], type: mrope}确保 MRoPE 位置编码行为与配置一致ask_gutenocr_3b.py。五、三个 PyPTO 融合算子的实现剖析PyPTO 算子库归档在 src/pypto_gym/ops/pypto_tensor/gutenocr_3b/包含rms_norm/、mrope/、swiglu_mlp/三个子包对外统一通过init.py 暴露三个开关与两个 wrapperUSE_PTO_SWIGLU_MLP False USE_PTO_MROPE False USE_PTO_RMS_NORM False5.1 RMSNormBF16 优化版替换对象Qwen2RMSNorm融合范围为mean(x^2) rsqrt(meaneps) x * rsqrt * weight实现方式pypto.tensor()无 shape 声明 pypto.rms_norm融合 APItile 根据 dim 动态设置见 rms_norm/README.md精度测试export TILE_FWK_DEVICE_ID0 python3 tests/ops/gutenocr_3b/rms_norm/test_rms_norm_gutenocr_3b.py。测试脚本 test_rms_norm_gutenocr_3b.py 的结构很典型从rms_norm_test_cases.json加载用例含 seed、shape、dtype、eps在 CPU 上生成 golden 参考将输入搬移到 NPU 后分别跑rms_norm_golden与rms_norm_pto_native最后用assert_allclose按 rtol/atol默认 1e-2断言同时校验输出 shape 与 dtype。测试用例来源是从模型打点采集的真实 shape/dtypeprefill 主 norm[1, seq_len, 2048]decode 主 norm[1, 1, 2048]q_norm[1, seq_len, 16, 128]k_norm[1, seq_len, 8, 128]产品支持情况Atlas A2 / A3 训练与推理系列支持Ascend 950PR 不支持。5.2 MRoPEMultimodal Rotary Position Embedding替换对象多模态三维旋转位置编码temporal height widthmrope_section[16, 24, 24]BF16实现方式目前为等价 torch 路径融合范围是concat(temporal, height, width) cos/sin rotation concat output核心调用为mrope_pto_correct与 Baseline 行为一致见 mrope/README.md精度测试export TILE_FWK_DEVICE_ID0 python3 tests/ops/gutenocr_3b/mrope/test_mrope.py。测试用例来源同样为模型打点采集q/k 形如[1, 16, seq_len, 128]/[1, 2, seq_len, 128]cos/sin 形如[3, 1, seq_len, 128]三个通道对应 temporal/height/width 三维。5.3 SwiGLU MLP融合版替换对象Qwen2MLP融合范围为gate_proj(x) SiLU up_proj(x) element-wise mul down_proj(hidden)D2048I11008BF16实现方式swiglu_mlp_fused/swiglu_mlp_fused_statickernel__init__.py中的swiglu_mlp_wrapper(mlp_module, hidden_states)完成nn.Module → kernel的桥接读取gate_proj/up_proj/down_proj权重并转置为连续内存、按需补零 bias、将输入展平为[batch*seq, hidden_size]后调用融合 kernel异常时回退到 torch 路径见init.py精度测试export TILE_FWK_DEVICE_ID0 python3 tests/ops/gutenocr_3b/swiglu_mlp/test_swiglu_mlp.py。从其 README 的状态栏可以看到当前进度为「内核适配中 / 整网集成fallback torch 路径」即整网已接入 wrapper 桥接、但融合 kernel 仍在适配中精度与性能调优尚未完成——这正是阅读源码时需要注意的边界开关置位后实际走的是 wrapper 内部 torch fallback 还是融合 kernel取决于swiglu_mlp_fused的可用性。六、环境信息与性能对比6.1 已验证环境模型集成在以下组件版本组合下完成验证来自 README 的环境信息表组件版本torch2.9.0cputorch_npu2.9.0.post2transformers5.8.1CANNAscend 910B6.2 单算子替换的性能观测README 记录了同机同输入的实测对比--prompt 你好 --device 0 --output_length 50模式命令推理耗时吞吐峰值显存baseline--prompt 你好 --device 0 --output_length 504.51s11.1 tok/s7182 MBpto (RMSNormMRoPE)--prompt 你好 --device 0 --output_length 50 --use_pto5.85s8.6 tok/s7182 MB需要特别强调的是 README 给出的结论性注记单算子替换时 PTO 比基线慢属正常现象根因是 kernel launch 开销大于单算子替换带来的收益PyPTO 的真正收益来自多算子融合如 SwiGLU MLP 将三个 Linear 融合为一个 kernel这解释了为何 pto_kernels__init__.py中只有 SwiGLU MLP 被标注为所有 batch 均推荐启用端到端 3%~16%而 MRoPE、RMSNorm 因固化开销问题需要按 batch 取舍。此外仓库还提供了批量基准脚本 bench_gutenocr_3b.sh它会依次运行 Baseline 与 PyPTO 两个阶段--output_length 100分别输出bench_baseline.json与bench_pypto.json最后用一段内嵌 Python 汇总模型加载、推理耗时、生成 token 数、吞吐与峰值显存的差值百分比适合做多轮迭代回归对比。七、仓库归档映射源码到哪里找为了便于在仓库内快速定位README 给出了「归档前 → pypto-gym」的映射关系结合仓库实际路径整理如下来源models/gutenocr_3b/目标pypto-gym 仓库scripts/modeling/transformers/gutenocr_3b/含ask_gutenocr_3b.py、bench_gutenocr_3b.sh、README.mdconfig.jsonsrc/pypto_gym/transformers/gutenocr_3b/config.jsonmodeling_qwen2_5_vl.py、configuration_qwen2_5_vl.pysrc/pypto_gym/transformers/gutenocr_3b/pto_kernels/src/pypto_gym/ops/pypto_tensor/gutenocr_3b/rms_norm/、mrope/、swiglu_mlp/配套的单算子精度测试则位于 tests/ops/gutenocr_3b/ 下按算子分子目录组织rms_norm/、mrope/、swiglu_mlp/每个目录包含 golden 参考实现、测试用例 JSON 与 pytest 测试脚本是理解算子行为与复现精度验证的最佳入口。八、小结GutenOCR-3B 整网集成样例完整展示了 PyPTO 在真实多模态模型上的落地路径通过trust_remote_code加载改造后的 Qwen2.5-VL 代码、以sys.modules别名注册机制将 PyPTO 算子库注入模型、用开关变量实现算子级 A/B 切换再以「打点采集真实 shape golden 对比」的方式逐算子验证精度。其中最有工程价值的经验是单算子替换阶段的性能回落不代表 PyPTO 无效收益验证必须放到多算子融合的端到端场景中而具体的算子取舍如 SwiGLU MLP 全开、MRoPE 限 batch≤4、RMSNorm 暂不建议应依据实测数据而非直觉判断。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考