簡介

接下來要來討論如何微調 (Finetune) 一個大型語言模型。微調 LLM 與微調其他模型其實很相似，但是因為 LLM 的參數量較大，所以訓練的最低需求會比一般的模型高的多。而且 LLM 是個很精密的模型，所以訓練時很容易失敗，無法生出理想的結果。今天就來分享一下，訓練 LLM 的方法與心得。

任務設計

一開始學習訓練模型時，我們可以透過單純的玩具任務 (Toy Task) 來練習。玩具任務通常很簡單，因此模型容易學起來，而且資料來源也很容易爬取跟操作。在這裡，筆者設計一個路名解析任務，例如：

輸入：臺北市中正區八德路
輸出：{"city": "臺北市", "town": "中正區", "road": "八德路"}

我們輸入一段帶有行政區的路名，並且請模型將他解析成 JSON 格式。要讓模型學會完成這個任務，最單純的想法是輸入原句，並直接輸出 JSON 結果，像是將上方的範例直接丟進 LLM 裡面做訓練。

這樣做不是不行，但這是傳統 Encoder-Decoder LM 的訓練方法。拿這類的資料訓練現在的 Decoder LM 效果通常並不好，不如直接用參數量更少的 Encoder-Decoder LM 來訓練。那對於 Decoder LM 而言，怎樣的訓練資料比較好呢？對 Instruct LLM 而言，我們需要放入更多 Instruction 的成份在裡面：

### USER:

請將以下路名解析為 JSON 格式。

輸入：臺北市中正區八德路
輸出：{"city": "臺北市", "town": "中正區", "road": "八德路"}

輸入：屏東縣佳冬鄉文化三路

### ASSISTANT:

{"city": "屏東縣", "town": "佳冬鄉", "road": "文化三路"}

這樣的資料對於一個指令微調過的語言模型而言會更加有意義，訓練上也比較不容易破壞 LLM 的其他能力。因為 LLM 是藉由因果建模的方式進行訓練，因此給定越多的「因」就能產生越穩定的「果」。

試想我們直接將「臺北市中正區八德路」輸入到 ChatGPT 裡面，一定沒辦法直接得到 JSON 的解析結果，需要明確的指示 ChatGPT 將其解析為 JSON 格式，甚至給個範例，他才能完成我們的目標。因此對於一個 Instruct LLM 而言，在資料中加入明確指示是相當重要的。

建立資料集

行政區與路名的資料可以從政府資料開放平台取得，其中的全國路名資料就有我們需要的資訊，這裡筆者使用 112 全國路名資料的 CSV 檔做示範，其格式大致如下：

city,site_id,road
縣市名稱,行政區域名稱,全國路名
宜蘭縣,宜蘭縣三星鄉,人和一路
宜蘭縣,宜蘭縣三星鄉,人和七路
宜蘭縣,宜蘭縣三星鄉,人和九路

全臺灣有近三萬五千條路，每一行都包含了行政區以及路名。在這裡，我們需要將直轄縣市與鄉鎮市區分開，各自放進 city 與 town 裡面，最後與 road 形成完整的資料。首先對資料進行處理：

import csv

dataset = list()
with open("opendata112road.csv", "rt", encoding="UTF-8") as fp:
    # 跳過前兩行說明用的資料列
    fp.readline()  # city,site_id,road
    fp.readline()  # 縣市名稱,行政區域名稱,全國路名

    for city, site, road in csv.reader(fp):
        town = site.replace(city, "")
        item = {"city": city, "town": town, "road": road}
        dataset.append(item)

接下來，我們將資料切分成訓練、驗證、測試集等三份。因為只是小規模驗證，為了加快訓練速度，所以訓練與驗證集各取 100 筆資料即可，而測試集則可以取多一點：

import random

from utils import dump_json

random.seed(2135)
random.shuffle(dataset)

train = dataset[:100]
dev = dataset[100:200]
test = dataset[200:700]

dump_json(train, "train.json")
dump_json(dev, "dev.json")
dump_json(test, "test.json")

因為會對 JSON 檔案頻繁的操作，所以我先將常用 Functions 放在 utils.py 裡面：

import gzip
import json

def load_json(file_path):
    with open(file_path, "rt", encoding="UTF-8") as fp:
        return json.load(fp)


def dump_json(data, file_path):
    with open(file_path, "wt", encoding="UTF-8") as fp:
        json.dump(data, fp, ensure_ascii=False, indent=2)


def dump_json_gz(data, file_path):
    with gzip.open(file_path, "wt", encoding="UTF-8") as fp:
        json.dump(data, fp)

其中也有 GZip 版本的 JSON 讀寫，對於純文字的資料集而言，是個節省硬碟空間的好方法，而且 Hugging Face Datasets 也能直接把 .json.gz 檔當成一般 JSON 格式的檔案來讀取。

接下來要對訓練集與驗證集進行斷詞，這裡使用 TinyLlama 1.1B 模型進行實驗，因此使用他的 Tokenizer 來進行斷詞：

from transformers import LlamaTokenizerFast as TkCls

model_id = "PY007/TinyLlama-1.1B-Chat-v0.3"
tk: TkCls = TkCls.from_pretrained(model_id)

在 Hugging Face 的 CLM Training 裡面，每筆訓練資料必須包含 input_ids 與 labels 兩個欄位，這兩個欄位的內容通常是完全一樣的，例如：

[
    {
        "input_ids": [ 1, 43, 92, 71,  2, 2],
        "labels":    [ 1, 43, 92, 71,  2, 2]
    },
    {
        "input_ids": [ 1, 97, 66, 92, 71, 2],
        "labels":    [ 1, 97, 66, 92, 71, 2]
    }
]

當文本量很大時，每次訓練模型都要重新斷詞一次會顯得相對沒效率。因此筆者習慣先將文本進行斷詞，並存成檔案，這樣每次跑實驗時就不用再重新斷詞和建立輸入格式了。

拜訪資料集的內容：

def iter_dataset(file_path):
    data = load_json(file_path)

    for item in data:
        city = item["city"]
        town = item["town"]
        road = item["road"]

        full = f"{city}{town}{road}"

        yield full, item

我們將 City, Town, Road 合併在一起當作輸入句，並另外回傳完整的資料本體，用來比對答案是否正確。接下來將每筆資料放進 TinyLlama 提供的 Template 裡面：

template = """<|im_start|>user
請將以下路名解析為 JSON 格式。

輸入：臺北市中正區八德路
輸出：{{"city": "臺北市", "town": "中正區", "road": "八德路"}}

輸入：{}
<|im_end|>
<|im_start|>assistant
{}"""


def build_prompt(inn, out=""):
    return template.format(inn, out)


ds_type = "train"
ds_tokens = list()
for full, item in iter_dataset(f"{ds_type}.json"):
    # 將字典轉換為字串
    output = json.dumps(item, ensure_ascii=False)
    prompt = build_prompt(full, output)

    # 轉換成 Token 並加上 EOS
    tokens = tk.encode(prompt) + [tk.eos_token_id]
    ds_tokens.append(tokens)

將文本轉為 Token 後，結尾記得加上一個 EOS Token 來確保模型完成回答後可以結束生成。完成 Token 轉換後，我們需要對資料進行 Padding 對齊。因為訓練模型時，每個 Batch 裡面的資料，長度必須相同才能進行訓練：

# 計算最大長度
maxlen = max(map(len, ds_tokens))
print(f"Max Length: {maxlen}")

# 對資料集進行 Padding
dataset = list()
for tokens in ds_tokens:
    delta = maxlen - len(tokens)

    # 將 EOS Token 當作 PAD Token 來用
    tokens += [tk.eos_token_id] * delta

    # 訓練用的 input_ids 與 labels 通常是完全一樣的序列
    dataset.append({"input_ids": tokens, "labels": tokens})

# 確認所有序列的長度都是一致的
for item in dataset:
    assert len(item["input_ids"]) == maxlen
    assert len(item["labels"]) == maxlen

完成 Padding 後，就可以將斷詞完的結果存下來：

# 將斷詞完的結果存下來
dump_json_gz(dataset, f"{ds_type}.tokens.json.gz")

對訓練資料的 Padding 與推論時 Padding 的方向並不相同，在訓練時通常將 PAD Token 放在右邊，而推論時會放在左邊。

最後將 ds_type 分別設定為 train 與 dev 各跑一次，得到斷詞完的訓練集與驗證集。

訓練流程

接下來開始訓練模型的流程，首先將模型與 Tokenizer 讀取起來：

import torch
from transformers import LlamaForCausalLM as ModelCls
from transformers import LlamaTokenizerFast as TkCls

# 讀取 Model & Tokenizer
model_name = "PY007/TinyLlama-1.1B-Chat-v0.3"
model: ModelCls = ModelCls.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
tk: TkCls = TkCls.from_pretrained(model_name)

之前使用 HF Transformers 讀取模型時，都會加上 Quantization 的參數，但是目前理論上是不能直接對量化模型進行訓練的！因此這邊使用 torch.bfloat16 資料型態，俗稱 BF16，與一般的 Float16 (FP16) 不同的地方在於，BF16 可以表達的數值範圍較廣，但是小數點後的精準度較低。所以 BF16 的廣度較高，但精度較低，相較於 FP32 而言也能節省記憶體。

（圖源：NVIDIA Blog）

對於看似要求精準度的深度學習模型而言，使用 BF16 有什麼好處呢？在訓練結構較為複雜且參數量大的模型時，可能會產生比較大的梯度，甚至可能大到 FP16 放不下，因此產生 NaN 的 Loss 結果。如果使用 BF16 的話，就比較不容易發生這樣的情況。

話雖這麼說，但其實選擇 BF16 是個滿自然的過程：

1. FP32 塞不進 GPU 記憶體。
2. 8-Bit 以下不能直接訓練。
3. FP16 的 Loss 會 NaN 炸開。

那我們只剩下 BF16 可以用了 😢

註：在 1B 參數量的情況下，可以將 FP32 的模型權重放進一張 24GB 記憶體的 GPU 裡面，但是開始計算梯度時就會爆開了。因此將模型權重放進 GPU 裡面通常不是問題，處理梯度的記憶體消耗才是關鍵。

接下來讀取資料集：

import datasets

# 讀取資料集
data_files = {
    "train": "train.tokens.json.gz",
    "dev": "dev.tokens.json.gz",
}

dataset = datasets.load_dataset(
    "json",
    data_files=data_files,
    cache_dir="cache",
)

這會將資料集讀取成一個類似字典的物件，可以透過以下程式碼檢查內容：

for data in dataset["train"]:
    print(data["input_ids"])
    print(data["labels"])

訓練效果不佳時，不妨檢查看看資料內容是否正常。另外將 cache_dir 設定為 cache 時，會在工作目錄底下產生一個 cache 資料夾來存放快取檔案。有時候就算更新了資料集，HF Datasets 也會判定你沒改，導致訓練出問題。這時就可以手動把 cache 資料夾刪掉，來強制 HF Datasets 重新讀取。

接下來設定訓練參數：

from transformers import TrainingArguments

# 設定訓練參數
output_dir = "Models/TinyLlama-1B-TwAddr"
train_args = TrainingArguments(
    output_dir,
    per_device_train_batch_size=4,
    evaluation_strategy="epoch",
    bf16=True,
)

evaluation_strategy 指定每個 Epoch 訓練完後就進行一次評估。因為我們是用 BF16 訓練，所以 bf16=True 要打開。在這個簡單的玩具任務裡面，我們只需要設定這幾個參數即可。但其實還有相當多訓練參數可以設定，以下稍微介紹一些筆者常用到的參數。

per_device_train_batch_size 與 per_device_eval_batch_size 分別設定訓練與評估時的 Batch Size 大小，這會大幅影響記憶體的用量。通常單卡在訓練 LLM 時，Batch Size 都沒辦法開很大，這樣其實容易造成訓練效果不好。這時就需要搭配 eval_accumulation_steps 參數，可以設定每累積幾個 Step 的梯度再計算一次。假設 Batch Size 為 4 而 Accumulation Steps 為 8，那就會計算 4 * 8 = 32 筆訓練資料後再做一次反向梯度，這樣可以達到類似 Batch Size 32 的效果。

有時候一個 Epoch 需要花相當多 Steps 來進行，因此可以將 evaluation_strategy 設定為 steps 來提高評估的頻率，搭配 eval_steps 參數來決定幾個 Step 要做一次評估。

透過 save_strategy="steps" 與 save_steps 來指定每幾個 Step 要存一個檢查點。save_strategy 設定為 "epoch" 時則是每個 Epoch 都會存一個檢查點。啟用檢查點時，會在輸出目錄底下產生 checkpoint-25, checkpoint-50, checkpoint-75 之類的資料夾，每個資料夾裡面都會存一個完整的模型權重在裡面。如果訓練意外中斷時，可以從這個 Checkpoint 恢復訓練狀態。

除此之外，使用評估資料集的目的，除了監測模型有沒有訓練壞掉以外，也能用來尋找模型效果的最佳時間點。搭配參數 load_best_model_at_end 使用，可以在訓練結束後，將最佳評估結果的模型讀取出來。另外，為了避免硬碟被這些檢查點模型塞爆，可以設定 save_total_limit 來決定只保留最近的幾個檢查點。

因此最後的訓練參數可能如下：

train_args = TrainingArguments(
    output_dir,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    eval_accumulation_steps=2,
    evaluation_strategy="steps",
    save_strategy="steps",
    eval_steps=25,
    save_steps=25,
    save_total_limit=3,
    num_train_epochs=3,
    load_best_model_at_end=True,
    bf16=True,
)

參數決定好之後，就可以開始訓練模型了！

from transformers import Trainer

# 開始訓練模型
trainer = Trainer(
    model=model,
    args=train_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["dev"],
)
trainer.train()

# 儲存訓練完的模型
trainer.save_model()

# 也另外存一份 Tokenizer 方便評估
tk.save_pretrained(output_dir)

訓練完之後，別忘了把模型跟 Tokenizer 都存下來。在一張 RTX 3090 上訓練這份資料集，約三分鐘內就可以完成。

最終評估

最終評估是相當重要的一環，用來確保模型並沒有被訓練壞。這個環節分成固定測試集的評估，與自由測試的評估。我們可以借助 vLLM 的力量，對模型進行大規模的評估：

import json

from vllm import LLM, SamplingParams

from utils import build_prompt, iter_dataset

# 建立測試集的 Prompt 列表
prompts, items = list(), list()
for full, item in iter_dataset("test.json"):
    prompt = build_prompt(full)
    prompts.append(prompt)
    items.append(item)

# 讀取模型
model_name = "PY007/TinyLlama-1.1B-Chat-v0.3"
llm = LLM(model_name, dtype="float16")

# temperature 設為 0.0 為 Greedy Decode
# 確保每次實驗的結果都是一樣的
sampling_params = SamplingParams(
    max_tokens=512,
    temperature=0.0,
    stop=["}"],
)

# 對所有 Prompt 同時進行推論
outputs = llm.generate(prompts, sampling_params)

# 評估生成結果
results = list()
for out, item in zip(outputs, items):
    text = out.outputs[0].text

    # 嘗試解析模型的輸出
    try:
        begin = text.index("{")
        text = text[begin:] + "}"
        pred = json.loads(text)
    except:
        pred = None

    results.append(pred == item)

# 輸出準確率
accuracy = sum(results) / len(results)
print(f"Accuracy: {accuracy:.2%}")

首先測量原本模型的效果當作基準：

Accuracy: 7.20%

準確率只有 7% 而已，顯然模型完全無法理解這個任務。

接下來看看 Finetune 過的模型效果：

Accuracy: 97.40%

登登登，效果大幅提昇 🎉

這時我們可以試試看，如果將 Prompt Template 裡面的範例移除會怎樣呢？

template = """<|im_start|>user
請將以下路名解析為 JSON 格式。

輸入：{}
<|im_end|>
<|im_start|>assistant
{}"""

得到的效果：

Accuracy: 89.00%

居然也有將近九成的準確率，很顯然我們讓模型留下了一個刻板印象：如果要將路名解析成 JSON 格式，則必定是這種 city-town-road 的 Key-Value 組合，而不會是什麼city_name 之類的。這樣的結果到底是好還是壞呢？在一些比較 Aggressive 的 Domain-Specific Training 裡面，通常是不太在意 Out-Domain 的效能下降的問題，但也不能破壞的太誇張。因此我們能用 Free Try 的評估方式，來看看模型原本的能力是否還健在：

from vllm import LLM, SamplingParams

model_name = "Models/TinyLlama-1B-TwAddr"
llm = LLM(model_name)

sampling_params = SamplingParams(
    max_tokens=512,
    temperature=0.0,
    stop=["###"],
)

template = "<|im_start|>\nuser\n{}<|im_end|>\n<|im_start|>assistant\n"
while True:
    message = input(" > ")
    prompt = template.format(message)
    outputs = llm.generate(
        [prompt],
        sampling_params,
        use_tqdm=False,
    )
    print(outputs[0].outputs[0].text)

測試結果如下：

看起來模型原本的問答能力還保留著！

因為這個玩具任務的難度很低，所以有機會出現這種皆大歡喜的場面。多數情況下，某個能力成長了，必然有另外一個能力會消退。如何在有限的資源與成本內做抉擇，是訓練的人必須去權衡的事情。

結論

今天介紹了如何微調一個語言模型，雖然只是個小規模的實驗，但能讓我們更加熟悉資料處理的流程。而且訓練速度相對快，能縮短實驗的週期，並快速的嘗試各種不同的參數，也比較容易累積一些基礎的訓練經驗。

不過今天都在訓練 TinyLlama 1.1B 模型，但現在外面的 LLM 少說都 7B 起跳啊！為什麼不來訓練 7B 的模型呢？於是我們將模型名稱改為 TheBloke/Llama-2-7b-chat-fp16 並嘗試訓練：

torch.cuda.OutOfMemoryError: CUDA out of memory.

事實便是如此殘酷，今天介紹的訓練方法為 Full Fine-Tuning (FFT)，是最傳統的訓練方法，但也是訓練成本最高的方法。隨著訓練資料的長度越長，GPU 記憶體的消耗還會平方成長上去。即便這份路名資料集的長度僅約 120 Tokens 左右，在單張 24GB 的 GPU 上依然無法進行訓練，更不用說是長度一兩千以上的資料集了。

但是不用擔心，那個神秘的笑臉將再度出手，拯救單顯卡的貧民玩家們。

參考

• 政府資料開放平台
• TinyLlama 1.1B
• NVIDIA Blog
• HF Docs: Floating Data Types