# 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. # These are declared by notebook install cells that the in-image runner # neutralises (deps are meant to be prebaked), so bake them here. All # pure-Python or self-contained wheels; none names torch, so the cu128 pin # is undisturbed: # soundfile TTS notebooks read/write audio (bundles libsndfile in its wheel) # evaluate + jiwer Whisper notebook's WER metric (evaluate.load("wer") -> jiwer) # tensorboard default TrainingArguments report_to backend # langid DeepSeek-R1 GRPO reward's language-id check # easydict some vision trust_remote_code modeling files # protobuf slow->fast tokenizer conversion for sentencepiece models # omegaconf TTS families + both NeMo-Gym RL notebooks' config objects # einx TTS codec tensor-rearrange (Llasa / Oute / Spark TTS) # librosa Whisper audio feature extraction (pairs with soundfile + torchcodec) # ftfy Oute TTS text normalisation # decord (ERNIE-VL video decode) is installed separately below: it ships no # aarch64 wheel, so a hard install here would break the arm64 build. # librosa pulls numba/soxr/audioread; numba is already pinned >=0.65 (numpy 2.4 # compatible) by the vLLM pass, so the resolve must NOT move torch/numpy/numba -- # the assertion below fails the build loudly if it did. # Pinned (==) to the resolved, tested versions for reproducible rebuilds -- the # same convention as the cu128 core (torch/torchvision/torchaudio). Bump these # deliberately, not silently on the next build. Transitive deps of these are # captured by the full venv lockfile (docker/freeze.sh -> requirements.lock.txt). RUN ${VENV}/bin/uv pip install \ --python ${VENV}/bin/python \ "jupyterlab==4.6.0" "notebook==7.6.0" "ipywidgets==8.1.8" "matplotlib==3.11.0" \ "soundfile==0.14.0" "evaluate==0.4.6" "jiwer==4.0.0" "tensorboard==2.20.0" \ "langid==1.1.6" "easydict==1.13" "protobuf==6.33.6" \ "omegaconf==2.3.1" "einx==0.4.3" "librosa==0.11.0" "ftfy==6.3.1" \ && ${VENV}/bin/python -c "import torch, numpy, numba; from packaging.version import Version; assert torch.__version__.startswith('2.10.0'), torch.__version__; assert Version(numpy.__version__) >= Version('2.3'), numpy.__version__; assert Version(numba.__version__) >= Version('0.65'), numba.__version__; print('notebook-deps pins OK:', torch.__version__, numpy.__version__, numba.__version__)" # decord (ERNIE-VL video decode) publishes wheels only for x86_64 / win_amd64, # so install it on its own and fail-soft: amd64 gets it; on arm64 the ERNIE-VL # video path is skipped rather than breaking the whole image build. RUN ${VENV}/bin/uv pip install --python ${VENV}/bin/python "decord==0.6.0" \ || echo ">> decord skipped (no matching wheel for ${TARGETARCH:-amd64}); ERNIE-VL video decode unavailable" # 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})" # Coherent transformers SIDECARS for per-notebook version activation (see # docker/unsloth_nb_compat.py). unslothai/notebooks pin many transformers # versions in their install cells; the base venv ships one (newest 5.x). Each # sidecar is `transformers==X` + its matched huggingface_hub/tokenizers/ # safetensors, installed --no-deps into its own --target dir under # ${VENV}/tf-sidecars (rides along in the COPY to runtime). Activating one # (prepend to sys.path before any ML import) swaps transformers for a model # WITHOUT touching the cu128 torch/vLLM/unsloth base stack -- verified: base # unsloth loads + generates under both a 4.57.6 and a 5.5.0 sidecar on B200. # Versions mirror Unsloth Studio's tiers (4.57.6 default + 5.3.0/5.5.0/5.10.2). # The companion versions are RESOLVED at build time (not hardcoded) so they # always satisfy each transformers' hard requirements. ~300MB total after the # __pycache__/tests strip in the cleanup RUN below. Fail-soft per arch/wheel. RUN set -eux \ && for TFV in 4.57.6 5.3.0 5.5.0 5.10.2; do \ SCRATCH="$(mktemp -d)"; \ if ! ${VENV}/bin/uv pip install --python ${VENV}/bin/python \ --target "$SCRATCH" "transformers==${TFV}" >/dev/null 2>&1; then \ echo ">> sidecar resolve failed for ${TFV}; skipping"; rm -rf "$SCRATCH"; continue; \ fi; \ pin() { ls -d "$SCRATCH/$1"-*.dist-info 2>/dev/null \ | sed -E "s@.*/$1-([0-9][0-9A-Za-z.]*)\.dist-info@\1@" | head -1; }; \ HFV="$(pin huggingface_hub)"; TKV="$(pin tokenizers)"; SFV="$(pin safetensors)"; \ rm -rf "$SCRATCH"; \ DEST="${VENV}/tf-sidecars/t_$(echo "${TFV}" | tr . _)"; \ ${VENV}/bin/uv pip install --python ${VENV}/bin/python --target "$DEST" --no-deps \ "transformers==${TFV}" \ ${HFV:+"huggingface_hub==${HFV}"} \ ${TKV:+"tokenizers==${TKV}"} \ ${SFV:+"safetensors==${SFV}"}; \ echo ">> sidecar transformers==${TFV} (hf_hub=${HFV} tokenizers=${TKV} safetensors=${SFV})"; \ done \ && { du -sh ${VENV}/tf-sidecars || true; } # 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 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. # Also (size reductions, all verified runtime-safe): # * npp: torchcodec dlopens only libnppicc + libnppc; the other ~10 npp libs # (libnppif 163M, libnppist, libnppig, ...) are dead weight (~388MB). Nothing # else in the venv links them. # * static .a archives (~143MB): xgrammar/triton-cupti/nvperf/nvshmem ship .a # alongside the .so they actually load at runtime; .a are link-time only and # nothing in the image links venv archives (nvcc JIT links /usr/local/cuda). # * nvshmem device-side bitcode (~30MB): host .so kept; .bc is device-relink only. # We deliberately do NOT strip headers (torch/include etc.): the leave-to-pip # kernels causal-conv1d / mamba-ssm build against torch headers at notebook time # with --no-build-isolation, so torch/include must survive. RUN set -eux \ && 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 \ && SP=${VENV}/lib/python${PYTHON_VERSION}/site-packages \ && if [ -d "$SP/nvidia/npp/lib" ]; then \ find "$SP/nvidia/npp/lib" -maxdepth 1 -name 'libnpp*.so.*' \ ! -name 'libnppicc.so.*' ! -name 'libnppc.so.*' -delete; \ fi \ && find ${VENV} -name '*.a' -delete \ && rm -f "$SP"/nvidia/nvshmem/lib/libnvshmem_device.bc \ && echo "venv size after prune:" && du -sh ${VENV} # 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 ARG CUDA_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. # cuda-nvcc + cudart-dev: flash-linear-attention's TileLang backend (Qwen3.5 # gated-delta-rule models in Studio) JIT-compiles CUDA kernels at runtime via # nvcc; the -base image ships only runtime libs, so without the compiler any # TileLang-allowlisted model dies with "[Errno 2] No such file or directory: # '/usr/local/cuda/bin/nvcc'" on its first backward pass. RUN CUDA_PKG="$(echo "${CUDA_VERSION}" | awk -F. '{print $1"-"$2}')" \ && 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 \ "cuda-nvcc-${CUDA_PKG}" "cuda-cudart-dev-${CUDA_PKG}" \ && 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 \ && test -x /usr/local/cuda/bin/nvcc \ && 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--linux-x64-cuda12-portable.tar.gz (sm_70..sm_120) # arm64 -> app--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. # Default "latest" -> fetch_llama_prebuilt.py resolves the newest # unslothai/llama.cpp release at build time (follows the /releases/latest # redirect, no API token). build.sh resolves this to a concrete tag before # invoking docker build so the layer cache busts only when upstream publishes; # pass --build-arg LLAMA_PREBUILT_TAG= for a frozen, reproducible build. ARG LLAMA_PREBUILT_TAG=latest 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} # --------------------------------------------------------------------------- # Per-notebook transformers version activation -- run unslothai/notebooks # UNCHANGED. See docker/unsloth_nb_compat.py for the full rationale. Pieces: # * unsloth_nb_compat.py -> site-packages (importable everywhere): tier # detection + sidecar resolution + activation + the IPython hook. # * pip/uv shim on a PATH dir AHEAD of the venv bin: a notebook's # `!pip install ...` / `!uv pip install ...` cell becomes SAFE + idempotent # (keeps the baked torch/vLLM stack; records the requested transformers so # its sidecar is activated for the model cells). # * IPython startup hook: activates the right sidecar before the first model # cell in manual JupyterLab. # * unsloth-run: headless `unsloth-run ` that auto-picks the # sidecar and executes every cell -- the robust driven path. # --------------------------------------------------------------------------- COPY unsloth_nb_compat.py unsloth_pip_shim.py unsloth_ipython_startup.py unsloth_run.py unsloth_sync_notebooks.sh unsloth_nb_content_sig.py unsloth_nb_view.py unsloth_nb_strip_colab.py unsloth_colab_compat.py /opt/unsloth-nb/ RUN set -eux \ && SP=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages \ && cp /opt/unsloth-nb/unsloth_nb_compat.py "$SP/unsloth_nb_compat.py" \ && cp /opt/unsloth-nb/unsloth_colab_compat.py "$SP/unsloth_colab_compat.py" \ && chmod +x /opt/unsloth-nb/unsloth_pip_shim.py /opt/unsloth-nb/unsloth_run.py /opt/unsloth-nb/unsloth_sync_notebooks.sh /opt/unsloth-nb/unsloth_nb_content_sig.py /opt/unsloth-nb/unsloth_nb_view.py /opt/unsloth-nb/unsloth_nb_strip_colab.py \ && mkdir -p /opt/unsloth-nb/bin \ && for t in pip pip3 uv; do ln -sf /opt/unsloth-nb/unsloth_pip_shim.py /opt/unsloth-nb/bin/$t; done \ && ln -sf /opt/unsloth-nb/unsloth_run.py /usr/local/bin/unsloth-run \ && ln -sf /opt/unsloth-nb/unsloth_sync_notebooks.sh /usr/local/bin/unsloth-sync-notebooks \ && ln -sf /opt/unsloth-nb/unsloth_nb_content_sig.py /usr/local/bin/unsloth-nb-content-sig \ && ln -sf /opt/unsloth-nb/unsloth_nb_view.py /usr/local/bin/unsloth-nb-view \ && ln -sf /opt/unsloth-nb/unsloth_nb_strip_colab.py /usr/local/bin/unsloth-nb-strip-colab \ && mkdir -p /root/.ipython/profile_default/startup \ && cp /opt/unsloth-nb/unsloth_ipython_startup.py /root/.ipython/profile_default/startup/00-unsloth-nb.py \ && /opt/unsloth-venv/bin/python -c "import sys, glob; sys.path.insert(0, '$SP'); import unsloth_nb_compat, unsloth_colab_compat; print('nb-compat OK; baked sidecars:', sorted(glob.glob('/opt/unsloth-venv/tf-sidecars/t_*')))" # Shim dir AHEAD of the venv bin so `!pip`/`!uv` resolve to the shim, not the real tool. ENV PATH=/opt/unsloth-nb/bin:${PATH} # Pre-clone unslothai/notebooks so JupyterLab opens with the notebooks already # present (no git clone or wget needed). Baked here as a READ-ONLY template # (~206MB, .git stripped); on boot the entrypoint copies it to # /workspace/unsloth-notebooks and best-effort refreshes from GitHub when # upstream has advanced. It never overwrites a notebook the user has touched, # and for untouched notebooks it skips the rewrite when only the install header # / announcements / footer moved upstream (the tutorial body is unchanged) -- # see unsloth_sync_notebooks.sh + unsloth_nb_content_sig.py. Inherited as-is by # the studio image (FROM base). RUN set -eux \ && git clone --depth 1 https://github.com/unslothai/notebooks /opt/unsloth-notebooks \ && git -C /opt/unsloth-notebooks rev-parse HEAD > /opt/unsloth-notebooks/.unsloth_template_commit \ && rm -rf /opt/unsloth-notebooks/.git \ && du -sh /opt/unsloth-notebooks # JupyterLab lives in the venv (see builder stage). The unslothai/notebooks # collection is pre-populated into /workspace/unsloth-notebooks on boot; mount a # volume on /workspace to persist your own notebooks and outputs across runs. 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"]