Press "Enter" to skip to content

使用QLoRA在Amazon SageMaker Studio笔记本上交互式地微调Falcon-40B和其他LLMs

微调大型语言模型(LLMs)可以让您调整开源基础模型,以在特定领域的任务上实现更好的性能。在本文中,我们讨论了使用Amazon SageMaker笔记本进行微调最先进的开源模型的优势。我们利用了Hugging Face的参数高效微调(PEFT)库和通过bitsandbytes支持使用单个笔记本实例进行交互微调极大模型的量化技术。具体来说,我们展示了如何使用单个ml.g5.12xlarge实例(4个A10G GPU)来微调Falcon-40B,但是同样的策略也适用于在p4d/p4de笔记本实例上微调更大的模型。

通常,这些非常大模型的完整精度表示无法在单个甚至多个GPU的内存中容纳。为了支持在这种大小的模型上进行交互式微调和推断的笔记本环境,我们使用了一种称为量化LLMs与低秩适配器(QLoRA)的新技术。 QLoRA是一种高效的微调方法,可以减少LLMs的内存使用量,同时保持较好的性能。 Hugging Face和所提到的论文的作者已经发布了一篇详细的博客文章,介绍了与Transformers和PEFT库的基本知识和集成。

使用笔记本进行微调LLMs

SageMaker提供两种选项来启动完全托管的笔记本,用于探索数据和构建机器学习(ML)模型。第一种选项是快速启动,即在Amazon SageMaker Studio中可访问的协作笔记本,这是一个完全集成的ML开发环境(IDE)。您可以快速启动SageMaker Studio中的笔记本,根据需要调整底层计算资源的规模,而不会中断您的工作,甚至可以实时共同编辑和协作笔记本。除了创建笔记本,您还可以在SageMaker Studio的单个视图中执行所有ML开发步骤,包括构建、训练、调试、跟踪、部署和监控模型。第二种选项是SageMaker笔记本实例,它是在云中运行笔记本的单个完全托管的ML计算实例,可提供更多对笔记本配置的控制。

在本文的其余部分,我们使用SageMaker Studio笔记本,因为我们希望利用SageMaker Studio的托管TensorBoard实验跟踪功能,该功能与Hugging Face Transformer对TensorBoard的支持兼容。但是,通过conda_pytorch_p310内核在笔记本实例上使用相同的概念示例代码也可以工作。值得注意的是,SageMaker Studio的Amazon Elastic File System(Amazon EFS)卷意味着您无需预先分配Amazon Elastic Block Store(Amazon EBS)卷大小,这在LLMs的模型权重非常大时很有用。

使用支持大型GPU实例的笔记本可以实现快速原型设计和调试,而无需启动冷启动容器。但是,这也意味着在使用完毕后需要关闭笔记本实例以避免额外的成本。其他选项,如Amazon SageMaker JumpStart和SageMaker Hugging Face容器,可用于微调,我们建议您参考以下有关上述方法的文章,以选择最适合您和您的团队的选项:

  • 在Amazon SageMaker JumpStart上对金融数据进行基础模型领域自适应微调
  • 使用Hugging Face和LoRA在单个Amazon SageMaker GPU上训练大型语言模型

先决条件

如果您是第一次使用SageMaker Studio,请首先创建一个SageMaker域。我们还使用托管的TensorBoard实例进行实验跟踪,尽管这对本教程来说是可选的。

此外,您可能需要请求相应的SageMaker Studio KernelGateway应用程序的服务配额增加。对于Falcon-40B的微调,我们使用ml.g5.12xlarge实例。

要请求服务配额增加,请在AWS服务配额控制台上导航到AWS服务Amazon SageMaker,然后选择运行在ml.g5.12xlarge实例上的Studio KernelGateway应用程序使用QLoRA在Amazon SageMaker Studio笔记本上交互式地微调Falcon-40B和其他LLMs 机器学习 第1张

开始

本文的代码示例可以在以下GitHub存储库中找到。为了开始,我们选择来自SageMaker Studio的Data Science 3.0镜像和Python 3内核,以便我们拥有一个最新的Python 3.10环境来安装我们的软件包。

使用QLoRA在Amazon SageMaker Studio笔记本上交互式地微调Falcon-40B和其他LLMs 机器学习 第2张

我们安装PyTorch和所需的Hugging Face和bitsandbytes库:

%pip install -q -U torch==2.0.1 bitsandbytes==0.39.1
%pip install -q -U datasets py7zr einops tensorboardX
%pip install -q -U git+https://github.com/huggingface/transformers.git@850cf4af0ce281d2c3e7ebfc12e0bc24a9c40714
%pip install -q -U git+https://github.com/huggingface/peft.git@e2b8e3260d3eeb736edf21a2424e89fe3ecf429d
%pip install -q -U git+https://github.com/huggingface/accelerate.git@b76409ba05e6fa7dfc59d50eee1734672126fdba

接下来,我们使用已安装的CUDA设置CUDA环境路径,这是PyTorch安装的依赖项。这是bitsandbytes库正确找到并加载正确的CUDA共享对象二进制文件的必要步骤。

# 将安装的cuda运行时添加到bitsandbytes的路径中
import os
import nvidia

cuda_install_dir = '/'.join(nvidia.__file__.split('/')[:-1]) + '/cuda_runtime/lib/'
os.environ['LD_LIBRARY_PATH'] =  cuda_install_dir

加载预训练的基础模型

我们使用bitsandbytes将Falcon-40B模型量化为4位精度,以便我们可以使用Hugging Face Accelerate的naive pipeline parallelism将模型加载到4个A10G GPU的内存中。如先前提到的Hugging Face文章所描述的,QLoRA调整在各种实验中显示出与16位微调方法相匹配的效果,因为模型权重存储为4位NormalFloat,但根据需要在前向和后向传递中解量化为计算bfloat16。

model_id = "tiiuae/falcon-40b"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

在加载预训练权重时,我们指定device_map="auto",以便Hugging Face Accelerate会自动确定将模型的每一层放在哪个GPU上。这个过程称为模型并行。

# Falcon要求您允许远程代码执行。这是因为该模型使用了一个尚未包含在transformers中的新架构。
# 代码在模型作者的repo中提供。
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, quantization_config=bnb_config, device_map="auto")

使用Hugging Face的PEFT库,您可以冻结大部分原始模型权重,并通过训练一组额外的、远小于原始模型的参数来替换或扩展模型的层。这使得训练所需的计算量大大减少。我们在LoRA配置中将要微调的Falcon模块设置为target_modules

from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=[
        "query_key_value",
        "dense",
        "dense_h_to_4h",
        "dense_4h_to_h",
    ],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, config)
print_trainable_parameters(model)
# 输出:可训练参数: 55541760 || 全部参数: 20974518272|| 可训练百分比: 0.2648058910327664

注意,我们只微调了模型参数的0.26%,这在合理的时间范围内是可行的。

加载数据集

我们在微调中使用samsum数据集。Samsum是一个包含16000个类似于消息的对话和带有标签摘要的集合。以下是数据集的示例:

{
    "id": "13818513",
    "summary": "阿曼达烤了饼干,明天会给杰瑞带一些。",
    "dialogue": "阿曼达:我烤了饼干。你要吃吗?\r\n杰瑞:当然!\r\n阿曼达:我明天会给你带来的 :-)"
}

在实践中,您希望使用具有与您希望调整模型的任务相关的特定知识的数据集。使用Amazon SageMaker Ground Truth Plus可以加速构建此类数据集的过程,详细信息请参见Amazon SageMaker Ground Truth Plus中的高质量人工反馈用于生成式AI应用。

微调模型

在进行微调之前,我们定义要使用的超参数并训练模型。我们还可以通过定义参数logging_dir并要求Hugging Face transformerreport_to="tensorboard"来将指标记录到TensorBoard中:

bucket = ”<YOUR-S3-BUCKET>”
log_bucket = f"s3://{bucket}/falcon-40b-qlora-finetune"

import transformers

# 我们将num_train_epochs=1仅用于演示

trainer = transformers.Trainer(
    model=model,
    train_dataset=lm_train_dataset,
    eval_dataset=lm_test_dataset,
    args=transformers.TrainingArguments(
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        logging_dir=log_bucket,
        logging_steps=2,
        num_train_epochs=1,
        learning_rate=2e-4,
        bf16=True,
        save_strategy = "no",
        output_dir="outputs",
         report_to="tensorboard",
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

监控微调过程

通过以上设置,我们可以实时监控微调过程。要实时监控GPU使用情况,我们可以直接从内核的容器中运行nvidia-smi。要在图像容器上启动一个终端,只需选择笔记本顶部的终端图标。

使用QLoRA在Amazon SageMaker Studio笔记本上交互式地微调Falcon-40B和其他LLMs 机器学习 第3张

从这里,我们可以使用Linux的watch命令每半秒钟运行一次nvidia-smi

watch -n 0.5 nvidia-smi

使用QLoRA在Amazon SageMaker Studio笔记本上交互式地微调Falcon-40B和其他LLMs 机器学习 第4张

在上述动画中,我们可以看到模型权重分布在4个GPU上,并且计算在层被串行处理时在它们之间分配。

要监控训练指标,我们利用我们写入指定的Amazon Simple Storage Service(Amazon S3)存储桶的TensorBoard日志。我们可以从AWS SageMaker控制台启动我们的SageMaker Studio域用户的TensorBoard:

使用QLoRA在Amazon SageMaker Studio笔记本上交互式地微调Falcon-40B和其他LLMs 机器学习 第5张

加载后,您可以指定您指示Hugging Face transformer记录以查看训练和评估指标的S3存储桶。

使用QLoRA在Amazon SageMaker Studio笔记本上交互式地微调Falcon-40B和其他LLMs 机器学习 第6张

使用QLoRA在Amazon SageMaker Studio笔记本上交互式地微调Falcon-40B和其他LLMs 机器学习 第7张

评估模型

在我们的模型训练完成后,我们可以进行系统评估或者生成回应:

tokens_for_summary = 30
output_tokens = input_ids.shape[1] + tokens_for_summary

outputs = model.generate(inputs=input_ids, do_sample=True, max_length=output_tokens)
gen_text = tokenizer.batch_decode(outputs)[0]
print(gen_text)
# 示例输出:
# 总结对话内容:
# Richie:Pogba
# Clay:Pogboom
# Richie:真是一脚漂亮的进球!
# Clay:当他将球传回自己的右脚时,我就离开座位了
# Richie:我也是,伙计
# Clay:希望他的状态能持续下去
# Richie:这个赛季他更加成熟了
# Clay:是的,约瑟有信任他
# Richie:每个人都有信任他
# Clay:是的,他在前60分钟后真的应该得分
# Richie:回报
# Clay:是的,伙计
# Richie:那就酷
# Clay:酷
# ---
# 总结:
# Richie和Clay讨论了Paul Pogba的进球。他本赛季的表现有所提高,他们希望这种状态能够持续下去。

当您对模型的性能感到满意后,您可以保存模型:

trainer.save_model("保存路径")

您还可以选择在专用的SageMaker端点上托管它。

清理

按照以下步骤清理资源:

  1. 关闭SageMaker Studio实例,以避免产生额外的费用。

  2. 关闭TensorBoard应用程序。

  3. 通过清除Hugging Face缓存目录来清理EFS目录:

    rm -R ~/.cache/huggingface/hub

结论

SageMaker笔记本允许您在交互式环境中快速高效地微调LLMs。在本文中,我们展示了如何使用Hugging Face PEFT和bitsandbytes来使用QLoRA在SageMaker Studio笔记本上微调Falcon-40B模型。试试看,并在评论区告诉我们您的想法!

我们还鼓励您通过探索SageMaker JumpStart、Amazon Titan模型和Amazon Bedrock来了解更多关于Amazon生成式AI能力的信息。

Leave a Reply

Your email address will not be published. Required fields are marked *