diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 41ca6c2aa8..fd6b91b76d 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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}", diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index d44ddc0d67..8030d49dc6 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -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