测试 Python 算子#

FlagGems 使用 pytest 来驱动算子精度测试和性能基准测试。 项目使用 Triton 的 triton.testing.do_bench 来执行内核层级的性能评估。

1. 算子精度测试#

要在特定的后端硬件(如 CUDA)上运行测试:

pytest tests/test_${name}.py

下面的命令执行在 CPU 上的精度测试:

pytest tests/test_${case}.py --ref cpu

2. 在具体模型下执行精度测试#

pytest examples/${name}_test.py

3. 测试算子的性能#

在 CUDA 平台上测试算子的性能:

pytest benchmark/test_foo.py -s

下面的命令对算子执行端到端的性能基准测试:

pytest benchmark/test_foo.py -s --ref cpu

4. 测试套件支持的命令行选项#

除了 pytest 自带的选项(例如用于关键字过滤的 -k、用于按标记过滤的 -m, 以及在首次失败时终止的 -x)之外,tests/conftest.py 还注册了若干 FlagGems 特有的选项,可以更精细地控制精度测试的运行方式。

选项取值描述
--ref<device>(默认)、cpu用于计算参照结果的设备,FlagGems 的输出将与其进行比对。
--quick开关精简 tests/accuracy_utils.py 中使用的形状/数据类型组合,以加快冒烟测试。
--recordnone(默认)、log、json将测试执行的详细信息记录到日志文件或 JSON 报告中。
--output路径使用 --record json 时的输出文件路径,默认为 accuracy_result.json。
--collect-marks路径不实际运行测试,而是将所有被收集到的测试用例的标记信息导出到指定的 YAML 文件。

例如,下面的命令会在 CPU 上快速检查某个算子,并保留一份 JSON 报告:

pytest tests/test_softmax.py --ref cpu --quick --record json --output softmax_result.json

基于标记的过滤#

每一个精度测试都会带有一个以其所覆盖算子命名的标记,例如 tests/test_fliplr.py 中的 @pytest.mark.fliplr。借助这一机制, 你可以按算子名称运行测试,而不必关心测试用例分布在哪些文件中:

pytest tests/ -m "fliplr or rms_norm"

报告与覆盖率集成#

FlagGems 的 CI 流程使用 pytest-md-report 插件,将一次测试运行转换为 Markdown 格式的摘要报告;同时使用 coverage 包裹整个测试过程以收集单元测试覆盖率:

coverage run -m pytest -s tests/test_foo.py \
    --md-report --md-report-verbose=1 --md-report-output=summary.md
coverage html

关于覆盖率报告的更多细节,请参阅单元测试覆盖率。