目录
快捷方式

安装说明

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

安装 FBGEMM_GPU 的一般步骤如下

  1. 建立一个隔离的构建环境。

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

  3. 安装 PyTorch。

  4. 安装 FBGEMM_GPU 软件包。

  5. 运行安装后检查。

在安装 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 安装中出现未定义符号的原因如下

  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 政策。