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

studio: pin packaging + triton with FLA --no-deps install

Browse source 
An end-to-end install simulation in a fresh venv caught a real
regression: `fla/utils.py` does `from packaging import version` and
`import triton` at module load, but fla-core's METADATA only declares
einops + torch. With `--no-deps` the worker would land FLA in any
runtime that lacks packaging (e.g. minimal torch builds) and the
post-install import probe would fall back to the torch GDN loop
silently.

Add `packaging` and `triton` to `_FLA_RUNTIME_DEPS` so the install
spec list always carries them. Tests updated to assert both are now in
the install command.
This commit is contained in:
danielhanchen 2026-05-16 11:24:48 +00:00
commit 66dface7d7
2 changed files with 18 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

21
studio/backend/core/training/worker.py
View file

@ -65,9 +65,13 @@ _TILELANG_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_TILELANG_INSTALL"
_FLA_PACKAGE_VERSION = "0.5.0"
_FLA_CORE_PACKAGE_VERSION = "0.5.0"
_FLA_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FLA_INSTALL"
# fla-core's runtime dep that --no-deps suppresses. Without einops,
# `import fla.modules` raises ModuleNotFoundError at startup.
_FLA_RUNTIME_DEPS = ("einops",)
# fla-core declares `einops` in its METADATA but `fla/utils.py`
# also imports `packaging` at module load; that one is NOT declared
# upstream (an FLA bug). triton is a torch dep but we list it
# defensively because some torch wheel builds skip it. With --no-deps
# we have to bring these in ourselves, otherwise `import fla.modules`
# raises ModuleNotFoundError at startup.
_FLA_RUNTIME_DEPS = ("einops", "packaging", "triton")
# Studio installer permits torch>=2.4,<2.11.0 but fla-core 0.5.0
# declares torch>=2.7.0; skip FLA on older torch to keep the
# fallback path clean.
@ -349,9 +353,9 @@ def _ensure_flash_linear_attention(event_queue: Any, model_name: str) -> None:
mamba_ssm path and never call FLA's GDN kernels, so we skip them.
Pinned ``flash-linear-attention``, ``fla-core`` and the runtime
deps we explicitly want (``einops``) are installed with ``--no-deps``
so pip never silently upgrades torch from fla-core's ``torch>=2.7.0``
requirement.
deps we explicitly want (``einops``, ``packaging``, ``triton``)
are installed with ``--no-deps`` so pip never silently upgrades
torch from fla-core's ``torch>=2.7.0`` requirement.
Set ``UNSLOTH_STUDIO_SKIP_FLA_INSTALL=1`` to bypass entirely.
"""
@ -392,8 +396,9 @@ def _ensure_flash_linear_attention(event_queue: Any, model_name: str) -> None:
)
# Install fla-core's required non-torch runtime deps explicitly
# because `--no-deps` suppresses them. Without einops, `import
# fla.modules` raises ModuleNotFoundError at runtime.
# because `--no-deps` suppresses them. Without einops/packaging
# (and triton, on minimal torch builds), `import fla.modules`
# raises ModuleNotFoundError at runtime.
specs = [
*_FLA_RUNTIME_DEPS,
f"fla-core=={_FLA_CORE_PACKAGE_VERSION}",

5
studio/backend/tests/test_training_worker_flash_attn.py
View file

@ -339,7 +339,12 @@ def test_flash_linear_attention_install_includes_einops(monkeypatch):
args = run_mock.call_args[0][0]
assert "--no-deps" in args
# einops is declared by fla-core; packaging and triton are pulled in
# because fla/utils.py imports them at module load but neither is
# declared in fla-core's METADATA (an upstream FLA gap).
assert "einops" in args
assert "packaging" in args
assert "triton" in args
assert f"flash-linear-attention=={worker._FLA_PACKAGE_VERSION}" in args
assert f"fla-core=={worker._FLA_CORE_PACKAGE_VERSION}" in args