CANN/xla-npu 示例指南
样例运行验证【免费下载链接】xla-npuXLA-NPU 是一个面向华为昇腾NPU硬件的 XLA后端实现。本项目通过接入OpenXLA/XLA开源项目将XLA开源生态与华为 CANN软件栈集成对接JAX框架。JAX框架运行时可以直接加载XLA-NPU使得基于JAX框架开发的模型可以运行在昇腾NPU上提供推理场景图编译加速能力。项目地址: https://gitcode.com/cann/xla-npu开发者完成 XLA-NPU 安装后可通过本目录下的样例验证 XLA-NPU 功能。本代码仓提供基础算子测试和模型推理样例具体如下。样例列表目录名称功能描述opsJAX 基础算子测试样例支持 CPU 与 NPU 执行结果对比包含 Add、Tanh、Multiply、Dot、Divide 等多种算子mmoeMMOE (Multi-gate Mixture-of-Experts) 模型推理样例展示多任务学习模型在 NPU 上的推理性能qwenQwen2 模型推理样例展示大语言模型在 NPU 上的推理性能common公共工具模块包含基础运行器、配置管理和结果对比工具output样例运行输出目录存放 CPU 和 NPU 的执行结果文件目录结构说明examples/ ├── common/ # 公共工具模块 │ ├── base_runner.py # 基础运行器类 │ ├── config.py # 公共配置 │ └── utils.py # 工具函数 ├── ops/ # 基础算子测试样例 │ ├── main.py # 主入口 │ ├── config.py # 算子配置 │ ├── run_cpu.py # CPU 运行脚本 │ ├── run_npu.py # NPU 运行脚本 │ └── utils.py # 工具函数 ├── mmoe/ # MMOE 模型样例 │ ├── main.py # 主入口 │ ├── model.py # 模型定义 │ ├── config.py # 模型配置 │ ├── init_params.py # 参数初始化 │ ├── run_cpu.py # CPU 运行脚本 │ ├── run_npu.py # NPU 运行脚本 │ ├── utils.py # 工具函数 │ └── show_result.py # 结果展示 ├── qwen/ # Qwen2 模型样例 │ ├── main.py # 主入口 │ ├── model.py # 模型定义 │ ├── config.py # 模型配置 │ ├── init_params.py # 参数初始化 │ ├── run_cpu.py # CPU 运行脚本 │ ├── run_npu.py # NPU 运行脚本 │ ├── utils.py # 工具函数 │ └── test.sh # 测试脚本 └── output/ # 输出目录 ├── ops/ # ops 样例输出 ├── mmoe/ # mmoe 样例输出 └── qwen/ # qwen 样例输出快速开始环境准备确保已完成 XLA-NPU 的安装 参考项目根目录 INSTALL_GUIDE.md设置运行所需环境变量 参考项目根目录 README.md修改common/config.py中的NPU_LIBRARY_NAME可选默认为NPU_LIBRARY_NAME通过环境变量XLA_NPU_ROOT_DIR自动拼接路径如需自定义请修改为xla-npu编译产物的实际路径在参考README.md中执行source ./build/xla_npu_env后可不用手动修改安装 MMOE 模型样例所需的 Python 依赖:pip install flax运行算子测试样例cd ops python main.py可选参数:--num-runs N: 设置运行次数默认为 1--skip-cpu: 跳过 CPU 执行仅运行 NPU--skip-npu: 跳过 NPU 执行仅运行 CPU--cleanup: 测试完成后清理临时文件运行 MMOE 模型样例cd mmoe # 首次运行时如果 examples/mmoe/ 目录下没有 model_params.npz 文件则需要初始化模型参数 python main.py --init # 后续运行 python main.py可选参数:--num-runs N: 设置运行次数默认为 1--skip-cpu: 跳过 CPU 执行仅运行 NPU--skip-npu: 跳过 NPU 执行仅运行 CPU--init: 初始化模型参数--cleanup: 测试完成后清理临时文件运行 Qwen2 模型样例cd qwen # 首次运行时如果 examples/qwen/ 目录下没有 model_params.npz 文件则需要初始化模型参数 python main.py --init # 后续运行 python main.py可选参数:--num-runs N: 设置运行次数默认为 1--skip-cpu: 跳过 CPU 执行仅运行 NPU--skip-npu: 跳过 NPU 执行仅运行 CPU--init: 初始化模型参数--cleanup: 测试完成后清理临时文件结果对比样例运行后会自动对比 CPU 和 NPU 的执行结果输出:执行时间统计结果差异分析(最大绝对误差、最大相对误差等)测试结果状态(通过/失败)注意事项运行 NPU 样例前请确保 NPU 设备可用模型执行的输入和输出数据默认保存在examples/output/目录下(执行测试用例后自动生成) 可通过--cleanup清理。【免费下载链接】xla-npuXLA-NPU 是一个面向华为昇腾NPU硬件的 XLA后端实现。本项目通过接入OpenXLA/XLA开源项目将XLA开源生态与华为 CANN软件栈集成对接JAX框架。JAX框架运行时可以直接加载XLA-NPU使得基于JAX框架开发的模型可以运行在昇腾NPU上提供推理场景图编译加速能力。项目地址: https://gitcode.com/cann/xla-npu创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考