c6d92160f6
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
46 lines
1.7 KiB
Bash
Executable file
46 lines
1.7 KiB
Bash
Executable file
#!/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"
|