Support flash attention flash-attn --no-build-isolation with uv sync

[vwxyzjn]

I get the following error despite installation flash attention was successful via uv add flash-attn --no-build-isolation. See #6402. Is there anyway to honor the no-build-isolation when doing uv sync? I imagine the idea is to mark the flash-attn dependencies somehow and run it with no-build-isolation after all the other dependencies have been installed.

pyproject.toml and uv.lock are here: https://gist.github.com/vwxyzjn/fba2c7ffac8e6c923ffe912a872ec9a8

  open-instruct git:(uv) ✗ uv sync
Using Python 3.12.5
Creating virtualenv at: .venv
⠦ fire==0.6.0                                                                                    error: Failed to download and build `flash-attn==2.6.3`
  Caused by: Build backend failed to determine extra requires with `build_wheel()` with exit status: 1
--- stdout:

--- stderr:
Traceback (most recent call last):
  File "<string>", line 14, in <module>
  File "/home/costa/.cache/uv/builds-v0/.tmpLR2n19/lib/python3.12/site-packages/setuptools/build_meta.py", line 332, in get_requires_for_build_wheel
    return self._get_build_requires(config_settings, requirements=[])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/costa/.cache/uv/builds-v0/.tmpLR2n19/lib/python3.12/site-packages/setuptools/build_meta.py", line 302, in _get_build_requires
    self.run_setup()
  File "/home/costa/.cache/uv/builds-v0/.tmpLR2n19/lib/python3.12/site-packages/setuptools/build_meta.py", line 502, in run_setup
    super().run_setup(setup_script=setup_script)
  File "/home/costa/.cache/uv/builds-v0/.tmpLR2n19/lib/python3.12/site-packages/setuptools/build_meta.py", line 318, in run_setup
    exec(code, locals())
  File "<string>", line 21, in <module>
ModuleNotFoundError: No module named 'torch'
---
  Caused by: This error likely indicates that flash-attn==2.6.3 depends on torch, but doesn't declare it as a build dependency. If flash-attn==2.6.3 is a first-party package, consider adding torch to its `build-system.requires`. Otherwise, `uv pip install torch` into the environment and re-run with `--no-build-isolation`.

[zanieb]

You'll need to use uv sync --no-build-isolation-package flash-attn — it'd be great if we tracked this though!

[vwxyzjn]

Wow amazing that this feature already exists!

[vwxyzjn]

I just tested it out, and it seems to fail:

uv sync --no-build-isolation-package flash-attn
⠹ flash-attn==2.6.3                                                                                 error: Failed to download and build `flash-attn==2.6.3`
  Caused by: Build backend failed to determine extra requires with `build_wheel()` with exit status: 1
--- stdout:

--- stderr:
Traceback (most recent call last):
  File "<string>", line 14, in <module>
  File "/home/costa/.cache/uv/builds-v0/.tmpEfmf64/lib/python3.12/site-packages/setuptools/build_meta.py", line 332, in get_requires_for_build_wheel
    return self._get_build_requires(config_settings, requirements=[])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/costa/.cache/uv/builds-v0/.tmpEfmf64/lib/python3.12/site-packages/setuptools/build_meta.py", line 302, in _get_build_requires
    self.run_setup()
  File "/home/costa/.cache/uv/builds-v0/.tmpEfmf64/lib/python3.12/site-packages/setuptools/build_meta.py", line 502, in run_setup
    super().run_setup(setup_script=setup_script)
  File "/home/costa/.cache/uv/builds-v0/.tmpEfmf64/lib/python3.12/site-packages/setuptools/build_meta.py", line 318, in run_setup
    exec(code, locals())
  File "<string>", line 21, in <module>
ModuleNotFoundError: No module named 'torch'
---
  Caused by: This error likely indicates that flash-attn==2.6.3 depends on torch, but doesn't declare it as a build dependency. If flash-attn==2.6.3 is a first-party package, consider adding torch to its `build-system.requires`. Otherwise, `uv pip install torch` into the environment and re-run with `--no-build-isolation`.

[charliermarsh]

I think to make this work with uv sync, sadly you need to do something like uv pip install torch prior to running uv sync. The build dependencies have to be available in the virtual environment before you run the install.

[kabouzeid]

I'm curious about what the recommended best practice is for including torch extensions in the project.

Here's my current approach:

uv sync  # torch is installed here
uv sync --extra build  # to keep things tidy: setuptools, wheel, and other stuff missing because of --no-build-isolation
uv sync --extra torchextensions --no-build-isolation  # torch extensions installed here

Previously, I did this, which is also not ideal:

uv sync  # torch is installed here
uv pip install setuptools wheel ...  # missing because of --no-build-isolation
uv pip install --no-build-isolation some_torch_extension

Your recommendation here is also not ideal, because you would need to manually make sure to pip install the exact torch version that is also in the uv.lock file:

uv pip install torch==2.3.1 --index-url https://download.pytorch.org/whl/cu121
uv pip install setuptools wheel ...  # missing because of --no-build-isolation
uv sync --no-build-isolation  # torch extensions installed here

[vwxyzjn]

You can include wheel and setup tools in uv add I think

[charliermarsh]

I added some docs for this here: #6607

[eladrich]

Hi, I was following the suggested solution that was also introduced in the docs of specifying flash-attn as an extra dependency and use the following two commands

uv sync --extra build
uv sync --extra build --extra compile

My main challenge with that solution is from now on this required me to always remember to specify these flags when using uv for that project with the flash-attn dependencies, otherwise if I accidentally use a uv sync command the package gets uninstalled.
A different suggested solution was to separately use a uv pip install torch beforehand but I think that requires verifying I'm using the right torch versions in both the uv pip commands and sync ones and I wanted to avoid the uv pip syntax completely.

Instead, I managed to build the environment by specifying flash-attn alongside the other dependencies (while still defining it as a no-build-isolation-package) and then using the --no-install-package to create the environment the first time.

uv sync --no-install-package flash-attn
uv sync

If I understand correctly this allows me to use the regular uv sync once the project was initialized. I was wondering whether that is a valid solution or whether it has other downsides?