Flash_attn not able to import with transformer==4.34.1
Author: Girrajjangid
Pull request: no

### System Info

!pip -q install auto-gptq==0.4.2
!pip -q install optimum==1.13.2
!pip -q install bitsandbytes==0.41.1
!pip -q install accelerate==0.23.0
!pip -q install transformers==4.34.1
!pip -q install mlflow==2.7.0

--------------------------------------------------------

- `transformers` version: 4.34.1
- Platform: Linux-5.15.0-1045-aws-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.2
- Accelerate version: 0.23.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 1.13.1+cu117 (True)
- Tensorflow version (GPU?): 2.11.1 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in a script?: No


### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

``` python
model_name_or_path = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ"
revision = "main"
model_path = snapshot_download(repo_id = model_name_or_path, revision = revision)

tokenizer = transformers.AutoTokenizer.from_pretrained(model_path, use_fast=True)
model = transformers.AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, device_map="auto").eval()
```

# Error that I got:
```python
---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-33e11ec6-0ebd-4ce4-bae6-a63d862b86d4/lib/python3.10/site-packages/transformers/utils/import_utils.py:1282, in _LazyModule._get_module(self, module_name)
   1281 try:
-> 1282     return importlib.import_module("." + module_name, self.__name__)
   1283 except Exception as e:

File /usr/lib/python3.10/importlib/__init__.py:126, in import_module(name, package)
    125         level += 1
--> 126 return _bootstrap._gcd_import(name[level:], package, level)

File <frozen importlib._bootstrap>:1050, in _gcd_import(name, package, level)

File <frozen importlib._bootstrap>:1027, in _find_and_load(name, import_)

File <frozen importlib._bootstrap>:1006, in _find_and_load_unlocked(name, import_)

File <frozen importlib._bootstrap>:688, in _load_unlocked(spec)

File <frozen importlib._bootstrap_external>:883, in exec_module(self, module)

File <frozen importlib._bootstrap>:241, in _call_with_frames_removed(f, *args, **kwds)

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-33e11ec6-0ebd-4ce4-bae6-a63d862b86d4/lib/python3.10/site-packages/transformers/models/mistral/modeling_mistral.py:45
     44 if is_flash_attn_available():
---> 45     from flash_attn import flash_attn_func, flash_attn_varlen_func
     46     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa

ImportError: cannot import name 'flash_attn_func' from 'flash_attn' (/databricks/python/lib/python3.10/site-packages/flash_attn/__init__.py)

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
File <command-236141353542329>, line 8
      5 model_path = snapshot_download(repo_id = model_name_or_path, revision = revision)
      7 tokenizer = transformers.AutoTokenizer.from_pretrained(model_path, use_fast=True)
----> 8 model = transformers.AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, device_map="auto").eval()

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-33e11ec6-0ebd-4ce4-bae6-a63d862b86d4/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py:564, in _BaseAutoModelClass.from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
    560     return model_class.from_pretrained(
    561         pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs
    562     )
    563 elif type(config) in cls._model_mapping.keys():
--> 564     model_class = _get_model_class(config, cls._model_mapping)
    565     return model_class.from_pretrained(
    566         pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs
    567     )
    568 raise ValueError(
    569     f"Unrecognized configuration class {config.__class__} for this kind of AutoModel: {cls.__name__}.\n"
    570     f"Model type should be one of {', '.join(c.__name__ for c in cls._model_mapping.keys())}."
    571 )

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-33e11ec6-0ebd-4ce4-bae6-a63d862b86d4/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py:387, in _get_model_class(config, model_mapping)
    386 def _get_model_class(config, model_mapping):
--> 387     supported_models = model_mapping[type(config)]
    388     if not isinstance(supported_models, (list, tuple)):
    389         return supported_models

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-33e11ec6-0ebd-4ce4-bae6-a63d862b86d4/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py:739, in _LazyAutoMapping.__getitem__(self, key)
    737 if model_type in self._model_mapping:
    738     model_name = self._model_mapping[model_type]
--> 739     return self._load_attr_from_module(model_type, model_name)
    741 # Maybe there was several model types associated with this config.
    742 model_types = [k for k, v in self._config_mapping.items() if v == key.__name__]

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-33e11ec6-0ebd-4ce4-bae6-a63d862b86d4/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py:753, in _LazyAutoMapping._load_attr_from_module(self, model_type, attr)
    751 if module_name not in self._modules:
    752     self._modules[module_name] = importlib.import_module(f".{module_name}", "transformers.models")
--> 753 return getattribute_from_module(self._modules[module_name], attr)

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-33e11ec6-0ebd-4ce4-bae6-a63d862b86d4/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py:697, in getattribute_from_module(module, attr)
    695 if isinstance(attr, tuple):
    696     return tuple(getattribute_from_module(module, a) for a in attr)
--> 697 if hasattr(module, attr):
    698     return getattr(module, attr)
    699 # Some of the mappings have entries model_type -> object of another model type. In that case we try to grab the
    700 # object at the top level.

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-33e11ec6-0ebd-4ce4-bae6-a63d862b86d4/lib/python3.10/site-packages/transformers/utils/import_utils.py:1272, in _LazyModule.__getattr__(self, name)
   1270     value = self._get_module(name)
   1271 elif name in self._class_to_module.keys():
-> 1272     module = self._get_module(self._class_to_module[name])
   1273     value = getattr(module, name)
   1274 else:

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-33e11ec6-0ebd-4ce4-bae6-a63d862b86d4/lib/python3.10/site-packages/transformers/utils/import_utils.py:1284, in _LazyModule._get_module(self, module_name)
   1282     return importlib.import_module("." + module_name, self.__name__)
   1283 except Exception as e:
-> 1284     raise RuntimeError(
   1285         f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its"
   1286         f" traceback):\n{e}"
   1287     ) from e

RuntimeError: Failed to import transformers.models.mistral.modeling_mistral because of the following error (look up to see its traceback):
cannot import name 'flash_attn_func' from 'flash_attn' (/databricks/python/lib/python3.10/site-packages/flash_attn/__init__.py)

```

### Expected behavior

The same code is working fine on colab. The difference is that on colab there is no `flash_attn`. But on my cluster `flash_attn==1.0.7` is installed.

The same code is also working fine when I install.

!pip -q install transformers==4.34.0                            # downgrade.
!pip -q install torch==2.1.0                                          # upgrade. 
!pip -q install flash-attn==2.3.2 --no-build-isolation  # upgrade.

But it doesnt work when I create serving endpoint. [Failed during container Image build] with this eror.

```txt
#11 0.283 channels:
#11 0.283 - conda-forge
#11 0.283 dependencies:
#11 0.283 - python=3.10.12
#11 0.283 - pip<=22.2.2
#11 0.283 - pip:
#11 0.283   - auto-gptq==0.4.2
#11 0.283   - bitsandbytes==0.41.1
#11 0.283   - boto3==1.24.28
#11 0.283   - cloudpickle==2.0.0
#11 0.283   - configparser==5.2.0
#11 0.283   - defusedxml==0.7.1
#11 0.283   - flash-attn==2.3.2
#11 0.283   - google-cloud-storage==2.10.0
#11 0.283   - gunicorn==20.1.0
#11 0.283   - ipython==8.10.0
#11 0.283   - numpy==1.21.5
#11 0.283   - optimum==1.13.2
#11 0.283   - packaging==21.3
#11 0.283   - pandas==1.4.4
#11 0.283   - protobuf==3.19.4
#11 0.283   - pyarrow==8.0.0
#11 0.283   - pyyaml==6.0
#11 0.283   - requests==2.28.1
#11 0.283   - scipy==1.9.1
#11 0.283   - sentencepiece==0.1.99
#11 0.283   - soundfile==0.12.1
#11 0.283   - tensorflow==2.11.1
#11 0.283   - auto-gptq==0.4.2
#11 0.283   - optimum==1.13.2
#11 0.283   - bitsandbytes==0.41.1
#11 0.283   - accelerate==0.23.0
#11 0.283   - transformers==4.34.0
#11 0.283   - mlflow==2.7.0
#11 0.283   - torch==2.1.0
#11 0.283 name: mlflow-env
#11 0.616 Collecting package metadata (repodata.json): ...working... done
#11 55.01 Solving environment: ...working... done
#11 58.25 
#11 58.25 
#11 58.25 ==> WARNING: A newer version of conda exists. <==
#11 58.25   current version: 4.10.3
#11 58.25   latest version: 23.9.0
#11 58.25 
#11 58.25 Please update conda by running
#11 58.25 
#11 58.25     $ conda update -n base -c defaults conda
#11 58.25 
#11 58.25 
#11 58.26 
#11 58.26 Downloading and Extracting Packages
#11 58.26 
tzdata-2023c         | 115 KB    |            |   0% 
tzdata-2023c         | 115 KB    | #3         |  14% 
tzdata-2023c         | 115 KB    | ########## | 100% 
#11 58.46 
libsqlite-3.43.2     | 820 KB    |            |   0% 
libsqlite-3.43.2     | 820 KB    | ########## | 100% 
#11 58.55 
pip-22.2.2           | 1.5 MB    |            |   0% 
pip-22.2.2           | 1.5 MB    | ########## | 100% 
pip-22.2.2           | 1.5 MB    | ########## | 100% 
#11 58.90 
libgomp-13.2.0       | 411 KB    |            |   0% 
libgomp-13.2.0       | 411 KB    | ########## | 100% 
#11 58.95 
xz-5.2.6             | 409 KB    |            |   0% 
xz-5.2.6             | 409 KB    | ########## | 100% 
xz-5.2.6             | 409 KB    | ########## | 100% 
#11 59.08 
wheel-0.41.2         | 56 KB     |            |   0% 
wheel-0.41.2         | 56 KB     | ########## | 100% 
#11 59.15 
libzlib-1.2.13       | 60 KB     |            |   0% 
libzlib-1.2.13       | 60 KB     | ########## | 100% 
#11 59.21 
bzip2-1.0.8          | 484 KB    |            |   0% 
bzip2-1.0.8          | 484 KB    | ########## | 100% 
bzip2-1.0.8          | 484 KB    | ########## | 100% 
#11 59.32 
libffi-3.4.2         | 57 KB     |            |   0% 
libffi-3.4.2         | 57 KB     | ########## | 100% 
#11 59.38 
ncurses-6.4          | 860 KB    |            |   0% 
ncurses-6.4          | 860 KB    | ########## | 100% 
ncurses-6.4          | 860 KB    | ########## | 100% 
#11 59.64 
openssl-3.1.3        | 2.5 MB    |            |   0% 
openssl-3.1.3        | 2.5 MB    | ########## | 100% 
openssl-3.1.3        | 2.5 MB    | ########## | 100% 
#11 59.77 
tk-8.6.13            | 3.1 MB    |            |   0% 
tk-8.6.13            | 3.1 MB    | ########## | 100% 
tk-8.6.13            | 3.1 MB    | ########## | 100% 
#11 59.90 
libuuid-2.38.1       | 33 KB     |            |   0% 
libuuid-2.38.1       | 33 KB     | ########## | 100% 
#11 59.96 
ca-certificates-2023 | 146 KB    |            |   0% 
ca-certificates-2023 | 146 KB    | ########## | 100% 
#11 60.01 
ld_impl_linux-64-2.4 | 688 KB    |            |   0% 
ld_impl_linux-64-2.4 | 688 KB    | ########## | 100% 
#11 60.09 
libgcc-ng-13.2.0     | 753 KB    |            |   0% 
libgcc-ng-13.2.0     | 753 KB    | ########## | 100% 
#11 60.16 
libnsl-2.0.1         | 33 KB     |            |   0% 
libnsl-2.0.1         | 33 KB     | ########## | 100% 
#11 60.21 
readline-8.2         | 275 KB    |            |   0% 
readline-8.2         | 275 KB    | ########## | 100% 
#11 60.26 
python-3.10.12       | 24.4 MB   |            |   0% 
python-3.10.12       | 24.4 MB   | ##9        |  29% 
python-3.10.12       | 24.4 MB   | #######3   |  73% 
python-3.10.12       | 24.4 MB   | ########## | 100% 
#11 60.95 
_libgcc_mutex-0.1    | 3 KB      |            |   0% 
_libgcc_mutex-0.1    | 3 KB      | ########## | 100% 
#11 60.99 
setuptools-68.2.2    | 454 KB    |            |   0% 
setuptools-68.2.2    | 454 KB    | ########## | 100% 
#11 61.07 
_openmp_mutex-4.5    | 23 KB     |            |   0% 
_openmp_mutex-4.5    | 23 KB     | ########## | 100% 
#11 61.12 Preparing transaction: ...working... done
#11 61.33 Verifying transaction: ...working... done
#11 62.26 Executing transaction: ...working... done
#11 64.36 Installing pip dependencies: ...working... Pip subprocess error:
#11 68.27   error: subprocess-exited-with-error
#11 68.27   
#11 68.27   × python setup.py egg_info did not run successfully.
#11 68.27   │ exit code: 1
#11 68.27   ╰─> [6 lines of output]
#11 68.27       Traceback (most recent call last):
#11 68.27         File "<string>", line 2, in <module>
#11 68.27         File "<pip-setuptools-caller>", line 34, in <module>
#11 68.27         File "/tmp/pip-install-xbx5x5ho/flash-attn_d51c6af2cbbd43a7ba38c9152f326ee4/setup.py", line 8, in <module>
#11 68.27           from packaging.version import parse, Version
#11 68.27       ModuleNotFoundError: No module named 'packaging'
#11 68.27       [end of output]
#11 68.27   
#11 68.27   note: This error originates from a subprocess, and is likely not a problem with pip.
#11 68.27 error: metadata-generation-failed
#11 68.27 
#11 68.27 × Encountered error while generating package metadata.
#11 68.27 ╰─> See above for output.
#11 68.27 
#11 68.27 note: This is an issue with the package mentioned above, not pip.
#11 68.27 hint: See above for details.
#11 68.27 
#11 68.27 Ran pip subprocess with arguments:
#11 68.27 ['/opt/conda/envs/mlflow-env/bin/python', '-m', 'pip', 'install', '-U', '-r', '/model/condaenv.s6k7436y.requirements.txt']
#11 68.27 Pip subprocess output:
#11 68.27 Collecting auto-gptq==0.4.2
#11 68.27   Downloading auto_gptq-0.4.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.8 MB)
#11 68.27      ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.8/1.8 MB 31.7 MB/s eta 0:00:00
#11 68.27 Collecting bitsandbytes==0.41.1
#11 68.27   Downloading bitsandbytes-0.41.1-py3-none-any.whl (92.6 MB)
#11 68.27      ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 92.6/92.6 MB 19.4 MB/s eta 0:00:00
#11 68.27 Collecting boto3==1.24.28
#11 68.27   Downloading boto3-1.24.28-py3-none-any.whl (132 kB)
#11 68.27      ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 132.5/132.5 kB 25.5 MB/s eta 0:00:00
#11 68.27 Collecting cloudpickle==2.0.0
#11 68.27   Downloading cloudpickle-2.0.0-py3-none-any.whl (25 kB)
#11 68.27 Collecting configparser==5.2.0
#11 68.27   Downloading configparser-5.2.0-py3-none-any.whl (19 kB)
#11 68.27 Collecting defusedxml==0.7.1
#11 68.27   Downloading defusedxml-0.7.1-py2.py3-none-any.whl (25 kB)
#11 68.27 Collecting flash-attn==2.3.2
#11 68.27   Downloading flash_attn-2.3.2.tar.gz (2.3 MB)
```



 

Comments:
Hi @Girrajjangid 
Thanks for the issue, I think that you might be having a package conflict, can you make sure to have flash attention version greater than 2.0?
Also, you might consider switching to main branch to avoid this issue e.g. this fix: https://github.com/huggingface/transformers/pull/26785 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.