Enhancing SFT Training Efficiency Using Packing and FlashAttention2 with Position IDs

[RhuiDih]

What does this PR do?

Improve throughput as well as train time and memory utilization for instruction tuning by enabling padding-free and attention mask-free attention.

Specifically, this PR adds the capability to utilize position_ids in FlashAttention2 _flash_attention_forward() in the case of packing (attention_mask=None) to models which use position_ids in their respective DecoderLayer implementations.

This PR also adds a new off-the-shelf data collator DataCollatorWithFlattening which packs the examples in a mini batch into one long sequence and return position_ids as well and turns the first token of labels to -100 to prevent the last token of previous example predicting the first token of next example.

This enables the following:

1. Use of packing for instruction tuning without incorrect cross-example attention
2. Significant increase in training throughput and reduction in memory utilization

Example Result 1
dataset: OrcaMath subset
setup: FSDP with 8 GPUs

Model	DataProcess	Time	Throughput (token/s)	Memory (MB)
Llama2-7B	Padding	790	1269	22305
Llama2-7B	ThisPR	574	1746	20950
Mistral-7B	Padding	812	1216	23603
Mistral-7B	ThisPR	596	1658	22409

Example Result 2
dataset: FLAN subset
setup: FSDP with 8 GPUs

Model	DataProcess	Time	Throughput (token/s)	Memory (MB)
Llama2-7B	Padding	1526	771	29234
Llama2-7B	ThisPR	809	1455	23854
Mistral-7B	Padding	742	742	30625
Mistral-7B	ThisPR	1408	1408	24549

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Models:

• text models: @ArthurZucker

Edit - add more description on data collator

[ArthurZucker]

Thanks! I think we can wait a tad bit, #31446 is almost ready for merge!

[ArthurZucker]

this is an input dependent control flow which won't be supported by compile, but compile is already not supported.
This needs a comment: we are checking that the input is not padded, and that we are doing prefill right?

[ArthurZucker]

we never import from another modeling code in transformers, we use "copied from" or we define the general function in falsh_attention_utils for example

[ArthurZucker]

this needs to be documented to explain what is happening

[RhuiDih]

@ArthurZucker #31446 has just merged, let me rebase and update this PR

[RhuiDih]

@ArthurZucker Rebase is done.

I have move the codes to modeling_flash_attention_utils.py and added required comments/description.

Regarding the control flow, ideally the cu_seq_len should be prepared in data collator and model should accept it as arg to trigger flash_attn_varlen_func without relying on position_ids, but this is the quick & fast way to do padding free.

[ArthurZucker]

This LGTM! Only thing missing is a test!
Testing this codepath in particuler, + datacollator in the trainer tests maybe?

[ArthurZucker]

I think we keep track of a list of models that support FA2, if needed for guidance here!

[ArthurZucker]

this is an input dependent control flow, but we already don't support compile with FA2 AFAIK.
Alright for now!

[mayank31398]

I think loss computation is wrong for everything
it needs to be:
https://github.com/ibm-granite/dolomite-engine/blob/7e2527528ac54cb3388040a3052ef3b6d845a40d/dolomite_engine/hf_models/models/gpt_dolomite/main.py#L185-L191

[RhuiDih]

@mayank31398
it can be done within the model or when labels preparation in data collator

transformers/src/transformers/data/data_collator.py

Lines 1634 to 1637 in 808fd63

	if is_labels_provided:
	ret["labels"] += [-100] + features[idx]["labels"][1:]
	else:
	ret["labels"] += [-100] + features[idx]["input_ids"][1:]

[wynterl]

Hello @mayank31398 We are providing the capability to define the labels as the user wishes in the data collator in this PR via transformers/src/transformers/data/data_collator.py

[RhuiDih]

@ArthurZucker
I have added tests for the new data collator and the models which already have FA2 tests in place. For the trainer, i can't find existing FA2 + Trainer testing combo, could you point me there ?

[ArthurZucker]

LGTM, I just want a second look from our @fxmarty .
This should improve performances for everyone. but let's make sure we don't break BC! WDYT @fxmarty ?

[ArthurZucker]

No worries I think the tests you added are enough

[fxmarty]

Looks good!

I strongly think this should be documented somewhere, maybe https://huggingface.co/docs/transformers/main/en/perf_train_gpu_one#flash-attention-2.

[fxmarty]

As this is already in test_modeling_common.py, adding this here is not needed.

[fxmarty]

same for all

[RhuiDih]

i have removed them in latest commit

[fxmarty]

I am not a big fan as having unpadded inputs as [1, total_seqlens]. To me, it would make more sense to have support for one dimensional tensors as input_ids ([total_tokens]), which becomes after embedding inputs_embeds ([total_tokens, hidden_size]), as anyway flash attention varlen frontend expects 3D tensors, not 4D.

Thus

    query = query.view(-1, query.size(-2), query.size(-1))
    key = key.view(-1, key.size(-2), key.size(-1))
    value = value.view(-1, value.size(-2), value.size(-1))

would not be needed, and neither calculating the cumulative seqlens at each layer. This maybe can be left for an other PR.

[RhuiDih]

very much agree it is less than ideal that cu_seq_len is calculated again in every layers. It can be done in data_collator or in the first forward() and passed along all DecoderLayer, not sure how it links to switching from [bs, seq_len, ... ] to just [seq_len, ...]. This removal of bs dimension requires substantial change to the model implementation, should HF decide to do this switch, those 3 lines can be just removed.

[fxmarty]

How should a user make sure?

[RhuiDih]

no easy way, the user will have to make sure the model supports FA2 and have position_ids
until the model accepting cu_seq_len which will allow every model remove padding once for all, i think this serve as a good warning

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

🚀 nice getting this done!

[olisicky]

Hi, thank you for this! I tried to use it without DataCollatorWithFlattening but with data prepared with input_ids and position_ids, but I encounter RuntimeError: CUDA error: an illegal memory access was encountered just after elif position_ids is not None and not (torch.diff(position_ids, dim=-1) >= 0).all() and query_length != 1: was called.

I am using:

• llama3
• flash_attn version: 2.5.8
• deepspeed - ZeRO1
• transformers 4.44.1

Thank you!

[RhuiDih]

@olisicky Hi, could you provide a minimal code to reproduce the error ? that would help greatly

[olisicky]

@olisicky Hi, could you provide a minimal code to reproduce the error ? that would help greatly

Hi @RhuiDih . I found a mistake in my data as I was preparing a minimal code for you:D. So thank you for encouraging me to do so! As I am preparing data without DataCollatorWithFlattening I kept padding there to context_length even for bs=1 and it resulted in plenty of 0 in position_ids tensor on indexes corresponding to PAD tokens.

However, I have an additional question. You are setting -100 in labels for the first token in each sequence in the batch. I wonder why. I thought that the use of position_ids itself should reduce cross-example attention but here you mentioned that the -100 should prevent the last token of previous example predicting the first token of next example.

If I would have a sequences which start with BOS token and then it will be ignored during training by setting -100 in labels, then I would lost the importance of this first token. Is it necessary? Or maybe shouldn't we use some dummy token there rather then a part of the sequence itself?

Thank you very much!

[RhuiDih]

@olisicky
The setting of -100 is due to causal loss implementation, if we were flattening all sequences into batch of 1 without setting -100 accordingly, it would incur unwanted loss.

First token of input_ids is still used to produce logits, hence no importance is lost. What we want to ignore, its the first token of labels due to the mentioned causal loss implementation.
Hope it is clear.

[olisicky]

@olisicky The setting of -100 is due to causal loss implementation, if we were flattening all sequences into batch of 1 without setting -100 accordingly, it would incur unwanted loss.

First token of input_ids is still used to produce logits, hence no importance is lost. What we want to ignore, its the first token of labels due to the mentioned causal loss implementation. Hope it is clear.

Sorry for late reply. Yes, thank you. I will try it.

[MahmoudAshraf97]

Is packing usable in transformers generate method? AFAIK only padding is supported

[ArthurZucker]

Packing is supported with the collator introduced. We don't pack when generating tho, PRs are welcome for this! 🤗

[qibao77]

What does this PR do?

Improve throughput as well as train time and memory utilization for instruction tuning by enabling padding-free and attention mask-free attention.

Specifically, this PR adds the capability to utilize position_ids in FlashAttention2 _flash_attention_forward() in the case of packing (attention_mask=None) to models which use position_ids in their respective DecoderLayer implementations.

This PR also adds a new off-the-shelf data collator DataCollatorWithFlattening which packs the examples in a mini batch into one long sequence and return position_ids as well and turns the first token of labels to -100 to prevent the last token of previous example predicting the first token of next example.

This enables the following:

1. Use of packing for instruction tuning without incorrect cross-example attention
2. Significant increase in training throughput and reduction in memory utilization

Example Result 1 dataset: OrcaMath subset setup: FSDP with 8 GPUs

Model DataProcess Time Throughput (token/s) Memory (MB)
Llama2-7B Padding 790 1269 22305
Llama2-7B ThisPR 574 1746 20950
Mistral-7B Padding 812 1216 23603
Mistral-7B ThisPR 596 1658 22409
Example Result 2 dataset: FLAN subset setup: FSDP with 8 GPUs

Model DataProcess Time Throughput (token/s) Memory (MB)
Llama2-7B Padding 1526 771 29234
Llama2-7B ThisPR 809 1455 23854
Mistral-7B Padding 742 742 30625
Mistral-7B ThisPR 1408 1408 24549

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Models:

• text models: @ArthurZucker

Edit - add more description on data collator

How to use this feature for pretraining?

[RhuiDih]

@qibao7 usually one would truncate and pack to max length during pretraining, this PR will not benefit pretraining, unless most of your pretraining data are short