建置說明

注意： 最新的建置說明已嵌入 FBGEMM 儲存庫中 setup_env.bash 內的一組腳本中。

目前可用的 FBGEMM_GPU 建置變體為

• 僅限 CPU

• CUDA

• GenAI (實驗性)

• ROCm

建置 FBGEMM_GPU 的一般步驟如下

1. 設定隔離的建置環境。

2. 為僅限 CPU、CUDA 或 ROCm 建置設定工具鏈。

3. 安裝 PyTorch。

4. 執行建置腳本。

設定隔離的建置環境

安裝 Miniconda

建議設定 Miniconda 環境，以實現可重現的建置

export PLATFORM_NAME="$(uname -s)-$(uname -m)"

# Set the Miniconda prefix directory
miniconda_prefix=$HOME/miniconda

# Download the Miniconda installer
wget -q "https://repo.anaconda.com/miniconda/Miniconda3-latest-${PLATFORM_NAME}.sh" -O miniconda.sh

# Run the installer
bash miniconda.sh -b -p "$miniconda_prefix" -u

# Load the shortcuts
. ~/.bashrc

# Run updates
conda update -n base -c defaults -y conda

從此處開始，所有安裝命令都將針對或在 Conda 環境內執行。

設定 Conda 環境

使用指定的 Python 版本建立 Conda 環境

env_name=<ENV NAME>
python_version=3.13

# Create the environment
conda create -y --name ${env_name} python="${python_version}"

# Upgrade PIP and pyOpenSSL package
conda run -n ${env_name} pip install --upgrade pip
conda run -n ${env_name} python -m pip install pyOpenSSL>22.1.0

設定僅限 CPU 建置

請依照 設定隔離的建置環境 中的 Conda 環境設定說明，然後依照 安裝建置工具 中的說明。

設定 CUDA / 僅限 GenAI 建置

FBGEMM_GPU 的 CUDA 建置需要最新版本的 nvcc，其支援運算能力 3.5+。設定機器以進行 FBGEMM_GPU 的 CUDA 建置可以透過預先建置的 Docker 映像檔或透過裸機上的 Conda 安裝來完成。請注意，GPU 或 NVIDIA 驅動程式都不需要存在即可進行建置，因為它們僅在執行階段使用。

CUDA Docker 映像檔

對於透過 Docker 設定，只需提取適用於所需 Linux 發行版本和 CUDA 版本的預先安裝 CUDA Docker 映像檔。

# Run for Ubuntu 22.04, CUDA 11.8
docker run -it --entrypoint "/bin/bash" nvidia/cuda:11.8.0-devel-ubuntu22.04

從這裡開始，其餘的建置環境可以透過 Conda 建構，因為它仍然是建立隔離且可重現的建置環境的建議機制。

安裝 CUDA

透過 Conda 安裝完整的 CUDA 套件，其中包括 NVML

# See https://anaconda.org/nvidia/cuda for all available versions of CUDA
cuda_version=12.4.1

# Install the full CUDA package
conda install -n ${env_name} -y cuda -c "nvidia/label/cuda-${cuda_version}"

驗證是否找到 cuda_runtime.h、libnvidia-ml.so 和 libnccl.so*

conda_prefix=$(conda run -n ${env_name} printenv CONDA_PREFIX)

find "${conda_prefix}" -name cuda_runtime.h
find "${conda_prefix}" -name libnvidia-ml.so
find "${conda_prefix}" -name libnccl.so*

安裝 cuDNN

cuDNN 是 FBGEMM_GPU 的 CUDA 變體的建置時期相依性。下載並解壓縮給定 CUDA 版本的 cuDNN 套件

# cuDNN package URLs for each platform and CUDA version can be found in:
# https://github.com/pytorch/builder/blob/main/common/install_cuda.sh
cudnn_url=https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.9.2.26_cuda12-archive.tar.xz

# Download and unpack cuDNN
wget -q "${cudnn_url}" -O cudnn.tar.xz
tar -xvf cudnn.tar.xz

安裝 CUTLASS

本節僅適用於建置實驗性的 FBGEMM_GPU GenAI 模組。CUTLASS 應已在儲存庫中以 git 子模組形式提供 (請參閱 準備建置)。以下包含路徑已新增至 CMake 設定

設定 ROCm 建置

FBGEMM_GPU 支援在 AMD (ROCm) 裝置上執行。設定機器以進行 FBGEMM_GPU 的 ROCm 建置可以透過預先建置的 Docker 映像檔或透過裸機來完成。

ROCm Docker 映像檔

對於透過 Docker 設定，只需提取適用於所需 ROCm 版本的預先安裝 ROCm 最小 Docker 映像檔

# Run for ROCm 6.2.0
docker run -it --entrypoint "/bin/bash" rocm/rocm-terminal:6.2.0

雖然 完整 ROCm Docker 映像檔 隨附所有預先安裝的 ROCm 套件，但它會產生非常大的 Docker 容器，因此基於此原因，建議使用最小映像檔來建置和執行 FBGEMM_GPU。

從這裡開始，其餘的建置環境可以透過 Conda 建構，因為它仍然是建立隔離且可重現的建置環境的建議機制。

安裝 ROCm

透過作業系統套件管理器安裝完整的 ROCm 套件。完整說明可在 ROCm 安裝指南 中找到

# [OPTIONAL] Disable apt installation prompts
export DEBIAN_FRONTEND=noninteractive

# Update the repo DB
apt update

# Download the installer
wget -q https://repo.radeon.com/amdgpu-install/6.3.1/ubuntu/focal/amdgpu-install_6.3.60301-1_all.deb -O amdgpu-install.deb

# Run the installer
apt install ./amdgpu-install.deb

# Install ROCm
amdgpu-install -y --usecase=hiplibsdk,rocm --no-dkms

安裝 MIOpen

MIOpen 是 FBGEMM_GPU 的 ROCm 變體的相依性，需要安裝

apt install hipify-clang miopen-hip miopen-hip-dev

安裝建置工具

本節中的說明適用於所有 FBGEMM_GPU 變體的建置。

C/C++ 編譯器 (GCC)

安裝支援 C++20 的 GCC 工具鏈版本。sysroot 套件也需要安裝，以避免在編譯 FBGEMM_CPU 時，GLIBCXX 缺少版本化符號的問題

# Set GCC to 10.4.0 to keep compatibility with older versions of GLIBCXX
#
# A newer versions of GCC also works, but will need to be accompanied by an
# appropriate updated version of the sysroot_linux package.
gcc_version=10.4.0

conda install -n ${env_name} -c conda-forge --override-channels -y \
  gxx_linux-64=${gcc_version} \
  sysroot_linux-64=2.17

雖然可以使用較新版本的 GCC，但在較新版本的 GCC 下編譯的二進位檔將與較舊的系統 (例如 Ubuntu 20.04 或 CentOS Stream 8) 不相容，因為編譯後的函式庫將參考系統 libstdc++.so.6 不支援的 GLIBCXX 版本符號。若要查看可用的 libstdc++.so.6 支援哪些版本的 GLIBC 和 GLIBCXX

libcxx_path=/path/to/libstdc++.so.6

# Print supported for GLIBC versions
objdump -TC "${libcxx_path}" | grep GLIBC_ | sed 's/.*GLIBC_\([.0-9]*\).*/GLIBC_\1/g' | sort -Vu | cat

# Print supported for GLIBCXX versions
objdump -TC "${libcxx_path}" | grep GLIBCXX_ | sed 's/.*GLIBCXX_\([.0-9]*\).*/GLIBCXX_\1/g' | sort -Vu | cat

C/C++ 編譯器 (Clang)

可以使用 Clang 作為主機編譯器來建置 FBGEMM 和 FBGEMM_GPU (僅限 CPU 和 CUDA 變體)。若要執行此操作，請安裝支援 C++20 的 Clang 工具鏈版本

# Minimum LLVM+Clang version required for FBGEMM_GPU
llvm_version=16.0.6

# NOTE: libcxx from conda-forge is outdated for linux-aarch64, so we cannot
# explicitly specify the version number
conda install -n ${env_name} -c conda-forge --override-channels -y \
    clangxx=${llvm_version} \
    libcxx \
    llvm-openmp=${llvm_version} \
    compiler-rt=${llvm_version}

# Append $CONDA_PREFIX/lib to $LD_LIBRARY_PATH in the Conda environment
ld_library_path=$(conda run -n ${env_name} printenv LD_LIBRARY_PATH)
conda_prefix=$(conda run -n ${env_name} printenv CONDA_PREFIX)
conda env config vars set -n ${env_name} LD_LIBRARY_PATH="${ld_library_path}:${conda_prefix}/lib"

# Set NVCC_PREPEND_FLAGS in the Conda environment for Clang to work correctly as the host compiler
conda env config vars set -n ${env_name} NVCC_PREPEND_FLAGS=\"-std=c++20 -Xcompiler -std=c++20 -Xcompiler -stdlib=libstdc++ -ccbin ${clangxx_path} -allow-unsupported-compiler\"

注意，對於 CUDA 程式碼編譯，即使 nvcc 支援 Clang 作為主機編譯器，但對於 nvcc 使用的任何主機編譯器，僅支援 libstd++ (GCC 的 C++ 標準函式庫實作)。

這表示 GCC 是 FBGEMM_GPU 的 CUDA 變體的必要相依性，無論它是否使用 Clang 建置。在這種情況下，建議先安裝 GCC 工具鏈，然後再安裝 Clang 工具鏈；請參閱 C/C++ 編譯器 (GCC) 中的說明。

編譯器符號連結

安裝編譯器工具鏈後，將 C 和 C++ 編譯器符號連結到 binpath (根據需要覆寫現有的符號連結)。在 Conda 環境中，binpath 位於 $CONDA_PREFIX/bin

conda_prefix=$(conda run -n ${env_name} printenv CONDA_PREFIX)

ln -sf "${path_to_either_gcc_or_clang}" "$(conda_prefix)/bin/cc"
ln -sf "${path_to_either_gcc_or_clang}" "$(conda_prefix)/bin/c++"

這些符號連結稍後將在 FBGEMM_GPU 建置設定階段中使用。

其他建置工具

安裝其他必要的建置工具，例如 ninja、cmake 等

conda install -n ${env_name} -c conda-forge --override-channels -y \
    click \
    cmake \
    hypothesis \
    jinja2 \
    make \
    ncurses \
    ninja \
    numpy \
    scikit-build \
    wheel

安裝 PyTorch

官方 PyTorch 首頁 包含關於如何透過 Conda 或透過 PIP 安裝 PyTorch 的最權威說明。

透過 Conda 安裝

# Install the latest nightly
conda install -n ${env_name} -y pytorch -c pytorch-nightly

# Install the latest test (RC)
conda install -n ${env_name} -y pytorch -c pytorch-test

# Install a specific version
conda install -n ${env_name} -y pytorch==2.0.0 -c pytorch

請注意，在未指定版本的情況下透過 Conda 安裝 PyTorch (如 nightly builds 的情況) 可能並非總是可靠。例如，已知 PyTorch nightly builds 的 GPU 建置比僅限 CPU 的建置晚 2 小時到達 Conda。因此，在該時間範圍內安裝 pytorch-nightly 的 Conda 將會靜默地回退到安裝僅限 CPU 的變體。

另請注意，由於 GPU 和僅限 CPU 版本的 PyTorch 都放置在相同的成品儲存區中，因此安裝期間選取的 PyTorch 變體將取決於系統上是否安裝了 CUDA。因此，對於 GPU 建置，在 PyTorch 之前先安裝 CUDA / ROCm 非常重要。

透過 PyTorch PIP 安裝

建議透過 PyTorch PIP 安裝 PyTorch，而不是 Conda，因為它更具確定性，因此更可靠

# Install the latest nightly, CPU variant
conda run -n ${env_name} pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu/

# Install the latest test (RC), CUDA variant
conda run -n ${env_name} pip install --pre torch --index-url https://download.pytorch.org/whl/test/cu126/

# Install a specific version, CUDA variant
conda run -n ${env_name} pip install torch==2.6.0+cu126 --index-url https://download.pytorch.org/whl/cu126/

# Install the latest nightly, ROCm variant
conda run -n ${env_name} pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/rocm6.3/

對於安裝 PyTorch 的 ROCm 變體，PyTorch PIP 是撰寫時唯一可用的管道。

安裝後檢查

使用 import 測試驗證 PyTorch 安裝 (版本和變體)

# Ensure that the package loads properly
conda run -n ${env_name} python -c "import torch.distributed"

# Verify the version and variant of the installation
conda run -n ${env_name} python -c "import torch; print(torch.__version__)"

對於 PyTorch 的 CUDA 變體，驗證是否至少找到 cuda_cmake_macros.h

conda_prefix=$(conda run -n ${env_name} printenv CONDA_PREFIX)
find "${conda_prefix}" -name cuda_cmake_macros.h

安裝 PyTorch-Triton

本節僅適用於建置實驗性的 FBGEMM_GPU Triton-GEMM 模組。Triton 應透過 pytorch-triton 安裝，這通常會安裝 torch，但也可以手動安裝

# pytorch-triton repos:
# https://download.pytorch.org/whl/nightly/pytorch-triton/
# https://download.pytorch.org/whl/nightly/pytorch-triton-rocm/

# The version SHA should follow the one pinned in PyTorch
# https://github.com/pytorch/pytorch/blob/main/.ci/docker/ci_commit_pins/triton.txt
conda run -n ${env_name} pip install --pre pytorch-triton==3.0.0+dedb7bdf33 --index-url https://download.pytorch.org/whl/nightly/

使用 import 測試驗證 PyTorch-Triton 安裝

# Ensure that the package loads properly
conda run -n ${env_name} python -c "import triton"

其他建置前設定

準備建置

複製儲存庫及其子模組，並安裝 requirements.txt

# !! Run inside the Conda environment !!

# Select a version tag
FBGEMM_VERSION=v1.0.0

# Clone the repo along with its submodules
git clone --recursive -b ${FBGEMM_VERSION} https://github.com/pytorch/FBGEMM.git fbgemm_${FBGEMM_VERSION}

# Install additional required packages for building and testing
cd fbgemm_${FBGEMM_VERSION}/fbgemm_gpu
pip install -r requirements.txt

建置流程

FBGEMM_GPU 建置流程使用基於 scikit-build CMake 的建置流程，並在安裝執行之間保持狀態。因此，建置可能會過時，並且在由於缺少相依性等原因導致建置失敗後嘗試重新執行時可能會導致問題。為了解決此問題，只需清除建置快取

# !! Run in fbgemm_gpu/ directory inside the Conda environment !!

python setup.py clean

設定 Wheel 建置變數

在建置 Python wheel 時，必須先正確設定套件名稱、Python 版本標籤和 Python 平台名稱

# Set the package name depending on the build variant
export package_name=fbgemm_gpu_{cpu, cuda, rocm}

# Set the Python version tag.  It should follow the convention `py<major><minor>`,
# e.g. Python 3.13 --> py313
export python_tag=py313

# Determine the processor architecture
export ARCH=$(uname -m)

# Set the Python platform name for the Linux case
export python_plat_name="manylinux_2_28_${ARCH}"
# For the macOS (x86_64) case
export python_plat_name="macosx_10_9_${ARCH}"
# For the macOS (arm64) case
export python_plat_name="macosx_11_0_${ARCH}"
# For the Windows case
export python_plat_name="win_${ARCH}"

僅限 CPU 建置

對於僅限 CPU 的建置，需要指定 --cpu_only 旗標。

# !! Run in fbgemm_gpu/ directory inside the Conda environment !!

# Build the wheel artifact only
python setup.py bdist_wheel \
    --package_variant=cpu \
    --python-tag="${python_tag}" \
    --plat-name="${python_plat_name}"

# Build and install the library into the Conda environment (GCC)
python setup.py install \
    --package_variant=cpu

若要使用 Clang + libstdc++ 而非 GCC 進行建置，只需附加 --cxxprefix 旗標

# !! Run in fbgemm_gpu/ directory inside the Conda environment !!

# Build the wheel artifact only
python setup.py bdist_wheel \
    --package_variant=cpu \
    --python-tag="${python_tag}" \
    --plat-name="${python_plat_name}" \
    --cxxprefix=$CONDA_PREFIX

# Build and install the library into the Conda environment (Clang)
python setup.py install \
    --package_variant=cpu
    --cxxprefix=$CONDA_PREFIX

請注意，這假定 Clang 工具鏈已與 GCC 工具鏈一起正確安裝，並且可作為 ${cxxprefix}/bin/cc 和 ${cxxprefix}/bin/c++ 使用。

若要啟用執行階段偵錯功能，例如 CUDA 和 HIP 中的裝置端判言，只需在叫用 setup.py 時附加 --debug 旗標。

CUDA 建置

為 CUDA 建置 FBGEMM_GPU 需要安裝 NVML 和 cuDNN，並透過環境變數提供給建置。但是，建置套件不需要 CUDA 裝置。

與僅限 CPU 的建置類似，可以透過將 --cxxprefix=$CONDA_PREFIX 附加至建置命令來啟用使用 Clang + libstdc++ 進行建置，假設工具鏈已正確安裝。

# !! Run in fbgemm_gpu/ directory inside the Conda environment !!

# [OPTIONAL] Specify the CUDA installation paths
# This may be required if CMake is unable to find nvcc
export CUDACXX=/path/to/nvcc
export CUDA_BIN_PATH=/path/to/cuda/installation

# [OPTIONAL] Provide the CUB installation directory (applicable only to CUDA versions prior to 11.1)
export CUB_DIR=/path/to/cub

# [OPTIONAL] Allow NVCC to use host compilers that are newer than what NVCC officially supports
nvcc_prepend_flags=(
  -allow-unsupported-compiler
)

# [OPTIONAL] If clang is the host compiler, set NVCC to use libstdc++ since libc++ is not supported
nvcc_prepend_flags+=(
  -Xcompiler -stdlib=libstdc++
  -ccbin "/path/to/clang++"
)

# [OPTIONAL] Set NVCC_PREPEND_FLAGS as needed
export NVCC_PREPEND_FLAGS="${nvcc_prepend_flags[@]}"

# [OPTIONAL] Enable verbose NVCC logs
export NVCC_VERBOSE=1

# Specify cuDNN header and library paths
export CUDNN_INCLUDE_DIR=/path/to/cudnn/include
export CUDNN_LIBRARY=/path/to/cudnn/lib

# Specify NVML filepath
export NVML_LIB_PATH=/path/to/libnvidia-ml.so

# Specify NCCL filepath
export NCCL_LIB_PATH=/path/to/libnccl.so.2

# Build for SM70/80 (V100/A100 GPU); update as needed
# If not specified, only the CUDA architecture supported by current system will be targeted
# If not specified and no CUDA device is present either, all CUDA architectures will be targeted
cuda_arch_list=7.0;8.0

# Unset TORCH_CUDA_ARCH_LIST if it exists, bc it takes precedence over
# -DTORCH_CUDA_ARCH_LIST during the invocation of setup.py
unset TORCH_CUDA_ARCH_LIST

# Build the wheel artifact only
python setup.py bdist_wheel \
    --package_variant=cuda \
    --python-tag="${python_tag}" \
    --plat-name="${python_plat_name}" \
    --nvml_lib_path=${NVML_LIB_PATH} \
    --nccl_lib_path=${NCCL_LIB_PATH} \
    -DTORCH_CUDA_ARCH_LIST="${cuda_arch_list}"

# Build and install the library into the Conda environment
python setup.py install \
    --package_variant=cuda \
    --nvml_lib_path=${NVML_LIB_PATH} \
    --nccl_lib_path=${NCCL_LIB_PATH} \
    -DTORCH_CUDA_ARCH_LIST="${cuda_arch_list}"

僅限 GenAI 建置

預設情況下，FBGEMM_GPU 的 CUDA 建置包含所有用於 GenAI 應用程式的實驗性模組。僅建置實驗性模組的說明與 CUDA 建置的說明相同，但需在建置叫用中指定 --package_variant=genai

# Build the wheel artifact only
python setup.py bdist_wheel \
    --package_variant=genai \
    --python-tag="${python_tag}" \
    --plat-name="${python_plat_name}" \
    --nvml_lib_path=${NVML_LIB_PATH} \
    --nccl_lib_path=${NCCL_LIB_PATH} \
    -DTORCH_CUDA_ARCH_LIST="${cuda_arch_list}"

# Build and install the library into the Conda environment
python setup.py install \
    --package_variant=genai \
    --nvml_lib_path=${NVML_LIB_PATH} \
    --nccl_lib_path=${NCCL_LIB_PATH} \
    -DTORCH_CUDA_ARCH_LIST="${cuda_arch_list}"

請注意，目前實驗性模組僅支援 CUDA。

ROCm 建置

對於 ROCm 建置，需要指定 ROCM_PATH 和 PYTORCH_ROCM_ARCH。但是，建置套件不需要 ROCm 裝置。

與僅限 CPU 和 CUDA 建置類似，可以透過將 --cxxprefix=$CONDA_PREFIX 附加至建置命令來啟用使用 Clang + libstdc++ 進行建置，假設工具鏈已正確安裝。

# !! Run in fbgemm_gpu/ directory inside the Conda environment !!

export ROCM_PATH=/path/to/rocm

# [OPTIONAL] Enable verbose HIPCC logs
export HIPCC_VERBOSE=1

# Build for the target architecture of the ROCm device installed on the machine (e.g. 'gfx908,gfx90a,gfx942')
# See https://rocm.docs.amd.com/en/latest/reference/gpu-arch-specs.html for list
export PYTORCH_ROCM_ARCH=$(${ROCM_PATH}/bin/rocminfo | grep -o -m 1 'gfx.*')

# Build the wheel artifact only
python setup.py bdist_wheel \
    --package_variant=rocm \
    --python-tag="${python_tag}" \
    --plat-name="${python_plat_name}" \
    -DAMDGPU_TARGETS="${PYTORCH_ROCM_ARCH}" \
    -DHIP_ROOT_DIR="${ROCM_PATH}" \
    -DCMAKE_C_FLAGS="-DTORCH_USE_HIP_DSA" \
    -DCMAKE_CXX_FLAGS="-DTORCH_USE_HIP_DSA"

# Build and install the library into the Conda environment
python setup.py install \
    --package_variant=rocm \
    -DAMDGPU_TARGETS="${PYTORCH_ROCM_ARCH}" \
    -DHIP_ROOT_DIR="${ROCM_PATH}" \
    -DCMAKE_C_FLAGS="-DTORCH_USE_HIP_DSA" \
    -DCMAKE_CXX_FLAGS="-DTORCH_USE_HIP_DSA"

建置後檢查 (適用於開發人員)

建置完成後，執行一些檢查以驗證建置是否確實正確很有用。

未定義符號檢查

由於 FBGEMM_GPU 包含許多 Jinja 和 C++ 範本實例化，因此務必確保在開發過程中沒有意外產生的未定義符號

# !! Run in fbgemm_gpu/ directory inside the Conda environment !!

# Locate the built .SO file
fbgemm_gpu_lib_path=$(find . -name fbgemm_gpu_py.so)

# Check that the undefined symbols don't include fbgemm_gpu-defined functions
nm -gDCu "${fbgemm_gpu_lib_path}" | sort

GLIBC 版本相容性檢查

驗證參考的 GLIBCXX 版本號以及某些函式符號的可用性也很有用

# !! Run in fbgemm_gpu/ directory inside the Conda environment !!

# Locate the built .SO file
fbgemm_gpu_lib_path=$(find . -name fbgemm_gpu_py.so)

# Note the versions of GLIBCXX referenced by the .SO
# The libstdc++.so.6 available on the install target must support these versions
objdump -TC "${fbgemm_gpu_lib_path}" | grep GLIBCXX | sed 's/.*GLIBCXX_\([.0-9]*\).*/GLIBCXX_\1/g' | sort -Vu | cat

# Test for the existence of a given function symbol in the .SO
nm -gDC "${fbgemm_gpu_lib_path}" | grep " fbgemm_gpu::merge_pooled_embeddings("
nm -gDC "${fbgemm_gpu_lib_path}" | grep " fbgemm_gpu::jagged_2d_to_dense("