Stacks a Colab-like JupyterLab and Studio experience on top of the
existing Blackwell image. Additive only: the training stack, CUDA/torch
pinning, and the Studio/JupyterLab/sshd service trio are unchanged.
JupyterLab labextension (prebuilt in a throwaway builder stage, so the
runtime image stays Node-free):
- Unsloth Dark (Monokai) theme, adaptive light/dark by system preference
- Colab-style ArrowDown/Up cell navigation
- top-bar Unsloth logo (stock Jupyter logo disabled and locked)
- #@title lines render as collapsible Heading-2 form bars
- Ctrl+A in a cell output selects only that output, not the whole
notebook (the old behaviour ran notebook:select-all and was laggy)
- right activity bar hidden by default
- overrides.json: per-cell run button without auto-advance, labeled
Restart and Run All, windowing off so collapsing an output does not
snap to the cell top, news/update prompts suppressed
Studio and login branding: Unsloth favicon, page logo, and a dark
Unsloth login page that rotates through the curated Studio sloth
stickers (fail-soft to the logo).
Notebook organization and Colab compatibility (base image):
- categorized folder view built from relative symlinks mirroring the
README sections, rebuilt each boot; real .ipynb files never moved,
and the symlink tree is invisible to the sync state machine
- AMD-* notebooks shown only on an AMD/HIP host (autodetected)
- Docker-only strip of the Colab "Run all on Colab" intro sentence
from unedited notebooks (upstream notebooks unchanged)
- hoist %%capture above a leading #@title form so the cell runs
- the per-cell transformers-sidecar log is silent unless
UNSLOTH_ENABLE_LOGGING=1
Dependency pinning and naming: the curated notebook extras are pinned to
their resolved versions for reproducible rebuilds; decord is split into
its own fail-soft install (no aarch64 wheel). The lean base image is
renamed from :base to :core.
Adds tests/validate_studio_features.py, a static self-test for the
labextension plugins, overrides keys, and branding wiring.
Two related changes to the baked llama.cpp prebuilt.
1. Dynamically follow the newest unslothai/llama.cpp release. build.sh resolves
the latest release tag (following the /releases/latest redirect, no API
token) to a concrete tag and passes it as LLAMA_PREBUILT_TAG, so the layer
cache busts only when upstream publishes. The Dockerfile default is now
"latest" and fetch_llama_prebuilt.py resolves it the same way, so a plain
`docker build .` also tracks latest. Pin LLAMA_PREBUILT_TAG to a concrete
tag for a reproducible, frozen build.
2. Make the in-app "newer llama.cpp available" banner work inside the image.
Studio's freshness check (utils.llama_cpp_freshness.check_prebuilt_freshness)
keys off tag / release_tag / published_repo in UNSLOTH_PREBUILT_INFO.json --
the schema install_llama_prebuilt.py writes. The image bakes the bundle
directly, so the marker was the release tarball's own, which only carries
upstream_tag / source_repo; the freshness check then bailed with
installed_tag=None and could never report "behind", hiding the banner.
fetch_llama_prebuilt.py now augments the baked marker with those keys
(setdefault, no build timestamp so the layer stays byte-identical). A fresh
build is on latest -> no banner; once upstream publishes a newer release the
banner appears, as verified against the real freshness backend.
The boot-time refresh now compares only the tutorial body (the
non-boilerplate cells) when deciding whether to update an untouched
notebook. If only the install header, announcements, or footer moved
upstream, the user's file is left as-is so it is not churned. Notebooks
the user has edited or run are still kept untouched, and non-notebook
files keep the whole-file refresh. Adds unsloth_nb_content_sig.py to
segment head/middle/tail and bakes it into the image.
JupyterLab now opens with the unslothai/notebooks collection already present,
so people can open and run a notebook directly without a git clone or wget.
- Bake the repo into the image as a read-only template at /opt/unsloth-notebooks
(~206MB, .git stripped, build commit recorded). Inherited by the studio image.
- On boot the entrypoint populates /workspace/unsloth-notebooks from the template
(instant, works offline) and best-effort refreshes from GitHub, but only when
upstream has actually advanced (cheap git ls-remote gate, no download otherwise).
- The user's edits always win. We record the content hash of every file we write;
on refresh a file whose hash differs from what we last wrote is treated as
user-modified and is left untouched, so the refresh only updates files the user
has not changed and adds new ones. It never overwrites an edited notebook and
never produces merge conflicts. Verified: an edited notebook stays the user's
version across repeated upstream changes.
- Fully best-effort and gated: UNSLOTH_SKIP_NOTEBOOK_SYNC=1 disables it,
UNSLOTH_SKIP_NOTEBOOK_REFRESH=1 keeps the baked copy and never hits the network.
Offline boots keep what is there and never error.
base 18.45 -> 18.67GB, studio 24.88 -> 25.10GB (+~206MB baked notebooks).
Notebook dependency coverage (base Dockerfile):
- Bake omegaconf, einx, librosa, decord, ftfy so the TTS/STT and vision
notebooks stop dying on a silent No module named X. Installed in the
notebook-deps layer (after the torch/vLLM resolve) with an assertion that
the resolve did not move torch 2.10.0 / numpy>=2.3 / numba>=0.65.
Image size (no functional change):
- Base: prune npp to the two libs torchcodec actually dlopens
(libnppicc + libnppc), drop link-time-only .a archives and the nvshmem
device bitcode. Headers (torch/include etc) are kept so causal-conv1d /
mamba-ssm still build at notebook time with --no-build-isolation.
- Studio: pin the Studio venv to Python 3.12 (matches base) so its
nvidia-*-cu12 wheels are byte-identical to the base venv's, then symlink
the heavy arch-independent CUDA libs (cudnn/cublas/nccl/...) into the base
venv copy. cuda_nvrtc and cuda_runtime are excluded (the arm64 nvrtc swap
mutates nvrtc in place). Also remove the build-only frontend node_modules
(runtime serves the committed dist). Studio image drops ~4.8GB.
Per-notebook transformers version, run notebooks unchanged:
- Bake coherent transformers sidecars (4.57.6 default + 5.3.0/5.5.0/5.10.2),
each transformers==X with its matched huggingface_hub/tokenizers/
safetensors installed --no-deps into its own dir. Companion versions are
resolved at build time so they satisfy each transformers' requirements.
- unsloth_nb_compat.py: pick the sidecar from the notebook's pin or the
model name and activate it (prepend to sys.path) before any ML import,
without touching the base cu128 torch/vLLM/unsloth stack.
- pip/uv shim on PATH: a notebook install cell becomes safe and idempotent
inside a kernel (keeps the baked stack, records the requested transformers
for its sidecar); passthrough to the real tool everywhere else.
- IPython startup hook for manual JupyterLab, and unsloth-run for the
headless driven path.
Continuation of the notebook-dep prebaking: the in-image notebook runner
neutralises pip cells, so declared deps must be prebaked. evaluate's WER
metric imports jiwer (Whisper), DeepSeek-R1 GRPO's reward uses langid,
some vision trust_remote_code files need easydict, and sentencepiece
tokenizer conversion needs protobuf. All pure-Python; torch pin intact.
The TTS notebooks (Sesame CSM, Orpheus) read audio via soundfile, the
Whisper notebook computes WER via evaluate, and TrainingArguments
defaults report_to to tensorboard. These are declared by notebook pip
cells that the in-image notebook runner neutralises (deps are meant to be
prebaked), so without them those notebooks die on import. All are
pure-Python or self-contained wheels and never name torch, so the cu128
pin set is undisturbed.
flash-linear-attention's TileLang backend JIT-compiles CUDA kernels via
nvcc at runtime for gated-delta-rule models (Qwen3.5 family). The -base
image only ships runtime libraries, so Studio vision training of
unsloth/Qwen3.5-2B died on the first backward pass with
[Errno 2] No such file or directory: /usr/local/cuda/bin/nvcc.
Install cuda-nvcc and cuda-cudart-dev matching the image CUDA version
and assert nvcc is executable at build time. Found by driving a real
Qwen3.5-2B training run through the Studio UI in the image.
The notebook validation matrix caught the synthetic-data notebook dying
because the vllm server SyntheticDataKit launches never came up. Two
layers to the failure:
1. flashinfer's cpp_ext JIT shells out to ninja. The pip ninja lives in
the venv bin, which subprocesses like vllm serve do not always
inherit on PATH, so the JIT failed with exit 127. Install ninja-build
so the binary is reachable from any PATH.
2. With ninja present the JIT still cannot succeed for device code: the
runtime image deliberately ships no nvcc. Bake flashinfer-jit-cache
(cu128) so ops missing from the cubin package (fmha_gen on sm_100a
was the repro) come precompiled. In-process GRPO never hit this
because unsloth-zoo blocks the FlashInfer JIT path; standalone
vllm serve gets no zoo patches.
Fail-soft on the jit-cache for arches without a wheel; the vLLM chain
itself stays fail-loud on amd64.
Notebooks fetch sample assets with !wget; without the binary the shell
prints not-found to stderr, the cell still exits zero from Jupyter's
perspective, and the next cell crashes confusingly on the missing file.
The Whisper notebook died exactly this way in the validation matrix.
PyPI has shipped aarch64 abi3 wheels for every vLLM release since 0.17,
so the arm64 skip rested on a stale premise. With torch held at 2.10.0
the resolver lands on vllm 0.19.1 (the release pinning torch==2.10.0)
on both arches; verified by cross-resolving the exact index set for
aarch64-unknown-linux-gnu.
amd64 keeps fail-loud semantics. arm64 is fail-soft because the aarch64
wheels are newer and their GPU kernels get validated on Spark hardware
via docker_confirm.sh rather than in CI; on failure the fallback
uninstalls vllm and restores the numpy/numba floor so a partial install
cannot break import unsloth (numpy 2.2.6 ships a broken numpy.testing).
The install steps form an explicit && chain instead of a set -e
subshell: POSIX shells disable errexit inside condition contexts
(verified on dash), so a (set -e; ...) condition would mask failures.
Both confirm scripts gain a 5b vLLM phase: ok on import, bad if missing
on x86_64, warn on other arches where fast_inference=True is best-effort.
The wheel install belongs in the builder (the venv copy carries it),
but the ld.so.conf.d registration and the import check belong in the
runtime stage: the conf file does not survive the stage copy and the
import needs ffmpeg, which only the runtime stage installs.
The TTS/STT notebooks decode datasets Audio features through torchcodec,
which fails three different ways on a fresh image: the PyPI wheel pairs
with the cu13 torch line and dlopens libnvrtc.so.13; builds newer than
0.10 reference torch 2.11+ symbols; and the matching +cu128 build dlopens
torch and NVIDIA runtime libraries that live inside the venv where the
loader cannot see them. Bake ffmpeg, torchcodec==0.10.0 from the cu128
channel, nvidia-npp-cu12, and register the venv lib dirs via ld.so.conf.d
(not LD_LIBRARY_PATH, so the llama.cpp bundle keeps winning through its
own RUNPATH). Verified in-container: AudioDecoder imports and llama-server
still resolves its bundled libraries.
Running the published unslothai/notebooks set inside the image surfaced
two gaps: the Ollama export notebook installs ollama in-container and
that installer needs zstd for extraction, and DeepSeek-OCR's
trust_remote_code modeling file imports matplotlib unconditionally
(plotting is also simply expected in a Jupyter image).
- Dockerfile: lift numba past vllm's 0.61.2 pin after the numpy>=2.4
re-upgrade; 0.61.2 refuses numpy 2.3+ at import time and the stack
cannot move numpy down. Verified numba 0.65 + numpy 2.4.6 + vllm
import cleanly together.
- docker-publish.yml: resolve UNSLOTH_ZOO_REF in a step that mirrors
the pushed tag only when the tag exists in unsloth-zoo (the zoo
currently cuts no tags, so blind mirroring broke every tag publish);
falls back to main.
- Dockerfile.studio: Studio venv stays on cu128 for arm64 too, matching
the base venv (cu130 wheels would lift the driver floor to 580+), and
gets the same NVRTC cu13 swap for DGX Spark / GB10 sm_121 support.
- docker_confirm.sh: do not drop to CPU mode when docker info lacks a
nvidia runtime entry; CDI installs and Docker Desktop WSL2 expose
GPUs without one. The phase 3 --gpus probe is now the authority.
- docker_confirm.ps1: GPU selector built as an args array; comma device
lists get version-aware CSV quoting (native arg passing changed in
PowerShell 7.3).
- studio_launch.sh: no fixed Jupyter default password; generate a
random one and print it when JUPYTER_PASSWORD is unset. Env snapshot
for SSH sessions now written via shlex.quote instead of sed so
values with quotes or command substitution cannot break or inject
into /etc/profile.d.
- install.ps1: honour UNSLOTH_TORCH_INDEX_FAMILY like install.sh does.
The first bake attempt reused studio/install_llama_prebuilt.py, but that
resolver selects a bundle for the CURRENT host: on a GPU build host
/proc/driver/nvidia leaks into docker build and the resolver goes down the
CUDA path with no readable driver runtime (chosen_asset=none, exit 2),
while on a GPU-less CI runner it would resolve a CPU bundle instead. Both
violate the image's build-host-independence rule.
fetch_llama_prebuilt.py pins by build target only: amd64 takes the
linux-x64-cuda12-portable bundle, arm64 the linux-arm64-cuda13-portable
bundle (DGX Spark / Grace), both sha256-verified against the release's
llama-prebuilt-sha256.json. convert_hf_to_gguf.py plus gguf-py/ are
hydrated from the same release's source tarball so the converter's tensor
mappings match the binaries, mirroring unsloth_zoo's
_hydrate_converter_sources layout. LLAMA_PREBUILT_TAG build-arg overrides
the pinned release.
docker_confirm.sh: one-command confirmation script for any machine
(Linux / WSL2 / macOS) following the staging confirm-script conventions:
host + docker + GPU detection with CPU-mode auto-fallback, image pulls,
in-container torch.cuda check, 5-step LoRA training smoke, baked llama.cpp
verification, full-image boot probing Studio /api/health and JupyterLab
/api, PASS/WARN/FAIL summary with RESULT line.
Base image (docker/Dockerfile):
- Install JupyterLab + notebook + ipywidgets in a separate pure-Python uv
pass so the cu128 pin set cannot move; EXPOSE 8888.
- Bake the prebuilt llama.cpp bundle into /opt/unsloth/llama.cpp at the
runtime stage using studio/install_llama_prebuilt.py from the same
UNSLOTH_REF (sha256-verified, portable CUDA bundle since the build host
has no GPU; arm64 resolves the linux-arm64-cuda13 bundle). Export
UNSLOTH_LLAMA_CPP_PATH so unsloth_zoo's save_pretrained_gguf finds it
and never reaches the interactive install prompt or a source build.
- Optional github_token BuildKit secret for the resolver's API calls on
shared CI runner IPs.
Entrypoint: UNSLOTH_ALLOW_CPU=1 degrades a missing GPU to a warning so
Docker Desktop on macOS / Windows-without-WSL2-GPU and plain CPU hosts can
run Jupyter, GGUF tooling and Studio chat; with a GPU visible the normal
pre-flight still runs.
Full image (docker/Dockerfile.studio): now mirrors the production service
set under supervisord - Studio on 8000, JupyterLab on 8888, key-only sshd
on 22 (enabled only when PUBLIC_KEY/SSH_KEY is set). Points Studio's
llama.cpp dir at the baked bundle to skip a duplicate download, accepts
any git ref via fetch+checkout (CI passes commit SHAs), and FROMs a
digest-pinned BASE_IMAGE.
Publish workflow: base image moves to the base-* tag namespace; new
build-studio/merge-studio jobs publish the full image as :latest (hub
parity with the previous production image, which shipped Studio + Jupyter
+ SSH). Studio builds FROM the exact base manifest digest published by the
same run. GPU smoke job now also boots the full image and probes Studio
/api/health and Jupyter /api.
run.sh: UNSLOTH_GPUS=none, UNSLOTH_ALLOW_CPU forwarding, UNSLOTH_PORTS
publish flags, CPU-mode and Jupyter usage examples.
- docker-publish.yml: add `concurrency: docker-publish-${{ github.ref }}`
(cancel-in-progress: false) so two pushes to main never race the
`:latest` retag. Don't cancel in-progress runs -- the build is
expensive and a half-built image left around is worse than a stale
:latest for a few minutes.
- Dockerfile: soften the requirements.lock.txt comment. `pip freeze`
captures versions but not wheel hashes, and several deps resolve
from VCS / nightly indexes that float, so the file is not actually
byte-reproducible. Reword as an "informational pin record".
1. Stop leaking secrets via docker run -e VAR=VALUE argv (run.sh, test_locally.sh)
`docker run ... -e HF_TOKEN=hf_xxx ...` puts the literal token in
the docker CLI's argv, which is visible to any user on the host
via `ps auxe` / `/proc/<pid>/cmdline` for the lifetime of the
process. Switch to the dash-only form `-e HF_TOKEN`, which tells
docker to read the value from the parent shell's env and never
appears in argv. Same fix for WANDB_API_KEY and UNSLOTH_LICENSE in
run.sh and HF_TOKEN in test_locally.sh.
2. Stop stripping numpy/tests/ in the runtime layer (Dockerfile)
The Dockerfile explicitly upgrades numpy >= 2.4 because numpy 2.2.6
shipped a stripped wheel where `from numpy._core.tests._natype
import pd_NA` fails. Numpy 2.4 restores `numpy/_core/tests/`, then
the existing `find ${VENV} -name tests -exec rm -rf {} +` deleted
it again -- re-introducing the same broken-import state on the
deployed image (the build-time verification at line 220 runs
BEFORE the strip so it passed). Whitelist numpy's tests directories
from the strip; keep stripping the rest.
3. Align :latest tag gate between merge and smoke-test jobs
(.github/workflows/docker-publish.yml)
merge job: enable = is-default-branch AND unsloth_ref == ''
smoke-test job: enable = is_default_branch only
On `workflow_dispatch`, `github.event.inputs.unsloth_ref` defaults to
"main" (not ""), so the merge step skipped `:latest` but the smoke
step still emitted `:latest` as tags[0]. The smoke step then
`docker pull`-ed a prior `:latest` from Docker Hub instead of the
image just merged -- so the smoke test verified the OLD image, not
the new one. Copy the merge step's exact `enable=` expression into
the smoke-test step so the two stay byte-identical and a workflow_
dispatch run validates whatever was actually merged.
Round-2 of the 12-persona reviewer.py pass found 17 issues. Address the
P1s + the regression-class P2s in this commit; the remaining nits are
left for a follow-up cleanup pass.
1. unsloth/_gpu_init.py: the `NVIDIA_VISIBLE_DEVICES in os.environ` check
triggered for every NVIDIA-runtime container including `--gpus all`
(NVIDIA_VISIBLE_DEVICES=all is the default). Gate strictly on a
non-special device list. Also drop the precondition that the env var
was absent: if the user already pinned TORCHINDUCTOR_COMPILE_THREADS=1
we should still plant the UNSLOTH_FORCE_SINGLE_COMPILE_WORKER sentinel
so the zoo-side patch knows to preserve the forcing.
2. unsloth/_gpu_init.py: after the post-`import unsloth_zoo` reassertion,
monkey-patch `unsloth_zoo.temporary_patches.common.determine_compile_threads`
to return 1, so any later `torch.compile` call that rebuilds the
options dict still sees the single-worker forcing even if a downstream
patch_torch_compile pops the env var again.
3. docker/Dockerfile: torchaudio==2.11.0 mismatched the torch==2.10.0
release pairing; pin to 2.10.0 so the ABI is correct and the audio
stack matches torch/cu128.
4. docker/Dockerfile: drop `12.1+PTX` from TORCH_CUDA_ARCH_LIST. The
cu128 toolkit compiler does not know about compute_121; the trailing
PTX entry forced nvcc to emit a `sm_121` gencode that breaks any
in-container source builds.
5. docker/smoke_test.py: the device-capability floor said `cap[0] < 8`,
rejecting Turing (sm_75) while the Dockerfile + entrypoint advertise
sm_75 as supported. Lower the smoke floor to sm_75 and print a hint
that bf16 is not available on Turing.
6. docker/run.sh: `-it` is unconditional; CI / non-TTY invocations died
with "the input device is not a TTY". Probe `[ -t 0 ] && [ -t 1 ]`
first. Also remove `set -x` which echoed the forwarded HF_TOKEN /
WANDB_API_KEY / UNSLOTH_LICENSE values to stdout.
7. docker/test_locally.sh: `-e HF_TOKEN="${HF_TOKEN:-}"` either pasted
the secret verbatim into the process arg list or shadowed any
in-container value with an empty string. Forward conditionally.
8. .github/workflows/docker-publish.yml: gate `latest` on default branch
AND on `unsloth_ref` not being overridden via workflow_dispatch.
Otherwise a maintainer testing a feature SHA from main could overwrite
`:latest` with non-main source.
9. docker/Dockerfile.studio: add an `UNSLOTH_STUDIO_REF` build-arg so
the Studio companion image is pinned to a known unsloth ref instead
of cloning `main` whenever it builds.
Two vision-notebook deps that ship by reference rather than via unsloth
extras: transformers' Gemma3N + TimmWrapperModel needs `timm`, and
DeepSeek-OCR's dynamic modeling file requires `addict`. Both are tiny
(~30MB combined). Including them in the base unified resolve avoids
hitting `ImportError: TimmWrapperModel requires the timm library`
or `ImportError: This modeling file requires the following packages
that were not found in your environment: addict` after the user has
already downloaded the model.
Repros: nb/Gemma3N_(4B)-Vision.ipynb (timm), nb/Deepseek_OCR_(3B).ipynb
(addict).
vLLM 0.19.1 pulls numpy down to 2.2.6 whose wheel ships numpy/_core/
without the tests/ subdir, but numpy/testing/_private/utils.py imports
`from numpy._core.tests._natype import pd_NA`. Anything that hits
`from numpy import *` (scipy._lib.array_api_compat does) then crashes.
unsloth_zoo's gemma patch does `from transformers.processing_utils import
Unpack` which touches that path, so `import unsloth` blew up on every
GRPO notebook in the vLLM image.
Bump numpy to >=2.4 right after the vllm install; vllm still imports
fine on numpy 2.4.6 (verified locally).
The case-arm `auto) [ "${TARGETARCH}" = "amd64" ] && WANT_VLLM=1` exits 1
on arm64 (the [ test ] is false and nothing follows ||), which with
`set -e` aborts the entire RUN. Replace with an explicit if/then/fi so
each arch's auto branch returns 0.
Caught by ubuntu-24.04-arm CI on the staging fork.
Unsloth's GRPO notebooks (Qwen3_4B-GRPO.ipynb, Qwen3_8B_FP8_GRPO.ipynb,
Llama_FP8_GRPO.ipynb, etc.) set `fast_inference=True` which requires
vLLM to be importable in the same venv. Install vllm pre-release wheels
from https://wheels.vllm.ai/nightly alongside the cu128 pytorch index,
holding torch==2.10.0 fixed so uv refuses any vLLM build that would
yank torch out from under unsloth.
amd64 only -- vLLM does not publish aarch64 wheels yet
(vllm-project/vllm#31128 is open). On arm64 the GRPO notebooks that
need fast_inference will fail to import vllm; non-GRPO and
fast_inference=False paths are unaffected.
Gated by ARG INSTALL_VLLM=auto so the install can be disabled for
contributors who want a smaller image or are blocked by vllm/torch
resolve conflicts during iteration.
The nvidia/cuda base image already registers the CUDA apt repo with its own
Signed-By keyring. Installing cuda-keyring_1.1-1_all.deb on top adds a
duplicate sources entry with a different Signed-By value, which makes
`apt-get update` refuse the entire repo:
E: Conflicting values set for option Signed-By regarding source
https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/sbsa/
The repo URL is monolithic (every CUDA version is served from the same
path), so we can install cuda-nvrtc-13-0 + cuda-nvcc-13-0 directly without
touching the keyring. Empirically reproduced on the ubuntu-24.04-arm
GitHub Actions runner (staging-fork CI run 26360461375); fix verified via
the same staging-fork after force-push.
torch wheels ship their own cuDNN/cuBLAS/cuSPARSE/cuRAND/cuSOLVER/cuFFT/
NCCL/cuSparseLt inside torch/lib/, and libtorch_cuda.so's RPATH
($ORIGIN/../../nvidia/cudnn/lib:$ORIGIN/../../nvidia/cublas/lib:...)
points at those wheel-bundled copies. The dynamic loader resolves through
the wheel, never the system, so the libcudnn/libcublas in the system
cudnn-runtime layer are unreachable code on every pull.
Verified empirically via `readelf -d torch/lib/libtorch_cuda.so` and
confirmed bitsandbytes' NEEDED list resolves against torch's bundled
libcudart/libcublas/libcublasLt/libcusparse/libnvJitLink before bnb
loads. Triton's .so files have zero CUDA NEEDED entries -- they dlopen
through the host driver.
Compressed image saving: ~2.7 GB (cudnn-runtime base 2.86 GB -> base
0.10 GB, on amd64; arm64 similar). Uncompressed: ~5 GB. Zero functional
impact.
Source: Fork 5 image-size audit, May 2026.
Empirically (cu128 wheel SASS list `sm_80;90;90a;100;100a;120;120a` on
aarch64) the cu128 wheel covers DGX Spark sm_121 via sm_120 binary
forward-compat. BUT two CPU-side compilers shipped at cu12.8 do not know
sm_121 and need a cu13 swap:
(1) torch's bundled libnvrtc.so.12 from CUDA 12.8 rejects sm_121 as a
--gpu-architecture. Symlinks libnvrtc.so.13 over it.
(2) Triton's nvidia backend runs ptxas. Wheels older than 3.6.0 bundled
cu12.8 ptxas which silently downgrades sm_121 to sm_80 (see
triton-lang/triton#8335). Bump pin triton>=3.6.0 (3.6 bundles cu13
ptxas) AND install cuda-nvcc-13-0 so the entrypoint can point
TRITON_PTXAS_PATH at it as defense in depth.
Both fixes are arm64-only (gated on TARGETARCH, ~400 MB on the arm64
image; amd64 is untouched, no sm_121 hardware exists on x86_64). Neither
component talks to libcuda, so this does NOT bump the toolkit driver
floor away from cu128's 570+.
TRITON_PTXAS_PATH is set from the entrypoint (only when the cu13 ptxas
actually exists in the image) rather than via a Dockerfile ENV, because
ENV is unconditional and Triton errors out if TRITON_PTXAS_PATH points
at a nonexistent file.
Sources: martimramos/dgx-spark-ml-guide Challenge 14; triton-lang/triton
issue #8335; ptrblck PyTorch forum thread on sm_121 fwd-compat from
sm_120.
Empirical reality (cuobjdump on the downloaded cu128 wheels):
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
Earlier comments claimed sm_89 native and a "+PTX JIT to sm_121" fallback;
both are wrong. cu128 wheels ship NO PTX. Ada (sm_89) runs on sm_86 SASS,
B300/GB300 (sm_103) on sm_100, DGX Spark (sm_121) on sm_120 -- all
forward-compat WITHIN a major architecture, which is the canonical CUDA
rule and ptrblck (PyTorch maintainer) confirmed it directly:
"the compatibility ... is also used for e.g. sm_89 with sm_86 and sm_80."
Build-time assertion was `any(a in ("sm_120", "sm_121"))` on arm64. Since
sm_121 is never in any cu128 wheel, the OR was misleading and could mask
a real wheel regression. Tightened to just `assert "sm_120" in arches`
on both arches.
GitHub announced free linux/arm64 hosted runners for public repos (GA Aug
2025) under labels `ubuntu-24.04-arm` / `ubuntu-22.04-arm`. Switching the
arm64 leg from QEMU-on-amd64 to a native arm64 matrix runner is ~3x
faster and avoids QEMU's occasional flakiness on long cu128 installs.
The workflow now:
* builds amd64 and arm64 in parallel on their native runners,
pushing each as a single-arch image *by digest* (no tag)
* stitches both digests into one multi-platform manifest in a
follow-up `merge` job, using `docker buildx imagetools create`
* keeps a separate buildx cache scope per platform to avoid
cross-arch cache collisions
Smoke-test job now needs `merge` (was `build`) so it only runs once the
final manifest is published.
Dockerfile header: replace the speculative aarch64 SASS list with the
verified one from pytorch/pytorch v2.10.0 .ci/manywheel/build_cuda.sh
(8.0;9.0;10.0;12.0 on aarch64), and note that sm_120 is forward-compatible
to sm_121 per PyTorch maintainers -- which is what makes DGX Spark work
without an explicit sm_121 SASS section in the wheel.
setup_qemu.sh / test_locally.sh --platform stay in place: they're for
the local-dev path on x86_64 boxes that don't have arm64 hardware.
Make the docker image multi-arch so DGX Spark (GB10, sm_121, aarch64) and
the Grace-Hopper / Grace-Blackwell SoCs (GH200 arm64, GB200 arm64) pull a
natively-built arm64 child from the same manifest. Runtime emulation is
NOT involved -- QEMU is used only for the cross-compile step on x86_64
CI runners; consumers on aarch64 hosts get a normal arm64 image and CUDA
works as on any other host.
Dockerfile:
* ARG TARGETARCH; switch unsloth extras between cu128-ampere-torch2100
(amd64, with xformers) and huggingface (arm64, no xformers -- there
is no cu128 aarch64 xformers wheel as of 0.0.34, so we fall back to
Unsloth's native SDPA path; ~5-10% slowdown but functionally complete).
* Build-time torch._C._cuda_getArchFlags() assertion: amd64 still
requires sm_120, arm64 accepts sm_120 or sm_121.
* Same TORCH_CUDA_ARCH_LIST on both arches; nvcc emits whatever's listed.
docker/setup_qemu.sh (new):
One-time host setup -- registers binfmt_misc handlers via
tonistiigi/binfmt and creates a 'unsloth-multiarch' docker-container
buildx builder. Required only on x86_64 build hosts targeting arm64.
docker/test_locally.sh:
--platform amd64|arm64 flag. Cross-builds verify QEMU is registered,
then build through the in-image arch-flags assertion. Smoke + notebook
blocks auto-skip when image arch != host arch (CUDA cannot run under
user-space QEMU + nvidia-container-toolkit cannot bridge a QEMU guest
to a real GPU).
.github/workflows/docker-publish.yml:
platforms: linux/amd64,linux/arm64 (single manifest, two children).
Timeout bumped 60 -> 150 min for the slower arm64-under-QEMU leg.
docker/setup-qemu-action@v3 with platforms: arm64 (was implicit before).
TORCH_CUDA_ARCH_LIST now covers the full set of compute capabilities
NVIDIA publishes on https://developer.nvidia.com/cuda/gpus for x86_64
hardware, from Turing onward:
sm_75 Turing T4, RTX 20-series, Quadro RTX
sm_80 Ampere DC A100, A30
sm_86 Ampere A40, RTX A6000, RTX 30-series
sm_89 Ada L4, L40, L40S, RTX 40-series
sm_90 Hopper H100, H200, GH200
sm_100 Blackwell DC B100, B200, GB200
sm_103 Blackwell DC B300, GB300
sm_120 Blackwell RTX 50-series, RTX PRO 6000 Blackwell
sm_121 Blackwell GB10 (DGX Spark)
with +PTX on the highest entry so future arch revisions can JIT.
Setting TORCH_CUDA_ARCH_LIST only affects nvcc invocations for any
source build the user adds on top of this image (e.g. flash-attn, a
custom CUDA op). The prebuilt cu128 wheels already include SASS for
sm_70/75/80/86/90/100/120 (verified at build time via
torch._C._cuda_getArchFlags()). Ada (sm_89), B300 (sm_103) and DGX
Spark (sm_121) GPUs run via JIT-PTX from the nearest available arch.
Jetson archs (sm_87 Orin, sm_110 Thor) are intentionally NOT included
-- they require aarch64 wheels and this image is linux/amd64 only.
Also lower the entrypoint's compute-capability gate from sm_80 to
sm_75. Turing GPUs work, with the caveat that bfloat16 is unavailable;
the entrypoint prints a NOTE in that case so Unsloth's fp16 fallback
isn't a surprise.
Triton's nvidia backend lazily JIT-compiles a small C extension
(CudaUtils, in triton/backends/nvidia/driver.py) on first GPU access.
Without a C compiler and Python headers in the runtime image, the
very first forward pass of any Unsloth model dies with:
RuntimeError: Failed to find C compiler.
Please specify via CC environment variable.
The builder stage has build-essential and python3.12-dev so this
worked during the build's verification step (no GPU = no Triton kernel
call = no C extension build). But the runtime stage stripped those
out for size, so the failure only surfaces when a real user runs
training inside the container.
Add gcc + g++ + python3.12-dev to the runtime stage. Increases the
runtime image by ~250MB, which is the cost of letting Triton JIT
correctly. Pre-compiling CudaUtils at build time would need a real
CUDA device (the constructor calls cuda runtime functions), so
shipping the toolchain is the right trade-off.
Ubuntu 24.04 (noble) marks the system Python interpreter as
externally-managed per PEP 668, so:
curl get-pip.py | python
python -m pip install -U pip uv
fails inside the builder image with:
error: externally-managed-environment
This environment is externally managed
The system-level pip and uv were never used: the very next RUN creates
the venv at /opt/unsloth-venv, which bootstraps its own pip via the
ensurepip module (provided by the python3.12-venv apt package). uv is
then installed INTO the venv with the venv's pip, and used from there.
Drop the two system-pip bootstrap lines. The venv path is unchanged.
Reproduces on any Docker build of the unsloth-blackwell image against
a noble base image (which our nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
is).
When someone launches the unsloth container, the common failure modes are not
unsloth bugs -- they're Docker / nvidia-container-toolkit / driver issues that
surface as cryptic CUDA errors deep in torch. The entrypoint catches the three
that cover ~95% of "it doesn't work" reports up front:
1. nvidia-smi inside the container sees no GPU
-> user forgot --gpus all, or host is missing nvidia-container-toolkit
-> entrypoint prints the exact docker run flag and the toolkit install URL
2. nvidia-smi works but torch.cuda.is_available() is False
-> host driver is older than CUDA 12.8 supports
-> entrypoint prints the minimum driver version per architecture
3. compute capability < sm_80
-> entrypoint prints the supported architecture table and exits
Each check fails with a clear, actionable message rather than a stack trace.
Set UNSLOTH_SKIP_GPU_CHECK=1 to bypass (for docs builds, offline tooling, CI).
run.sh wraps `docker run` with the flags people most often forget:
--gpus all (without it, the new entrypoint refuses to start)
--ipc=host (DataLoader workers need >64MB shm)
--ulimit memlock=-1 (NCCL + CUDA pinned host buffers)
--ulimit stack=64MB (some torch kernels OOM the default 8MB stack)
Plus it mounts the host HF cache + Triton JIT cache so model downloads and
compiled kernels persist across container runs, and forwards HF_TOKEN /
WANDB_API_KEY / UNSLOTH_LICENSE only when they are set on the host.
Usage:
bash docker/run.sh # interactive python REPL
bash docker/run.sh bash # shell in container
bash docker/run.sh python /workspace/smoke_test.py
bash docker/run.sh python /workspace/host/train.py # $PWD mounted at /workspace/host
Verified locally:
- No GPU visible: entrypoint refuses with driver-version message, exit 1
- B200 sm_100 visible: entrypoint prints GPU banner, exits cleanly into the
user command (rc=0)
Adds a multi-stage Dockerfile producing an image that works on Ampere through
Blackwell (sm_80 through sm_120: A100, RTX 30/40, H100, B100/B200, RTX 50-series,
RTX 6000 Pro Blackwell). The build itself requires no GPU at all and runs on a
free GitHub-hosted ubuntu-latest runner.
How the GPU-less build works:
1. cu128 PyTorch wheels are fat binaries. torch._C._cuda_getArchFlags() returns
'sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120' regardless of which GPU
compiled the image, because the wheels are cross-compiled upstream by the
PyTorch team.
2. All deps resolve in a single uv pip install pass with explicit pins
(torch==2.10.0, --extra-index-url cu128, no --torch-backend=auto, no
install.sh). This prevents the silent cu cascade where bitsandbytes'
transitive cuda-toolkit==13 dep upgrades torch to 2.12+cu130 in a later
resolver pass, leaving xformers and other cu128 wheels stranded.
3. Build-time verification uses package metadata (importlib.metadata.version)
and the raw torch._C._cuda_getArchFlags() accessor. We deliberately avoid
import unsloth at build time because unsloth.__init__ calls
torch.cuda.get_device_properties(0), which requires an actual CUDA device
and is not bypassable. Import-time correctness is exercised at deploy time
by smoke_test.py with --gpus all.
4. UNSLOTH_COMPILE_DISABLE=1 and CUDA_VISIBLE_DEVICES="" during the build stage
prevent any code path from JIT-compiling kernels for the build host's
compute capability and baking the resulting cache into the image. The
deploy GPU produces its own cache on first use.
Other notes:
- --index-strategy unsafe-best-match is needed because the PyTorch wheel index
serves an old requests==2.28.1 that conflicts with datasets>=2.32.2, which
the default first-index-wins strategy rejects.
- Extra is cu128-ampere-torch2100 (ampere precedes the torch version in the
pyproject ordering).
- No flash-attn in the base image. FA3 is hard-refused on Blackwell upstream
and unsloth gracefully falls back to xformers + SDPA. Users on Ampere /
Ada / Hopper who want FA2 can pip install flash-attn on top.
- Two stages: nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 for the build,
-cudnn-runtime for the deploy image. No nvcc in the published image.
- A lockfile is emitted at /opt/unsloth-venv/requirements.lock.txt inside
the image and can be extracted with docker/freeze.sh for byte-identical
rebuilds even after PyPI moves on.
CI workflow .github/workflows/docker-publish.yml:
- Builds on ubuntu-latest on every push to main, every tag, weekly via cron,
and manually via workflow_dispatch. Pushes to docker.io/unsloth/unsloth
with cache via type=gha.
- Optional smoke-test job runs on a self-hosted GPU runner if vars.HAS_GPU_RUNNER
is set; skipped otherwise. End-to-end verification on sm_120 hardware is a
nice-to-have, not a publish blocker.
Validation:
- Install path validated on a B200 host with CUDA_VISIBLE_DEVICES="" set
(simulating the GPU-less CI runner): torch 2.10.0+cu128 holds, xformers
0.0.34, bitsandbytes 0.49.2, triton 3.6.0, transformers 5.5.0, trl 0.24.0,
peft 0.19.1, accelerate 1.13.0. Arch flags include sm_100 and sm_120.
- Runtime path validated end-to-end on B200: smoke_test.py imports unsloth,
loads Llama-3.2-1B-Instruct-bnb-4bit in 4-bit, completes 5 LoRA steps with
loss decreasing 4.11 -> 3.75. xformers fallback active as designed.
Files:
- docker/Dockerfile multi-stage cu128 build
- docker/build.sh local build wrapper
- docker/freeze.sh extract lockfile from a built image
- docker/smoke_test.py runtime verification, run with --gpus all
- docker/.dockerignore
- .github/workflows/docker-publish.yml
