安装说明¶

注意： 最新的安装说明已嵌入到 FBGEMM 仓库中的一组脚本中，位于 setup_env.bash。

安装 FBGEMM_GPU 的一般步骤如下：

建立一个隔离的构建环境。
为仅 CPU、CUDA 或 ROCm 运行时设置工具链。
安装 PyTorch。
安装 FBGEMM_GPU 包。
运行安装后检查。

在安装 FBGEMM_GPU 之前，请检查 FBGEMM 发行版兼容性表，以确保您使用的先决硬件和软件与您计划安装的 FBGEMM_GPU 版本兼容。

设置仅 CPU 环境¶

按照设置独立的构建环境中的说明设置 Conda 环境，然后按照安装 Python 库进行操作。

设置 CUDA 环境¶

FBGEMM_GPU 的 CUDA 版本需要在计算机上安装 NVIDIA GPU，并安装有效的 NVIDIA 驱动程序；否则，该库将回退到运行运算符的 CPU 版本。

并非所有硬件架构都得到 FBGEMM_GPU CUDA 包的原生支持；请参阅 FBGEMM 发行版兼容性表了解详情。可以通过从头开始构建包来实现对其他架构的支持，但这不能保证有效（尤其是对于较旧的架构）。

安装 NVIDIA 驱动程序¶

在所有其他环境设置之前，必须在系统上安装 NVIDIA 显示驱动程序。由 NVIDIA 和 PyTorch 提供的步骤是进行此操作最权威的说明。可以使用 nvidia-smi 命令验证驱动程序设置。

nvidia-smi

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.76       Driver Version: 515.76       CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA A10G         Off  | 00000000:00:1E.0 Off |                    0 |
|  0%   31C    P0    59W / 300W |      0MiB / 23028MiB |      2%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

设置 CUDA Docker 容器和 Conda 环境¶

建议（但不强制）通过 Docker 设置安装和运行 FBGEMM_GPU，以实现 CUDA 环境的隔离和可重复性。

需要安装 NVIDIA-Docker 运行时才能将驱动程序暴露给容器。由 PyTorch 提供的安装步骤详细介绍了如何实现此目的。

完成此操作后，请按照 CUDA Docker 镜像中的说明拉取 CUDA Docker 镜像并启动容器。

从那里，可以使用 Conda 构建其余的运行时环境。按照设置独立的构建环境中的说明设置 Conda 环境，然后按照安装 Python 库进行操作。

安装 CUDA 运行时¶

如果操作系统/Docker 环境尚未包含完整的 CUDA 运行时，请按照安装 CUDA 中的说明在 Conda 环境中安装 CUDA 工具包。

设置 ROCm 环境¶

FBGEMM_GPU 的 ROCm 版本需要在计算机上安装 AMD GPU，并安装有效的 AMDGPU 驱动程序；否则，该库将回退到运行运算符的 CPU 版本。

安装 AMDGPU 驱动程序¶

在所有其他环境设置之前，必须在系统上安装 AMDGPU 显示驱动程序。由 AMD 提供的步骤是进行此操作最权威的说明。可以使用 rocm-smi 命令验证驱动程序设置。

rocm-smi

======================= ROCm System Management Interface =======================
================================= Concise Info =================================
GPU  Temp (DieEdge)  AvgPwr  SCLK    MCLK     Fan  Perf  PwrCap  VRAM%  GPU%
0    33.0c           37.0W   300Mhz  1200Mhz  0%   auto  290.0W    0%   0%
1    32.0c           39.0W   300Mhz  1200Mhz  0%   auto  290.0W    0%   0%
2    33.0c           37.0W   300Mhz  1200Mhz  0%   auto  290.0W    0%   0%
================================================================================
============================= End of ROCm SMI Log ==============================

设置 ROCm Docker 容器和 Conda 环境¶

建议（但不强制）通过 Docker 设置安装和运行 FBGEMM_GPU，以实现 ROCm 环境的隔离和可重复性，而 ROCm 环境可能难以设置。

按照 ROCm Docker 镜像中的说明拉取完整的 ROCm Docker 镜像并启动容器。

从那里，可以使用 Conda 构建其余的运行时环境。按照安装 ROCm 中的说明设置 Conda 环境，然后按照安装 Python 库进行操作。

安装 Python 库¶

安装与 FBGEMM_GPU 配合使用的相关 Python 库。

conda install -n ${env_name} -c conda-forge --override-channels -y \
    hypothesis \
    numpy \
    scikit-build

安装 PyTorch¶

按照安装 PyTorch 中的说明在 Conda 环境中安装 PyTorch。

安装 Triton¶

Triton 通常已随 PyTorch 安装包一起提供。可以通过以下方式验证：

conda run -n ${env_name} python -c "import triton"

如果 Triton 不可用，则可以通过 PyTorch PIP 进行安装：

# Most recent version used can be found in the build scripts
TRITON_VERSION=3.0.0+45fff310c8

conda run -n ${env_name} pip install \
  --pre pytorch-triton==${TRITON_VERSION} \
  --index-url https://download.pytorch.org/whl/nightly/

有关 PyTorch-Triton 发行版的信息可以在此处找到。

安装 FBGEMM_GPU 包¶

通过 PyTorch PIP 安装¶

PyTorch PIP 是安装 FBGEMM_GPU 的首选渠道。

# !! Run inside the Conda environment !!

# CPU-only Nightly
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu/
pip install --pre fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cpu/

# CPU-only Release
pip install torch --index-url https://download.pytorch.org/whl/cpu/
pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/cpu/

# CUDA Nightly
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126/
pip install --pre fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu126/

# CUDA Release
pip install torch --index-url https://download.pytorch.org/whl/cu126/
pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/cu126/

# ROCm Nightly
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/rocm6.3/
pip install --pre fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/rocm6.3/

# Test the installation
python -c "import torch; import fbgemm_gpu"

通过公共 PyPI 安装¶

# !! Run inside the Conda environment !!

# CPU-Only Nightly
pip install fbgemm-gpu-nightly-cpu

# CPU-Only Release
pip install fbgemm-gpu-cpu

# CUDA Nightly
pip install fbgemm-gpu-nightly

# CUDA Release
pip install fbgemm-gpu

在撰写本文时，FBGEMM_GPU 的 ROCm 版本的软件包尚未发布到公共 PyPI。

安装后检查¶

安装后，运行导入测试以确保库已正确链接和设置。

# !! Run inside the Conda environment !!

python -c "import torch; import fbgemm_gpu; print(torch.ops.fbgemm.merge_pooled_embeddings)"

未定义符号¶

遇到的常见错误是 Python 中导入 FBGEMM_GPU 失败，其错误签名如下：

Traceback (most recent call last):
  File "/root/miniconda/envs/mycondaenv/lib/python3.10/site-packages/torch/_ops.py", line 565, in __getattr__
    op, overload_names = torch._C._jit_get_operation(qualified_op_name)
RuntimeError: No such operator fbgemm::jagged_2d_to_dense
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/root/miniconda/envs/mycondaenv/lib/python3.10/site-packages/fbgemm_gpu-0.4.1.post47-py3.10-linux-aarch64.egg/fbgemm_gpu/__init__.py", line 21, in <module>
    from . import _fbgemm_gpu_docs  # noqa: F401, E402
  File "/root/miniconda/envs/mycondaenv/lib/python3.10/site-packages/fbgemm_gpu-0.4.1.post47-py3.10-linux-aarch64.egg/fbgemm_gpu/_fbgemm_gpu_docs.py", line 18, in <module>
    torch.ops.fbgemm.jagged_2d_to_dense,
  File "/root/miniconda/envs/mycondaenv/lib/python3.10/site-packages/torch/_ops.py", line 569, in __getattr__
    raise AttributeError(
AttributeError: '_OpNamespace' 'fbgemm' object has no attribute 'jagged_2d_to_dense'
ERROR conda.cli.main_run:execute(47): `conda run python -c import fbgemm_gpu` failed. (See above for error)
/root/miniconda/envs/mycondaenv/lib/python3.10/site-packages/fbgemm_gpu-0.4.1.post47-py3.10-linux-aarch64.egg/fbgemm_gpu/fbgemm_gpu_py.so: undefined symbol: _ZN6fbgemm48FloatOrHalfToFusedNBitRowwiseQuantizedSBHalfAvx2ItLi2EEEvPKT_miPh

一般来说，FBGEMM_GPU 安装中出现未定义符号的原因可能包括：

FBGEMM_GPU 所依赖的运行时库，例如 libnvidia-ml.so 或 libtorch.so，要么未正确安装，要么在 LD_LIBRARY_PATH 中不可见。
FBGEMM_GPU 包构建不正确，包含未链接的声明（例如，请参阅 PR 1618）。

在前一种情况下，可以通过重新安装相关软件包和/或手动更新 LD_LIBRARY_PATH 来解决。

在后一种情况下，这是一个严重的构建和/或打包问题，应报告给 FBGEMM 开发人员。