目录
快捷方式

构建说明

注意：最新的构建说明包含在 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 环境运行。

设置 Conda 环境

使用指定的 Python 版本创建一个 Conda 环境

env_name=<ENV NAME>
python_version=3.12

# 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 及更高版本）**。可以通过预构建的 Docker 镜像或在裸机上通过 Conda 安装来设置 CUDA 构建的 FBGEMM_GPU 的机器。请注意，构建不需要 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.1.0

# 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) 设备上运行。可以通过预构建的 Docker 镜像或在裸机上设置 CUDA 构建的 FBGEMM_GPU 的机器。

ROCm Docker 镜像

对于通过 Docker 进行的设置，只需提取所需的 ROCm 版本的预安装 ROCm 最小 Docker 镜像。

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

虽然 完整的 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/5.6.1/ubuntu/focal/amdgpu-install_5.6.50601-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 -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 -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++ 标准库实现）。

这意味着无论是否使用 Clang 构建，GCC 都是 FBGEMM_GPU 的 CUDA 变体的必需依赖项。在这种情况下，建议在此方案中先安装 GCC 工具链，然后再安装 Clang 工具链；有关说明，请参阅 C/C++ 编译器 (GCC)。

其他构建工具

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

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

安装 PyTorch

PyTorch 官方 主页 提供了安装 PyTorch 的最权威指南，可以通过 Conda 或 PIP 进行安装。

通过 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 版本会在 Conda 中比仅 CPU 版本晚 2 个小时出现。因此，在此时间窗口内，Conda 安装 pytorch-nightly 将静默回退到安装仅 CPU 版本。

另请注意，由于 PyTorch 的 GPU 和仅 CPU 版本都放置在同一个工件存储桶中，因此在安装过程中选择的 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/cu121/

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

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

截至撰写本文时，对于安装 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=v0.8.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 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.12 --> py312
export python_tag=py312

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

# Set the Python platform name for the Linux case
export python_plat_name="manylinux2014_${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[@]}"

# 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

# Build for the target architecture of the ROCm device installed on the machine (e.g. 'gfx906;gfx908;gfx90a')
# See https://wiki.gentoo.org/wiki/ROCm 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}" \
    -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 \
    -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("

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取适合初学者和高级开发人员的深入教程

查看教程

资源

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

查看资源

为了分析流量并优化您的体验，我们在此网站上提供 Cookie。点击或导航即表示您同意我们使用 Cookie。作为本网站的当前维护方，Facebook 的 Cookie 政策适用。了解更多信息，包括有关可用控件的信息：Cookie 政策。