使用 LoRA 和 Hugging Face 高效訓練大語言模型

在本文中，我們將展示如何使用 大語言模型低秩適配 (Low-Rank Adaptation of Large Language Models，LoRA) 技術在單 GPU 上微調 110 億引數的 FLAN-T5 XXL 模型。在此過程中，我們會使用到 Hugging Face 的 Transformers、Accelerate 和 PEFT 庫。

透過本文，你會學到:

1. 如何搭建開發環境
2. 如何載入並準備資料集
3. 如何使用 LoRA 和 bnb (即 bitsandbytes) int-8 微調 T5
4. 如何評估 LoRA FLAN-T5 並將其用於推理
5. 如何比較不同方案的價效比

另外，你可以 點選這裡 線上檢視此博文對應的 Jupyter Notebook。

快速入門: 輕量化微調 (Parameter Efficient Fine-Tuning，PEFT)

PEFT 是 Hugging Face 的一個新的開源庫。使用 PEFT 庫，無需微調模型的全部引數，即可高效地將預訓練語言模型 (Pre-trained Language Model，PLM) 適配到各種下游應用。PEFT 目前支援以下幾種方法:

• LoRA: LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS
• Prefix Tuning: P-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks
• P-Tuning: GPT Understands, Too
• Prompt Tuning: The Power of Scale for Parameter-Efficient Prompt Tuning

注意: 本教程是在 g5.2xlarge AWS EC2 例項上建立和執行的，該例項包含 1 個 NVIDIA A10G。

1. 搭建開發環境

在本例中，我們使用 AWS 預置的 PyTorch 深度學習 AMI，其已安裝了正確的 CUDA 驅動程式和 PyTorch。在此基礎上，我們還需要安裝一些 Hugging Face 庫，包括 transformers 和 datasets。執行下面的程式碼就可安裝所有需要的包。

# install Hugging Face Libraries
!pip install git+https://github.com/huggingface/peft.git
!pip install "transformers==4.27.1" "datasets==2.9.0" "accelerate==0.17.1" "evaluate==0.4.0" "bitsandbytes==0.37.1" loralib --upgrade --quiet
# install additional dependencies needed for training
!pip install rouge-score tensorboard py7zr

2. 載入並準備資料集

這裡，我們使用 samsum 資料集，該資料集包含大約 16k 個含摘要的聊天類對話資料。這些對話由精通英語的語言學家制作。

{
  "id": "13818513",
  "summary": "Amanda baked cookies and will bring Jerry some tomorrow.",
  "dialogue": "Amanda: I baked cookies. Do you want some?\r\nJerry: Sure!\r\nAmanda: I'll bring you tomorrow :-)"
}

我們使用 ? Datasets 庫中的 load_dataset() 方法來載入 samsum 資料集。

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'])}")

# Train dataset size: 14732
# Test dataset size: 819

為了訓練模型，我們要用 ? Transformers Tokenizer 將輸入文字轉換為詞元 ID。如果你需要了解這一方面的知識，請移步 Hugging Face 課程的 第 6 章。

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model_id="google/flan-t5-xxl"

# Load tokenizer of FLAN-t5-XL
tokenizer = AutoTokenizer.from_pretrained(model_id)

在開始訓練之前，我們還需要對資料進行預處理。生成式文字摘要屬於文字生成任務。我們將文字輸入給模型，模型會輸出摘要。我們需要了解輸入和輸出文字的長度資訊，以利於我們高效地批次處理這些資料。

from datasets import concatenate_datasets
import numpy as np
# The maximum total input sequence length after tokenization.
# Sequences longer than this will be truncated, sequences shorter will be padded.
tokenized_inputs = concatenate_datasets([dataset["train"], dataset["test"]]).map(lambda x: tokenizer(x["dialogue"], truncation=True), batched=True, remove_columns=["dialogue", "summary"])
input_lenghts = [len(x) for x in tokenized_inputs["input_ids"]]
# take 85 percentile of max length for better utilization
max_source_length = int(np.percentile(input_lenghts, 85))
print(f"Max source length: {max_source_length}")

# The maximum total sequence length for target text after tokenization.
# Sequences longer than this will be truncated, sequences shorter will be padded."
tokenized_targets = concatenate_datasets([dataset["train"], dataset["test"]]).map(lambda x: tokenizer(x["summary"], truncation=True), batched=True, remove_columns=["dialogue", "summary"])
target_lenghts = [len(x) for x in tokenized_targets["input_ids"]]
# take 90 percentile of max length for better utilization
max_target_length = int(np.percentile(target_lenghts, 90))
print(f"Max target length: {max_target_length}")

我們將在訓練前統一對資料集進行預處理並將預處理後的資料集儲存到磁碟。你可以在本地機器或 CPU 上執行此步驟並將其上傳到 Hugging Face Hub。

def preprocess_function(sample,padding="max_length"):
    # add prefix to the input for t5
    inputs = ["summarize: " + item for item in sample["dialogue"]]

    # tokenize inputs
    model_inputs = tokenizer(inputs, max_length=max_source_length, padding=padding, truncation=True)

    # Tokenize targets with the `text_target` keyword argument
    labels = tokenizer(text_target=sample["summary"], max_length=max_target_length, padding=padding, truncation=True)

    # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
    # padding in the loss.
    if padding == "max_length":
        labels["input_ids"] = [
            [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
        ]

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=["dialogue", "summary", "id"])
print(f"Keys of tokenized dataset: {list(tokenized_dataset['train'].features)}")

# save datasets to disk for later easy loading
tokenized_dataset["train"].save_to_disk("data/train")
tokenized_dataset["test"].save_to_disk("data/eval")

3. 使用 LoRA 和 bnb int-8 微調 T5

除了 LoRA 技術，我們還使用 bitsanbytes LLM.int8() 把凍結的 LLM 量化為 int8。這使我們能夠將 FLAN-T5 XXL 所需的記憶體降低到約四分之一。

訓練的第一步是載入模型。我們使用 philschmid/flan-t5-xxl-sharded-fp16 模型，它是 google/flan-t5-xxl 的分片版。分片可以讓我們在載入模型時不耗盡記憶體。

from transformers import AutoModelForSeq2SeqLM

# huggingface hub model id
model_id = "philschmid/flan-t5-xxl-sharded-fp16"

# load model from the hub
model = AutoModelForSeq2SeqLM.from_pretrained(model_id, load_in_8bit=True, device_map="auto")

現在，我們可以使用 peft 為 LoRA int-8 訓練作準備了。

from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, TaskType

# Define LoRA Config
lora_config = LoraConfig(
 r=16,
 lora_alpha=32,
 target_modules=["q", "v"],
 lora_dropout=0.05,
 bias="none",
 task_type=TaskType.SEQ_2_SEQ_LM
)
# prepare int-8 model for training
model = prepare_model_for_int8_training(model)

# add LoRA adaptor
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# trainable params: 18874368 || all params: 11154206720 || trainable%: 0.16921300163961817

如你所見，這裡我們只訓練了模型引數的 0.16%！這個巨大的記憶體增益讓我們安心地微調模型，而不用擔心記憶體問題。

接下來需要建立一個 DataCollator，負責對輸入和標籤進行填充，我們使用 ? Transformers 庫中的 DataCollatorForSeq2Seq 來完成這一環節。

from transformers import DataCollatorForSeq2Seq

# we want to ignore tokenizer pad token in the loss
label_pad_token_id = -100
# Data collator
data_collator = DataCollatorForSeq2Seq(
    tokenizer,
    model=model,
    label_pad_token_id=label_pad_token_id,
    pad_to_multiple_of=8
)

最後一步是定義訓練超參 ( TrainingArguments)。

from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments

output_dir="lora-flan-t5-xxl"

# Define training args
training_args = Seq2SeqTrainingArguments(
    output_dir=output_dir,
        auto_find_batch_size=True,
    learning_rate=1e-3, # higher learning rate
    num_train_epochs=5,
    logging_dir=f"{output_dir}/logs",
    logging_strategy="steps",
    logging_steps=500,
    save_strategy="no",
    report_to="tensorboard",
)

# Create Trainer instance
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=tokenized_dataset["train"],
)
model.config.use_cache = False # silence the warnings. Please re-enable for inference!

執行下面的程式碼，開始訓練模型。請注意，對於 T5，出於收斂穩定性考量，某些層我們仍保持 float32 精度。

# train model
trainer.train()

訓練耗時約 10 小時 36 分鐘，訓練 10 小時的成本約為 13.22 美元。相比之下，如果 在 FLAN-T5-XXL 上進行全模型微調 10 個小時，我們需要 8 個 A100 40GB，成本約為 322 美元。

我們可以將模型儲存下來以用於後面的推理和評估。我們暫時將其儲存到磁碟，但你也可以使用 model.push_to_hub 方法將其上傳到 Hugging Face Hub。

# Save our LoRA model & tokenizer results
peft_model_id="results"
trainer.model.save_pretrained(peft_model_id)
tokenizer.save_pretrained(peft_model_id)
# if you want to save the base model to call
# trainer.model.base_model.save_pretrained(peft_model_id)

最後生成的 LoRA checkpoint 檔案很小，僅需 84MB 就包含了從 samsum 資料集上學到的所有知識。

4. 使用 LoRA FLAN-T5 進行評估和推理

我們將使用 evaluate 庫來評估 rogue 分數。我們可以使用 PEFT 和 transformers 來對 FLAN-T5 XXL 模型進行推理。對 FLAN-T5 XXL 模型，我們至少需要 18GB 的​​ GPU 視訊記憶體。

import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# Load peft config for pre-trained checkpoint etc.
peft_model_id = "results"
config = PeftConfig.from_pretrained(peft_model_id)

# load base LLM model and tokenizer
model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model_name_or_path, load_in_8bit=True, device_map={"":0})
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

# Load the Lora model
model = PeftModel.from_pretrained(model, peft_model_id, device_map={"":0})
model.eval()

print("Peft model loaded")

我們用測試資料集中的一個隨機樣本來試試摘要效果。

from datasets import load_dataset
from random import randrange

# Load dataset from the hub and get a sample
dataset = load_dataset("samsum")
sample = dataset['test'][randrange(len(dataset["test"]))]

input_ids = tokenizer(sample["dialogue"], return_tensors="pt", truncation=True).input_ids.cuda()
# with torch.inference_mode():
outputs = model.generate(input_ids=input_ids, max_new_tokens=10, do_sample=True, top_p=0.9)
print(f"input sentence: {sample['dialogue']}\n{'---'* 20}")

print(f"summary:\n{tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]}")

不錯！我們的模型有效！現在，讓我們仔細看看，並使用 test 集中的全部資料對其進行評估。為此，我們需要實現一些工具函式來幫助生成摘要並將其與相應的參考摘要組合到一起。評估摘要任務最常用的指標是 rogue_score)，它的全稱是 Recall-Oriented Understudy for Gisting Evaluation。與常用的準確率指標不同，它將生成的摘要與一組參考摘要進行比較。

import evaluate
import numpy as np
from datasets import load_from_disk
from tqdm import tqdm

# Metric
metric = evaluate.load("rouge")

def evaluate_peft_model(sample,max_target_length=50):
    # generate summary
    outputs = model.generate(input_ids=sample["input_ids"].unsqueeze(0).cuda(), do_sample=True, top_p=0.9, max_new_tokens=max_target_length)
    prediction = tokenizer.decode(outputs[0].detach().cpu().numpy(), skip_special_tokens=True)
    # decode eval sample
    # Replace -100 in the labels as we can't decode them.
    labels = np.where(sample['labels']!= -100, sample['labels'], tokenizer.pad_token_id)
    labels = tokenizer.decode(labels, skip_special_tokens=True)

    # Some simple post-processing
    return prediction, labels

# load test dataset from distk
test_dataset = load_from_disk("data/eval/").with_format("torch")

# run predictions
# this can take ~45 minutes
predictions, references = [], []
for sample in tqdm(test_dataset):
    p,l = evaluate_peft_model(sample)
    predictions.append(p)
    references.append(l)

# compute metric
rogue = metric.compute(predictions=predictions, references=references, use_stemmer=True)

# print results
print(f"Rogue1: {rogue['rouge1']* 100:2f}%")
print(f"rouge2: {rogue['rouge2']* 100:2f}%")
print(f"rougeL: {rogue['rougeL']* 100:2f}%")
print(f"rougeLsum: {rogue['rougeLsum']* 100:2f}%")

# Rogue1: 50.386161%
# rouge2: 24.842412%
# rougeL: 41.370130%
# rougeLsum: 41.394230%

我們 PEFT 微調後的 FLAN-T5-XXL 在測試集上取得了 50.38% 的 rogue1 分數。相比之下，flan-t5-base 的全模型微調獲得了 47.23 的 rouge1 分數。rouge1 分數提高了 3%。

令人難以置信的是，我們的 LoRA checkpoint 只有 84MB，而且效能比對更小的模型進行全模型微調後的 checkpoint 更好。

你可以 點選這裡 線上檢視此博文對應的 Jupyter Notebook。

英文原文: https://www.philschmid.de/fine-tune-flan-t5-peft

原文作者：Philipp Schmid

譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

排版/審校: zhongdongy (阿東)