1. 前言

使用梯度下降算法通過下一個token預測任務預訓練大語言模型GPTModel，前向傳播流程每次會輸入一個batch的長度均為context_len的訓練樣本，執行$\text{batch\_size}\times \text{context\_len}$次下一個token預測任務，共預測輸出$\text{batch\_size}\times \text{context\_len}$個tokens。后向傳播流程首先會使用交叉熵(Cross Entropy)損失函數計算大語言模型GPTModel的預測輸出與訓練樣本標簽之間的損失(loss)，再通過后向傳播算法計算大語言模型參數梯度，最后使用梯度下降算法更新大語言模型的參數。

本文使用交叉熵損失函數計算生成大語言模型GPTModel的預測輸出與訓練樣本標簽之間的loss，介紹大語言模型的預訓練流程，并實現預訓練大語言模型的函數pretrain_model。

2. 損失函數Cross Entropy

交叉熵損失函數可以度量大語言模型的預測輸出與訓練樣本標簽之間的差異。損失函數計算生成的loss值越大，表明大語言模型的預測輸出與訓練樣本標簽之間的差異越大，loss值越小，表明大語言模型的預測輸出與訓練樣本標簽之間的差異越小。

對輸入文本做tokenization，將輸入文本轉換成包含context_len個token ID的列表，并輸入大語言模型GPTModel，可以得到context_len個維度為vocabulary_size的logits向量，第$i$個logits向量是大語言模型根據前$i$個token預測生成的下一個token的概率分數向量，logits向量中的第$k$個概率分數值越大，表明大語言模型預測生成的下一個token的ID為$k$的概率越高。

使用softmax函數將大語言模型預測生成的logits向量歸一化，得到大語言模型預測生成的下一個token的概率分布，概率分布中對應樣本標簽位置的概率值表示大語言模型預測輸出的token為相應訓練樣本標簽的概率。對應樣本標簽位置的概率值越接近1，表明大語言模型預測輸出的token為相應訓練樣本標簽的概率越高，大語言模型的預測輸出與訓練樣本標簽之間的差異越小。

使用梯度下降算法預訓練大語言模型GPTModel的前向傳播流程中，大語言模型每次會預測生成$\text{batch\_size}\times \text{context\_len}$個下一個token的概率分布。如下圖所示，交叉熵損失函數會分別獲取$\text{batch\_size}\times \text{context\_len}$個概率分布中對應樣本標簽位置的概率值，使用對數函數計算這些概率值的對數，并計算所有對數值的均值，最后將對數均值的相反數作為大語言模型GPTModel的預測輸出與訓練樣本標簽之間的損失loss。

如下面的代碼所示，使用torch.tensor函數創建訓練樣本inputs及訓練樣本標簽targets，將訓練樣本inputs輸入大語言模型gpt2_small，并使用softmax函數將大語言模型的輸出張量logits歸一化，得到$2\times 3$個下一個token的概率分布，其中每個概率分布的維度均等于詞匯表的大小50257。分別獲取$2\times 3$個下一個token的概率分布中對應樣本標簽位置的概率值，使用torch.log函數計算這些概率值的對數，并計算所有對數值均值的相反數，可以得到大語言模型gpt2_small的預測輸出與樣本標簽targets之間的交叉熵損失：

import torch
# from [從零開始實現大語言模型（七）：多頭注意力機制] import MultiHeadAttention
# from [從零開始實現大語言模型（八）：Layer Normalization] import LayerNorm
# from [從零開始實現大語言模型（九）：前饋神經網絡與GELU激活函數] import GELU, FeedForward
# from [從零開始實現大語言模型（十一）：構建大語言模型GPTModel] import TransformerBlock, GPTModeltorch.manual_seed(123)embedding_dim = 768
num_layers = 12
num_heads = 12
context_len = 1024
vocabulary_size = 50257
dropout = 0.1
qkv_bias = Falsegpt2_small = GPTModel(embedding_dim=embedding_dim,num_layers=num_layers,num_heads=num_heads,context_len=context_len,vocabulary_size=vocabulary_size,dropout=dropout,qkv_bias=qkv_bias
)inputs = torch.tensor([[16833, 3626, 6100],  # [["every effort moves"],[40, 1107, 588]]      # ["I really like"]]
)targets = torch.tensor([[3626, 6100, 345],  # [[" effort moves you"],[588, 428, 11311]]  # [" really like chocolate"]]
)with torch.no_grad():logits = gpt2_small(inputs)
probas = torch.softmax(logits, dim=-1)target_probas_1 = probas[0, [0, 1, 2], targets[0]]
target_probas_2 = probas[1, [0, 1, 2], targets[1]]log_probas = torch.log(torch.cat((target_probas_1, target_probas_2)))
avg_log_probas = torch.mean(log_probas)
neg_avg_log_probas = avg_log_probas * -1print("probas shape:", probas.shape)
print("target_probas_1:", target_probas_1)
print("target_probas_2:", target_probas_2)
print("log_probas:", log_probas)
print("avg_log_probas:", avg_log_probas)
print("cross entropy loss:", neg_avg_log_probas)

執行上面代碼，打印結果如下：

probas shape: torch.Size([2, 3, 50257])
target_probas_1: tensor([2.6369e-05, 1.5997e-05, 1.6926e-05])
target_probas_2: tensor([1.5638e-05, 8.9422e-06, 1.7967e-05])
log_probas: tensor([-10.5433, -11.0431, -10.9867, -11.0658, -11.6247, -10.9270])
avg_log_probas: tensor(-11.0318)
cross entropy loss: tensor(11.0318)

可以直接使用PyTorch內置的cross_entropy函數執行上述計算流程，得到大語言模型gpt2_small的預測輸出logits與樣本標簽targets之間的交叉熵損失：

loss = torch.nn.functional.cross_entropy(logits.flatten(0, 1), targets.flatten())
print(loss)

執行上面代碼，打印結果如下：

tensor(11.0318)

根據打印結果可知，上述6個步驟計算得到的損失值與PyTorch內置的cross_entropy函數計算得到的交叉熵損失完全相同。交叉熵損失本質上就是大語言模型預測生成的$\text{batch\_size}\times \text{context\_len}$個下一個token的概率分布中對應樣本標簽位置概率值的對數均值的相反數。

3. 大語言模型預訓練流程

預訓練大語言模型的流程與訓練普通神經網絡模型本質上并沒有任何不同。如下圖所示，預訓練大語言模型可以把整個訓練數據集掃幾個epoch，每個epoch會把整個訓練數據集掃一遍，每次會使用訓練數據集中一個batch的訓練樣本訓練一次大語言模型。前向傳播流程會將一個batch的訓練樣本輸入大語言模型，得到大語言模型的預測輸出logits。后向傳播流程首先會使用交叉熵損失函數計算大語言模型的預測輸出logits與訓練樣本標簽targets之間的損失loss，再通過后向傳播算法計算大語言模型參數梯度，最后使用梯度下降算法更新大語言模型的參數。

可以使用如下代碼定義計算一個batch樣本數據交叉熵損失的函數calc_loss_batch，以及計算整個數據集上所有樣本數據交叉熵損失的函數calc_loss_loader。函數calc_loss_batch將一個batch的樣本數據輸入大語言模型，得到大語言模型的預測輸出logits，并使用torch.nn.functional.cross_entropy函數計算大語言模型的預測輸出logits與樣本標簽targets之間的損失loss。函數calc_loss_loader每次取數據集中一個batch的樣本數據，使用calc_loss_batch函數計算該batch樣本數據的交叉熵損失，并返回數據集上所有樣本數據損失loss的均值：

def calc_loss_batch(input_batch, target_batch, model, device):input_batch = input_batch.to(device)target_batch = target_batch.to(device)logits = model(input_batch)loss = torch.nn.functional.cross_entropy(logits.flatten(0, 1), target_batch.flatten())return lossdef calc_loss_loader(data_loader, model, device, num_batches=None):model.eval()total_loss = 0.0if num_batches is None:num_batches = len(data_loader)else:num_batches = min(num_batches, len(data_loader))with torch.no_grad():for i, (input_batch, target_batch) in enumerate(data_loader):if i < num_batches:loss = calc_loss_batch(input_batch, target_batch, model, device)total_loss += loss.item()else:breakreturn total_loss / num_batches

實現預訓練大語言模型的函數pretrain_model，可以使用for循環將整個訓練數據集掃num_epochs遍，并在每次訓練大語言模型的循環中，首先使用optimizer.zero_grad函數將大語言模型所有參數的梯度置為0，然后使用函數calc_loss_batch計算一個batch訓練樣本的交叉熵損失loss。使用loss.backward函數可以執行后向傳播流程，計算大語言模型所有參數的梯度，并通過optimizer.step函數使用梯度下降算法更新大語言模型參數。具體代碼如下所示：

# from [從零開始實現大語言模型（十二）：文本生成策略] import generate_textdef pretrain_model(model, optimizer, train_loader, num_epochs, device,eval_freq, eval_iter, tokenizer, start_context,save_freq, checkpoint_dir, checkpoint=None, val_loader=None
):if not os.path.exists(checkpoint_dir):os.makedirs(checkpoint_dir, exist_ok=True)if checkpoint is not None:model_checkpoint_path = os.path.join(checkpoint_dir, f"model_{checkpoint:06d}.pth")optimizer_checkpoint_path = os.path.join(checkpoint_dir, f"optimizer_{checkpoint:06d}.pth")model.load_state_dict(torch.load(model_checkpoint_path))optimizer.load_state_dict(torch.load(optimizer_checkpoint_path))else:checkpoint = -1train_losses, val_losses, track_tokens_seen = [], [], []tokens_seen, global_step = 0, -1for epoch in range(num_epochs):model.train()for i, (input_batch, target_batch) in enumerate(train_loader):if global_step % eval_freq == 0:model.train()optimizer.zero_grad()loss = calc_loss_batch(input_batch, target_batch, model, device)loss.backward()optimizer.step()tokens_seen += input_batch.numel()global_step += 1print(f"Epoch {epoch + 1} (Batch {i:06d}): Train loss {loss.item():.3f}")checkpoint, train_loss, val_loss = val_and_save(model, optimizer, train_loader, val_loader, epoch, global_step, eval_freq,eval_iter, start_context, tokenizer, save_freq, checkpoint_dir, checkpoint, device)if train_loss is not None:train_losses.append(train_loss)val_losses.append(val_loss)track_tokens_seen.append(tokens_seen)checkpoint, _, _ = val_and_save(model, optimizer, train_loader, val_loader, epoch, global_step, 1,eval_iter, start_context, tokenizer, 1, checkpoint_dir, checkpoint, device)print(f"Epoch {epoch + 1} finished, checkpoint: {checkpoint:06d}")return train_losses, val_losses, track_tokens_seendef val_and_save(model, optimizer, train_loader, val_loader, epoch, global_step, eval_freq,eval_iter, start_context, tokenizer, save_freq, checkpoint_dir, checkpoint, device
):train_loss, val_loss = None, Noneif global_step % eval_freq == 0:if val_loader is not None:train_loss = calc_loss_loader(train_loader, model, device, eval_iter)val_loss = calc_loss_loader(val_loader, model, device, eval_iter)print(f"Epoch {epoch + 1} (Step {global_step:06d}): Train loss {train_loss:.3f}, Val loss {val_loss:.3f}")generated_sample_text = generate_text(model, start_context, max_new_tokens=50, tokenizer=tokenizer,context_size=model.pos_emb.weight.shape[0], top_k=1, compact_format=True)print(f"Generated Sample Text: {generated_sample_text}")print("=====================================================================")if global_step % save_freq == 0:checkpoint += 1model_checkpoint_path = os.path.join(checkpoint_dir, f"model_{checkpoint:06d}.pth")optimizer_checkpoint_path = os.path.join(checkpoint_dir, f"optimizer_{checkpoint:06d}.pth")torch.save(model.state_dict(), model_checkpoint_path)torch.save(optimizer.state_dict(), optimizer_checkpoint_path)return checkpoint, train_loss, val_loss

PyTorch中神經網絡模型的state_dict是一個字典對象，字典中的key為神經網絡模型中參數的名稱，value為相應的參數。使用.state_dict函數可以一次性獲取神經網絡模型中的所有參數，并通過torch.save函數將所有參數保存為一個checkpoint。torch.load函數可以讀取指定checkpoint，通過.load_state_dict函數可以將神經網絡模型中的參數修改為checkpoint中的記錄值。

所有具有自適應能力的優化器（如AdamW可以根據歷史梯度信息動態調整學習率）都需要記錄每個神經網絡參數的歷史梯度等信息，同樣可以使用.state_dict一次性獲取優化器中的所有數據記錄，以及通過.load_state_dict函數從指定checkpoint中還原這些記錄數據。

如下面的代碼所示，使用從零開始實現大語言模型（二）：文本數據處理中構建的Dataset創建訓練集train_dataset及驗證集val_dataset，并通過PyTorch內置的torch.utils.data.DataLoader類創建訓練集及驗證集對應的DataLoader。使用torch.optim.AdamW實例化訓練大語言模型的優化器optimizer，最后使用函數pretrain_model預訓練大語言模型gpt2_small：

import os
import random
import tiktoken
from torch.utils.data import Dataset, DataLoader# from [從零開始實現大語言模型（二）：文本數據處理] import LLMDatasettrain_data_path = "train_data"
val_data_path = "val_data"
vocabulary = "gpt2"
special_token_id = 50256
context_len = 1024
stride = 1024
batch_size = 2num_epochs = 10
eval_freq = 5
eval_iter = 1
save_freq = 5
checkpoint_dir = "checkpoint"
start_context = "蕭炎，斗之力，三段"
tokenizer = tiktoken.encoding_for_model(vocabulary)
device = torch.device("cpu")
gpt2_small.to(device)
optimizer = torch.optim.AdamW(gpt2_small.parameters(), lr=0.0006, weight_decay=0.1)train_dataset = LLMDataset(train_data_path, vocabulary, special_token_id, context_len, stride)
val_dataset = LLMDataset(val_data_path, vocabulary, special_token_id, context_len, stride)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False, drop_last=False)
print(f"train_loader len: {len(train_loader)}")train_losses, val_losses, tokens_seen = pretrain_model(gpt2_small, optimizer, train_loader, num_epochs, device,eval_freq, eval_iter, tokenizer, start_context,save_freq, checkpoint_dir, val_loader=val_loader
)

執行上面代碼，打印結果如下：

train_loader len: 7
Epoch 1 (Batch 000000): Train loss 11.034
Epoch 1 (Step 000000): Train loss 9.827, Val loss 9.784
Generated Sample Text: 蕭炎，斗之力，三段 Knowledge�緦�緦緦�703 clashes�緦 longest，，緦，���，緦緦�
=====================================================================
Epoch 1 (Batch 000001): Train loss 9.940
Epoch 1 (Batch 000002): Train loss 8.811
Epoch 1 (Batch 000003): Train loss 7.954
Epoch 1 (Batch 000004): Train loss 7.286
Epoch 1 (Batch 000005): Train loss 6.629
Epoch 1 (Step 000005): Train loss 5.980, Val loss 6.003
Generated Sample Text: 蕭炎，斗之力，三段，，，，，，，，，，，，，，，，�
=====================================================================
Epoch 1 (Batch 000006): Train loss 6.027
Epoch 1 (Step 000006): Train loss 5.390, Val loss 5.479
Generated Sample Text: 蕭炎，斗之力，三段，，，�，，，��，，，，，�，�，，，
=====================================================================
Epoch 1 finished, checkpoint: 000002
Epoch 2 (Batch 000000): Train loss 5.401
Epoch 2 (Batch 000001): Train loss 5.028
Epoch 2 (Batch 000002): Train loss 4.788
Epoch 2 (Batch 000003): Train loss 4.616
Epoch 2 (Step 000010): Train loss 4.511, Val loss 4.526
Generated Sample Text: 蕭炎，斗之力，三段，�，�，�，�，�，�，�，�，�，��，�，
=====================================================================[...]Epoch 9 (Step 000060): Train loss 2.561, Val loss 3.470
Generated Sample Text: 蕭炎，斗之力，三段���是在臉�的�，�炣�殸廢是蕭炣也是曰�，蕭�
=====================================================================
Epoch 9 (Batch 000005): Train loss 2.560
Epoch 9 (Batch 000006): Train loss 2.558
Epoch 9 (Step 000062): Train loss 2.456, Val loss 3.455
Generated Sample Text: 蕭炎，斗之力，三段���，臉庿，炎�，蕭炎蕭�炎�蕭�，蕭�的�
=====================================================================
Epoch 9 finished, checkpoint: 000021
Epoch 10 (Batch 000000): Train loss 2.525
Epoch 10 (Batch 000001): Train loss 2.388
Epoch 10 (Batch 000002): Train loss 2.663
Epoch 10 (Step 000065): Train loss 2.270, Val loss 3.468
Generated Sample Text: 蕭炎，斗之力，三段��技蕭�的蕭炣也�，蕭�詎��更中著曰蕭�著�
=====================================================================
Epoch 10 (Batch 000003): Train loss 2.464
Epoch 10 (Batch 000004): Train loss 2.602
Epoch 10 (Batch 000005): Train loss 2.511
Epoch 10 (Batch 000006): Train loss 2.557
Epoch 10 (Step 000069): Train loss 2.117, Val loss 3.474
Generated Sample Text: 蕭炎，斗之力，三段��，這的�法的蕭�煉�蕭�法，蕭�級級父了�
=====================================================================
Epoch 10 finished, checkpoint: 000023

從上面的打印結果可知，使用梯度下降算法訓練大語言模型gpt2_small，可以減小大語言模型的預測輸出與樣本標簽之間的交叉熵損失，并顯著提升大語言模型的文本生成能力。在訓練剛開始時，將蕭炎，斗之力，三段輸入大語言模型gpt2_small，生成的是Knowledge�緦�緦緦�703 clashes�緦 longest，，緦，���，緦緦�或者，，，，，，，，，，，，，，，，�這樣不包含任何有效信息的自然語言文本序列。在僅包含7個batch訓練樣本的數據集上訓練10個epoch，大語言模型gpt2_small已經可以生成���，臉庿，炎�，蕭炎蕭�炎�蕭�，蕭�的�以及��，這的�法的蕭�煉�蕭�法，蕭�級級父了�這樣與訓練數據集存在一定關聯的自然語言文本了。

可以使用如下代碼，分別繪制大語言模型gpt2_small在訓練集及驗證集上交叉熵損失的變化情況圖像：

import matplotlib.pyplot as pltdef plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):fig, ax1 = plt.subplots(figsize=(5, 3))ax1.plot(epochs_seen, train_losses, label="Training loss")ax1.plot(epochs_seen, val_losses, linestyle="-.", label="Validation loss")ax1.set_xlabel("Epochs")ax1.set_ylabel("Loss")ax1.legend(loc="upper right")ax2 = ax1.twiny()ax2.plot(tokens_seen, train_losses, alpha=0)ax2.set_xlabel("Tokens seen")fig.tight_layout()plt.show()epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))
plot_losses(epochs_tensor, tokens_seen, train_losses, val_losses)

執行上面代碼，生成交叉熵損失的變化情況圖像如下：

從上面的交叉熵損失變化情況圖像可知，在訓練剛開始時，訓練集及驗證集上的交叉熵損失都非常大。使用梯度下降算法訓練大語言模型gpt2_small，可以減小大語言模型的預測輸出與樣本標簽之間的交叉熵損失，使大語言模型的預測輸出與樣本標簽之間的差異性更小。

隨著訓練的進行，訓練集和驗證集上交叉熵損失的差異會越來越大，訓練集上的交叉熵損失值會比驗證集小的越來越明顯，表明大語言模型在訓練數據集上的過擬合情況越來越嚴重。在工業界的預訓練大語言模型實踐中，并不會在一個很小的訓練數據集上訓練多個epoch，而是會在一個非常大的訓練數據集上訓練少數幾個甚至只訓練一個epoch，這種訓練策略可以很大程度上解決預訓練大語言模型時的過擬合問題。

4. 結束語

前向傳播流程將一個batch的訓練樣本輸入大語言模型，共預測輸出$\text{batch\_size}\times \text{context\_len}$個維度為vocabulary_size的logits向量。后向傳播流程首先使用交叉熵損失函數計算大語言模型的預測輸出與訓練樣本標簽之間的損失loss，并通過后向傳播算法計算大語言模型參數梯度，最后使用梯度下降算法更新大語言模型的參數。

預訓練大語言模型就是不斷從訓練數據集中獲取一個batch的訓練樣本，然后執行這個操作直至收斂的過程。