ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

docker: tighten comments across the Blackwell image and helpers

Browse source 
Condense the verbose explanatory comments added by this branch to their essential
points without dropping any load-bearing rationale. Touches comments and
docstrings only, no code changes. Leaves the stable-tag gate rationale, the
byte-identical enable= sync notes, and the update-alternatives pin comment as is.
This commit is contained in:
Daniel Han 2026-07-18 11:49:15 +00:00
commit a26ead4957
33 changed files with 761 additions and 1300 deletions
Download patch file Download diff file Expand all files Collapse all files

178
.github/workflows/docker-publish.yml vendored
View file

@ -1,30 +1,16 @@
# Builds and publishes the Blackwell-compatible Unsloth Docker image.
#
# The build runs on free GitHub-hosted Ubuntu runners with NO GPU attached.
# This is possible because:
# 1. cu128 PyTorch wheels are fat binaries -- they already ship sm_70 through
# sm_120 SASS on amd64 (and sm_80;90;100;120 on aarch64), 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.
# Runs on free GPU-less GitHub Ubuntu runners: cu128 wheels are fat binaries
# (sm_70..sm_120 amd64, sm_80;90;100;120 aarch64), the Dockerfile pins explicit
# wheel URLs, the build-time check uses torch._C._cuda_getArchFlags() (no CUDA
# device needed), and UNSLOTH_COMPILE_DISABLE=1 blocks GPU-keyed JIT.
#
# Multi-arch: build amd64 and arm64 in parallel on NATIVE GitHub runners
# (`ubuntu-latest` and `ubuntu-24.04-arm`, both free on public repos since
# Aug-2025), then merge the per-arch digests into a single multi-platform
# manifest. Native arm64 is ~3x faster than building aarch64 under QEMU,
# and avoids QEMU's occasional flakiness on long-running cu* installs.
# End users on DGX Spark / Grace pull the arm64 child natively; CUDA works
# as normal (no runtime emulation).
# Multi-arch: amd64 + arm64 build in parallel on native runners (ubuntu-latest +
# ubuntu-24.04-arm), then merge per-arch digests into one manifest. Native arm64
# is ~3x faster and less flaky than QEMU; DGX Spark / Grace pull the arm64 child.
#
# 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
# Required secrets: DOCKERHUB_USERNAME, DOCKERHUB_TOKEN
# Optional variable HAS_GPU_RUNNER='true' gates the smoke-test job.
name: Publish Blackwell Docker image
@ -37,10 +23,8 @@ on:
workflow_dispatch:
inputs:
unsloth_ref:
# Blank means "the dispatched branch" (the resolver below falls back to
# the triggering sha, then main). The stable-tag gates (:core/:latest/
# :studio) require this input to be EMPTY -- stable tags only when the
# operator did not override the source ref -- so a non-blank default
# Blank means "the dispatched branch" (resolver falls back to sha, then
# main). The stable-tag gates require this EMPTY, so a non-blank default
# would make every UI-default dispatch publish SHA tags only.
description: 'unsloth git ref override (blank = dispatched branch + stable tags)'
required: false
@ -62,33 +46,26 @@ env:
REGISTRY: docker.io
IMAGE_NAME: unsloth/unsloth
# Serialise per-ref runs so two pushes to main (or two scheduled
# fires racing a manual dispatch) don't both retag `:latest` from
# different commits. Don't cancel in-progress runs -- the build is
# expensive and a half-built image left around in Docker Hub is
# worse than a slightly stale `:latest` for a few minutes.
# Serialise per-ref runs so two pushes don't both retag :latest from different
# commits. Don't cancel in-progress -- the build is expensive and a half-built
# image is worse than a briefly stale :latest.
concurrency:
group: docker-publish-${{ github.ref }}
cancel-in-progress: false
# Least-privilege default for the GITHUB_TOKEN across every job (CodeQL: set an
# explicit permissions block). Pushes go to Docker Hub via registry creds, not
# GITHUB_TOKEN, so read is enough as the default; the merge jobs that need it
# already declare `packages: write` in their own permissions block.
# Least-privilege default for GITHUB_TOKEN. Pushes use Docker Hub registry creds,
# not GITHUB_TOKEN, so read is enough; jobs needing more declare packages: write.
permissions:
contents: read
jobs:
# ---------------------------------------------------------------------------
# Resolve every upstream ref ONCE, up front -- the llama.cpp prebuilt tag plus
# one unsloth sha, one zoo sha and one notebooks commit -- so both arch legs
# of the base build AND the Studio build bake identical bits. Resolving
# per-leg would let upstream advance between the amd64 and arm64 builds (or
# between the base and Studio builds), putting different content under one
# published tag. An explicit dispatch input pins a frozen value; otherwise a
# branch/tag is frozen to a sha via ls-remote (falling back to the bare ref
# on a lookup miss so the Dockerfile can still fetch it by name), and the
# llama "latest" follows the /releases/latest redirect (mirrors build.sh).
# Resolve every upstream ref ONCE (llama tag + unsloth/zoo shas + notebooks
# commit) so both arch legs and the Studio build bake identical bits; resolving
# per-leg would let upstream advance mid-run under one tag. A dispatch input
# pins a frozen value; else a branch/tag is frozen to a sha via ls-remote
# (falling back to the bare ref on a miss), and llama "latest" follows the
# /releases/latest redirect (mirrors build.sh).
# ---------------------------------------------------------------------------
prepare:
runs-on: ubuntu-latest
@ -116,10 +93,8 @@ jobs:
echo "tag=${TAG:-latest}" >> "$GITHUB_OUTPUT"
echo "llama.cpp prebuilt tag: ${TAG:-latest}"
# Requested-ref precedence (same as the old inline build-arg): the
# dispatch input wins (blank by default, so stable tags stay enabled),
# else the pushed tag, else the triggering commit sha, else main --
# then frozen to one sha per the job header.
# Requested-ref precedence: dispatch input, else pushed tag, else trigger
# sha, else main -- then frozen to one sha per the job header.
- name: Resolve unsloth ref
id: unsloth_ref
env:
@ -140,10 +115,9 @@ jobs:
echo "ref=${SHA}" >> "$GITHUB_OUTPUT"
echo "unsloth ref: ${SHA}"
# Mirror the unsloth tag into the zoo ONLY when that tag actually exists
# there. unsloth's v* tags are Studio releases the zoo never cuts (the zoo
# repo currently has no tags at all), so blindly mirroring github.ref_name
# made every tag publish fail inside the Dockerfile's zoo install.
# Mirror the unsloth tag into the zoo ONLY when that tag exists there:
# unsloth's v* tags are Studio releases the zoo never cuts, so blindly
# mirroring github.ref_name made every tag publish fail at zoo install.
- name: Resolve unsloth-zoo ref
id: zoo_ref
run: |
@ -165,9 +139,8 @@ jobs:
echo "ref=${SHA}" >> "$GITHUB_OUTPUT"
echo "unsloth-zoo ref: ${SHA}"
# Freeze unslothai/notebooks to ONE commit per the job header, so the
# baked templates + .unsloth_template_commit are identical across legs
# and release reruns.
# Freeze notebooks to ONE commit per the job header, so baked templates +
# .unsloth_template_commit are identical across legs and reruns.
- name: Resolve unsloth/notebooks commit
id: notebooks
env:
@ -184,12 +157,10 @@ jobs:
echo "notebooks commit: ${SHA}"
# ---------------------------------------------------------------------------
# Per-arch build. The matrix fans out two parallel jobs on the matching
# native runner. Each pushes a single-arch image *by digest* (no human-
# readable tag), and the merge job below stitches the two digests into one
# multi-arch manifest under the real tags. This is the canonical pattern
# from docker/build-push-action's docs and avoids the "last push wins" race
# that you get when two jobs push the same tag separately.
# Per-arch build. The matrix fans out two parallel jobs on native runners;
# each pushes a single-arch image by digest (no tag), and the merge job
# stitches the digests into one multi-arch manifest. Canonical build-push-action
# pattern; avoids the "last push wins" race of two jobs pushing the same tag.
# ---------------------------------------------------------------------------
build:
needs: prepare
@ -210,15 +181,13 @@ jobs:
steps:
- uses: actions/checkout@v4
# Free up ~20GB on the runner so cu128 wheels + cudnn fit. Layout is
# similar between the amd64 and arm64 runners but not identical -- the
# arm64 image lacks /usr/share/dotnet, hence `|| true`.
# Free up ~20GB so cu128 wheels + cudnn fit. Runner layouts differ (arm64
# lacks /usr/share/dotnet), hence `|| true`.
- name: Reclaim disk
run: |
# Hosted runners keep only ~14-20 GB free -- not enough for the image
# plus buildkit state (Studio install hit ENOSPC before this list grew).
# None of these toolchains are used here; paths differ across the amd64
# and arm64 runners, hence `|| true`.
# Hosted runners keep only ~14-20 GB free -- not enough for the image +
# buildkit state. None of these toolchains are used; paths differ across
# runners, hence `|| true`.
sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc \
/opt/hostedtoolcache "$AGENT_TOOLSDIRECTORY" \
/usr/local/.ghcup /usr/share/swift \
@ -236,9 +205,8 @@ jobs:
username: ${{ secrets.DOCKERHUB_USERNAME }}
password: ${{ secrets.DOCKERHUB_TOKEN }}
# Pull the image label/annotation set we'll attach to the FINAL manifest.
# We don't apply tags at this layer because each per-arch build pushes by
# digest only; tags get attached by the merge job.
# Labels/annotations for the FINAL manifest. No tags here -- each per-arch
# build pushes by digest only; tags are attached by the merge job.
- name: Resolve labels
id: meta
uses: docker/metadata-action@v5
@ -253,16 +221,13 @@ jobs:
file: ./docker/Dockerfile
platforms: ${{ matrix.platform }}
labels: ${{ steps.meta.outputs.labels }}
# Per-arch build cache. Keying on the platform suffix lets the two
# matrix legs reuse their own caches without colliding.
# Per-arch build cache: the platform suffix keeps the two legs from colliding.
cache-from: type=gha,scope=build-${{ matrix.platform }}
cache-to: type=gha,scope=build-${{ matrix.platform }},mode=max
outputs: type=image,name=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }},push-by-digest=true,name-canonical=true,push=true
# NOTE: keep prose OUT of build-args -- docker/build-push-action
# forwards every non-empty line verbatim, so a leading-# line would be
# passed as a bogus --build-arg. All four values come from the prepare
# job: resolved once so both arch legs and the Studio build bake
# identical bits (precedence rules live on prepare's steps).
# Keep prose OUT of build-args -- build-push-action forwards every
# non-empty line verbatim, so a #-line becomes a bogus --build-arg. All
# four values come from the prepare job (resolved once).
build-args: |
CUDA_VERSION=12.8.1
UBUNTU_VERSION=24.04
@ -272,9 +237,8 @@ jobs:
LLAMA_PREBUILT_TAG=${{ needs.prepare.outputs.llama_tag }}
UNSLOTH_NOTEBOOKS_REF=${{ needs.prepare.outputs.notebooks_commit }}
# Stash the per-arch digest as an artifact for the merge job to pick up.
# Filenames need to be unique across the matrix; `platform` contains a
# slash so substitute it for a dash.
# Stash the per-arch digest as an artifact for the merge job. `platform`
# has a slash, so substitute a dash for a unique filename.
- name: Export digest
run: |
mkdir -p /tmp/digests
@ -290,9 +254,8 @@ jobs:
retention-days: 1
# ---------------------------------------------------------------------------
# Merge the two per-arch digests into a multi-platform manifest under the
# real, user-facing tag(s). This job runs only after both `build` matrix
# legs finish successfully.
# Merge the two per-arch digests into a multi-platform manifest under the real
# user-facing tag(s). Runs only after both build legs succeed.
# ---------------------------------------------------------------------------
merge:
runs-on: ubuntu-latest
@ -302,10 +265,9 @@ jobs:
contents: read
packages: write
outputs:
# Multi-arch manifest digest of the just-published base image. The
# build-studio job FROMs this exact digest so the Studio image always
# layers on the bits published by THIS run, not whatever `base`
# happens to point at when the job is scheduled.
# Manifest digest of the just-published base image; build-studio FROMs this
# exact digest so Studio layers on THIS run's bits, not whatever `base`
# points at later.
digest: ${{ steps.manifest_digest.outputs.digest }}
steps:
- uses: actions/download-artifact@v4
@ -371,11 +333,9 @@ jobs:
# ---------------------------------------------------------------------------
# Full image: base + Unsloth Studio + JupyterLab + sshd (Dockerfile.studio).
# This is what :latest points at, matching the service set of the previous
# production image. Same by-digest build + manifest-merge pattern as the
# base. FROMs the exact base manifest digest published by the merge job.
# The arm64 leg builds Studio's vite frontend natively on the arm runner;
# that is the long pole, hence the larger timeout.
# This is :latest. Same by-digest build + merge pattern as the base, FROMing the
# base manifest digest from the merge job. The arm64 leg builds Studio's vite
# frontend natively (the long pole), hence the larger timeout.
# ---------------------------------------------------------------------------
build-studio:
# `merge` for the freshly-published base manifest digest; `prepare` for the
@ -399,10 +359,9 @@ jobs:
- name: Reclaim disk
run: |
# Hosted runners keep only ~14-20 GB free -- not enough for the image
# plus buildkit state (Studio install hit ENOSPC before this list grew).
# None of these toolchains are used here; paths differ across the amd64
# and arm64 runners, hence `|| true`.
# Hosted runners keep only ~14-20 GB free -- not enough for the image +
# buildkit state. None of these toolchains are used; paths differ across
# runners, hence `|| true`.
sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc \
/opt/hostedtoolcache "$AGENT_TOOLSDIRECTORY" \
/usr/local/.ghcup /usr/share/swift \
@ -434,17 +393,14 @@ jobs:
file: ./docker/Dockerfile.studio
platforms: ${{ matrix.platform }}
labels: ${{ steps.meta.outputs.labels }}
# mode=min (final layers only): a mode=max cache of this ~24GB
# image would blow straight through the 10GB per-repo GHA cache
# quota and evict the base build's cache for zero hit-rate gain.
# mode=min (final layers only): mode=max on this ~24GB image would blow
# the 10GB GHA cache quota and evict the base build's cache for no gain.
cache-from: type=gha,scope=studio-${{ matrix.platform }}
cache-to: type=gha,scope=studio-${{ matrix.platform }},mode=min
outputs: type=image,name=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }},push-by-digest=true,name-canonical=true,push=true
# All three pins are the SAME resolved values the base build baked
# (prepare job), so the Studio tree, its zoo overlay AND its llama.cpp
# bundle match the base image even if a branch or upstream release
# moved mid-run. (Prose stays out of build-args -- forwarded lines
# must be KEY=VALUE only.)
# All three pins are the SAME values the base build baked (prepare job),
# so Studio, its zoo overlay and its llama.cpp match the base even if
# upstream moved mid-run. (build-args must be KEY=VALUE only.)
build-args: |
BASE_IMAGE=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}@${{ needs.merge.outputs.digest }}
UNSLOTH_STUDIO_REF=${{ needs.prepare.outputs.unsloth_ref }}
@ -522,9 +478,8 @@ jobs:
done
# ---------------------------------------------------------------------------
# Optional: pull the freshly published image onto a self-hosted GPU runner
# and run smoke_test.py. Skipped automatically when no GPU runner is
# registered. Architecture matches whatever the runner is.
# Optional: pull the freshly published image onto a self-hosted GPU runner and
# run smoke_test.py. Skipped when no GPU runner is registered.
# ---------------------------------------------------------------------------
smoke-test:
needs: [merge, merge-studio]
@ -601,9 +556,8 @@ jobs:
ok_studio=0; ok_jupyter=0
for i in $(seq 1 60); do
if curl -fsS http://localhost:18000/api/health >/dev/null 2>&1; then ok_studio=1; fi
# Probe /login, not /api: the launcher always sets a Jupyter password
# hash, so /api returns 403 (curl -f would never flip ok_jupyter).
# /login is the unauthenticated page and 200s once the server is up.
# Probe /login, not /api: the launcher sets a password hash so /api
# returns 403; /login is unauthenticated and 200s once up.
if curl -fsS http://localhost:18888/login >/dev/null 2>&1; then ok_jupyter=1; fi
[ "$ok_studio" = 1 ] && [ "$ok_jupyter" = 1 ] && break
sleep 5

625
docker/Dockerfile
View file

@ -1,48 +1,32 @@
# syntax=docker/dockerfile:1.7
# -----------------------------------------------------------------------------
# Unsloth + unsloth-zoo for every current NVIDIA arch (Turing -> Blackwell),
# on both linux/amd64 and linux/arm64.
# on 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:
# Why it works:
# * cu128 wheels ship native SASS (no PTX), verified via `cuobjdump --list-elf`:
# 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 runs the PRECOMPILED SASS (torch,
# llama.cpp, source-built ops per the arch list below).
# * Unsloth's runtime kernels are Triton, which JIT-compiles per device at
# first run. JIT targets the ACTUAL device cap, and the bundled cu12.8
# ptxas/NVRTC cannot emit compute_103 (sm_103) or compute_121 (sm_121).
# Both are handled by the cu13 NVRTC/ptxas override below: amd64 sm_103
# (B300/GB300) and arm64 sm_121 (DGX Spark / GB10). Precompiled SASS also
# runs on sm_103 via sm_100 forward-compat and on sm_121 via sm_120
# forward-compat, so only JIT-heavy paths depend on the override.
# * 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;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.
# * SASS is forward-compatible within a major: sm_86->sm_89 (Ada),
# sm_100->sm_103 (B300/GB300), sm_120->sm_121 (DGX Spark/GB10), so every
# non-Jetson GPU on https://developer.nvidia.com/cuda/gpus runs precompiled
# SASS (torch, llama.cpp, source-built ops).
# * Triton kernels JIT per-device at first run; the bundled cu12.8 ptxas/NVRTC
# cannot emit compute_103/compute_121, so the cu13 override below handles
# amd64 sm_103 and arm64 sm_121 (SASS still runs there via forward-compat, so
# only JIT-heavy paths need it).
# * Rare source builds compile against
# TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;12.0+PTX"; the host GPU is
# irrelevant, 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:
# docker run --privileged --rm tonistiigi/binfmt --install all # 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).
# Cross-arch build (arm64 / sm_121): built via QEMU binfmt on an x86_64 host
# (`docker run --privileged --rm tonistiigi/binfmt --install all` once, then
# `docker buildx build --platform linux/arm64 ...`). QEMU is build-time only;
# the image runs natively on aarch64. xformers has no cu128 aarch64 wheel, so
# arm64 falls back to Unsloth's SDPA (~5-10% slower, 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 (the one-time setup above)
# * A GPU is NOT required at build time.
# Build host needs Docker buildkit + buildx, and QEMU binfmt for arm64-on-x86_64;
# nvidia-container-toolkit only for test-time `--gpus all`. No GPU at build time.
# -----------------------------------------------------------------------------
ARG CUDA_VERSION=12.8.1
@ -54,9 +38,8 @@ ARG PYTHON_VERSION=3.12
# =============================================================================
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 (the xformers aarch64 gap -- see header).
# TARGETARCH (buildx: amd64/arm64) selects the unsloth extras matching the
# wheels available for the platform (xformers aarch64 gap -- see header).
ARG TARGETARCH
ARG PYTHON_VERSION
ENV DEBIAN_FRONTEND=noninteractive \
@ -64,42 +47,27 @@ ENV DEBIAN_FRONTEND=noninteractive \
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 -- covered by sm_100 SASS (see above),
# NOT a separate target here: the bundled CUDA 12.8 nvcc cannot
# compile compute_103 (added in CUDA 12.9), so listing 10.3 would
# break any source / JIT build that honors TORCH_CUDA_ARCH_LIST.
# 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).
# Cross-compile for every current NVIDIA arch (developer.nvidia.com/cuda/gpus):
# sm_75 Turing (T4, RTX 20xx) | sm_80 A100/A30 | sm_86 A40/RTX 30xx
# sm_89 Ada (L4/L40/RTX 40xx) | sm_90 Hopper (H100/H200/GH200)
# sm_100 Blackwell DC (B100/B200/GB200) | sm_120 Blackwell (RTX 50xx, RTX PRO 6000)
# sm_103 (B300/GB300) and sm_121 (GB10) are omitted: CUDA 12.8 nvcc cannot
# compile compute_103/121 and sm_100/sm_120 SASS cover them via forward-compat.
# +PTX on the highest lets future revisions JIT. Same list on both arches:
# nvcc emits archs absent from the host, irrelevant extras only cost a little
# compile time (no source builds on this pin set).
TORCH_CUDA_ARCH_LIST="7.5;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/.
# Build-host-independence guards: the build must NEVER introspect a GPU (host
# may be a B200, RTX 6000, or GPU-less CI) so all yield byte-identical images.
# 1) no JIT-compiled sm_NNN blob into unsloth_compiled_cache/ at import.
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.
# 2) don't probe torch.cuda.is_available() at setup (would silently skip wheels).
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`).
# 3) empty CUDA_VISIBLE_DEVICES so stray torch.cuda calls see no devices, not
# host-specific paths (re-enabled at runtime via `docker run --gpus all`).
CUDA_VISIBLE_DEVICES=""
RUN apt-get update && apt-get install -y --no-install-recommends \
@ -112,48 +80,34 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
&& 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.
# Isolated prefix. We do NOT touch the system Python (Ubuntu 24.04 marks it
# externally-managed, PEP 668); the venv bootstraps pip via ensurepip (from the
# python3.12-venv apt package) and gets uv 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.)
# Unified install: torch + triton + bitsandbytes + unsloth + unsloth_zoo in a
# SINGLE uv pass. Mandatory -- splitting it lets bnb's transitive `cuda-toolkit`
# silently upgrade torch to 2.12.0+cu130 in a later pass, breaking the pinned
# cu128 xformers wheel (the cu cascade hits after xformers is on disk).
#
# 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 `cu128onlytorch2110` extra.
# Flags:
# --index-strategy unsafe-best-match: the PyTorch index serves an old
# requests==2.28.1 that conflicts with datasets>=2.32.2; both indexes
# (pytorch cu128 + pypi) are equally trusted, so override uv's first-wins.
# --extra-index-url .../cu128: torch +cu128 wheels + the xformers/cu128 URLs
# from unsloth's cu128onlytorch2110 extra.
#
# Why the plain `huggingface` extra plus an EXPLICIT xformers pin (amd64):
# The cu128 CUDA extras on main stop at the torch2100 family, which pins
# xformers 0.0.34 -> torch==2.10.0 and would conflict with the torch 2.11.0
# line held below; a not-yet-existing extra name would only WARN (pip/uv
# install plain unsloth) and silently drop xformers until the required-
# package check below failed the build. Pinning `xformers==0.0.35` directly
# (it does not pin torch, pairing with 2.11.0) keeps this build
# self-contained on today's main; arm64 stays xformers-less (no cu128
# aarch64 wheel -- see header).
# Plain `huggingface` extra + explicit xformers pin (amd64): the cu128 extras on
# main stop at torch2100 (xformers 0.0.34 -> torch 2.10.0), conflicting with the
# torch 2.11.0 held below; a missing extra name would only WARN and drop xformers
# until the check below failed. Pinning xformers==0.0.35 (untied to torch) keeps
# this self-contained; arm64 stays xformers-less (no cu128 aarch64 wheel).
#
# Why no `flash-attn` here:
# - FA3 is hard-refused on Blackwell (Dao-AILab/flash-attention#1810).
# - FA2 has no prebuilt wheel for cu128+torch2.11+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.
# No flash-attn: FA3 is hard-refused on Blackwell (Dao-AILab/flash-attention#1810);
# FA2 has no cu128+torch2.11+cp312 wheel (~30min fragile source build on CI) and
# Unsloth falls back to xformers/SDPA anyway. Users on Ampere/Ada/Hopper can
# `pip install flash-attn` on top at deploy time.
ARG UNSLOTH_REF=main
ARG UNSLOTH_ZOO_REF=main
RUN set -eux \
@ -176,22 +130,12 @@ RUN set -eux \
"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.11.0 for vLLM's
# choice. Splitting the install lets the unified pass settle on torch
# 2.11.0 first, then vLLM bolts on top: with torch held at 2.11.0 the
# resolver lands on the newest compatible vLLM (0.20+, which pins torch
# 2.11.0) by itself, and tracks our torch pin when it moves.
# * 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 manually on Spark hardware, not in CI.
#
# vLLM: required by Unsloth's GRPO path (fast_inference=True). Installed as a
# SECOND uv pass so the unified pass settles on torch 2.11.0 first, then vLLM
# bolts on top: with torch held, uv picks the newest compatible vLLM (0.20+ pins
# torch 2.11.0) and tracks our pin. PyPI ships x86_64 + aarch64 wheels since 0.17,
# so this runs on arm64 too; amd64 failures abort, arm64 is fail-soft (aarch64
# kernels are validated on Spark hardware, not CI).
# https://docs.vllm.ai/en/latest/getting_started/installation/gpu/
# https://wheels.vllm.ai/nightly
ARG INSTALL_VLLM=auto
@ -204,25 +148,17 @@ RUN set -eux \
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.11.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.
# Explicit && chain, not `set -e` -- POSIX shells disable errexit inside a
# condition context (verified on dash), masking install failures.
# Step 1: uv resolves vLLM's deps with torch==2.11.0 held, landing on the
# newest matching vLLM (fails loudly if none); unsafe-best-match picks the
# best wheel across the three indexes.
# Step 2: vLLM pulls numpy down to 2.2.6 whose wheel ships a broken
# numpy.testing (tests/ stripped), crashing `from numpy import *` and thus
# `import unsloth`; upgrade numpy back to a self-consistent release.
# Step 3: vLLM pins numba 0.61.2 which hard-refuses numpy>=2.3, but the
# stack needs numpy>=2.3; lift numba to one supporting numpy 2.4 (0.65
# imports cleanly, vllm still imports).
{ ${VENV}/bin/uv pip install \
--python ${VENV}/bin/python \
--pre \
@ -240,12 +176,10 @@ RUN set -eux \
&& ${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.
# flashinfer-jit-cache: precompiled cubins so flashinfer ops don't hit the
# JIT path (standalone `vllm serve` dies there for fmha_gen on sm_100a;
# in-process GRPO survives since zoo blocks the FlashInfer JIT). Removes
# the runtime-compile failure class for ~1.5 GB.
&& { ${VENV}/bin/uv pip install \
--python ${VENV}/bin/python \
--index-url https://flashinfer.ai/whl/cu128 \
@ -255,10 +189,9 @@ RUN set -eux \
} || { \
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
# 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`). arm64 staging CI
# 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 \
@ -272,37 +205,27 @@ RUN set -eux \
echo ">> vLLM skipped (INSTALL_VLLM=${INSTALL_VLLM}, TARGETARCH=${TARGETARCH:-amd64})"; \
fi
# JupyterLab so the published image runs unslothai/notebooks out of the box:
# JupyterLab so the 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
# Separate pass AFTER the torch-pinned resolves: this closure is pure-Python
# and never names torch, so uv cannot disturb the cu128 pin set (naming torch
# without the cu128 index 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)
# Separate pass AFTER the torch-pinned resolves: pure-Python, never names torch,
# so uv can't disturb the cu128 pin set. These are declared by notebook install
# cells (neutralised by the in-image runner), so bake them here:
# matplotlib plotting; some trust_remote_code files import it (DeepSeek-OCR)
# soundfile TTS audio read/write (bundles libsndfile)
# evaluate+jiwer Whisper WER metric
# 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 is installed separately below (no aarch64 wheel; see that block).
# 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 in-image pin record (/opt/unsloth-venv/requirements.lock.txt).
# langid DeepSeek-R1 GRPO reward language-id check
# easydict some vision trust_remote_code modeling files
# protobuf slow->fast tokenizer conversion for sentencepiece
# omegaconf TTS + NeMo-Gym RL notebook configs
# einx TTS codec tensor-rearrange (Llasa/Oute/Spark)
# librosa Whisper audio features (pulls numba, already pinned >=0.65)
# ftfy Oute TTS text normalisation
# decord is separate below (no aarch64 wheel). Pinned (==) for reproducible
# rebuilds; bump deliberately. The resolve must NOT move torch/numpy/numba --
# the assertion below fails the build if it did. Transitive deps land in
# /opt/unsloth-venv/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" \
@ -311,10 +234,8 @@ RUN ${VENV}/bin/uv pip install \
"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.11.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.
# Install it on its own: HARD on amd64 (a missing/incompatible wheel is a real
# regression there and must fail the build), fail-soft on arm64/other (no wheel
# exists, so drop the ERNIE-VL video path rather than break the image build).
# decord (ERNIE-VL video decode) has wheels only for x86_64. Installed alone:
# HARD on amd64 (a missing wheel is a real regression), fail-soft elsewhere.
RUN if [ "${TARGETARCH:-amd64}" = "amd64" ]; then \
${VENV}/bin/uv pip install --python ${VENV}/bin/python "decord==0.6.0"; \
else \
@ -322,17 +243,12 @@ RUN if [ "${TARGETARCH:-amd64}" = "amd64" ]; then \
|| echo ">> decord skipped (no matching wheel for ${TARGETARCH:-}); ERNIE-VL video decode unavailable"; \
fi
# 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.11 pairs with torch 2.11 (a mismatched
# build references other torch symbols and fails 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).
# Audio decode out of the box: TTS/STT notebooks decode datasets' Audio features
# through torchcodec. Three traps: (1) torchcodec 0.11 must pair with torch 2.11;
# (2) the wheel must come from the cu128 channel, not the PyPI default (cu13,
# dlopens libnvrtc.so.13); (3) its libs dlopen venv torch/NVIDIA libs that ld.so
# can't see, registered via ld.so.conf.d in the runtime stage (this layer only
# installs the wheels; the import check needs ffmpeg from the runtime stage).
# Fail-soft on arches without a matching wheel.
RUN set -eux \
&& { ${VENV}/bin/uv pip install \
@ -342,19 +258,15 @@ RUN set -eux \
&& ${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.
# transformers SIDECARS for per-notebook version activation (see
# unsloth_nb_compat.py). unslothai/notebooks pin many transformers versions; the
# base venv ships the newest 5.x. Each sidecar is transformers==X + matched
# huggingface_hub/tokenizers/safetensors, installed --no-deps into its own
# --target under ${VENV}/tf-sidecars. Activating one (prepend to sys.path before
# any ML import) swaps transformers WITHOUT touching the cu128 torch/vLLM/unsloth
# base -- verified: base unsloth loads + generates under 4.57.6 and 5.5.0 on B200.
# Versions mirror Studio's tiers (4.57.6 + 5.3.0/5.5.0/5.10.2); companions are
# resolved at build time. ~300MB after the strip 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)"; \
@ -376,35 +288,22 @@ RUN set -eux \
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 <image> cat /opt/unsloth-venv/requirements.lock.txt`
# if you need to forensic-diff two images.
# 5) Informational pin record (NOT byte-reproducible: pip freeze omits wheel
# hashes and unsloth/unsloth_zoo/vllm --pre float from VCS/nightly). Read via
# `docker run --rm <image> cat /opt/unsloth-venv/requirements.lock.txt`.
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.
# The `-name tests` strip excludes numpy's tests dirs: numpy 2.4 ships
# numpy/_core/tests/ back, and removing it re-triggers the
# `from numpy._core.tests._natype import pd_NA` ImportError the earlier upgrade
# fixed. Other verified-safe size cuts:
# * npp: torchcodec dlopens only libnppicc + libnppc; drop the other ~10 libs (~388MB).
# * static .a archives (~143MB): link-time only, nothing links venv archives.
# * nvshmem device .bc (~30MB): device-relink only; host .so kept.
# Do NOT strip headers (torch/include): causal-conv1d / mamba-ssm build against
# them at notebook time with --no-build-isolation.
RUN set -eux \
&& find ${VENV} -depth -type d -name __pycache__ -exec rm -rf {} + \
&& find ${VENV} -depth -type d -name tests \
@ -423,17 +322,11 @@ RUN set -eux \
&& 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.
# (1) arch-list check uses the RAW C++ accessor: torch.cuda.get_arch_list()
# returns [] with no GPU visible (CUDA_VISIBLE_DEVICES is empty here).
# (2) required packages verified via metadata only -- we do NOT import unsloth/
# unsloth_zoo (their __init__ needs a real CUDA device). Import 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")
@ -447,15 +340,13 @@ print("arches:", arches)
assert torch.__version__.startswith("2.11.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.
# cu128 wheels ship sm_120 native SASS on both amd64 and aarch64. On arm64 DGX
# Spark (sm_121) runs it via forward-compat; sm_121 is never in a cu128 wheel.
assert "sm_120" in arches, f"sm_120 missing: {arches}"
print(f"OK: torch 2.11.0+cu128 with sm_100 + sm_120 native SASS intact ({target})")
from importlib.metadata import version, PackageNotFoundError
# xformers is amd64-only (the aarch64 wheel gap -- see Dockerfile header).
# xformers is amd64-only (aarch64 wheel gap -- see header).
REQUIRED = ["torch", "triton", "bitsandbytes", "unsloth",
"unsloth_zoo", "transformers", "trl", "peft", "accelerate"]
if target == "amd64":
@ -484,19 +375,15 @@ 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.
# The "-base-" variant (vs "-cudnn-runtime-") drops ~2.7 GB of system CUDA libs
# we never load: torch wheels bake their own cuDNN/cuBLAS/etc into torch/lib/ and
# libtorch_cuda.so's RPATH resolves through the wheel, not the system (verified
# via `readelf -d torch/lib/libtorch_cuda.so`). The base still provides
# nvidia-smi + libcuda stubs + libnvidia-ml, all the entrypoint + 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.
# The base manifest is multi-arch; buildx picks the right one for
# TARGETPLATFORM at this FROM line, no conditional needed.
ARG TARGETARCH
ARG PYTHON_VERSION
ARG CUDA_VERSION
@ -507,29 +394,20 @@ ENV DEBIAN_FRONTEND=noninteractive \
PATH=/opt/unsloth-venv/bin:${PATH} \
HF_HOME=/workspace/.cache/huggingface \
TRITON_CACHE_DIR=/workspace/.cache/triton \
# Keep the arch list visible at runtime so an in-container source build of a
# custom CUDA op gets the same SASS coverage as the builder stage. 10.3 is
# omitted for the same cu12.8-cannot-emit-compute_103 reason as the builder
# list + header (sm_103 runs sm_100 SASS via forward-compat).
# Keep the arch list at runtime so an in-container source build gets the same
# SASS coverage as the builder; 10.3 omitted (cu12.8 can't emit compute_103;
# sm_103 runs sm_100 SASS via forward-compat).
TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;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.
# System packages needed by the notebooks:
# zstd Ollama installer (`curl ollama.com/install.sh | sh`) extracts a zstd tarball
# ffmpeg torchcodec dlopens system FFmpeg libs (not bundled in the wheel)
# wget notebooks fetch assets with `!wget URL`
# ninja-build flashinfer cpp_ext JIT shells out to ninja; subprocesses may not
# inherit the venv pip ninja on PATH
# cuda-nvcc + cudart-dev flash-linear-attention's TileLang backend (Qwen3.5 in
# Studio) JIT-compiles CUDA kernels via /usr/local/cuda/bin/nvcc,
# absent from the -base image
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 \
@ -542,49 +420,33 @@ RUN CUDA_PKG="$(echo "${CUDA_VERSION}" | awk -F. '{print $1"-"$2}')" \
&& 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.)
# gcc + g++ + python3.12-dev in runtime: Triton's nvidia backend lazily compiles
# a C extension (CudaUtils) on first GPU access; without a C compiler + Python
# headers the first forward pass dies with "Failed to find C compiler". ~250MB,
# the cost of correct JIT (pre-compiling would need a GPU).
COPY --from=builder /opt/unsloth-venv /opt/unsloth-venv
# Blackwell JIT fix for sm_103 (amd64) and sm_121 (arm64) -- the cu12.8 JIT
# gap described in the header. Two JIT paths need the cu13 override:
#
# (1) torch's bundled libnvrtc.so.12 is CUDA 12.8. The jiterator C++ side
# queries the device cap directly, so any NVRTC JIT path (e.g.
# torch.fft.rfft(complex).abs(), used inside mel-spectrogram code)
# errors out on sm_103/sm_121. Fix: stage a cu13 NVRTC alias beside the
# immutable cu12.8 default (mechanics at the staging step below).
#
# (2) Triton's nvidia backend invokes its OWN bundled ptxas, which in the
# triton 3.6.0 we pin is still CUDA 12.8 (V12.8.93): it tops out at
# sm_120, rejects sm_103, and silently downgrades sm_121 to sm_80 per
# triton-lang/triton#8335. Fix: install cu13 ptxas and point Triton at
# it with TRITON_PTXAS_PATH.
#
# Both cu13 tools are CPU-side compilers (no driver-floor bump at INSTALL
# time), but their OUTPUT cubin needs a >= 580 driver to LOAD, so neither is
# baked as a global ENV/symlink default -- that would break the Ampere/Ada/
# Hopper/Turing GPUs this image still supports on 570-579 drivers. Instead
# select_cuda_jit_tools in entrypoint.sh activates them per device, only for
# sm_103/sm_121 (which only ship on >= 580 drivers, so the gate is always
# safe). Both arches carry the ~400 MB: amd64 for sm_103, arm64 for sm_121.
# Blackwell JIT fix for sm_103 (amd64) and sm_121 (arm64) -- the cu12.8 JIT gap
# from the header. Two JIT paths need the cu13 override:
# (1) torch's bundled libnvrtc.so.12 (CUDA 12.8): any NVRTC JIT path (e.g.
# torch.fft.rfft(complex).abs() in mel-spectrogram code) errors on
# sm_103/sm_121. Fix: stage a cu13 NVRTC alias beside the cu12.8 default.
# (2) Triton's bundled ptxas (still 12.8 in triton 3.6.0) tops out at sm_120,
# rejects sm_103, silently downgrades sm_121 to sm_80 (triton-lang/triton#8335).
# Fix: install cu13 ptxas, point Triton at it via TRITON_PTXAS_PATH.
# Both cu13 tools are CPU-side compilers (no install-time driver bump), but their
# output cubin needs a >=580 driver to LOAD, so neither is a global default (would
# break Ampere/Ada/Hopper/Turing on 570-579 drivers). select_cuda_jit_tools in
# entrypoint.sh activates them per device, only for sm_103/sm_121 (>=580 drivers,
# always safe). Both arches carry the ~400 MB.
RUN set -eux; \
# The nvidia/cuda base already configures the CUDA apt repo (x86_64 or
# sbsa) 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.
# Verified on the ubuntu-24.04 (x86_64) and ubuntu-24.04-arm runners.
# The nvidia/cuda base already configures the CUDA apt repo with its own
# Signed-By keyring; adding cuda-keyring_1.1-1_all.deb would add a second
# sources file with a different Signed-By and make `apt-get update` refuse
# the repo ("Conflicting values set for option Signed-By"). The base repo
# is monolithic and serves 13.x too, so install cu13 packages directly
# without touching the keyring. Verified on x86_64 and arm runners.
apt-get update; \
apt-get install -y --no-install-recommends \
cuda-nvrtc-13-0 \
@ -598,13 +460,11 @@ RUN set -eux; \
# to manual mode so later apt operations cannot flip it again.
update-alternatives --set cuda /usr/local/cuda-12.8; \
rm -rf /var/lib/apt/lists/*; \
# (1) NVRTC staging: keep the wheel's real cu12.8 lib as .cu128.orig,
# point libnvrtc.so.12 at it (relative symlink), and stage
# .cu13 -> the cu13 lib; select_cuda_jit_tools retargets the
# symlink ONLY on sm_103/sm_121 hosts. The default needs no
# runtime write, so a non-root `docker run --user` container
# (which cannot rewrite the symlink) keeps cu12.8, loadable on
# every supported 570+ driver.
# (1) NVRTC staging: keep the wheel's cu12.8 lib as .cu128.orig, point
# libnvrtc.so.12 at it (relative symlink), stage .cu13 -> the cu13 lib;
# select_cuda_jit_tools retargets the symlink only on sm_103/sm_121.
# The default needs no runtime write, so a non-root --user container
# keeps cu12.8, loadable on every supported 570+ driver.
NVRTC_DIR=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages/nvidia/cuda_nvrtc/lib; \
if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ] && [ ! -L "${NVRTC_DIR}/libnvrtc.so.12" ]; then \
mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \
@ -614,12 +474,10 @@ RUN set -eux; \
# (2) ptxas: the cu13 nvcc package above provides it; TRITON_PTXAS_PATH is set
# per device at boot (select_cuda_jit_tools) for the same driver-floor reason.
# 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.
# Register the venv's torch + NVIDIA lib dirs with the loader so torchcodec can
# dlopen them. ld.so.conf.d, NOT LD_LIBRARY_PATH: the cache is consulted after
# DT_RUNPATH, so llama.cpp keeps resolving its own $ORIGIN libs first. Import
# check runs here (ffmpeg lives in this stage); fail-soft if the wheel is absent.
RUN set -eux \
&& SP=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages \
&& printf "%s\n" "$SP/torch/lib" "$SP/nvidia/cuda_nvrtc/lib" \
@ -630,34 +488,25 @@ RUN set -eux \
"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.
# Prebuilt llama.cpp so GGUF export works out of the box. unsloth_zoo's
# save_pretrained_gguf() -> check_llama_cpp() looks for llama-quantize +
# convert_hf_to_gguf.py + gguf-py/ in $UNSLOTH_LLAMA_CPP_PATH; without a baked
# install the first export hits install_llama_cpp()'s prompt + slow source build.
#
# 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):
# NOT studio/install_llama_prebuilt.py: it selects a bundle for the CURRENT host
# (nvidia-smi etc., which leak through a GPU build host). The build must never
# introspect the host, so pin release + asset by build target (see
# fetch_llama_prebuilt.py):
# * amd64 -> app-<tag>-linux-x64-cuda12-portable.tar.gz (sm_70..sm_120)
# arm64 -> app-<tag>-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
# arm64 -> app-<tag>-linux-arm64-cuda13-portable.tar.gz (sm_90..sm_121)
# * portable bundles carry their own CUDA libs, so they 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=<tag> for a frozen, reproducible build.
# * converter + gguf-py from the SAME release's source tarball (mappings match)
# /opt (not /root) so it survives `docker run --user`; UNSLOTH_LLAMA_CPP_PATH lets
# zoo find it regardless of $HOME. Default "latest" resolves the newest
# unslothai/llama.cpp release at build time; build.sh pins it to a concrete tag so
# the cache busts only on new upstream releases. --build-arg LLAMA_PREBUILT_TAG=<tag>
# for a frozen 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 \
@ -671,20 +520,16 @@ 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).
# * unsloth_nb_pip_magic.py -> site-packages: re-points the IPython `%pip` /
# `%uv` line magics and the `!python -m pip` form at the same shim, so the
# in-process / module install paths cannot bypass PATH and clobber the stack.
# * IPython startup hook: activates the right sidecar before the first model
# cell in manual JupyterLab.
# * unsloth-run: headless `unsloth-run <notebook|url>` that auto-picks the
# sidecar and executes every cell -- the robust driven path.
# UNCHANGED (see unsloth_nb_compat.py). Pieces:
# * unsloth_nb_compat.py -> site-packages: tier detection + sidecar resolution
# + activation + IPython hook.
# * pip/uv shim on a PATH dir AHEAD of the venv bin: makes `!pip install` /
# `!uv pip install` cells safe + idempotent (keeps the baked stack, records
# the requested transformers to activate its sidecar).
# * unsloth_nb_pip_magic.py -> site-packages: re-points `%pip`/`%uv` magics and
# `!python -m pip` at the same shim so in-process installs can't bypass PATH.
# * IPython startup hook: activates the right sidecar before the first model cell.
# * unsloth-run: headless `unsloth-run <notebook|url>`, the robust driven path.
# ---------------------------------------------------------------------------
COPY unsloth_nb_compat.py unsloth_pip_shim.py unsloth_nb_pip_magic.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 \
@ -708,31 +553,22 @@ RUN set -eux \
# 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}
# Load the notebook startup hook (sidecar activation + %pip/%uv magic re-point)
# for EVERY kernel, whatever uid runs it. IPYTHONDIR (inherited by any user via
# ENV) points IPython at this shared profile, so the hook still loads when the
# container is started with `--user <uid>` and $HOME is not /root -- unlike a
# /root/.ipython startup dir, which only a root kernel reads. Kernel-writable
# state (history.sqlite) still lands under each user's own path, so a read-only
# profile dir is fine.
# for EVERY kernel, any uid: IPYTHONDIR (via ENV) points IPython at this shared
# profile, so it loads under `--user <uid>` too -- unlike /root/.ipython, which
# only a root kernel reads. Writable state (history.sqlite) still lands per-user.
ENV IPYTHONDIR=/opt/unsloth-nb/ipython
# 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).
# Pre-clone unslothai/notebooks so JupyterLab opens with them present. Baked as a
# READ-ONLY template (~206MB, .git stripped); on boot the entrypoint copies it to
# /workspace/unsloth-notebooks and best-effort refreshes from GitHub, never
# overwriting a user-touched notebook and skipping rewrites when only the install
# header/announcements/footer moved (see unsloth_sync_notebooks.sh +
# unsloth_nb_content_sig.py). Inherited as-is by the studio image.
#
# UNSLOTH_NOTEBOOKS_REF pins ONE commit/branch/tag so a multi-arch publish bakes
# identical templates into both arch legs even if unslothai/notebooks advances
# mid-build. The publish workflow resolves the live HEAD sha once (like
# LLAMA_PREBUILT_TAG) and passes it here; the default "main" keeps a plain
# `docker build` tracking the tip. We fetch the single resolved ref (a full
# 40-char sha fetches by object; a branch/tag fetches by name) at depth 1, so
# the bake stays a shallow one-commit pull.
# identical templates into both legs. The publish workflow resolves the live HEAD
# sha once (like LLAMA_PREBUILT_TAG); default "main" tracks the tip. Fetched at
# depth 1 (sha by object, branch/tag by name).
ARG UNSLOTH_NOTEBOOKS_REF=main
RUN set -eux \
&& git init -q /opt/unsloth-notebooks \
@ -743,9 +579,8 @@ RUN set -eux \
&& 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.
# JupyterLab lives in the venv. notebooks are pre-populated into
# /workspace/unsloth-notebooks on boot; mount a volume on /workspace to persist.
EXPOSE 8888
COPY smoke_test.py /workspace/smoke_test.py

213
docker/Dockerfile.studio
View file

@ -1,44 +1,30 @@
# Full Unsloth image: base training stack + Studio + JupyterLab + sshd.
#
# This is the image published as docker.io/unsloth/unsloth:studio (and the
# default :latest). It layers Unsloth Studio on top of the lean core image
# (Dockerfile, published under the `core` tags) and runs the same service trio as
# the previous production image: Studio on 8000, JupyterLab on 8888, sshd on 22.
# Published as unsloth/unsloth:studio (and default :latest); layers Studio on the
# lean core image and runs Studio:8000, JupyterLab:8888, sshd:22.
#
# Build (local):
# docker buildx build \
# --build-arg BASE_IMAGE=unsloth-blackwell:test \
# -f docker/Dockerfile.studio \
# -t unsloth-blackwell:studio docker/
#
# docker buildx build --build-arg BASE_IMAGE=unsloth-blackwell:test \
# -f docker/Dockerfile.studio -t unsloth-blackwell:studio docker/
# Run:
# docker run --rm --gpus all -p 8000:8000 -p 8888:8888 \
# -v $HOME/.cache/huggingface:/workspace/.cache/huggingface \
# unsloth-blackwell:studio
# -v $HOME/.cache/huggingface:/workspace/.cache/huggingface unsloth-blackwell:studio
#
# Open http://localhost:8000 for Studio (first-boot admin password is printed
# in the container logs and persisted under /opt/unsloth-studio/auth/) and
# http://localhost:8888 for JupyterLab (password: JUPYTER_PASSWORD env; when
# unset a random one is generated and printed in the container logs). On
# hosts without GPU passthrough (Docker Desktop on macOS, Windows without
# WSL2 GPU) add -e UNSLOTH_ALLOW_CPU=1: training is unavailable but Studio
# chat / Data Recipes / GGUF tooling / Jupyter work.
#
# CI pins BASE_IMAGE to the just-published multi-arch base digest so the two
# images always ship the same stack.
# Studio on :8000 (first-boot admin password in the logs, persisted under
# /opt/unsloth-studio/auth/); JupyterLab on :8888 (JUPYTER_PASSWORD env, else a
# random one is printed). Without GPU passthrough add -e UNSLOTH_ALLOW_CPU=1:
# training is unavailable but Studio chat / Data Recipes / GGUF / Jupyter work.
# CI pins BASE_IMAGE to the published base digest so both images ship the same stack.
ARG BASE_IMAGE=unsloth-blackwell:test
# --- builder stage: prebuild the Unsloth JupyterLab extension -----------------
# Builds the named "Unsloth Dark" (Monokai) theme + the Colab-style Down/Up
# cell-navigation keymap. Node lives ONLY in this throwaway stage; the final
# image copies just the prebuilt static labextension, so the runtime stays
# Node-free. Uses the base image's bundled jlpm + jupyterlab (version-matched).
# Builds the "Unsloth Dark" (Monokai) theme + Colab-style cell-nav keymap. Node
# lives only in this throwaway stage; the final image copies just the prebuilt
# labextension (runtime stays Node-free). Uses the base's bundled jlpm+jupyterlab.
FROM ${BASE_IMAGE} AS labext-builder
ENV DEBIAN_FRONTEND=noninteractive
# JupyterLab 4.6's build tooling declares a Node >=20 engine; Ubuntu 24.04's
# distro nodejs is 18, so pull Node 20 LTS from NodeSource (it bundles npm).
# This stage is thrown away, so the extra apt sources never reach the runtime.
# JupyterLab 4.6 needs Node >=20; Ubuntu 24.04 ships 18, so pull Node 20 LTS from
# NodeSource. This stage is thrown away, so the apt sources never reach runtime.
RUN apt-get update \
&& apt-get install -y --no-install-recommends ca-certificates curl gnupg git \
&& curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \
@ -51,32 +37,23 @@ RUN cd /opt/labext-src \
FROM ${BASE_IMAGE}
# Studio source ref to clone. Defaults to `main`, but a CI publish pipeline
# that pins BASE_IMAGE to a digest should pin this too (same UNSLOTH_REF as
# the base) so the published image is reproducible against a known ref.
# Studio source ref to clone. Defaults to main; CI pins it (same UNSLOTH_REF as
# the base) so the published image is reproducible.
ARG UNSLOTH_STUDIO_REF=main
# unsloth-zoo ref overlaid into the Studio venv by install.sh --local. The
# publish workflow resolves ONE zoo ref and passes it to both the base and
# Studio builds, so the Studio backend runs the same zoo as the base image and
# the operator-requested ref instead of always tracking main.
# unsloth-zoo ref overlaid into the Studio venv by install.sh --local. The publish
# workflow passes ONE zoo ref to both builds, so Studio runs the same zoo as the
# base and the operator-requested ref instead of always main.
ARG UNSLOTH_STUDIO_ZOO_REF=main
# The SAME llama.cpp tag the base image baked (the prepare job resolves it
# once). install.sh -> setup.sh honours UNSLOTH_LLAMA_TAG; without this pin a
# dispatch override, or an upstream release landing between the base and
# Studio jobs, lets the no-GPU Studio build re-resolve "latest" and replace
# the base's pinned CUDA bundle instead of reusing it.
# The SAME llama.cpp tag the base baked (prepare resolves it once). install.sh ->
# setup.sh honours UNSLOTH_LLAMA_TAG; without the pin the Studio build could
# re-resolve "latest" and replace the base's pinned bundle instead of reusing it.
ARG LLAMA_PREBUILT_TAG=latest
ARG TARGETARCH
# Services run as root in this revision (the base image is root-only by
# design); the previous production image ran them as a dedicated uid-1001
# user. Non-root parity is a tracked follow-up. sshd is key-only and stays
# disabled unless a PUBLIC_KEY/SSH_KEY is provided, and no secrets are
# persisted to disk (see studio_launch.sh).
#
# The JUPYTER_PORT / UNSLOTH_ENABLE_SSHD defaults exist so supervisord's
# %(ENV_*)s expansions still resolve when someone bypasses the launcher
# and runs supervisord directly.
# Services run as root here (base is root-only; non-root parity is a follow-up).
# sshd is key-only and stays disabled unless PUBLIC_KEY/SSH_KEY is set; no secrets
# are persisted (see studio_launch.sh). The JUPYTER_PORT / UNSLOTH_ENABLE_SSHD
# defaults let supervisord's %(ENV_*)s resolve when run directly (bypassing the launcher).
USER root
ENV UNSLOTH_STUDIO_HOME=/opt/unsloth-studio \
JUPYTER_PORT=8888 \
@ -91,39 +68,24 @@ RUN apt-get update \
&& rm -rf /var/lib/apt/lists/*
# Clone + install Studio into a dedicated venv under $UNSLOTH_STUDIO_HOME.
# --local makes install.sh use the just-cloned source tree (editable
# install), so the source dir MUST persist for the venv's `unsloth_cli`
# entrypoint to keep resolving. Move it under $UNSLOTH_STUDIO_HOME/src
# (already inside the persistent layer) instead of deleting it. Strip
# .git to save ~120MB.
# --local uses the cloned tree (editable install), so the source MUST persist for
# the venv's unsloth_cli entrypoint -- move it to $STUDIO_HOME/src, strip .git (~120MB).
#
# The llama.cpp symlink BEFORE install.sh points Studio's prebuilt dir at
# the bundle already baked into the base image (validated, sha256-checked,
# UNSLOTH_PREBUILT_INFO.json present), so the installer's prebuilt step
# recognises it and skips a second ~400MB download. The
# .unsloth-studio-owned marker satisfies setup.sh's ownership assertion for
# custom STUDIO_HOMEs -- the dir IS provisioned exclusively for Studio.
# The llama.cpp symlink BEFORE install.sh points Studio's prebuilt dir at the
# base image's baked bundle so the installer skips a second ~400MB download; the
# .unsloth-studio-owned marker satisfies setup.sh's ownership assertion.
#
# UNSLOTH_TORCH_INDEX_FAMILY pins the torch wheel index for the Studio
# venv: at build time there is no GPU and no nvidia-smi, so install.sh's
# probing would land on cpu or cu126 wheels depending on which host built
# the image. cu128 on BOTH arches, mirroring the base venv: cu130 wheels
# would silently lift the arm64 driver floor to 580+ while the base venv
# keeps the documented 570+ floor. Blackwell JIT (amd64 sm_103 B300/GB300 and
# arm64 sm_121 DGX Spark / GB10) support comes from the same NVRTC cu13 swap
# the base image applies to its venv -- repeated below for the Studio venv's
# own bundled libnvrtc, on BOTH arches (the base cu13 layer installed
# cuda-nvrtc-13-0 on both, so the cu13 .so exists here regardless of arch).
# UNSLOTH_TORCH_INDEX_FAMILY pins the Studio venv's torch index: no GPU/nvidia-smi
# at build time would land install.sh on cpu/cu126 wheels. cu128 on both arches,
# mirroring the base (cu130 would lift the arm64 floor to 580+). Blackwell JIT
# (sm_103/sm_121) comes from the same cu13 NVRTC swap the base applies, repeated
# below for the Studio venv on both arches.
#
# UNSLOTH_PYTHON=3.12 pins the Studio venv to the SAME Python minor as the base
# venv (install.sh defaults Linux to 3.13). Matching minors makes the two venvs'
# nvidia-*-cu12 CUDA wheels byte-identical, which lets the dedup RUN further down
# replace the Studio venv's ~3.7GB of CUDA .so with symlinks into the base venv's
# copies (cudnn/cublas/nccl/... are plain C libs, Python-minor independent).
# UNSLOTH_PYTHON=3.12 pins the Studio venv to the base's Python minor (install.sh
# defaults to 3.13), making the nvidia-*-cu12 wheels byte-identical so the dedup
# below can symlink the Studio venv's ~3.7GB of CUDA .so into the base venv's.
#
# fetch+checkout FETCH_HEAD instead of `clone --branch` because the CI
# pipeline passes a commit SHA as the ref (clone --branch only accepts
# branch/tag names).
# fetch+checkout FETCH_HEAD, not `clone --branch`: CI passes a commit SHA.
RUN set -eux \
&& case "${TARGETARCH:-amd64}" in \
amd64|arm64) TORCH_FAMILY="cu128" ;; \
@ -143,33 +105,26 @@ RUN set -eux \
UNSLOTH_LLAMA_TAG="${LLAMA_PREBUILT_TAG}" \
UNSLOTH_PYTHON=3.12 \
bash install.sh --local \
# Fail loud unless the Studio venv torch EXACTLY matches the base venv
# torch (version AND CUDA family) before the dedup below symlinks the two
# venvs' CUDA libs together. A family-only check is not enough: a studio
# install that ignored UNSLOTH_TORCH_INDEX_FAMILY (falling back to
# build-time nvidia-smi probing -> cu126, no sm_100/sm_120 kernels) OR that
# capped torch below the base's version (e.g. 2.10.0+cu128 vs the base's
# 2.11.0+cu128) would slip a mismatched torch past `endswith('+cu128')` and
# make the dedup link incompatible CUDA libs. Comparing to the base venv's
# own torch also avoids hardcoding the version here. metadata check only:
# importing torch needs native libs, which QEMU arm64 builds cannot load.
# Fail loud unless the Studio venv torch EXACTLY matches the base (version AND
# CUDA family) before the dedup symlinks their CUDA libs. A family-only check
# would miss a torch that ignored UNSLOTH_TORCH_INDEX_FAMILY (cu126 probe) or
# capped below the base's version, linking incompatible libs. Compare to the
# base's own torch (no hardcoded version); metadata only, since importing torch
# needs native libs QEMU arm64 can't load.
&& BASE_TORCH="$(/opt/unsloth-venv/bin/python -c "from importlib.metadata import version; print(version('torch'))")" \
&& "${UNSLOTH_STUDIO_HOME}/unsloth_studio/bin/python" -c "import sys; from importlib.metadata import version; assert sys.version_info[:2] == (3, 12), 'Studio venv python %d.%d is not 3.12 (UNSLOTH_PYTHON pin ignored) -- CUDA dedup below depends on it' % sys.version_info[:2]; v = version('torch'); assert v == '${BASE_TORCH}', 'Studio venv torch ' + v + ' does not match base venv torch ${BASE_TORCH} (CUDA dedup would link mismatched libs)'; print('Studio venv python %d.%d torch' % sys.version_info[:2], v, '== base', '${BASE_TORCH}')" \
# setup.sh may relink the root llama-quantize into build/bin; prove the
# relinked quantizer still resolves its libraries, or GGUF export breaks
# at runtime with "No working quantizer found". Content check, not rc:
# llama-quantize exits nonzero on --help, while a loader failure prints
# "error while loading shared libraries" and no usage text.
# setup.sh may relink llama-quantize into build/bin; prove it still resolves
# its libraries or GGUF export breaks with "No working quantizer found".
# Content check, not rc: --help exits nonzero but prints usage; a loader
# failure prints "error while loading shared libraries" and no usage.
&& { "${UNSLOTH_STUDIO_HOME}/llama.cpp/llama-quantize" --help 2>&1 || true; } | grep -q "usage" \
&& rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" \
"${UNSLOTH_STUDIO_HOME}/src/studio/frontend/node_modules" \
/root/.cache \
# Stage the Studio venv's NVRTC exactly like the base venv (see
# docker/Dockerfile: immutable .cu128.orig default + staged .cu13 alias,
# retargeted per device by select_cuda_jit_tools). Run on BOTH arches:
# amd64 sm_103 needs cu13 NVRTC exactly as arm64 sm_121 does, the CUDA
# dedup below never touches cuda_nvrtc, and the base cu13 layer installs
# cuda-nvrtc-13-0 on both arches so libnvrtc.so.13 always exists.
# Stage the Studio venv's NVRTC like the base venv (.cu128.orig default +
# staged .cu13 alias, retargeted per device by select_cuda_jit_tools). Both
# arches: sm_103 needs cu13 NVRTC as much as sm_121, the dedup never touches
# cuda_nvrtc, and the base layer installed cuda-nvrtc-13-0 on both arches.
&& for NVRTC_DIR in "${UNSLOTH_STUDIO_HOME}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do \
if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ] && [ ! -L "${NVRTC_DIR}/libnvrtc.so.12" ]; then \
mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \
@ -203,41 +158,30 @@ RUN set -eux \
COPY supervisord.conf /etc/supervisor/supervisord.conf
COPY studio_launch.sh /usr/local/bin/unsloth-studio-launch
# In-place updaters (no image pull):
# unsloth-studio-update refreshes the Studio packages (backend + baked
# frontend) and restarts the service.
# unsloth-llama-update swaps the baked llama.cpp prebuilt to the latest
# release (the same swap the in-app banner performs).
# unsloth-studio-update refresh Studio packages (backend + frontend) and restart
# unsloth-llama-update swap the baked llama.cpp prebuilt to the latest release
COPY unsloth_studio_update.sh /usr/local/bin/unsloth-studio-update
COPY unsloth_llama_update.sh /usr/local/bin/unsloth-llama-update
# unsloth-llama-update reuses the build-time fetcher (redirect-based, no GitHub
# API, so it is not rate-limited; deterministic portable bundle that runs on CPU
# and every supported GPU) rather than the host-probing installer.
# unsloth-llama-update reuses the build-time fetcher (redirect-based, not rate-
# limited; deterministic portable bundle) rather than the host-probing installer.
COPY fetch_llama_prebuilt.py /usr/local/lib/unsloth/fetch_llama_prebuilt.py
# Optional public Cloudflare tunnel for JupyterLab (UNSLOTH_JUPYTER_CLOUDFLARE=1,
# or `unsloth-jupyter-tunnel --force`); supervisord runs it as jupyter-cloudflare.
COPY unsloth_jupyter_tunnel.sh /usr/local/bin/unsloth-jupyter-tunnel
# JupyterLab defaults baked for every container: the named "Unsloth Dark"
# (Monokai) theme with adaptive light/dark by system preference, a per-cell run
# button that does NOT auto-advance, a labeled "Restart & Run All", windowing
# disabled so collapsing a long output does not snap to the cell top,
# ArrowDown/Up jumping to the TOP of the next/previous cell, and the official
# Jupyter "get notified about news" prompt suppressed (fetchNews/checkForUpdates
# off). overrides.json is the system-wide settings override (read from the base
# venv's share/jupyter/lab/settings); the theme + keymap + Unsloth top-bar logo
# ship as the prebuilt labextension built in the labext-builder stage above.
# JupyterLab defaults baked for every container: "Unsloth Dark" (Monokai) theme
# with adaptive light/dark, a per-cell run button that doesn't auto-advance, a
# labeled "Restart & Run All", windowing off (collapsing output won't snap to
# top), ArrowDown/Up to the top of the next/prev cell, and the "news" prompt off.
# overrides.json is the system-wide settings override; the theme + keymap + logo
# ship as the prebuilt labextension from labext-builder above.
COPY jupyter/overrides.json /opt/unsloth-venv/share/jupyter/lab/settings/overrides.json
COPY --from=labext-builder /opt/labext-src/unsloth-jupyterlab/labextension /opt/unsloth-venv/share/jupyter/labextensions/unsloth-jupyterlab
# Unsloth branding (all served by jupyter_server, so applied to its site-packages
# the same way): replace the browser-tab favicon and the page logo with the
# Unsloth logo, and brand the login screen (dark Unsloth-themed login.html).
# Also disable + lock the stock top-left Jupyter logo plugin so the Unsloth logo
# widget shipped by the labextension is the only one rendered in the top bar
# (lock keeps users from re-enabling it in the UI).
# The sloth-sticker install is the ONLY fail-soft branding step: it is scoped to
# its own { ...; } group with a `|| echo` fallback below, so a missing Studio
# "Sloth emojis" folder does not break the build, while the REQUIRED steps above
# it (JS resolve, favicon/logo/login copy) stay fatal. login.html's onerror falls
# back to the Unsloth logo if the sticker dir is ever absent.
# Unsloth branding (served by jupyter_server, applied to its site-packages):
# replace the favicon + page logo, brand the login screen (login.html), and
# disable+lock the stock top-left Jupyter logo so only the labextension's Unsloth
# logo renders. The sloth-sticker install is the ONLY fail-soft step (own { }
# group with `|| echo`), so a missing "Sloth emojis" folder doesn't break the
# build; the required steps above (JS resolve, favicon/logo/login copy) stay fatal.
COPY jupyter/favicon.ico /tmp/unsloth-branding/favicon.ico
COPY jupyter/logo.png /tmp/unsloth-branding/logo.png
COPY jupyter/login.html /tmp/unsloth-branding/login.html
@ -258,13 +202,11 @@ RUN JS="$(/opt/unsloth-venv/bin/python -c 'import os, jupyter_server; print(os.p
&& /opt/unsloth-venv/bin/jupyter labextension disable @jupyterlab/apputils-extension:splash \
&& /opt/unsloth-venv/bin/jupyter labextension lock @jupyterlab/apputils-extension:splash \
&& /opt/unsloth-venv/bin/jupyter labextension lock unsloth-jupyterlab
# Branding integrity guard: the canonical attribution checker (also a
# jupyter_server extension), the full AGPLv3 license text, and the config that
# enables the extension. Installed into the base venv so it is on the jupyter
# process's import + config search path. The stock @apputils-extension:splash is
# disabled+locked above so the labextension's spinning-logo splash is the sole
# ISplashScreen provider. The build-time --verify FAILS the image build if any
# Unsloth attribution / license asset is missing or altered.
# Branding integrity guard: the attribution checker (a jupyter_server extension),
# the AGPLv3 license text, and its enabling config, installed into the base venv
# (on the jupyter import + config path). The stock splash is disabled+locked
# above so the labextension's splash is the sole provider. --verify FAILS the
# build if any attribution / license asset is missing or altered.
COPY jupyter/unsloth_branding.py /tmp/unsloth-branding-guard/unsloth_branding.py
COPY jupyter/jupyter_server_config.d/unsloth_branding_guard.json /tmp/unsloth-branding-guard/unsloth_branding_guard.json
RUN SP="$(/opt/unsloth-venv/bin/python -c 'import sysconfig; print(sysconfig.get_path("purelib"))')" \
@ -281,8 +223,7 @@ RUN chmod +x /usr/local/bin/unsloth-studio-launch \
/usr/local/bin/unsloth-llama-update \
/usr/local/bin/unsloth-jupyter-tunnel
# Studio web UI, JupyterLab, sshd. All bind 0.0.0.0 inside the container's
# network namespace; the operator publishes them explicitly with -p.
# Studio, JupyterLab, sshd. All bind 0.0.0.0 in the container; publish with -p.
EXPOSE 8000 8888 22
# The base ENTRYPOINT (unsloth-entrypoint) still runs its GPU pre-flight

143
docker/entrypoint.sh
View file

@ -1,70 +1,45 @@
#!/usr/bin/env bash
# Container startup checks for Unsloth.
#
# Fails fast with actionable error messages when the host GPU isn't reachable,
# instead of letting torch crash deep with cryptic CUDA errors. Catches the
# three failure modes that cover ~95% of "it doesn't work" tickets:
#
# 1. nvidia-smi inside the container can't see any GPU
# - User forgot --gpus all
# - Host missing nvidia-container-toolkit
# 2. nvidia-smi works but torch.cuda.is_available() is False
# - Host driver too old for CUDA 12.8
# 3. GPU attaches but is older than Ampere (sm < 80)
# - Unsloth requires sm_80+
#
# Bypass for offline tooling / docs / CI:
# docker run -e UNSLOTH_SKIP_GPU_CHECK=1 ...
# Container startup checks for Unsloth. Fails fast with actionable errors when the
# host GPU isn't reachable, catching the three modes behind ~95% of tickets:
# 1. nvidia-smi sees no GPU (missing --gpus all or nvidia-container-toolkit)
# 2. nvidia-smi works but torch.cuda.is_available() is False (driver too old)
# 3. GPU older than Ampere (sm < 80; Unsloth requires sm_80+)
# Bypass for offline tooling/docs/CI: docker run -e UNSLOTH_SKIP_GPU_CHECK=1 ...
set -euo pipefail
# --- CUDA JIT toolchain selection (device-gated) ----------------------------
# The image bakes CUDA 13 ptxas + NVRTC ONLY so the two Blackwell datacenter
# arches the cu12.8 tools cannot target -- sm_103 (B300 / GB300) and sm_121
# (GB10 / DGX Spark) -- can JIT Triton and torch/NVRTC kernels. Both launched
# AFTER cu12.8, so any host carrying them runs a >= 580 driver, which is exactly
# what a cu13-produced cubin needs to LOAD.
#
# Every OTHER supported arch (Turing..sm_120) works with the bundled cu12.8
# tools and is allowed on a 570-579 driver (the documented floor). A cu13 cubin
# CANNOT load on a 570-579 driver even when it targets an old arch like sm_80
# (CUDA has forward, not backward, driver compatibility across major versions),
# so routing those hosts' JIT through the cu13 tools would break ordinary
# training. ptxas/NVRTC are host-side compilers (they never link libcuda), so
# they RUN under any driver -- it is only their OUTPUT the older driver rejects.
#
# Pick per DEVICE at boot (the compute capability is unknown at build time):
# cu12.8 is the immutable baked default (loadable on every supported 570+
# driver), and only sm_103 / sm_121 -- which ship on >= 580 drivers -- switch
# Triton to cu13 ptxas and retarget the venv NVRTC symlink to the staged cu13
# alias. Runs before every early-exit below so the selection always applies.
# Best-effort: because the safe default needs no write, a non-root / read-only
# rootfs is always fine; only the rare non-root datacenter host cannot switch.
# The image bakes CUDA 13 ptxas + NVRTC only for the two Blackwell datacenter
# arches cu12.8 can't target -- sm_103 (B300/GB300) and sm_121 (GB10/DGX Spark).
# Both launched after cu12.8, so their hosts run a >=580 driver, exactly what a
# cu13 cubin needs to load. Every other arch (Turing..sm_120) uses the cu12.8
# tools on the documented 570-579 floor; a cu13 cubin can't load there (CUDA
# driver compat is forward-only), so routing their JIT through cu13 would break
# training. ptxas/NVRTC are host-side compilers, so they RUN under any driver --
# only their output the old driver rejects.
# Pick per DEVICE at boot (cap unknown at build time): cu12.8 is the immutable
# default (loadable on 570+), only sm_103/sm_121 switch Triton to cu13 ptxas and
# retarget the NVRTC symlink. Runs before every early-exit. Best-effort: the safe
# default needs no write (non-root/read-only fine); only a non-root datacenter
# host can't switch.
select_cuda_jit_tools() {
local caps="" cc nvrtc_dir need_cu13=0
if command -v nvidia-smi >/dev/null 2>&1; then
caps="$( { nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true; } )"
fi
# Scan EVERY visible GPU, not just the first: a Blackwell datacenter part
# (sm_103 B300/GB300 or sm_121 GB10/DGX Spark) can sit behind an H100/B200 in
# the nvidia-smi ordering, so keying off only the first compute_cap would miss
# it. If ANY visible GPU needs cu13, switch to it for the whole process --
# those parts only ship on >= 580 drivers, so the host tolerates cu13 cubins
# for every arch present.
# Scan EVERY visible GPU: a sm_103/sm_121 part can sit behind an H100/B200 in
# nvidia-smi ordering. If ANY needs cu13, switch for the whole process -- those
# parts ship on >=580 drivers, so the host tolerates cu13 cubins for all archs.
while IFS= read -r cc || [[ -n "${cc}" ]]; do
cc="$(printf '%s' "${cc}" | tr -d '[:space:]')"
case "${cc}" in
10.3|12.1) need_cu13=1 ;;
esac
done <<< "${caps}"
# Non-datacenter / undetectable / CPU host: cu12.8 is the immutable baked
# default (libnvrtc.so.12 -> .cu128.orig, Triton on its bundled cu12.8
# ptxas), loadable on every supported 570+ driver, and needs NO write -- so
# a non-root `docker run --user` container is never left on a cu13 NVRTC a
# 570-579 driver cannot load. One exception needs a write: an earlier boot
# of this SAME container on sm_103/sm_121 left libnvrtc.so.12 -> .cu13 in
# the writable layer, and the container now runs on a GPU whose 570-579
# driver cannot load cu13 output -- deterministically reverse exactly that
# selection (best-effort, same non-root caveat as the forward switch).
# Non-datacenter / undetectable / CPU host: keep cu12.8 (libnvrtc.so.12 ->
# .cu128.orig, Triton on bundled cu12.8 ptxas), loadable on 570+ and needs no
# write. One exception needs a write: an earlier boot on sm_103/sm_121 left
# libnvrtc.so.12 -> .cu13 and this GPU's 570-579 driver can't load it --
# reverse that selection (best-effort, same non-root caveat).
if [[ "${need_cu13}" -ne 1 ]]; then
for nvrtc_dir in \
/opt/unsloth-venv/lib/python*/site-packages/nvidia/cuda_nvrtc/lib \
@ -75,12 +50,10 @@ select_cuda_jit_tools() {
done
return 0
fi
# Blackwell datacenter present: cu12.8 cannot emit compute_103/121, so point
# Triton at cu13 ptxas and retarget each venv's libnvrtc.so.12 -> the staged
# cu13 alias. -z guard leaves an explicit `docker run -e TRITON_PTXAS_PATH`
# win. Best-effort: a read-only / --user rootfs that cannot rewrite the
# symlink simply keeps cu12.8 (a rare non-root datacenter case). Covers the
# base venv and, on the Studio image, the Studio venv.
# Blackwell datacenter present: point Triton at cu13 ptxas and retarget each
# venv's libnvrtc.so.12 -> the staged cu13 alias. -z guard lets an explicit
# TRITON_PTXAS_PATH win. Best-effort: a read-only/--user rootfs keeps cu12.8.
# Covers the base venv and the Studio venv.
if [[ -x /usr/local/cuda-13.0/bin/ptxas && -z "${TRITON_PTXAS_PATH:-}" ]]; then
export TRITON_PTXAS_PATH=/usr/local/cuda-13.0/bin/ptxas
fi
@ -94,10 +67,9 @@ select_cuda_jit_tools() {
# Best-effort: never let JIT-tool selection block container startup.
select_cuda_jit_tools || true
# Make the unslothai/notebooks collection available under /workspace before the
# user command runs (JupyterLab, unsloth-run, or a shell). Best-effort: it is
# fully gated by UNSLOTH_SKIP_NOTEBOOK_SYNC and never blocks or fails the
# container (see unsloth_sync_notebooks.sh).
# Make unslothai/notebooks available under /workspace before the user command.
# Best-effort, gated by UNSLOTH_SKIP_NOTEBOOK_SYNC, never blocks the container
# (see unsloth_sync_notebooks.sh).
sync_notebooks() {
if [[ -x /usr/local/bin/unsloth-sync-notebooks ]]; then
/usr/local/bin/unsloth-sync-notebooks || true
@ -112,16 +84,12 @@ fi
err() { printf "\033[1;31mERROR:\033[0m %s\n" "$*" >&2; }
warn() { printf "\033[1;33mWARN:\033[0m %s\n" "$*" >&2; }
# CPU mode for hosts that cannot pass a GPU into a Linux container at all:
# Docker Desktop on macOS (no Metal passthrough), Docker Desktop on Windows
# without WSL2 GPU support, plain CPU Linux boxes, and CI runners. CPU mode
# covers Jupyter, the GGUF tooling and llama.cpp-backed Studio chat (llama.cpp
# runs on CPU), and Data Recipes. It does NOT cover training or loading an
# Unsloth model for chat (FastLanguageModel.from_pretrained runs CUDA probes
# like torch.cuda.get_device_properties and raises without a GPU). With
# UNSLOTH_ALLOW_CPU=1 a missing GPU degrades to a warning instead of the hard
# pre-flight failure; when a GPU IS visible the normal checks below still run so
# a broken GPU setup is not silently ignored.
# CPU mode for hosts that can't pass a GPU into a container (Docker Desktop on
# macOS/Windows-without-WSL2, CPU Linux, CI). Covers Jupyter, GGUF tooling,
# llama.cpp Studio chat and Data Recipes; NOT training or loading an Unsloth
# model (FastLanguageModel runs CUDA probes and raises without a GPU). With
# UNSLOTH_ALLOW_CPU=1 a missing GPU warns instead of failing pre-flight; a
# visible GPU still runs the checks below.
if [[ "${UNSLOTH_ALLOW_CPU:-0}" == "1" ]]; then
if ! command -v nvidia-smi >/dev/null 2>&1 || ! nvidia-smi -L 2>/dev/null | grep -q '^GPU'; then
warn "UNSLOTH_ALLOW_CPU=1 and no GPU visible -- continuing on CPU."
@ -132,11 +100,9 @@ if [[ "${UNSLOTH_ALLOW_CPU:-0}" == "1" ]]; then
fi
fi
# --- Check 1: nvidia-smi present and can enumerate at least one GPU ---------
# nvidia-smi is injected by nvidia-container-toolkit when the container is
# started with a GPU request; it is NOT baked into the image. A missing
# binary therefore means "no GPU was attached", the same failure class as
# an empty -L listing, not a broken image.
# --- Check 1: nvidia-smi present and enumerates at least one GPU ------------
# nvidia-smi is injected by nvidia-container-toolkit on a GPU request, not baked
# in; a missing binary means "no GPU attached", same class as an empty -L.
if ! command -v nvidia-smi >/dev/null 2>&1 || ! nvidia-smi -L 2>/dev/null | grep -q '^GPU'; then
err "No GPU visible inside the container."
cat >&2 <<'MSG'
@ -174,8 +140,7 @@ MSG
fi
# --- Check 2: torch can actually use the GPU --------------------------------
# This catches host-driver-too-old (the GPU enumerates via nvidia-smi but
# the kernel module rejects CUDA contexts).
# Catches host-driver-too-old (nvidia-smi enumerates but CUDA contexts fail).
python - >&2 <<'PY' || exit 1
import sys
import torch
@ -205,8 +170,7 @@ name = torch.cuda.get_device_name(0)
n = torch.cuda.device_count()
print(f"Unsloth container: {n} GPU(s). Primary: {name} sm_{major}{minor} bf16={torch.cuda.is_bf16_supported()}")
# Image targets every current x86_64 NVIDIA arch from Turing onward, per
# https://developer.nvidia.com/cuda/gpus.
# Image targets every current NVIDIA arch from Turing onward.
SUPPORTED = (
("sm_75", "Turing", "T4, RTX 20-series, Quadro RTX"),
("sm_80", "Ampere DC", "A100, A30"),
@ -230,10 +194,9 @@ if major < 8:
print(f"NOTE: {name} is Turing (sm_{major}{minor}) -- bfloat16 is not supported.")
print(" Unsloth will fall back to fp16. Training works but is slightly slower.")
# Secondary devices: the launcher exposes ALL GPUs by default, so on a mixed
# rig an unsupported later device would only surface once a job pins to it or
# a multi-GPU launch fans out. Device 0 stays fatal above; secondaries warn
# now, at startup, while the fix (excluding the device) is still cheap.
# Secondary devices: all GPUs are exposed by default, so an unsupported later
# device only surfaces when a job pins to it. Device 0 is fatal above;
# secondaries warn now while excluding them is still cheap.
for d in range(1, n):
dmaj, dmin = torch.cuda.get_device_capability(d)
if dmaj < 7 or (dmaj == 7 and dmin < 5):
@ -244,12 +207,10 @@ for d in range(1, n):
PY
# --- arm64 note: baked llama.cpp is a CUDA 13 build -------------------------
# Upstream publishes no CUDA 12 arm64 llama.cpp bundle (only arm64-cpu and
# arm64-cuda13), so the arm64 image bakes the cu13 build while the torch stack
# (cu128) runs fine on a 570-series driver. A CUDA 13 cubin cannot load on a
# 570-579 driver, so on GH200/GB200-class hosts below 580 GGUF export and
# Studio chat would fail even though training works -- say so up front instead
# of letting llama-server fail mysteriously later.
# Upstream ships no CUDA 12 arm64 llama.cpp (only arm64-cpu/arm64-cuda13), so the
# arm64 image bakes cu13 while the torch stack (cu128) runs on 570+. A cu13 cubin
# can't load on 570-579, so below 580 GGUF export / Studio chat fail even though
# training works -- say so up front instead of failing mysteriously later.
if [ "$(uname -m)" = "aarch64" ]; then
_drv="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1)"
_drv_major="${_drv%%.*}"

46
docker/fetch_llama_prebuilt.py
View file

@ -45,9 +45,8 @@ RELEASE_REPO = "unslothai/llama.cpp"
def resolve_latest_tag(repo: str) -> str:
# Follow the /releases/latest redirect to /releases/tag/<tag>. This needs no
# API token and is not subject to the GitHub API rate limit, so it works on
# any build host (CI, laptop, B200) without configuration.
# Follow the /releases/latest redirect: no API token, no rate limit, works on
# any build host.
url = f"https://github.com/{repo}/releases/latest"
request = urllib.request.Request(url, headers = {"User-Agent": "unsloth-docker-build"})
with urllib.request.urlopen(request, timeout = 60) as response:
@ -149,18 +148,12 @@ def main() -> None:
if os.path.isdir(conversion):
shutil.copytree(conversion, os.path.join(install_dir, "conversion"), dirs_exist_ok = True)
# Make the baked marker readable by Studio's llama.cpp freshness check
# (utils.llama_cpp_freshness.check_prebuilt_freshness) so the in-app
# "newer llama.cpp available" banner works inside the Docker image.
# The release tarball's UNSLOTH_PREBUILT_INFO.json carries upstream_tag /
# source_repo, but the freshness reader keys off tag / release_tag /
# published_repo -- the schema Studio's install_llama_prebuilt.py writes,
# which the image bypasses by baking the bundle directly. Without these
# keys the freshness check bails and can never report "behind", so the
# banner stays hidden even when a newer release exists. setdefault() so a
# future tarball that already ships these keys is left untouched, and we
# add no build timestamp -- behind/update_available do not need one, and
# omitting it keeps the layer byte-identical across build hosts.
# Make the baked marker readable by Studio's freshness check so the in-app
# "newer llama.cpp available" banner works. The tarball's marker carries
# upstream_tag/source_repo, but the reader keys off tag/release_tag/
# published_repo (the schema install_llama_prebuilt.py writes). setdefault()
# leaves a future tarball that already has these keys untouched; no build
# timestamp, so the layer stays byte-identical across build hosts.
marker_path = os.path.join(install_dir, "UNSLOTH_PREBUILT_INFO.json")
try:
with open(marker_path) as f:
@ -175,15 +168,12 @@ def main() -> None:
f.write("\n")
print(f"marker augmented for freshness: tag={tag} published_repo={RELEASE_REPO}")
# Mirror the install into build/bin/ via hardlinks (zero extra bytes).
# Studio's setup.sh treats an executable build/bin/llama-server +
# build/bin/llama-quantize as a complete local build and skips its
# source-build fallback -- which would otherwise fire inside the image
# build, where the host-probing prebuilt updater cannot succeed, and
# compile a CPU-only llama.cpp over the baked CUDA bundle. Hardlinks
# (not symlinks) keep $ORIGIN rpath resolution working from build/bin
# and avoid a cycle when setup.sh later relinks the root quantizer to
# build/bin/llama-quantize.
# Mirror the install into build/bin/ via hardlinks (zero extra bytes). Studio's
# setup.sh treats executable build/bin/llama-server + llama-quantize as a
# complete local build and skips its source-build fallback (which would
# otherwise compile a CPU-only llama.cpp over the baked CUDA bundle). Hardlinks
# (not symlinks) keep $ORIGIN rpath resolution and avoid a cycle when setup.sh
# relinks the root quantizer to build/bin/llama-quantize.
build_bin = os.path.join(install_dir, "build", "bin")
os.makedirs(build_bin, exist_ok = True)
for entry in os.listdir(install_dir):
@ -202,11 +192,9 @@ def main() -> None:
if "/" not in target and not os.path.lexists(dest):
os.symlink(target, dest)
# Sanity: the server binary must execute on a GPU-less host (the CUDA
# backend is a dlopen'd plugin, so --version works anywhere). Check the
# quantizer from BOTH roots: Studio's setup.sh relinks the root
# llama-quantize to build/bin/llama-quantize, so the build/bin copy must
# resolve its libraries standalone.
# Sanity: the server must execute on a GPU-less host (the CUDA backend is a
# dlopen'd plugin). Check the quantizer from BOTH roots: setup.sh relinks the
# root llama-quantize to build/bin, so the build/bin copy must resolve standalone.
checks = (
# llama-quantize has no --version; a healthy run prints usage with
# rc 0, while a loader failure prints to stderr with rc 127.

15
docker/jupyter/unsloth_branding.py
View file

@ -80,10 +80,9 @@ def resolve_paths(
jupyter_server_dir = os.path.dirname(jupyter_server.__file__)
labext_dir = os.path.join(venv_share, "labextensions", LABEXT_NAME)
# Every page_config.json JupyterLab merges to compute disabledExtensions: the
# app-settings file plus a labconfig/ file under each jupyter config dir
# (where `jupyter labextension disable` writes). Tests pass config_dirs=[] for
# a hermetic tree; live resolution scans the real jupyter config path.
# Every page_config.json JupyterLab merges for disabledExtensions: the
# app-settings file plus a labconfig/ file per jupyter config dir. Tests pass
# config_dirs=[] for a hermetic tree; live resolution scans the real path.
if config_dirs is None:
try:
from jupyter_core.paths import jupyter_config_path
@ -201,11 +200,9 @@ def verify_branding(paths = None):
problems.append("missing or empty logo: " + paths["logo"])
# 7. No page_config.json disables the Unsloth extension or its plugins.
# Disabling via `disabledExtensions` leaves the static bundle on disk (so
# check 5 still passes) yet strips the logo / About / splash at load. Since
# the guard exists to refuse stripped attribution, reject that too. Stock
# plugins we disable ourselves (logo/splash) are unaffected -- we only flag
# ids belonging to unsloth-jupyterlab.
# Disabling via disabledExtensions leaves the bundle on disk (check 5 passes)
# yet strips the logo/About/splash at load, so reject it too. We only flag
# ids belonging to unsloth-jupyterlab (our own stock disables are fine).
for pc_path in paths.get("page_configs", []):
text = _read(pc_path)
if not text:

20
docker/jupyter/unsloth_labext/src/cellNav.ts
View file

@ -8,20 +8,16 @@ import {
import { INotebookTracker } from '@jupyterlab/notebook';
/**
* Colab-style cell navigation that works in BOTH command and edit mode.
* Colab-style cell navigation in BOTH command and edit mode.
*
* Pressing ArrowDown on the last line of a cell (edit mode) or while a cell is
* selected (command mode) moves to the next cell and aligns its TOP to the
* viewport; ArrowUp is the mirror. JupyterLab's built-in selection scroll uses
* `scrollIntoViewIfNeeded`, which CENTERS any cell taller than the viewport --
* so moving onto a cell with a long output (e.g. `trainer.train()`) drops the
* view in the middle of the output instead of at the cell top.
* ArrowDown on a cell's last line (edit) or while selected (command) moves to the
* next cell and aligns its TOP to the viewport; ArrowUp mirrors it. JupyterLab's
* built-in scroll CENTERS cells taller than the viewport, dropping the view in
* the middle of a long output (e.g. `trainer.train()`).
*
* Settings cannot fix this: since JupyterLab 4.1 the editor handles keydown in
* the bubbling phase, and the command-mode arrows are owned by Lumino. So we
* listen in the CAPTURE phase (before CodeMirror or Lumino see the key), decide
* whether we are at a cell boundary, and when we are we move the active cell and
* scroll its top into view ourselves.
* Settings can't fix this (JupyterLab 4.1 handles keydown in the bubbling phase,
* command-mode arrows are Lumino's), so we listen in the CAPTURE phase, detect a
* cell boundary, and move + scroll-to-top ourselves.
*/
const cellNavPlugin: JupyterFrontEndPlugin<void> = {
id: 'unsloth-jupyterlab:cell-nav',

16
docker/jupyter/unsloth_labext/src/colabTitle.ts
View file

@ -10,17 +10,11 @@ import { Cell } from '@jupyterlab/cells';
/**
* Colab "#@title" form cells. In Colab a code cell whose first line is
* `#@title Some Title` renders as a titled, collapsed form: the title shows as a
* clickable header, the code is hidden by default ("Show code"), and the output
* stays visible. JupyterLab has no equivalent, so this plugin reproduces it.
*
* For each code cell whose first line matches `#@title <text>` we inject a small
* clickable title bar at the top of the cell and hide the cell input by default
* via a CSS class on the cell node (we toggle visibility with CSS rather than
* the model's source_hidden so we never mutate/persist notebook metadata and the
* output area is untouched). Clicking the bar shows/hides the code. Windowing is
* disabled image-wide (overrides.json), so cell nodes are stable and the
* injected bar persists.
* `#@title Some Title` renders as a titled, collapsed form (clickable header,
* code hidden by default, output visible). JupyterLab has no equivalent, so this
* reproduces it: inject a clickable title bar and hide the input via a CSS class
* (not the model's source_hidden, so notebook metadata is never mutated).
* Clicking toggles the code. Windowing is disabled image-wide, so the bar persists.
*/
const TITLE_RE = /^\s*#\s*@title\b[ \t]*(.*)$/;

10
docker/jupyter/unsloth_labext/src/index.ts
View file

@ -40,12 +40,10 @@ const themePlugin: JupyterFrontEndPlugin<void> = {
};
/**
* Replace the top-left Jupyter logo with the Unsloth logo. The stock
* `@jupyterlab/application-extension:logo` plugin is disabled + locked at image
* build time (jupyter labextension disable/lock), so this is the only logo
* widget added to the top bar. We render an <img> with inline styles rather than
* a LabIcon/CSS so the branding shows identically regardless of the active theme
* (the theme CSS is only loaded while Unsloth Dark is selected).
* Replace the top-left Jupyter logo with the Unsloth logo. The stock logo plugin
* is disabled + locked at build time, so this is the only logo widget. Rendered
* as an <img> with inline styles (not a LabIcon/CSS) so branding shows identically
* in any theme (the theme CSS loads only while Unsloth Dark is selected).
*/
const logoPlugin: JupyterFrontEndPlugin<void> = {
id: 'unsloth-jupyterlab:logo',

39
docker/jupyter/unsloth_labext/src/outputSelect.ts
View file

@ -9,29 +9,16 @@ import {
/**
* Colab-style Ctrl/Cmd+A inside a cell output.
*
* In JupyterLab, clicking a cell's output leaves the notebook in command mode
* (an output area is not an editor), so Ctrl/Cmd+A fires `notebook:select-all`
* which selects EVERY cell in the notebook. On a large notebook that is both
* surprising and laggy. Colab instead selects only the text of the output you
* clicked. This plugin reproduces that: when the keystroke originates from
* within an output area we select just that output's text and stop the event so
* the notebook-wide select-all command never runs.
* Clicking a cell's output leaves the notebook in command mode, so Ctrl/Cmd+A
* fires `notebook:select-all` (selects EVERY cell). Colab instead selects only
* the clicked output's text; this reproduces that and stops the event so the
* notebook-wide select-all never runs.
*
* We listen in the CAPTURE phase (before Lumino's command keybindings) and only
* act when:
* - the chord is exactly Ctrl/Cmd+A (no Alt; Shift ignored), and
* - focus is NOT in a text editor / input / contenteditable (so editing a
* code cell with Ctrl+A still selects within that editor), and
* - the keystroke target OR the last pointer-down landed inside an output area.
*
* We deliberately do NOT use the text selection anchor to decide ownership: a
* stale selection inside an output survives a later click onto a command-mode
* cell or the file browser (clicking a non-text region does not always move the
* anchor), which would make Ctrl/Cmd+A keep re-selecting that old output instead
* of doing the normal select-all in the new context. The last pointer-down is
* reset on every click (to null when the click is outside any output), so it
* tracks the user's current intent; in every other case we do nothing and
* JupyterLab keeps its default behaviour.
* Listens in the CAPTURE phase and acts only when the chord is exactly Ctrl/Cmd+A
* (no Alt), focus is NOT in an editor/input/contenteditable, and the keystroke
* target or last pointer-down landed in an output area. We use the last
* pointer-down, not the text selection anchor, because a stale anchor survives a
* click away and would hijack select-all elsewhere.
*/
// Output containers, widest first. `.jp-OutputArea-output` is a single output;
@ -102,11 +89,9 @@ const outputSelectPlugin: JupyterFrontEndPlugin<void> = {
if (inEditableContext()) {
return;
}
// Own the chord only when the user is actually in an output right now:
// the keystroke target, else the last place they clicked. We do NOT trust
// the text selection anchor -- it goes stale after clicking away from a
// previously selected output (see the file header), which would otherwise
// hijack select-all in the notebook / file browser.
// Own the chord only when in an output now: the keystroke target, else the
// last click. Not the selection anchor -- it goes stale after clicking away
// (see the header) and would hijack select-all elsewhere.
const output =
closestOutput(event.target as Node | null) ?? lastPointerOutput;
if (!output) {

11
docker/jupyter/unsloth_labext/src/uiChrome.ts
View file

@ -10,13 +10,10 @@ import {
/**
* Colab-like chrome tweaks applied image-wide.
*
* Hide the right activity bar (the vertical strip that carries the Property
* Inspector / Debugger tabs) by default. JupyterLab has no settings key to hide
* a side activity bar outright -- `@jupyterlab/application-extension:shell` only
* exposes `activityBarPosition` (move it) and `layout` (reposition widgets) --
* so we hide the strip with always-on CSS (independent of the active theme) and
* collapse the right panel once on startup. Panels can still be reopened from
* the View menu / command palette; nothing is removed, only hidden by default.
* Hide the right activity bar (Property Inspector / Debugger tabs) by default.
* JupyterLab has no settings key to hide a side activity bar, so hide the strip
* with always-on CSS and collapse the right panel once on startup. Panels can
* still be reopened from the View menu; nothing is removed, only hidden.
*/
const STYLE_ID = 'unsloth-ui-chrome-style';

63
docker/run.sh
View file

@ -1,19 +1,11 @@
#!/usr/bin/env bash
# Convenience wrapper for `docker run unsloth/unsloth`. Sets the flags that
# people most often forget and that cause the most confusing failures:
#
# --gpus all Without this, no GPU is attached and the container's
# entrypoint will refuse to start.
# --ipc=host PyTorch DataLoader workers need ample /dev/shm. The
# default 64MB causes "DataLoader worker (pid X) exited
# unexpectedly" on any non-trivial dataset.
# --ulimit memlock=-1 Unlimited pinned memory for NCCL / CUDA pinned host
# buffers. Without this, multi-GPU training stalls.
# --ulimit stack=64MB Larger thread stack for libtorch (some kernels OOM
# the default 8MB stack).
#
# Plus mounts the host Hugging Face cache and Triton JIT cache so model
# downloads and compiled kernels persist across container runs.
# Convenience wrapper for `docker run unsloth/unsloth`. Sets the easily-forgotten
# flags behind the most confusing failures:
# --gpus all attach a GPU (entrypoint refuses to start without one)
# --ipc=host ample /dev/shm; the default 64MB crashes DataLoader workers
# --ulimit memlock=-1 unlimited pinned memory (else multi-GPU training stalls)
# --ulimit stack=64MB larger libtorch thread stack (some kernels OOM the 8MB default)
# Plus mounts the host HF + Triton caches so downloads and kernels persist.
#
# Usage:
# bash docker/run.sh # interactive python REPL
@ -50,12 +42,10 @@ set -euo pipefail
IMAGE="${UNSLOTH_IMAGE:-unsloth/unsloth:latest}"
GPUS="${UNSLOTH_GPUS:-all}"
# Translate index selectors to Docker's `device=` form. The header docstring
# advertises UNSLOTH_GPUS values like "0" and "0,1" but Docker reads a bare
# integer for --gpus as a COUNT, not an INDEX, so `UNSLOTH_GPUS=0` would
# expose zero GPUs and the entrypoint would refuse to start. `all` and
# already-quoted `device=...` / `"device=..."` selectors pass through.
# "none" omits --gpus entirely (CPU mode; pair with UNSLOTH_ALLOW_CPU=1).
# Translate index selectors to Docker's `device=` form: Docker reads a bare
# integer for --gpus as a COUNT not an INDEX, so `UNSLOTH_GPUS=0` would expose
# zero GPUs. `all` and already-quoted `device=...` selectors pass through;
# "none" omits --gpus (CPU mode; pair with UNSLOTH_ALLOW_CPU=1).
GPU_FLAG=(--gpus "$GPUS")
case "$GPUS" in
none) GPU_FLAG=() ;;
@ -72,33 +62,26 @@ WORK_DIR="${UNSLOTH_WORKDIR:-$PWD}"
mkdir -p "$HF_CACHE" "$TRITON_CACHE"
# Warn early if the host doesn't have the nvidia runtime registered.
# We let `docker run` fail loudly rather than abort here -- some setups
# (rootless docker, custom runtimes) report runtimes differently.
# Warn early if the host has no nvidia runtime registered. Let `docker run` fail
# loudly rather than abort -- some setups report runtimes differently.
if ! docker info 2>/dev/null | grep -qi 'Runtimes:.*nvidia'; then
printf "\033[1;33mWARN:\033[0m 'docker info' does not list 'nvidia' as a runtime.\n" >&2
printf " If --gpus all fails below, install nvidia-container-toolkit:\n" >&2
printf " https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html\n\n" >&2
fi
# Forward common secrets only if they're set in the host environment.
# Empty strings would shadow whatever is already inside the image.
# IMPORTANT: use the dash-only form `-e VAR` (no `=VALUE`). Docker reads
# the value from the parent shell, so the literal secret never lands in
# argv where it would be visible to any user on the host via
# `ps auxe` / `/proc/<pid>/cmdline` for the lifetime of the docker CLI
# process.
# Forward common secrets only if set (empty strings would shadow the image's).
# Use the dash-only `-e VAR` form: Docker reads the value from the parent shell,
# so the secret never lands in argv (visible via `ps auxe` / /proc/<pid>/cmdline).
declare -a ENV_FORWARD=(-e HF_HUB_ENABLE_HF_TRANSFER=1)
[[ -n "${HF_TOKEN:-}" ]] && ENV_FORWARD+=(-e HF_TOKEN)
[[ -n "${WANDB_API_KEY:-}" ]] && ENV_FORWARD+=(-e WANDB_API_KEY)
[[ -n "${UNSLOTH_LICENSE:-}" ]] && ENV_FORWARD+=(-e UNSLOTH_LICENSE)
[[ -n "${UNSLOTH_ALLOW_CPU:-}" ]] && ENV_FORWARD+=(-e UNSLOTH_ALLOW_CPU)
# Studio/Jupyter service config read by studio_launch.sh. Same dash-only -e VAR
# form as the secrets above: the value comes from the parent env, so even
# JUPYTER_PASSWORD never lands in argv (ps auxe / /proc/<pid>/cmdline). Without
# these, `JUPYTER_PASSWORD=... bash docker/run.sh` silently got a random
# password, PUBLIC_KEY/SSH_KEY never enabled sshd, and UNSLOTH_JUPYTER_CLOUDFLARE
# never started the tunnel when using the bundled launcher.
# form so even JUPYTER_PASSWORD never lands in argv. Without these, the bundled
# launcher got a random password and never enabled sshd (PUBLIC_KEY/SSH_KEY) or
# the tunnel (UNSLOTH_JUPYTER_CLOUDFLARE).
[[ -n "${JUPYTER_PASSWORD:-}" ]] && ENV_FORWARD+=(-e JUPYTER_PASSWORD)
[[ -n "${PUBLIC_KEY:-}" ]] && ENV_FORWARD+=(-e PUBLIC_KEY)
[[ -n "${SSH_KEY:-}" ]] && ENV_FORWARD+=(-e SSH_KEY)
@ -118,11 +101,9 @@ if [ -t 0 ] && [ -t 1 ]; then
TTY_FLAG=(-it)
fi
# Avoid `set -x` here so the literal HF_TOKEN / WANDB_API_KEY / UNSLOTH_LICENSE
# values do not get echoed to stdout/CI logs. The forwarded env vars are
# already in ENV_FORWARD; printing them again was a secret leak.
# The ${arr[@]+"${arr[@]}"} form keeps empty arrays nounset-safe on
# bash 3.2 (macOS /bin/bash), where a bare "${empty[@]}" trips set -u.
# No `set -x` here: it would echo HF_TOKEN / WANDB_API_KEY / UNSLOTH_LICENSE to
# CI logs. The ${arr[@]+"${arr[@]}"} form keeps empty arrays nounset-safe on
# bash 3.2 (macOS), where a bare "${empty[@]}" trips set -u.
exec docker run --rm ${TTY_FLAG[@]+"${TTY_FLAG[@]}"} \
${GPU_FLAG[@]+"${GPU_FLAG[@]}"} \
--ipc=host \

22
docker/smoke_test.py
View file

@ -45,11 +45,9 @@ def check_torch() -> tuple[int, int]:
cap = torch.cuda.get_device_capability(0)
name = torch.cuda.get_device_name(0)
print(f"device 0 {name} sm_{cap[0]}{cap[1]}")
# The cu128 wheels ship SASS down to sm_75 (Turing), and the runtime
# entrypoint allows the same floor. Match here so the post-publish
# smoke job does not false-fail on a Turing-only self-hosted runner.
# Turing falls back to fp16 since bf16 isn't supported -- that's a
# capability hint, not a hard failure.
# cu128 wheels ship SASS down to sm_75 (Turing); match the runtime entrypoint's
# floor so the smoke job doesn't false-fail on a Turing-only runner. Turing
# falls back to fp16 (a capability hint, not a hard failure).
if cap[0] < 7 or (cap[0] == 7 and cap[1] < 5):
sys.exit(f"FAIL: pre-Turing GPU {name} is not supported by this image")
if cap[0] < 8:
@ -62,21 +60,17 @@ def check_imports() -> None:
import triton
print(f"triton {triton.__version__}")
# Import order matters: unsloth must be imported BEFORE transformers / trl /
# peft so its monkey-patches land, and BEFORE unsloth_zoo so the latter sees
# the UNSLOTH_IS_PRESENT env marker that unsloth/__init__.py sets. Doing it
# otherwise trips an explicit guard in unsloth_zoo/__init__.py with
# "ImportError: Please install Unsloth via `pip install unsloth`!".
# Import order matters: unsloth BEFORE transformers/trl/peft (so its patches
# land) and BEFORE unsloth_zoo (which needs the UNSLOTH_IS_PRESENT marker,
# else its __init__ guard raises "Please install Unsloth via pip install unsloth").
import unsloth
print(f"unsloth {unsloth.__version__}")
import unsloth_zoo
print(f"unsloth_zoo {unsloth_zoo.__version__}")
# xformers is not built for aarch64 cu128 as of this writing; the arm64
# variant of this image installs unsloth with `[huggingface]` extras
# which omits it. Treat the import as best-effort so the same script
# smoke-tests both arches.
# xformers has no aarch64 cu128 wheel, so the arm64 image omits it
# ([huggingface] extras). Best-effort import so one script covers both arches.
try:
import xformers
print(f"xformers {xformers.__version__}")

43
docker/studio_launch.sh
View file

@ -21,12 +21,9 @@ export UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}"
# resolves; set to 1 (docker run -e) to expose JupyterLab on a trycloudflare URL.
export UNSLOTH_JUPYTER_CLOUDFLARE="${UNSLOTH_JUPYTER_CLOUDFLARE:-0}"
# Make the runtime env visible to SSH sessions, which get a fresh login shell
# without the `docker run -e` vars. Secrets are excluded on purpose: tokens,
# API keys and passwords stay in process env only, never on disk where an
# SSH session (or anything reading /etc/profile.d) could pick them up.
# shlex.quote() each value: env vars can contain quotes, $, backticks etc,
# and this file is sourced by every login shell.
# Make the runtime env visible to SSH login shells (which lack the `docker run -e`
# vars). Secrets are excluded on purpose -- they stay in process env, never on
# disk. shlex.quote() each value since this file is sourced by every login shell.
python - > /etc/profile.d/unsloth_env.sh <<'PY' || true
import os, re, shlex
keep = re.compile(r"^(HF_|CUDA_|NCCL_|JUPYTER_|UNSLOTH_|WANDB_|TRITON_)|^PATH$")
@ -37,9 +34,8 @@ for key, value in sorted(os.environ.items()):
PY
# --- Jupyter -----------------------------------------------------------------
# Hash the password with jupyter's own helper; never store the plaintext.
# No fixed default password: when JUPYTER_PASSWORD is unset we generate a
# random one and print it once in the boot banner (docker logs).
# Hash the password with jupyter's helper; never store plaintext. No fixed
# default: when JUPYTER_PASSWORD is unset, generate a random one and print it once.
JUPYTER_CONFIG_DIR=/root/.jupyter
JUPYTER_NOTE="password from JUPYTER_PASSWORD env"
if [[ -f "${JUPYTER_CONFIG_DIR}/jupyter_lab_config.py" ]]; then
@ -63,23 +59,20 @@ c.ServerApp.open_browser = False
c.ServerApp.root_dir = "/workspace"
c.PasswordIdentityProvider.hashed_password = "${HASH}"
EOF
# Land straight in the categorized notebook view, but only when it is enabled
# AND lives under root_dir (so it is expressible as a /lab/tree path). Mirror
# unsloth_sync_notebooks.sh's gating -- UNSLOTH_NOTEBOOKS_VIEW_DIR plus both
# UNSLOTH_SKIP_NOTEBOOK_VIEW (no view built) and UNSLOTH_SKIP_NOTEBOOK_SYNC
# (entrypoint skips sync entirely, so nothing under the view dir exists) -- so
# a relocated, disabled, or unsynced view never points JupyterLab at a missing
# dir; in those cases JupyterLab just opens on its default (/lab) over /workspace.
# Land in the categorized notebook view, but only when it's enabled AND under
# root_dir (expressible as /lab/tree). Mirror unsloth_sync_notebooks.sh's
# gating (UNSLOTH_NOTEBOOKS_VIEW_DIR + SKIP_NOTEBOOK_VIEW + SKIP_NOTEBOOK_SYNC)
# so a relocated/disabled/unsynced view never points at a missing dir;
# otherwise JupyterLab opens on its default /lab over /workspace.
_root_dir="/workspace"
_view_dir="${UNSLOTH_NOTEBOOKS_VIEW_DIR:-/workspace/Unsloth Notebooks}"
if [[ "${UNSLOTH_SKIP_NOTEBOOK_VIEW:-0}" != "1" \
&& "${UNSLOTH_SKIP_NOTEBOOK_SYNC:-0}" != "1" \
&& "${_view_dir}" == "${_root_dir}/"* ]]; then
_view_rel="${_view_dir#${_root_dir}/}"
# default_url must be set on BOTH ServerApp and LabApp -- the lab
# extension app otherwise overrides ServerApp's value back to /lab.
# preferred_dir points the file browser at that folder. A literal space
# is URL-encoded to %20 in the redirect itself.
# default_url must be set on BOTH ServerApp and LabApp (the lab app
# otherwise overrides ServerApp back to /lab). preferred_dir points the
# file browser at that folder; a literal space is URL-encoded to %20.
cat >> "${JUPYTER_CONFIG_DIR}/jupyter_lab_config.py" <<EOF
c.ServerApp.default_url = "/lab/tree/${_view_rel}"
c.LabApp.default_url = "/lab/tree/${_view_rel}"
@ -105,12 +98,10 @@ fi
mkdir -p /workspace
# --- Branding / AGPLv3 attribution integrity gate (whole container) -----------
# This image is built by Unsloth and ships under the GNU AGPLv3. Refuse to start
# anything if the Unsloth attribution (Help/About, spinning-logo splash, branded
# login, theme, AGPLv3 license text + source links) has been stripped or altered.
# The same checker also runs as a jupyter_server extension (refuses JupyterLab on
# its own) and at image-build time. Bypass for local development if ever needed:
# UNSLOTH_SKIP_BRANDING_CHECK=1 (intended for Unsloth's own debugging, not resale).
# This image ships under the GNU AGPLv3. Refuse to start if the Unsloth
# attribution (Help/About, splash, login, theme, AGPLv3 license + source links)
# is stripped or altered. The same checker runs as a jupyter_server extension and
# at build time. Bypass for local dev: UNSLOTH_SKIP_BRANDING_CHECK=1 (not resale).
if [[ "${UNSLOTH_SKIP_BRANDING_CHECK:-0}" != "1" ]]; then
if ! /opt/unsloth-venv/bin/python -m unsloth_branding --verify; then
echo "Refusing to start the container." >&2

11
docker/unsloth_ipython_startup.py
View file

@ -17,13 +17,10 @@ try:
# the safe-install behaviour. Unset everywhere else => shim is a passthrough.
os.environ["UNSLOTH_NB_SHIM"] = "1"
# Scope the transformers-request marker to THIS kernel so two notebooks
# running concurrently in the same container (each its own kernel process)
# do not read each other's pin. The pip/uv shim runs as a child of this
# kernel and inherits UNSLOTH_NB_TF_MARKER, so writer (shim) and reader
# (unsloth_nb_compat pre_run_cell hook, same process tree) agree on the
# path. Falls back to the shared default when unset (e.g. `unsloth-run`,
# which drives a single notebook per process).
# Scope the transformers-request marker to THIS kernel so concurrent notebooks
# don't read each other's pin. The pip/uv shim (a child of this kernel)
# inherits UNSLOTH_NB_TF_MARKER, so writer and reader agree on the path. Falls
# back to the shared default when unset (e.g. `unsloth-run`, one notebook/process).
if not os.environ.get("UNSLOTH_NB_TF_MARKER"):
# A kernel id that is stable for the kernel's lifetime and unique per
# kernel: the ipykernel connection file name, else the kernel PID.

21
docker/unsloth_llama_update.sh
View file

@ -99,12 +99,10 @@ fi
# an atomic rename), then swap. On any failure the existing install is untouched.
parent="$(dirname "$INSTALL_DIR")"
# The documented persistence recipe mounts a named volume AT the install dir
# (-v unsloth_llama:/opt/unsloth/llama.cpp). A mount point cannot be renamed --
# rename(2) fails EBUSY -- so the whole-dir swap below would always fail there.
# Detect the mount and swap the CONTENTS inside the mounted tree instead, which
# also keeps the update IN the volume (persistent across a recreate).
# UNSLOTH_LLAMA_UPDATE_IN_PLACE=1/0 overrides the autodetection.
# The persistence recipe mounts a named volume AT the install dir. A mount point
# can't be renamed (rename(2) EBUSY), so the whole-dir swap below would fail
# there; detect the mount and swap the CONTENTS inside the tree (also keeps the
# update in the volume). UNSLOTH_LLAMA_UPDATE_IN_PLACE=1/0 overrides autodetection.
IN_PLACE="${UNSLOTH_LLAMA_UPDATE_IN_PLACE:-}"
if [ -z "$IN_PLACE" ]; then
IN_PLACE=0
@ -125,17 +123,14 @@ else
fi
swap_done=0
# The exit handler must never delete $backup while it is the ONLY copy of the
# install (signal between the two renames, or a failed swap whose restore also
# failed): put the old tree back first, and remove it only after the new tree
# is verifiably active. The signal traps make bash run the EXIT trap on
# HUP/INT/TERM too.
# install: put the old tree back first, and remove it only after the new tree is
# verifiably active. The signal traps run the EXIT trap on HUP/INT/TERM too.
cleanup() {
if [ "$swap_done" -ne 1 ]; then
if [ "$IN_PLACE" = "1" ]; then
# Contents-swap restore. Every old entry lives in exactly one of
# $backup / $INSTALL_DIR, so a same-named entry in the install dir
# can only be a half-moved NEW one: drop it, then move the old one
# back. Never deletes anything that is not shadowed by the backup.
# $backup / $INSTALL_DIR, so a same-named entry in the install dir is a
# half-moved NEW one: drop it, then move the old one back.
if [ -d "$backup" ]; then
_restore_fail=0
for _e in "$backup"/* "$backup"/.[!.]* "$backup"/..?*; do

10
docker/unsloth_nb_content_sig.py
View file

@ -48,12 +48,10 @@ def _is_install_code(cell):
low = t.lower()
if any(m in low for m in _INSTALL_MARKERS):
return True
# A %%capture / %%bash cell is boilerplate ONLY when it also carries an
# install command. A bare %%capture (e.g. wrapping training to silence
# output) or a %%bash cell doing real tutorial setup is substantive: hashing
# it keeps the boot refresh from silently skipping an upstream fix to that
# cell (a false SAME). The install markers above already catch the generated
# install cell, which begins with %%capture.
# A %%capture / %%bash cell is boilerplate ONLY when it also carries an install
# command. A bare %%capture or a %%bash doing real setup is substantive: hash
# it so the boot refresh doesn't skip an upstream fix (a false SAME). The
# install markers above already catch the generated install cell.
return False

11
docker/unsloth_nb_pip_magic.py
View file

@ -24,13 +24,10 @@ subprocess, so the shim applies. Safe no-op outside IPython.
import re
# Only the explicit `!<python> -m pip|uv ...` shell form (the `!` makes it a
# shell escape). Matched against the line with its trailing newline stripped.
# Input transformers see the RAW cell text -- IPython expands `{sys.executable}`
# later, inside the system() execution path -- so the braced form notebooks use
# to target the running kernel (`!{sys.executable} -m pip install ...`) and
# absolute interpreter paths (`!/opt/unsloth-venv/bin/python -m pip ...`),
# quoted or bare, must be matched here too or module-pip bypasses the PATH shim.
# Only the explicit `!<python> -m pip|uv ...` shell form. Input transformers see
# the RAW cell text (IPython expands `{sys.executable}` later), so the braced form
# (`!{sys.executable} -m pip install ...`) and absolute interpreter paths, quoted
# or bare, must be matched here too or module-pip bypasses the PATH shim.
_PY_M_PIP = re.compile(
r"""^(\s*)!\s*
(?:

50
docker/unsloth_nb_strip_colab.py
View file

@ -4,36 +4,23 @@
# Remove the Colab-only "how to run" sentence from Unsloth notebooks for Docker.
#
# Every generated notebook opens with a first markdown cell whose first line is a
# Colab instruction, e.g.
#
# To run this, press "*Runtime*" and press "*Run all*" on a **free** Tesla T4
# Google Colab instance!
# ... (and the A100 / L4 / "AMD Dev Cloud" variants) ...
#
# Inside the Docker image there is no "Runtime > Run all" menu and no Colab GPU,
# so the sentence is wrong/confusing. This strips ONLY that leading sentence; the
# rest of the cell (the Unsloth badge row, the "install on your local device"
# guide link, the "You will learn how to do ..." line) is kept untouched.
#
# This is a Docker-only transform applied at notebook-sync time. It is NOT pushed
# upstream: on Colab the sentence is correct, so the public notebooks keep it.
# Every generated notebook's first markdown cell opens with a Colab instruction
# ("To run this, press Runtime > Run all on a free Tesla T4 ...", plus A100/L4/AMD
# variants). Inside Docker there is no such menu or Colab GPU, so it is wrong;
# strip ONLY that leading sentence and keep the rest of the cell (badge row,
# local-install link, "You will learn ..." line). Docker-only, applied at sync
# time; NOT pushed upstream (on Colab the sentence is correct).
#
# Two modes:
# unsloth_nb_strip_colab.py <a.ipynb> [b.ipynb ...]
# strip the listed notebooks in place (idempotent).
# unsloth_nb_strip_colab.py <a.ipynb> [b.ipynb ...] strip in place (idempotent)
# unsloth_nb_strip_colab.py --state <STATE> --dest <DEST>
# STATE-aware sync migration. For each .ipynb in the "<sha256> <relpath>"
# STATE file (written by unsloth_sync_notebooks.sh) that still hashes to its
# recorded value (WE own it, unedited), strip the intro and update the
# recorded hash in place; user-edited notebooks (hash != recorded) are left
# untouched. Runs after every STATE write, covering first-boot populate,
# deleted-file restore, GitHub refresh and in-place image upgrades.
#
# Safe with refresh decisions: unsloth_nb_content_sig.py already classifies the
# intro cell as boilerplate, so the body digest used to detect "only boilerplate
# moved upstream" is identical whether or not the sentence is present.
# STATE-aware sync migration: for each .ipynb in the STATE file that still
# hashes to its recorded value (owned + unedited), strip the intro and update
# the hash; user-edited notebooks are left untouched. Runs after every STATE
# write (populate, restore, refresh, in-place upgrade).
#
# Safe with refresh decisions: content_sig already classifies the intro cell as
# boilerplate, so the body digest is identical with or without the sentence.
# Exit code is always 0.
import argparse
import hashlib
@ -44,12 +31,11 @@ import sys
# The stable identifier for the offending line (covers every GPU/Cloud variant).
_INTRO_PREFIX = "to run this, press"
# The baked notebooks ship example tqdm/progress-bar widget outputs plus a
# metadata.widgets state block; JupyterLab's ipywidgets manager cannot always
# rebuild the Colab-saved state, so they render as a stuck "Loading widget..."
# placeholder. Dropping the widget outputs + orphan state removes it; running the
# cell still creates a fresh widget. Outputs are not part of the refresh signature
# (content_sig hashes only cell type+source), so this is safe for edit detection.
# The baked notebooks ship example tqdm widget outputs + a metadata.widgets state
# block; JupyterLab can't always rebuild the Colab-saved state, so they render as
# a stuck "Loading widget..." placeholder. Dropping the widget outputs + orphan
# state removes it (running the cell recreates a fresh widget). Outputs aren't in
# the refresh signature (content_sig hashes cell type+source), so this is safe.
_WIDGET_VIEW_MIME = "application/vnd.jupyter.widget-view+json"

65
docker/unsloth_nb_view.py
View file

@ -4,39 +4,27 @@
# Build a categorized, Colab-like folder VIEW of the Unsloth notebooks.
#
# The canonical notebooks live under DEST/nb/<file>.ipynb (a mirror of
# unslothai/notebooks, populated + refreshed by unsloth_sync_notebooks.sh). That
# flat tree is great for syncing but poor for browsing. This builds a sibling
# directory of *relative symlinks* grouped into folders that mirror the README
# section headers, e.g.
#
# The canonical notebooks live flat under DEST/nb/<file>.ipynb (mirror of
# unslothai/notebooks, kept by unsloth_sync_notebooks.sh). This builds a sibling
# dir of *relative symlinks* grouped into folders mirroring the README headers:
# <VIEW>/01 Main Notebooks/Llama3_2_(1B_and_3B)_Conversational.ipynb
# <VIEW>/02 Gemma 4 Notebooks/...
# ...
# <VIEW>/99 Other Notebooks/<anything on disk not linked from the README>
#
# Why symlinks: the real .ipynb files are never moved or renamed, so the sync
# state machine (which walks `find -type f`, skipping symlinks) and the
# edit/refresh logic are completely unaffected. The VIEW is a sibling of DEST
# (outside it), rebuilt from scratch on every boot, and disposable.
# Symlinks so the real files never move (the sync state machine skips symlinks);
# the VIEW is a disposable sibling of DEST, rebuilt from scratch on every boot.
#
# Categorization rules:
# * Section = the nearest preceding `### ` header in DEST/README.md. The same
# topic header repeats across the Fine-tuning / Kaggle / AMD domains; those
# merge into one folder (first appearance fixes the order).
# * Folder names are cleaned: dashes and slashes -> spaces, whitespace
# collapsed, numbered `NN ` by first appearance so JupyterLab's alpha sort
# preserves README order. "Other Notebooks" is always last.
# * A notebook linked under several sections lands in its first (README order).
# * AMD-*.ipynb are hidden unless --amd (an AMD/HIP GPU was detected).
# * Any on-disk nb/*.ipynb not linked from the README goes to "Other Notebooks".
# * Section = nearest preceding `###` header in README.md; a header repeated
# across Fine-tuning/Kaggle/AMD domains merges into one folder (first order).
# * Folder names cleaned (dashes/slashes -> spaces) and numbered `NN ` by first
# appearance so JupyterLab's alpha sort keeps README order; "Other" is last.
# * A notebook linked under several sections lands in its first.
# * AMD-*.ipynb hidden unless --amd; unlinked nb/*.ipynb go to "Other Notebooks".
#
# Usage:
# unsloth_nb_view.py <DEST> <VIEW> [--amd] build the symlink view
# unsloth_nb_view.py <DEST> --print [--amd] print "section\tfile" rows
#
# Exit code is 0 on success; on any error it prints a diagnostic to stderr and
# exits non-zero so the caller can fall back to the raw tree.
# Exits 0 on success; on error prints to stderr and exits nonzero so the caller
# can fall back to the raw tree.
import argparse
import os
import re
@ -51,10 +39,9 @@ _OTHER = "Other Notebooks"
def clean_section(title):
"""README header text -> a filesystem-friendly folder label."""
# Drop trailing '#' and surrounding whitespace.
title = title.strip().strip("#").strip()
# Strip a leading run of emoji / symbols some domain headers lead with (e.g.
# "🐧 AMD Notebooks", "📒 Kaggle Notebooks") so the folder label is clean text.
# Strip a leading run of emoji/symbols some domain headers lead with so the
# folder label is clean text.
title = re.sub(r"^[^\w]+", "", title)
title = title.replace("-", " ").replace("/", " ")
title = re.sub(r"\s+", " ", title).strip()
@ -79,11 +66,9 @@ def parse_readme(readme_path):
rows = []
seen_pairs = set() # (section, filename) already emitted
section = None
# Reset on ANY markdown heading, not just `###`. The catalog uses `#`/`##`
# domain headers (e.g. "# AMD Notebooks", "# Kaggle Notebooks") that carry
# their own `nb/*.ipynb` link tables directly, with no intervening `###`.
# Matching only `###` left `section` stale, so those links were mis-filed
# under the previous section instead of getting their own folder.
# Reset on ANY markdown heading, not just `###`: `#`/`##` domain headers carry
# their own nb/*.ipynb tables with no intervening `###`, so matching only `###`
# left `section` stale and mis-filed those links under the previous section.
for line in text.splitlines():
m = re.match(r"^#{1,6}\s+(.*)$", line)
if m:
@ -204,16 +189,10 @@ def _points_into(link, dest_real):
def _clear_view(path, dest_real):
# Tear down a previously built VIEW in place. VIEW is also JupyterLab's
# landing directory, so a user may have saved real notebooks (or their own
# symlinks) here -- those MUST survive a rebuild. We therefore unlink only
# the symlinks we own (they resolve into DEST, see _points_into) and rmdir
# only folders that end up empty; any regular file and any user symlink is
# left untouched, and a non-empty folder simply stays.
#
# The VIEW root itself is never unlinked: build_view already resolved a
# symlinked root to its target, and an operator's routing symlink must
# survive. isdir on a non-link root is safe to walk.
# Tear down a previously built VIEW in place. It is also JupyterLab's landing
# dir, so user files/symlinks MUST survive: unlink only the symlinks we own
# (resolve into DEST, see _points_into) and rmdir only emptied folders. The
# VIEW root is never unlinked (build_view already resolved a symlinked root).
if os.path.islink(path) or not os.path.isdir(path):
return
for root, dirs, files in os.walk(path, topdown = False):

217
docker/unsloth_pip_shim.py
View file

@ -79,12 +79,9 @@ _VALUE_FLAGS = {
"--implementation",
"-e",
"--editable",
# Every remaining value-taking flag of `uv pip install` / `pip install`
# (generated from both tools' --help). A value flag missing here makes the
# scanner misread its VALUE: `uv pip install --torch-backend cu128 torch`
# dropped the protected torch but then exec'd uv with no install target at
# all (uv hard-errors) instead of no-oping like the attached `=` form.
# uv:
# Every remaining value-taking flag of pip/uv install (from both --help). A
# missing one makes the scanner misread its VALUE: `--torch-backend cu128 torch`
# dropped torch then exec'd uv with no target (hard-error). uv:
"--allow-insecure-host",
"--build-constraints",
"-b",
@ -142,49 +139,35 @@ _VALUE_FLAGS = {
"--requirements-from-script",
"--uploaded-prior-to",
}
# Of those value-flags, the ones whose VALUE is itself an install target: a
# requirements file pulls real requirements. An index-url / find-links /
# constraint / target value is an option, not something to install.
# uv spells the long forms in the PLURAL (`--requirements`, `--constraints`);
# include both so a `uv pip install --requirements reqs.txt` is filtered too.
# Value-flags whose VALUE is itself an install target: a requirements file pulls
# real requirements (index-url/find-links/constraint/target values are options).
# uv spells the long forms plural (--requirements/--constraints); include both.
_REQ_FILE_FLAGS = {"-r", "--requirement", "--requirements"}
# Constraint files are not install targets, but pip applies their pins during
# resolution, so a `-c constraints.txt` that pins torch/transformers/etc. can
# still downgrade or reinstall a baked package when another target pulls it in.
# Filter protected packages out of them the same way as requirement files.
# (uv's long form is the plural `--constraints`.)
# Constraint files aren't install targets, but pip applies their pins during
# resolution, so a -c that pins torch/transformers can still downgrade a baked
# package. Filter them like requirement files. (uv's long form is --constraints.)
_CONSTRAINT_FILE_FLAGS = {"-c", "--constraint", "--constraints"}
# -e/--editable <path|url|vcs> takes the NEXT token as its target (pip:
# `-e, --editable <path/url>`), and that target is a real install target. A
# protected editable (e.g. `-e git+https://.../unsloth.git#egg=unsloth`) must
# drop BOTH the flag and its value; dropping the value alone leaves pip a
# dangling `-e` that swallows the next kept package and fails the whole cell.
# -e/--editable <path|url|vcs> takes the NEXT token as a real install target. A
# protected editable must drop BOTH flag and value; dropping only the value
# leaves pip a dangling -e that swallows the next kept package and fails the cell.
_EDITABLE_FLAGS = {"-e", "--editable"}
# -P/--upgrade-package and --reinstall-package are uv's selective upgrade/reinstall
# flags: naming a baked package (`uv pip install -P torch peft`) lets an ordinary
# target refresh/reinstall it and clobber the pinned stack. Filter the value through
# _KEEP too, dropping the flag+value pair for a protected name so no dangling
# selector swallows the next target. Unlike -e, none is itself an install target.
# -P/--upgrade-package and --reinstall-package are uv's selective upgrade flags:
# naming a baked package lets an ordinary target refresh it. Filter the value
# through _KEEP, dropping the flag+value pair for a protected name. Unlike -e,
# none is itself an install target.
_UPGRADE_PKG_FLAGS = {"-P", "--upgrade-package", "--reinstall-package"}
# Short value-flags pip/uv accept in the ATTACHED form, i.e. the 2-char flag
# glued to its value in one token: `-rreqs.txt`, `-cconstraints.txt`, `-epath`,
# `-Pname`. The scanner splits the flag from the value so the value is filtered
# (requirement/constraint file) or classified (-e/-P) instead of falling through
# as an opaque option -- otherwise an attached `-r`-only cell no-ops and an
# attached `-c`/`-e`/`-P` value bypasses _KEEP.
# Short value-flags accepted ATTACHED (flag glued to value): -rreqs.txt, -cX,
# -epath, -Pname. The scanner splits flag from value so it is filtered/classified,
# else an attached -r-only cell no-ops and -c/-e/-P bypasses _KEEP.
_ATTACHED_SHORT_FLAGS = {"-r", "-c", "-e", "-P"}
# Resolver-wide reinstall / ignore-installed switches (pip --force-reinstall,
# --ignore-installed, -I; uv --reinstall) REINSTALL already-satisfied packages,
# including the baked torch/transformers a kept target pulls in as deps. Drop them
# so an unprotected install cannot rebuild the pinned stack; the kept target still
# installs. Per-package selectors (-P / --reinstall-package) go via _UPGRADE_PKG_FLAGS.
# uv's --exact is destructive the other way: an exact SYNC that REMOVES everything
# outside the kept target's closure (vLLM, bitsandbytes, NVIDIA libs), so drop it too.
# Resolver-wide reinstall/ignore-installed switches (pip --force-reinstall,
# --ignore-installed, -I; uv --reinstall) rebuild already-satisfied baked deps;
# drop them (the kept target still installs). uv's --exact is destructive the
# other way (SYNC removes everything outside the target's closure), so drop it too.
_REINSTALL_FLAGS = {"--force-reinstall", "--ignore-installed", "-I", "--reinstall", "--exact"}
# Value-flags whose flag+value pair is dropped outright in shim mode.
# `--upgrade-strategy eager` makes pip upgrade EVERY dependency of a kept target,
# refreshing the baked torch/transformers. Dropping it falls back to pip's default
# `only-if-needed`, so the target still installs but satisfied protected deps stay.
# Value-flags whose flag+value pair is dropped outright. --upgrade-strategy eager
# would upgrade EVERY dep of a kept target; dropping it falls back to pip's
# only-if-needed default so satisfied protected deps stay.
_DROP_VALUE_FLAGS = {"--upgrade-strategy"}
@ -215,12 +198,9 @@ def _canon(token):
if the token is not a plain pkg spec (url / path / vcs / option)."""
if token.startswith("-"):
return None
# PEP 508 direct reference: "name [extras] @ <url>" (e.g.
# "torch @ https://.../torch.whl", "unsloth @ git+https://..."). The name is
# at the front, so pull it out BEFORE the url/vcs guard below -- otherwise a
# protected package pinned through a URL slips past _KEEP and reinstalls into
# the base venv. A non-protected direct reference still returns its name and
# is kept by the caller exactly as before (treated as an install target).
# PEP 508 direct reference: "name [extras] @ <url>". The name is at the front,
# so pull it out BEFORE the url/vcs guard below, or a protected package pinned
# through a URL slips past _KEEP. Non-protected refs still return their name.
_dref = re.match(
r"^([A-Za-z0-9][A-Za-z0-9._-]*)\s*(?:\[[^\]]*\])?\s*@(?:\s|git\+|hg\+|bzr\+|svn\+|[a-z]+://)",
token,
@ -228,54 +208,37 @@ def _canon(token):
if _dref:
return _dref.group(1).lower().replace("_", "-") or None
if re.match(r"^[a-z]+\+", token) or "://" in token or token.startswith((".", "/")):
# A VCS / URL install can still name a protected package via the legacy
# `#egg=NAME` (or `&egg=NAME`) fragment, e.g.
# `git+https://github.com/unslothai/unsloth.git#egg=unsloth`. Pull that
# name out so _KEEP can drop it; otherwise the shim would exec the URL
# and reinstall a baked package into the venv. A non-protected egg name
# is returned too, but the caller keeps it as a normal target either way.
# A VCS/URL install can name a protected package via the legacy #egg=NAME
# (or &egg=NAME) fragment; pull it out so _KEEP can drop it, else the shim
# execs the URL and reinstalls a baked package.
_egg = re.search(r"[#&]egg=([A-Za-z0-9][A-Za-z0-9._-]*)", token)
if _egg:
return _egg.group(1).lower().replace("_", "-") or None
# A direct wheel URL or local wheel path still names its distribution in
# the PEP 427 filename ({distribution}-{version}-...-...-....whl), so a
# bare `pip install https://.../torch-2.11.0+cu128-...whl` would slip a
# protected package past _KEEP as an opaque positional and reinstall the
# baked torch. Dashes cannot appear inside the distribution component (a
# run of -_. normalises to a single -), so the leading dash-split of the
# basename is the distribution name; pull it so _KEEP can drop it. A
# non-protected wheel returns its name and the caller keeps the token.
# A wheel URL/path names its distribution in the PEP 427 filename, so a
# bare `pip install .../torch-2.11.0+cu128-...whl` would slip torch past
# _KEEP. Dashes can't appear in the distribution component, so the leading
# dash-split of the basename is the name; pull it so _KEEP can drop it.
_whl = re.search(r"([^/\\#?]+)\.whl(?:[#?]|$)", token)
if _whl:
dist = _whl.group(1).split("-", 1)[0].strip().lower().replace("_", "-")
if dist:
return dist
# A source archive (sdist / zip) URL or path names its distribution the
# same way ({name}-{version}.tar.gz etc.), so `pip install
# https://files.pythonhosted.org/.../unsloth-2026.7.1.tar.gz` or
# `./torch-2.11.0.tar.gz` must be matched against _KEEP too, not passed
# through as an opaque positional that reinstalls the baked package.
# A source archive URL/path names its distribution the same way
# ({name}-{version}.tar.gz), so match it against _KEEP too instead of
# passing it through as an opaque positional.
_arch = _sdist_name(token.split("#", 1)[0].split("?", 1)[0].rstrip("/").rsplit("/", 1)[-1])
if _arch:
return _arch
# A VCS URL without an #egg= fragment still installs a named project:
# pip/uv derive the distribution from the repo, and for the packages we
# protect the repo basename equals the distribution
# (huggingface/transformers.git -> transformers,
# unslothai/unsloth-zoo.git -> unsloth-zoo). Infer it from the last path
# segment so a bare `pip install git+https://github.com/huggingface/
# transformers.git` -- an egg-less form this repo itself recommends in
# unsloth/models/loader.py -- cannot reinstall the baked package past
# _KEEP. A non-protected repo returns its basename and the caller keeps
# the token as a normal target either way.
# A VCS URL without #egg= still installs a named project: the repo
# basename equals the distribution for the packages we protect
# (huggingface/transformers.git -> transformers). Infer it from the last
# path segment so a bare egg-less git+ URL can't reinstall past _KEEP.
if re.match(r"^[a-z]+\+", token):
_rest = token.split("#", 1)[0].split("?", 1)[0]
# Drop the @ref from the PATH portion BEFORE taking the last path
# segment: a ref may itself contain a slash (@feature/foo), which
# would otherwise become the "basename" and dodge _KEEP. Split the
# path off the authority first so an SSH userinfo @ (git+ssh://
# git@github.com/...) is never mistaken for the ref separator;
# like pip's own parser, the ref is everything after the LAST @.
# Drop the @ref from the PATH before taking the basename: a ref may
# contain a slash (@feature/foo) and dodge _KEEP. Split path from
# authority first so an SSH userinfo @ isn't mistaken for the ref;
# like pip, the ref is everything after the LAST @.
if "://" in _rest:
_authority, _slash, _path = _rest.partition("://")[2].partition("/")
if "@" in _path:
@ -288,31 +251,23 @@ def _canon(token):
_seg = _seg.strip().lower().replace("_", "-")
if _seg:
return _seg
# A local project DIRECTORY (`pip install ./transformers`,
# `pip install -e ./unsloth`) installs the project it contains, and a
# same-version dev build slips past even the protected constraints file
# (constraints only reject a version MISMATCH), silently swapping the
# baked, tested wheel for a local build. Resolve the project name from
# its metadata so _KEEP applies to this form like every other artifact
# form (wheel/sdist/VCS/egg). Non-directories and metadata-less dirs
# pass through as before.
# A local project DIRECTORY installs the project it contains; a same-
# version dev build slips past even the constraints file (which only
# rejects a MISMATCH). Resolve the name from its metadata so _KEEP applies
# like every other artifact form. Metadata-less dirs pass through.
_local = _local_project_name(token)
if _local:
return _local
return None # plain url / metadata-less local path -> let it pass through
# A local project dir referenced without ./ or / (`pip install subdir/proj`)
# is still a path target to pip when it exists on disk; classify it the same
# way before the spec parse below mangles the separator.
# A local project dir referenced without ./ or / is still a path target when
# it exists on disk; classify it before the spec parse mangles the separator.
if "/" in token or os.sep in token:
_local = _local_project_name(token)
if _local:
return _local
# A bare wheel filename (no ./ or / prefix and no scheme) is still a valid
# pip target from the CWD: `pip install torch-2.11.0-cp312-...-linux.whl`.
# It reaches here because it starts with neither `.`/`/` nor a scheme, so
# without this it would fall through as the whole filename and miss _KEEP,
# reinstalling the baked torch. Parse its PEP 427 distribution the same way
# as the URL/path wheel case above.
# A bare wheel filename from the CWD (no ./ or scheme) is still a valid pip
# target; without this it falls through and misses _KEEP. Parse its PEP 427
# distribution like the URL/path wheel case above.
if token.lower().endswith(".whl"):
dist = token.rsplit("/", 1)[-1][:-4].split("-", 1)[0].strip().lower().replace("_", "-")
if dist:
@ -375,10 +330,9 @@ def _version_pin(token):
return m.group(1) if m else None
# pip expands ${UPPERCASE_NAME} in requirements files AFTER we classify the
# literal text (pip's ENV_VAR_RE; uv matches it), so `${PKG}==...` with
# PKG=torch would slip a protected package past _KEEP. Expand with the same
# syntax for CLASSIFICATION only; kept lines are forwarded verbatim.
# pip expands ${UPPERCASE_NAME} in requirements files after we classify the text,
# so `${PKG}==...` with PKG=torch would slip past _KEEP. Expand for CLASSIFICATION
# only; kept lines are forwarded verbatim.
_ENV_REF_RE = re.compile(r"\$\{([A-Z0-9_]+)\}")
@ -650,10 +604,9 @@ def main():
if argv[:1] == ["--unsloth-selfcheck-value-flags"]:
_selfcheck_value_flags()
# Only intercept inside a notebook kernel (UNSLOTH_NB_SHIM is set by the baked
# IPython startup and by `unsloth-run`). EVERYWHERE else -- install.sh during
# the image build, internal tooling, an interactive shell -- behave exactly
# like the real tool, so we never disturb the build or system package mgmt.
# Only intercept inside a notebook kernel (UNSLOTH_NB_SHIM set by the baked
# IPython startup and unsloth-run). Everywhere else (build install.sh, shells)
# behave exactly like the real tool.
if os.environ.get("UNSLOTH_NB_SHIM") != "1":
os.execv(REAL[tool], [REAL[tool]] + argv)
return
@ -705,19 +658,16 @@ def main():
keep_args.append(_c_path)
dropped.extend(_c_drp)
elif prev_flag in _DROP_VALUE_FLAGS:
# --upgrade-strategy (eager): the flag was appended when we saw
# it; pop it and drop the flag+value pair so pip falls back to
# its safe only-if-needed default.
# --upgrade-strategy (eager): pop the appended flag and drop the
# pair so pip falls back to only-if-needed.
if keep_args and keep_args[-1] == prev_flag:
keep_args.pop()
dropped.append(prev_flag + " " + tok)
elif prev_flag in _EDITABLE_FLAGS or prev_flag in _UPGRADE_PKG_FLAGS:
# The flag was held back (not appended yet): its value is an
# install target (-e path/url/vcs) or an upgrade selector
# (-P name), both filtered through _KEEP. Dropping a protected
# value drops the flag with it, so pip/uv is never left a
# dangling `-e`/`-P` that fails the cell or refreshes a baked
# package. A kept editable target sets has_target; -P does not.
# The flag was held back: its value is an install target (-e) or
# upgrade selector (-P), both filtered through _KEEP. Dropping a
# protected value drops the flag too (no dangling -e/-P). A kept
# editable sets has_target; -P does not.
_action, _ver = _classify_flag_target(tok)
if _action == "drop":
if _ver and not recorded:
@ -733,11 +683,9 @@ def main():
skip_next = False
prev_flag = None
continue
# --flag=value form: pip accepts --requirement=reqs.txt / --index-url=URL
# as a single token. Without this the token starts with "-", so it is kept
# as an opaque option and a `-r` file is never filtered -- and worse, it
# never counts as a target, so a cell whose only target is that file
# silently no-ops and installs nothing.
# --flag=value form (--requirement=reqs.txt / --index-url=URL as one
# token). Without this it is kept as an opaque option, the -r file is never
# filtered, and a file-only cell silently installs nothing.
if tok.startswith("--") and "=" in tok:
_flag, _, _val = tok.partition("=")
if _flag in _VALUE_FLAGS:
@ -775,12 +723,10 @@ def main():
else:
keep_args.append(tok) # option with inline value, not a target
continue
# Attached short value-flag form: pip/uv accept `-rreqs.txt`,
# `-cconstraints.txt`, `-epath` and `-Pname` as ONE token. Without this
# the token starts with "-" and falls through as an opaque option, so an
# `-r`-only cell no-ops (has_target stays False) and an attached
# `-c`/`-e`/`-P` value bypasses _KEEP. Split the 2-char flag from its
# value and reuse the separated-form handling.
# Attached short value-flag form (-rreqs.txt, -cX, -epath, -Pname as ONE
# token). Without this it falls through as an opaque option: an -r-only
# cell no-ops and -c/-e/-P bypasses _KEEP. Split the flag from its value
# and reuse the separated-form handling.
if len(tok) > 2 and tok[0] == "-" and tok[1] != "-" and tok[:2] in _ATTACHED_SHORT_FLAGS:
_sflag, _sval = tok[:2], tok[2:]
if (_sflag in _REQ_FILE_FLAGS or _sflag in _CONSTRAINT_FILE_FLAGS) and "://" in _sval:
@ -862,18 +808,15 @@ def main():
if dropped:
print("[unsloth-nb] kept baked versions, skipped: " + " ".join(dropped))
# Anything left to actually install? `has_target` was set during the scan for
# a kept package spec, a positional url / path / vcs / editable target, or a
# -r/--requirement file. A line carrying only baked packages plus option flags
# (e.g. `--extra-index-url <url> torch`) leaves no target, so no-op instead of
# exec'ing a bare `pip install --extra-index-url <url>` that would fail.
# Anything left to install? has_target was set for a kept spec, a positional
# url/path/vcs/editable, or a -r file. A line with only baked packages + option
# flags leaves no target, so no-op instead of exec'ing a bare install that fails.
if not has_target:
print("[unsloth-nb] nothing to install after keeping the baked stack; ok.")
return
cmd = [REAL[tool]] + head + keep_args
# Constrain the resolver too: without this an allowed target could pull an
# incompatible torch/transformers/etc. in as a DEPENDENCY and replace the
# baked wheel even though the argument filter kept it off the command line.
# Constrain the resolver too: an allowed target could pull an incompatible
# torch/transformers in as a DEPENDENCY and replace the baked wheel.
constraints = _protected_constraints_file()
if constraints:
cmd += ["--constraint", constraints]

10
docker/unsloth_run.py
View file

@ -71,12 +71,10 @@ def main():
want = args.tf or pin or (compat.tier_for_model(model) if compat else None)
sidecar = compat.sidecar_for(want) if (compat and want) else None
# Materialise the notebook locally for nbconvert. With --out, stage both the
# input copy and the executed result as temp files NEXT TO the destination
# (same dir, so the kernel cwd matches and the publish is one atomic
# os.replace) and only publish over an existing --out file when execution
# succeeded -- a timeout / failed cell / missing kernel must not destroy the
# previous output.
# Materialise the notebook for nbconvert. With --out, stage the input copy and
# the result as temp files NEXT TO the destination (same dir, so kernel cwd
# matches and publish is one atomic os.replace) and only publish on success --
# a timeout / failed cell / missing kernel must not destroy the previous output.
tmp_dir = None
tmp_files = []
publish_from = None

8
docker/unsloth_studio_update.sh
View file

@ -65,11 +65,9 @@ echo "[studio-update] before: unsloth $(version_of)"
# (or any branch/tag/sha); otherwise take the latest PyPI release.
if [ -n "$REF" ]; then
SPECS="git+https://github.com/unslothai/unsloth.git@${REF}#egg=unsloth"
# unsloth-zoo does NOT track unsloth's tags/SHAs (its release cadence differs;
# the publish workflow resolves the zoo ref separately for the same reason).
# Use --zoo-ref if given; else use the unsloth ref only when the zoo repo
# actually has it, falling back to main so `--ref <unsloth-tag>` does not fail
# on a tag/SHA that simply does not exist in unsloth-zoo.
# unsloth-zoo does NOT track unsloth's tags/SHAs (different cadence). Use
# --zoo-ref if given; else the unsloth ref only when the zoo repo has it,
# falling back to main so `--ref <unsloth-tag>` doesn't fail on a missing ref.
_zoo_ref="$ZOO_REF"
if [ -z "$_zoo_ref" ]; then
if git ls-remote --exit-code https://github.com/unslothai/unsloth-zoo.git \

73
docker/unsloth_sync_notebooks.sh
View file

@ -1,16 +1,14 @@
#!/usr/bin/env bash
# Populate and refresh /workspace/unsloth-notebooks from unslothai/notebooks.
#
# The image bakes a read-only template at /opt/unsloth-notebooks so the
# notebooks are present in JupyterLab instantly and offline. On boot this script
# copies the template into /workspace/unsloth-notebooks (first run only) and then
# best-effort refreshes from GitHub when upstream has actually advanced.
# The image bakes a read-only template at /opt/unsloth-notebooks so notebooks are
# present instantly and offline. On boot this copies the template into
# /workspace/unsloth-notebooks (first run) then best-effort refreshes from GitHub
# when upstream advances.
#
# The user's edits ALWAYS win. We remember the content hash of every file we
# wrote; on refresh a file whose current hash differs from what we last wrote is
# treated as user-modified and is left untouched. So a refresh only updates files
# the user has not changed and adds new ones -- it never clobbers an edited
# notebook and never produces merge conflicts.
# The user's edits ALWAYS win: we record each written file's hash; on refresh a
# file whose hash differs is treated as user-modified and left untouched. So a
# refresh only updates unchanged files and adds new ones, never clobbering edits.
#
# Opt-out / tuning (all optional):
# UNSLOTH_SKIP_NOTEBOOK_SYNC=1 do nothing (no populate, no refresh)
@ -37,11 +35,9 @@ STATE="$DEST/.unsloth_sync_state" # "sha256 relpath" of what we last wrote
SYNCED="$DEST/.unsloth_sync_commit" # upstream commit we last synced to
TIMEOUT="${UNSLOTH_NOTEBOOK_FETCH_TIMEOUT:-60}"
# Resolve a helper script ($1 explicit override, $2 PATH command name, $3
# sibling filename next to this script), echoing the resolved path or nothing.
# An empty result leaves the caller's guard to degrade gracefully. Used for the
# content-sig comparator (SIG), categorized-view builder (VIEW) and Docker-only
# Colab-intro stripper (STRIP).
# Resolve a helper script ($1 override, $2 PATH command, $3 sibling filename),
# echoing the path or nothing (empty lets the caller degrade). Used for SIG,
# VIEW and STRIP helpers.
PYBIN="$(command -v python3 2>/dev/null || command -v python 2>/dev/null || true)"
_self_dir="$(cd "$(dirname "$0")" 2>/dev/null && pwd)"
resolve_helper() {
@ -54,11 +50,10 @@ SIG_HELPER="$(resolve_helper "${UNSLOTH_NB_SIG_HELPER:-}" unsloth-nb-content-sig
VIEW_HELPER="$(resolve_helper "${UNSLOTH_NB_VIEW_HELPER:-}" unsloth-nb-view unsloth_nb_view.py)"
STRIP_HELPER="$(resolve_helper "${UNSLOTH_NB_STRIP_HELPER:-}" unsloth-nb-strip-colab unsloth_nb_strip_colab.py)"
# True only when BOTH are .ipynb, the SIG helper is usable, and it reports the
# non-boilerplate middle (ignoring the auto-generated install header /
# announcements / footer) is identical -- so a refresh doesn't rewrite an
# untouched notebook when only that boilerplate moved upstream. Any failure
# returns false, so the caller falls back to a normal refresh.
# True only when both are .ipynb, the SIG helper is usable, and it reports the
# non-boilerplate middle (ignoring install header/announcements/footer) identical,
# so a refresh doesn't rewrite an untouched notebook when only boilerplate moved.
# Any failure returns false (caller falls back to a normal refresh).
middle_unchanged() {
case "$1" in *.ipynb) : ;; *) return 1 ;; esac
[ -n "$PYBIN" ] && [ -n "$SIG_HELPER" ] || return 1
@ -93,7 +88,7 @@ nb_gpu_is_amd() {
# Rebuild the sibling symlink VIEW (categorized folders mirroring the README
# headers) from scratch. Symlinks live OUTSIDE $DEST, so the sync state machine
# (which walks `find -type f`, skipping symlinks) never sees them.
# (find -type f) never sees them.
build_categorized_view() {
[ "${UNSLOTH_SKIP_NOTEBOOK_VIEW:-0}" = "1" ] && return 0
[ -n "$PYBIN" ] && [ -n "$VIEW_HELPER" ] || return 0
@ -140,11 +135,10 @@ if [ ! -f "$STATE" ]; then
rel="${rel#./}"
case "$rel" in .unsloth_template_commit) continue ;; esac
mkdir -p "$DEST/$(dirname "$rel")" 2>/dev/null || true
# A pre-existing file here (bind-mounted or hand-created before first boot)
# is user data: keep it, and do NOT record it in the sync state -- if
# recorded, the GitHub refresh below would see the hash match, treat it as
# pristine and overwrite it. Only files we lay down (or that already match
# the template byte-for-byte) are recorded as managed.
# A pre-existing file (bind-mounted or hand-created) is user data: keep it
# and do NOT record it -- if recorded, the refresh below would see a hash
# match, treat it as pristine and overwrite it. Only files we lay down (or
# that match the template byte-for-byte) are recorded as managed.
if [ -e "$DEST/$rel" ] \
&& [ "$(hash_of "$DEST/$rel")" != "$(hash_of "$TEMPLATE/$rel")" ]; then
echo "[unsloth-nb] kept existing user file: $DEST/$rel"
@ -159,12 +153,11 @@ if [ ! -f "$STATE" ]; then
echo "[unsloth-nb] notebooks ready at $DEST"
fi
# 1b) Every-boot OFFLINE restore of deleted notebooks: a file we previously
# wrote that the user has since DELETED comes back from the baked template (no
# network needed). Files that still exist (edited or not) are never touched, so
# this cannot clobber an edit; the restored file's recorded hash is reset to
# the template's so the GitHub refresh below treats it as pristine and bumps it
# to latest. Opt out with UNSLOTH_KEEP_DELETED_NOTEBOOKS=1.
# 1b) Every-boot OFFLINE restore of deleted notebooks: a file we wrote that the
# user has since DELETED comes back from the baked template (no network). Existing
# files are never touched (can't clobber an edit); the restored hash is reset to
# the template's so the refresh below bumps it to latest. Opt out with
# UNSLOTH_KEEP_DELETED_NOTEBOOKS=1.
if [ -f "$STATE" ] && [ "${UNSLOTH_KEEP_DELETED_NOTEBOOKS:-0}" != "1" ]; then
restored=0
RS_TMP="$(mktemp)"
@ -217,9 +210,8 @@ while IFS= read -r -d '' f; do
if [ -e "$dst" ]; then
rec="${LAST[$rel]:-}"
if [ -z "$rec" ]; then
# File exists in DEST but the sync state never recorded it -> it is a
# pre-existing user / bind-mounted file. Treat it as user-owned: keep
# it and do not adopt it into the state (so it stays protected).
# In DEST but never recorded -> a pre-existing user/bind-mounted file.
# Keep it and don't adopt it into the state (stays protected).
kept=$((kept + 1))
continue
fi
@ -230,19 +222,16 @@ while IFS= read -r -d '' f; do
continue
fi
if [ -n "$rec" ] && middle_unchanged "$dst" "$f"; then
# Untouched notebook whose only upstream change is the install
# header / announcements / footer. The tutorial body is identical,
# so don't churn the user's file -- keep it and its marker as-is.
# Untouched notebook whose only upstream change is the install header/
# announcements/footer. Body identical, so keep it and its marker.
printf '%s %s\n' "$rec" "$rel" >> "$TMPSTATE"
unchanged=$((unchanged + 1))
continue
fi
elif [ -n "${LAST[$rel]:-}" ] && [ "${UNSLOTH_KEEP_DELETED_NOTEBOOKS:-0}" = "1" ]; then
# We previously wrote this notebook and the user has since DELETED it.
# With the opt-out set, honor the deletion instead of restoring it from
# the fresh clone when upstream advances (otherwise the deletion only
# held until the next remote refresh). Keep the record so it stays known
# as managed-but-deleted.
# We wrote this notebook and the user DELETED it. With the opt-out set,
# honor the deletion instead of restoring it from the fresh clone. Keep
# the record so it stays known as managed-but-deleted.
printf '%s %s\n' "${LAST[$rel]}" "$rel" >> "$TMPSTATE"
kept=$((kept + 1))
continue

5
install.ps1
View file

@ -1965,9 +1965,8 @@ exit 0
function Get-TorchIndexUrl {
$baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
# Explicit override (parity with install.sh):
# UNSLOTH_TORCH_INDEX_FAMILY=cu128|cu130|cu126|cpu|... pins the wheel
# index when probing is wrong or impossible (no GPU on the build host,
# containerised installs, CI).
# UNSLOTH_TORCH_INDEX_FAMILY=cu128|cu130|cu126|cpu|... pins the wheel index
# when probing is wrong or impossible (no GPU, containers, CI).
if ($env:UNSLOTH_TORCH_INDEX_FAMILY) { return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY)" }
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {

21
install.sh
View file

@ -1996,12 +1996,10 @@ fi
_REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)"
# ── unsloth-zoo overlay ref (for --local installs) ──
# --local installs overlay unsloth-zoo straight from git so the Studio venv
# tracks the same zoo as the editable unsloth checkout. Honor UNSLOTH_ZOO_REF
# (the Docker publish workflow resolves one ref and forwards it to BOTH the base
# and Studio builds) so the published image runs the operator-requested zoo, not
# whatever main happens to be at build time. Unset -> main, byte-identical to the
# previous bare git URL (pip treats no @ref as the repo's default branch).
# --local overlays unsloth-zoo from git so the Studio venv tracks the same zoo as
# the editable unsloth checkout. Honor UNSLOTH_ZOO_REF (the Docker publish
# workflow forwards one ref to both builds) so the image runs the requested zoo.
# Unset -> main, byte-identical to the previous bare git URL.
_ZOO_REF="${UNSLOTH_ZOO_REF:-main}"
_ZOO_GIT_SPEC="unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo@${_ZOO_REF}"
@ -2069,13 +2067,10 @@ _has_amd_rocm_gpu() {
get_torch_index_url() {
_base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
_base="${_base%/}"
# Explicit pin for hosts where probing is impossible or must not happen
# (Docker image builds, CI runners). Names the index path leaf directly:
# UNSLOTH_TORCH_INDEX_FAMILY=cu128|cu130|cu126|rocm7.2|cpu|...
# The Blackwell Docker image build uses this: at build time there is no
# GPU and no nvidia-smi, but the image targets CUDA, so probing would
# land on the cpu (CI) or cu126 (GPU build hosts leak /proc/driver/nvidia
# but not nvidia-smi) wheels depending on which host built the image.
# Explicit pin for hosts where probing is impossible (Docker builds, CI).
# Names the index leaf: UNSLOTH_TORCH_INDEX_FAMILY=cu128|cu130|cu126|rocm7.2|cpu|...
# The Blackwell build uses this: no GPU/nvidia-smi at build time, but the image
# targets CUDA, so probing would land on cpu (CI) or cu126 wheels.
if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then
echo "$_base/${UNSLOTH_TORCH_INDEX_FAMILY}"; return
fi

7
studio/install_llama_prebuilt.py
View file

@ -4995,10 +4995,9 @@ def activate_staged_dir(staging_dir: Path, dst: Path) -> None:
try:
os.replace(staging_dir, dst)
except OSError as exc:
# Busy/in-use (Windows AV holding a DLL) OR cross-device (overlayfs in a
# Docker build): both are safe to complete by copying the freshly
# extracted staging tree and removing it. Anything else (disk full,
# missing path) re-raises so we never leave a partial install behind.
# Busy/in-use (Windows AV) OR cross-device (overlayfs in a Docker build):
# both are safe to complete by copying the staging tree and removing it.
# Anything else (disk full, missing path) re-raises.
if not (is_busy_lock_error(exc) or is_cross_device_error(exc)):
raise
log(f"os.replace failed ({exc!r}); falling back to file-by-file copy of staging tree")

9
studio/install_python_stack.py
View file

@ -1028,8 +1028,8 @@ def _detect_cuda_torch_index_url() -> str:
(e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback).
"""
# Explicit override (parity with install.sh / install.ps1):
# UNSLOTH_TORCH_INDEX_FAMILY=cu128|cu130|cu126|cpu|... pins the wheel
# index when probing is wrong or impossible (no GPU at build time, CI).
# UNSLOTH_TORCH_INDEX_FAMILY=cu128|cu130|cu126|cpu|... pins the wheel index
# when probing is wrong or impossible (no GPU at build time, CI).
family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY")
if family:
return f"{_PYTORCH_WHL_BASE}/{family}"
@ -2070,9 +2070,8 @@ def install_python_stack() -> int:
# --local overlays a local repo checkout after updating deps.
local_repo = os.environ.get("STUDIO_LOCAL_REPO", "")
# unsloth-zoo git ref for the --local overlay. Honor UNSLOTH_ZOO_REF (the
# Docker publish workflow / unsloth-studio-update resolve one ref and forward
# it) so the Studio venv can track the operator-requested zoo instead of
# always main. Unset -> main, byte-identical to the previous bare git URL.
# publish workflow / unsloth-studio-update forward one ref) so the Studio venv
# tracks the requested zoo, not always main. Unset -> main.
zoo_ref = os.environ.get("UNSLOTH_ZOO_REF", "").strip() or "main"
zoo_git_spec = "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo@" + zoo_ref
base_total = 11 if IS_WINDOWS else 12 # +1 for the anyio repair check (step 8b)

10
tests/python/test_unsloth_pip_shim.py
View file

@ -84,12 +84,10 @@ def _run(shim, tool, args):
execd = None
except _Exec as exc:
# main() builds [REAL[tool]] + head + keep_args + the protected
# constraints pair; head ends with the `install` verb, so everything
# after it is what we asserted on. The trailing
# `--constraint <unsloth-nb-protected-*.txt>` pair is injected on
# EVERY forwarded install (resolver-level protection); strip it here
# so each test asserts on its own arguments -- the dedicated
# constraint-injection tests below cover the pair itself.
# constraints pair; head ends with `install`, so everything after it
# is what we assert on. The trailing `--constraint <...>` pair is
# injected on EVERY install; strip it here so each test asserts on its
# own args (dedicated tests below cover the pair).
i = exc.argv.index("install")
execd = exc.argv[i + 1 :]
if (

25
tests/sh/test_select_cuda_jit_tools.sh
View file

@ -3,12 +3,10 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Unit tests for select_cuda_jit_tools() from docker/entrypoint.sh.
#
# cu12.8 is the immutable baked default (libnvrtc.so.12 -> .cu128.orig); the cu13
# tools are switched on ONLY for sm_103 (B300 / GB300) and sm_121 (GB10 / DGX
# Spark), which ship on >= 580 drivers -- see the rationale in docker/entrypoint.sh.
# The function picks per device via nvidia-smi compute_cap: those two arches
# retarget libnvrtc.so.12 -> the staged .cu13 alias (and point Triton at cu13
# ptxas); every other arch keeps the cu12.8 default and leaves ptxas unset.
# cu12.8 is the immutable baked default; the cu13 tools switch on ONLY for sm_103
# and sm_121 (>= 580 drivers). The function picks per device via nvidia-smi
# compute_cap: those two arches retarget libnvrtc.so.12 -> the .cu13 alias (and
# point Triton at cu13 ptxas); every other arch keeps cu12.8 and leaves ptxas unset.
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
@ -32,16 +30,11 @@ assert_eq() {
}
# $1 = compute_cap(s) the mock nvidia-smi reports, ONE PER LINE ("none" -> no
# nvidia-smi on PATH). A multi-line value models a mixed-GPU host so we can check
# that every visible cap is scanned, not just the first.
# $2 (optional) = the target libnvrtc.so.12 starts on; defaults to the baked
# cu12.8 default, and "libnvrtc.so.12.cu13" models the stale link an earlier
# sm_103/sm_121 boot left in the same container's writable layer.
# Builds a fake Studio venv NVRTC dir exactly as the build stages it: the real
# cu12.8 lib as .cu128.orig, libnvrtc.so.12 -> it (the immutable default), and a
# .cu13 alias pointing at a stand-in cu13 lib. Runs the function against it via
# UNSLOTH_STUDIO_HOME. The hardcoded base venv path does not exist on the test
# host, so its glob is skipped. Prints "<PTXAS_STATE> <NVRTC_TARGET>".
# nvidia-smi). A multi-line value models a mixed-GPU host (checks every cap is
# scanned). $2 (optional) = the target libnvrtc.so.12 starts on; defaults to the
# cu12.8 default, "libnvrtc.so.12.cu13" models a stale link from an earlier boot.
# Builds a fake Studio venv NVRTC dir as the build stages it and runs the function
# via UNSLOTH_STUDIO_HOME. Prints "<PTXAS_STATE> <NVRTC_TARGET>".
run_select() {
_cap="$1"
_init="${2:-libnvrtc.so.12.cu128.orig}"

43
unsloth/_gpu_init.py
View file

@ -114,18 +114,13 @@ del maybe_set_windows_rocm_bnb_version
# Fixes https://github.com/unslothai/unsloth/issues/1266
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
# Containers launched with `docker --gpus '"device=N"'` only set
# NVIDIA_VISIBLE_DEVICES to specific device ids/UUIDs and leave
# CUDA_VISIBLE_DEVICES absent. Inductor's compile worker subprocess pool
# then fails to enumerate the cgroup-pinned GPU and raises
# `Could not find an active GPU backend` from
# torch/_inductor/runtime/triton_helpers.py::set_driver_to_gpu. Force a
# single in-process compile thread so the pool is never spawned.
#
# Gate only on the cgroup-pinned fingerprint -- specific device ids in
# NVIDIA_VISIBLE_DEVICES. NVIDIA_VISIBLE_DEVICES in {"all","none","void",""}
# (the default in `--gpus all` runs) must NOT trigger this.
# Set UNSLOTH_FORCE_SINGLE_COMPILE_WORKER=0 to opt out.
# `docker --gpus '"device=N"'` sets only NVIDIA_VISIBLE_DEVICES to specific ids
# and leaves CUDA_VISIBLE_DEVICES absent, so Inductor's compile-worker pool can't
# enumerate the cgroup-pinned GPU and raises "Could not find an active GPU
# backend". Force a single in-process compile thread so the pool never spawns.
# Gate only on the cgroup-pinned fingerprint (specific ids); "all"/"none"/"void"/""
# (the `--gpus all` default) must NOT trigger it. Opt out with
# UNSLOTH_FORCE_SINGLE_COMPILE_WORKER=0.
_nvd = os.environ.get("NVIDIA_VISIBLE_DEVICES", "").strip().lower()
_cgroup_pinned = _nvd not in ("", "all", "none", "void")
if (
@ -133,9 +128,8 @@ if (
and _cgroup_pinned
and "CUDA_VISIBLE_DEVICES" not in os.environ
):
# Either set the env var if absent, or honour the user's existing
# value -- but always plant the sentinel so the zoo-side patch knows
# to preserve the forcing.
# Set the env var if absent (honour an existing value), but always plant the
# sentinel so the zoo-side patch preserves the forcing.
if os.environ.get("TORCHINDUCTOR_COMPILE_THREADS") in (None, "", "1"):
os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
os.environ["UNSLOTH_FORCE_SINGLE_COMPILE_WORKER"] = "1"
@ -185,15 +179,11 @@ except ModuleNotFoundError:
except:
raise
# Re-assert the single-compile-worker policy after unsloth_zoo has had a
# chance to run its patch_torch_compile (which historically popped
# TORCHINDUCTOR_COMPILE_THREADS in non-debug mode). Force the Inductor
# config value directly so the Docker --gpus '"device=N"' subprocess-pool
# bug is fixed even when the installed unsloth_zoo predates the
# corresponding zoo-side patch. Also monkey-patch the zoo's
# `determine_compile_threads` so the Inductor options dict (rebuilt per
# `torch.compile` call) always sees 1 even if a later import path pops the
# env var again. No-op when the user opted out.
# Re-assert the single-compile-worker policy after unsloth_zoo's
# patch_torch_compile (which historically popped TORCHINDUCTOR_COMPILE_THREADS).
# Force the Inductor config directly so the bug is fixed even against an older
# unsloth_zoo, and monkey-patch the zoo's determine_compile_threads so the
# per-call options dict always sees 1. No-op when the user opted out.
if os.environ.get("UNSLOTH_FORCE_SINGLE_COMPILE_WORKER", "0") == "1":
try:
torch._inductor.config.compile_threads = 1
@ -313,9 +303,8 @@ del patch_accelerate_recursively_apply
# Torch 2.4 has including_emulation
if DEVICE_TYPE == "cuda" and not torch.cuda.is_available():
# UNSLOTH_ALLOW_CPU=1 keeps DEVICE_TYPE "cuda" on driverless hosts (CPU
# CI, Docker Desktop without GPU passthrough); probing the device would
# raise. bf16 stays on: CPU bf16 kernels exist, fp16 ones largely do not.
# UNSLOTH_ALLOW_CPU=1 keeps DEVICE_TYPE "cuda" on driverless hosts; probing
# would raise. bf16 stays on (CPU bf16 kernels exist, fp16 largely don't).
SUPPORTS_BFLOAT16 = True
torch.cuda.is_bf16_supported = lambda *args, **kwargs: True
elif DEVICE_TYPE == "cuda":

12
unsloth/dataprep/synthetic.py
View file

@ -285,19 +285,17 @@ class SyntheticDataKit:
keep_lines = 2000,
echo = False,
name = "vLLM STDERR",
# vLLM >= 0.19 emits "Starting vLLM API server ... on ..." (and
# the uvicorn startup lines) through the logging module, which
# writes to STDERR. Watching stdout alone makes a healthy server
# look like a startup timeout, after which we kill it.
# vLLM >= 0.19 emits the startup lines through logging, which writes
# to STDERR; watching stdout alone makes a healthy server look like a
# timeout and get killed.
ready_regex = ready_re,
text = False,
)
# we don't print stderr to console but self.stderr_capture.tail(200) will print the last 200 lines
ready = False
# timeout = None (or 0) preserves the previous Event.wait(None) escape
# hatch: wait indefinitely for the readiness message (useful for large
# models or slow first-time downloads). Any positive value is a deadline.
# timeout None/0 keeps the previous Event.wait(None): wait indefinitely
# for readiness (large models / slow downloads). A positive value is a deadline.
deadline = (time.monotonic() + timeout) if timeout else None
while deadline is None or time.monotonic() < deadline:
if self.stdout_capture.wait_for_ready(timeout = 1) or self.stderr_capture.wait_for_ready(

9
unsloth/models/vision.py
View file

@ -398,11 +398,10 @@ def unsloth_base_fast_generate(self, *args, **kwargs):
):
kwargs.pop("mm_token_type_ids", None)
# VLMs do not allow logits_to_keep.
# transformers >= 5.0 sets logits_to_keep=1 itself in GenerationMixin.generate
# (utils.py:2527) AFTER _validate_model_kwargs runs, so pre-injecting it here
# makes the strict validator raise ValueError on PEFT-wrapped models. Skip on
# v5+ and let HF handle it. Strip any leaked kwarg defensively.
# VLMs do not allow logits_to_keep. transformers >= 5.0 sets logits_to_keep=1
# itself in GenerationMixin.generate AFTER _validate_model_kwargs, so pre-
# injecting it makes the strict validator raise on PEFT models. Skip on v5+
# and strip any leaked kwarg defensively.
if Version(transformers_version) < Version("5.0.0.dev0"):
global NUM_LOGITS_TO_KEEP
if arch not in NUM_LOGITS_TO_KEEP: