快捷方式

安裝說明

注意: 最新的安裝說明嵌入在 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 顯示驅動程式。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

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 的全面開發者文件

檢視文件

教程

獲取面向初學者和高階開發者的深入教程

檢視教程

資源

查詢開發資源並獲得問題解答

檢視資源