安装说明¶

注意： 最新的安装说明嵌入在 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 运行时¶

如果 OS / 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 开发者。