https://github.com/NVIDIA/apex

为了可以离线安装，现编译出whl文件：

bash
展开代码
#!/bin/bash
# 构建 apex whl 文件的脚本

# 创建输出目录
mkdir -p /workspace/wheels

# 配置 pip 源（避免代理问题，使用国内镜像）
export PIP_INDEX=${PIP_INDEX:-https://mirrors.aliyun.com/pypi/simple/}
export PIP_TRUSTED_HOST=${PIP_TRUSTED_HOST:-mirrors.aliyun.com}

# 先安装构建依赖（避免构建时下载失败）
pip install --no-cache-dir -i ${PIP_INDEX} --trusted-host ${PIP_TRUSTED_HOST} \
    packaging wheel setuptools pyproject-metadata

# 克隆 apex 仓库
git clone https://github.com/NVIDIA/apex.git /workspace/apex

# 进入 apex 目录
cd /workspace/apex

# 设置环境变量（与 Dockerfile 中保持一致）
export TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6;8.9;9.0"
export CUDA_HOME=/usr/local/cuda
export MAX_JOBS=${MAX_JOBS:-16}

# 使用 pip wheel 构建 whl 文件（不安装）
# -w 指定输出目录
# --no-build-isolation: 使用已安装的构建依赖，不创建隔离环境
# --no-deps: 不安装运行时依赖（只构建 wheel）
NVCC_APPEND_FLAGS="--threads 4" APEX_PARALLEL_BUILD=8 APEX_CPP_EXT=1 APEX_CUDA_EXT=1 \
pip wheel -v --no-build-isolation --no-deps -i ${PIP_INDEX} --trusted-host ${PIP_TRUSTED_HOST} \
    -w /workspace/wheels .

echo "apex whl 文件已构建完成，保存在 /workspace/wheels 目录"
ls -lh /workspace/wheels/*.whl

重新安装：

bash
展开代码
pip uninstall -y apex || true && APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation /workspace/apex-0.1-cp312-cp312-linux_x86_64.whl