透過使用 Flash Attention 提高 Hugging Face 訓練效率

釋出於 2024 年 8 月 21 日

摘要

引言

吞吐量提升高達 2 倍

工作原理

快速入門

如何使用

結論

簡而言之

透過最近的 PR 和新的 DataCollatorWithFlattening，現在 Hugging Face 中使用打包指令調優示例（無填充）的訓練與 Flash Attention 2 相容。

它可以在保持收斂質量的同時，將訓練吞吐量提高多達 2 倍。請繼續閱讀詳細資訊！

引言

在訓練期間，通常會採用在迷你批次中填充輸入序列的方法來整理輸入。然而，這會導致效率低下，因為存在不相關的填充標記。在不填充的情況下打包示例，並利用標記位置資訊，是一種更有效的替代方案。但是，以前的打包實現在使用 Flash Attention 2 時沒有考慮示例邊界，導致出現不必要的跨示例注意力，從而降低質量和收斂性。

Hugging Face Transformers 現在透過一項新功能解決了這個問題，該功能在打包過程中保持了邊界感知，同時引入了新的資料整理器 DataCollatorWithFlattening。

透過選擇 DataCollatorWithFlattening，Hugging Face Trainer 使用者現在可以無縫地將序列連線成單個張量，同時在 Flash Attention 2 計算期間考慮序列邊界。這是透過 flash_attn_varlen_func 實現的，它計算每個迷你批次中累積的序列長度（cu_seqlens）。

Hugging Face SFTTrainer 使用者在呼叫資料整理器 DataCollatorForCompletionOnlyLM 時，透過設定一個新標誌 padding_free=True，可以在 TRL 庫中使用相同的功能。

吞吐量提升高達 2 倍

我們發現，使用此功能和新的 DataCollatorWithFlattening 後，訓練吞吐量顯著提高。下圖顯示了訓練期間以 tokens/second 衡量的吞吐量。在此示例中，吞吐量是 8 個 A100-80 GPU 在一個 epoch 中，從兩個不同的指令調優資料集（FLAN 和 OrcaMath）中隨機選擇 2 萬個樣本的每 GPU 平均值。

FLAN 的序列平均較短，但序列長度的方差較大，因此每個批次中示例的長度可能差異很大。這意味著填充的 FLAN 批次可能會因未使用的填充標記而產生顯著的開銷。使用新的 DataCollatorWithFlattening 在 FLAN 資料集上進行訓練，在提高吞吐量方面顯示出顯著的優勢。我們看到此處顯示的模型：llama2-7B、mistral-7B 和 granite-8B-code 的吞吐量提高了 2 倍。

OrcaMath 具有更長的示例和更低的示例長度方差。因此，打包帶來的改進較低。我們的實驗表明，在 OrcaMath 資料集上使用這種形式的打包訓練，這三種模型的吞吐量提高了 1.4 倍。

使用新的 DataCollatorWithFlattening 進行打包也改善了記憶體使用。下圖顯示了相同三個模型在相同兩個資料集上訓練的峰值記憶體使用情況。FLAN 資料集上的峰值記憶體減少了 20%，這得益於打包的顯著優勢。

OrcaMath 資料集因其更同質的示例長度，峰值記憶體減少了 6%。

當減少最佳化步驟時，打包示例可能會損害訓練收斂。然而，新功能保留了迷你批次，因此最佳化步驟的數量與使用填充示例時相同。因此，對訓練收斂沒有影響，正如我們在下圖中所看到的，該圖顯示了相同三個模型在相同兩個資料集上訓練時，使用新的 DataCollatorWithFlattening 進行打包或使用填充訓練的模型具有相同的驗證損失。

工作原理

考慮一個批次大小為 4 的資料批次，其中四個序列如下：

連線示例後，無填充整理器會返回每個示例的 input_ids、labels 和 position_ids。因此，對於此資料批次，整理器提供：

所需的修改很輕量，僅限於為 Flash Attention 2 提供 position_ids。

然而，這依賴於模型公開 position_ids。截至本文撰寫之時，有 14 個模型公開了它們並受此解決方案支援。具體來說，Llama 2 和 3、Mistral、Mixtral、Granite、DBRX、Falcon、Gemma、OLMo、Phi 1、2 和 3、phi3、Qwen 2 和 2 MoE、StableLM 和 StarCoder 2 都受此解決方案支援。

入門

利用 position_ids 打包的好處很容易實現。

如果您使用的是 Hugging Face Transformers 中的 Trainer，只需兩個步驟：

使用 Flash Attention 2 例項化模型
使用新的 DataCollatorWithFlattening

如果您使用的是 TRL 中的 Hugging Face SFTTrainer 和 DataCollatorForCompletionOnlyLM，則需要兩個步驟：

使用 Flash Attention 2 例項化模型
呼叫 DataCollatorForCompletionOnlyLM 時，將 padding_free 設定為 True，如下所示：collator = DataCollatorForCompletionOnlyLM(response_template_ids, tokenizer=tokenizer, padding_free=True)

如何使用

對於 Trainer 使用者，以下示例說明了如何使用新功能。

# Example using DataCollatorWithFlattening
 
import torch

# load model as usual
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "instructlab/merlinite-7b-lab",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2"
)

# read dataset as usual
from datasets import load_dataset
train_dataset = load_dataset("json", data_files="path/to/my/dataset")["train"]

# use DataCollatorWithFlattening
from transformers import DataCollatorWithFlattening
data_collator = DataCollatorWithFlattening()

# train
from transformers import TrainingArguments, Trainer
train_args = TrainingArguments(output_dir="/save/path")
trainer = Trainer(
    args=train_args,
    model=model,
    train_dataset=train_dataset,
    data_collator=data_collator
)
trainer.train()

對於 TRL 使用者，以下示例展示瞭如何將新功能與 SFTTrainer 結合使用。

# SFTTrainer example using DataCollatorForCompletionOnlyLM

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM

dataset = load_dataset("lucasmccabe-lmi/CodeAlpaca-20k", split="train")

model = AutoModelForCausalLM.from_pretrained(
    "instructlab/merlinite-7b-lab",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2")
tokenizer = AutoTokenizer.from_pretrained("instructlab/merlinite-7b-lab")
tokenizer.pad_token = tokenizer.eos_token 

def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        text = f"### Question: {example['instruction'][i]}\n ### Answer: {example['output'][i]}"
        output_texts.append(text)
    return output_texts

response_template = " ### Answer:"
response_template_ids = tokenizer.encode(response_template, add_special_tokens=False)[2:]
collator = DataCollatorForCompletionOnlyLM(response_template_ids, tokenizer=tokenizer, padding_free=True)

trainer = SFTTrainer(
    model,
    train_dataset=dataset,
    args=SFTConfig(
        output_dir="./tmp",
        gradient_checkpointing=True,
        per_device_train_batch_size=8
    ),
    formatting_func=formatting_prompts_func,
    data_collator=collator,
)

trainer.train()

結論

得益於最近的 PR 和新的 DataCollatorWithFlattening，打包指令調優示例（而不是填充）現在與 Flash Attention 2 完全相容。該方法與使用 position_ids 的模型相容。在訓練過程中，吞吐量和峰值記憶體使用方面可以看到好處，而訓練收斂性沒有下降。實際的吞吐量和記憶體改進取決於模型和訓練資料中示例長度的分佈。透過使用 DataCollatorWithFlattening，訓練資料示例長度變化較大的模型將獲得相對於填充的最大收益。TRL 庫中的 SFTTrainer 使用者可以透過在呼叫 DataCollatorForCompletionOnlyLM 時設定新標誌 padding_free=True 來使用相同的功能。

有關更詳細的分析，請參閱論文：https://huggingface.co/papers/2407.09105