ImportError using Llama 2 with BetterTransformer

[jpgard]

System Info

I'm attempting to fine-tune a Llama2 model. My training loop works fine without BetterTransformers, but when I attempt to import the module it looks like there is an ImportError: ImportError: cannot import name '_expand_mask' from 'transformers.models.llama.modeling_llama' .

I already installed Transformers from git as described in this issue.

Version info:

• transformers @ git+https://github.com/huggingface/transformers@32f799db0d625ec5cf82624ff2604c5a891ebf61
• optimum==1.13.2
• Python 3.8.16

Full stack trace is below. Any suggestions? Thanks!

Traceback (most recent call last):
  File "scripts/train.py", line 235, in <module>
    main(model_arguments=model_args,
  File "scripts/train.py", line 81, in main
    from optimum.bettertransformer import BetterTransformer
  File "/path/to/miniconda3/envs/myenv/lib/python3.8/site-packages/optimum/bettertransformer/__init__.py", line 14, in <module>
    from .models import BetterTransformerManager
  File "/path/to/miniconda3/envs/myenv/lib/python3.8/site-packages/optimum/bettertransformer/models/__init__.py", line 17, in <module>
    from .decoder_models import (
  File "/path/to/miniconda3/envs/myenv/lib/python3.8/site-packages/optimum/bettertransformer/models/decoder_models.py", line 53, in <module>
    from .attention import (
  File "/path/to/miniconda3/envs/myenv/lib/python3.8/site-packages/optimum/bettertransformer/models/attention.py", line 24, in <module>
    from transformers.models.llama.modeling_llama import _expand_mask as _llama_expand_mask
ImportError: cannot import name '_expand_mask' from 'transformers.models.llama.modeling_llama' (/path/to/miniconda3/envs/myenv/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py)
Traceback (most recent call last):
  File "scripts/train.py", line 235, in <module>
    main(model_arguments=model_args,
  File "scripts/train.py", line 81, in main
    from optimum.bettertransformer import BetterTransformer
  File "/path/to/miniconda3/envs/myenv/lib/python3.8/site-packages/optimum/bettertransformer/__init__.py", line 14, in <module>
    from .models import BetterTransformerManager
  File "/path/to/miniconda3/envs/myenv/lib/python3.8/site-packages/optimum/bettertransformer/models/__init__.py", line 17, in <module>
    from .decoder_models import (
  File "/path/to/miniconda3/envs/myenv/lib/python3.8/site-packages/optimum/bettertransformer/models/decoder_models.py", line 53, in <module>
    from .attention import (
  File "/path/to/miniconda3/envs/myenv/lib/python3.8/site-packages/optimum/bettertransformer/models/attention.py", line 24, in <module>
    from transformers.models.llama.modeling_llama import _expand_mask as _llama_expand_mask
ImportError: cannot import name '_expand_mask' from 'transformers.models.llama.modeling_llama' (/path/to/miniconda3/envs/myenv/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py)

This is probably a question for @ArthurZucker based on his answer to the post linked above!

Who can help?

@ArthurZucker

Information

• The official example scripts
• My own modified scripts

Tasks

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

Reproduction

Installing the versions above and importing BetterTransformers should reproduce the issue but I haven't tried it.

Expected behavior

I expect to import BetterTransformer without errors.

[ArthurZucker]

cc @fxmarty this function is private and was removed in huggingface/transformers#26792!
Let's transfer this to optimum library

[jpgard]

@ArthurZucker thanks for the reply (and for all your and the HF team's work on the code!!). Does this mean the solution is to reinstall from git? (not sure how long PRs take to percolate into the version pip installs from git, or if it is instant)

[patrickvonplaten]

Worst case we can also leave and deprecate it in transformers if it's too breaking. Otherwise we force users to install optimum from "main" no? Happy to open a PR in transformers to leave it

[patrickvonplaten]

We'll deprecate it in Transformers: huggingface/transformers#27074 (comment)

[fxmarty]

It is private, I'll fix.

[patrickvonplaten]

For convenience sake, I deprecated the functions now in huggingface/transformers#27074 (comment)

[fxmarty]

Please use directly transformers>=4.36 with torch>=2.1.1 to benefit from PyTorch SDPA optimizations by default. BetterTransformer for Llama is deprecated: https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-and-memory-efficient-attention-through-pytorchs-scaleddotproductattention