flash-attention icon indicating copy to clipboard operation
flash-attention copied to clipboard

Can't be installed using "uv" due to an issue in the setup script

Open halflings opened this issue 1 year ago • 7 comments

See https://github.com/astral-sh/uv/issues/1582#issuecomment-1950219890 for full context, but flash-attn can't be installed because it references packaging at the top of the setup script, then lists it as a dependency lower in the script: https://github.com/Dao-AILab/flash-attention/blob/5cdabc2809095b98c311283125c05d222500c8ff/setup.py#L9C6-L9C15

Instead packaging cannot be used and be a dependency at the same time.

halflings avatar Feb 17 '24 19:02 halflings

I'm not sure how this works, do you have a suggestion?

tridao avatar Feb 17 '24 20:02 tridao

Perhaps packaging can be moved to setup_requires and not install_requires. If the pkg is only used for setup.py, it should be moved to setup_requires.

My PR uses this to avoid having psutil installed to the final pkg since it is only part of setup.

https://github.com/Dao-AILab/flash-attention/pull/832/files

Qubitium avatar Feb 18 '24 00:02 Qubitium

@tridao If flash-attn does not JIT but only precompile during setup, it may be best to move all non-runtime depends in setup.py to setup_requires

Qubitium avatar Feb 18 '24 00:02 Qubitium

Related to (or almost a dupe of) #453 I believe.

Taytay avatar Feb 19 '24 17:02 Taytay

I've been investigating this for the past hour or two, and here's what I have learned:

Short version:

The best flash-attn can do is detect that it's running in build-isolated mode, and that its required modules (like packaging, torch, etc) are missing, and give a nice error message directing people to installation instructions. This isn't really a bug with its setup script, and there isn't a way to solve it with proper setup config!

Long version:

(Background : This is all due to build isolation, which is enabled by pip most of the time when installing python modules. flash-attn needs you to pip install with --no-build-isolation primarily so that things like torch can just be referenced during setup/compilation. If you don't use --no-build-isolation, setup.py runs in a totally clean environment, devoid of things like setuptools, packaging, pytorch, etc. This is fine for most packages, because they can just declare (in pyproject.toml) all of the modules they need, and those will get installed before the setup.py is run. Cool! However, flash-attn can't be expected to just list torch as one of its setup dependencies. That means that pytorch would get reinstalled as part of flash-attn's installation process. But flash-attn just wants to reference whatever Pytorch is already installed.)

You could put packaging inside of setup_requires, but by the time that runs, I believe it's too late! setup.py is already running! You can't list modules there that you need installed at the beginning of setup.py (from what I read anyway). I experimented with this by creating a pyproject.toml file at the root of the flash-attn folder, (which forces pip install -e . to run in isolated mode), and sure enough, whether I have packaging in setup_requires or not, setup.py can't find the packaging module.

You can move stuff like packaging into the pyproject.toml file:

[build-system]
requires = ["wheel", "packaging", "ninja"]
build-backend = "setuptools.build_meta"

But that doesn't work either because you won't have torch or other large modules available to you at setup time.

So I think the best course of action for flash-attn is to have some nice code at the top that does stuff like this:


try:
    from packaging.version import parse, Version
except ImportError:
    # The inability to import packaging.version is likely due to running with build isolation
    print(
        """Could not import the 'packaging' module. You likely forgot to install without the "--build-isolation" flag. Follow the installation instructions here: https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#installation-and-features"""
    )
    exit(1)

There are a few places you'd need to do that conditional import or check for a module: packaging, torch, psutil, ninja. But nice checks and error messages like this would get rid of confusion and common bug reports like this (that bit me too!): https://github.com/Dao-AILab/flash-attention/issues/453

I wrote up similar findings for uv here when I requested they add support for --no-build-isolation: https://github.com/astral-sh/uv/issues/1715

Taytay avatar Feb 19 '24 20:02 Taytay

It's wild that https://github.com/pypa/pip/issues/8437 is locked.

fxmarty avatar Mar 28 '24 10:03 fxmarty

FWIW, a pyproject.toml like the following works for me in uv's project flow (uv lock, uv sync, uv run).

[project]
name = "asdf"
version = "0.1.0"
description = "asdf"
authors = [
    { name = "me", email = "[email protected]" }
]
requires-python = "~=3.11"
dependencies = [
    "bitsandbytes~=0.43",
    "datasets~=2.20",
    "devtools>=0.12.2",
    "numpy~=1.26",
    "optuna~=3.6",
    "pandas>=2.2.2",
    "peft~=0.11",
    "pillow~=10.4",
    "pydantic~=2.8",
    "schedulefree>=1.2.6",
    "scikit-learn~=1.5",
    "scipy~=1.14",
    "timm~=1.0",
    "torch==2.3.0",
    "torchvision~=0.18.0",
    "transformers~=4.42",
    "triton~=2.3.0",
    "wandb>=0.17.4",
    "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.3/flash_attn-2.6.3+cu123torch2.3cxx11abiFALSE-cp311-cp311-linux_x86_64.whl"
]

tyler-romero avatar Aug 22 '24 23:08 tyler-romero