软件公司网站 手机版建博客网站

基于GitHub项目：https://github.com/datawhalechina/llms-from-scratch-cn

设计 LLM 的架构

GPT 模型基于 Transformer 的 decoder-only 架构，其主要特点包括：

• 顺序生成文本

• 参数数量庞大（而非代码量复杂）

• 大量重复的模块化组件

以 GPT-2 small 模型（124M 参数）为例，其配置如下：

GPT_CONFIG_124M = {"vocab_size": 50257,  # BPE 分词器词表大小"ctx_len": 1024,      # 最大上下文长度"emb_dim": 768,       # 嵌入维度"n_heads": 12,        # 注意力头数量"n_layers": 12,       # Transformer 块层数"drop_rate": 0.1,     # Dropout 比例"qkv_bias": False     # QKV 计算是否使用偏置
}

GPT 模型基本结构

cfg是配置实例

import torch.nn as nnclass GPTModel(nn.Module):def __init__(self, cfg):super().__init__()# Token 嵌入层self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])# 位置嵌入层self.pos_emb = nn.Embedding(cfg["ctx_len"], cfg["emb_dim"])# Dropout 层self.drop_emb = nn.Dropout(cfg["drop_rate"])# 堆叠n_layers相同的Transformer 块self.trf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])# 最终层归一化self.final_norm = LayerNorm(cfg["emb_dim"])# 输出层self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)def forward(self, in_idx):batch_size, seq_len = in_idx.shape# Token 嵌入tok_embeds = self.tok_emb(in_idx)# 位置嵌入pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))# 组合嵌入x = tok_embeds + pos_embedsx = self.drop_emb(x)# 通过 Transformer 块x = self.trf_blocks(x)# 最终归一化x = self.final_norm(x)# 输出 logitslogits = self.out_head(x)return logits

 层归一化 (Layer Normalization)

层归一化将激活值规范化为均值为 0、方差为 1 的分布，加速模型收敛：

class LayerNorm(nn.Module):def __init__(self, emb_dim):super().__init__()self.eps = 1e-5    # 防止除零错误的标准设定值self.scale = nn.Parameter(torch.ones(emb_dim))  #可学习缩放参数，初始化为全1向量self.shift = nn.Parameter(torch.zeros(emb_dim)) #可学习平移参数，初始化为全0向量def forward(self, x):mean = x.mean(dim=-1, keepdim=True)    #计算均值 μ，沿最后一维，保持维度var = x.var(dim=-1, keepdim=True, unbiased=False)    #计算方差 σ²，同均值维度，有偏估计（分母n）norm_x = (x - mean) / torch.sqrt(var + self.eps)    #标准化计算，分母添加ε防溢出return self.scale * norm_x + self.shift    #仿射变换，恢复模型表达能力

GELU 激活函数与前馈网络

GPT 使用 GELU（高斯误差线性单元）激活函数：

场景	ReLU 的行为	GELU 的行为
处理微弱负信号	直接丢弃（可能丢失细节）	部分保留（如：保留 30% 的信号强度）
遇到强烈正信号	完全放行	几乎完全放行（保留 95% 以上）
训练稳定性	容易在临界点卡顿	平滑过渡，减少训练震荡
应对复杂模式	需要堆叠更多层数	单层就能捕捉更细腻的变化

class GELU(nn.Module):def __init__(self):super().__init__()def forward(self, x):return 0.5 * x * (1 + torch.tanh(torch.sqrt(torch.tensor(2.0 / torch.pi)) * (x + 0.044715 * torch.pow(x, 3))))

前馈神经网络实现：

class FeedForward(nn.Module):def __init__(self, cfg):super().__init__()self.layers = nn.Sequential(nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),GELU(),nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),nn.Dropout(cfg["drop_rate"]))def forward(self, x):return self.layers(x)

Shortcut 连接

Shortcut 连接（残差连接）解决深度网络中的梯度消失问题：

class TransformerBlock(nn.Module):def __init__(self, cfg):super().__init__()self.att = MultiHeadAttention(d_in=cfg["emb_dim"],d_out=cfg["emb_dim"],block_size=cfg["ctx_len"],num_heads=cfg["n_heads"], dropout=cfg["drop_rate"],qkv_bias=cfg["qkv_bias"])self.ff = FeedForward(cfg)self.norm1 = LayerNorm(cfg["emb_dim"])self.norm2 = LayerNorm(cfg["emb_dim"])self.drop_resid = nn.Dropout(cfg["drop_rate"])def forward(self, x):# 注意力块的残差连接shortcut = xx = self.norm1(x)x = self.att(x)x = self.drop_resid(x)x = x + shortcut# 前馈网络的残差连接shortcut = xx = self.norm2(x)x = self.ff(x)x = self.drop_resid(x)x = x + shortcutreturn x

Transformer 块整合

将多头注意力与前馈网络整合为 Transformer 块：

class TransformerBlock(nn.Module):def __init__(self, cfg):super().__init__()self.att = MultiHeadAttention(d_in=cfg["emb_dim"],d_out=cfg["emb_dim"],block_size=cfg["ctx_len"],num_heads=cfg["n_heads"], dropout=cfg["drop_rate"],qkv_bias=cfg["qkv_bias"])self.ff = FeedForward(cfg)self.norm1 = LayerNorm(cfg["emb_dim"])self.norm2 = LayerNorm(cfg["emb_dim"])self.drop_resid = nn.Dropout(cfg["drop_rate"])def forward(self, x):# 注意力块的残差连接shortcut = xx = self.norm1(x)x = self.att(x)x = self.drop_resid(x)x = x + shortcut# 前馈网络的残差连接shortcut = xx = self.norm2(x)x = self.ff(x)x = self.drop_resid(x)x = x + shortcutreturn x

完整 GPT 模型实现

class GPTModel(nn.Module):def __init__(self, cfg):super().__init__()self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])self.pos_emb = nn.Embedding(cfg["ctx_len"], cfg["emb_dim"])self.drop_emb = nn.Dropout(cfg["drop_rate"])self.trf_blocks = nn.Sequential(*[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])self.final_norm = LayerNorm(cfg["emb_dim"])self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)def forward(self, in_idx):batch_size, seq_len = in_idx.shapetok_embeds = self.tok_emb(in_idx)pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))x = tok_embeds + pos_embedsx = self.drop_emb(x)x = self.trf_blocks(x)x = self.final_norm(x)logits = self.out_head(x)return logits

文本生成

使用贪婪解码生成文本：

def generate_text_simple(model, idx, max_new_tokens, context_size):for _ in range(max_new_tokens):# 截断超过上下文长度的部分idx_cond = idx[:, -context_size:]with torch.no_grad():logits = model(idx_cond)# 获取最后一个 token 的 logitslogits = logits[:, -1, :]  probas = torch.softmax(logits, dim=-1)idx_next = torch.argmax(probas, dim=-1, keepdim=True)idx = torch.cat((idx, idx_next), dim=1)return idx

使用示例：

# 初始化模型
model = GPTModel(GPT_CONFIG_124M)# 设置评估模式
model.eval()# 生成文本
start_context = "Every effort moves you"
encoded = tokenizer.encode(start_context)
encoded_tensor = torch.tensor(encoded).unsqueeze(0)generated = generate_text_simple(model=model,idx=encoded_tensor,max_new_tokens=10,context_size=GPT_CONFIG_124M["ctx_len"]
)decoded_text = tokenizer.decode(generated.squeeze(0).tolist())
print(decoded_text)