快捷方式

安装说明

注意: 最新的安装说明嵌入在 FBGEMM 代码库的 setup_env.bash 脚本集中。

安装 FBGEMM_GPU 的一般步骤如下

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

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

  3. 安装 PyTorch。

  4. 安装 FBGEMM_GPU 包。

  5. 运行安装后检查。

FBGEMM 版本兼容性表

FBGEMM 根据 PyTorch 发布计划发布,每个版本均不保证与早于 FBGEMM 版本对应的 PyTorch 版本协同工作。

FBGEMM 版本

对应的 PyTorch 版本

支持的 Python 版本

支持的 CUDA 版本

支持的 CUDA 架构

(实验性)支持的 ROCm 版本

(实验性)支持的 ROCm 架构

1.1.0

2.6.x

3.9, 3.10, 3.11, 3.12, 3.13

11.8, 12.4, 12.6

7.0、8.0、9.0、9.0a

6.1, 6.2.4, 6.3

gfx908、gfx90a、gfx942

1.0.0

2.5.x

3.9, 3.10, 3.11, 3.12

11.8, 12.1, 12.4

7.0、8.0、9.0、9.0a

6.0, 6.1

gfx908、gfx90a

0.8.0

2.4.x

3.8, 3.9, 3.10, 3.11, 3.12

11.8, 12.1, 12.4

7.0、8.0、9.0、9.0a

6.0, 6.1

gfx908、gfx90a

0.7.0

2.3.x

3.8, 3.9, 3.10, 3.11, 3.12

11.8, 12.1

7.0, 8.0, 9.0

6.0

gfx908、gfx90a

0.6.0

2.2.x

3.8, 3.9, 3.10, 3.11, 3.12

11.8, 12.1

7.0, 8.0, 9.0

5.7

gfx90a

0.5.0

2.1.x

3.8, 3.9, 3.10, 3.11

11.8, 12.1

7.0, 8.0, 9.0

5.5, 5.6

gfx90a

0.4.0

2.0.x

3.8, 3.9, 3.10

11.7, 11.8

7.0, 8.0

5.3, 5.4

gfx90a

请注意,支持的 CUDA 和 ROCm 架构列表指的是默认安装包中可用的目标支持,并且为其他架构构建可能是可能的,但不保证。

有关更多信息,请访问

设置仅 CPU 环境

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

设置 CUDA 环境

FBGEMM_GPU 的 CUDA 变体需要机器上安装 NVIDIA GPU 以及可用的 NVIDIA 驱动程序;否则,库将回退到运行 CPU 版本的运算符。

FBGEMM_GPU CUDA 包目前仅针对 SM70 和 SM80 架构(分别为 V100 和 A100 GPU)构建。可以通过从头开始构建包来实现对其他架构的支持,但不保证能够正常工作(特别是对于较旧的架构)。

安装 NVIDIA 驱动程序

在所有其他环境设置之前,必须在系统上安装 NVIDIA 显示驱动程序。NVIDIAPyTorch 提供的步骤是执行此操作的最权威说明。可以使用 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

本节仅适用于使用实验性 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/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 安装中出现未定义符号的原因如下

  1. FBGEMM_GPU 依赖的运行时库(例如 libnvidia-ml.solibtorch.so)未正确安装或在 LD_LIBRARY_PATH 中不可见。

  2. FBGEMM_GPU 包构建不正确,并且包含未链接的声明(例如,请参阅 PR 1618)。

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

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

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

查看资源