从零到一构建一个小型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：数据预处理与表示

1. 文本数据准备:
   • 加载小型文本数据集（例如，某个小说或诗歌子集）。
   • 数据清洗与标准化: 转换为小写，去除标点符号，处理特殊字符等。
2. 分词器实现 (Simple Tokenizer):
   • 实现一个基于字符级 (Character-level) 或简单词级 (Simple Word-level) 的分词器，以简化复杂的分词算法（如BPE）的初期实现。
   • 构建词汇表 (Vocabulary) 和token到ID的映射 (Token-to-ID mapping)。
   • 实现encode和decode方法。
   • Padding和截断 (Padding and Truncation): 将所有输入序列统一到max_len。
3. Token嵌入层:
   • 实现一个PyTorch的nn.Embedding层，将token ID转换为d_model维的向量。
4. 位置编码层:
   • 实现正弦位置编码 (Sinusoidal Positional Encoding)，它可以在不知道序列最大长度的情况下推广到更长的序列。

**步骤2：注意力机制与Transformer块 **

1. 实现Scaled Dot-Product Attention:
   • 定义scaled_dot_product_attention(Q, K, V, mask)函数，包含矩阵乘法、缩放、掩码应用和softmax。
   • 关键是实现mask的应用: 在解码器中，未来信息是不可见的，所以需要一个下三角矩阵 (Lower Triangular Matrix) 形式的掩码，将未来token的注意力权重设为负无穷，使其在softmax后变为0。
2. 实现Multi-Head Attention:
   • 定义MultiHeadAttention模块，包含多个线性变换（用于Q, K, V的投影），将输入分为多个头，并行计算Scaled Dot-Product Attention，然后拼接结果，最后再通过一个线性投影。
3. 实现Feed-Forward Network (FFN):
   • 定义FeedForward模块，包含两个nn.Linear层和一个激活函数。
4. 实现Decoder Block:
   • 定义DecoderBlock模块，组合掩码多头自注意力层、FFN、层归一化和残差连接。
   • 注意残差连接和层归一化的正确顺序（例如，Post-LN或Pre-LN）。

步骤3：构建MiniGPT模型

1. 组合Decoder Blocks:
   • 在MiniGPT主模型类中，实例化Token嵌入层和位置编码层。
   • 堆叠多个DecoderBlock实例。
2. 实现输出层:
   • 一个nn.Linear层将最终解码器输出映射到词汇表维度。
   • 不需要显式Softmax层，因为CrossEntropyLoss在内部包含了Softmax。

步骤4：训练与优化

1. 数据加载器 (DataLoader):
   • 创建数据集类和数据加载器，批量处理输入序列和对应的目标序列（下一个token）。
2. 损失函数与优化器:
   • 使用nn.CrossEntropyLoss作为损失函数。
   • 选择torch.optim.AdamW作为优化器。
3. 训练循环 (Training Loop):
   • 实现一个基本的训练循环，包括前向传播、损失计算、反向传播和参数更新。
   • 因果语言建模任务: 输入序列X，目标是预测X的每个token的下一个token。例如，如果输入是"hello world"，模型会尝试从"hello"预测"world"，并从"hello world"预测下一个token。
4. 梯度裁剪 (Gradient Clipping):
   • 为防止梯度爆炸，在反向传播后应用梯度裁剪。

**步骤5：文本生成 (Inference) **

1. 实现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)

---

模型实现的挑战与考虑

要想继续深入，需要解决以下问题，：

1. 大规模数据处理: 如何高效地读取、预处理和批量化TB级的数据。
2. 分布式训练: 单个GPU无法承载大模型训练，需要数据并行、模型并行（张量并行、管道并行）等技术。
3. 内存优化: KV Cache优化、量化、混合精度训练等。
4. 模型评估: 除了损失值，还需要针对生成质量、忠实度、一致性等进行定性和定量评估。
5. 生产部署: 模型推理优化、模型服务框架的选择和使用。
6. 超参数调优: 系统化的超参数搜索策略（如网格搜索、随机搜索、贝叶斯优化）。

通过上述的MiniGPT设计和逐步实现过程，希望读者将能够从底层理解LLM的工作原理，为后续深入学习和构建更复杂的大模型打下坚实的基础。