在 AWS Trainium 上微调和测试 Llama-3 8B

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

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

您将了解如何

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

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

Mistral 模型，如 Mistral 7b (mistralai/Mistral-7B-Instruct-v0.3)
Llama-2 模型，如 Llama-2 7b (meta-llama/Llama-2-7b-hf)

以及许多其他模型！

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

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

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

除了格式化我们的样本外，我们还希望将多个样本打包成一个序列，以实现更有效的训练。换句话说，我们将多个样本堆叠到一个序列中，并用 EOS Token 将它们分开。打包/堆叠样本可以在训练过程中完成，也可以在训练之前完成。

以下 pack_dataset 函数接受一个 dataset 和一个 chunk_length，并返回一个打包后的数据集

from functools import partial
from itertools import chain

# empty list to save remainder from batches to use in next batch
remainder = {"input_ids": [], "attention_mask": [], "token_type_ids": []}

def pack_dataset(dataset, chunk_length=2048):
    print(f"Chunking dataset into chunks of {chunk_length} tokens.")

    def chunk(sample, chunk_length=chunk_length):
        # define global remainder variable to save remainder from batches to use in next batch
        global remainder
        # Concatenate all texts and add remainder from previous batch
        concatenated_examples = {k: list(chain(*sample[k])) for k in sample.keys()}
        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]])

        # 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()
        }
        # add remainder to global variable for next batch
        remainder = {k: concatenated_examples[k][batch_chunk_length:] for k in concatenated_examples.keys()}
        # prepare labels
        result["labels"] = result["input_ids"].copy()
        return result

    # tokenize and chunk dataset
    lm_dataset = dataset.map(
        partial(chunk, chunk_length=chunk_length),
        batched=True,
    )
    print(f"Total number of samples: {len(lm_dataset)}")
    return lm_dataset

为了总结，准备我们的数据集，我们将

使用模板方法格式化我们的样本，并在每个样本末尾添加 EOS 标记
标记化我们的数据集，将其从文本转换为标记
将我们的数据集打包成 2048 个标记

from transformers import AutoTokenizer
from random import randint

# Hugging Face Hub model id 
# model_id = "meta-llama/Meta-Llama-3-8B" # gated
model_id = "NousResearch/Meta-Llama-3-8B" # ungated

tokenizer = AutoTokenizer.from_pretrained(model_id)

# 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))

# print random sample
print(dataset[randint(0, len(dataset))]["text"])

# tokenize dataset
dataset = dataset.map(
    lambda sample: tokenizer(sample["text"]), batched=True, remove_columns=list(dataset.features)
)

# chunk dataset
lm_dataset = pack_dataset(dataset, chunk_length=2048) # We use 2048 as the maximum length for packing

3. 使用 NeuronTrainer 在 AWS Trainium 上微调 Llama

通常，您会使用 Trainer 和 TrainingArguments 类来微调基于 PyTorch 的转换器模型。

但是，我们与 AWS 合作开发了 [~optimum.neuron.NeuronTrainer]，以便在 Trainium 实例上训练时提高性能、稳健性和易用性。它可以作为 Trainer 的 1 对 1 替换。

由于 Llama-3 8B 是一个大型模型，它无法完全适应单个 Neuron 内核，因此我们需要分布式训练。在 Optimum Neuron 中，我们支持以下几种方式：

ZeRO-1: 它是一种数据并行性的优化，它将优化器状态（通常占设备所需内存的一半或更多）在数据并行秩之间进行分片。
张量并行: 它是一种将每个模型矩阵乘法沿着给定轴（行或列）在多个设备上进行分片的技术。它也被称为层内模型并行。用于分片参数的设备数量称为 tensor_parallel_size。
序列并行: 它是在张量并行之上进行的优化，它将激活在张量并行区域之外的序列轴上进行分片。它很有用，因为它通过对激活进行分片来节省内存。
管道并行: 它将模型块层在多个设备上进行分片。它也被称为层间模型并行。用于分片层的设备数量称为 pipeline_parallel_size。

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

这里，由于我们要微调一个 8B 模型，因此不需要使用管道并行。我们的训练代码如下所示：

from optimum.neuron import NeuronTrainer as Trainer
from optimum.neuron.distributed import lazy_load_for_parallelism

# Define the tensor_parallel_size
tensor_parallel_size = 8

# Load model from the Hugging face Hub 
with lazy_load_for_parallelism(tensor_parallel_size=tensor_parallel_size):
    model = AutoModelForCausalLM.from_pretrained(model_id)

trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=dataset,
    data_collator=default_data_collator,  # no special collator needed since we stacked the dataset
)

# Start training
trainer.train()

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

这里需要注意的要点是：

我们使用 lazy_load_for_parallelism 上下文管理器来延迟加载模型。这样不会在每个 worker 上加载完整的模型权重，而是只加载所需的权重（分片或完整）。这在内存效率方面要高得多，通常是必须使用的。
我们使用 [~optimum.neuron.NeuronTrainer] 来执行训练。它将接收延迟加载的模型以及 training_args（它是 [~optimum.neuron.NeuronTrainingArguments] 的实例），并将处理 Neuron 内核上的所有并行化和训练。

4. 启动训练

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

该脚本是我们的官方示例训练脚本的简化版本，用于运行因果语言建模微调，名为 run_clm.py。为了本教程的目的，我们尝试删除了所有不必要的内容，并添加了微调所需的格式化步骤，但是如果您想做更多自定义操作，可能 run_clm.py 中已经实现了解决方案！

此外，这些脚本更像是模板而不是最终脚本。您可以随意使用 finetune_llm.py 或 run_clm.py，并根据自己的需要进行调整！

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

预编译

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

为了克服这个问题，我们添加了一个模型缓存库，它允许我们使用 Hugging Face Hub 中的预编译模型来跳过编译步骤。但请注意：模型配置的任何更改都可能导致新的编译，这会导致一些缓存丢失。

注意：如果您的模型配置没有被缓存，请在 Github 上提交问题，我们很乐意将其包含在内。

编译命令只需将您的脚本作为输入调用 neuron_parallel_compile 工具即可

MALLOC_ARENA_MAX=64 XLA_USE_BF16=1 neuron_parallel_compile torchrun --nproc_per_node=32 finetune_llm.py \
 --model_id meta-llama/Meta-Llama-3-8B \
 --bf16 True \
 --learning_rate 5e-5 \
 --output_dir dolly_llama \
 --overwrite_output_dir True \
 --per_device_train_batch_size 1 \
 --gradient_accumulation_steps 16 \
 --gradient_checkpointing True \
 --tensor_parallel_size 8 \
 --max_steps 10 \
 --logging_steps 10

请确保在约 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。

使用以下命令启动训练。

MALLOC_ARENA_MAX=64 XLA_USE_BF16=1 torchrun --nproc_per_node=32 finetune_llm.py \
 --model_id meta-llama/Meta-Llama-3-8B \
 --bf16 True \
 --learning_rate 5e-5 \
 --output_dir dolly_llama \
 --overwrite_output_dir True \
 --skip_cache_push True \
 --per_device_train_batch_size 1 \
 --gradient_accumulation_steps 16 \
 --gradient_checkpointing True \
 --tensor_parallel_size 8 \
 --num_train_epochs 3 \
 --logging_steps 10

就是这样，我们成功地在 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 进行推理。

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