在单个 AWS Trainium 实例上对 Llama 3 8B 进行监督微调

注意：本教程的完整脚本可以从这里下载。

本教程将教你如何使用 AWS Trainium 对开源 LLM（如 Llama 3）进行微调。在我们的示例中，我们将利用 Optimum Neuron、Transformers 和 Datasets 库。

您将学习如何

设置 AWS 环境
加载和处理数据集
使用 NeuronSFTTrainer 在 AWS Trainium 上对 Llama 进行监督微调
启动训练
评估和测试微调后的 Llama 模型

虽然我们将在本教程中使用 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

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 blizarre 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
from optimum.neuron.distributed import lazy_load_for_parallelism

# 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

with lazy_load_for_parallelism(tensor_parallel_size=tensor_parallel_size):
    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

这里的关键点是

我们使用 lazy_load_for_parallelism 上下文管理器来延迟加载模型。这不会在每个工作进程上加载完整的模型权重，而是只加载所需的权重（分片或完整）。**这在内存方面效率更高，并且通常是必须使用的。**
我们定义一个 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
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

NUM_EPOCHS=1
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=output-$SLURM_JOB_ID

if [ "$NEURON_EXTRACT_GRAPHS_ONLY" = "1" ]; then
    MAX_STEPS=$((LOGGING_STEPS + 5))
else
    MAX_STEPS=-1
fi


XLA_USE_BF16=1 neuron_parallel_compile 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 \
  --zero_1 false \
  --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

确保此预编译阶段运行大约 10 个训练步骤。这通常足以累积和编译在实际训练期间将需要的所有图。

注意：在没有缓存的情况下进行编译可能需要一段时间。它还会在 dolly_llama_sharded 中创建虚拟文件，在编译后您需要将其删除。我们还需要添加 MALLOC_ARENA_MAX=64 以限制 CPU 分配以避免潜在的崩溃，暂时不要将其删除。

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

实际训练

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

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

与编译命令的区别在于，我们从 max_steps=10 更改为 num_train_epochs=3。

使用以下命令启动训练。

#!/bin/bash
set -ex

export NEURON_FUSE_SOFTMAX=1
export NEURON_RT_ASYNC_EXEC_MAX_INFLIGHT_REQUESTS=3
export MALLOC_ARENA_MAX=64
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

NUM_EPOCHS=1
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=output-$SLURM_JOB_ID

if [ "$NEURON_EXTRACT_GRAPHS_ONLY" = "1" ]; then
    MAX_STEPS=$((LOGGING_STEPS + 5))
else
    MAX_STEPS=-1
fi


XLA_USE_BF16=1 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 \
  --zero_1 false \
  --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

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

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

合并检查点

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

optimum-cli neuron consolidate dolly_llama dolly_llama

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 分钟。幸运的是，您只需要运行一次。因为您可以在之后保存模型。如果您打算在 Inferentia2 上运行，则需要重新编译。编译是特定于参数和硬件的。

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

我们现在可以测试推理，但必须确保我们将输入格式化为我们用于微调的提示格式。因此，我们创建了一个辅助方法，该方法接受一个包含我们的 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):]

让我们测试推理。首先，我们在没有上下文的情况下进行测试。

注意：使用 2 个核心在 AWS Trainium 上进行推理预计不会非常快。对于推理，我们建议使用 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