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

Fix Colab training failing with "Please install unsloth_zoo"

Browse source 
The Colab no-venv path in setup.sh (added in #4601, modified in #4618)
only installs studio.txt (fastapi, uvicorn, structlog, etc.) into
system Python. It skips install_python_stack entirely via
_SKIP_PYTHON_DEPS=true, which means base.txt (unsloth + unsloth-zoo),
extras.txt, extras-no-deps.txt (trl, transformers, sentence_transformers),
and overrides.txt (torchao) are never installed.

When the training worker subprocess spawns and tries to import unsloth,
unsloth/__init__.py calls importlib.metadata.version("unsloth_zoo")
which raises PackageNotFoundError, surfaced as:
  "Failed to import ML libraries: Please install unsloth_zoo"

Even if unsloth_zoo were installed, the next failure would be trl
(imported at trainer.py:63) since it lives in extras-no-deps.txt.

Fix: install all four requirement groups in the Colab block, matching
what install_python_stack.py does for the normal venv path:
  - base.txt (unsloth + unsloth-zoo) -- fail-fast on error
  - extras.txt (version-stripped to avoid Colab conflicts)
  - extras-no-deps.txt (trl, transformers, sentence_transformers)
  - overrides.txt (torchao)

Only the Colab code path is changed. Non-Colab paths are untouched.
This commit is contained in:
Daniel Han 2026-03-26 14:18:49 +00:00
commit be4c57b052
1 changed files with 14 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

19
studio/setup.sh
View file

@ -304,15 +304,24 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then
| grep -v '^#' | grep -v '^$' \
| pip install -q -r /dev/stdin 2>/dev/null || true
# Install core ML packages (unsloth + unsloth-zoo) and training
# extras (trl, transformers, etc.) into system Python. Without
# these the training/inference/export workers fail at import time
# with "Failed to import ML libraries: Please install unsloth_zoo".
# Install the same ML package groups that install_python_stack.py
# uses for the normal (venv) path. Without these the training,
# inference, and export workers fail at import time with:
# "Failed to import ML libraries: Please install unsloth_zoo"
# base.txt must succeed -- it has unsloth + unsloth-zoo.
echo " Installing core ML packages for training..."
pip install -q -r "$SCRIPT_DIR/backend/requirements/base.txt" 2>/dev/null || true
pip install -q -r "$SCRIPT_DIR/backend/requirements/base.txt" || {
echo "Failed to install base ML packages (unsloth + unsloth-zoo)"
exit 1
}
sed 's/[><=!~;].*//' "$SCRIPT_DIR/backend/requirements/extras.txt" \
| grep -v '^#' | grep -v '^$' \
| pip install -q -r /dev/stdin 2>/dev/null || true
# trl, transformers, sentence_transformers live in extras-no-deps.txt
pip install -q --no-deps \
-r "$SCRIPT_DIR/backend/requirements/extras-no-deps.txt" 2>/dev/null || true
pip install -q --force-reinstall \
-r "$SCRIPT_DIR/backend/requirements/overrides.txt" 2>/dev/null || true
_COLAB_NO_VENV=true
else