Skip to content

Add flash-attention patch for falcon-7b - #3580

Merged
andreaskoepf merged 5 commits into
mainfrom
falcon7b_flash_attn_patch
Jul 19, 2023
Merged

Add flash-attention patch for falcon-7b#3580
andreaskoepf merged 5 commits into
mainfrom
falcon7b_flash_attn_patch

Conversation

@andreaskoepf

@andreaskoepf andreaskoepf commented Jul 17, 2023

Copy link
Copy Markdown
Collaborator

Enable the use_flash_attention configuration flag for Falcon models. When use_flash_attention is set to true the FalconAttention.forwad() method is replaced with a variant that uses Tri Dao's flash_attention instead of pytorch's scaled_dot_product_attention function.

At the moment the patch works only for falcon-7b but technically it will also work for falcon-40b with the right configuration. The falcon model situation is currently a bit messy: The Falcon model was recently added to Huggingface transformers (see PR transformers#24523) but the falcon models on the hugginface hub use still the code which is shipped together with the weights (a PR to change this was reverted). Falcon-7b and 40b use both slightly different code (which was unified in the HF transformers impl and can there be controlled via a configuration member called new_decoder_architecture see configuration_falcon.py#L65-L67). The HF Falcon impl uses different names in the configuration class, e.g. compare new configuration_falcon.py and old configuration_RW.py

HF Falcon implementation compatible model configurations can be found here:
7B: config.json
40B: config.json

@jordiclive jordiclive left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM, minor clarification

Comment thread model/model_training/models/__init__.py
Comment thread model/model_training/models/patching_utils.py
Comment thread model/pyproject.toml
Comment thread model/model_training/models/patching_falcon.py

@jordiclive jordiclive left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM

@andreaskoepf
andreaskoepf merged commit 1e6e569 into main Jul 19, 2023
@andreaskoepf
andreaskoepf deleted the falcon7b_flash_attn_patch branch July 19, 2023 14:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants