目录
快捷方式

安装说明

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

安装 FBGEMM_GPU 的一般步骤如下：

  1. 设置隔离的构建环境。

  2. 为 CPU 仅、CUDA 或 ROCm 运行时设置工具链。

  3. 安装 PyTorch。

  4. 安装 FBGEMM_GPU 包。

  5. 运行安装后检查。

设置 CPU 仅环境

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

设置 CUDA 环境

CUDA 版本的 FBGEMM_GPU 需要在机器上安装 NVIDIA GPU，并安装工作正常的 NVIDIA 驱动程序；否则库将退回到运行运算符的 CPU 版本。

FBGEMM_GPU CUDA 包目前仅针对 SM70 和 SM80 架构（分别为 V100 和 A100 GPU）构建。可以通过从头开始构建包来支持其他架构，但不能保证可以工作（特别是对于旧架构）。

安装 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 环境

ROCm 版本的 FBGEMM_GPU 需要在机器上安装 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} -y \
    hypothesis \
    numpy \
    scikit-build

安装 PyTorch

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

安装 Triton

本节仅适用于处理实验性的 FBGEMM_GPU GenAI 模块。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/cu121/
pip install --pre fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu121/

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

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

# 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

截至撰写本文时，ROCm 版本的 FBGEMM_GPU 包尚未发布到公共 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 安装中出现，原因如下：

  1. FBGEMM_GPU 依赖的运行时库（例如 libnvidia-ml.so 或 libtorch.so）要么未正确安装，要么在 LD_LIBRARY_PATH 中不可见。

  2. FBGEMM_GPU 包构建错误，包含未链接的声明（例如，请参见 PR 1618）。

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

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

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源

为了分析流量并优化您的体验，我们在本网站上提供 Cookie。点击或导航，即表示您同意我们使用 Cookie。作为本网站的当前维护者，Facebook 的 Cookie 政策适用。了解更多信息，包括有关可用控件的信息：Cookie 政策。