首页 > 其他分享 >通过打包 Flash Attention 来提升 Hugging Face 训练效率

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

时间:2024-09-12 12:24:45浏览次数:10  
标签:DataCollatorWithFlattening 示例 Flash Hugging Face train 使用 dataset 打包

简单概述

现在,在 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_idslabelsposition_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

标签:DataCollatorWithFlattening,示例,Flash,Hugging,Face,train,使用,dataset,打包
From: https://www.cnblogs.com/huggingface/p/18409936

相关文章

  • <<编码>> 第 4 章 手电筒剖析(Anatomy of a Flashlight) 示例电路
    简单灯泡电路info::操作说明鼠标单击按钮开关切换开合状态另:黄色小点为电流,后同.可通过“菜单–选项–显示电流”控制是否显示primary::在线交互操作链接https://cc.xiaogd.net/?startCircuitLink=https://book.xiaogd.net/code-hlchs-examples/assets/circ......
  • 如何使用huggingface下载数据集和预训练模型
            如果各位在下载huggingface上的模型和数据库也会出现“connectclosed/failed”等错误,不妨试试下面的解决方案,思路大致是,通过设置镜像的方式来解决。下载数据集1.找到你要下载的数据库名称,并复制2.打开终端,并选择需要使用的conda环境,编写bash文件(或者......
  • “Interface 和 Type 区别”深度解析
    “Interface和Type区别”深度解析文章目录一、Interface和Type是什么二、如何使用Interface和Type1.定义Interface2.定义Type3.使用Interface和Type4.区别与联系三、Interface和Type二者有哪些区别,分别在哪些场景使用1.区别2.场景......
  • 墙裂推荐:《Transformer自然语言处理实战:使用Hugging-Face-Transformers库构建NLP应用
    大家好,今天给大家推荐一本大模型神书——《Transformer自然语言处理实战:使用Hugging-Face-Transformers库构建NLP应用》。近年来,Transformer模型在NLP领域取得了显著成果。为了让广大开发者更好地掌握这一技术,给大家推荐一本实战教程——《Transformer自然语言处理实战:使用......
  • Android中SurfaceView的双缓冲机制和普通View叠加问题解决办法
    本文首发于公众号“AntDream”,欢迎微信搜索“AntDream”或扫描文章底部二维码关注,和我一起每天进步一点点SurfaceView是Android平台上用于高效渲染图形的视图控件。它将内容绘制在一个独立的Surface上,可以直接由渲染线程访问,从而提高性能,尤其是在需要频繁刷新和更新......
  • 解锁Python中的人脸识别:Face Recognition库详解与应用
    在当今的人工智能时代,人脸识别技术已经成为了计算机视觉领域的一项重要应用。无论是在安全监控、社交媒体还是智能设备中,人脸识别都扮演着不可或缺的角色。在众多的人脸识别工具和库中,Python的FaceRecognition库以其简单易用和高效性而备受青睐。本文将深入探讨FaceRecogniti......
  • Hugging Face 的应用
    大纲Hugging-Face介绍Hugging-Face大语言模型LLM管理Transformers机器学习框架文本生成推理(TGI)HuggingFaceHugging-Face--大语言模型界的GithubHuggingFace专门开发用于构建机器学习应用的工具。该公司的代表产品是其为自然语言处理应用构建的transformers库,以及允......
  • Hugging Face 的应用
    大纲Hugging-Face介绍Hugging-Face大语言模型LLM管理Transformers机器学习框架文本生成推理(TGI)HuggingFaceHugging-Face--大语言模型界的GithubHuggingFace专门开发用于构建机器学习应用的工具。该公司的代表产品是其为自然语言处理应用构建的transformer......
  • OpenGL ES通过缩小GLSurfaceView来解决纹理贴图变形的问题
    一、概述在使用OpenGLES做纹理贴图的时候,图片有小有大。默认情况下纹理是撑满整个屏幕的。这就导致大图会被压扁、小图会被拉伸。这种体验相当不好。解决此问题的其中一种方式是:通过缩小GLSurfaceView的宽或高来解决问题。ps:公式可以看做是固定的,直接使用即可。......
  • PLC结构化文本(ST)——接口引用转换(__QUERYINTERFACE)
    PLCStructuredTextObjectOrientedProgrammingPLC结构化文本(ST)——接口引用转换(__QUERYINTERFACE)__QUERYINTERFACE运算符__QUERYPOINTER是IEC61131-3的扩展,该运算符允许在运行时将一个接口引用转换成另一个接口的引用。返回值BOOL类型:True表示转换成功,Flase表示转换失败。......