diff --git a/builder/manywheel/README.md b/builder/manywheel/README.md index 1d936f9879..8cef3b19ee 100644 --- a/builder/manywheel/README.md +++ b/builder/manywheel/README.md @@ -3,7 +3,8 @@ ## Building lmdeploy builder images To build all lmdeploy builder images, such as -`lmdeploy-builder:manylinux_2_28-cuda12.8`, execute: +`lmdeploy-builder:manylinux_2_28-cuda12.8` and +`lmdeploy-builder:manylinux_2_28-cuda13.0`, execute: ```bash ./build_all_lmdeploy_builders.sh @@ -15,7 +16,7 @@ WITH_PUSH=true ./build_all_lmdeploy_builders.sh For custom builds with specific versions: ```bash -MANY_LINUX_VERSION=2_28 GPU_ARCH_VERSION=12.8 ./build_lmdeploy_builder.sh +MANY_LINUX_VERSION=2_28 GPU_ARCH_VERSION=13.0 ./build_lmdeploy_builder.sh ``` ## Build lmdeploy wheels diff --git a/docker/Dockerfile b/docker/Dockerfile index 099ef5ab3a..12636635e2 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,15 +1,20 @@ # Base images ARG IMAGE_TYPE=final -ARG CUDA_VERSION=cu12 +ARG CUDA_VERSION=cu128 -FROM nvidia/cuda:13.0.2-devel-ubuntu22.04 AS cu13 +FROM nvidia/cuda:13.0.3-devel-ubuntu24.04 AS cu130 ENV CUDA_VERSION_SHORT=cu130 +ENV PYTORCH_CUDA_VERSION=cu130 -FROM nvidia/cuda:12.8.1-devel-ubuntu22.04 AS cu12.8 +FROM nvidia/cuda:12.8.1-devel-ubuntu24.04 AS cu128 ENV CUDA_VERSION_SHORT=cu128 +# PyTorch 2.12 no longer publishes cu128 wheels; cu126 is its CUDA 12 build. +ENV PYTORCH_CUDA_VERSION=cu126 -FROM nvidia/cuda:12.6.3-devel-ubuntu22.04 AS cu12 -ENV CUDA_VERSION_SHORT=cu126 +# Backward-compatible stage aliases. +FROM cu130 AS cu13 +FROM cu128 AS cu12.8 +FROM cu128 AS cu12 # Builder image FROM ${CUDA_VERSION} AS dev @@ -20,7 +25,13 @@ ENV DEBIAN_FRONTEND=noninteractive ENV TZ=Etc/UTC RUN --mount=type=cache,target=/root/.cache \ - sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list && \ + if [ -f /etc/apt/sources.list ]; then \ + sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list; \ + fi && \ + if [ -d /etc/apt/sources.list.d ]; then \ + find /etc/apt/sources.list.d -type f -name '*.sources' -exec \ + sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' {} +; \ + fi && \ apt-get update -y && \ apt-get install -y --no-install-recommends \ tzdata wget curl openssh-server ssh sudo git-core \ @@ -40,7 +51,9 @@ ENV PYTHON_VERSION=${PYTHON_VERSION} COPY requirements/common.txt requirements/runtime_cuda.txt /tmp/requirements/ RUN --mount=type=cache,target=/root/.cache \ - pip install -r /tmp/requirements/runtime_cuda.txt --extra-index-url https://download.pytorch.org/whl/${CUDA_VERSION_SHORT} + grep -E '^torch(vision)?([<>=]|$)' /tmp/requirements/runtime_cuda.txt > /tmp/requirements/torch.txt && \ + pip install -r /tmp/requirements/torch.txt --index-url https://download.pytorch.org/whl/${PYTORCH_CUDA_VERSION} && \ + pip install -r /tmp/requirements/runtime_cuda.txt RUN mkdir -p /wheels @@ -67,24 +80,30 @@ COPY --from=thirdparty-wheels /wheels /wheels COPY --from=lmdeploy-wheel /wheels /wheels # Runtime image -FROM nvidia/cuda:13.0.2-base-ubuntu22.04 AS cu13-base +FROM nvidia/cuda:13.0.3-base-ubuntu24.04 AS cu130-base ENV CUDA_VERSION_SHORT=cu130 +ENV PYTORCH_CUDA_VERSION=cu130 -FROM nvidia/cuda:12.8.1-base-ubuntu22.04 AS cu12.8-base +FROM nvidia/cuda:12.8.1-base-ubuntu24.04 AS cu128-base ENV CUDA_VERSION_SHORT=cu128 +ENV PYTORCH_CUDA_VERSION=cu126 -FROM nvidia/cuda:12.6.3-base-ubuntu22.04 AS cu12-base -ENV CUDA_VERSION_SHORT=cu126 +# Backward-compatible stage aliases. +FROM cu130-base AS cu13-base +FROM cu128-base AS cu12.8-base +FROM cu128-base AS cu12-base FROM ${CUDA_VERSION}-base AS final ARG PYTHON_VERSION=3.12 ENV PYTHON_VERSION=${PYTHON_VERSION} +# Install the CUDA runtime requirements explicitly so torch is selected from +# the matching PyTorch CUDA index before local wheels resolve dependencies. # Some dependencies such as timm (required by InternVL models) are listed in requirements/serve.txt. -COPY requirements/serve.txt /tmp/requirements/serve.txt +COPY requirements/common.txt requirements/runtime_cuda.txt requirements/serve.txt /tmp/requirements/ COPY docker/install.sh /tmp/install.sh RUN --mount=type=cache,target=/root/.cache \ - --mount=type=cache,target=/wheels,from=wheels,source=/wheels \ + --mount=type=bind,target=/wheels,from=wheels,source=/wheels \ /tmp/install.sh # explicitly set ptxas path for triton diff --git a/docker/Dockerfile.jetson b/docker/Dockerfile.jetson index 1cca19a920..ac7d2f5871 100644 --- a/docker/Dockerfile.jetson +++ b/docker/Dockerfile.jetson @@ -27,12 +27,19 @@ RUN --mount=type=cache,target=/root/.cache \ COPY . /opt/lmdeploy WORKDIR /opt/lmdeploy -# The Jetson AI Lab index may lag xgrammar releases. Force xgrammar from PyPI, -# while keeping the LMDeploy wheel and its CUDA/Torch dependencies on the Jetson index. +# Install the Jetson-specific Torch stack from the Jetson AI Lab index first. +# Resolve generic LMDeploy dependencies from PyPI so packages missing from the +# Jetson index, such as pydantic, do not fail dependency resolution. RUN --mount=type=cache,target=/root/.cache \ - --mount=type=cache,target=/opt/pytorch \ + pip install --upgrade pip && \ pip install build change-wheel-version && \ python -m build -w -o /wheels -v . && \ change_wheel_version --local-version cu126 --delete-old-wheel /wheels/lmdeploy*.whl && \ - pip install -v --no-deps --only-binary=:all: xgrammar==0.1.33 --index-url https://pypi.org/simple/ && \ - pip install -v /wheels/lmdeploy*.whl --index-url https://pypi.jetson-ai-lab.io/jp6/cu126/+simple/ + grep -E '^torch(vision)?([<>=]|$)' requirements/runtime_cuda.txt > /tmp/requirements_torch.txt && \ + mkdir -p /tmp/pytorch-wheels && \ + pip download --no-deps --only-binary=:all: -r /tmp/requirements_torch.txt -d /tmp/pytorch-wheels \ + --index-url https://pypi.jetson-ai-lab.io/jp6/cu126/+simple/ && \ + pip install -v /tmp/pytorch-wheels/*.whl --index-url https://pypi.org/simple/ && \ + pip install -v --no-deps --only-binary=:all: xgrammar==0.1.33 \ + --index-url https://pypi.org/simple/ && \ + pip install -v /wheels/lmdeploy*.whl --index-url https://pypi.org/simple/ diff --git a/docker/Dockerfile_dev b/docker/Dockerfile_dev index 3b8874442f..3b6f0a6d5f 100644 --- a/docker/Dockerfile_dev +++ b/docker/Dockerfile_dev @@ -1,16 +1,23 @@ -FROM nvidia/cuda:12.8.1-devel-ubuntu22.04 AS cu12.8 +FROM nvidia/cuda:12.8.1-devel-ubuntu24.04 AS cu128 # environment variables ENV DEBIAN_FRONTEND=noninteractive \ TZ=Etc/UTC \ PATH=/opt/py3/bin:/root/.local/bin:${PATH} \ TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas \ - CUDA_VERSION_SHORT=cu128 + CUDA_VERSION_SHORT=cu128 \ + PYTORCH_CUDA_VERSION=cu126 # Install dependencies and create python virtual environment RUN --mount=type=cache,target=/var/cache/apt \ --mount=type=cache,target=/root/.cache \ - sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list && \ + if [ -f /etc/apt/sources.list ]; then \ + sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list; \ + fi && \ + if [ -d /etc/apt/sources.list.d ]; then \ + find /etc/apt/sources.list.d -type f -name '*.sources' -exec \ + sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' {} +; \ + fi && \ apt-get update -y && \ apt-get install -y --no-install-recommends \ tzdata wget curl openssh-server ssh sudo git-core \ @@ -33,7 +40,9 @@ RUN mkdir -p /wheels # install lmdeploy and its dependencies RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements_cuda.txt --extra-index-url https://download.pytorch.org/whl/cu128 && \ + grep -E '^torch(vision)?([<>=]|$)' requirements/runtime_cuda.txt > /tmp/requirements_torch.txt && \ + uv pip install -r /tmp/requirements_torch.txt --index-url https://download.pytorch.org/whl/${PYTORCH_CUDA_VERSION} && \ + uv pip install -r requirements_cuda.txt && \ uv pip install -e . RUN --mount=type=cache,target=/root/.cache/uv \ diff --git a/docker/build.sh b/docker/build.sh index ec492e6643..8b43d11a4e 100755 --- a/docker/build.sh +++ b/docker/build.sh @@ -2,7 +2,7 @@ mkdir -p /wheels -if [[ "${CUDA_VERSION_SHORT}" = "cu130" ]]; then +if [[ "${CUDA_VERSION_SHORT}" == cu13* ]]; then pip install nvidia-nccl-cu13 else pip install nvidia-nccl-cu12 diff --git a/docker/install.sh b/docker/install.sh index 5bd9cef2b0..6142d6da31 100755 --- a/docker/install.sh +++ b/docker/install.sh @@ -4,20 +4,40 @@ if [ ! -f "/opt/py3/bin/python" ]; then # install system packages export DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC - sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list + if [ -f /etc/apt/sources.list ]; then + sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list + fi + if [ -d /etc/apt/sources.list.d ]; then + find /etc/apt/sources.list.d -type f -name '*.sources' -exec \ + sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' {} + + fi apt-get update -y apt-get install -y --no-install-recommends \ tzdata wget curl ssh sudo git-core vim libibverbs1 ibverbs-providers ibverbs-utils librdmacm1 libibverbs-dev rdma-core libmlx5-1 - if [[ ${PYTHON_VERSION} != "3.10" ]]; then + # Ubuntu may expose the interpreter package without the matching -dev and + # -venv packages. Check the complete set before deciding whether the PPA + # is needed. + python_packages=( + "python${PYTHON_VERSION}" + "python${PYTHON_VERSION}-dev" + "python${PYTHON_VERSION}-venv" + ) + need_python_ppa=0 + for package in "${python_packages[@]}"; do + if ! apt-cache show "${package}" >/dev/null 2>&1; then + need_python_ppa=1 + break + fi + done + if ((need_python_ppa)); then apt-get install -y --no-install-recommends software-properties-common add-apt-repository -y ppa:deadsnakes/ppa apt-get update -y fi # install python, create virtual env - apt-get install -y --no-install-recommends \ - python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv + apt-get install -y --no-install-recommends "${python_packages[@]}" pushd /opt >/dev/null python${PYTHON_VERSION} -m venv py3 @@ -46,13 +66,35 @@ export PATH=/opt/py3/bin:$PATH pip install -U pip wheel setuptools -if [[ "${CUDA_VERSION_SHORT}" = "cu130" ]]; then +# In the production image, install torch/torchvision from the matching CUDA +# index first. Installing all local wheels directly would otherwise let pip +# resolve torch from the default PyPI index. +if [ -f /tmp/requirements/runtime_cuda.txt ]; then + grep -E '^torch(vision)?([<>=]|$)' \ + /tmp/requirements/runtime_cuda.txt > /tmp/requirements/torch.txt + pip install -r /tmp/requirements/torch.txt \ + --index-url "https://download.pytorch.org/whl/${PYTORCH_CUDA_VERSION}" + pip install -r /tmp/requirements/runtime_cuda.txt +fi + +if [[ "${CUDA_VERSION_SHORT}" == cu13* ]]; then pip install nvidia-nvshmem-cu13==3.4.5 else pip install nvidia-nvshmem-cu12==3.4.5 fi +for package in deep_ep deep_gemm flash_attn_3; do + if ! compgen -G "/wheels/${package}-*.whl" >/dev/null; then + echo "${package} wheel is missing from /wheels" >&2 + exit 1 + fi +done pip install /wheels/*.whl +python - <<'PY' +import deep_ep +import deep_gemm +import lmdeploy.pytorch.third_party.flash_attn_interface +PY pip install dlslime==0.0.2.post1 pip install ninja einops packaging @@ -62,9 +104,48 @@ if [ -f /tmp/requirements/serve.txt ]; then pip install -r /tmp/requirements/serve.txt fi +# PyTorch cu130 pins an older NCCL release in its wheel metadata. Check the +# declared dependencies before intentionally upgrading NCCL for GIN support. +pip check + if [[ "${CUDA_VERSION_SHORT}" = "cu128" ]]; then - # As described in https://github.com/InternLM/lmdeploy/pull/4313, - # window registration may cause memory leaks in NCCL 2.27, NCCL 2.28+ resolves the issue, - # but turbomind engine will use nccl GIN for EP in future, which is brought in since 2.29 - pip install "nvidia-nccl-cu12>2.29" + pip install "nvidia-nccl-cu12>=2.30.4" +elif [[ "${CUDA_VERSION_SHORT}" == cu13* ]]; then + pip install "nvidia-nccl-cu13>=2.30.4" fi + +python - <<'PY' +import os +from importlib.metadata import distributions, version + +import torch +from packaging.version import Version + +expected = os.environ['PYTORCH_CUDA_VERSION'] +actual = torch.version.cuda +installed = { + dist.metadata['Name'].lower().replace('_', '-') + for dist in distributions() + if dist.metadata['Name'] +} +if expected == 'cu126': + assert actual is not None and actual.startswith('12.6'), actual + nccl_package = 'nvidia-nccl-cu12' + nccl_minimum = Version('2.29') + unexpected = sorted(name for name in installed if name.endswith('-cu13')) +elif expected == 'cu130': + assert actual is not None and actual.startswith('13.0'), actual + nccl_package = 'nvidia-nccl-cu13' + nccl_minimum = Version('2.30.4') + unexpected = sorted(name for name in installed if name.endswith('-cu12')) +else: + raise AssertionError(f'unsupported PYTORCH_CUDA_VERSION: {expected}') +assert not unexpected, f'unexpected CUDA packages: {unexpected}' +nccl_version = version(nccl_package) +if expected == 'cu126': + assert Version(nccl_version) > nccl_minimum, nccl_version +else: + assert Version(nccl_version) >= nccl_minimum, nccl_version +print(f'Validated torch {torch.__version__} with CUDA {actual} and ' + f'{nccl_package} {nccl_version}') +PY diff --git a/docker/prepare_3rdparty_wheel.sh b/docker/prepare_3rdparty_wheel.sh index a5a2eb2dcf..cb319fd3da 100644 --- a/docker/prepare_3rdparty_wheel.sh +++ b/docker/prepare_3rdparty_wheel.sh @@ -33,10 +33,16 @@ pip wheel -v --no-build-isolation --no-deps -w /wheels "git+https://github.com/d # sm100 compilation for Flash MLA requires NVCC 12.9 or higher FLASH_MLA_DISABLE_SM100=1 pip wheel -v --no-build-isolation --no-deps -w /wheels "git+https://github.com/deepseek-ai/FlashMLA.git@${FLASH_MLA_VERSION}" -# flash_attn_3 (prebuilt wheels; CUDA + torch must match this image) -TORCH_VER=$(python3 -c "import torch; print(''.join(torch.__version__.split('+')[0].split('.')))") -FA3_WHEELS_URL="https://windreamer.github.io/flash-attention3-wheels/${CUDA_VERSION_SHORT}_torch${TORCH_VER}" -pip download --no-deps -d /wheels "flash_attn_3" --find-links "${FA3_WHEELS_URL}" +# flash_attn_3 (official, LibTorch ABI-stable wheels for torch >= 2.9). +# Select the FA3 channel matching the PyTorch binary. CUDA 13 FA3 wheels are +# published in the cu130 channel. +FA3_CUDA_VERSION="${PYTORCH_CUDA_VERSION:-${CUDA_VERSION_SHORT}}" +if [[ "${FA3_CUDA_VERSION}" == cu13* ]]; then + FA3_CUDA_VERSION=cu130 +fi +pip download --no-deps --only-binary=:all: -d /wheels \ + "flash-attn-3==3.0.0" \ + --index-url "https://download.pytorch.org/whl/${FA3_CUDA_VERSION}" # GDRCopy debs apt-get update -y \ diff --git a/lmdeploy/pytorch/check_env/triton.py b/lmdeploy/pytorch/check_env/triton.py index 11661cb5bd..796576b502 100644 --- a/lmdeploy/pytorch/check_env/triton.py +++ b/lmdeploy/pytorch/check_env/triton.py @@ -3,7 +3,7 @@ from .base import BaseChecker -MAX_TRITON_VERSION = '3.6.0' +MAX_TRITON_VERSION = '3.7.1' MIN_TRITON_VERSION = '3.0.0' diff --git a/requirements/runtime_cuda.txt b/requirements/runtime_cuda.txt index eb091e0f66..21bbc98c92 100644 --- a/requirements/runtime_cuda.txt +++ b/requirements/runtime_cuda.txt @@ -8,6 +8,6 @@ opencv-python-headless peft<=0.14.0 prometheus_client tilelang==0.1.11; sys_platform == "linux" and "aarch64" not in platform_machine and "arm" not in platform_machine -torch<=2.10.0,>=2.0.0 -torchvision<=0.25.0,>=0.15.0 -triton<=3.6.0,>=3.0.0; sys_platform == "linux" and "aarch64" not in platform_machine and "arm" not in platform_machine +torch<=2.12.1,>=2.0.0 +torchvision<=0.27.1,>=0.15.0 +triton<=3.7.1,>=3.0.0; sys_platform == "linux" and "aarch64" not in platform_machine and "arm" not in platform_machine