仅用250美元，Hugging Face技术主管手把手教你微调Llama 3

机器之心报道

大语言模型的微调，说起来容易做起来难。最近 Hugging Face 技术主管 Philipp Schmid 发了一篇博客，手把手教大家如何用 Hugging Face 的库、FSDP 和 Q-Lora 来微调大模型。这一套组合拳下来，哪怕资源有限，也能搞定 Llama 3 70B 这样的大家伙。

Meta 的 Llama 3、Mistral AI 的 Mistral 和 Mixtral、AI21 的 Jamba 等开源模型，已经在和 OpenAI 掰手腕了。不过，想让这些模型真正派上用场，多数时候得靠自己的数据去微调。过去，在单张 GPU 上用 Q-Lora 微调像 Mistral 这样的小模型不算难，但轮到 Llama 3 70b 或 Mixtral 这种量级的，高效微调一直是个坎儿。

Philipp Schmid 这次在博客里详细介绍了怎么用 PyTorch FSDP 和 Q-Lora，配合 Hugging Face 的 TRL、Transformers、peft、datasets 等库，对 Llama 3 进行微调。除了 FSDP，他还适配了 PyTorch 2.2 更新后的 Flash Attention v2。整个微调流程大致分这么几步：

• 设置开发环境
• 创建并加载数据集
• 用 PyTorch FSDP、Q-Lora 和 SDPA 微调大语言模型
• 测试模型并进行推理

需要说明的是，实验是在英伟达 H100 和 A10G GPU 上做的。配置文件和代码针对 4 块 A10G GPU（每块 24GB 内存）做了优化。如果手头算力更足，第 3 步的 yaml 配置文件需要相应调整。

FSDP+Q-Lora 背景知识

这项工作是 Answer.AI、Q-Lora 创建者 Tim Dettmers 和 Hugging Face 合作项目的成果。FSDP 和 Q-Lora 结合，能在 2 块消费级 GPU（24GB）上就微调 Llama 2 70b 或 Mixtral 8x7B，细节可以参考下文。其中 Hugging Face 的 PEFT 库发挥了关键作用。

文章地址：https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html

简单说，PyTorch FSDP 是一种数据/模型并行技术，能把模型拆开跨 GPU 放，减少内存需求，训练更大的模型更高效。Q-LoRA 则是一种微调方法，靠量化和低秩适配器，有效降低计算和内存占用。两者一结合，门槛就降下来了。

设置开发环境

第一步自然是装 Hugging Face Libraries 和 PyTorch，包括 trl、transformers、datasets 等。trl 这个新库，建立在 transformers 和 datasets 之上，让微调、RLHF、对齐都变得更简单。

# Install Pytorch for FSDP and FA/SDPA
%pip install "torch==2.2.2" tensorboard

# Install Hugging Face libraries
%pip install  --upgrade "transformers==4.40.0" "datasets==2.18.0" "accelerate==0.29.3" "evaluate==0.4.1" "bitsandbytes==0.43.1" "huggingface_hub==0.22.2" "trl==0.8.6" "peft==0.10.0"

装好后，登录 Hugging Face 获取 Llama 3 70b 模型的访问权限。

创建和加载数据集

环境搞定，接下来就是准备数据了。微调用的数据集最好包含你任务的那些示例样本。关于怎么创建数据集，可以看看《如何在 2024 年使用 Hugging Face 微调 LLM》这篇文章（链接：https://www.philschmid.de/fine-tune-llms-in-2024-with-trl#3-create-and-prepare-the-dataset）。

作者用了 HuggingFaceH4/no_robots 数据集，包含 10,000 条高质量指令和样本，经过精细标注。这个数据集适合有监督微调（SFT），能让模型更好地遵循人类指令。它是以 OpenAI 的 InstructGPT 论文描述的人类指令数据集为原型，主要由单句指令构成。

{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}

数据集里的 10,000 个样本，分成 9,500 个训练和 500 个测试，有些样本缺了 system 信息。作者用 datasets 库加载数据，补上缺失的 system 信息，然后存成单独的 json 文件。

from datasets import load_dataset

# Convert dataset to OAI messages
system_message = """You are Llama, an AI assistant created by Philipp to be helpful and honest. Your knowledge spans a wide range of topics, allowing you to engage in substantive conversations and provide analysis on complex subjects."""

def create_conversation(sample):
  if sample["messages"][0]["role"] == "system":
    return sample
  else:
    sample["messages"] = [{"role": "system", "content": system_message}] + sample["messages"]
    return sample

# Load dataset from the hub
dataset = load_dataset("HuggingFaceH4/no_robots")

# Add system message to each conversation
columns_to_remove = list(dataset["train"].features)
columns_to_remove.remove("messages")
dataset = dataset.map(create_conversation, remove_columns=columns_to_remove,batched=False)

# Filter out conversations which are corrupted with wrong turns, keep which ha ve even number of turns after adding system message
dataset["train"] = dataset["train"].filter(lambda x: len(x["messages"][1:]) % 2 == 0)
dataset["test"] = dataset["test"].filter(lambda x: len(x["messages"][1:]) % 2 == 0)

# sa ve datasets to disk
dataset["train"].to_json("train_dataset.json", orient="records", force_ascii=False)
dataset["test"].to_json("test_dataset.json", orient="records", force_ascii=False)

使用 PyTorch FSDP、Q-Lora 和 SDPA 来微调 LLM

接下来是重头戏——用 PyTorch FSDP、Q-Lora 和 SDPA 开始微调。因为是在分布式设备上跑，所以得用 torchrun 和 python 脚本来启动训练。

作者写了一个 run_fsdp_qlora.py 脚本，负责从磁盘加载数据集、初始化模型和分词器，然后开始训练。脚本用的是 trl 库里的 SFTTrainer，这个 trainer 让有监督微调变得更顺手，具体体现在：

• 格式化的数据集（包括多轮会话和指令，已使用）
• 只训练完整内容，忽略只有 prompts 的情况（未使用）
• 打包数据集，提高训练效率（已使用）
• 支持参数高效微调技术，包括 Q-LoRA（已使用）
• 为会话级任务初始化模型和分词器（未使用，见下文）

有一点要提一下：作者用的是类似 Anthropic/Vicuna 的聊天模板，设置了「用户」和「助手」角色。这是因为基础 Llama 3 里的特殊分词器（<|begin_of_text|> 和 <|reserved_special_token_XX|>）没有经过训练。如果要在模板里用它们，就得额外训练，更新嵌入层和 lm_head，会增加内存消耗。如果算力足够，可以改 run_fsdp_qlora.py 脚本里的 LLAMA_3_CHAT_TEMPLATE 环境变量。

配置方面，作者用了新的 TrlParser 变量，既能在 yaml 文件里写超参数，也能通过 CLI 传参数覆盖，比如 --num_epochs 10。下面是在 4 块 A10G GPU（或 4 块 24GB GPU）上微调 Llama 3 70B 的配置文件：

%%writefile llama_3_70b_fsdp_qlora.yaml
# script parameters
model_id: "meta-llama/Meta-Llama-3-70b" # Hugging Face model id
dataset_path: "."  # path to dataset
max_seq_len:  3072 # 2048# max sequence length for model and packing of the dataset
# training parameters
output_dir: "./llama-3-70b-hf-no-robot" # Temporary output directory for model checkpoints
report_to: "tensorboard" # report metrics to tensorboard
learning_rate: 0.0002# learning rate 2e-4
lr_scheduler_type: "constant"# learning rate scheduler
num_train_epochs: 3# number of training epochs
per_device_train_batch_size: 1 # batch size per device during training
per_device_eval_batch_size: 1# batch size for evaluation
gradient_accumulation_steps: 2 # number of steps before performing a backward/update pass
optim: adamw_torch # use torch adamw optimizer
logging_steps: 10  # log every 10 steps
sa ve_strategy: epoch # sa ve checkpoint every epoch
evaluation_strategy: epoch   # evaluate every epoch
max_grad_norm: 0.3 # max gradient norm
warmup_ratio: 0.03 # warmup ratio
bf16: true # use bfloat16 precision
tf32: true # use tf32 precision
gradient_checkpointing: true # use gradient checkpointing to sa ve memory
# FSDP parameters: https://huggingface.co/docs/transformers/main/en/fsdp
fsdp: "full_shard auto_wrap offload" # remove offload if enough GPU memory
fsdp_config:
  backward_prefetch: "backward_pre"
  forward_prefetch: "false"
  use_orig_params: "false"

注意：训练结束时，GPU 内存使用量会略有上升（大约 10%），这是模型保存的开销。所以跑之前要确保 GPU 有富余内存来保存模型。

启动训练时，作者用了 torchrun，这样更灵活，也方便调整，像 Amazon SageMaker 和 Google Cloud Vertex AI 一样。同时还要设置环境变量 ACCELERATE_USE_FSDP 和 FSDP_CPU_RAM_EFFICIENT_LOADING，告诉 transformers/accelerate 用 FSDP 并以节省内存的方式加载模型。

如果不打算用 CPU offloading，得改 fsdp 设置，这招只适用于内存大于 40GB 的 GPU。

训练命令长这样：

!ACCELERATE_USE_FSDP=1 FSDP_CPU_RAM_EFFICIENT_LOADING=1 torchrun --nproc_per_node=4 ./scripts/run_fsdp_qlora.py --config llama_3_70b_fsdp_qlora.yaml

内存使用情况预期如下：

• 用 FSDP 全微调：约 16 块 80GB GPU
• FSDP+LoRA：约 8 块 80GB GPU
• FSDP+Q-Lora：约 2 块 40GB GPU
• FSDP+Q-Lora+CPU offloading：4 块 24GB GPU，外加一块 22GB GPU 和 127GB CPU RAM，序列长度 3072，batch 大小 1

作者在 g5.12xlarge 服务器上，用 1 万个样本的数据集，配合 Flash Attention，对 Llama 3 70B 训练了 3 个 epoch，总耗时 45 小时。每小时成本 5.67 美元，总计 255.15 美元。听起来不便宜，但能在比较小的 GPU 资源上搞定 Llama 3 70B，也算值了。

如果换成 4 块 H100 GPU，训练时间会缩短到大约 125 小时。按 1 块 H100 每小时 5-10 美元算，总成本在 25-50 美元之间。这里需要权衡易用性和性能。算力越多，训练时间和成本越低；但就算资源有限，也能微调 Llama 3 70B。不过对于 4 块 A10G 来说，得把模型加载到 CPU 上，这会拉低整体 flops，所以成本和性能会有些不同。

值得一提的是，作者在评估测试中发现，大约 40 个最大步长（80 个样本堆叠成 3000 长度的序列）就能得到初步结果。40 个步长训练大约 1 小时，成本约 5 美元。这个性价比相当可观。

可选步骤：将 LoRA 的适配器融入原始模型

用 QLoRA 时，只训练适配器，不碰整个模型。所以训练过程中保存模型时，只存适配器权重，而不是完整模型。

如果想存完整模型，方便和文本生成推理器一起用，可以用 merge_and_unload 方法把适配器权重合并到模型权重里，然后用 sa ve_pretrained 保存。这样会存出一个默认模型，直接用来推理。

注意：CPU 内存需要大于 192GB。

#### COMMENT IN TO MERGE PEFT AND BASE MODEL ####
# from peft import AutoPeftModelForCausalLM
# # Load PEFT model on CPU
# model = AutoPeftModelForCausalLM.from_pretrained(
#   args.output_dir,
#   torch_dtype=torch.float16,
#   low_cpu_mem_usage=True,
# )
# # Merge LoRA and base model and sa ve
# merged_model = model.merge_and_unload()
# merged_model.sa ve_pretrained(args.output_dir,safe_serialization=True, max_shard_size="2GB")

模型测试和推理

训练完了，当然要测一测。作者从原始数据集中挑了不同的样本，手工评估。评估生成式 AI 模型不是件容易的事，因为一个输入可能有多个正确答案。可以看看《评估 LLMs 和 RAG，一个使用 Langchain 和 Hugging Face 的实用案例》（链接：https://www.philschmid.de/evaluate-llm）。

import torch
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer

peft_model_id = "./llama-3-70b-hf-no-robot"

# Load Model with PEFT adapter
model = AutoPeftModelForCausalLM.from_pretrained(
  peft_model_id,
  torch_dtype=torch.float16,
  quantization_config= {"load_in_4bit": True},
  device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(peft_model_id)

接着加载测试数据集，试试生成指令：

from datasets import load_dataset
from random import randint

# Load our test dataset
eval_dataset = load_dataset("json", data_files="test_dataset.json", split="train")
rand_idx = randint(0, len(eval_dataset))
messages = eval_dataset[rand_idx]["messages"][:2]

# Test on sample
input_ids = tokenizer.apply_chat_template(messages,add_generation_prompt=True,return_tensors="pt").to(model.device)
outputs = model.generate(
  input_ids,
  max_new_tokens=512,
  eos_token_id= tokenizer.eos_token_id,
  do_sample=True,
  temperature=0.6,
  top_p=0.9,
)
response = outputs[0][input_ids.shape[-1]:]

print(f"**Query:**n{eval_dataset[rand_idx]['messages'][1]['content']}n")
print(f"**Original Answer:**n{eval_dataset[rand_idx]['messages'][2]['content']}n")
print(f"**Generated Answer:**n{tokenizer.decode(response,skip_special_tokens=True)}")

# **Query:**
# How long was the Revolutionary War?
# **Original Answer:**
# The American Revolutionary War lasted just over seven years. The war started on April 19, 1775, and ended on September 3, 1783.
# **Generated Answer:**
# The Revolutionary War, also known as the American Revolution, was an 18th-century war fought between the Kingdom of Great Britain and the Thirteen Colonies. The war lasted from 1775 to 1783.

到这里，整个流程就走完了。心动不如行动，赶紧从第一步开始试试吧。

原文链接：https://www.philschmid.de/fsdp-qlora-llama3?continueFlag=7e3b3f9059405e4318552e99bd128514