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

Installer: add UNSLOTH_TORCH_INDEX_FAMILY override for torch wheel index

Browse source 
get_torch_index_url() (install.sh) and Get-TorchIndexUrl (install.ps1) pick
the PyTorch wheel index by probing the local GPU with nvidia-smi. That works
on a normal machine, but it has no good answer when there is no GPU to probe
at the time the index is chosen, for example inside a Docker image build or a
CI runner that builds wheels for a different target than it runs on. In that
situation the bash path falls back to cu126 (build hosts that expose
/proc/driver/nvidia but cannot run nvidia-smi) or cpu, baking the wrong
wheels into an image that is meant to run CUDA 12.8.

Add an optional UNSLOTH_TORCH_INDEX_FAMILY environment variable that names the
index path leaf directly (cu128, cu130, cu126, cpu, rocm6.4, ...). When set,
it short-circuits probing and pins the index deterministically. When unset or
empty the function behaves exactly as before, so native installs are
unaffected.
This commit is contained in:
Daniel Han 2026-06-24 01:09:40 +00:00
commit 43e4090ebe
2 changed files with 16 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

5
install.ps1
View file

@ -1939,6 +1939,11 @@ exit 0
# Mirrors Get-PytorchCudaTag in setup.ps1.
function Get-TorchIndexUrl {
$baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
# Explicit override for hosts where GPU probing is impossible or must not
# happen (image builds, CI runners that build for a different target than
# they run on). Names the wheel index path leaf directly, e.g.
# UNSLOTH_TORCH_INDEX_FAMILY=cu128 | cu130 | cu126 | cpu | rocm6.4 | ...
if ($env:UNSLOTH_TORCH_INDEX_FAMILY) { return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY)" }
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {
$output = Invoke-NvidiaSmiBounded $NvidiaSmiExe

11
install.sh
View file

@ -1962,6 +1962,17 @@ _has_usable_nvidia_gpu() {
get_torch_index_url() {
_base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
_base="${_base%/}"
# Explicit override for hosts where GPU probing is impossible or must not
# happen (Docker image builds, CI runners that build for a different
# target than they run on). Names the wheel index path leaf directly:
# UNSLOTH_TORCH_INDEX_FAMILY=cu128 | cu130 | cu126 | cu124 | cu118 | cpu | rocm6.4 | ...
# Example: a Docker image that targets CUDA 12.8 is built on a host with no
# GPU and no nvidia-smi, so probing would otherwise fall back to cpu (or, on
# build hosts that leak /proc/driver/nvidia but cannot run nvidia-smi, cu126)
# and bake the wrong wheels. Setting this pins the index deterministically.
if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then
echo "$_base/${UNSLOTH_TORCH_INDEX_FAMILY}"; return
fi
# macOS: always CPU (no CUDA support)
case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac
# Try nvidia-smi -- require the binary to actually list a usable GPU.