Feat: Add support for upstream FA2

[NanoCode012]

Untested!

Only works for

• llama
• falcon

TODO:

• Test

Ref Upstream PR: https://github.com/huggingface/transformers/pull/25598/files

[mhenrichsen]

Maybe add is_falcon_derived_model to the falcon examples?

[winglian]

there is some falcon patching that also already happens here https://github.com/OpenAccess-AI-Collective/axolotl/blob/f3d939016a6ec681be40e83c6a57c682cb60e2b1/src/axolotl/utils/models.py#L117-L127

[NanoCode012]

Does sample packing work for falcon?

I'll consolidate the check for falcon derived model.

[winglian]

Does sample packing work for falcon?

I'll consolidate the check for falcon derived model.

no, it's basic flash attention support only

[NanoCode012]

Tested with

CUDA_VISIBLE_DEVICES=0 accelerate launch -m axolotl.cli.train examples/openllama-3b/lora.yml

using sample_packing off.

With main 5e5296a, it gives: TypeError: flashattn_forward() got an unexpected keyword argument 'padding_mask'
With this PR,

[2023-09-26 13:00:15,139] [INFO] [axolotl.train.train:84] [PID:751] [RANK:0] Pre-saving adapter config to ./lora-out
[2023-09-26 13:00:15,141] [INFO] [axolotl.train.train:108] [PID:751] [RANK:0] Starting trainer...
{'loss': 1.3152, 'learning_rate': 1e-05, 'epoch': 0.0}

---

With sample_packing on, it runs as usual. main has no issues here either.