Enhancing SFT Training Efficiency Using Packing and FlashAttention2 with Position IDs
Author: RhuiDih
Pull request: yes

# 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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? 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](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

## Who can review?
Models:
- text models: @ArthurZucker

Edit - add more description on data collator

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

@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.


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

@mayank31398 
it can be done within the model or when labels preparation in data collator
https://github.com/huggingface/transformers/blob/808fd638182289cf1903a1d4c536f00c48809fb8/src/transformers/data/data_collator.py#L1634-L1637


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](https://github.com/huggingface/transformers/blob/808fd638182289cf1903a1d4c536f00c48809fb8/src/transformers/data/data_collator.py#L1634-L1637)

@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 ?

No worries I think the tests you added are enough

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_31629). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

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!

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

> @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!

@olisicky 
The setting of `-100` is due to [causal loss implementation](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1221-L1229), 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 The setting of `-100` is due to [causal loss implementation](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1221-L1229), 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. 

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

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

> # 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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
>   Pull Request section?
> * [ ]  Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? 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](https://github.com/huggingface/transformers/tree/main/docs), and
>   [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
> * [ ]  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?

@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

@RhuiDih This is a neat feature! Thank you! There is one thing I am confused about. 

The doc in [1] says "Using DataCollatorWithFlattening will flatten the entire mini batch into single long sequence".

Suppose that I have 15 examples in the dataset, with 10 1k length examples, and 5 1.5k length examples. If I used max-seq-length=2k and used per-device-batch-size=2. What happens if the mini-batch uses 2 1.5k examples? Their combined length becomes 3k. Will the algorithm just truncate 3k to 2k or it might lead to OOM?

[1]: DataCollatorWithFlattening