Generative AI 新世界 | Falcon 40B 大模型微调和量化实践

在上一期的文章中，我们一起梳理了[大模型参数高效微调（PEFT）和 QLoRA 量化技术背后的理论基础](https://dev.amazoncloud.cn/column/article/64e74d0d1defb92aa89afc06)。与标准的 16 位模型微调相比，QLoRA 减少了大模型微调的内存使用量，而无需权衡性能。探索完基本理论之后，我们就要开始动手实践了。 本期文章，我们将探讨使用 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio 在交互式环境中，快速高效地微调大语言模型。我们将运用 QLoRA 和 4-bits的bitsandbtyes  量化技术原理，在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio 上使用 Hugging Face PEFT 来微调 Falcon-40B 模型。 ### **实践方法概述** [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 提供了两个选项方法，用于启动完全托管的 notebook，用于探索数据和构建[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)（ML）模型。 第一种选项是 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio。这是一个完全集成的[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)开发环境（IDE），用户可以在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio 中快速启动 notebook，在不中断工作的情况下向上或向下伸缩底层计算资源，甚至可以在 notebook 上实时共同编辑和协作。用户可以在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio 的单一管理面板中执行所有[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)开发步骤，包括：构建、训练、调试、跟踪、部署和监控模型等。 第二个选项是 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Notebook 实例。这是一个在云端运行 notebook 的完全托管的 ML 计算实例，这个方法可以帮助用户更好地控制 notebook 配置。 本例中我们将使用** [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio**。主要选择的原因有两点： 1） 可以利用 SageMaker Studio 的托管 TensorBoard 实验跟踪，以及 Hugging Face Transformer 对 TensorBoard 的支持； 2） [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio 的 [Amazon EFS](https://aws.amazon.com/cn/efs/?trk=cndc-detail) 容量，可以无需预先预置 EBS 卷大小。鉴于 LLM 中模型权重较大，这在实践中很有帮助。 如果你想选择第二个选项，也是可行的。示例代码将同样适用于使用 conda_pytorch_p310 内核的 notebook 实例。 另外，在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio 使用完 notebook 实例后，请将其关闭，以避免产生不必要的额外费用。本文最后有删除和清理资源的示例代码，可供参考。 本文的主要参考文档来自以下亚马逊云科技的官方博客。为阐述清楚其中的细节，本文增加了较多的章节扩展分析和代码对照的讲解： https://aws.amazon.com/cn/blogs/machine-learning/interactively-fine-tune-falcon-40b-and-other-llms-on-amazon-sagemaker-studio-notebooks-using-qlora/?trk=cndc-detail ### **模型微调过程的拆解分析** #### **1**.**启动 Amazon SageMaker JumpStart 环境** 本实验的完整示例代码可参考： https://github.com/aws-samples/amazon-sagemaker-generativeai/blob/main/studio-notebook-fine-tuning/falcon-40b-qlora-finetune-summarize.ipynb?trk=cndc-detail 示例代码的 notebook 在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio 测试通过，内核为 Python 3（Data Science 3.0），实例为一台 ml.g5.12xlarge 实例。 #### **2.安装模型微调所需的库** 首先，安装所需的库，包括 Hugging Face 库，然后重新启动内核。 ```js %pip install -q -U torch==2.0.1 bitsandbytes==0.40.2 %pip install -q -U transformers==4.31.0 peft==0.4.0 accelerate==0.21.0 %pip install -q -U datasets py7zr einops tensorboardX # Add installed cuda runtime to path for bitsandbytes import os import nvidia cuda_install_dir = '/'.join(nvidia.__file__.split('/')[:-1]) + '/cuda_runtime/lib/' os.environ['LD_LIBRARY_PATH'] = cuda_install_dir print("cuda_install_dir: ",cuda_install_dir) ``` 我自己环境的 cuda_install_dir 路径如下输出所示：  #### **3.输入文本的 Tokenizer 和大模型的量化** 要训练模型，我们需要将输入文本转换为 token ID，这个工作可以交给 Hugging Face Transformers Tokenizer 完成。除了 QLoRA 之外，我们还将使用 bitsandbytes 的 4 位精度方法将 LLM 量化为 4 位，并遵照在上篇文章中介绍过的 QLoRA 论文中阐述的方法，在量化后的 LLM 上面配接 LoRA adapter。 ```js import torch from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig 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 ) tokenizer = AutoTokenizer.from_pretrained(model_id) # Falcon requires you to allow remote code execution. This is because the model uses a new architecture that is not part of transformers yet. # The code is provided by the model authors in the repo. model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, quantization_config=bnb_config, device_map="auto") ```   这个 Falcon 40B LLM 量化到 4-bit 精度的过程，我自己实测大约需要用时 5 分钟左右。如上图所示。 另外，我们还需要设置代表句子结尾的特殊标记，如下图输出的 “|endoftext|”，就是代表句子结尾的特殊标记。关于 Tokenizer 的设置详情，可参考 [Tokenizer的类说明文档](https://huggingface.co/docs/transformers/main_classes/tokenizer?trk=cndc-detail)。 ```js # Set the Falcon tokenizer tokenizer.pad_token = tokenizer.eos_token print("tokenizer.pad_token: ", tokenizer.pad_token) ``` 输出如下图所示：  #### **4.为 LoRA 微调训练方法准备模型** 接下来要开始为 PEFT 工作准备模型了。 ```js from peft import prepare_model_for_kbit_training model.gradient_checkpointing_enable() model = prepare_model_for_kbit_training(model) print("model: ", model) ```  定义打印模型中可训练参数的函数： ```js def print_trainable_parameters(model): """ Prints the number of trainable parameters in the model. """ trainable_params = 0 all_param = 0 for _, param in model.named_parameters(): all_param += param.numel() if param.requires_grad: trainable_params += param.numel() print( f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" ) ``` 使用 PEFT 为 LoRA 微调训练方法准备模型和 LoRA 参数配置细节： ```js model, preprocess = clip.load("ViT-B/32") model.cuda().eval() input_resolution = model.visual.input_resolution context_length = model.context_length vocab_size = model.vocab_size ``` 注意看输出的最后一行，即可训练参数的数量：  由以上输出可见，所有参数数量是 209 亿以上，可训练参数数量是 5500 万，只占全部参数的 0.26%。由于训练参数的大幅减少，可预计的训练时间也会大大缩短。 #### 5.加载微调模型的数据集 要加载 [samsum](https://huggingface.co/datasets/samsum?trk=cndc-detail) 数据集，我们使用 Hugging Face 数据集库中的 load_dataset () 方法。 ```js from datasets import load_dataset # Load dataset from the hub dataset = load_dataset("samsum") print(f"Train dataset size: {len(dataset['train'])}") print(f"Test dataset size: {len(dataset['test'])}") ```  由以上输出可见，这个数据集并不大。训练数据 14732 条，测试数据 819 条。 接下来，我们需要创建提示词模板，并使用随机样本加载数据集做汇总测试。 ```js from random import randint # custom instruct prompt start prompt_template = f"Summarize the chat dialogue:\\n{{dialogue}}\\n---\\nSummary:\\n{{summary}}{{eos_token}}" # template dataset to add prompt to each sample def template_dataset(sample): sample["text"] = prompt_template.format(dialogue=sample["dialogue"], summary=sample["summary"], eos_token=tokenizer.eos_token) return sample # apply prompt template per sample train_dataset = dataset["train"].map(template_dataset, remove_columns=list(dataset["train"].features)) print(train_dataset[randint(0, len(dataset))]["text"]) ``` 提示词模版示例输出如下图所示。  以下代码将提示词模版应用到每一个 sample 里： ```js # apply prompt template per sample test_dataset = dataset["test"].map(template_dataset, remove_columns=list(dataset["test"].features)) ``` 对数据集做 tokenize 和 chunk： ```js # tokenize and chunk dataset lm_train_dataset = train_dataset.map( lambda sample: tokenizer(sample["text"]), batched=True, batch_size=24, remove_columns=list(train_dataset.features) ) lm_test_dataset = test_dataset.map( lambda sample: tokenizer(sample["text"]), batched=True, remove_columns=list(test_dataset.features) ) # Print total number of samples print(f"Total number of train samples: {len(lm_train_dataset)}") ```  为了完成实时监控的任务，我们首先需要把训练过程的各项指标记录下来，比如：用一个 S3 桶来存放记录。因此，我们还需要为这个实验创建一个 S3 桶，以方便我们将训练中的各项指标完整地记录到 TensorBoard： ```js # bucket = <YOUR-S3-BUCKET> bucket = "llm-demo-xxxxxx" log_bucket = f"s3://{bucket}/falcon-40b-qlora-finetune" log_bucket ``` 如果你试着打印 log_bucket 的输出，将类似是如下这样的，它定义了一个 S3 桶： ```js 's3://llm-demo-xxxxxx/falcon-40b-qlora-finetune' ``` #### 6.微调训练过程指标的记录 然后将使用 Hugging Face Trainer 类对模型进行微调，定义要使用的超参数。我们还创建了一个 DataCollator 来填充我们的输入和标签。另外，因为要考虑做模型微调训练过程的监测，可以通过定义参数 logging_dir 和设置report_to="tensorboard"，来请求 Hugging Face Transformer 把微调训练的日志记录到 TensorBoard。 ```js import transformers # We set num_train_epochs=1 simply to run a demonstration 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), ) model.config.use_cache = False # silence the warnings. Please re-enable for inference! print("model:", model) ``` 完成以上配置，就可以开始正式启动微调的训练了。启动的代码很简单，如下所示： ### **模型微调监控的拆解分析** #### **1.实时监控 GPU 使用情况** 前面的设置完成后，我们就可以实时监控微调过程了。为了实时监控 GPU 使用情况，我们可以直接从内核的容器运行 nvidia-smi 命令。要启动在镜像容器上运行的终端，只需选择笔记本顶部的终端图标即可。  我们可以使用 Linux watch 命令每半秒钟重复运行 nvidia-smi： ```js watch -n 0.5 nvidia-smi ```  在上面的动图中，我们可以看到模型权重分布在 4 个 GPU 上，随着层的串行处理，计算负载在这些 GPU 之间做分布计算。 #### 2.实时监控模型微调的训练指标 为了监控模型微调过程中的训练指标，我们将把 TensorBoard 日志写入之前已经配置好的 S3 桶。可以从 SageMaker 控制台启动 SageMaker Studio 域用户的 TensorBoard，如下截图所示：  TensorBoard 加载完成后，可以指定 Hugging Face transformer 把训练日志写入指定的 S3 桶，以便查看训练和评估指标，如下图所示。  以上配置完成后，就可以通过 TensorBoard 的 “Time Series” 菜单，监控微调模型训练过程中各项指标随着训练时间的变化，包括：回合（epoch）、学习率（learning rate）、损失函数（loss）等等。  ### **模型微调后的评估和生成结果** 模型完成微调训练后，我们就可以对微调后的大模型进行系统评估或直接生成结果了。由于模型的评估又是另一个宏大的话题，本文会暂时略过这个话题，留待以后专题讨论；本文之后将主要聚焦微调后模型的结果生成示例。 首先，加载之前分拆出来的 samsum 测试数据集，并尝试使用随机样本进行 LLM 总结（Summary）测试： ```js # Load dataset from the hub test_dataset = load_dataset("samsum", split="test") # select a random test sample sample = test_dataset[randint(0, len(test_dataset))] # format sample prompt_template = f"Summarize the chat dialogue:\\n{{dialogue}}\\n---\\nSummary:\\n" test_sample = prompt_template.format(dialogue=sample["dialogue"]) print(test_sample) ```  接下来，把输入数据 tokenizer 化： ```js input_ids = tokenizer(test_sample, return_tensors="pt").input_ids ``` 把 tokenizer 化后的输入数据传给微调后的 LLM，获取 LLM 总结（Summary）的输出结果： ```js #set the tokens for the summary evaluation 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) ```  有心的同学，可以通过模型微调前后，LLM 输出的总结（Summary）结果比较，来对比微调后的模型，是否在准确性和完整性上有了一定的改善。 如果你对模型的性能感到满意，可以把模型保存下来，如下代码所示： ```js trainer.save_model("path_to_save") ``` 或者把模型部署到一个专门的 SageMaker 终端节点。部署终端节点的文档可参考： https://aws.amazon.com/blogs/machine-learning/deploy-falcon-40b-with-large-model-inference-dlcs-on-amazon-sagemaker/?trk=cndc-detail ### **资源的删除和清理** 实验完成后，请记得删除和清理资源，以避免不必要的额外费用。需要清理的资源分为三个部分，如下所示： 1）关闭 SageMaker Studio 实例 https://docs.aws.amazon.com/sagemaker/latest/dg/notebooks-run-and-manage-shut-down.html?trk=cndc-detail 2）关闭你的 TensorBoard 应用程序 https://docs.aws.amazon.com/sagemaker/latest/dg/tensorboard-on-sagemaker.html#debugger-htb-delete-app?trk=cndc-detail 3）清除 Hugging Face 缓存目录，参考命令如下所示： ```js rm -R ~/.cache/huggingface/hub ``` ### **总结** 在本文中，我们探讨了使用 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio在交互式环境中，快速高效地微调 Falcon 40B 大语言模型。我们运用了 QLoRA和 4-bits 的 bitsandbtyes 量化技术原理，在 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) Studio 上使用 Hugging Face PEFT 微调了 Falcon-40B 模型。 本文做为 “Generative AI 新世界”的第十二篇文章，在不知不觉中已经伴随着各位热爱生成式 AI 领域知识的读者们，从初春走到了盛夏。随着我们一起对生成式 AI 知识的逐步深入学习，这个系列后面的文章内容会往更深度的专业领域做拓展。 目前计划有三个大方向： 1）代码深度实践方向。例如用代码完整诠释 Diffusion 模型的工作原理，或者 Transformer 的完整架构等； 2）模型部署和训练优化方向。例如尝试解读 LMI、DeepSpeed、Accelerate、FlashAttention 等不同模型优化方向的最新进展； 3）模型量化实践方向。例如 GPTQ、bitsandbtyes 等前沿模型量化原理和实践等。敬请期待。 请持续关注 Build On Cloud 专栏，了解更多面向开发者的技术分享和云开发动态！