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)
This commit is contained in:
Daniel Han
parent
a75aef063c
commit
58693c4c73
4 changed files with 193 additions and 2 deletions
|
|
@ -1,3 +1,4 @@
|
|||
**
|
||||
!Dockerfile
|
||||
!entrypoint.sh
|
||||
!smoke_test.py
|
||||
|
|
|
|||
|
|
@ -195,6 +195,19 @@ WORKDIR /workspace
|
|||
RUN mkdir -p ${HF_HOME} ${TRITON_CACHE_DIR}
|
||||
|
||||
COPY smoke_test.py /workspace/smoke_test.py
|
||||
COPY entrypoint.sh /usr/local/bin/unsloth-entrypoint
|
||||
RUN chmod +x /usr/local/bin/unsloth-entrypoint
|
||||
|
||||
# Default entry: drop into python; override with `docker run ... bash` for a shell.
|
||||
# Entrypoint runs three fast pre-flight checks before user code:
|
||||
# 1. nvidia-smi sees at least one GPU (catches missing --gpus all)
|
||||
# 2. torch.cuda.is_available() is True (catches host driver too old)
|
||||
# 3. compute capability >= sm_80 (catches pre-Ampere GPUs)
|
||||
# Each check fails with an actionable error pointing to the fix.
|
||||
# Bypass for offline tooling: docker run -e UNSLOTH_SKIP_GPU_CHECK=1 ...
|
||||
ENTRYPOINT ["/usr/local/bin/unsloth-entrypoint"]
|
||||
|
||||
# Default command: interactive python REPL.
|
||||
# Override examples:
|
||||
# docker run --gpus all unsloth/unsloth:latest python /workspace/smoke_test.py
|
||||
# docker run --gpus all -it unsloth/unsloth:latest bash
|
||||
CMD ["python"]
|
||||
|
|
|
|||
108
docker/entrypoint.sh
Executable file
108
docker/entrypoint.sh
Executable file
|
|
@ -0,0 +1,108 @@
|
|||
#!/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 ...
|
||||
set -euo pipefail
|
||||
|
||||
if [[ "${UNSLOTH_SKIP_GPU_CHECK:-0}" == "1" ]]; then
|
||||
exec "$@"
|
||||
fi
|
||||
|
||||
err() { printf "\033[1;31mERROR:\033[0m %s\n" "$*" >&2; }
|
||||
warn() { printf "\033[1;33mWARN:\033[0m %s\n" "$*" >&2; }
|
||||
|
||||
# --- Check 1: nvidia-smi present and can enumerate at least one GPU ---------
|
||||
if ! command -v nvidia-smi >/dev/null 2>&1; then
|
||||
err "nvidia-smi not found inside the container."
|
||||
err "The CUDA runtime in this image is broken. Re-pull the image."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! nvidia-smi -L 2>/dev/null | grep -q '^GPU'; then
|
||||
err "No GPU visible to nvidia-smi from inside the container."
|
||||
cat >&2 <<'MSG'
|
||||
|
||||
Likely causes (in order of frequency):
|
||||
|
||||
1. You started the container without --gpus all.
|
||||
Re-launch with:
|
||||
docker run --gpus all <other-flags> unsloth/unsloth:latest <cmd>
|
||||
Or use the bundled wrapper:
|
||||
bash docker/run.sh <cmd>
|
||||
|
||||
2. Host is missing nvidia-container-toolkit.
|
||||
Install: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html
|
||||
Then: sudo systemctl restart docker
|
||||
|
||||
3. nvidia-container-toolkit is installed but the Docker daemon was not
|
||||
restarted after install. Run:
|
||||
sudo systemctl restart docker
|
||||
|
||||
4. You are using Podman / Kubernetes / a managed container service that
|
||||
needs a different GPU flag than --gpus all. See the relevant docs:
|
||||
podman: --device nvidia.com/gpu=all
|
||||
k8s: nvidia.com/gpu resource request + GPU operator
|
||||
|
||||
To bypass this check (e.g. offline tooling), set UNSLOTH_SKIP_GPU_CHECK=1.
|
||||
MSG
|
||||
exit 1
|
||||
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).
|
||||
python - >&2 <<'PY' || exit 1
|
||||
import sys
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
sys.exit(0)
|
||||
print("ERROR: torch.cuda.is_available() is False despite nvidia-smi working.")
|
||||
print()
|
||||
print("Most likely the host NVIDIA driver is too old for CUDA 12.8.")
|
||||
print("Required host driver versions for this image:")
|
||||
print(" >= 570 RTX 50-series, RTX 6000 Pro Blackwell (sm_120)")
|
||||
print(" >= 555 B100 / B200 (sm_100)")
|
||||
print(" >= 535 H100 / H200 (sm_90)")
|
||||
print(" >= 525 Ada / Ampere (sm_80 / sm_86 / sm_89)")
|
||||
print()
|
||||
print("Check the host (NOT the container) with: nvidia-smi")
|
||||
print("Then upgrade the driver to match your GPU.")
|
||||
sys.exit(1)
|
||||
PY
|
||||
|
||||
# --- Check 3: compute capability is supported -------------------------------
|
||||
python - >&2 <<'PY' || exit 1
|
||||
import sys
|
||||
import torch
|
||||
major, minor = torch.cuda.get_device_capability(0)
|
||||
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()}")
|
||||
if major < 8:
|
||||
print()
|
||||
print(f"ERROR: Unsloth requires Ampere or newer (sm_80+). Got {name} sm_{major}{minor}.")
|
||||
print()
|
||||
print("Supported architectures baked into this image:")
|
||||
print(" sm_80 Ampere (A100, A40, A30)")
|
||||
print(" sm_86 Ampere (RTX 30-series, A10)")
|
||||
print(" sm_89 Ada (RTX 40-series, L40)")
|
||||
print(" sm_90 Hopper (H100, H200)")
|
||||
print(" sm_100 Blackwell DC (B100, B200)")
|
||||
print(" sm_120 Blackwell (RTX 50-series, RTX 6000 Pro Blackwell)")
|
||||
sys.exit(1)
|
||||
PY
|
||||
|
||||
exec "$@"
|
||||
69
docker/run.sh
Executable file
69
docker/run.sh
Executable file
|
|
@ -0,0 +1,69 @@
|
|||
#!/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.
|
||||
#
|
||||
# Usage:
|
||||
# bash docker/run.sh # interactive python REPL
|
||||
# bash docker/run.sh bash # shell in the container
|
||||
# bash docker/run.sh python /workspace/smoke_test.py # run the smoke test
|
||||
# bash docker/run.sh python /workspace/host/train.py # run your training script
|
||||
# ($PWD is mounted at
|
||||
# /workspace/host)
|
||||
#
|
||||
# Overridable env:
|
||||
# UNSLOTH_IMAGE=unsloth/unsloth:latest image and tag to pull/run
|
||||
# UNSLOTH_GPUS=all GPUs to expose ("all" | "0" | "0,1")
|
||||
# HF_HOME=$HOME/.cache/huggingface host HF cache dir to mount
|
||||
# TRITON_CACHE_DIR=$HOME/.cache/unsloth-triton
|
||||
# host Triton cache dir to mount
|
||||
# UNSLOTH_WORKDIR=$PWD host dir mounted at /workspace/host
|
||||
set -euo pipefail
|
||||
|
||||
IMAGE="${UNSLOTH_IMAGE:-unsloth/unsloth:latest}"
|
||||
GPUS="${UNSLOTH_GPUS:-all}"
|
||||
HF_CACHE="${HF_HOME:-$HOME/.cache/huggingface}"
|
||||
TRITON_CACHE="${TRITON_CACHE_DIR:-$HOME/.cache/unsloth-triton}"
|
||||
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.
|
||||
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.
|
||||
declare -a ENV_FORWARD=(-e HF_HUB_ENABLE_HF_TRANSFER=1)
|
||||
[[ -n "${HF_TOKEN:-}" ]] && ENV_FORWARD+=(-e "HF_TOKEN=${HF_TOKEN}")
|
||||
[[ -n "${WANDB_API_KEY:-}" ]] && ENV_FORWARD+=(-e "WANDB_API_KEY=${WANDB_API_KEY}")
|
||||
[[ -n "${UNSLOTH_LICENSE:-}" ]] && ENV_FORWARD+=(-e "UNSLOTH_LICENSE=${UNSLOTH_LICENSE}")
|
||||
|
||||
set -x
|
||||
exec docker run --rm -it \
|
||||
--gpus "$GPUS" \
|
||||
--ipc=host \
|
||||
--ulimit memlock=-1 \
|
||||
--ulimit stack=67108864 \
|
||||
-v "$HF_CACHE":/workspace/.cache/huggingface \
|
||||
-v "$TRITON_CACHE":/workspace/.cache/triton \
|
||||
-v "$WORK_DIR":/workspace/host \
|
||||
"${ENV_FORWARD[@]}" \
|
||||
"$IMAGE" "$@"
|
||||
Loading…
Add table
Add a link
Reference in a new issue