使用Hugging Face和LoRA在单个Amazon SageMaker GPU上训练大型语言模型

本文与来自Hugging Face的Philipp Schmid共同撰写。

我们都听说过在大型语言模型(LLMs)领域取得的进展,以及LLMs在提供有价值的见解的问题集数量不断增长。大型模型在训练大规模的数据集和多个任务时,也能够很好地泛化到它们没有专门训练的任务上。这样的模型被称为基础模型,这个术语最初是由斯坦福人机中心人工智能研究所首创的。尽管这些基础模型能够很好地泛化,特别是在提示工程技术的帮助下,但往往用例是如此特定于领域,或者任务是如此不同,以至于模型需要进一步定制。一种提高大型模型特定领域或任务性能的方法是使用更小的任务特定数据集进一步对模型进行训练。虽然这种称为微调的方法成功地提高了LLMs的准确性,但它需要修改所有模型权重。微调比模型的预训练要快得多,因为数据集大小要小得多,但仍需要大量的计算能力和内存。微调修改了原始模型的所有参数权重,这使得它很昂贵,并且得到的模型与原始模型的大小相同。

为了解决这些挑战,Hugging Face推出了参数高效微调库(PEFT)。该库允许您冻结大多数原始模型权重,通过训练另一个更小的参数集来替换或扩展模型层。这使得训练在所需的计算和内存方面更加廉价。

在本文中,我们将向您展示如何使用Amazon SageMaker上的单个图形处理器(GPU)训练70亿参数的BloomZ模型,Amazon的机器学习(ML)平台,用于准备、构建、训练和部署高质量的ML模型。BloomZ是一个通用的自然语言处理(NLP)模型。我们使用PEFT优化此模型,使其适用于摘要类似于Messenger的对话的具体任务。我们使用的单个GPU实例是AWS提供的许多实例类型的低成本示例。在单个GPU上训练此模型突显了AWS致力于成为最具成本效益的AI / ML服务提供商的承诺。

此演练的代码可以在Hugging Face笔记本GitHub存储库的sagemaker / 24_train_bloom_peft_lora文件夹中找到。

先决条件

为了跟随本文,您应具备以下先决条件:

  • AWS帐户。
  • Amazon SageMaker Studio或SageMaker笔记本实例中的Jupyter笔记本。
  • 您将需要访问包含单个NVIDIA A10G GPU的SageMaker ml.g5.2xlarge实例类型。在AWS管理控制台上,导航到SageMaker的服务配额,并请求以下配额的1个实例增加:ml.g5.2xlarge用于培训作业使用 ml.g5.2xlarge用于端点使用
  • 在将请求的配额应用于您的帐户后,您可以使用默认的Studio Python 3(数据科学)映像和ml.t3.medium实例来运行笔记本代码片段。有关可用内核的完整列表,请参见Available Amazon SageMaker Kernels。

设置SageMaker会话

使用以下代码设置SageMaker会话:

import sagemaker
import boto3
sess = sagemaker.Session()
# sagemaker会话存储桶 ->用于上传数据、模型和日志
# 如果不存在,则sagemaker将自动创建此存储桶
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # 如果未给出存储桶名称,则设置为默认存储桶
    sagemaker_session_bucket = sess.default_bucket()

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

print(f"sagemaker角色arn: {role}")
print(f"sagemaker存储桶: {sess.default_bucket()}")
print(f"sagemaker会话区域: {sess.boto_region_name}")

加载并准备数据集

我们使用samsum数据集,这是一个包含16000个类似于Messenger的对话和摘要的集合。这些对话是由英语流利的语言学家创建和书写的。以下是数据集的示例:

{
  "id": "13818513",
  "summary": "Amanda 烤了饼干,明天会给 Jerry 一些。",
  "dialogue": "Amanda: 我烤了饼干,你想要吗?\r\nJerry: 当然!\r\nAmanda: 我明天给你 :-)"
}

为了训练模型,您需要将输入的文本转换为标记 ID。这是通过 Hugging Face Transformers tokenizer 完成的。有关更多信息,请参阅 Hugging Face NLP 课程第 6 章。

使用以下代码转换输入:

from transformers import AutoTokenizer

model_id="bigscience/bloomz-7b1"

# 加载 BLOOMZ 的 tokenizer
tokenized = AutoTokenizer.from_pretrained(model_id)
tokenizer.model_max_length = 2048 # 修正错误值

在开始训练之前,您需要处理数据。一旦训练完成,模型将以一组文本消息作为输入,并生成摘要作为输出。您需要将数据格式化为提示(消息)及其正确响应(摘要)。您还需要将示例分成更长的输入序列,以优化模型训练。见下面的代码:

from random import randint
from itertools import chain
from functools import partial

# 自定义提示开始
prompt_template = f"总结聊天对话:\n{{dialogue}}\n---\n摘要:\n{{summary}}{{eos_token}}"

# 模板数据集以将提示添加到每个样本中
def template_dataset(sample):
    sample["text"] = prompt_template.format(dialogue=sample["dialogue"],
                                            summary=sample["summary"],
                                            eos_token=tokenizer.eos_token)
    return sample


# 对每个样本应用提示模板
dataset = dataset.map(template_dataset, remove_columns=list(dataset.features))

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

# 空列表以保存批次中的剩余部分以在下一批次中使用
remainder = {"input_ids": [], "attention_mask": []}


def chunk(sample, chunk_length=2048):
    # 定义全局 remainder 变量以保存批次中的剩余部分以在下一批次中使用
    global remainder
    # 合并所有文本并添加先前批次的剩余部分
    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()}
    # 获取批次的总令牌数
    batch_total_length = len(concatenated_examples[list(sample.keys())[0]])

    # 获取批次的最大块数
    if batch_total_length >= chunk_length:
        batch_chunk_length = (batch_total_length // chunk_length) * chunk_length

    # 按 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()
    }
    # 将剩余部分添加到全局变量以供下一批次使用
    remainder = {k: concatenated_examples[k][batch_chunk_length:] for k in concatenated_examples.keys()}
    # 准备标签
    result["labels"] = result["input_ids"].copy()
    return result


# 对数据集进行标记化和分块处理
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(f"样本总数:{len(lm_dataset)}")

现在,您可以使用 FileSystem 集成将数据集上传到 Amazon Simple Storage Service (Amazon S3):

# 保存 train_dataset 到 s3
training_input_path = f's3://{sess.default_bucket()}/processed/samsum-sagemaker/train'
lm_dataset.save_to_disk(training_input_path)

print("uploaded data to:")
print(f"training dataset to: {training_input_path}")

In [ ]:
training_input_path="s3://sagemaker-us-east-1-558105141721/processed/samsum-sagemaker/train"

使用 LoRA 和 bitsandbytes int-8 在 SageMaker 上微调 BLOOMZ-7B

Hugging Face BLOOMZ-7B 模型卡片表示其初始训练是在 8 个节点上进行的,每个节点都有 8 个 A100 80 GB GPU 和 512 GB 内存 CPU。这种计算配置不容易获得,对消费者来说成本高昂,并需要分布式训练性能优化方面的专业知识。SageMaker 通过其分布式训练库降低了复制此设置的门槛;但是,相当于八个按需 ml.p4de.24xlarge 实例的成本每小时为 376.88 美元。此外,完全训练的模型占用约 40 GB 的内存,超过了许多消费者可用的个人 GPU 的内存,并需要针对大型模型推理采取策略。因此,为了多次模型运行和部署的全面微调,需要在不容易获得的硬件上支付显着的计算、内存和存储成本。

我们的目标是找到一种更加可访问和成本效益高的方法,使BLOOMZ-7B适应我们的聊天摘要用例,同时保持精度。为了让我们的模型能够在SageMaker ml.g5.2xlarge实例上进行微调,该实例仅配备一款消费级NVIDIA A10G GPU,我们采用了两种技术来降低微调的计算和内存需求:LoRA和量化。

LoRA(低秩适应)是一种技术,它显著减少了微调到新任务所需的模型参数和相关计算,而不会损失预测性能。首先,它冻结您的原始模型权重,而是优化较小的秩分解权重矩阵以适应您的新任务,而不是更新全部权重,然后将这些调整后的权重注入到原始模型中。因此,更少的权重梯度更新意味着在微调期间需要更少的计算和GPU内存。这种方法背后的直觉是,LoRA允许LLMs专注于最重要的输入和输出令牌,同时忽略冗余和不重要的令牌。要加深您对LoRA技术的理解,可以参考原始论文LoRA: Low-Rank Adaptation of Large Language Models。

除了LoRA技术外,您还使用了bitsanbytes Hugging Face集成LLM.int8()方法来量化冻结的BloomZ模型,即将权重和偏差值的精度从float16舍入为int8。量化将BloomZ的所需内存减少了约四倍,这使您可以在A10G GPU实例上安装模型,而不会显著损失预测性能。要深入了解int8量化的工作原理,它在bitsandbytes库中的实现以及它与Hugging Face Transformers库的集成,请参见A Gentle Introduction to 8-bit Matrix Multiplication for transformers at scale using Hugging Face Transformers, Accelerate and bitsandbytes。

Hugging Face通过PEFT库及其与bitsandbytes库的集成,使LoRA和量化可在广泛范围的Transformer模型中使用。准备脚本run_clm.py中的create_peft_config()函数说明了它们在准备模型进行训练时的使用:

def create_peft_config(model):
    from peft import (
        get_peft_model,
        LoraConfig,
        TaskType,
        prepare_model_for_int8_training,
    )

    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        inference_mode=False,
        r=8, # Lora attention dimension.
        lora_alpha=32, # the alpha parameter for Lora scaling.
        lora_dropout=0.05, # the dropout probability for Lora layers.
        target_modules=["query_key_value"],
    )

    # prepare int-8 model for training
    model = prepare_model_for_int8_training(model)
    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()
    return model

使用LoRA,print_trainable_parameters()的输出表明,我们能够将模型参数数量从70亿减少到390万。这意味着只有原始模型参数的5.6%需要进行更新。这种显著的计算和内存需求的降低使我们可以在GPU上安装和训练模型而不会出现问题。

要创建SageMaker训练作业,您需要一个Hugging Face估计器。估计器处理端到端的SageMaker训练和部署任务。SageMaker会为您启动和管理所有必需的Amazon Elastic Compute Cloud(Amazon EC2)实例。此外,它提供正确的Hugging Face培训容器,上传提供的脚本,并将数据从我们的S3存储桶下载到容器的路径/opt/ml/input/data。然后,它启动训练作业。请参见以下代码:

import time
# 定义训练作业名称 
job_name = f'huggingface-peft-{time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())}'

from sagemaker.huggingface import HuggingFace

# 超参数,这些超参数将传递给训练作业
hyperparameters ={
  'model_id': model_id,                                # 预训练模型
  'dataset_path': '/opt/ml/input/data/training', # sagemaker将保存训练数据集的路径
  'epochs': 3,                                         # 训练轮数
  'per_device_train_batch_size': 1,                    # 训练批次大小
  'lr': 2e-4,                                          # 训练期间使用的学习率
}

# 创建估计器
huggingface_estimator = HuggingFace(
    entry_point          = 'run_clm.py',      # 训练脚本
    source_dir           = 'scripts',         # 包含所有用于训练所需的文件的目录
    instance_type        = 'ml.g5.2xlarge', # 用于训练作业的实例类型
    instance_count       = 1,                 # 用于训练的实例数量
    base_job_name        = job_name,          # 训练作业名称
    role                 = role,              # 用于训练作业访问AWS资源(例如S3)的IAM角色
    volume_size          = 300,               # EBS卷的大小(以GB为单位)
    transformers_version = '4.26',            # 在训练作业中使用的transformers版本
    pytorch_version      = '1.13',            # 在训练作业中使用的pytorch版本
    py_version           = 'py39',            # 在训练作业中使用的python版本
    hyperparameters      =  hyperparameters
)

现在您可以使用.fit()方法开始训练作业,并传递S3路径到训练脚本:

# 使用我们上传的S3链接定义数据输入字典
data = {'training': training_input_path}

# 使用我们上传的数据集作为输入启动训练作业
huggingface_estimator.fit(data, wait=True)

使用LoRA和量化使得将BLOOMZ-7B微调到我们的任务变得经济高效。当使用SageMaker训练作业时,您只需为模型训练期间的GPU付费。在我们的示例中,SageMaker训练作业需要20,632秒,约为5.7小时。我们使用的ml.g5.2xlarge实例的按需使用费用为每小时1.515美元。因此,微调我们的BLOOMZ-7B模型的总成本仅为8.63美元。相比之下,在原始计算配置上进行模型的完全微调需要约600美元,或者每次训练运行高出6900%,假设Hugging Face模型卡片中概述的线性GPU缩放。实际上,这将进一步因您的训练策略、实例选择和实例定价而有所不同。

我们还可以通过使用SageMaker托管的Spot实例进一步降低培训成本。但是,由于Spot实例中断的可能性,总训练时间可能会增加。有关实例定价详细信息,请参见Amazon SageMaker价格。

将模型部署到SageMaker端点以进行推理

使用LoRA,您之前已将较小的一组权重适应于您的新任务。您需要一种方法将这些任务特定的权重与原始模型的预训练权重结合起来。在run_clm.py脚本中,PEFT库的merge_and_unload()方法负责将基本的BLOOMZ-7B模型与更新的适配器权重合并,以便在不引入与原始模型相比的任何推理延迟的情况下更容易部署它们。

在本节中,我们将介绍使用微调模型工件创建SageMaker模型并将其部署到SageMaker端点以进行推理的步骤。首先,您可以使用新的微调模型工件创建Hugging Face模型,以便将其部署到SageMaker端点进行推理。由于之前使用SageMaker Hugging Face estimator训练了模型,因此您可以立即部署模型。您可以将训练好的模型上传到S3存储桶,稍后再使用它们创建模型包。请参见以下代码:

from sagemaker.huggingface import HuggingFaceModel

# 1. 创建Hugging Face模型类
huggingface_model = HuggingFaceModel(
   model_data=huggingface_estimator.model_data,
   #model_data="s3://hf-sagemaker-inference/model.tar.gz",  # 更改为您的模型路径
   role=role, 
   transformers_version="4.26", 
   pytorch_version="1.13", 
   py_version="py39",
   model_server_workers=1
)

与任何SageMaker estimator一样,您可以使用Hugging Face estimator对象的deploy()方法部署模型,传递所需的实例数量和类型。在此示例中,我们使用与先前步骤中微调的模型在上面相同的G5实例类型,该实例类型配备了一张NVIDIA A10g GPU:

# 2. 将模型部署到SageMaker推理
predictor = huggingface_model.deploy(
   initial_instance_count=1,
   instance_type= "ml.g5.4xlarge"
)

可能需要5-10分钟才能使SageMaker端点上线并下载您的模型以准备接受推理请求。

当端点运行时,您可以通过发送数据集测试拆分中的示例对话来测试它。首先使用Hugging Face Datasets库加载测试拆分。接下来,为索引切片从数据集数组中选择一个随机整数以获取单个测试样例。使用字符串格式化,将测试样例与提示模板组合成结构化输入以指导我们模型的响应。此结构化输入然后可以与其他模型输入参数组合成格式化的样本JSON负载。最后,使用格式化的样本调用SageMaker端点,并打印模型的输出,总结样本对话。请参见以下代码:

from random import randint
from datasets import load_dataset

# 1. 从Hub加载数据集
test_dataset = load_dataset("samsum", split="test")

# 2. 选择一个随机测试样例
sample = test_dataset[randint(0,len(test_dataset))]

# 3. 格式化样本
prompt_template = f"Summarize the chat dialogue:\n{{dialogue}}\n---\nSummary:\n"

fomatted_sample = {
  "inputs": prompt_template.format(dialogue=sample["dialogue"]),
  "parameters": {
    "do_sample": True, # sample output predicted probabilities
    "top_p": 0.9, # sampling technique Fan et. al (2018)
    "temperature": 0.1, # increasing the likelihood of high probability words and decreasing the likelihood of low probability words
    "max_new_tokens": 100, # 
  }
}

# 4. 使用格式化的样本调用SageMaker端点
res = predictor.predict(fomatted_sample)


# 5. 打印模型输出
print(res[0]["generated_text"].split("Summary:")[-1])
# 示例模型输出:Kirsten and Alex are going bowling this Friday at 7 pm. They will meet up and then go together。

现在让我们将模型总结的对话输出与测试样本摘要进行比较:

print(sample["summary"])
# 样本模型输入:Kirsten提醒Alex,青年团将在本周五晚上7点去打保龄球。

清理

现在您已经测试了模型,请确保清理相关的SageMaker资源以避免持续收费:

predictor.delete_model()
predictor.delete_endpoint()

总结

在本文中,您使用了Hugging Face Transformer、PEFT和bitsandbytes库,利用SageMaker对BloomZ大型语言模型进行了微调,使用单个GPU花费了8美元,然后将模型部署到SageMaker端点,对测试样本进行推理。SageMaker提供多种使用Hugging Face模型的方式;有关更多示例,请查看AWS Samples GitHub。

要继续使用SageMaker对基础模型进行微调,请尝试在文章”在Amazon SageMaker上构建个性化生成AI SaaS应用”中尝试一些技术。我们还鼓励您通过探索JumpStart、Amazon Titan models和Amazon Bedrock来了解更多关于Amazon生成AI功能的信息。