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

Expand arch list to every current x86_64 NVIDIA CC per developer.nvidia.com/cuda/gpus

Browse source 
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.
This commit is contained in:
Daniel Han 2026-05-24 08:31:15 +00:00
commit dde5170e7a
2 changed files with 40 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

23
docker/Dockerfile
View file

@ -6,8 +6,11 @@
# * cu128 wheels are fat binaries: SASS for sm_80;86;89;90;100;120.
# * Unsloth's runtime kernels are Triton, which JIT-compiles per device at first run.
# * Anything that DOES need to be source-built (rare on this pin set) compiles
# against TORCH_CUDA_ARCH_LIST="10.0;12.0+PTX" -- the host GPU is irrelevant
# for compilation; nvcc emits whatever the arch list says.
# against TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;10.3;12.0;12.1+PTX",
# covering every current x86_64 NVIDIA compute capability per
# https://developer.nvidia.com/cuda/gpus.
# The host GPU is irrelevant for compilation; nvcc emits whatever the arch
# list says.
#
# Build host requirements:
# * Docker with buildkit (default since 23.x)
@ -30,9 +33,19 @@ ENV DEBIAN_FRONTEND=noninteractive \
PIP_DISABLE_PIP_VERSION_CHECK=1 \
PYTHONDONTWRITEBYTECODE=1 \
PYTHONUNBUFFERED=1 \
# Cross-compile for: Ampere, Ada, Hopper, B100/B200 (sm_100), RTX 50x / 6000 Pro (sm_120).
# +PTX on the highest arch lets future Blackwell SKUs run via JIT-PTX.
TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0;10.0;12.0+PTX" \
# Cross-compile for every current x86_64 NVIDIA arch per
# https://developer.nvidia.com/cuda/gpus:
# 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)
# +PTX on the highest lets future arch revisions run via JIT-PTX.
TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0;10.0;10.3;12.0;12.1+PTX" \
MAX_JOBS=4 \
CUDA_HOME=/usr/local/cuda \
# Build-host-independence guards. The build must NEVER introspect a GPU,

31
docker/entrypoint.sh
View file

@ -91,18 +91,31 @@ 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:
# Image targets every current x86_64 NVIDIA arch from Turing onward, per
# https://developer.nvidia.com/cuda/gpus.
SUPPORTED = (
("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)"),
)
if major < 7 or (major == 7 and minor < 5):
print()
print(f"ERROR: Unsloth requires Ampere or newer (sm_80+). Got {name} sm_{major}{minor}.")
print(f"ERROR: Unsloth image requires Turing or newer (sm_75+). 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)")
print("Supported architectures in this image:")
for arch, fam, ex in SUPPORTED:
print(f" {arch:7s} {fam:13s} ({ex})")
sys.exit(1)
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.")
PY
exec "$@"