透過打包 Flash Attention 來提升 Hugging Face 訓練效率

HuggingFace發表於2024-09-12
Hugging Face

簡單概述

現在，在 Hugging Face 中，使用打包的指令調整示例 (無需填充) 進行訓練已與 Flash Attention 2 相容，這要歸功於一個 最近的 PR 以及新的 DataCollatorWithFlattening。

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

簡介

在訓練期間，對小批次輸入進行填充是一種常見的整理輸入資料的方法。然而，由於無關的填充 token ，這引入了效率低下的問題。不進行填充而是打包示例，並使用 token 位置資訊，是一種更高效的選擇。然而，之前打包的實現並沒有在使用 Flash Attention 2 時考慮示例邊界，導致出現不希望的跨示例注意力，這降低了質量和收斂性。

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

透過選擇 DataCollatorWithFlattening ，Hugging Face Trainer 的使用者現在可以無縫地將序列連線成一個單一的張量，同時在 Flash Attention 2 計算過程中考慮到序列邊界。這是透過 flash_attn_varlen_func 實現的，它計算每個小批次的累積序列長度 ( cu_seqlens )。
同樣的功能也適用於 TRL 庫中的 Hugging Face SFTTrainer 使用者，透過在呼叫資料整理器 DataCollatorForCompletionOnlyLM 時設定一個新的標誌 padding_free=True 來實現。

吞吐量提高多達 2 倍

我們使用帶有新 DataCollatorWithFlattening 的此功能在訓練過程中看到了顯著的處理吞吐量提升。下圖顯示了在訓練期間測量的吞吐量，單位為 token /秒。在這個例子中，吞吐量是在 8 個 A100-80 GPU 上對一個 epoch 內的 20K 個隨機選自兩個不同指令調整資料集 (FLAN 和 OrcaMath) 的樣本的平均值。

throughput

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

OrcaMath 資料集的示例較長，且示例長度差異較小。因此，從打包中獲得的改進較低。我們的實驗顯示，在使用這種打包方式在 OrcaMath 資料集上訓練時，這三個模型的吞吐量增加了 1.4 倍。

memory

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

在 OrcaMath 資料集上，由於其示例長度更為均勻，峰值記憶體減少了 6%。

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

ValLoss

工作原理

考慮一個批處理資料，其中批次大小 (batchsize) 為 4，四個序列如下:

batch

在將示例連線之後，無填充整理器返回每個示例的 input_ids 、 labels 和 position_ids 。因此，對於這批資料，整理器提供了以下內容:

example

所需的修改是輕量級的，僅限於向 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 ，只需兩個步驟:

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

如果你正在使用 TRL 中的 Hugging Face SFTTrainer 配合 DataCollatorForCompletionOnlyLM ，那麼所需的兩個步驟是:

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

如何使用它

對於 Trainer 使用者，下面的例子展示瞭如何使用這個新功能。

# 使用 DataCollatorWithFlattening 的示例
 
import torch

# 載入模型
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "instructlab/merlinite-7b-lab",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2"
)

# 讀取資料集
from datasets import load_dataset
train_dataset = load_dataset("json", data_files="path/to/my/dataset")["train"]

# 使用 DataCollatorWithFlattening
from transformers import DataCollatorWithFlattening
data_collator = DataCollatorWithFlattening()

# 訓練
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 時使用這個新功能。

# 使用 DataCollatorForCompletionOnlyLM SFTTrainer 示例

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。

英文原文: https://hf.co/blog/packing-with-FA2

原文作者: Rhui Dih Lee, Arthur Zucker, Achintya Kundu, Laura Wynter, Raghu Ganti, Mayank Mishra

譯者: innovation64

相關文章