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

studio: honor UNSLOTH_TORCH_INDEX_FAMILY in CUDA repair path, assert torch CUDA family at studio image build

Browse source 
_detect_cuda_torch_index_url now respects the explicit family override
before probing nvidia-smi, matching install.sh get_torch_index_url and
install.ps1 Get-TorchIndexUrl. Without it, a GPU-less environment falls
back to cu126 wheels which lack sm_100/sm_120 kernels and break training
on Blackwell. ROCm repair path is intentionally unchanged.

Dockerfile.studio now fails the build if the Studio venv torch local
version tag does not match the pinned TORCH_FAMILY, so a studio ref whose
installer ignores the override can never ship a silently wrong image.
Metadata-only check so QEMU arm64 builds do not need to load torch.
This commit is contained in:
Daniel Han 2026-06-12 14:09:41 +00:00
commit 9d39aeec2b
2 changed files with 12 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

6
docker/Dockerfile.studio
View file

@ -101,6 +101,12 @@ RUN set -eux \
&& UNSLOTH_STUDIO_HOME="${UNSLOTH_STUDIO_HOME}" \
UNSLOTH_TORCH_INDEX_FAMILY="${TORCH_FAMILY}" \
bash install.sh --local \
# Fail loud if the Studio venv torch missed the pinned CUDA family (an
# install.sh that ignores UNSLOTH_TORCH_INDEX_FAMILY falls back to
# nvidia-smi probing, which cannot work at build time and lands on cu126
# wheels with no sm_100/sm_120 kernels). metadata check only: importing
# torch needs native libs, which QEMU arm64 builds cannot load.
&& "${UNSLOTH_STUDIO_HOME}/unsloth_studio/bin/python" -c "from importlib.metadata import version; v = version('torch'); assert v.endswith('+${TORCH_FAMILY}'), 'Studio venv torch ' + v + ' does not match ${TORCH_FAMILY}'; print('Studio venv torch', v)" \
&& rm -rf "${UNSLOTH_STUDIO_HOME}/src/.git" /root/.cache \
&& if [ "${TARGETARCH:-amd64}" = "arm64" ]; then \
for NVRTC_DIR in "${UNSLOTH_STUDIO_HOME}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do \

6
studio/install_python_stack.py
View file

@ -816,6 +816,12 @@ def _detect_cuda_torch_index_url() -> str:
Defaults to cu126 when nvidia-smi is missing or the version is unreadable
(e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback).
"""
# Explicit override (parity with install.sh / install.ps1):
# UNSLOTH_TORCH_INDEX_FAMILY=cu128|cu130|cu126|cpu|... pins the wheel
# index when probing is wrong or impossible (no GPU at build time, CI).
family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY")
if family:
return f"{_PYTORCH_WHL_BASE}/{family}"
exe = shutil.which("nvidia-smi")
if not exe and os.path.isfile("/usr/bin/nvidia-smi"):
exe = "/usr/bin/nvidia-smi"