Skip to content

update llama flash attention#646

Merged
research4pan merged 4 commits into
OptimalScale:mainfrom
yaoguany:main
Sep 19, 2023
Merged

update llama flash attention#646
research4pan merged 4 commits into
OptimalScale:mainfrom
yaoguany:main

Conversation

@yaoguany

Copy link
Copy Markdown
Collaborator

1.update llama flash attention code(past_key_value, attention_mask)
2.support flash attention + kv_cache now.

@yaoguany

yaoguany commented Sep 19, 2023

Copy link
Copy Markdown
Collaborator Author

a40 train

dataset:alpaca

llama2-7b, 2048 blocksize, 10 batchsize, gradient checkpoint
with flash attention: memory:19672mb
without flash attention: memory:42700mb

a40 eval

dataset:MedQA-USMLE validation, block_size:4096, batch size:10

acc

same acc in first 200 samples:0.41

memory:

with flash attention:39752mb
without flash attention: 34876mb

@yaoguany

yaoguany commented Sep 19, 2023

Copy link
Copy Markdown
Collaborator Author

bugs need to be fixed in the future

  • eval with flash attention still occupies much memory(the eval result is same)
  • use_cache should be set to True when use flash attention, otherwise the eval process may get stuck.

@research4pan research4pan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

LGTM, update llama flash attention implementation to make it compatible with latest version of transformers and flash_attention. Major changes are

  • Support past_key_value cache optimization of reusing key value tensor caches during token-by-token decoding
  • Flash attention is not enabled during token-by-token decoding or incremental state decoding

@research4pan
research4pan merged commit c44175b into OptimalScale:main Sep 19, 2023
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