僅用250美元,Hugging Face技術主管手把手教你微調Llama 3

机器之心發表於2024-05-06
語言模型的微調一直是說起來容易做起來難的事兒。近日 Hugging Face 技術主管 Philipp Schmid 發表了一篇部落格,詳細講解了如何利用 Hugging Face 上的庫和 fsdp 以及 Q-Lora 對大模型進行微調。

圖片

我們知道,Meta 推出的 Llama 3、Mistral AI 推出的 Mistral 和 Mixtral 模型以及 AI21 實驗室推出的 Jamba 等開源大語言模型已經成為 OpenAI 的競爭對手。

不過,大多數情況下,使用者需要根據自己的資料對這些開源模型進行微調,才能充分釋放模型的潛力。

雖然在單個 GPU 上使用 Q-Lora 對較小的大語言模型(如 Mistral)進行微調不是難事,但對像 Llama 3 70b 或 Mixtral 這樣的大模型的高效微調直到現在仍是一個難題。

因此,Hugging Face 技術主管 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 微調大語言模型
  • 測試模型並進行推理

注:本文進行的實驗是在英偉達(NVIDIA)H100 和英偉達(NVIDIA)A10G GPU 上建立和驗證的。配置檔案和程式碼針對 4xA10G GPU 進行了最佳化,每個 GPU 均配備 24GB 記憶體。如果使用者有更多的算力,第 3 步提到的配置檔案(yaml 檔案)需要做相應的修改。

FSDP+Q-Lora 背景知識

基於一項由 Answer.AI、Q-Lora 建立者 Tim Dettmers 和 Hugging Face 共同參與的合作專案,作者對 Q-Lora 和 PyTorch FSDP(完全共享資料並行)所能提供的技術支援進行了總結。

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 以及 Pyroch,包括 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),使語言模型更好地遵循人類指令。no_robots 資料集以 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": "..."}]}

no_robots 資料集中的 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 have 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)

# save 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 來對模型進行微調。

SFTTrainer 能夠讓對開源大語言模型的有監督微調更加容易上手,具體來說有以下幾點:

  • 格式化的資料集,包括格式化的多輪會話和指令(已使用)
  • 只對完整的內容進行訓練,忽略只有 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。以下是在 4x A10G GPU 或 4x24GB 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

save_strategy: epoch                   # save 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 save 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 一樣。

對於 torchrun 和 FSDP,作者需要對環境變數 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,以及一塊具備 22 GB 記憶體的 GPU 和 127 GB 的 CPU RAM,序列長度為 3072、batch 大小為 1。

在 g5.12xlarge 伺服器上,基於包含 1 萬個樣本的資料集,作者使用 Flash Attention 對 Llama 3 70B 進行 3 個 epoch 的訓練,總共需要 45 小時。每小時成本為 5.67 美元,總成本為 255.15 美元。這聽起來很貴,但可以讓你在較小的 GPU 資源上對 Llama 3 70B 進行微調。

如果我們將訓練擴充套件到 4x H100 GPU,訓練時間將縮短至大約 125 小時。如果假設 1 臺 H100 的成本為 5-10 美元 / 小時,那麼總成本將在 25-50 美元之間。

我們需要在易用性和效能之間做出權衡。如果能獲得更多更好的計算資源,就能減少訓練時間和成本,但即使只有少量資源,也能對 Llama 3 70B 進行微調。對於 4x A10G GPU 而言,需要將模型載入到 CPU 上,這就降低了總體 flops,因此成本和效能會有所不同。

注意:在作者進行的評估和測試過程中,他注意到大約 40 個最大步長(將 80 個樣本堆疊為長度為三千的序列)就足以獲得初步結果。40 個步長的訓練時間約為 1 小時,成本約合 5 美元。

可選步驟:將 LoRA 的介面卡融入原始模型

使用 QLoRA 時,作者只訓練介面卡而不對整個模型做出修改。這意味著在訓練過程中儲存模型時,只儲存介面卡權重,而不儲存完整模型。

如果使用者想儲存完整的模型,使其更容易與文字生成推理器一起使用,則可以使用 merge_and_unload 方法將介面卡權重合併到模型權重中,然後使用 save_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 save

# merged_model = model.merge_and_unload()

# merged_model.save_pretrained(args.output_dir,safe_serialization=True, max_shard_size="2GB")

模型測試和推理

訓練完成後,我們要對模型進行評估和測試。作者從原始資料集中載入不同的樣本,並手動評估模型。評估生成式人工智慧模型並非易事,因為一個輸入可能有多個正確的輸出。閱讀《評估 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

相關文章