From e728eeda6f3fad26cb04e1e5fa3b420985895def Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 24 May 2026 11:35:31 +0000
Subject: [PATCH] Dockerfile: switch runtime base cudnn-runtime -> base (~2.7
 GB lighter)

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.
---
 docker/Dockerfile | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index f74cde4560..8839b7a552 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -237,11 +237,19 @@ print(f"OK: {' + '.join(LIGHT_IMPORTS)} import cleanly on no-GPU host")
 PY
 
 # =============================================================================
-# Stage 2: runtime  -- slim runtime image, no nvcc, no headers
+# Stage 2: runtime  -- slim runtime image, no nvcc, no cuDNN/cuBLAS layers
 # =============================================================================
-FROM nvidia/cuda:${CUDA_VERSION}-cudnn-runtime-ubuntu${UBUNTU_VERSION} AS runtime
+# 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.
+FROM nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION} AS runtime
 
-# The nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04 manifest is multi-arch
+# 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.
 ARG TARGETARCH