構建說明

注意： 最新的構建說明已嵌入 FBGEMM 程式碼倉庫中位於 setup_env.bash 下的一組指令碼中。

當前可用的 FBGEMM_GPU 構建變體有

• 僅限 CPU

• CUDA

• 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 環境執行。

設定 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 構建的設定

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 配置中

set(THIRDPARTY ${FBGEMM}/external)

${THIRDPARTY}/cutlass/include
${THIRDPARTY}/cutlass/tools/util/include

ROCm 構建的設定

FBGEMM_GPU 支援在 AMD (ROCm) 裝置上執行。為 FBGEMM_GPU 的 ROCm 構建設定機器可以透過預構建的 Docker 映象或在裸機上完成。

ROCm Docker 映象

對於透過 Docker 進行的設定，只需拉取適用於所需 ROCm 版本的預安裝 Minimal Docker image for ROCm 即可

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

儘管 full ROCm Docker image 附帶所有預安裝的 ROCm 包，但這會導致 Docker 容器非常大，因此，出於這個原因，建議使用 minimal image 來構建和執行 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 時不指定版本（如夜間構建）可能並非總是可靠。例如，已知 PyTorch 夜間構建的 GPU 版本比僅限 CPU 的版本晚 2 小時到達 Conda。因此，在該時間視窗內安裝 pytorch-nightly 的 Conda 安裝將默默地回退到安裝僅限 CPU 的變體。

另請注意，由於 PyTorch 的 GPU 和僅限 CPU 版本都放置在同一個 artifact bucket 中，安裝過程中選擇的 PyTorch 變體將取決於系統是否安裝了 CUDA。因此，對於 GPU 構建，在安裝 PyTorch 之前先安裝 CUDA / ROCm 非常重要。

透過 PyTorch PIP 安裝

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

# 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.2.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

# NOTE: To build the package as part of generating the documentation, use
# `--package_variant=docs` flag instead!

要使用 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}"

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("