使用 LoRA 微调 Llama 3 8B

注意：本教程的完整脚本可在此处下载。

本教程将教你如何在 AWS Trainium 上微调像 Llama 3 这样的开源 LLM。在我们的示例中，我们将利用 Optimum Neuron、Transformers 和 Datasets 库。

您将学习如何

在一台 AWS Trainium 实例上对 Llama 3 8B 进行监督式微调

虽然在本教程中我们将使用 Llama-3 8B，但完全可以使用其他模型，只需切换 model_id 即可。

1. 设置 AWS 环境

在开始本教程之前，你需要设置你的环境

创建一个 AWS Trainium 实例。你需要一个 trn1.32xlarge 实例，它包含 16 个 Neuron 设备。 你可以按照这篇指南来创建一个。
确保你已经登录到 Hugging Face Hub

huggingface-cli login --token YOUR_TOKEN

检查你是否可以访问该模型。一些开源模型是受限的，这意味着用户需要向模型所有者申请才能使用模型权重。这里我们将训练 Llama-3 8B，有两种可能性

官方受限仓库：meta-llama/Meta-Llama-3-8B
非官方不受限仓库：NousResearch/Meta-Llama-3-8B

克隆 Optimum Neuron 仓库，其中包含本教程中描述的完整脚本：

git clone https://github.com/huggingface/optimum-neuron.git

确保你安装了 training 附加包，以获取所有必要的依赖项

python -m pip install .[training]

2. 加载并准备数据集

在本教程中，我们将使用 Dolly，这是一个开源的指令遵循记录数据集，其类别在 InstructGPT 论文中有概述，包括头脑风暴、分类、封闭式问答、生成、信息提取、开放式问答和摘要。

示例

{
  "instruction": "What is world of warcraft",
  "context": "",
  "response": (
        "World of warcraft is a massive online multi player role playing game. "
        "It was released in 2004 by bizarre entertainment"
    )
}

我们可以使用 🤗 Datasets 库中的 load_dataset() 方法非常轻松地加载 dolly 数据集。

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

为了对我们的模型进行指令微调，我们需要

将我们的结构化示例转换为通过指令描述的任务集合
（可选）将多个示例打包到一个序列中以进行更高效的训练。换句话说，我们将多个示例堆叠成一个示例，并用 EOS 标记分隔它们。

我们可以手动完成这些操作，但我们将改用 NeuronSFTTrainer。

3. 使用 NeuronSFTTrainer 在 AWS Trainium 上对 Llama 进行监督式微调

通常，你会使用 SFTConfig 和 SFTTrainer 类来对基于 PyTorch 的 Transformer 模型进行监督式微调。

在这里，我们将改用 NeuronSFTConfig 和 NeuronSFTTrainer。这些类复制了 trl 库中的类，同时确保它们在 Neuron 核心上正常工作。

格式化我们的数据集

有多种方法可以将数据集提供给 NeuronSFTTrainer，其中一种是提供一个格式化函数。对于不打包示例的 dolly，它看起来如下

def format_dolly(examples):
    output_text = []
    for i in range(len(examples["instruction"])):
        instruction = f"### Instruction\n{examples['instruction'][i]}"
        context = f"### Context\n{examples['context'][i]}" if len(examples["context"][i]) > 0 else None
        response = f"### Answer\n{examples['response'][i]}"
        prompt = "\n\n".join([i for i in [instruction, context, response] if i is not None])
        output_text.append(prompt)
    return output_text

准备模型

由于 Llama-3 8B 是一个大模型，即使使用分布式训练，它也无法在一个 trn1.32xlarge 实例上容纳。要仅使用一个 Trainium 实例实际微调一个 8B 模型，我们需要同时使用 LoRA 和分布式训练。

如果你想了解更多关于分布式训练的信息，可以查看文档。

在这里，我们将结合 LoRA 使用张量并行。我们的训练代码将如下所示

from peft import LoraConfig
from optimum.neuron import NeuronSFTConfig, NeuronSFTTrainer

# Define the tensor_parallel_size
tensor_parallel_size = 2

dataset = load_dataset("databricks/databricks-dolly-15k", split="train")

model_id = "meta-llama/Meta-Llama-3-8B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(model_id)

config = LoraConfig(
    r=16,
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=[
        "q_proj",
        "gate_proj",
        "v_proj",
        "o_proj",
        "k_proj",
        "up_proj",
        "down_proj"
    ],
    bias="none",
    task_type="CAUSAL_LM",
)

# training_args is an instance of NeuronTrainingArguments
args = training_args.to_dict()
sft_config = NeuronSFTConfig(
    max_seq_length=1024,
    packing=False,
    **args,
)

trainer = NeuronSFTTrainer(
    args=sft_config,
    model=model,
    peft_config=config,
    tokenizer=tokenizer,
    train_dataset=dataset,
    formatting_func=format_dolly,
)

# Start training
trainer.train()

trainer.save_model()  # Saves the tokenizer too for easy upload

这里的关键点是

我们定义了一个 LoraConfig，它指定了哪些层应该有适配器，以及这些适配器的超参数。
我们从常规的 NeuronTrainingArguments 创建一个 NeuronSFTConfig。这里我们指定不打包我们的示例，并且最大序列长度应为 1024，这意味着每个示例都将被填充或截断到 1024 的长度。
我们使用 NeuronSFTTrainer 来进行训练。它将接收模型以及 lora_config、sft_config 和 format_dolly，并准备用于监督式微调的数据集和模型。

4. 启动训练

我们准备了一个名为 sft_lora_finetune_llm.py 的脚本，总结了本教程中提到的所有内容。

PyTorch Neuron 使用 torch_xla。它在训练循环执行期间延迟评估操作，这意味着它在后台构建一个符号图，并且只有在打印张量、将其传输到 CPU 或调用 xm.mark_step() 时，图才会在硬件上执行。在执行过程中，根据控制流可能会构建多个图，按顺序编译每个图可能需要时间。为了缓解这个问题，Neuron SDK 提供了 neuron_parallel_compile，这是一个执行快速试运行的工具，它会构建所有图并并行编译它们。这个步骤通常被称为预编译。

预编译

在 AWS Trainium 上训练模型时，我们首先需要用我们的训练参数来编译模型。

为了简化这一步，我们添加了一个模型缓存仓库，它允许我们使用来自 Hugging Face Hub 的预编译模型来跳过编译步骤。这很有用，因为它能比实际训练时更快地编译模型，因为编译可以并行化。但要注意：模型配置的任何更改都可能导致新的编译，这可能会导致一些缓存未命中。

要了解更多关于缓存系统以及如何创建自己的私有缓存仓库的信息，请查看这篇指南。

编译命令只需将你的脚本作为输入传递给 neuron_parallel_compile 工具即可

#!/bin/bash
set -ex

export NEURON_FUSE_SOFTMAX=1
export NEURON_RT_ASYNC_EXEC_MAX_INFLIGHT_REQUESTS=3
export MALLOC_ARENA_MAX=64 # limit the CPU allocation to avoid potential crashes
export NEURON_CC_FLAGS="--model-type=transformer --distribution-strategy=llm-training --enable-saturate-infinity --cache_dir=/home/ubuntu/cache_dir_neuron/"

PROCESSES_PER_NODE=8

TP_DEGREE=2
PP_DEGREE=1
BS=1
GRADIENT_ACCUMULATION_STEPS=8
LOGGING_STEPS=1
MODEL_NAME="meta-llama/Meta-Llama-3-8B"
OUTPUT_DIR=dolly_llama_output

if [ "$NEURON_EXTRACT_GRAPHS_ONLY" = "1" ]; then
    MAX_STEPS=10
    NUM_EPOCHS=1
else
    MAX_STEPS=-1
    NUM_EPOCHS=3
fi


torchrun --nproc_per_node $PROCESSES_PER_NODE docs/source/training_tutorials/sft_lora_finetune_llm.py \
  --model_id $MODEL_NAME \
  --num_train_epochs $NUM_EPOCHS \
  --do_train \
  --learning_rate 5e-5 \
  --warmup_ratio 0.03 \
  --max_steps $MAX_STEPS \
  --per_device_train_batch_size $BS \
  --per_device_eval_batch_size $BS \
  --gradient_accumulation_steps $GRADIENT_ACCUMULATION_STEPS \
  --gradient_checkpointing true \
  --bf16 \
  --tensor_parallel_size $TP_DEGREE \
  --pipeline_parallel_size $PP_DEGREE \
  --logging_steps $LOGGING_STEPS \
  --save_total_limit 1 \
  --output_dir $OUTPUT_DIR \
  --lr_scheduler_type "constant" \
  --overwrite_output_dir

为方便起见，我们将这个 shell 脚本保存到一个文件中，sft_lora_finetune_llm.sh。你现在可以将其传递给 neuron_parallel_compile 工具来触发编译

neuron_parallel_compile bash docs/source/training_tutorials/sft_lora_finetune_llm.sh

注意：在编译结束时，可能会出现一个 `FileNotFoundError` 消息。你可以安全地忽略它，因为一些编译缓存已经创建好了。

这个预编译阶段会运行 10 个训练步骤，以确保编译器已经编译了所有必要的图。这通常足以累积和编译实际训练期间需要的所有图。

注意：没有缓存的编译可能需要一段时间。它还会在编译期间在 `dolly_llama_output` 目录中创建一些虚拟文件，你之后需要将它们删除。

# remove dummy artifacts which are created by the precompilation command
rm -rf dolly_llama_output

实际训练

编译完成后，我们可以用类似的命令开始我们的实际训练，我们只需要移除对 `neuron_parallel_compile` 的使用。

我们将使用 `torchrun` 来启动我们的训练脚本。`torchrun` 是一个可以自动将 PyTorch 模型分布到多个加速器上的工具。我们可以将加速器的数量作为 `nproc_per_node` 参数与我们的超参数一起传递。

与编译命令的区别在于，我们更改了变量 `max_steps=10` 和 `num_train_epochs=3`。

使用与预编译步骤中相同的命令启动训练，但不带 `neuron_parallel_compile`

bash docs/source/training_tutorials/sft_lora_finetune_llm.sh

就是这样，我们成功地在 AWS Trainium 上训练了 Llama-3 8B！

但在我们分享和测试我们的模型之前，我们需要整合我们的模型。由于我们在训练期间使用了张量并行，我们保存了检查点的分片版本。我们现在需要将它们整合起来。

整合检查点并合并模型

Optimum CLI 提供了一种非常简单的方法，通过 `optimum neuron consolidate [sharded_checkpoint] [output_dir]` 命令来实现这一点。

optimum-cli neuron consolidate dolly_llama_output dolly_llama_output

这将创建一个 `adapter_model.safetensors` 文件，即我们在上一步中训练的 LoRA 适配器权重。我们现在可以重新加载模型并将其合并，以便可以加载它进行评估。

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel, PeftConfig

MODEL_NAME = 'meta-llama/Meta-Llama-3-8B'
ADAPTER_PATH = 'dolly_llama_output'
MERGED_MODEL_PATH = 'dolly_llama'

# Load base odel
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Load adapter configuration and model
adapter_config = PeftConfig.from_pretrained(ADAPTER_PATH)
finetuned_model = PeftModel.from_pretrained(model, ADAPTER_PATH, config=adapter_config)

print("Saving tokenizer")
tokenizer.save_pretrained(MERGED_MODEL_PATH)
print("Saving model")
finetuned_model = finetuned_model.merge_and_unload()
finetuned_model.save_pretrained(MERGED_MODEL_PATH)

这个步骤可能需要几分钟。我们现在有了一个目录，其中包含评估微调模型所需的所有文件。

5. 评估和测试微调后的 Llama 模型

与训练一样，为了能够在 AWS Trainium 或 AWS Inferentia2 上运行推理，我们需要编译我们的模型。在这种情况下，我们将使用我们的 Trainium 实例进行推理测试，但你可以切换到 Inferentia2（`inf2.24xlarge`）进行推理。

Optimum Neuron 实现了类似于 Transformers AutoModel 的类，以便于推理使用。我们将使用 `NeuronModelForCausalLM` 类来加载我们原始的 transformers 检查点，并将其转换为 neuron。

from optimum.neuron import NeuronModelForCausalLM
from transformers import AutoTokenizer

compiler_args = {"num_cores": 2, "auto_cast_type": 'fp16'}
input_shapes = {"batch_size": 1, "sequence_length": 2048}

tokenizer = AutoTokenizer.from_pretrained("dolly_llama")
model = NeuronModelForCausalLM.from_pretrained(
        "dolly_llama",
        export=True,
        **compiler_args,
        **input_shapes)

注意：推理编译可能需要长达 25 分钟。幸运的是，你只需要运行一次。与训练前完成的预编译步骤一样，如果你更改运行推理的硬件，例如从 Trainium 移动到 Inferentia2，你也需要运行此编译步骤。编译是参数和硬件特定的。

# COMMENT IN if you want to save the compiled model
# model.save_pretrained("compiled_dolly_llama_output")

我们现在可以测试推理了，但必须确保我们将输入格式化为我们用于微调的提示格式。因此，我们创建了一个辅助方法，它接受一个包含 `instruction` 和可选的 `context` 的 `dict`。

def format_dolly_inference(sample):
    instruction = f"### Instruction\n{sample['instruction']}"
    context = f"### Context\n{sample['context']}" if "context" in sample else None
    response = f"### Answer\n"
    prompt = "\n\n".join([i for i in [instruction, context, response] if i is not None])
    return prompt


def generate(sample):
    prompt = format_dolly_inference(sample)
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.9,
        top_k=50,
        top_p=0.9
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=False)[len(prompt):]

让我们测试一下推理。首先我们测试没有上下文的情况。

注意：在 AWS Trainium 上使用 2 个核心进行推理，速度预计不会非常快。对于推理，我们建议使用 Inferentia2。

prompt = {
  "instruction": "Can you tell me something about AWS?"
}
res = generate(prompt)

print(res)

AWS 代表 Amazon Web Services。AWS 是亚马逊提供的一套远程计算服务。其中最广泛使用的包括 Amazon Elastic Compute Cloud（Amazon EC2），它在云中提供可调整大小的计算容量；Amazon Simple Storage Service（Amazon S3），它是一种对象存储服务；以及 Amazon Elastic Block Store（Amazon EBS），它旨在为 AWS 实例提供高性能、持久的块存储卷。AWS 还提供其他服务，例如 AWS Identity and Access Management（IAM），一项使组织能够控制对其 AWS 资源访问的服务，以及 AWS Key Management Service（AWS KMS），它帮助客户创建和控制加密密钥的使用。

这看起来是正确的。现在，让我们添加一些上下文，例如在 RAG 应用中你会做的那样

prompt = {
  "instruction": "How can I train models on AWS Trainium?",
  "context": "🤗 Optimum Neuron is the interface between the 🤗 Transformers library and AWS Accelerators including [AWS Trainium](https://aws.amazon.com/machine-learning/trainium/?nc1=h_ls) and [AWS Inferentia](https://aws.amazon.com/machine-learning/inferentia/?nc1=h_ls). It provides a set of tools enabling easy model loading, training and inference on single- and multi-Accelerator settings for different downstream tasks."
}
res = generate(prompt)

print(res)

你可以使用 Optimum Neuron 接口在 AWS Trainium 上训练模型。

太棒了，我们的模型也正确地使用了提供的上下文。我们完成了。恭喜你在 AWS Trainium 上微调了 Llama。

AWS Trainium 和 Inferentia