從零到一構建一個小型LLM (Small Language Model)暫時起名為MiniGPT。這個模型將專注于因果語言建模 (Causal Language Modeling),這是許多現代LLM(如GPT系列)的核心預訓練任務。
模型設計:
我們設計的模型是一個僅包含解碼器 (Decoder-only) 的Transformer架構,專注于生成式任務。這里將簡化其規模,以使其更易于理解和從頭實現,但保留核心的Transformer組件。
模型架構概覽
- 輸入層:
- Token嵌入 (Token Embeddings): 將輸入的離散token轉換為連續向量表示。
- 位置編碼 (Positional Encoding): 捕獲token在序列中的順序信息。我們將使用絕對位置編碼(例如,正弦位置編碼)。
- 解碼器塊 (Decoder Block): MiniGPT將由多個相同的解碼器塊堆疊而成。每個解碼器塊包含:
- 掩碼多頭自注意力層 (Masked Multi-Head Self-Attention Layer): 允許模型關注序列中當前及之前的token,以預測下一個token。這是解碼器獨有的關鍵部分,確保了生成過程的因果性。
- 層歸一化 (Layer Normalization): 在每個子層操作后應用,有助于訓練穩定。
- 前饋網絡 (Feed-Forward Network, FFN): 包含兩個線性變換和一個激活函數(例如ReLU或GELU),用于處理注意力層的輸出。
- 殘差連接 (Residual Connections): 每個子層(注意力層和FFN)的輸出都會與輸入相加,再進行歸一化,有助于解決深度網絡的梯度消失問題。
- 輸出層:
- 線性層 (Linear Layer): 將解碼器最終輸出的向量映射回詞匯表大小的維度。
- Softmax層 (Softmax Layer): 將線性層的輸出轉換為概率分布,表示下一個token是詞匯表中每個詞的概率。
模型參數設定(示例)
為了簡化實現,我們將采用較小的參數:
- 詞匯表大小 (Vocab Size): 例如,50000(覆蓋常見詞匯)。
- 嵌入維度 (Embedding Dimension,
d_model
): 例如,256。 - 序列最大長度 (Max Sequence Length,
max_len
): 例如,256。 - 解碼器塊數量 (Number of Decoder Blocks): 例如,4。
- 注意力頭數量 (Number of Attention Heads): 例如,4。
- 前饋網絡維度 (FFN Dimension): 例如,
d_model * 4
(1024)。
實現步驟
步驟1:數據預處理與表示
- 文本數據準備:
- 加載小型文本數據集(例如,某個小說或詩歌子集)。
- 數據清洗與標準化: 轉換為小寫,去除標點符號,處理特殊字符等。
- 分詞器實現 (Simple Tokenizer):
- 實現一個基于字符級 (Character-level) 或簡單詞級 (Simple Word-level) 的分詞器,以簡化復雜的分詞算法(如BPE)的初期實現。
- 構建詞匯表 (Vocabulary) 和token到ID的映射 (Token-to-ID mapping)。
- 實現
encode
和decode
方法。 - Padding和截斷 (Padding and Truncation): 將所有輸入序列統一到
max_len
。
- Token嵌入層:
- 實現一個PyTorch的
nn.Embedding
層,將token ID轉換為d_model
維的向量。
- 實現一個PyTorch的
- 位置編碼層:
- 實現正弦位置編碼 (Sinusoidal Positional Encoding),它可以在不知道序列最大長度的情況下推廣到更長的序列。
**步驟2:注意力機制與Transformer塊 **
- 實現Scaled Dot-Product Attention:
- 定義
scaled_dot_product_attention(Q, K, V, mask)
函數,包含矩陣乘法、縮放、掩碼應用和softmax。 - 關鍵是實現
mask
的應用: 在解碼器中,未來信息是不可見的,所以需要一個下三角矩陣 (Lower Triangular Matrix) 形式的掩碼,將未來token的注意力權重設為負無窮,使其在softmax后變為0。
- 定義
- 實現Multi-Head Attention:
- 定義
MultiHeadAttention
模塊,包含多個線性變換(用于Q, K, V的投影),將輸入分為多個頭,并行計算Scaled Dot-Product Attention,然后拼接結果,最后再通過一個線性投影。
- 定義
- 實現Feed-Forward Network (FFN):
- 定義
FeedForward
模塊,包含兩個nn.Linear
層和一個激活函數。
- 定義
- 實現Decoder Block:
- 定義
DecoderBlock
模塊,組合掩碼多頭自注意力層、FFN、層歸一化和殘差連接。 - 注意殘差連接和層歸一化的正確順序(例如,Post-LN或Pre-LN)。
- 定義
步驟3:構建MiniGPT模型
- 組合Decoder Blocks:
- 在
MiniGPT
主模型類中,實例化Token嵌入層和位置編碼層。 - 堆疊多個
DecoderBlock
實例。
- 在
- 實現輸出層:
- 一個
nn.Linear
層將最終解碼器輸出映射到詞匯表維度。 - 不需要顯式Softmax層,因為
CrossEntropyLoss
在內部包含了Softmax。
- 一個
步驟4:訓練與優化
- 數據加載器 (DataLoader):
- 創建數據集類和數據加載器,批量處理輸入序列和對應的目標序列(下一個token)。
- 損失函數與優化器:
- 使用
nn.CrossEntropyLoss
作為損失函數。 - 選擇
torch.optim.AdamW
作為優化器。
- 使用
- 訓練循環 (Training Loop):
- 實現一個基本的訓練循環,包括前向傳播、損失計算、反向傳播和參數更新。
- 因果語言建模任務: 輸入序列
X
,目標是預測X
的每個token的下一個token。例如,如果輸入是"hello world",模型會嘗試從"hello"預測"world",并從"hello world"預測下一個token。
- 梯度裁剪 (Gradient Clipping):
- 為防止梯度爆炸,在反向傳播后應用梯度裁剪。
**步驟5:文本生成 (Inference) **
- 實現
generate
方法:- 給定一個起始prompt,模型循環預測下一個token。
- 將預測的token添加到序列中,作為下一個時間步的輸入。
- 循環直到達到最大生成長度或生成結束符。
- 采樣策略: 可以實現簡單的貪婪采樣 (Greedy Sampling) 或溫度采樣 (Temperature Sampling)。
核心代碼 (代碼)
import torch
import torch.nn as nn
import torch.nn.functional as F
import math# --- 步驟1: 數據預處理與表示 ---
class SimpleTokenizer:def __init__(self, text):# 簡化版:從文本構建詞匯表self.vocab = sorted(list(set(text)))self.char_to_idx = {ch: i for i, ch in enumerate(self.vocab)}self.idx_to_char = {i: ch for i, ch in enumerate(self.vocab)}self.vocab_size = len(self.vocab)def encode(self, text):return [self.char_to_idx[ch] for ch in text if ch in self.char_to_idx]def decode(self, indices):return "".join([self.idx_to_char[idx] for idx in indices])class PositionalEncoding(nn.Module):def __init__(self, d_model, max_len=5000):super().__init__()pe = torch.zeros(max_len, d_model)position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))pe[:, 0::2] = torch.sin(position * div_term)pe[:, 1::2] = torch.cos(position * div_term)self.register_buffer('pe', pe.unsqueeze(0))def forward(self, x):# x: (batch_size, seq_len, d_model)return x + self.pe[:, :x.size(1)]# --- 步驟2: 注意力機制與Transformer塊 ---
def scaled_dot_product_attention(Q, K, V, mask=None):# Q, K, V: (..., seq_len, d_k)d_k = Q.size(-1)scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)if mask is not None:scores = scores.masked_fill(mask == 0, float('-inf'))attention_weights = F.softmax(scores, dim=-1)output = torch.matmul(attention_weights, V)return output, attention_weightsclass MultiHeadAttention(nn.Module):def __init__(self, d_model, num_heads):super().__init__()self.d_model = d_modelself.num_heads = num_headsself.d_k = d_model // num_headsself.wq = nn.Linear(d_model, d_model)self.wk = nn.Linear(d_model, d_model)self.wv = nn.Linear(d_model, d_model)self.wo = nn.Linear(d_model, d_model)def forward(self, x, mask=None):batch_size = x.size(0)Q = self.wq(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # (batch_size, num_heads, seq_len, d_k)K = self.wk(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)V = self.wv(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) # Concat headsoutput = self.wo(output)return output, attn_weightsclass FeedForward(nn.Module):def __init__(self, d_model, d_ff):super().__init__()self.linear1 = nn.Linear(d_model, d_ff)self.gelu = nn.GELU() # Or nn.ReLU()self.linear2 = nn.Linear(d_ff, d_model)def forward(self, x):return self.linear2(self.gelu(self.linear1(x)))class DecoderBlock(nn.Module):def __init__(self, d_model, num_heads, d_ff, dropout=0.1):super().__init__()self.self_attn = MultiHeadAttention(d_model, num_heads)self.ffn = FeedForward(d_model, d_ff)self.norm1 = nn.LayerNorm(d_model)self.norm2 = nn.LayerNorm(d_model)self.dropout1 = nn.Dropout(dropout)self.dropout2 = nn.Dropout(dropout)def forward(self, x, tgt_mask):# Masked Multi-Head Self-Attentionattn_output, _ = self.self_attn(x, mask=tgt_mask)x = self.norm1(x + self.dropout1(attn_output)) # Add & Norm# Feed Forwardffn_output = self.ffn(x)x = self.norm2(x + self.dropout2(ffn_output)) # Add & Normreturn x# --- 步驟3: 構建MiniGPT模型 ---
class MiniGPT(nn.Module):def __init__(self, vocab_size, d_model, num_heads, num_layers, d_ff, max_len, dropout=0.1):super().__init__()self.token_embedding = nn.Embedding(vocab_size, d_model)self.positional_encoding = PositionalEncoding(d_model, max_len)self.dropout = nn.Dropout(dropout)self.decoder_layers = nn.ModuleList([DecoderBlock(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])self.final_norm = nn.LayerNorm(d_model)self.output_linear = nn.Linear(d_model, vocab_size)def generate_square_subsequent_mask(self, sz):mask = torch.triu(torch.ones(sz, sz), diagonal=1).transpose(0, 1)return mask.bool() # Boolean mask for masked_filldef forward(self, src):# src: (batch_size, seq_len)seq_len = src.size(1)tgt_mask = self.generate_square_subsequent_mask(seq_len).to(src.device)x = self.token_embedding(src) # (batch_size, seq_len, d_model)x = self.positional_encoding(x)x = self.dropout(x)for decoder_layer in self.decoder_layers:x = decoder_layer(x, tgt_mask)x = self.final_norm(x)logits = self.output_linear(x) # (batch_size, seq_len, vocab_size)return logitsdef generate(self, tokenizer, prompt, max_new_tokens):self.eval() # Set model to evaluation modeinput_ids = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0) # Add batch dimfor _ in range(max_new_tokens):# If sequence length exceeds max_len, truncate (for simplicity)current_input_ids = input_ids if input_ids.size(1) <= tokenizer.max_len else input_ids[:, -tokenizer.max_len:]with torch.no_grad():logits = self(current_input_ids) # Get logits for the current sequence# Predict the next token based on the last token's logitslast_token_logits = logits[:, -1, :] # (batch_size, vocab_size)# Simple greedy sampling: take the token with the highest probabilitynext_token_id = torch.argmax(last_token_logits, dim=-1).unsqueeze(0) # (1, 1)input_ids = torch.cat((input_ids, next_token_id), dim=1)# Break if generated token is special end token (e.g., <EOS>)# For this simple example, we don't have explicit EOS, just max_new_tokensreturn tokenizer.decode(input_ids[0].tolist())# --- 步驟4: 訓練循環示例 ---
def train_minigpt(model, tokenizer, data, epochs, batch_size, learning_rate, device):optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)criterion = nn.CrossEntropyLoss()model.train() # Set model to training mode# Simplified data preparation for training# For a real scenario, you'd have a DataLoader and iterate batches# Here, we'll just process a long string into input/target pairs# Create input-target pairs for causal language modeling# Example: "hello world" -> input: "hello worl", target: "ello world"input_ids_full = torch.tensor(tokenizer.encode(data), dtype=torch.long)for epoch in range(epochs):total_loss = 0num_batches = 0for i in range(0, len(input_ids_full) - model.token_embedding.max_len, batch_size):batch_input = input_ids_full[i : i + model.token_embedding.max_len]batch_target = input_ids_full[i+1 : i + model.token_embedding.max_len + 1] # Shifted targetif len(batch_input) < model.token_embedding.max_len + 1: # Ensure we have input and targetcontinue# For simplicity, if batch_size is 1, and we're processing char by char for a small model# This needs to be adapted for proper batching with padding if sequence lengths vary.# Here, we're assuming fixed max_len for each input chunk.# Reshape for single batchbatch_input = batch_input[:-1].unsqueeze(0).to(device) # Remove last token, add batch dimbatch_target = batch_target.unsqueeze(0).to(device) # Add batch dimoptimizer.zero_grad()logits = model(batch_input) # (batch_size, seq_len, vocab_size)# Reshape logits and targets for CrossEntropyLossloss = criterion(logits.view(-1, logits.size(-1)), batch_target.view(-1))loss.backward()torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # Gradient clippingoptimizer.step()total_loss += loss.item()num_batches += 1avg_loss = total_loss / num_batchesprint(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")if __name__ == "__main__":# 示例用法text = "hello world, this is a test string for my minigpt model. " * 100 # Make it longertext += "let's see if it can learn to complete sentences. " * 50tokenizer = SimpleTokenizer(text)# Model parametersvocab_size = tokenizer.vocab_sized_model = 128num_heads = 4num_layers = 2d_ff = d_model * 4max_len = 64 # Max sequence length for our MiniGPT# Adjust tokenizer's max_len based on model's max_lentokenizer.max_len = max_len model = MiniGPT(vocab_size, d_model, num_heads, num_layers, d_ff, max_len)device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model.to(device)print(f"Using device: {device}")print(f"MiniGPT Model created with {sum(p.numel() for p in model.parameters() if p.requires_grad)} trainable parameters.")# Train the model (simplified)# This training setup is very basic for demonstration and would need significant improvements# for actual meaningful learning (e.g., proper dataset iteration, more epochs, larger data)print("\nStarting training...")train_minigpt(model, tokenizer, text, epochs=5, batch_size=16, learning_rate=1e-3, device=device)print("Training complete.")# Generate textprint("\nGenerating text:")prompt = "hello wor"generated_text = model.generate(tokenizer, prompt, max_new_tokens=50)print(generated_text)
模型實現的挑戰與考慮
要想繼續深入,需要解決以下問題,:
- 大規模數據處理: 如何高效地讀取、預處理和批量化TB級的數據。
- 分布式訓練: 單個GPU無法承載大模型訓練,需要數據并行、模型并行(張量并行、管道并行)等技術。
- 內存優化: KV Cache優化、量化、混合精度訓練等。
- 模型評估: 除了損失值,還需要針對生成質量、忠實度、一致性等進行定性和定量評估。
- 生產部署: 模型推理優化、模型服務框架的選擇和使用。
- 超參數調優: 系統化的超參數搜索策略(如網格搜索、隨機搜索、貝葉斯優化)。
通過上述的MiniGPT設計和逐步實現過程,希望讀者將能夠從底層理解LLM的工作原理,為后續深入學習和構建更復雜的大模型打下堅實的基礎。