11430aaab5
The notebook validation matrix caught the synthetic-data notebook dying because the vllm server SyntheticDataKit launches never came up. Two layers to the failure: 1. flashinfer's cpp_ext JIT shells out to ninja. The pip ninja lives in the venv bin, which subprocesses like vllm serve do not always inherit on PATH, so the JIT failed with exit 127. Install ninja-build so the binary is reachable from any PATH. 2. With ninja present the JIT still cannot succeed for device code: the runtime image deliberately ships no nvcc. Bake flashinfer-jit-cache (cu128) so ops missing from the cubin package (fmha_gen on sm_100a was the repro) come precompiled. In-process GRPO never hit this because unsloth-zoo blocks the FlashInfer JIT path; standalone vllm serve gets no zoo patches. Fail-soft on the jit-cache for arches without a wheel; the vLLM chain itself stays fail-loud on amd64.
567 lines
31 KiB
Docker
567 lines
31 KiB
Docker
# syntax=docker/dockerfile:1.7
|
|
# -----------------------------------------------------------------------------
|
|
# Unsloth + unsloth-zoo for every current NVIDIA arch (Turing -> Blackwell),
|
|
# on both linux/amd64 and linux/arm64.
|
|
#
|
|
# Why this image works:
|
|
# * cu128 wheels ship native SASS (no PTX), empirically verified via
|
|
# `cuobjdump --list-elf` against the downloaded wheels:
|
|
# amd64: sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120
|
|
# arm64: sm_80 sm_90 sm_90a sm_100 sm_100a sm_120 sm_120a
|
|
# * SASS is binary-compatible UPWARDS within a major (per ptrblck on the
|
|
# PyTorch forum, May 2026): sm_86 SASS runs on sm_89 hardware (Ada);
|
|
# sm_100 SASS runs on sm_103 (B300/GB300); sm_120 SASS runs on sm_121
|
|
# (DGX Spark / GB10). So every non-Jetson NVIDIA GPU on
|
|
# https://developer.nvidia.com/cuda/gpus is covered.
|
|
# * Unsloth's runtime kernels are Triton, which JIT-compiles per device at first run.
|
|
# * Anything that DOES need to be source-built (rare on this pin set) compiles
|
|
# against TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;10.3;12.0+PTX",
|
|
# covering every current NVIDIA compute capability per
|
|
# https://developer.nvidia.com/cuda/gpus.
|
|
# The host GPU is irrelevant for compilation; nvcc emits whatever the arch
|
|
# list says.
|
|
#
|
|
# Cross-arch build (DGX Spark / GB10 / sm_121):
|
|
# The arm64 image is built via QEMU binfmt emulation on an x86_64 host:
|
|
# bash docker/setup_qemu.sh # one-time host setup
|
|
# docker buildx build --platform linux/arm64 -t unsloth-blackwell:arm64 .
|
|
# The resulting arm64 image runs NATIVELY on aarch64 hosts (DGX Spark, Grace).
|
|
# QEMU is only used at build time -- runtime emulation does NOT work for CUDA.
|
|
# xformers has no cu128 aarch64 wheel; on arm64 we fall back to Unsloth's
|
|
# built-in SDPA path (~5-10% slower than xformers but functionally complete).
|
|
#
|
|
# Build host requirements:
|
|
# * Docker with buildkit (default since 23.x)
|
|
# * docker buildx (mandatory for multi-platform; install: apt install docker-buildx)
|
|
# * nvidia-container-toolkit (only needed for `docker run --gpus all` at test time)
|
|
# * For arm64 builds on x86_64 hosts: QEMU binfmt (see docker/setup_qemu.sh)
|
|
# * A GPU is NOT required at build time.
|
|
# -----------------------------------------------------------------------------
|
|
|
|
ARG CUDA_VERSION=12.8.1
|
|
ARG UBUNTU_VERSION=24.04
|
|
ARG PYTHON_VERSION=3.12
|
|
|
|
# =============================================================================
|
|
# Stage 1: builder -- toolkit + dev headers, builds any source extensions
|
|
# =============================================================================
|
|
FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu${UBUNTU_VERSION} AS builder
|
|
|
|
# TARGETARCH is auto-populated by buildx ("amd64" or "arm64"). We use it to
|
|
# select an unsloth extras set that matches the wheels actually available for
|
|
# the target platform (xformers has no cu128 aarch64 wheel as of 0.0.34).
|
|
ARG TARGETARCH
|
|
ARG PYTHON_VERSION
|
|
ENV DEBIAN_FRONTEND=noninteractive \
|
|
PIP_NO_CACHE_DIR=1 \
|
|
PIP_DISABLE_PIP_VERSION_CHECK=1 \
|
|
PYTHONDONTWRITEBYTECODE=1 \
|
|
PYTHONUNBUFFERED=1 \
|
|
# Cross-compile for every current NVIDIA arch per
|
|
# https://developer.nvidia.com/cuda/gpus:
|
|
# sm_75 Turing T4, RTX 20-series, Quadro RTX (amd64 only)
|
|
# sm_80 Ampere DC A100, A30 (amd64 only)
|
|
# sm_86 Ampere A40, RTX A6000, RTX 30-series (amd64 only)
|
|
# sm_89 Ada L4, L40, L40S, RTX 40-series (amd64 only)
|
|
# sm_90 Hopper H100, H200, GH200 (Grace-Hopper is arm64)
|
|
# sm_100 Blackwell DC B100, B200, GB200 (GB200 is arm64)
|
|
# sm_103 Blackwell DC B300, GB300
|
|
# sm_120 Blackwell RTX 50-series, RTX PRO 6000 Blackwell (amd64 only)
|
|
# sm_121 Blackwell GB10 (DGX Spark) (arm64 only)
|
|
# +PTX on the highest lets future arch revisions run via JIT-PTX.
|
|
# We keep the same list on both arches: nvcc happily emits archs that
|
|
# don't exist on the build host, and any extras that aren't relevant for
|
|
# the target just bloat compile time slightly (not size, since we don't
|
|
# source-build any extension at install time on the pin set).
|
|
TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;10.3;12.0+PTX" \
|
|
MAX_JOBS=4 \
|
|
CUDA_HOME=/usr/local/cuda \
|
|
# Build-host-independence guards. The build must NEVER introspect a GPU,
|
|
# because the build host may have a B200, RTX 6000, or no GPU at all
|
|
# (GitHub Actions ubuntu-latest). All three must yield byte-identical images.
|
|
#
|
|
# 1) Stop unsloth from JIT-compiling kernels at import time and writing a
|
|
# sm_NNN-specific blob into /opt/unsloth-venv/.../unsloth_compiled_cache/.
|
|
UNSLOTH_COMPILE_DISABLE=1 \
|
|
UNSLOTH_COMPILE_OVERWRITE=0 \
|
|
# 2) Stop unsloth-zoo / vllm from probing torch.cuda.is_available() during
|
|
# setup. There's no GPU here, and we don't want it to silently skip a wheel.
|
|
UNSLOTH_DISABLE_GPU_PROBE=1 \
|
|
# 3) Force CUDA_VISIBLE_DEVICES empty so any stray torch.cuda call during
|
|
# `pip install` returns "no devices" rather than triggering host-specific
|
|
# code paths (we re-enable at runtime via `docker run --gpus all`).
|
|
CUDA_VISIBLE_DEVICES=""
|
|
|
|
RUN apt-get update && apt-get install -y --no-install-recommends \
|
|
software-properties-common ca-certificates curl git build-essential \
|
|
ninja-build cmake pkg-config \
|
|
&& add-apt-repository -y ppa:deadsnakes/ppa \
|
|
&& apt-get update && apt-get install -y --no-install-recommends \
|
|
python${PYTHON_VERSION} python${PYTHON_VERSION}-venv python${PYTHON_VERSION}-dev \
|
|
&& ln -sf /usr/bin/python${PYTHON_VERSION} /usr/local/bin/python \
|
|
&& ln -sf /usr/bin/python${PYTHON_VERSION} /usr/local/bin/python3 \
|
|
&& rm -rf /var/lib/apt/lists/*
|
|
|
|
# Build into an isolated prefix. NOTE: we do NOT install pip or uv into the
|
|
# system Python -- on Ubuntu 24.04 the system interpreter is marked
|
|
# externally-managed (PEP 668) and `pip install` is refused. Instead, the new
|
|
# venv bootstraps its own pip via ensurepip (provided by the python3.12-venv
|
|
# apt package), and we install uv into the venv a few lines below.
|
|
ENV VENV=/opt/unsloth-venv
|
|
RUN python -m venv ${VENV} && ${VENV}/bin/pip install -U pip wheel setuptools
|
|
|
|
# Unified install: torch + triton + bitsandbytes + unsloth + unsloth_zoo
|
|
# resolve in a SINGLE uv pip pass. This is mandatory -- splitting it across
|
|
# multiple `pip install` calls causes bnb's transitive `cuda-toolkit` dep to
|
|
# silently upgrade torch to 2.12.0+cu130 in a later pass, breaking the cu128
|
|
# xformers wheel that was pinned earlier. (Empirically discovered; the cu cascade
|
|
# happens AFTER xformers is already on disk, leaving a working-but-mismatched env.)
|
|
#
|
|
# uv-specific flags explained:
|
|
# --index-strategy unsafe-best-match
|
|
# The PyTorch index serves an old `requests==2.28.1` which conflicts with
|
|
# `datasets>=2.32.2`. uv's default is "first index wins per package" to
|
|
# prevent dependency confusion; we override here because both indexes
|
|
# (pytorch.org/whl/cu128 + pypi.org) are equally trusted.
|
|
# --extra-index-url https://download.pytorch.org/whl/cu128
|
|
# Where torch's +cu128 wheels live, plus the xformers/cu128 URLs referenced
|
|
# by unsloth's `cu128onlytorch2100` extra.
|
|
#
|
|
# Why the extra is `cu128-ampere-torch2100` (not `cu128-torch2100-ampere`):
|
|
# See unsloth_src/pyproject.toml:835. The ordering is ampere-then-torch-ver.
|
|
#
|
|
# Why arm64 uses a different extra:
|
|
# `cu128-ampere-torch2100` transitively pulls `cu128onlytorch2100` whose
|
|
# xformers wheel URL is hardcoded to manylinux_2_28_x86_64. There is no
|
|
# cu128 aarch64 wheel for xformers as of 0.0.34. We use the plain
|
|
# `huggingface` extra on arm64 -- Unsloth falls back to its native SDPA
|
|
# kernels (a ~5-10% slowdown vs xformers; functionally complete).
|
|
#
|
|
# Why no `flash-attn` here:
|
|
# - FA3 is hard-refused on Blackwell (Dao-AILab/flash-attention#1810).
|
|
# - FA2 has no prebuilt wheel for cu128+torch2.10+cp312 -> would require
|
|
# a ~30min source build, fragile on the 16GB ubuntu-latest CI runner.
|
|
# - Unsloth gracefully falls back to xformers/SDPA on Blackwell anyway.
|
|
# - Users on Ampere/Ada/Hopper who want FA2 can `pip install flash-attn`
|
|
# on top of this image at deploy time.
|
|
ARG UNSLOTH_REF=main
|
|
ARG UNSLOTH_ZOO_REF=main
|
|
RUN set -eux \
|
|
&& case "${TARGETARCH:-amd64}" in \
|
|
amd64) UNSLOTH_EXTRA="cu128-ampere-torch2100" ;; \
|
|
arm64) UNSLOTH_EXTRA="huggingface" ;; \
|
|
*) echo "ERROR: unsupported TARGETARCH=${TARGETARCH}" >&2; exit 1 ;; \
|
|
esac \
|
|
&& echo ">> TARGETARCH=${TARGETARCH:-amd64}, unsloth extra=[${UNSLOTH_EXTRA}]" \
|
|
&& ${VENV}/bin/pip install uv \
|
|
&& ${VENV}/bin/uv pip install \
|
|
--python ${VENV}/bin/python \
|
|
--index-strategy unsafe-best-match \
|
|
--extra-index-url https://download.pytorch.org/whl/cu128 \
|
|
"torch==2.10.0" "torchvision==0.25.0" "torchaudio==2.10.0" \
|
|
"triton>=3.6.0" \
|
|
"bitsandbytes>=0.49.2,!=0.46.0,!=0.48.0" \
|
|
"unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo@${UNSLOTH_ZOO_REF}" \
|
|
"unsloth[${UNSLOTH_EXTRA}] @ git+https://github.com/unslothai/unsloth@${UNSLOTH_REF}" \
|
|
"timm>=1.0.11" "addict"
|
|
|
|
# vLLM. Required by Unsloth's GRPO path when the notebook sets
|
|
# fast_inference=True. We install it as a SECOND uv pass rather than
|
|
# appending to the unified one because:
|
|
# * vLLM releases pin an exact torch (0.19.1 -> torch==2.10.0; 0.20+
|
|
# moved to torch==2.11.0). Adding it to the unified resolve forces uv
|
|
# to consider whether to swap our pinned torch 2.10.0 for vLLM's
|
|
# choice. Splitting the install lets the unified pass settle on torch
|
|
# 2.10.0 first, then vLLM bolts on top: with torch held at 2.10.0 the
|
|
# resolver lands on the newest compatible vLLM (0.19.1) by itself, and
|
|
# moves forward automatically when we bump torch.
|
|
# * PyPI ships both x86_64 AND aarch64 abi3 wheels for every release
|
|
# since 0.17, so the SAME pass now runs on the arm64 leg (DGX Spark /
|
|
# GB10 class). amd64 failures abort the build; arm64 is fail-soft
|
|
# because aarch64 wheels are newer and the GPU-side kernels there are
|
|
# validated on Spark hardware via docker_confirm.sh, not in CI.
|
|
#
|
|
# https://docs.vllm.ai/en/latest/getting_started/installation/gpu/
|
|
# https://wheels.vllm.ai/nightly
|
|
ARG INSTALL_VLLM=auto
|
|
RUN set -eux \
|
|
&& WANT_VLLM=0 \
|
|
&& case "${INSTALL_VLLM}" in \
|
|
auto|1|true|yes) WANT_VLLM=1 ;; \
|
|
0|false|no) WANT_VLLM=0 ;; \
|
|
*) echo "ERROR: invalid INSTALL_VLLM=${INSTALL_VLLM}" >&2; exit 1 ;; \
|
|
esac \
|
|
&& if [ "${WANT_VLLM}" = "1" ]; then \
|
|
echo ">> installing vLLM (TARGETARCH=${TARGETARCH:-amd64})"; \
|
|
# NOTE: an explicit && chain, not `set -e` -- POSIX shells disable
|
|
# errexit inside any condition context (verified on dash), so a
|
|
# (set -e; ...) subshell here would mask install failures.
|
|
# Step 1: let uv resolve vLLM's transitive deps. We pin
|
|
# torch==2.10.0 so uv MUST hold our torch fixed; with torch held,
|
|
# the resolver lands on the newest vLLM whose pin matches and the
|
|
# build fails loudly if none does. `unsafe-best-match` lets uv
|
|
# pull from whichever of the three indexes has a better wheel.
|
|
# Step 2: vLLM pulls numpy down to 2.2.6 whose wheel ships a
|
|
# broken numpy.testing (`from numpy._core.tests._natype import
|
|
# pd_NA` -- tests/ is stripped from the wheel). Any path that hits
|
|
# `from numpy import *` (e.g. scipy.optimize) then crashes, taking
|
|
# `import unsloth` with it. Upgrade numpy back to a release with a
|
|
# self-consistent testing module.
|
|
# Step 3: vLLM pins numba==0.61.2, which hard-refuses numpy >= 2.3
|
|
# at import time -- and the stack needs numpy >= 2.3, so the numpy
|
|
# ceiling cannot move down. Lift numba to a release that supports
|
|
# numpy 2.4 (verified: numba 0.65 imports cleanly, vllm still
|
|
# imports). Same intentional-override class as the numpy bump.
|
|
{ ${VENV}/bin/uv pip install \
|
|
--python ${VENV}/bin/python \
|
|
--pre \
|
|
--index-strategy unsafe-best-match \
|
|
--extra-index-url https://wheels.vllm.ai/nightly \
|
|
--extra-index-url https://download.pytorch.org/whl/cu128 \
|
|
"torch==2.10.0" \
|
|
vllm \
|
|
&& ${VENV}/bin/uv pip install \
|
|
--python ${VENV}/bin/python \
|
|
--upgrade "numpy>=2.4" \
|
|
&& ${VENV}/bin/uv pip install \
|
|
--python ${VENV}/bin/python \
|
|
--upgrade "numba>=0.62" \
|
|
&& ${VENV}/bin/python -c "import vllm; print('vllm', vllm.__version__)" \
|
|
&& ${VENV}/bin/python -c "import numpy.testing, numpy; print('numpy', numpy.__version__, 'testing ok')" \
|
|
&& ${VENV}/bin/python -c "import numba; print('numba', numba.__version__, 'imports ok')" \
|
|
# flashinfer-jit-cache: the runtime image ships no nvcc, so any
|
|
# flashinfer op missing from the cubin package would hit the JIT
|
|
# path and die (standalone `vllm serve` does exactly this for
|
|
# fmha_gen on sm_100a; in-process GRPO survives because zoo blocks
|
|
# the FlashInfer JIT). The precompiled cache removes the entire
|
|
# runtime-compile failure class for ~1.5 GB.
|
|
&& { ${VENV}/bin/uv pip install \
|
|
--python ${VENV}/bin/python \
|
|
--index-url https://flashinfer.ai/whl/cu128 \
|
|
"flashinfer-jit-cache==0.6.6" \
|
|
|| echo ">> flashinfer-jit-cache unavailable for ${TARGETARCH:-amd64}; vllm serve may require nvcc for uncached ops"; } \
|
|
&& echo ">> vLLM installed (numpy + numba re-upgraded post-vllm)"; \
|
|
} || { \
|
|
if [ "${TARGETARCH:-amd64}" != "amd64" ]; then \
|
|
echo ">> vLLM skipped on ${TARGETARCH}: install or import check failed (fail-soft on non-amd64)"; \
|
|
# A partial install must not poison the base stack: drop
|
|
# vllm and restore the numpy/numba floor it may have moved
|
|
# (numpy 2.2.6's broken numpy.testing breaks `import
|
|
# unsloth` outright). The arm64 staging CI contract probe
|
|
# re-verifies `import unsloth` after this.
|
|
${VENV}/bin/uv pip uninstall --python ${VENV}/bin/python vllm || true; \
|
|
${VENV}/bin/uv pip install --python ${VENV}/bin/python \
|
|
--upgrade "numpy>=2.4" "numba>=0.62"; \
|
|
${VENV}/bin/python -c "import numpy.testing, numba; print('numpy/numba restored')"; \
|
|
else \
|
|
echo "ERROR: vLLM install failed on amd64" >&2; exit 1; \
|
|
fi; \
|
|
}; \
|
|
else \
|
|
echo ">> vLLM skipped (INSTALL_VLLM=${INSTALL_VLLM}, TARGETARCH=${TARGETARCH:-amd64})"; \
|
|
fi
|
|
|
|
# JupyterLab so the published image runs unslothai/notebooks out of the box:
|
|
# docker run --gpus all -p 8888:8888 unsloth/unsloth \
|
|
# jupyter lab --ip 0.0.0.0 --port 8888 --allow-root --no-browser
|
|
# Installed as a separate pass AFTER the torch-pinned resolves on purpose:
|
|
# jupyterlab's dependency closure is pure-Python (tornado, jinja2, nbconvert,
|
|
# nbclient, ipykernel, ...) and never names torch, so uv cannot disturb the
|
|
# cu128 pin set here. Naming torch in this pass would be actively dangerous:
|
|
# without the cu128 extra index uv could swap in the PyPI CPU wheel.
|
|
# matplotlib rides along for the notebook crowd: plotting is table stakes in
|
|
# a Jupyter image, and several model repos' trust_remote_code modeling files
|
|
# (e.g. DeepSeek-OCR) import it unconditionally.
|
|
RUN ${VENV}/bin/uv pip install \
|
|
--python ${VENV}/bin/python \
|
|
jupyterlab notebook ipywidgets matplotlib
|
|
|
|
# Audio decode out of the box: the TTS/STT notebooks feed datasets' Audio
|
|
# features, which decode through torchcodec. Three traps, all defended:
|
|
# * version pairing: torchcodec 0.10 pairs with torch 2.10 (newer builds
|
|
# reference torch 2.11+ symbols and fail to dlopen);
|
|
# * CUDA line: the PyPI default wheel pairs with the cu13 torch line and
|
|
# dlopens libnvrtc.so.13 -- it must come from the cu128 channel;
|
|
# * its libraries dlopen torch + NVIDIA runtime libs that live inside the
|
|
# venv where ld.so cannot see them. Registered via ld.so.conf.d in the
|
|
# RUNTIME stage (this builder layer only installs the wheels, which ride
|
|
# along in the venv copy; the import check needs ffmpeg, which only the
|
|
# runtime stage installs).
|
|
# Fail-soft on arches without a matching wheel.
|
|
RUN set -eux \
|
|
&& { ${VENV}/bin/uv pip install \
|
|
--python ${VENV}/bin/python \
|
|
--index-url https://download.pytorch.org/whl/cu128 \
|
|
"torchcodec==0.10.0" \
|
|
&& ${VENV}/bin/uv pip install --python ${VENV}/bin/python nvidia-npp-cu12; } \
|
|
|| echo ">> torchcodec bake skipped (no matching wheel for ${TARGETARCH:-amd64})"
|
|
|
|
# 5) Emit an informational pin record so downstream consumers can see exactly
|
|
# what was resolved. This is NOT a byte-reproducible lockfile -- `pip freeze`
|
|
# captures version strings but not wheel hashes, and several deps (unsloth,
|
|
# unsloth_zoo, vllm --pre) resolve from VCS / nightly indexes that float.
|
|
# Read it with `docker run --rm <image> cat /opt/unsloth-venv/requirements.lock.txt`
|
|
# if you need to forensic-diff two images.
|
|
RUN ${VENV}/bin/pip freeze --exclude-editable > ${VENV}/requirements.lock.txt \
|
|
&& head -50 ${VENV}/requirements.lock.txt
|
|
|
|
# 6) Strip pip cache & __pycache__ to shrink the layer copied to runtime.
|
|
#
|
|
# Note on the `-name tests` strip: numpy 2.4 ships `numpy/_core/tests/`
|
|
# back into the wheel (numpy 2.2.6 had stripped it, which is what the
|
|
# explicit upgrade earlier in this Dockerfile was meant to fix). Blowing
|
|
# away every `tests/` directory under the venv would re-introduce the
|
|
# same `from numpy._core.tests._natype import pd_NA` ImportError on the
|
|
# deployed image. Exclude numpy's tests directories explicitly so the
|
|
# upgrade fix stays in effect; keep stripping the rest.
|
|
RUN find ${VENV} -depth -type d -name __pycache__ -exec rm -rf {} + \
|
|
&& find ${VENV} -depth -type d -name tests \
|
|
! -path "*numpy/_core/tests*" \
|
|
! -path "*numpy/tests*" \
|
|
! -path "*numpy/ma/tests*" \
|
|
-exec rm -rf {} + \
|
|
&& rm -rf /root/.cache/pip /root/.cache/uv
|
|
|
|
# Build-time verification.
|
|
#
|
|
# (1) arch-list check uses the RAW C++ accessor (not torch.cuda.get_arch_list()).
|
|
# The Python wrapper checks torch.cuda.is_available() first and returns []
|
|
# when no GPU is visible -- which is always the case here because
|
|
# CUDA_VISIBLE_DEVICES is empty by design.
|
|
#
|
|
# (2) We verify required packages via package metadata only -- we do NOT import
|
|
# unsloth or unsloth_zoo here. Their __init__ calls torch.cuda.get_device_
|
|
# properties(0) which requires an actual CUDA device (UNSLOTH_ALLOW_CPU=1
|
|
# only bypasses the first gate, not the deeper init). Import-time
|
|
# correctness is exercised at deploy time by smoke_test.py with --gpus all.
|
|
RUN TARGETARCH="${TARGETARCH:-amd64}" ${VENV}/bin/python - <<'PY'
|
|
import os, platform
|
|
target = os.environ.get("TARGETARCH", "amd64")
|
|
mach = platform.machine()
|
|
print(f"build target: TARGETARCH={target} platform.machine()={mach}")
|
|
|
|
import torch
|
|
arches = torch._C._cuda_getArchFlags().split()
|
|
print("torch", torch.__version__, "cuda", torch.version.cuda)
|
|
print("arches:", arches)
|
|
assert torch.__version__.startswith("2.10.0"), f"torch silently moved: {torch.__version__}"
|
|
assert "+cu128" in torch.__version__, f"cu build silently changed: {torch.__version__}"
|
|
assert "sm_100" in arches, f"sm_100 (B200/GB200) missing: {arches}"
|
|
# cu128 wheels ship sm_120 native SASS on BOTH amd64 (RTX 5090 / RTX PRO 6000
|
|
# Blackwell) and aarch64 (Grace systems). On arm64 sm_120 is what DGX Spark
|
|
# (sm_121) runs via minor-forward-compat within major 12; sm_121 itself is
|
|
# never in any cu128 wheel.
|
|
assert "sm_120" in arches, f"sm_120 missing: {arches}"
|
|
print(f"OK: torch 2.10.0+cu128 with sm_100 + sm_120 native SASS intact ({target})")
|
|
|
|
from importlib.metadata import version, PackageNotFoundError
|
|
# xformers has no cu128 aarch64 wheel as of 0.0.34, so we only require it
|
|
# on amd64. Everything else is platform-agnostic.
|
|
REQUIRED = ["torch", "triton", "bitsandbytes", "unsloth",
|
|
"unsloth_zoo", "transformers", "trl", "peft", "accelerate"]
|
|
if target == "amd64":
|
|
REQUIRED.insert(2, "xformers")
|
|
missing = []
|
|
for pkg in REQUIRED:
|
|
try:
|
|
v = version(pkg.replace("_", "-"))
|
|
print(f" {pkg:14s} {v}")
|
|
except PackageNotFoundError:
|
|
missing.append(pkg)
|
|
if missing:
|
|
raise SystemExit(f"FAIL: missing wheels: {missing}")
|
|
print("OK: all required wheels present")
|
|
|
|
# Lightweight imports: these init without touching CUDA, unlike unsloth.
|
|
import importlib
|
|
LIGHT_IMPORTS = ["bitsandbytes", "triton"]
|
|
if target == "amd64":
|
|
LIGHT_IMPORTS.insert(0, "xformers")
|
|
for pkg in LIGHT_IMPORTS:
|
|
importlib.import_module(pkg)
|
|
print(f"OK: {' + '.join(LIGHT_IMPORTS)} import cleanly on no-GPU host")
|
|
PY
|
|
|
|
# =============================================================================
|
|
# Stage 2: runtime -- slim runtime image, no nvcc, no cuDNN/cuBLAS layers
|
|
# =============================================================================
|
|
# The "-base-" variant (vs "-cudnn-runtime-") drops ~2.7 GB of system CUDA
|
|
# libraries we never load. torch wheels bake their OWN cuDNN/cuBLAS/cuSPARSE/
|
|
# cuRAND/cuSOLVER/cuFFT/NCCL/cuSparseLt inside `torch/lib/`, and libtorch_cuda.so
|
|
# has RPATH `$ORIGIN/../../nvidia/cudnn/lib:$ORIGIN/../../nvidia/cublas/lib:...`
|
|
# so the dynamic loader resolves through the wheel, never the system. Empirical
|
|
# verification via `readelf -d torch/lib/libtorch_cuda.so` (Fork 5 audit). The
|
|
# base image still provides nvidia-smi, libcuda stubs, libnvidia-ml -- which
|
|
# is everything our entrypoint pre-flight + torch.cuda need.
|
|
FROM nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION} AS runtime
|
|
|
|
# The nvidia/cuda:12.8.1-base-ubuntu24.04 manifest is multi-arch
|
|
# (linux/amd64 + linux/arm64). docker/buildx picks the right one for
|
|
# TARGETPLATFORM at this FROM line; no conditional needed.
|
|
ARG TARGETARCH
|
|
ARG PYTHON_VERSION
|
|
ENV DEBIAN_FRONTEND=noninteractive \
|
|
PIP_NO_CACHE_DIR=1 \
|
|
PYTHONDONTWRITEBYTECODE=1 \
|
|
PYTHONUNBUFFERED=1 \
|
|
PATH=/opt/unsloth-venv/bin:${PATH} \
|
|
HF_HOME=/workspace/.cache/huggingface \
|
|
TRITON_CACHE_DIR=/workspace/.cache/triton \
|
|
# Keep the arch list visible at runtime in case the user source-builds anything
|
|
# extra inside the container (e.g. a custom CUDA op). Same list as the builder
|
|
# stage so a `pip install some-cuda-ext` inside the container gets a SASS
|
|
# blob that covers every supported arch.
|
|
TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;10.3;12.0+PTX"
|
|
|
|
# zstd: the official Ollama notebooks run `curl ollama.com/install.sh | sh`
|
|
# inside the container, and that installer extracts a zstd tarball -- without
|
|
# it the Llama3 Ollama-export notebook dies at the install cell.
|
|
# ffmpeg: torchcodec (datasets' audio decode path, pip-installed by the
|
|
# TTS/STT notebooks) dlopens the system FFmpeg libraries; the wheel does not
|
|
# bundle them, so without ffmpeg every audio notebook dies at load_dataset.
|
|
# wget: notebooks fetch sample assets with `!wget URL`; without it the cell
|
|
# "succeeds" with sh's not-found on stderr and the next cell crashes on the
|
|
# missing file (the Whisper notebook died exactly this way).
|
|
# ninja-build: flashinfer's cpp_ext JIT shells out to ninja; subprocesses
|
|
# (e.g. `vllm serve` launched by unsloth.dataprep) do not always inherit the
|
|
# venv bin on PATH, so the pip ninja alone is not reachable there.
|
|
RUN apt-get update && apt-get install -y --no-install-recommends \
|
|
software-properties-common ca-certificates curl wget git libgomp1 \
|
|
gcc g++ zstd ffmpeg ninja-build \
|
|
&& add-apt-repository -y ppa:deadsnakes/ppa \
|
|
&& apt-get update && apt-get install -y --no-install-recommends \
|
|
python${PYTHON_VERSION} python${PYTHON_VERSION}-venv python${PYTHON_VERSION}-dev \
|
|
&& ln -sf /usr/bin/python${PYTHON_VERSION} /usr/local/bin/python \
|
|
&& ln -sf /usr/bin/python${PYTHON_VERSION} /usr/local/bin/python3 \
|
|
&& rm -rf /var/lib/apt/lists/*
|
|
# Why gcc + g++ + python3.12-dev in the RUNTIME stage:
|
|
# Triton's nvidia backend lazily compiles a small C extension (CudaUtils) on
|
|
# first GPU access. Without a C compiler + Python headers the very first
|
|
# forward pass of any Unsloth model dies with:
|
|
# RuntimeError: Failed to find C compiler. Please specify via CC env var.
|
|
# Adds ~250MB to the runtime image, which is the cost of letting every kernel
|
|
# JIT correctly. (Pre-compiling CudaUtils at build time would need a GPU, so
|
|
# shipping the toolchain is the right trade-off.)
|
|
|
|
COPY --from=builder /opt/unsloth-venv /opt/unsloth-venv
|
|
|
|
# DGX Spark / GB10 (sm_121) fix, arm64 ONLY.
|
|
#
|
|
# Two cu13 components need to override what the cu128 stack ships, because
|
|
# nothing in CUDA 12.8 -- toolkit or wheel -- knows about sm_121:
|
|
#
|
|
# (1) torch's bundled libnvrtc.so.12 (from CUDA 12.8) does not accept
|
|
# sm_121 as --gpu-architecture. The jiterator C++ side queries the
|
|
# device cap directly, so any path that JIT-compiles a kernel (e.g.
|
|
# torch.fft.rfft(complex).abs(), used inside mel-spectrogram code)
|
|
# errors out. Fix: symlink libnvrtc.so.13 over the bundled .so.12.
|
|
#
|
|
# (2) Triton's nvidia backend invokes ptxas. Triton wheels older than
|
|
# 3.6.0 bundled cu12.8 ptxas which tops out at sm_120 and refuses
|
|
# sm_121, silently downgrading to sm_80 per triton-lang/triton#8335.
|
|
# Triton 3.6.0 (which we pin above) bundles cu13 ptxas, but for
|
|
# defense in depth we ALSO install cuda-nvcc-13-0 and point Triton
|
|
# at it via TRITON_PTXAS_PATH.
|
|
#
|
|
# NVRTC and ptxas are CPU-side compilers; they do NOT call into libcuda,
|
|
# so we can install cu13 alongside the cu128 runtime without any driver
|
|
# requirement bump (toolkit driver floor stays 570+).
|
|
#
|
|
# amd64 image is untouched: no sm_121 hardware exists on amd64, and the
|
|
# extra ~400 MB would be dead weight.
|
|
RUN if [ "${TARGETARCH:-amd64}" = "arm64" ]; then \
|
|
set -eux; \
|
|
# The nvidia/cuda base already configures the CUDA apt repo
|
|
# (sbsa for arm64) with its own Signed-By keyring at
|
|
# /usr/share/keyrings/cuda-archive-keyring.gpg. Installing
|
|
# cuda-keyring_1.1-1_all.deb on top adds a second sources file
|
|
# with a different Signed-By, which makes `apt-get update` refuse
|
|
# the entire repo ("Conflicting values set for option Signed-By").
|
|
# The base's repo URL is monolithic and serves every CUDA version
|
|
# including 13.x, so we install cu13 packages directly without
|
|
# touching the keyring at all. Empirically verified on the
|
|
# ubuntu-24.04-arm GitHub Actions runner.
|
|
apt-get update; \
|
|
apt-get install -y --no-install-recommends \
|
|
cuda-nvrtc-13-0 \
|
|
cuda-nvcc-13-0; \
|
|
rm -rf /var/lib/apt/lists/*; \
|
|
# (1) NVRTC swap. torch's wheel-bundled cu128 NVRTC -> cu13 NVRTC.
|
|
NVRTC_DIR=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages/nvidia/cuda_nvrtc/lib; \
|
|
if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ]; then \
|
|
mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \
|
|
ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12"; \
|
|
fi; \
|
|
fi
|
|
|
|
# Register the venv's torch + NVIDIA lib dirs with the loader so torchcodec
|
|
# (installed in the builder, see the bake comment there) can dlopen them.
|
|
# ld.so.conf.d, NOT LD_LIBRARY_PATH: the cache is consulted only after
|
|
# DT_RUNPATH, so the llama.cpp bundle keeps resolving its own $ORIGIN
|
|
# libraries first. The import check runs here because ffmpeg lives in this
|
|
# stage; fail-soft where the torchcodec wheel was unavailable.
|
|
RUN set -eux \
|
|
&& SP=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages \
|
|
&& printf "%s\n" "$SP/torch/lib" "$SP/nvidia/cuda_nvrtc/lib" \
|
|
"$SP/nvidia/cuda_runtime/lib" "$SP/nvidia/npp/lib" \
|
|
> /etc/ld.so.conf.d/zz-unsloth-venv.conf \
|
|
&& ldconfig \
|
|
&& { /opt/unsloth-venv/bin/python -c \
|
|
"import torchcodec; print('torchcodec', torchcodec.__version__)" \
|
|
|| echo ">> torchcodec unavailable on this arch (audio decode falls back)"; }
|
|
|
|
# Prebuilt llama.cpp so GGUF export works out of the box.
|
|
#
|
|
# unsloth_zoo's save_pretrained_gguf() calls check_llama_cpp(), which looks
|
|
# for llama-quantize + convert_hf_to_gguf.py + gguf-py/ in
|
|
# $UNSLOTH_LLAMA_CPP_PATH (default ~/.unsloth/llama.cpp). Without a baked
|
|
# install the first GGUF export inside the container would hit
|
|
# install_llama_cpp()'s interactive prompt and then a slow source build.
|
|
#
|
|
# Why NOT studio/install_llama_prebuilt.py here: that resolver selects a
|
|
# bundle for the CURRENT host (nvidia-smi, /proc/driver/nvidia -- which
|
|
# leaks through from a GPU build host even inside docker build -- and the
|
|
# installed CUDA runtime). The build must never introspect the host, so we
|
|
# pin release + asset by build target instead (see fetch_llama_prebuilt.py):
|
|
# * amd64 -> app-<tag>-linux-x64-cuda12-portable.tar.gz (sm_70..sm_120)
|
|
# arm64 -> app-<tag>-linux-arm64-cuda13-portable.tar.gz (sm_90..sm_121,
|
|
# DGX Spark / Grace)
|
|
# * portable bundles carry their own CUDA runtime libs and dlopen the
|
|
# CUDA backend, so the same binaries also run CPU-only
|
|
# * sha256-verified against the release's llama-prebuilt-sha256.json
|
|
# * converter + gguf-py hydrated from the SAME release's source tarball
|
|
# so the python-side tensor mappings match the binaries
|
|
# /opt (not /root) so the install survives a `docker run --user` override;
|
|
# UNSLOTH_LLAMA_CPP_PATH makes zoo find it regardless of $HOME.
|
|
ARG LLAMA_PREBUILT_TAG=b9596-mix-e6f2453
|
|
COPY fetch_llama_prebuilt.py /tmp/fetch_llama_prebuilt.py
|
|
RUN /opt/unsloth-venv/bin/python /tmp/fetch_llama_prebuilt.py \
|
|
"${LLAMA_PREBUILT_TAG}" "${TARGETARCH:-amd64}" /opt/unsloth/llama.cpp \
|
|
&& rm -f /tmp/fetch_llama_prebuilt.py \
|
|
&& cat /opt/unsloth/llama.cpp/UNSLOTH_PREBUILT_INFO.json
|
|
ENV UNSLOTH_LLAMA_CPP_PATH=/opt/unsloth/llama.cpp
|
|
|
|
WORKDIR /workspace
|
|
RUN mkdir -p ${HF_HOME} ${TRITON_CACHE_DIR}
|
|
|
|
# JupyterLab lives in the venv (see builder stage). Persistent notebooks
|
|
# should be bind-mounted onto /workspace.
|
|
EXPOSE 8888
|
|
|
|
COPY smoke_test.py /workspace/smoke_test.py
|
|
COPY entrypoint.sh /usr/local/bin/unsloth-entrypoint
|
|
RUN chmod +x /usr/local/bin/unsloth-entrypoint
|
|
|
|
# Entrypoint runs three fast pre-flight checks before user code:
|
|
# 1. nvidia-smi sees at least one GPU (catches missing --gpus all)
|
|
# 2. torch.cuda.is_available() is True (catches host driver too old)
|
|
# 3. compute capability >= sm_80 (catches pre-Ampere GPUs)
|
|
# Each check fails with an actionable error pointing to the fix.
|
|
# Bypass for offline tooling: docker run -e UNSLOTH_SKIP_GPU_CHECK=1 ...
|
|
ENTRYPOINT ["/usr/local/bin/unsloth-entrypoint"]
|
|
|
|
# Default command: interactive python REPL.
|
|
# Override examples:
|
|
# docker run --gpus all unsloth/unsloth:latest python /workspace/smoke_test.py
|
|
# docker run --gpus all -it unsloth/unsloth:latest bash
|
|
CMD ["python"]
|