Press "Enter" to skip to content

使用@remote装饰器在Amazon SageMaker上微调Falcon 7B和其他LLMs

今天,生成式AI模型涵盖了各种任务,包括文本摘要、问答和图像与视频生成。为了提高输出质量,使用了一些方法,如n-short学习、Prompt工程、检索增强生成(RAG)和微调。微调允许您调整这些生成式AI模型,以在特定领域的任务上实现改进的性能。

借助Amazon SageMaker,您现在可以通过在Python代码中使用@remote装饰器来运行SageMaker训练作业。SageMaker Python SDK会自动将您现有的工作空间环境以及任何相关的数据处理代码和数据集转换成在训练平台上运行的SageMaker训练作业。这样做的优点是以更自然、面向对象的方式编写代码,同时利用SageMaker的能力在远程集群上运行训练作业,几乎不需要进行任何修改。

在本文中,我们展示了如何使用SageMaker Python SDK的@remote装饰器来微调Falcon-7B基础模型(FM)。它还使用了Hugging Face的参数高效微调(PEFT)库和通过bitsandbytes进行量化技术来支持微调。本博客中呈现的代码也可用于微调其他FM,例如Llama-2 13b。

这个模型的全精度表示可能在单个甚至多个图形处理单元(GPUs)上无法适应内存,甚至可能需要更大的实例。因此,为了在不增加成本的情况下进行微调,我们使用了一种称为带低秩适配器的量化LLM(QLoRA)的技术。QLoRA是一种高效的微调方法,它减少了LLM的内存使用,同时保持了非常好的性能。

使用@remote装饰器的优势

在进一步深入之前,让我们了解一下在使用SageMaker时,远程装饰器如何提高开发者的生产力:

  • @remote装饰器直接使用本机Python代码触发训练作业,无需显式调用SageMaker Estimators和SageMaker输入通道
  • 降低了开发者在SageMaker上训练模型的门槛。
  • 无需切换集成开发环境(IDE)。继续在您选择的IDE中编写代码并调用SageMaker训练作业。
  • 无需了解容器。继续在requirements.txt中提供依赖项,并将其提供给远程装饰器。

先决条件

需要一个具有AWS Identity and Access Management(AWS IAM)角色权限的AWS帐户来管理解决方案的资源。有关详细信息,请参阅创建AWS帐户。

在本文中,我们使用具有Data Science 3.0镜像和ml.t3.medium快速启动实例的Amazon SageMaker Studio。但是,您可以使用任何集成开发环境(IDE)。您只需正确设置AWS命令行接口(AWS CLI)凭据即可。有关更多信息,请参阅配置AWS CLI。

对于微调,本文中使用Falcon-7B的ml.g5.12xlarge实例。请确保AWS帐户中有足够的容量来使用此实例。

您需要克隆此Github存储库,以复制本文中演示的解决方案。

解决方案概述

  1. 安装预备条件以微调Falcon-7B模型
  2. 配置远程装饰器
  3. 预处理包含AWS服务FAQ的数据集
  4. 在AWS服务FAQ上微调Falcon-7B
  5. 在与AWS服务相关的示例问题上测试微调模型

1. 安装预备条件以微调Falcon-7B模型

通过选择ImageData ScienceKernelPython 3,在SageMaker Studio中启动笔记本falcon-7b-qlora-remote-decorator_qa.ipynb。安装requirements.txt中提到的所有所需库。其中一些库需要在笔记本实例本身上安装。执行其他需要的数据集处理和触发SageMaker训练作业的操作。

%pip install -r requirements.txt

%pip install -q -U transformers==4.31.0
%pip install -q -U datasets==2.13.1
%pip install -q -U peft==0.4.0
%pip install -q -U accelerate==0.21.0
%pip install -q -U bitsandbytes==0.40.2
%pip install -q -U boto3
%pip install -q -U sagemaker==2.154.0
%pip install -q -U scikit-learn

2. 设置远程装饰器配置

创建一个配置文件,其中指定了与Amazon SageMaker训练作业相关的所有配置。在运行训练作业时,@remote装饰器会读取此文件。此文件包含依赖项、训练镜像、实例和用于训练作业的执行角色等设置。有关配置文件支持的所有设置的详细参考,请查看使用SageMaker Python SDK配置和使用默认值。

SchemaVersion: '1.0'
SageMaker:
  PythonSDK:
    Modules:
      RemoteFunction:
        Dependencies: ./requirements.txt
        ImageUri: '{aws_account_id}.dkr.ecr.{region}.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04'
        InstanceType: ml.g5.12xlarge
        RoleArn: arn:aws:iam::111122223333:role/ExampleSageMakerRole

为了使用@remote装饰器,不一定要使用config.yaml文件。这只是一种更清晰的方式,将所有配置提供给@remote装饰器。这样一来,SageMaker和AWS相关的参数就可以与代码分离,并且只需要设置一次配置文件,就可以在团队成员之间共享。所有配置也可以直接在装饰器参数中提供,但这会降低代码的可读性和可维护性。此外,配置文件可以由管理员创建并与环境中的所有用户共享。

3. 预处理包含AWS服务FAQ的数据集

下一步是加载和预处理数据集,使其准备好进行训练作业。首先,让我们来看一下数据集:

使用@remote装饰器在Amazon SageMaker上微调Falcon 7B和其他LLMs 四海 第1张

它显示了一个AWS服务的常见问题解答。除了QLoRA之外,还使用bitsanbytes将其转换为4位精度以对冻结的LLM进行量化,并在其上附加LoRA适配器。

创建一个提示模板,将每个FAQ样本转换为提示格式:

from random import randint

# 自定义提示起始符
prompt_template = f"{{question}}\n---\nAnswer:\n{{answer}}{{eos_token}}"

# 使用模板将提示添加到每个样本中
def template_dataset(sample):
    sample["text"] = prompt_template.format(question=sample["question"],
                                            answer=sample["answers"],
                                            eos_token=tokenizer.eos_token)
    return sample

下一步是将输入(文本)转换为令牌ID。这是通过使用Hugging Face Transformers Tokenizer完成的。

from transformers import AutoTokenizer

model_id = "tiiuae/falcon-7b"

tokenizer = AutoTokenizer.from_pretrained(model_id)
# 设置Falcon tokenizer
tokenizer.pad_token = tokenizer.eos_token

现在只需使用prompt_template函数将所有FAQ转换为提示格式,并设置训练和测试数据集。

使用@remote装饰器在Amazon SageMaker上微调Falcon 7B和其他LLMs 四海 第2张

4. 在AWS服务FAQ上对Falcon-7B进行微调

现在,您可以准备训练脚本并定义训练函数train_fn,然后在函数上放置@remote装饰器。

训练函数执行以下操作:

  • 对数据集进行标记化和分块
  • 设置BitsAndBytesConfig,指定模型应以4位加载,但计算时应转换为bfloat16
  • 加载模型
  • 使用实用方法find_all_linear_names找到目标模块并更新必要的矩阵
  • 创建LoRA配置,指定更新矩阵的排序(s),缩放因子(lora_alpha),应用LoRA更新矩阵的模块(target_modules),Lora层的丢弃概率(lora_dropout),任务类型(task_type)等。
  • 开始训练和评估
import bitsandbytes as bnb

def find_all_linear_names(hf_model):
    lora_module_names = set()
    for name, module in hf_model.named_modules():
        if isinstance(module, bnb.nn.Linear4bit):
            names = name.split(".")
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])

    if "lm_head" in lora_module_names:
        lora_module_names.remove("lm_head")
    return list(lora_module_names)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from sagemaker.remote_function import remote
import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import transformers

# 开始训练
@remote(volume_size=50)
def train_fn(
        model_name,
        train_ds,
        test_ds,
        lora_r=8,
        lora_alpha=32,
        lora_dropout=0.05,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        learning_rate=2e-4,
        num_train_epochs=1
):
    # 标记化和分块数据集
    lm_train_dataset = train_ds.map(
        lambda sample: tokenizer(sample["text"]), batched=True, batch_size=24, remove_columns=list(train_dataset.features)
    )


    lm_test_dataset = test_ds.map(
        lambda sample: tokenizer(sample["text"]), batched=True, remove_columns=list(test_dataset.features)
    )

    # 打印样本总数
    print(f"训练样本总数:{len(lm_train_dataset)}")

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
    # Falcon要求您允许远程代码执行。这是因为该模型使用的是尚未包含在transformers中的新架构。
    # 代码由模型作者在repo中提供。
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        quantization_config=bnb_config,
        device_map="auto")

    model.gradient_checkpointing_enable()
    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)

    # 获取lora目标模块
    modules = find_all_linear_names(model)
    print(f"找到{len(modules)}个需要量化的模块:{modules}")

    config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=modules,
        lora_dropout=lora_dropout,
        bias="none",
        task_type="CAUSAL_LM"
    )

    model = get_peft_model(model, config)
    print_trainable_parameters(model)

    trainer = transformers.Trainer(
        model=model,
        train_dataset=lm_train_dataset,
        eval_dataset=lm_test_dataset,
        args=transformers.TrainingArguments(
            per_device_train_batch_size=per_device_train_batch_size,
            per_device_eval_batch_size=per_device_eval_batch_size,
            logging_steps=2,
            num_train_epochs=num_train_epochs,
            learning_rate=learning_rate,
            bf16=True,
            save_strategy="no",
            output_dir="outputs"
        ),
        data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
    )
    model.config.use_cache = False

    trainer.train()
    trainer.evaluate()

    model.save_pretrained("/opt/ml/model")

并调用train_fn()

train_fn(model_id, train_dataset, test_dataset)

调优作业将在Amazon SageMaker训练集群上运行。等待调优作业完成。

现在,是时候对模型进行一些测试了。首先,让我们加载模型:

from peft import PeftModel, PeftConfig
import torch
from transformers import AutoModelForCausalLM

device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'

config = PeftConfig.from_pretrained("./model")
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, trust_remote_code=True)
model = PeftModel.from_pretrained(model, "./model")
model.to(device)

现在从训练数据集中加载一个示例问题,以查看原始答案,然后向微调模型提问相同的问题,以查看答案的比较。

以下是训练集中的一个问题样本和原始答案:

使用@remote装饰器在Amazon SageMaker上微调Falcon 7B和其他LLMs 四海 第3张

现在,同一个问题被问到调优后的Falcon-7B模型:

使用@remote装饰器在Amazon SageMaker上微调Falcon 7B和其他LLMs 四海 第4张

这样就完成了使用Amazon SageMaker Python SDK的@remote装饰器在AWS服务FAQ数据集上对Falcon-7B进行微调的实现。

清理

按照以下步骤清理资源:

  • 关闭Amazon SageMaker Studio实例以避免产生额外的费用。

  • 清理Amazon Elastic File System(Amazon EFS)目录,清除Hugging Face缓存目录:

    rm -R ~/.cache/huggingface/hub

结论

在本文中,我们展示了如何有效地使用@remote装饰器的功能来使用QLoRA、Hugging Face PEFT和bitsandbtyes微调Falcon-7B模型,而不需要对训练笔记本进行重大更改,并使用Amazon SageMaker的能力在远程集群上运行训练作业。

本文中展示的所有代码都可以在GitHub存储库中找到。该存储库还包含一个笔记本,展示了如何微调Llama-13B。

作为下一步,我们鼓励您查看@remote装饰器的功能和Python SDK API,并在您选择的环境和IDE中使用它。amazon-sagemaker-examples存储库中还有其他示例可以帮助您快速入门。您还可以查看以下文章:

  • 使用最少的代码更改将您的本地机器学习代码作为Amazon SageMaker训练作业运行
  • 使用@remote装饰器访问私有存储库以进行Amazon SageMaker训练工作负载
  • 使用QLoRA在Amazon SageMaker Studio笔记本上交互式微调Falcon-40B和其他LLM模型
Leave a Reply

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