通过打包 Flash Attention 来提升 Hugging Face 训练效率

简单概述

现在，在 Hugging Face 中，使用打包的指令调整示例 (无需填充) 进行训练已与 Flash Attention 2 兼容，这要归功于一个最近的 PR以及新的DataCollatorWithFlattening。

• 最近的 PRhttps://github.com/huggingface/transformers/pull/31629

• DataCollatorWithFlatteninghttps://hf.co/docs/transformers/main/en/main_classes/data_collator#transformers.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) 的样本的平均值。

FLAN 数据集的平均序列较短，但序列长度差异较大，因此每个批次中的示例长度可能会有很大差异。这意味着填充的 FLAN 批次可能会因为未使用的填充 token 而产生显著的开销。在 FLAN 数据集上进行训练时，使用新的 DataCollatorWithFlattening 在提高吞吐量方面显示出显著的优势。我们在这里展示的模型中看到了 2 倍的吞吐量提升: llama2-7B、mistral-7B 和 granite-8B-code。

OrcaMath 数据集的示例较长，且示例长度差异较小。因此，从打包中获得的改进较低。我们的实验显示，在使用这种打包方式在 OrcaMath 数据集上训练时，这三个模型的吞吐量增加了 1.4 倍。

通过使用新的 DataCollatorWithFlattening 进行打包，内存使用也有所改善。下图显示了相同的三个模型在相同的两个数据集上训练时的峰值内存使用情况。在 FLAN 数据集上，峰值内存减少了 20%，这得益于打包的显著好处。

在 OrcaMath 数据集上，由于其示例长度更为均匀，峰值内存减少了 6%。

当打包示例减少了优化步骤的数量时，可能会损害训练的收敛性。然而，这个新功能保留了小批量，因此与使用填充示例相同的优化步骤数量。因此，对训练收敛性没有影响，正如我们在下一个图中看到的那样，该图显示了相同的三个模型在相同的两个数据集上训练时，无论是使用新的 DataCollatorWithFlattening 进行打包还是使用填充，模型的验证损失是相同的。

工作原理

考虑一个批处理数据，其中批量大小 (batchsize) 为 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 ，只需两个步骤:

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://hf.co/papers/2407.09105。

---

如何学习AI大模型？

？

大模型时代，火爆出圈的LLM大模型让程序员们开始重新评估自己的本领。 “AI会取代那些行业？”“谁的饭碗又将不保了？”等问题热议不断。

不如成为「掌握AI工具的技术人」，毕竟AI时代，谁先尝试，谁就能占得先机！

想正式转到一些新兴的 AI 行业，不仅需要系统的学习AI大模型。同时也要跟已有的技能结合，辅助编程提效，或上手实操应用，增加自己的职场竞争力。

但是LLM相关的内容很多，现在网上的老课程老教材关于LLM又太少。所以现在小白入门就只能靠自学，学习成本和门槛很高

那么针对所有自学遇到困难的同学们，我帮大家系统梳理大模型学习脉络，将这份 LLM大模型资料 分享出来：包括LLM大模型书籍、640套大模型行业报告、LLM大模型学习视频、LLM大模型学习路线、开源大模型学习教程等, 😝有需要的小伙伴，可以 扫描下方二维码领取🆓↓↓↓

👉[CSDN大礼包🎁：全网最全《LLM大模型入门+进阶学习资源包》免费分享（安全链接，放心点击）]()👈

学习路线

第一阶段： 从大模型系统设计入手，讲解大模型的主要方法；

第二阶段： 在通过大模型提示词工程从Prompts角度入手更好发挥模型的作用；

第三阶段： 大模型平台应用开发借助阿里云PAI平台构建电商领域虚拟试衣系统；

第四阶段： 大模型知识库应用开发以LangChain框架为例，构建物流行业咨询智能问答系统；

第五阶段： 大模型微调开发借助以大健康、新零售、新媒体领域构建适合当前领域大模型；

第六阶段： 以SD多模态大模型为主，搭建了文生图小程序案例；

第七阶段： 以大模型平台应用与开发为主，通过星火大模型，文心大模型等成熟大模型构建大模型行业应用。

👉学会后的收获：👈

• 基于大模型全栈工程实现（前端、后端、产品经理、设计、数据分析等），通过这门课可获得不同能力；

• 能够利用大模型解决相关实际项目需求： 大数据时代，越来越多的企业和机构需要处理海量数据，利用大模型技术可以更好地处理这些数据，提高数据分析和决策的准确性。因此，掌握大模型应用开发技能，可以让程序员更好地应对实际项目需求；

• 基于大模型和企业数据AI应用开发，实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能， 学会Fine-tuning垂直训练大模型（数据准备、数据蒸馏、大模型部署）一站式掌握；

• 能够完成时下热门大模型垂直领域模型训练能力，提高程序员的编码能力： 大模型应用开发需要掌握机器学习算法、深度学习框架等技术，这些技术的掌握可以提高程序员的编码能力和分析能力，让程序员更加熟练地编写高质量的代码。

1.AI大模型学习路线图
2.100套AI大模型商业化落地方案
3.100集大模型视频教程
4.200本大模型PDF书籍
5.LLM面试题合集
6.AI产品经理资源合集

👉获取方式：
😝有需要的小伙伴，可以保存图片到wx扫描二v码免费领取【保证100%免费】🆓