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
This commit is contained in:
Daniel Han
parent
56e9046b2f
commit
c6d92160f6
6 changed files with 530 additions and 0 deletions
118
.github/workflows/docker-publish.yml
vendored
Normal file
118
.github/workflows/docker-publish.yml
vendored
Normal file
|
|
@ -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
|
||||
3
docker/.dockerignore
Normal file
3
docker/.dockerignore
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
**
|
||||
!Dockerfile
|
||||
!smoke_test.py
|
||||
200
docker/Dockerfile
Normal file
200
docker/Dockerfile
Normal file
|
|
@ -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"]
|
||||
46
docker/build.sh
Executable file
46
docker/build.sh
Executable file
|
|
@ -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"
|
||||
26
docker/freeze.sh
Executable file
26
docker/freeze.sh
Executable file
|
|
@ -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}")"
|
||||
137
docker/smoke_test.py
Normal file
137
docker/smoke_test.py
Normal file
|
|
@ -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())
|
||||
Loading…
Add table
Add a link
Reference in a new issue