docker/build.sh
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env bash
# Build the unsloth-blackwell image on any Linux host with Docker. The build host's
# GPU is NOT used: nvcc cross-compiles.
#
# Usage:
# ./build.sh # builds unsloth-blackwell:latest pinned to unsloth main
# ./build.sh --rocm # builds unsloth-rocm:latest for AMD GPUs
# TAG=2026.05.1 ./build.sh # custom tag
# UNSLOTH_REF=v2026.5.6 UNSLOTH_ZOO_REF=v2026.5.4 ./build.sh # pin git refs
#
# ROCm: the default is the ROCm 7.2.4 base with the pytorch.org rocm7.2 wheels
# (RDNA2, RDNA3, RDNA4, CDNA). Strix APUs and RDNA4 cards get AMD's per-arch
# wheels with --gfx (or ROCM_GFX=...), which switches the index to
# repo.amd.com/rocm/whl/<family>/ and the torch 2.11 line, as install.sh does:
# ./build.sh --rocm --gfx gfx1151 # Strix Halo
# ROCM_GFX=gfx1201 ./build.sh --rocm # RX 9070 XT
# gfx906 (Radeon VII / MI50) needs the last base that carries it, and --gfx gfx906
# leaves out bitsandbytes (no prebuilt gfx906 kernels):
# ROCM_VERSION=6.3.4 ./build.sh --rocm --gfx gfx906
set -euo pipefail
cd "$(dirname "$0")"
ROCM=0
ROCM_GFX="${ROCM_GFX:-}"
while [[ $# -gt 0 ]]; do
case "$1" in
--rocm) ROCM=1 ;;
--gfx) [[ $# -ge 2 ]] || { echo "--gfx needs an arch (e.g. gfx1151)" >&2; exit 2; }
ROCM_GFX="$2"; shift ;;
--gfx=*) ROCM_GFX="${1#--gfx=}" ;;
*) echo "unknown option: $1 (build.sh takes --rocm [--gfx <arch>])" >&2; exit 2 ;;
esac
shift
done
if [[ -n "$ROCM_GFX" && $ROCM -eq 0 ]]; then
echo "--gfx / ROCM_GFX only applies to --rocm builds" >&2
exit 2
fi
TAG="${TAG:-latest}"
PYTHON_VERSION="${PYTHON_VERSION:-3.12}"
UNSLOTH_REF="${UNSLOTH_REF:-main}"
UNSLOTH_ZOO_REF="${UNSLOTH_ZOO_REF:-main}"
UNSLOTH_NOTEBOOKS_REF="${UNSLOTH_NOTEBOOKS_REF:-main}"
if [[ $ROCM -eq 1 ]]; then
IMAGE_NAME="${IMAGE_NAME:-unsloth-rocm}"
# 7.2 is the floor: rocm6.4 tops out below what unsloth-zoo needs, and RDNA4 needs 7.x.
ROCM_VERSION="${ROCM_VERSION:-7.2.4}"
# the index follows the base unless named (7.2.4 -> rocm7.2)
TORCH_INDEX_URL="${TORCH_INDEX_URL:-https://download.pytorch.org/whl/rocm${ROCM_VERSION%.*}}"
else
IMAGE_NAME="${IMAGE_NAME:-unsloth-blackwell}"
CUDA_VERSION="${CUDA_VERSION:-12.8.1}"
UBUNTU_VERSION="${UBUNTU_VERSION:-24.04}"
fi
# Resolve refs to a commit: docker caches RUN layers by command string, so "main" would
# never rebuild.
resolve_git_ref() {
local repo="$1" ref="$2" ls_out sha
printf '%s' "$ref" | grep -Eq '^[0-9a-f]{40}$' && { printf '%s' "$ref"; return 0; }
command -v git >/dev/null 2>&1 || { printf '%s' "$ref"; return 0; }
# ls-remote fails only when the remote is unreachable; warn and pass the name through.
# Also ask for ${ref}^{}: an annotated tag's own sha is not a commit, so uv refuses it.
if ! ls_out="$(git ls-remote "$repo" "$ref" "${ref}^{}" 2>/dev/null)"; then
echo "warning: ${repo} unreachable; passing mutable ref '${ref}' (docker may reuse a cached layer)" >&2
printf '%s' "$ref"
return 0
fi
sha="$(printf '%s\n' "$ls_out" | awk 'NR==1{s=$1; r=$2} $2==r"^{}"{s=$1} END{print s}')"
[ -n "$sha" ] || sha="$ref" # not a branch/tag: a short sha or already-gone ref
printf '%s' "$sha"
}
UNSLOTH_REF="$(resolve_git_ref https://github.com/unslothai/unsloth "$UNSLOTH_REF")"
UNSLOTH_ZOO_REF="$(resolve_git_ref https://github.com/unslothai/unsloth-zoo "$UNSLOTH_ZOO_REF")"
# Freeze the notebooks ref for the same layer-cache reason. Dockerfile.rocm carries neither
# the notebooks nor the llama.cpp prebuilt.
if [[ $ROCM -eq 0 ]]; then
UNSLOTH_NOTEBOOKS_REF="$(resolve_git_ref https://github.com/unslothai/notebooks "$UNSLOTH_NOTEBOOKS_REF")"
# Resolved to a concrete tag so layer caching is correct; pin with LLAMA_PREBUILT_TAG.
resolve_latest_llama_tag() {
curl -fsSL -o /dev/null -w '%{url_effective}' \
"https://github.com/unslothai/llama.cpp/releases/latest" 2>/dev/null \
| sed -n 's#.*/releases/tag/##p'
}
if [ -z "${LLAMA_PREBUILT_TAG:-}" ]; then
LLAMA_PREBUILT_TAG="$(resolve_latest_llama_tag || true)"
if [ -n "$LLAMA_PREBUILT_TAG" ]; then
echo "Resolved latest llama.cpp release: ${LLAMA_PREBUILT_TAG}"
else
LLAMA_PREBUILT_TAG="latest"
echo "Could not resolve latest llama.cpp tag here; passing 'latest' (resolved inside the build)"
fi
fi
fi
if [[ $ROCM -eq 1 ]]; then
echo "Building ${IMAGE_NAME}:${TAG} [AMD ROCm]"
echo " ROCm ${ROCM_VERSION} Python ${PYTHON_VERSION}"
if [[ "$ROCM_GFX" == gfx906 ]]; then
echo " torch index ${TORCH_INDEX_URL} (gfx906: no bitsandbytes)"
elif [[ -n "$ROCM_GFX" ]]; then
echo " torch index AMD per-arch wheels for ${ROCM_GFX} (repo.amd.com/rocm/whl)"
else
echo " torch index ${TORCH_INDEX_URL}"
fi
echo " unsloth @${UNSLOTH_REF}"
echo " unsloth-zoo @${UNSLOTH_ZOO_REF}"
echo
DOCKER_BUILDKIT=1 docker build \
--progress=plain \
-f Dockerfile.rocm \
--build-arg ROCM_VERSION="${ROCM_VERSION}" \
--build-arg PYTHON_VERSION="${PYTHON_VERSION}" \
--build-arg TORCH_INDEX_URL="${TORCH_INDEX_URL}" \
--build-arg ROCM_GFX="${ROCM_GFX}" \
--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 an AMD host (run.sh defaults to the published image, so name this one):"
echo " UNSLOTH_IMAGE=${IMAGE_NAME}:${TAG} UNSLOTH_WORKSPACE_VOLUME= bash run.sh --rocm python /workspace/smoke_test_rocm.py"
exit 0
fi
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 " llama.cpp ${LLAMA_PREBUILT_TAG}"
echo " notebooks @${UNSLOTH_NOTEBOOKS_REF}"
# Read the arch list from the Dockerfile rather than repeating it.
ARCH_LIST="$(sed -n 's/^[[:space:]]*TORCH_CUDA_ARCH_LIST="\([^"]*\)".*/\1/p' \
Dockerfile | head -n1)"
echo " arch list ${ARCH_LIST:-unknown}"
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}" \
--build-arg LLAMA_PREBUILT_TAG="${LLAMA_PREBUILT_TAG}" \
--build-arg UNSLOTH_NOTEBOOKS_REF="${UNSLOTH_NOTEBOOKS_REF}" \
-t "${IMAGE_NAME}:${TAG}" \
.
echo
echo "Built ${IMAGE_NAME}:${TAG}"
echo
# A failing nvidia-smi prints its error on stdout.
HOST_GPU=""
if HOST_GPU_QUERY="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null)"; then
HOST_GPU="$(printf '%s\n' "$HOST_GPU_QUERY" | head -1)"
fi
echo "Smoke test on this host${HOST_GPU:+ (${HOST_GPU})}:"
echo " docker run --rm --gpus all ${IMAGE_NAME}:${TAG} python /workspace/smoke_test.py"
echo
echo "On another host, the same command after:"
echo " docker pull ${IMAGE_NAME}:${TAG} # or load .tar"