几种推理后端的安装及配置

TVM

TVM 的官方文档[1]给出的配置是 apache-tvm v0.11，这个版本压根跑不起来。实测需要带 --pre 参数安装。

1
2
pip install apache-tvm --pre
Successfully installed apache-tvm-0.14.dev273 ml-dtypes-0.5.0

需补充安装如下库

1
2
pip install tlcpack-nightly-cu116 --pre -f https://tlcpack.ai/wheels
pip install pytest

XLA

PyTorch 与 TensorFlow 都有 XLA 支持。这里吃了大亏，误认为 PyTorch 生态更好，于是先尝试了 PyTorch-XLA，一步一坑！

相关教程见 [1:1] [2] [3] [4]，这里只粗略记录笔者的配置过程。

  1. 为兼容低版本 torch 安装旧版 xla。
  2. 样例报 ImportError: libpython3.10.so.1.0: cannot open shared object file
  3. 检查 /usr/local/lib，只有 libpython3.10.a，原来是编译时只生成了静态库，于是重新编译安装 Python 并更新动态链接器缓存 sudo ldconfig。
  4. 重跑，仍然提示 Undefined Symbol。
  5. 到了喜闻乐见的试版本开盲盒环节，尝试换 xla 版本，但是根本找不到仓库，无奈去翻 xla 的 GCS Bucket 自行拼链接下载，反复测试数个版本，不得行！最后还是会报各种 Attribute Error 等。

被折磨习惯了，但为了把样例跑起来，无奈转投 Tensorflow-XLA。

Triton 算子库

Triton 的轻度使用，通过 torch.compile() 即可。[5]

POFO

这玩意一眼写论文凑数用的，使用 Rotor[6] 实现。

  1. (Homepage) openxla/xla: A machine learning compiler for GPUs, CPUs, and ML accelerators ↩︎ ↩︎

  2. How to run with PyTorch/XLA:GPU — PyTorch/XLA master documentation ↩︎

  3. pytorch/xla: Enabling PyTorch on XLA Devices (e.g. Google TPU) ↩︎

  4. Releases · pytorch/pytorch ↩︎

  5. 在 AI Infra 体系里，Triton 实际上并不是与本文介绍的其他后端同一层的概念，但是在当前语境下，混淆一下，问题不大。 ↩︎

  6. https://gitlab.inria.fr/hiepacs/rotor/-/tree/offload ↩︎