【LLM】从零到一构建一个小型LLM--MiniGPT

原创已于 2025-07-17 10:41:35 修改 · 769 阅读

12 ·

CC 4.0 BY-SA版权

本文为博主原创文章，未经博主允许不得转载。

文章标签：

#pytorch #人工智能 #python

于 2025-07-17 10:40:59 首次发布

从零到一构建一个小型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维的向量。
位置编码层:
- 实现正弦位置编码 (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_weights

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.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 heads
        output = self.wo(output)
        return output, attn_weights

class 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-Attention
        attn_output, _ = self.self_attn(x, mask=tgt_mask)
        x = self.norm1(x + self.dropout1(attn_output)) # Add & Norm

        # Feed Forward
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout2(ffn_output)) # Add & Norm
        return 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_fill

    def 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 logits

    def generate(self, tokenizer, prompt, max_new_tokens):
        self.eval() # Set model to evaluation mode
        input_ids = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0) # Add batch dim

        for _ 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 logits
            last_token_logits = logits[:, -1, :] # (batch_size, vocab_size)
            
            # Simple greedy sampling: take the token with the highest probability
            next_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_tokens
            
        return 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 = 0
        num_batches = 0
        
        for 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 target
            
            if len(batch_input) < model.token_embedding.max_len + 1: # Ensure we have input and target
                continue
            
            # 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 batch
            batch_input = batch_input[:-1].unsqueeze(0).to(device) # Remove last token, add batch dim
            batch_target = batch_target.unsqueeze(0).to(device) # Add batch dim

            optimizer.zero_grad()
            
            logits = model(batch_input) # (batch_size, seq_len, vocab_size)
            
            # Reshape logits and targets for CrossEntropyLoss
            loss = criterion(logits.view(-1, logits.size(-1)), batch_target.view(-1))
            
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # Gradient clipping
            optimizer.step()
            
            total_loss += loss.item()
            num_batches += 1
        
        avg_loss = total_loss / num_batches
        print(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 longer
    text += "let's see if it can learn to complete sentences. " * 50

    tokenizer = SimpleTokenizer(text)
    
    # Model parameters
    vocab_size = tokenizer.vocab_size
    d_model = 128
    num_heads = 4
    num_layers = 2
    d_ff = d_model * 4
    max_len = 64 # Max sequence length for our MiniGPT
    
    # Adjust tokenizer's max_len based on model's max_len
    tokenizer.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 text
    print("\nGenerating text:")
    prompt = "hello wor"
    generated_text = model.generate(tokenizer, prompt, max_new_tokens=50)
    print(generated_text)