From c6d92160f6d20d01b3774db492ecda1383cae5c1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 24 May 2026 06:52:58 +0000 Subject: [PATCH] Add Docker build for Blackwell that runs on any NVIDIA GPU host Adds a multi-stage Dockerfile producing an image that works on Ampere through Blackwell (sm_80 through sm_120: A100, RTX 30/40, H100, B100/B200, RTX 50-series, RTX 6000 Pro Blackwell). The build itself requires no GPU at all and runs on a free GitHub-hosted ubuntu-latest runner. How the GPU-less build works: 1. cu128 PyTorch wheels are fat binaries. torch._C._cuda_getArchFlags() returns 'sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120' regardless of which GPU compiled the image, because the wheels are cross-compiled upstream by the PyTorch team. 2. All deps resolve in a single uv pip install pass with explicit pins (torch==2.10.0, --extra-index-url cu128, no --torch-backend=auto, no install.sh). This prevents the silent cu cascade where bitsandbytes' transitive cuda-toolkit==13 dep upgrades torch to 2.12+cu130 in a later resolver pass, leaving xformers and other cu128 wheels stranded. 3. Build-time verification uses package metadata (importlib.metadata.version) and the raw torch._C._cuda_getArchFlags() accessor. We deliberately avoid import unsloth at build time because unsloth.__init__ calls torch.cuda.get_device_properties(0), which requires an actual CUDA device and is not bypassable. Import-time correctness is exercised at deploy time by smoke_test.py with --gpus all. 4. UNSLOTH_COMPILE_DISABLE=1 and CUDA_VISIBLE_DEVICES="" during the build stage prevent any code path from JIT-compiling kernels for the build host's compute capability and baking the resulting cache into the image. The deploy GPU produces its own cache on first use. Other notes: - --index-strategy unsafe-best-match is needed because the PyTorch wheel index serves an old requests==2.28.1 that conflicts with datasets>=2.32.2, which the default first-index-wins strategy rejects. - Extra is cu128-ampere-torch2100 (ampere precedes the torch version in the pyproject ordering). - No flash-attn in the base image. FA3 is hard-refused on Blackwell upstream and unsloth gracefully falls back to xformers + SDPA. Users on Ampere / Ada / Hopper who want FA2 can pip install flash-attn on top. - Two stages: nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 for the build, -cudnn-runtime for the deploy image. No nvcc in the published image. - A lockfile is emitted at /opt/unsloth-venv/requirements.lock.txt inside the image and can be extracted with docker/freeze.sh for byte-identical rebuilds even after PyPI moves on. CI workflow .github/workflows/docker-publish.yml: - Builds on ubuntu-latest on every push to main, every tag, weekly via cron, and manually via workflow_dispatch. Pushes to docker.io/unsloth/unsloth with cache via type=gha. - Optional smoke-test job runs on a self-hosted GPU runner if vars.HAS_GPU_RUNNER is set; skipped otherwise. End-to-end verification on sm_120 hardware is a nice-to-have, not a publish blocker. Validation: - Install path validated on a B200 host with CUDA_VISIBLE_DEVICES="" set (simulating the GPU-less CI runner): torch 2.10.0+cu128 holds, xformers 0.0.34, bitsandbytes 0.49.2, triton 3.6.0, transformers 5.5.0, trl 0.24.0, peft 0.19.1, accelerate 1.13.0. Arch flags include sm_100 and sm_120. - Runtime path validated end-to-end on B200: smoke_test.py imports unsloth, loads Llama-3.2-1B-Instruct-bnb-4bit in 4-bit, completes 5 LoRA steps with loss decreasing 4.11 -> 3.75. xformers fallback active as designed. Files: - docker/Dockerfile multi-stage cu128 build - docker/build.sh local build wrapper - docker/freeze.sh extract lockfile from a built image - docker/smoke_test.py runtime verification, run with --gpus all - docker/.dockerignore - .github/workflows/docker-publish.yml --- .github/workflows/docker-publish.yml | 118 ++++++++++++++++ docker/.dockerignore | 3 + docker/Dockerfile | 200 +++++++++++++++++++++++++++ docker/build.sh | 46 ++++++ docker/freeze.sh | 26 ++++ docker/smoke_test.py | 137 ++++++++++++++++++ 6 files changed, 530 insertions(+) create mode 100644 .github/workflows/docker-publish.yml create mode 100644 docker/.dockerignore create mode 100644 docker/Dockerfile create mode 100755 docker/build.sh create mode 100755 docker/freeze.sh create mode 100644 docker/smoke_test.py diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml new file mode 100644 index 0000000000..5c0c0786b8 --- /dev/null +++ b/.github/workflows/docker-publish.yml @@ -0,0 +1,118 @@ +# Builds and publishes the Blackwell-compatible Unsloth Docker image. +# +# The build runs on a free GitHub-hosted Ubuntu runner with NO GPU attached. +# This is possible because: +# 1. cu128 PyTorch wheels are fat binaries -- they already ship sm_70 through +# sm_120 SASS, cross-compiled upstream by the PyTorch team. +# 2. The Dockerfile pins explicit wheel URLs (no --torch-backend=auto, no +# install.sh that introspects the host driver). +# 3. The build-time sanity check uses torch._C._cuda_getArchFlags(), which +# reads compiled wheel metadata and does NOT require a CUDA device. +# 4. UNSLOTH_COMPILE_DISABLE=1 prevents Unsloth from JIT-compiling a Triton +# kernel cache keyed to the (non-existent) build-host GPU. +# +# Required repository secrets: +# DOCKERHUB_USERNAME, DOCKERHUB_TOKEN +# +# Optional repository variable (gates the smoke-test job): +# HAS_GPU_RUNNER = 'true' if a self-hosted GPU runner is available + +name: Publish Blackwell Docker image + +on: + push: + branches: [main] + tags: ['v*'] + schedule: + - cron: '17 4 * * 1' # weekly Mon 04:17 UTC (off-the-hour on purpose) + workflow_dispatch: + inputs: + unsloth_ref: + description: 'unsloth git ref to bake in' + required: false + default: 'main' + unsloth_zoo_ref: + description: 'unsloth-zoo git ref to bake in' + required: false + default: 'main' + +env: + REGISTRY: docker.io + IMAGE_NAME: unsloth/unsloth + +jobs: + build: + runs-on: ubuntu-latest # no GPU, 16GB RAM, 4 vCPU + timeout-minutes: 60 + permissions: + contents: read + packages: write + + steps: + - uses: actions/checkout@v4 + + # Free up ~20GB on the runner so cu128 wheels + cudnn fit. + - name: Reclaim disk + run: | + sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc \ + /opt/hostedtoolcache/CodeQL "$AGENT_TOOLSDIRECTORY" + df -h / + + - uses: docker/setup-qemu-action@v3 + - uses: docker/setup-buildx-action@v3 + + - name: Log in to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Resolve tags + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} + tags: | + type=raw,value=latest,enable={{is_default_branch}} + type=ref,event=tag + type=schedule,pattern=nightly + type=sha,prefix=sha-,format=short + + - name: Build and push + uses: docker/build-push-action@v6 + with: + context: ./docker + file: ./docker/Dockerfile + platforms: linux/amd64 + push: true + tags: ${{ steps.meta.outputs.tags }} + labels: ${{ steps.meta.outputs.labels }} + cache-from: type=gha + cache-to: type=gha,mode=max + build-args: | + CUDA_VERSION=12.8.1 + UBUNTU_VERSION=24.04 + PYTHON_VERSION=3.12 + UNSLOTH_REF=${{ github.event.inputs.unsloth_ref || 'main' }} + UNSLOTH_ZOO_REF=${{ github.event.inputs.unsloth_zoo_ref || 'main' }} + + - name: Image digest + run: echo "${{ steps.meta.outputs.tags }} -> ${{ steps.meta.outputs.digest }}" + + # Optional: pull the freshly published image onto a self-hosted GPU runner + # and run smoke_test.py. Keeps "did the image actually work" decoupled from + # "was a GPU available at build time". Skipped automatically when no GPU + # runner is registered. + smoke-test: + needs: build + if: ${{ vars.HAS_GPU_RUNNER == 'true' }} + runs-on: [self-hosted, gpu] + timeout-minutes: 20 + steps: + - uses: actions/checkout@v4 + - name: Pull and smoke-test + run: | + docker pull ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest + docker run --rm --gpus all \ + ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest \ + python /workspace/smoke_test.py diff --git a/docker/.dockerignore b/docker/.dockerignore new file mode 100644 index 0000000000..eea1fc6e99 --- /dev/null +++ b/docker/.dockerignore @@ -0,0 +1,3 @@ +** +!Dockerfile +!smoke_test.py diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000000..b11e0877f4 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,200 @@ +# syntax=docker/dockerfile:1.7 +# ----------------------------------------------------------------------------- +# Unsloth + unsloth-zoo for Blackwell (sm_100 B200 + sm_120 RTX 50-series / 6000 Pro) +# +# Why this image works: +# * cu128 wheels are fat binaries: SASS for sm_80;86;89;90;100;120. +# * 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="10.0;12.0+PTX" -- the host GPU is irrelevant +# for compilation; nvcc emits whatever the arch list says. +# +# Build host requirements: +# * Docker with buildkit (default since 23.x) +# * nvidia-container-toolkit (only needed for `docker run --gpus all` at test time) +# * 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 + +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: Ampere, Ada, Hopper, B100/B200 (sm_100), RTX 50x / 6000 Pro (sm_120). + # +PTX on the highest arch lets future Blackwell SKUs run via JIT-PTX. + TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0;10.0;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 \ + && curl -fsSL https://bootstrap.pypa.io/get-pip.py | python \ + && python -m pip install -U pip uv \ + && rm -rf /var/lib/apt/lists/* + +# Build into an isolated prefix that we copy into the runtime stage. +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 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 ${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.11.0" \ + "triton>=3.3.1" \ + "bitsandbytes>=0.49.2,!=0.46.0,!=0.48.0" \ + "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo@${UNSLOTH_ZOO_REF}" \ + "unsloth[cu128-ampere-torch2100] @ git+https://github.com/unslothai/unsloth@${UNSLOTH_REF}" + +# 5) Emit a lockfile so the next rebuild can be byte-identical even if PyPI +# has moved on. Bake it into the image at /opt/unsloth-venv/requirements.lock.txt +# so `docker run ... cat /opt/unsloth-venv/requirements.lock.txt > pins.txt` +# gives you the input to a fully-pinned rebuild. +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. +RUN find ${VENV} -depth -type d -name __pycache__ -exec rm -rf {} + \ + && find ${VENV} -depth -type d -name 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 ${VENV}/bin/python - <<'PY' +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) missing: {arches}" +assert "sm_120" in arches, f"sm_120 (RTX 5090) missing: {arches}" +print("OK: torch 2.10.0+cu128 with sm_100 + sm_120 fat binary intact") + +from importlib.metadata import version, PackageNotFoundError +REQUIRED = ("torch", "triton", "xformers", "bitsandbytes", "unsloth", + "unsloth_zoo", "transformers", "trl", "peft", "accelerate") +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 (xformers, bnb, unsloth metadata visible)") + +# Lightweight imports: these init without touching CUDA, unlike unsloth. +import importlib +for pkg in ("xformers", "bitsandbytes", "triton"): + importlib.import_module(pkg) +print("OK: xformers + bitsandbytes + triton import cleanly on no-GPU host") +PY + +# ============================================================================= +# Stage 2: runtime -- slim runtime image, no nvcc, no headers +# ============================================================================= +FROM nvidia/cuda:${CUDA_VERSION}-cudnn-runtime-ubuntu${UBUNTU_VERSION} AS runtime + +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). + TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0;10.0;12.0+PTX" + +RUN apt-get update && apt-get install -y --no-install-recommends \ + software-properties-common ca-certificates curl git libgomp1 \ + && add-apt-repository -y ppa:deadsnakes/ppa \ + && apt-get update && apt-get install -y --no-install-recommends \ + python${PYTHON_VERSION} python${PYTHON_VERSION}-venv \ + && 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/* + +COPY --from=builder /opt/unsloth-venv /opt/unsloth-venv + +WORKDIR /workspace +RUN mkdir -p ${HF_HOME} ${TRITON_CACHE_DIR} + +COPY smoke_test.py /workspace/smoke_test.py + +# Default entry: drop into python; override with `docker run ... bash` for a shell. +CMD ["python"] diff --git a/docker/build.sh b/docker/build.sh new file mode 100755 index 0000000000..592a52da49 --- /dev/null +++ b/docker/build.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# Build the unsloth-blackwell image on this B200 host (or any Linux host with Docker). +# The build host's GPU is NOT used -- nvcc cross-compiles for sm_100 + sm_120. +# +# Usage: +# ./build.sh # builds unsloth-blackwell:latest pinned to unsloth main +# TAG=2026.05.1 ./build.sh # custom tag +# UNSLOTH_REF=v2026.5.6 UNSLOTH_ZOO_REF=v2026.5.4 ./build.sh # pin git refs +set -euo pipefail + +cd "$(dirname "$0")" + +IMAGE_NAME="${IMAGE_NAME:-unsloth-blackwell}" +TAG="${TAG:-latest}" +CUDA_VERSION="${CUDA_VERSION:-12.8.1}" +UBUNTU_VERSION="${UBUNTU_VERSION:-24.04}" +PYTHON_VERSION="${PYTHON_VERSION:-3.12}" +UNSLOTH_REF="${UNSLOTH_REF:-main}" +UNSLOTH_ZOO_REF="${UNSLOTH_ZOO_REF:-main}" + +echo "Building ${IMAGE_NAME}:${TAG}" +echo " CUDA ${CUDA_VERSION} Ubuntu ${UBUNTU_VERSION} Python ${PYTHON_VERSION}" +echo " unsloth @${UNSLOTH_REF}" +echo " unsloth-zoo @${UNSLOTH_ZOO_REF}" +echo " arch list 8.0;8.6;8.9;9.0;10.0;12.0+PTX" +echo + +DOCKER_BUILDKIT=1 docker build \ + --progress=plain \ + --build-arg CUDA_VERSION="${CUDA_VERSION}" \ + --build-arg UBUNTU_VERSION="${UBUNTU_VERSION}" \ + --build-arg PYTHON_VERSION="${PYTHON_VERSION}" \ + --build-arg UNSLOTH_REF="${UNSLOTH_REF}" \ + --build-arg UNSLOTH_ZOO_REF="${UNSLOTH_ZOO_REF}" \ + -t "${IMAGE_NAME}:${TAG}" \ + . + +echo +echo "Built ${IMAGE_NAME}:${TAG}" +echo +echo "Smoke test on this host (B200, sm_100):" +echo " docker run --rm --gpus all ${IMAGE_NAME}:${TAG} python /workspace/smoke_test.py" +echo +echo "Smoke test on an RTX 5090 host (sm_120):" +echo " docker pull ${IMAGE_NAME}:${TAG} # or load .tar" +echo " docker run --rm --gpus all ${IMAGE_NAME}:${TAG} python /workspace/smoke_test.py" diff --git a/docker/freeze.sh b/docker/freeze.sh new file mode 100755 index 0000000000..9089ae287c --- /dev/null +++ b/docker/freeze.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# Pull the lockfile out of a built image so the next rebuild can be byte-identical. +# +# ./freeze.sh # extracts to requirements.lock.txt next to Dockerfile +# ./freeze.sh some-tag-or-digest # custom source +# +# To rebuild against the frozen lockfile later, replace the `pip install` lines +# in the Dockerfile with `pip install -r /tmp/requirements.lock.txt --no-deps` +# (mounted via `docker build --build-context lock=./requirements.lock.txt`). +set -euo pipefail + +cd "$(dirname "$0")" + +SRC="${1:-unsloth-blackwell:latest}" +DEST="${2:-./requirements.lock.txt}" + +CID=$(docker create "${SRC}") +trap 'docker rm -f "${CID}" >/dev/null' EXIT + +docker cp "${CID}:/opt/unsloth-venv/requirements.lock.txt" "${DEST}" +echo "Wrote ${DEST}" +echo +echo "Top of lockfile:" +head -20 "${DEST}" +echo +echo "Lines: $(wc -l < "${DEST}")" diff --git a/docker/smoke_test.py b/docker/smoke_test.py new file mode 100644 index 0000000000..b4ca903746 --- /dev/null +++ b/docker/smoke_test.py @@ -0,0 +1,137 @@ +""" +Smoke test for the unsloth-blackwell image. + +What this checks (in order, fail-fast): + 1. torch sees the GPU and the arch list contains sm_100 + sm_120. + 2. The runtime device's compute capability is supported. + 3. xformers / bitsandbytes / triton import without ImportError. + 4. unsloth imports and exposes FastLanguageModel. + 5. A 5-step LoRA train on a tiny model actually runs forward + backward. + +Run inside the container: + docker run --rm --gpus all unsloth-blackwell:latest python /workspace/smoke_test.py + +Skip step 5 (faster, no model download): + docker run --rm --gpus all unsloth-blackwell:latest python /workspace/smoke_test.py --skip-train +""" +from __future__ import annotations + +import argparse +import sys + + +def banner(title: str) -> None: + print(f"\n=== {title} ===", flush=True) + + +def check_torch() -> tuple[int, int]: + banner("torch + arch list") + import torch + # Use the raw C++ accessor so this works even when CUDA isn't available + # (lets us run a partial smoke test on a no-GPU host). + arches = torch._C._cuda_getArchFlags().split() + print(f"torch {torch.__version__}") + print(f"cuda build {torch.version.cuda}") + print(f"arches {arches}") + assert "sm_100" in arches, f"sm_100 missing: {arches}" + assert "sm_120" in arches, f"sm_120 missing: {arches}" + + assert torch.cuda.is_available(), "CUDA not visible -- did you pass --gpus all?" + cap = torch.cuda.get_device_capability(0) + name = torch.cuda.get_device_name(0) + print(f"device 0 {name} sm_{cap[0]}{cap[1]}") + if cap[0] < 8: + sys.exit(f"FAIL: pre-Ampere GPU {name} is not supported by this image") + return cap + + +def check_imports() -> None: + banner("dep imports") + import triton; print(f"triton {triton.__version__}") + import xformers; print(f"xformers {xformers.__version__}") + import bitsandbytes as bnb; print(f"bnb {bnb.__version__}") + import transformers; print(f"transformers {transformers.__version__}") + import trl; print(f"trl {trl.__version__}") + import peft; print(f"peft {peft.__version__}") + import unsloth_zoo; print(f"unsloth_zoo {unsloth_zoo.__version__}") + + +def check_unsloth_import() -> None: + banner("unsloth import") + # Unsloth must be imported BEFORE transformers in real training scripts, + # but here we already imported transformers above for the version check. + # That's fine for this smoke -- we're not training Unsloth-patched models yet. + import unsloth + from unsloth import FastLanguageModel + print(f"unsloth {unsloth.__version__}") + print(f"FastLanguageModel {FastLanguageModel}") + + +def check_tiny_train(cap: tuple[int, int]) -> None: + banner("tiny LoRA train (5 steps)") + import os + # Unsloth must be imported first. + import unsloth # noqa: F401 + from unsloth import FastLanguageModel + import torch + + # Small, public, no-gate. ~125M params. + model_name = "unsloth/Llama-3.2-1B-Instruct-bnb-4bit" + print(f"loading {model_name}") + model, tokenizer = FastLanguageModel.from_pretrained( + model_name=model_name, + max_seq_length=512, + dtype=None, + load_in_4bit=True, + ) + model = FastLanguageModel.get_peft_model( + model, + r=8, + lora_alpha=16, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], + lora_dropout=0.0, + bias="none", + use_gradient_checkpointing="unsloth", + random_state=0, + ) + + prompts = [ + "Q: What is the capital of France?\nA:", + "Q: 2 + 2 = ?\nA:", + "Q: Name a primary color.\nA:", + "Q: Hello, who are you?\nA:", + ] * 2 + enc = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True, max_length=64) + enc = {k: v.cuda() for k, v in enc.items()} + labels = enc["input_ids"].clone() + + model.train() + optim = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-4) + for step in range(5): + out = model(**enc, labels=labels) + out.loss.backward() + optim.step() + optim.zero_grad(set_to_none=True) + print(f"step {step} loss={out.loss.item():.4f}", flush=True) + + print("OK: 5 LoRA steps completed") + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--skip-train", action="store_true", + help="Skip the tiny LoRA training step (no HF download).") + args = ap.parse_args() + + cap = check_torch() + check_imports() + check_unsloth_import() + if not args.skip_train: + check_tiny_train(cap) + + banner("all checks passed") + return 0 + + +if __name__ == "__main__": + sys.exit(main())