SuperAPI AI 工具集
合作ENAI 行业资讯收藏
Skill 广场ort-transformers-gpu-pytest

ort-transformers-gpu-pytest

测试与质量

查看项目

更新于 10-16 12:51

技能介绍

ort-transformers-gpu-pytest 是 ONNX Runtime 仓库中一项可复用的测试操作技能，目标是在同时安装了 torch 的机器上，针对 GPU 编译的 wheel 运行 onnxruntime/test/python/transformers 下的 Python transformers 测试，并验证真实的 cuDNN/flash SDPA 调度。技能内容侧重三个容易被反复踩到的坑，以及一个防止回归被隐藏的测试开关。

一、用途

该技能用于在已安装 GPU 版 ONNX Runtime wheel 的环境中，正确运行 transformers 相关 Python 测试，并确认 SDPA（scaled dot-product attention）实际走的是 cuDNN/flash 路径而非 MATH 回退路径。测试环境可能同时安装了 torch，这会引入库版本遮蔽问题。

二、主要能力

  1. 解决源码目录遮蔽问题。从仓库根目录运行 pytest 时，仓库根下的 ./onnxruntime/ 源码目录会出现在 sys.path[0]，导致 import onnxruntime 解析到没有编译扩展的源码包，而非虚拟环境 site-packages 中安装的 wheel，表现为 ModuleNotFoundError: No module named 'onnxruntime.capi' 或 onnxruntime 导入过程中的 AttributeError。修复方式是使用 mktemp -d 创建一个中立的、私有的临时工作目录，不要使用仓库根目录，也不要用裸的共享 /tmp，然后通过绝对路径指向测试文件，并将 transformers 测试辅助目录加入 PYTHONPATH。使用私有 mktemp -d 而非 bare /tmp 的原因是 pytest 会把当前工作目录加入 sys.path，且 Python 启动时会自动导入其中的 sitecustomize.py/usercustomize.py，共享机器上他人预置的脚本可能被当作任意代码执行。

  2. 解决库遮蔽问题（LD_PRELOAD 固定库版本）。pip 安装的 torch 会自带一套 CUDA runtime 与 cuDNN，位于 site-packages/nvidia/*/lib 下。如果 ONNX Runtime 是针对不同 CUDA/cuDNN 版本构建的，动态加载器解析到哪一套就先加载哪一套，torch 的旧库可能遮蔽 ONNX Runtime dlopen 的库，导致 SDPA 解码层静默路由到 MATH 而非 CUDNN_FLASH_ATTENTION，或出现 libcudnn.so.9 无法打开、未定义符号、cuDNN 版本不匹配等错误，或加载到与构建时不符的 CUDA runtime。修复方式是用 LD_PRELOAD 预先加载 ONNX Runtime 构建所针对的系统 CUDA runtime 与 cuDNN，并把这些目录加入 LD_LIBRARY_PATH。示例中使用 libcudart.so.12 与 libcudnn.so.9，前者对 CUDA 12.4 和 12.9 的 SONAME 均为主版本号，固定到 12.9 文件可强制正确的次版本。

  3. 确认真实 cuDNN SDPA 调度。数值正确并不能证明走了 cuDNN SDPA 路径，MATH 回退能产生相同结果。正确做法是观察 ONNX Runtime 自身的路由决策：在创建 InferenceSession 之前设置 ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1，ONNX 域 Attention 内核会输出形如 SdpaKernel=CUDNN_FLASH_ATTENTION 的调试行。需要在文件描述符级别捕获 stdout，而不是使用 contextlib.redirect_stdout，因为该行由 C++ 写入原生 fd-1，Python 级别的重定向无法拦截，会得到 dispatched=None 的静默假阴性。建议参考 onnxruntime 自身的 _CaptureStdout 做法，用 os.dup2 重定向 fd-1 到临时文件。解析时应使用 re.search（单节点解码探针适用），多注意力节点图应使用 re.findall 检查每个 token，避免后续节点的 MATH 回退被前面节点的 cuDNN 命中掩盖。相比读取 torch.backends.cudnn.version() 或任何库版本检查，观察调度读取的是 ONNX Runtime 实际的路由决策，更可靠。

  4. 不可跳过的金丝雀开关 ORT_TEST_REQUIRE_CUDNN_SDPA。仅靠观察调度来门控解码测试存在失败模式：若该层静默回归不再选择 cuDNN，观察结果为“未调度”，所有解码测试会被跳过并显示绿色，掩盖回归。用环境变量 ORT_TEST_REQUIRE_CUDNN_SDPA=1 关闭该漏洞：设置时调度断言变为不可跳过，在最小已知良好配置上出现 MATH 回退或不调度会直接失败而非跳过；未设置时回退到普通跳过保护，不会误报。该变量供操作者在存在已知良好 GPU CI 环节时导出使用，但当前 ONNX Runtime 流水线定义中没有任何环节导出它，因此在本项目 CI 中无效果，仅在开发者手动设置的本地运行中生效，不应在测试文档字符串中描述为 CI 已实施的强制手段。验证金丝雀有效性的方式是：通过 CUDA provider 的 sdpa_kernel provider 选项强制 MATH 位掩码（16），在 ORT_TEST_REQUIRE_CUDNN_SDPA=1 下确认失败信息为 AssertionError: 'CUDNN_FLASH_ATTENTION' != 'MATH'。

三、使用方式

完整运行块示例：

WORKDIR=$(mktemp -d); cd "$WORKDIR" source /abs/repo/.venv/bin/activate export LD_PRELOAD=/abs/cuda12.9/lib64/libcudart.so.12:/abs/cudnn9.8/lib/libcudnn.so.9 export LD_LIBRARY_PATH=/abs/cuda12.9/lib64:/abs/cudnn9.8/lib export PYTHONPATH=/abs/repo/onnxruntime/test/python/transformers F=/abs/repo/onnxruntime/test/python/transformers/<file>.py

python -m pytest "$F" -v ORT_TEST_REQUIRE_CUDNN_SDPA=1 python -m pytest "$F" -v

注意检查通过数量而非仅看退出码。pytest -v 即使所有测试都被跳过（无 CUDA、@skipUnless(ml_dtypes) 条件不满足等）也会退出 0，日志看起来像通过证据但什么也没证明。要求通过数量非零且无非预期跳过，并注意 pytest 退出码 5 表示“未收集到测试”（通常是路径错误或 -k 过滤问题，而非成功）。建议将输出重定向到日志（2>&1 | tee "$WORKDIR/gpu_run.log"），并写入 mktemp -d 创建的目录内，而非可预测的 /tmp 路径。

四、适用场景与限制

适用于在本地或指定 GPU 环境中验证 ONNX Runtime transformers 测试真实调度到 cuDNN SDPA；适用于排查源码遮蔽、库遮蔽、假绿等问题。限制在于：该技能属于操作经验总结，具体库路径、CUDA/cuDNN 版本需按实际环境替换；ORT_TEST_REQUIRE_CUDNN_SDPA 在当前 CI 中无效果，仅本地手动设置时生效；文中涉及的测试文件、支持判定函数（如 cudnn_decode_supported）为示意，需结合自身测试套件实现。