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

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

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

您将学习如何：

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

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

1. 设置 AWS 环境

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

创建一个 AWS Trainium 实例。您将需要一个包含 16 个 Neuron 设备的 trn1.32xlarge 实例。 您可以按照此指南创建一个。
确保您已登录 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 论文中概述的类别上的指令遵循记录，包括头脑风暴、分类、封闭式 QA、生成、信息提取、开放式 QA 和摘要。

示例

{
  "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. 在 AWS Trainium 上使用 NeuronSFTTrainer 对 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 上下文管理器来延迟加载模型。这不会在每个 worker 上加载完整的模型权重，而是仅加载所需的权重（分片或完整）。这大大提高了内存效率，并且通常是强制使用的。
我们定义了一个 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 类来加载我们的 vanilla 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):]

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

注意：预计在使用 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