统一成8.0;8.6;8.9

This commit is contained in:
hailin 2025-09-20 16:52:21 +08:00
parent 2e8ba323e1
commit 2f5ee47d74
1 changed files with 11 additions and 3 deletions

View File

@ -15,7 +15,7 @@ ENV USE_CUDA=1 \
ARG MAX_JOBS=90 ARG MAX_JOBS=90
ENV DEBIAN_FRONTEND=noninteractive \ ENV DEBIAN_FRONTEND=noninteractive \
PYTHONUNBUFFERED=1 LANG=C.UTF-8 LC_ALL=C.UTF-8 \ PYTHONUNBUFFERED=1 LANG=C.UTF-8 LC_ALL=C.UTF-8 \
TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0" TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9"
RUN apt-get update && apt-get install -y --no-install-recommends \ RUN apt-get update && apt-get install -y --no-install-recommends \
python3 python3-dev python3-pip python3-distutils git cmake ninja-build \ python3 python3-dev python3-pip python3-distutils git cmake ninja-build \
@ -27,6 +27,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
RUN python3 -m pip install --no-cache-dir numpy requests packaging build RUN python3 -m pip install --no-cache-dir numpy requests packaging build
# 在 PyTorch 要求cmake >=3.27,ubuntu 22.04默认是cmake 3.22.1,所以现在需要安装新的:
RUN python3 -m pip install --no-cache-dir "cmake>=3.29,<4.0" "ninja>=1.11" && \
cmake --version && ninja --version
WORKDIR /opt WORKDIR /opt
# RUN git clone --recursive -b v2.8.0 https://github.com/pytorch/pytorch.git # RUN git clone --recursive -b v2.8.0 https://github.com/pytorch/pytorch.git
COPY ./pytorch_2.8.0/ /opt/pytorch COPY ./pytorch_2.8.0/ /opt/pytorch
@ -42,7 +46,7 @@ RUN echo "Building PyTorch with USE_DISTRIBUTED=$USE_DISTRIBUTED" && \
ARG CUDA_VERSION=12.6.1 ARG CUDA_VERSION=12.6.1
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS builder-extras FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS builder-extras
ENV TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9" ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9"
ENV DEBIAN_FRONTEND=noninteractive PYTHONUNBUFFERED=1 LANG=C.UTF-8 LC_ALL=C.UTF-8 ENV DEBIAN_FRONTEND=noninteractive PYTHONUNBUFFERED=1 LANG=C.UTF-8 LC_ALL=C.UTF-8
RUN apt-get update && apt-get install -y --no-install-recommends \ RUN apt-get update && apt-get install -y --no-install-recommends \
@ -54,6 +58,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
curl xz-utils \ curl xz-utils \
&& python3 -m pip install --no-cache-dir --upgrade pip wheel setuptools && python3 -m pip install --no-cache-dir --upgrade pip wheel setuptools
# for torch vision以及sglang 0.5.2
RUN python3 -m pip install --no-cache-dir "cmake>=3.29,<4.0" "ninja>=1.11" && \
cmake --version && ninja --version
# ── 安装自编 torch 轮子 ────────────────────────────────────────────────────── # ── 安装自编 torch 轮子 ──────────────────────────────────────────────────────
COPY --from=builder-torch /opt/pytorch/dist /tmp/torch_dist COPY --from=builder-torch /opt/pytorch/dist /tmp/torch_dist
RUN set -e && \ RUN set -e && \
@ -83,7 +91,7 @@ WORKDIR /opt/flashinfer
# 覆盖你的目标算力3090=8.64090=8.9H100=9.0a;可按需增/减 # 覆盖你的目标算力3090=8.64090=8.9H100=9.0a;可按需增/减
ENV FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9" ENV FLASHINFER_CUDA_ARCH_LIST="8.0;8.6;8.9"
# 先做 AOT 预编译,再直接打 wheel不隔离使用同一份自编 torch # 先做 AOT 预编译,再直接打 wheel不隔离使用同一份自编 torch
RUN python3 -m pip install --no-cache-dir numpy requests build "cuda-python>=12.0,<13" "nvidia-nvshmem-cu12" ninja pynvml && \ RUN python3 -m pip install --no-cache-dir numpy requests build "cuda-python>=12.0,<13" "nvidia-nvshmem-cu12" ninja pynvml && \