安装 TVM Unity 编译器

TVM Unity，Apache TVM 的最新开发成果，是构建 MLC LLM 所必需的。其功能包括：

  • 无需调优即可即时生成高性能 CPU/GPU 代码；

  • 通过设计实现动态形状和符号形状跟踪；

  • 支持推理和训练；

  • 高效的 Python 优先编译器实现。作为具体示例，MLC LLM 编译是使用其 API 在纯 Python 中实现的。

TVM Unity 可以直接从预构建的开发包安装，也可以从源代码构建。

选项 1：预构建包

提供了 Apache TVM Unity 的 nightly 预构建 Python 包。

备注

❗ 每当使用 Python 时，强烈建议使用 conda 来管理独立的 Python 环境，以避免缺少依赖、版本不兼容和包冲突的问题。

conda activate your-environment
python -m pip install --pre -U -f https://mlc.ai/wheels mlc-ai-nightly-cpu

备注

如果遇到 GLIBC 未找到的问题，请在 conda 中安装最新版本的 glibc：

conda install -c conda-forge libgcc-ng

选项 2：从源代码构建

虽然通常建议始终使用预构建的 TVM Unity，但如果您需要更多自定义，则可能需要从源代码构建。注意： 只有在您熟悉 C++、CMake、LLVM、Python 和其他相关系统的复杂性时，才应尝试此操作。

Details

步骤 1：设置构建依赖项。 要从源代码构建，您需要确保满足以下构建依赖项：

  • CMake >= 3.24

  • LLVM >= 15 - 对于 ROCm 6.1，请安装 LLVM>=17；对于 ROCm 6.2，请安装 LLVM>=18。

  • Git

  • （可选）CUDA >= 11.8（针对 NVIDIA GPU）

  • （可选）Metal（针对 Apple GPU，如 M1 和 M2）

  • （可选）Vulkan（针对 NVIDIA、AMD、Intel 和移动 GPU）

  • （可选）OpenCL（针对 NVIDIA、AMD、Intel 和移动 GPU）

备注

  • 要针对 NVIDIA GPU，需要 CUDA 或 Vulkan（推荐使用 CUDA）；

  • 对于 AMD 和 Intel GPU，Vulkan 是必需的；

  • 当针对 Apple（macOS、iOS、iPadOS）时，Metal 是必需的依赖项；

  • 一些 Android 设备仅支持 OpenCL，但大多数设备支持 Vulkan。

管理依赖项的最简单方法是通过 conda，它维护了一组跨平台的工具链，包括 LLVM。要创建这些构建依赖项的环境，只需使用：

在 conda 中设置构建依赖项
# make sure to start with a fresh environment
conda env remove -n tvm-build-venv
# create the conda environment with build dependency
conda create -n tvm-build-venv -c conda-forge \
    "llvmdev>=15" \
    "cmake>=3.24" \
    git \
    python=3.11
# enter the build environment
conda activate tvm-build-venv

步骤 2：配置和构建。 建议使用基于 git 的标准工作流程来下载 Apache TVM Unity，然后在 config.cmake 中指定构建要求：

从 GitHub 下载 TVM Unity
# clone from GitHub
git clone --recursive https://github.com/mlc-ai/relax.git tvm-unity && cd tvm-unity
# create the build directory
rm -rf build && mkdir build && cd build
# specify build requirements in `config.cmake`
cp ../cmake/config.cmake .

备注

暂时使用 mlc-ai/relax，它包含一些我们将上游到 Apache TVM 的 unity 分支 的临时重要更改。

希望通过将它们附加到配置文件的末尾来专门调整以下标志：

在 config.cmake 中配置构建
# controls default compilation flags
echo "set(CMAKE_BUILD_TYPE RelWithDebInfo)" >> config.cmake
# LLVM is a must dependency
echo "set(USE_LLVM \"llvm-config --ignore-libllvm --link-static\")" >> config.cmake
echo "set(HIDE_PRIVATE_SYMBOLS ON)" >> config.cmake
# GPU SDKs, turn on if needed
echo "set(USE_CUDA   OFF)" >> config.cmake
echo "set(USE_METAL  OFF)" >> config.cmake
echo "set(USE_VULKAN OFF)" >> config.cmake
echo "set(USE_OPENCL OFF)" >> config.cmake
# FlashInfer related, requires CUDA w/ compute capability 80;86;89;90
echo "set(USE_FLASHINFER OFF)" >> config.cmake
echo "set(FLASHINFER_CUDA_ARCHITECTURES YOUR_CUDA_COMPUTE_CAPABILITY_HERE)" >> config.cmake
echo "set(CMAKE_CUDA_ARCHITECTURES YOUR_CUDA_COMPUTE_CAPABILITY_HERE)" >> config.cmake

备注

HIDE_PRIVATE_SYMBOLS 配置选项，用于启用 -fvisibility=hidden 标志。此标志有助于防止 TVM 和 PyTorch 之间潜在的符号冲突。这些冲突是由于框架附带了不同版本的 LLVM 而引起的。

CMAKE_BUILD_TYPE 控制默认编译标志：

  • Debug 设置 -O0 -g

  • RelWithDebInfo 设置 -O2 -g -DNDEBUG （推荐）

  • Release 设置 -O3 -DNDEBUG

备注

如果您使用的是 CUDA 且计算能力高于 80，则需要使用 set(USE_FLASHINFER ON) 进行构建。否则，在运行时可能会遇到 Cannot find PackedFunc 问题。

要检查您的 CUDA 计算能力，可以使用 nvidia-smi --query-gpu=compute_cap --format=csv。

一旦 config.cmake 被相应地编辑，请使用以下命令启动构建：

使用 cmake 和 cmake 构建 libtvm
cmake .. && cmake --build . --parallel $(nproc)

成功的构建应该在 /path-tvm-unity/build/ 目录下生成 libtvm 和 libtvm_runtime。

离开构建环境 tvm-build-venv 后，有两种方法可以将成功的构建安装到您的环境中：

export PYTHONPATH=/path-to-tvm-unity/python:$PYTHONPATH

验证 TVM 安装

使用具有多种语言绑定的编译器基础结构可能会容易出错。因此，强烈建议在使用前验证 TVM Unity 安装。

步骤 1：定位 TVM Python 包。 以下命令可以帮助确认 TVM 是否已正确安装为 Python 包，并提供 TVM Python 包的位置：

>>> python -c "import tvm; print(tvm.__file__)"
/some-path/lib/python3.11/site-packages/tvm/__init__.py

步骤 2：确认使用的 TVM 库。 当维护多个 TVM 构建或安装时，双重检查 Python 包是否使用正确的 libtvm 变得非常重要，可以使用以下命令：

>>> python -c "import tvm; print(tvm._ffi.base._LIB)"
<CDLL '/some-path/lib/python3.11/site-packages/tvm/libtvm.dylib', handle 95ada510 at 0x1030e4e50>

步骤 3：反射 TVM 构建选项。 有时当下游应用程序失败时，很可能是由于错误的 TVM 提交或错误的构建标志造成的。要找出原因，以下命令将很有帮助：

>>> python -c "import tvm; print('\n'.join(f'{k}: {v}' for k, v in tvm.support.libinfo().items()))"
... # Omitted less relevant options
GIT_COMMIT_HASH: 4f6289590252a1cf45a4dc37bce55a25043b8338
HIDE_PRIVATE_SYMBOLS: ON
USE_LLVM: llvm-config --link-static
LLVM_VERSION: 15.0.7
USE_VULKAN: OFF
USE_CUDA: OFF
CUDA_VERSION: NOT-FOUND
USE_OPENCL: OFF
USE_METAL: ON
USE_ROCM: OFF

备注

GIT_COMMIT_HASH 指示 TVM 构建的确切提交，可以通过 https://github.com/mlc-ai/relax/commit/$GIT_COMMIT_HASH 在 GitHub 上找到。

步骤 4：检查设备检测。 有时了解 TVM 是否能够检测到您的设备可能会有所帮助，可以使用以下命令：

>>> python -c "import tvm; print(tvm.metal().exist)"
True # or False
>>> python -c "import tvm; print(tvm.cuda().exist)"
False # or True
>>> python -c "import tvm; print(tvm.vulkan().exist)"
False # or True

请注意，上述命令验证了本地机器上是否存在实际设备，以便 TVM 运行时（而不是编译器）能够正确执行。但是，TVM 编译器可以在不需要物理设备的情况下执行编译任务。只要必要的工具链（例如 NVCC）可用，即使没有实际设备，TVM 也支持交叉编译。