使用 DeepSpeed ZeRO、LoRA 和 Flash Attention 微调 Falcon 180B

Falcon 180B是Falcon LLM家族的最新版本。它是最大的开源模型，拥有180B参数，并在更多的数据上进行训练 - 3.5T个令牌，上下文长度窗口最多为4K个令牌。在这个示例中，我们将展示如何在多GPU机器上使用DeepSpeed、Hugging Face Transformers、LoRA和Flash Attention对Falcon 180B进行微调。

详细内容中，您将学习如何：

1. 设置开发环境
2. 加载并准备数据集
3. 使用 DeepSpeed、Hugging Face Transformers、LoRA 与 Flash Attention 对 Falcon 180B 进行微调

在我们深入代码之前，让我们快速了解一下我们将要使用的技术和方法：

什么是DeepSpeed ZeRO？

DeepSpeed ZeRO 专注于高效的大规模训练转换器。ZeRO，或零冗余优化器，通过在设备之间分割模型状态而不是基本的数据并行，来减少内存占用。这节省了大量内存 - ZeRO-Infinity 可以将数据并行ism 的使用量减少 100 倍。ZeRO-Offload 进一步通过将模型和优化器的部分转移到 CPU 来减少内存，使 1 GPU 上可以运行 10B+ 参数的模型。ZeRO通过配置文件与 HuggingFace Transformers 集成。

什么是LoRA？

LoRA 使大型语言模型的高效微调成为可能。它将权重矩阵分解为更小的、可训练的更新矩阵，这些矩阵在保持原始权重冻结的同时进行适应。这大大减少了可训练的参数，从而实现更快、更低内存的微调。LoRA 通过 Hugging Face 的 PEFT 集成到 Transformers 中。它与 DeepSpeed 等方法结合得很好。主要优点是高效微调、模型便携，并且在合并训练权重时没有推理延迟。LoRA 允许在有限的资源下训练大规模模型。

什么是Flash Attention？

Flash Attention 是一种通过重新结构化计算来加速 Transformer 语言模型中核心注意力机制的算法。它使用了平铺和重新计算等技术来降低注意力的高内存成本，使模型能够处理更长的文本序列。Flash Attention 2 通过优化并行性和工作分区，使性能提高到前一版本的 2 倍，在 A100 GPU 上达到 230 TFLOPS/s。

访问falcon-180B

在我们开始训练之前，我们必须确保已经接受了许可证tiiuae/falcon-180B才能使用它。您可以通过点击模型页面上的“同意并访问存储库”按钮来接受许可证：

• tiiuae/falcon-180B

该示例是在DGX A100 8-GPU机器上创建和运行的，每块GPU具有80GB GPU内存。

1. 设置开发环境

conda create --name hf python=3.10 -c conda-forge

# install torch with the correct cuda version, check nvcc --version
!pip install torch --extra-index-url https://download.pytorch.org/whl/cu118 --upgrade
# install Hugging Face Libraries and additional dependencies
!pip install "transformers==4.34.0" "datasets==2.14.5" "accelerate==0.22.0" "evaluate==0.4.0" "peft==0.5.0" tensorboard packaging --upgrade
# install deepspeed and ninja for jit compilations of kernels
!pip install "deepspeed==0.10.3" ninja --upgrade
# install additional Flash Attention
!pip install flash-attn --no-build-isolation --upgrade

2. 加载并准备数据集

我们将使用 dolly ，一个由数千名Databricks员工在 InstructGPT论文中概述的几个行为类别中生成的指令遵循记录的开源数据集，包括头脑风暴、分类、闭合问答、生成、信息提取、开放问答和总结。

{"instruction":"魔兽世界是什么","context":"","response":"魔兽世界是一款大型多人在线角色扮演游戏。它由奇异娱乐公司于2004年发布"
}   

加载 samsum 数据集，我们使用 load_dataset() 🤗 Datasets 库中的方法。

from datasets import load_dataset
from random import randrange# Load dataset from the hub
dataset = load_dataset("databricks/databricks-dolly-15k", split="train")print(f"dataset size: {len(dataset)}")
print(dataset[randrange(len(dataset))])
# dataset size: 15011

为了指导我们的模型进行调整，我们需要将结构化的示例转换为通过指令描述的任务集合。我们定义了一个formatting_function，它接受一个样本并返回一个符合我们格式说明的字符串。

def format_dolly(sample):instruction = f"### Instruction\n{sample['instruction']}"context = f"### Context\n{sample['context']}" if len(sample["context"]) > 0 else Noneresponse = f"### Answer\n{sample['response']}"# join all the parts togetherprompt = "\n\n".join([i for i in [instruction, context, response] if i is not None])return prompt

让我们用一个随机的例子来测试我们的格式化函数。

from random import randrange
print(format_dolly(dataset[randrange(len(dataset))]))

此外，为了格式化我们的样本，我们还希望将多个样本打包到一个序列中，以进行更有效的训练。

from transformers import AutoTokenizermodel_id = "tiiuae/falcon-180B" 
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token

我们定义了一些辅助函数，将我们的样本包装成指定长度的序列，然后对它们进行标记化。

from random import randint
from itertools import chain
from functools import partial# template dataset to add prompt to each sample
def template_dataset(sample):sample["text"] = f"{format_dolly(sample)}{tokenizer.eos_token}"return sample# apply prompt template per sample
dataset = dataset.map(template_dataset, remove_columns=list(dataset.features))
'''
dataset.map() 是 datasets 库中用于数据处理的核心方法，它会对数据集中的每个样本应用指定的函数（这里是 template_dataset）
remove_columns=list(dataset.features) 表示在处理完成后，删除原始数据集中的所有列
dataset.features 包含了数据集中所有特征（列）的信息
list(dataset.features) 会得到所有列名的列表
'''
# print random sample
print(dataset[randint(0, len(dataset))]["text"])# empty list to save remainder from batches to use in next batch
remainder = {"input_ids": [], "attention_mask": [], "token_type_ids": []}def chunk(sample, chunk_length=2048):'''分块函数chunk的核心逻辑：目的：将长文本按chunk_length=2048拆分，确保每个样本长度不超过模型最大输入长度；关键处理：用remainder保存上一批次未用完的 token（如一批次处理后剩余 50 个 token，下一批次会先拼接这 50 个 token，避免文本被截断破坏语义）；步骤：拼接当前批次的所有 token 序列（如input_ids），并加上remainder中保存的上一批次剩余 token；计算总长度，按2048拆分出完整的块（如总长度 3000，则拆分为 1 个 2048 长度的块，剩余 952 个 token 存入remainder）；生成labels列（自回归语言模型训练中，labels 与 input_ids 相同，因为目标是预测下一个 token）。'''# define global remainder variable to save remainder from batches to use in next batch# 声明全局变量remainder，用于存储前一批次的剩余文本片段, 这使得我们可以将前一批次的剩余文本与当前批次的文本连接起来global remainder# Concatenate all texts and add remainder from previous batch# 将批次中的所有文本连接成一个长列表，并添加前一批次的剩余文本 # chain(*sample[k]) 将每个键k对应的所有子列表连接成一个迭代器# list(...) 将迭代器转换为列表concatenated_examples = {k: list(chain(*sample[k])) for k in sample.keys()}'''concatenated_examples,{'input_ids':        ['i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i'], 'attention_mask':   ['a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a'], 'token_type_ids':   ['t', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't', 't']}'''# 将前一批次的剩余文本添加到当前批次的开头concatenated_examples = {k: remainder[k] + concatenated_examples[k] for k in concatenated_examples.keys()}# get total number of tokens for batchbatch_total_length = len(concatenated_examples[list(sample.keys())[0]])# # print(batch_total_length) 30# get max number of chunks for batchif batch_total_length >= chunk_length:batch_chunk_length = (batch_total_length // chunk_length) * chunk_length# Split by chunks of max_len.result = {k: [t[i : i + chunk_length] for i in range(0, batch_chunk_length, chunk_length)]for k, t in concatenated_examples.items()}'''result,{'input_ids':        [['i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i'], ['i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i', 'i']], 'attention_mask':   [['a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a'], ['a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a']], 'token_type_ids':   [['t', 't', 't', 't', 't', 't', 't', 't', 't', 't'], ['t', 't', 't', 't', 't', 't', 't', 't', 't', 't']]}'''# add remainder to global variable for next batchremainder = {k: concatenated_examples[k][batch_chunk_length:] for k in concatenated_examples.keys()}'''remainder,{'input_ids': ['i'], 'attention_mask': ['a'], 'token_type_ids': ['t']}'''# prepare labelsresult["labels"] = result["input_ids"].copy()return result# tokenize and chunk dataset
'''
执行分词与分块
第一次map：对dataset中的text列进行分词（tokenizer(sample["text"])），将文本转换为input_ids（token 的整数编码）、attention_mask（标识哪些 token 是有效文本，非填充）等；batched=True表示批量处理（效率更高），并删除原始的text列；
第二次map：用partial(chunk, chunk_length=2048)固定分块长度为 2048，对分词后的结果进行分块（调用chunk函数），batched=True表示按批次处理；
打印处理后的样本总数（分块后的总块数）；
将处理后的数据集lm_dataset保存到磁盘（dolly-processed文件夹），方便后续训练加载。
'''
lm_dataset = dataset.map(lambda sample: tokenizer(sample["text"]), batched=True, remove_columns=list(dataset.features)
).map(partial(chunk, chunk_length=2048),batched=True,
)# Print total number of samples
print(f"Total number of samples: {len(lm_dataset)}")

输出

### Instruction
Identify which instrument is string or percussion: Xylophone, Ramkie### Answer
Ramkie is string, Xylophone is percussion.<|im_end|>
Total number of samples: 1359

在我们处理完数据集后，我们需要将其保存到磁盘上，以便在训练过程中稍后使用处理后的数据集。

lm_dataset.save_to_disk("dolly-processed")

3. 使用DeepSpeed、Hugging Face Transformers和LoRA with Flash Attention微调Falcon 180B

DeepSpeed ZeRO 本地集成到 Hugging Face Transformers 训练器 中。这种集成使得通过提供一个 DeepSpeed 配置文件即可利用 ZeRO，并且训练器会处理其余部分。我们为运行的实验创建了 2 个 deepspeed 配置，包括 CPU offloading：

• ds_falcon_180b_z3.json
• ds_falcon_180b_z3_offload.json

正如开头所提到的，我们使用8个NVIDIA A100 80GB运行了这些示例。这意味着我们可以利用bf16，这将模型的内存足迹减少近2倍，使我们能够在不进行卸载的情况下进行训练。我们将使用ds_falcon_180b_z3.json。如果你对auto值感到恼火，请查看文档。

除了 deepspeed 配置，我们还需要一个训练脚本，该脚本实现了 LoRA 并修补我们的模型以使用 flash-attention。我们创建了一个run_ds_lora.py脚本，该脚本使用falcon_patch.py工具修补 falcon 模型，并使用peft_utils.py实现 LoRA。

运行make时，请确保你有相同的文件夹结构和utils/configs。最简单的方法是克隆整个仓库。进入training目录并开始训练。

一旦我们确保我们有正确的配置和训练脚本，我们就可以使用torchrun开始训练。