大語言模型的訓練成本高昂，有效載入和使用預訓練模型權重至關重要。本文介紹瞭如何使用 PyTorch 載入和儲存模型的 state_dict，確保模型訓練的延續性。同時，詳細說明瞭如何從 OpenAI 載入預訓練的 GPT-2 模型權重，包含必要的環境設定、程式碼範例和步驟說明。此外，本文也探討了控制隨機性的解碼策略，包含 Top-k 取樣和溫度取樣，並提供了結合兩種策略的程式碼實作，讓開發者能更精確地控制生成文字的多樣性和準確性，並根據不同的應用場景調整引數。
5.3 控制隨機性的解碼策略
在生成文字的過程中，模型會根據輸入的上下文預測下一個詞彙的機率分佈。為了控制生成文字的多樣性和準確性，我們可以採用不同的解碼策略。
5.3.1 Top-k 取樣
Top-k 取樣是一種常見的解碼策略，旨在限制模型在預測下一個詞彙時的選擇範圍。具體來說，就是隻考慮機率最高的 k 個詞彙，而不是所有可能的詞彙。
首先，我們取得 logits 值最高的 top 三個 token 的 logits 值和 token ID，按照降序排列：
Top logits: tensor([6.7500, 6.2800, 4.5100])
Top positions: tensor([3, 7, 0])
接著，我們使用 PyTorch 的 where 函式，將不在 top 三個的 token 的 logit 值設為負無窮大（-inf）：
new_logits = torch.where(
condition=next_token_logits < top_logits[-1],
input=torch.tensor(float('-inf')),
other=next_token_logits
)
輸出結果如下：
tensor([4.5100, -inf, -inf, 6.7500, -inf, -inf, -inf, 6.2800, -inf])
最後，我們對這些 logit 值應用 softmax 函式，將它們轉換為機率：
topk_probas = torch.softmax(new_logits, dim=0)
得到的結果是三個非零機率值：
tensor([0.0615, 0.0000, 0.0000, 0.5775, 0.0000, 0.0000, 0.0000, 0.3610, 0.0000])
內容解密：
- 取得 Top Logits 和 Token ID：首先，我們根據模型的輸出 logits 值，找出機率最高的 top 三個 token 的 logits 值和對應的 token ID。
- 設定非 Top Token 的 Logits 為 -inf：使用 PyTorch 的
where函式，將不在 top 三個的 token 的 logit 值設為負無窮大，以確保它們在 softmax 後的機率為零。 - 應用 Softmax 函式：將處理後的 logits 值透過 softmax 函式轉換為機率，得到最終的機率分佈。
5.3.3 修改文字生成函式
現在，我們結合溫度取樣（temperature sampling）和 top-k 取樣，修改之前的 generate_text_simple 函式，建立一個新的 generate 函式。
def generate(model, idx, max_new_tokens, context_size,
temperature=0.0, top_k=None, eos_id=None):
for _ in range(max_new_tokens):
idx_cond = idx[:, -context_size:]
with torch.no_grad():
logits = model(idx_cond)
logits = logits[:, -1, :]
if top_k is not None:
top_logits, _ = torch.topk(logits, top_k)
min_val = top_logits[:, -1]
logits = torch.where(
logits < min_val,
torch.tensor(float('-inf')).to(logits.device),
logits
)
if temperature > 0.0:
logits = logits / temperature
probs = torch.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
else:
idx_next = torch.argmax(logits, dim=-1, keepdim=True)
if idx_next == eos_id:
break
idx = torch.cat((idx, idx_next), dim=1)
return idx
內容解密：
generate函式流程：該函式根據輸入的idx（token IDs）和模型，生成指定數量（max_new_tokens）的文字。- Top-k 取樣實作：如果指定了
top_k，則只保留機率最高的 k 個 token 的 logits 值，其他設為 -inf。 - 溫度取樣實作：如果溫度大於 0，則對 logits 值進行溫度縮放後，透過 softmax 取得機率分佈，並進行多項式取樣；否則，直接選擇機率最高的 token。
- 結束標記處理：如果生成的 token 是結束標記（
eos_id），則提前終止生成。
測試新的 generate 函式：
torch.manual_seed(123)
token_ids = generate(
model=model,
idx=text_to_token_ids("Every effort moves you", tokenizer),
max_new_tokens=15,
context_size=GPT_CONFIG_124M["context_length"],
top_k=25,
temperature=1.4
)
print("Output text:\n", token_ids_to_text(token_ids, tokenizer))
生成的文字如下：
Output text:
Every effort moves you stand to work on surprise, a one of us had gone with random-
相比之前的生成結果，這次的輸出更加多樣化。
練習題 5.2
嘗試不同的溫度和 top-k 設定，觀察生成的文字變化。思考在哪些應用場景下需要較低或較高的溫度和 top-k 設定。
5.4 在 PyTorch 中載入和儲存模型權重
為何儲存模型權重？
由於訓練大語言模型（LLM）非常耗時耗資源，因此儲存模型權重以便後續使用或繼續訓練至關重要。
如何儲存和載入模型權重？
PyTorch 提供簡單的方法來儲存和載入模型的 state_dict，即模型的權重引數。
儲存模型權重：
torch.save(model.state_dict(), "model.pth")
載入模型權重到新的 GPTModel 例項：
model.load_state_dict(torch.load("model.pth"))
內容解密：
- 儲存模型權重：使用
torch.save將模型的state_dict儲存到檔案中。 - 載入模型權重：使用
torch.load載入儲存的權重，並透過load_state_dict方法將權過載入到模型中。
練習題 5.3
思考如何設定 generate 函式的引數，以實作確定性的文字生成（即停用隨機取樣，使輸出固定）。
載入與使用預訓練的GPT-2模型權重
在前一章中，我們訓練了一個小型GPT-2模型，使用的是有限的資料集，即一本短篇故事書。這種方法使我們能夠專注於基本原理，而無需耗費大量時間和計算資源。
儲存與載入模型及最佳化器狀態
為了在稍後繼續預訓練模型，儲存最佳化器狀態是必要的。像AdamW這樣的自適應最佳化器會為每個模型權重儲存額外的引數。AdamW利用歷史資料動態調整每個模型引數的學習率。如果不儲存最佳化器狀態，最佳化器將重置，模型可能會學習得不理想，甚至無法正常收斂，這意味著它將失去生成連貫文字的能力。
torch.save({
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
}, "model_and_optimizer.pth")
然後，我們可以透過首先使用torch.load載入已儲存的資料，然後使用load_state_dict方法來還原模型和最佳化器的狀態：
checkpoint = torch.load("model_and_optimizer.pth", map_location=device)
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train()
內容解密：
- 儲存模型和最佳化器狀態：使用
torch.save儲存模型的狀態字典和最佳化器的狀態字典。 - 載入模型和最佳化器狀態：使用
torch.load載入已儲存的檢查點，然後分別載入模型和最佳化器的狀態字典。 - 設定模型為訓練模式：透過呼叫
model.train()，確保模型在載入狀態後處於訓練模式。
從OpenAI載入預訓練權重
幸運的是，OpenAI公開分享了他們的GPT-2模型的權重，從而消除了需要在大型語料函式庫上重新訓練模型的需要。首先，我們需要安裝TensorFlow和tqdm來載入GPT-2模型的權重：
pip install tensorflow>=2.15.0 tqdm>=4.66
下載GPT-2模型的權重
我們可以直接從本章的線上倉函式庫下載gpt_download.py Python模組：
import urllib.request
url = ("https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/ch05/01_main-chapter-code/gpt_download.py")
filename = url.split('/')[-1]
urllib.request.urlretrieve(url, filename)
然後，我們可以匯入download_and_load_gpt2函式，該函式將載入GPT-2架構設定（settings）和權重引數（params）：
from gpt_download import download_and_load_gpt2
settings, params = download_and_load_gpt2(model_size="124M", models_dir="gpt2")
內容解密：
- 安裝必要的函式庫：TensorFlow用於載入GPT-2模型的權重，tqdm用於顯示下載進度。
- 下載GPT-2模型的權重：透過執行
download_and_load_gpt2函式，下載並載入GPT-2模型的設定和權重。 - 檢查下載的內容：列印
settings和params的內容，以驗證下載是否成功。
檢視模型的設定與權重
下載完成後，我們可以檢視settings和params的內容：
print("Settings:", settings)
print("Parameter dictionary keys:", params.keys())
輸出結果如下：
Settings: {'n_vocab': 50257, 'n_ctx': 1024, 'n_embd': 768, 'n_head': 12, 'n_layer': 12}
Parameter dictionary keys: dict_keys(['blocks', 'b', 'g', 'wpe', 'wte'])
內容解密：
settings字典：包含了GPT-2模型的架構設定，如詞彙量、上下文視窗大小、嵌入維度等。params字典：包含了模型的實際權重張量，可以透過鍵值存取特定的權重。
載入OpenAI預訓練權重於GPT-2模型
在前面的章節中，我們已經下載並載入了最小的GPT-2模型權重，接下來我們需要將這些權重轉移到我們的GPTModel例項中。首先，我們需要定義一個字典來列出不同GPT模型大小之間的差異，如圖5.17所示。
GPT-2模型組態
GPT-2模型有多種不同的大小，包括124M、355M、774M和1558M。這些模型的架構相同，但嵌入維度和重複次數不同。
model_configs = {
"gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
"gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
"gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
"gpt2-xl (1558M)": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
}
更新GPTModel組態
假設我們要載入最小的模型"gpt2-small (124M)"，我們可以使用model_configs中的對應設定來更新我們的GPT_CONFIG_124M。
model_name = "gpt2-small (124M)"
NEW_CONFIG = GPT_CONFIG_124M.copy()
NEW_CONFIG.update(model_configs[model_name])
NEW_CONFIG.update({"context_length": 1024})
NEW_CONFIG.update({"qkv_bias": True})
組態更新說明：
- 上下文長度：原始GPT-2模型使用1024作為上下文長度，因此我們更新
NEW_CONFIG以比對此設定。 - qkv_bias：OpenAI的原始實作中使用了偏置向量，因此我們將
qkv_bias設為True以保持一致性。
初始化GPTModel例項
使用更新後的NEW_CONFIG，我們可以初始化一個新的GPTModel例項。
gpt = GPTModel(NEW_CONFIG)
gpt.eval()
載入權重
預設情況下，GPTModel例項會以隨機權重初始化。為了使用OpenAI的預訓練權重，我們需要覆寫這些隨機權重。
定義assign函式
首先，我們定義一個小型的assign函式，用於檢查兩個張量或陣列的維度是否相同，並傳回右側張量作為可訓練的PyTorch引數。
def assign(left, right):
if left.shape != right.shape:
raise ValueError(f"Shape mismatch. Left: {left.shape}, Right: {right.shape}")
return torch.nn.Parameter(torch.tensor(right))
載入權重函式
接下來，我們定義一個load_weights_into_gpt函式，用於將權重從params字典載入到GPTModel例項中。
import numpy as np
def load_weights_into_gpt(gpt, params):
gpt.pos_emb.weight = assign(gpt.pos_emb.weight, params['wpe'])
gpt.tok_emb.weight = assign(gpt.tok_emb.weight, params['wte'])
for b in range(len(params["blocks"])):
q_w, k_w, v_w = np.split((params["blocks"][b]["attn"]["c_attn"])["w"], 3, axis=-1)
gpt.trf_blocks[b].att.W_query.weight = assign(gpt.trf_blocks[b].att.W_query.weight, q_w.T)
gpt.trf_blocks[b].att.W_key.weight = assign(gpt.trf_blocks[b].att.W_key.weight, k_w.T)
gpt.trf_blocks[b].att.W_value.weight = assign(gpt.trf_blocks[b].att.W_value.weight, v_w.T)
q_b, k_b, v_b = np.split((params["blocks"][b]["attn"]["c_attn"])["b"], 3, axis=-1)
gpt.trf_blocks[b].att.W_query.bias = assign(gpt.trf_blocks[b].att.W_query.bias, q_b)
gpt.trf_blocks[b].att.W_key.bias = assign(gpt.trf_blocks[b].att.W_key.bias, k_b)
gpt.trf_blocks[b].att.W_value.bias = assign(gpt.trf_blocks[b].att.W_value.bias, v_b)
gpt.trf_blocks[b].att.out_proj.weight = assign(gpt.trf_blocks[b].att.out_proj.weight, params["blocks"][b]["attn"]["c_proj"]["w"].T)
載入權重說明：
- 位置嵌入和標記嵌入：首先，我們將位置嵌入和標記嵌入的權過載入到GPTModel中。
- 多頭注意力機制：對於每個變壓器區塊，我們將查詢、鍵和值的權重和偏置載入到相應的層中。
- 輸出投影：最後，我們載入輸出投影層的權重。
透過這些步驟，我們成功地將OpenAI的預訓練權過載入到我們的GPTModel例項中，使其能夠利用預訓練模型的知識進行進一步的微調或推斷。