Skip to content

Commit

Permalink
make properly pip installable without needing build isolation (uv pip…
Browse files Browse the repository at this point in the history 
… install now works too) (#402)

Co-authored-by: Tri Dao <[email protected]>
  • Loading branch information
@kszlim @tridao
kszlim and tridao authored Aug 6, 2024
1 parent a07ff1b commit 323db26
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 21 deletions.
2 changes: 2 additions & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@ with an efficient hardware-aware design and implementation in the spirit of [Fla

- [Option] `pip install causal-conv1d>=1.4.0`: an efficient implementation of a simple causal Conv1d layer used inside the Mamba block.
- `pip install mamba-ssm`: the core Mamba package.
- `pip install mamba-ssm[causal-conv1d]`: To install core Mamba package and causal-conv1d.
- `pip install mamba-ssm[dev]`: To install core Mamba package and dev depdencies.

It can also be built from source with `pip install .` from this repository.

Expand Down
46 changes: 46 additions & 0 deletions pyproject.toml
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
[project]
name = "mamba_ssm"
description = "Mamba state-space model"
readme = "README.md"
authors = [
{ name = "Tri Dao", email = "[email protected]" },
{ name = "Albert Gu", email = "[email protected]" }
]
requires-python = ">= 3.7"
dynamic = ["version"]
license = { file = "LICENSE" } # Include a LICENSE file in your repo
keywords = ["cuda", "pytorch", "state-space model"]
classifiers = [
"Programming Language :: Python :: 3",
"License :: OSI Approved :: BSD License",
"Operating System :: Unix"
]
dependencies = [
"torch",
"ninja",
"einops",
"triton",
"transformers",
"packaging",
"setuptools>=61.0.0",
]
urls = { name = "Repository", url = "https://github.com/state-spaces/mamba"}

[project.optional-dependencies]
causal-conv1d = [
"causal-conv1d>=1.2.0"
]
dev = [
"pytest"
]


[build-system]
requires = [
"setuptools>=61.0.0",
"wheel",
"torch",
"packaging",
"ninja",
]
build-backend = "setuptools.build_meta"
23 changes: 2 additions & 21 deletions setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
import torch
from torch.utils.cpp_extension import (
BuildExtension,
CppExtension,
CUDAExtension,
CUDA_HOME,
HIP_HOME
Expand Down Expand Up @@ -349,31 +348,13 @@ def run(self):
"mamba_ssm.egg-info",
)
),
author="Tri Dao, Albert Gu",
author_email="[email protected], [email protected]",
description="Mamba state-space model",
long_description=long_description,
long_description_content_type="text/markdown",
url="https://github.com/state-spaces/mamba",
classifiers=[
"Programming Language :: Python :: 3",
"License :: OSI Approved :: BSD License",
"Operating System :: Unix",
],

ext_modules=ext_modules,
cmdclass={"bdist_wheel": CachedWheelsCommand, "build_ext": BuildExtension}
if ext_modules
else {
"bdist_wheel": CachedWheelsCommand,
},
python_requires=">=3.8",
install_requires=[
"torch",
"packaging",
"ninja",
"einops",
"triton",
"transformers",
# "causal_conv1d>=1.4.0",
],
}
)

8 comments on commit 323db26

@amoskvic
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One of the latest three commits (as of the time of writing) broke installation on ROCm machines. I believe that it is this one. We'll be working on a fix, but it might be beneficial to revert it, at least temporarily. @kszlim please let me know if you're open to collaborate on fixing this together.

@kszlim
Copy link
Contributor Author

@kszlim kszlim commented on 323db26 Aug 13, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@amoskvic what error are you seeing?

Is this with or without the optional dependencies?

Is this for sdists or when you're trying to install a wheel?

@kszlim
Copy link
Contributor Author

@kszlim kszlim commented on 323db26 Aug 13, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you share how you're installing mamba? Is it with build isolation on or off?

@amoskvic
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is using the previously working method: "pip install ." to build from source. I give the full error output below.

I now also tried all other options (using pip with and without --no-build-isolation flag).
Installing from source with the "--no-build-isolation" flag does work. Perhaps we could make it a default version for builds on AMD? I am not 100% sure what the right approach would be here.

To clarify: the pip-based installation did not work previously either, so basically the downside is that this commit introduced the need to use "no-build-isolation" when building from source on AMD. It might be helpful if we make it a default option or clarify that in case of any error (not only torch version), it's worth trying the "no-build-isolation" approach.

` ERROR: Command errored out with exit status 1:
command: /opt/conda/envs/py_3.9/bin/python /opt/conda/envs/py_3.9/lib/python3.9/site-packages/pip/_vendor/pep517/_in_process.py get_requires_for_build_wheel /tmp/tmp1ezb7sh5
cwd: /tmp/pip-req-build-83tx_ki8
Complete output (24 lines):
/tmp/pip-build-env-o5311wb9/overlay/lib/python3.9/site-packages/torch/_subclasses/functional_tensor.py:258: UserWarning: Failed to initialize NumPy: No module named 'numpy' (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:84.)
cpu = _conversion_method_template(device=torch.device("cpu"))
No ROCm runtime is found, using ROCM_HOME='/opt/rocm'
:118: UserWarning: mamba_ssm was requested, but nvcc was not found. Are you sure your environment has nvcc available? If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.

torch.version = 2.4.0+cu121

Traceback (most recent call last):
File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/pip/_vendor/pep517/_in_process.py", line 280, in
main()
File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/pip/_vendor/pep517/_in_process.py", line 263, in main
json_out['return_val'] = hook(**hook_input['kwargs'])
File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/pip/_vendor/pep517/_in_process.py", line 114, in get_requires_for_build_wheel
return hook(config_settings)
File "/tmp/pip-build-env-o5311wb9/overlay/lib/python3.9/site-packages/setuptools/build_meta.py", line 327, in get_requires_for_build_wheel
return self._get_build_requires(config_settings, requirements=[])
File "/tmp/pip-build-env-o5311wb9/overlay/lib/python3.9/site-packages/setuptools/build_meta.py", line 297, in _get_build_requires
self.run_setup()
File "/tmp/pip-build-env-o5311wb9/overlay/lib/python3.9/site-packages/setuptools/build_meta.py", line 313, in run_setup
exec(code, locals())
File "", line 188, in
NameError: name 'bare_metal_version' is not defined

ERROR: Command errored out with exit status 1: /opt/conda/envs/py_3.9/bin/python /opt/conda/envs/py_3.9/lib/python3.9/site-packages/pip/_vendor/pep517/_in_process.py get_requires_for_build_wheel /tmp/tmp1ezb7sh5 Check the logs for full command output.`

@kszlim
Copy link
Contributor Author

@kszlim kszlim commented on 323db26 Aug 13, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah I think perhaps the best "fix" is to mention that if you have troubles pip installing that you should try without build isolation in the readme?

It seems weird that you also get the error on line 188, because with ROCm, you shouldn't hit that branch based off the setup.py.

Ah, it seems like your torch version is a CUDA version? That's probably why?

@amoskvic
Copy link
Contributor

@amoskvic amoskvic commented on 323db26 Aug 13, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@kszlim yes, I think the "no-build-isolation" is mentioned in the readme already, I apologize I did not realize this line was there since it was not relevant before. But I agree that rephrasing it from "If pip complains about PyTorch versions" to something more general might help.

Yes, that's the more strange bug. It actually not a CUDA torch version. It's run on a ROCm - specific docker image for pytorch (from here https://hub.docker.com/r/rocm/pytorch/tags). And torch.version.hip appropriately returns a hip version, and so on.
Still, it looks like something in the pipeline is failing to realize that and trying to build for nvidia. I'm not sure what exactly is going on there because '/opt/rocm' is a valid ROCm install directory in that docker, but somehow something fails to detect it.

@kszlim
Copy link
Contributor Author

@kszlim kszlim commented on 323db26 Aug 13, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you want to make a PR to clarify the documentation? Imo you should try fiddling with the setup.py, it's possible you could get it to work.

@amoskvic
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll try to see if there are some ways of fixing things through setup.py etc., but if that fails I'll submit a small documentation change PR sometime soon.

Please sign in to comment.