使用 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进行微调。

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

设置开发环境
加载并准备数据集
使用 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 None
    response = f"### Answer\n{sample['response']}"
    # join all the parts together
    prompt = "\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 AutoTokenizer

model_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 batch
    batch_total_length = len(concatenated_examples[list(sample.keys())[0]])
    # # print(batch_total_length) 30

    # get max number of chunks for batch
    if 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 batch
    remainder = {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 labels
    result["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：

正如开头所提到的，我们使用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开始训练。

!torchrun --nproc_per_node 8 run_ds_lora.py \
  --model_id tiiuae/falcon-180B \
  --dataset_path dolly-processed \
  --output_dir falcon-180b-lora-fa \
  --num_train_epochs 3 \
  --per_device_train_batch_size 1 \
  --learning_rate 4e-3 \
  --gradient_checkpointing True \
  --gradient_accumulation_steps 8 \
  --bf16 True \
  --tf32 True \
  --use_flash_attn True \
  --lr_scheduler_type "constant_with_warmup" \
  --logging_steps 25 \
  --save_steps 100 \
  --save_total_limit 3 \
  --deepspeed configs/ds_falcon_180b_z3.json

注意：由于我们使用了LoRA，我们只保存“训练好的”适配器权重，以节省一些存储空间。如果你想将适配器重新合并到基础模型并保存合并后的模型，可以添加--merge_adapters True或使用merge_adapter_weights.py脚本。

在我们的Falcon 180B示例中，训练时间是153 minutes或约2小时，进行了3个时期。相比之下，Falcon 180B的预训练成本约为700万GPU小时，是微调的350万倍。

结论

在博客文章中，您将学习如何在多GPU机器上使用DeepSpeed、Hugging Face Transformers和LoRA与Flash Attention微调Falcon 180B模型。我们使用了：

DeepSpeed ZeRO 用于内存优化，使在有限的 GPU 内存上训练具有数万亿参数的模型成为可能。我们使用了 stage 3 (ZeRO-Infinity) 来优化内存使用。
拥抱面的变换器和数据集，用于轻松加载和准备文本数据集，并提供直观的训练器API。
LoRA，一种通过在每次迭代中仅更新一小部分参数来高效微调大型语言模型的方法。这大大减少了内存使用和计算成本。
Flash Attention - 一种高度优化的注意力实现，进一步减少了内存占用。

将所有这些方法结合起来，使我们能够在有限的资源下微调拥有超过1000亿参数的LLM。这个例子提供了一个高效微调最大公开模型的模板。

posted @ 2025-07-29 11:13 有何m不可阅读(129) 评论(0) 收藏举报

刷新页面返回顶部

gongzb