这份 `README.md` 是专门为你提供的 **海光 DCU (DTK 25.04.2)** 环境定制的。它整合了你提供的 Docker 镜像信息、编译步骤,以及我们之前解决的 **OneTBB 路径修复**和 **HIP 类型适配**等关键补丁。 --- # FBGEMM_GPU 安装指南 (DTK 25.04.2) 本指南用于在海光加速卡环境下,基于 **DTK 25.04.2** 和 **PyTorch 2.5.1/2.7** 编译并安装 `fbgemm_gpu`。 ## 1. 环境启动 使用包含 DTK 25.04.2 的 PyTorch 镜像启动容器: ```bash # 拉取镜像 docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.5.1-ubuntu22.04-dtk25.04.2-py3.10 # 启动容器 docker run -it \ --network=host \ --ipc=host \ --shm-size=16G \ --device=/dev/kfd \ --device=/dev/mkfd \ --device=/dev/dri \ -v /opt/hyhal:/opt/hyhal \ --group-add video \ --cap-add=SYS_PTRACE \ --security-opt seccomp=unconfined \ image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.5.1-ubuntu22.04-dtk25.04.2-py3.10 ``` ## 2. 获取源码 下载并初始化 `FBGEMM` 仓库(推荐版本 v1.3.0): ```bash git clone https://github.com/pytorch/FBGEMM.git --branch v1.3.0 cd FBGEMM/fbgemm_gpu git submodule update --init --recursive ``` ## 3. 关键补丁与修复 (必须执行) 在编译前,需要针对 DTK 25.04.2 的类型定义和 TBB 库结构进行手动修复: ### a. 系统环境与代码清理 ```bash # 建立系统库链接 ln -s /lib/x86_64-linux-gnu/librt.so.1 /usr/lib/x86_64-linux-gnu/librt.so ``` ### b. 移除断言检查 CUDA_KERNEL_ASSERT2 ![cuda_kernel_assert2](./imgs/cuda_kernel_assert2.png) ### c. 移除 ./include/fbgemm_gpu/utils/kernel_launcher.cuh 中的 TORCH_CHECK ![torch_check](./imgs/torch_check.png) ## 4. 编译与安装 ```bash # 1. 清理缓存 rm -rf _skbuild build # 2. 设置环境变量 export PYTORCH_ROCM_ARCH="gfx936" export ROCM_PATH=/opt/dtk export BUILD_VERSION="1.3.0" export TORCH_PATH="/usr/local/lib/python3.10/dist-packages/torch/share/cmake/Torch" # 3. 执行安装命令 # 注意:通过 --no-as-needed 强制链接 TBB 并烧录 RPATH 以免运行时报错 python setup.py bdist_wheel \ --package_channel=release \ -DFBGEMM_BUILD_TARGET=default \ -DFBGEMM_BUILD_VARIANT=rocm \ -DCMAKE_PREFIX_PATH="${TORCH_PATH};/opt/dtk" \ -DAMDGPU_TARGETS="${PYTORCH_ROCM_ARCH}" \ -DTORCH_DIR="${TORCH_PATH}" \ -DCMAKE_C_FLAGS="-Wno-return-type -Wno-ignored-attributes -O1" \ -DCMAKE_CXX_FLAGS="-Wno-return-type -Wno-ignored-attributes -O1 -DhipDeviceProp_t=hipDeviceProp_t_v2" \ -DCMAKE_HIP_FLAGS="-Wno-return-type -O1 -DhipDeviceProp_t=hipDeviceProp_t_v2" \ -DUSE_AVX512=on \ -DCOPY_VISIBLE_LIBRARIES=ON \ -DCMAKE_SHARED_LINKER_FLAGS="-Wl,--no-as-needed -ltbb -Wl,-rpath,/usr/lib/x86_64-linux-gnu" \ -DTBB_INCLUDE_DIR=/usr/include ``` ## 5. 验证安装 安装完成后,在 Python 中验证: ```python import torch import fbgemm_gpu print("FBGEMM_GPU 导入成功!") ```