Commit graph

5 commits

Author SHA1 Message Date
danielhanchen
f1a63db6fa docker: ship Jupyter, Studio and prebuilt llama.cpp out of the box
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.
2026-06-12 05:06:51 +00:00
Daniel Han
c463d58277 entrypoint.sh: correct driver-floor message (570+ unconditionally on cu128)
The earlier message had per-arch driver minimums (525/535/555/570) that
came from when each chip first got driver support. That's not how CUDA
toolkit floors work -- cu128 imposes 570.26+ on EVERY GPU regardless of
arch. Only B300 (sm_103) and DGX Spark (sm_121) need a newer driver
(580+), and they ship factory with those drivers anyway.

External HF README has the same correction applied in temp/hf_readme.md
(updated separately when published).
2026-05-24 11:36:36 +00:00
Daniel Han
1769204ade Dockerfile: arm64 DGX Spark NVRTC + ptxas fix (cu13 alongside cu128)
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.
2026-05-24 11:35:07 +00:00
Daniel Han
dde5170e7a Expand arch list to every current x86_64 NVIDIA CC per developer.nvidia.com/cuda/gpus
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.
2026-05-24 08:31:15 +00:00
Daniel Han
58693c4c73 Add entrypoint with GPU pre-flight checks + opinionated run.sh wrapper
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)
2026-05-24 07:04:48 +00:00