尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

从零手写GPT:斯坦福CS336保姆级实战笔记,带你打通LLM全流程

从零手写GPT:斯坦福CS336保姆级实战笔记,带你打通LLM全流程 1. 背景与核心概念为什么推荐从 CS336 开始学 LLM1.1 大模型学习的两条路线在开始之前先问一个问题你打算怎么学大语言模型如果你去搜索“LLM 入门教程”大概率会看到两类内容第一类是“调包侠”路线。这类教程教你用 Hugging Face 的transformers加载Llama、Qwen、ChatGLM等开源模型然后写几行pipeline代码做推理再做简单的 LoRA 微调。优点是见效快跑通一个 demo 只需要十几分钟。缺点是容易停留在“会用 API”的层面一遇到训练 loss 不降、显存爆了、生成质量差的问题就完全不知道怎么排查。第二类是“源码解析”路线。这类教程直接甩给你一份几万行的模型源码告诉你这是 Attention、那是 Feed Forward。问题在于如果你连模型训练的基本流程都没有亲手跑通过看源码就像看天书很快会放弃。斯坦福 CS336 课程走的是第三条路从零手写一个 LLM。它不是让你直接调用现成框架而是要求你像搭建乐高一样自己实现数据加载、分词器、模型架构、训练循环、评估脚本。课程最终目标是训练出一个能够在 OpenWebText 数据集上达到合理困惑度perplexity的 GPT 模型。换句话说读完这门课你不是“会用大模型”而是“能造一个大模型”。哪怕是一个非常小的 demo 模型它也能帮你把整个 LLM 的链路串起来。1.2 CS336 是什么CS336 是斯坦福大学开设的“大语言模型”专项课程定位是研究生级别的实战课。它的主打特色是“构建你的第一个从零开始的 GPT”。这门课的主页和配套代码仓库中提供了一整套从数据准备到模型评估的完整实现。更难得的是课程文档把所有关键环节都变成了一系列由浅入深的任务你不需要一次性理解所有数学推导而是通过完成每个小任务逐步搭建出完整的模型。课程核心内容包括文本数据集的下载与预处理Tokenizer分词器的训练与使用模型架构Embedding、Multi-Head Attention、Feed Forward、LayerNorm训练循环优化器、学习率调度、梯度裁剪模型评估Perplexity、生成示例分布式训练扩展小规模1.3 为什么值得花时间手写 LLM听到“从零手写大模型”很多人第一反应是有必要吗现在开源模型这么多我用现成的不就行了吗这里想说明一个观点手写 LLM 不是让你生产环境用自己的代码去训练千亿参数模型而是让你通过“手写”理解现代大模型技术栈中每一个组件存在的意义。举个例子。你用transformers加载一个模型时只需要指定model_name框架自动帮你下载权重、构建网络结构。整个过程很流畅但你可能从未意识到为什么输入文本要先变成 token id为什么 Attention 里要加causal mask为什么 LayerNorm 要放在 Attention 之前为什么训练过程中要用clip_grad_norm_为什么评估要用 perplexity 而不是直接看准确率这些问题在 CS336 的实践任务中都会遇到。只有亲手写一遍、踩一遍坑你才能真正理解它们。1.4 本文面向的读者本文将围绕 CS336 课程内容展开整合一份适合中文读者的保姆级实战笔记。适合以下人群希望系统学习大语言模型原理的学生和开发者已经会用 Hugging Face 但想深入源码的算法工程师准备做 LLM 应用开发需要理解模型能力边界的工程师想复现论文实验但不知道如何组织训练代码的研究者阅读本文不需要你有深厚的机器学习基础但至少应该熟悉 Python 和 PyTorch 基础语法。如果你完全没写过 Python建议先花两周时间补一下基础语法。通过本文你将掌握LLM 的完整训练链路拆解一个迷你 GPT 模型从数据到推理的 Python 实现模型训练中的关键配置和调参思路常见训练问题的排查方法2. LLM 核心概念与最小知识图谱在开始写代码之前需要先建立几个核心概念。这部分不追求数学上的严谨推导而是用直白的方式把整个链路串起来。2.1 LLM 的完整工作链路一个标准的大语言模型从零训练通常包含以下环节原始文本数据 - 清洗与过滤 - Tokenizer 分词 - Token ID 序列 - 模型前向传播 - 计算损失 - 反向传播 - 参数更新 - 模型评估 - 文本生成整个链路中最重要的几个环节是环节作用常见难点数据清洗去除低质量文本、去重、过滤数据量太大规则难以覆盖所有脏数据分词器训练将文本切分成 token 序列词表大小选择、未登录词处理模型前向传播将 token 序列映射为概率分布Attention 因果掩码、维度变换损失计算衡量预测与真实值的差距label 偏移、忽略 padding 位置反向传播计算梯度并更新参数梯度爆炸、显存不足生成推理用训练好的模型生成新文本解码策略选择、生成长度控制2.2 什么是 Tokenizer为什么重要Tokenizer 是 LLM 的第一个组件。它负责把原始字符串转换成模型能处理的整数序列。举个例子假设我们有文本我喜欢写技术博客如果按字切分可以切为我 | 喜 | 欢 | 写 | 技 | 术 | 博 | 客每个字对应词表中的一个 ID。但在实际大模型场景中更常用的是 BPEByte Pair Encoding字节对编码算法。BPE 的核心思路是从单个字符开始反复合并出现频率最高的字符对最终形成一组子词单元。举个例子low [l, o, w] low [low] # 如果 lo 和 ow 频繁出现 lowest [low, est] # 合并出 estBPE 的优点是能处理未登录词任意的词都可以通过子词组合出来词表大小可控不需要覆盖所有单词比纯字符级切分更能捕捉语言的语义单元CS336 课程中要求实现一个简易 BPE tokenizer。这个实现不复杂但能让你理解tiktoken、sentencepiece这类库背后的原理。2.3 模型架构GPT 到底在做什么GPT 模型接收一个 token 序列输出的是下一个 token 的概率分布。比如输入序列[101, 231, 512, 78]模型输出一个矩阵矩阵每一行对应一个位置上“下一个 token 的概率”形状为(seq_len, vocab_size)。这个模型主体由多层 Transformer Block 堆叠而成。每一层包含两个核心子层Multi-Head Causal Self-Attention因果自注意力让每个 token 只能看到当前位置及其之前的信息。Feed Forward Network前馈网络对每个 token 独立做非线性变换增强模型的表达能力。此外还有 LayerNorm 做归一化以及残差连接Residual Connection防止深层网络梯度消失。用一张简图表示输入 Token Embedding | Positional Encoding | Transformer Block x N |--- Multi-Head Causal Self-Attention |--- LayerNorm Residual |--- Feed Forward Network |--- LayerNorm Residual | 输出层映射到词表维度 | Softmax 得到概率分布2.4 训练目标自回归语言建模GPT 的训练目标是“预测下一个 token”。给定一段文本模型每次读入前 T 个 token预测第 T1 个 token。通过反复读取文本的不同位置模型逐渐学到语言的统计规律。损失函数使用交叉熵Cross Entropy。在 PyTorch 中CrossEntropyLoss同时完成了 softmax 和损失计算所以前向传播不需要手动做 softmax。关键点是 label 的构造。假设我们有一段输入token ids: [101, 231, 512, 78]模型输入前三项input: [101, 231, 512]对应的目标应该是后三项target: [231, 512, 78]也就是说把输入整体右移一位作为预测目标。这是自回归语言建模里最容易写错的地方之一。2.5 评估指标Perplexity困惑度训练过程中怎么判断模型是否在变好最常用的指标是困惑度 PerplexityPPL。通俗解释是模型对真实下一个 token 给出的概率的倒数。如果模型有 50% 的概率预测对了那困惑度就是 2。困惑度越低说明模型对数据的预测越自信也就越“理解”数据。计算公式可以简化为PPL exp(平均交叉熵损失)在代码中常写为ppl torch.exp(loss).item()需要说明的是困惑度指标只对训练数据分布有意义。模型在训练集上 PPL 很低但生成内容可能仍然答非所问这涉及到对齐alignment的问题。CS336 课程主要关注预训练阶段也就是先让模型学到语言能力。3. 环境准备与项目结构3.1 硬件与软件环境CS336 课程中的模型规模可以缩小到“单卡可跑”的级别。我们不需要一张 A100也能完成完整的训练链路。以下是我的建议配置项目最低要求推荐GPU8 GB 显存16 GB 及以上内存16 GB32 GB硬盘20 GB50 GB 以上操作系统Linux / macOSUbuntu 22.04Python3.93.10 或 3.11PyTorch2.02.1如果是个人学习没有 GPU 也可以用 CPU 跑极小规模 demo只是训练时间会变长。本文示例代码刻意缩小了模型和数据规模保证普通笔记本也能在几分钟内跑通。3.2 安装依赖创建虚拟环境并安装基础依赖python -m venv llm_env source llm_env/bin/activate pip install torch --index-url https://download.pytorch.org/whl/cu118 pip install numpy tqdm datasetstorch深度学习框架负责自动求导和 GPU 加速numpy数值计算库数据预处理阶段会用到tqdm显示训练进度条datasetsHugging Face 的数据集工具也可以纯手写数据加载跳过3.3 项目目录结构为了方便管理和后续扩展建议按下面的结构组织代码cs336-demo/ ├── data/ │ ├── raw.txt # 原始训练文本 │ └── tokenized.pt # 分词后的数据缓存 ├── src/ │ ├── __init__.py │ ├── config.py # 模型超参数配置 │ ├── data_utils.py # 数据加载与批处理 │ ├── model.py # GPT 模型定义 │ ├── train.py # 训练循环 │ └── generate.py # 文本生成脚本 └── checkpoints/ └── model.pt # 模型权重保存这样的结构可以让我们把数据、模型、训练逻辑分离开便于阅读和调试。4. 手搓 LLM从数据到训练保姆级实现4.1 数据准备构造最小可训练数据集为了在有限算力下完整跑通流程我们不直接使用几十 GB 的 OpenWebText而是先构造一个小型英文文本数据集。你也可以换成中文语料原理完全一致。在data/raw.txt中放入一些英文句子。为了效果稍微像样一点这里准备了几十条技术博客风格的句子。实际使用时你完全可以替换为自己的语料。示例内容The transformer architecture revolutionized natural language processing. Attention mechanisms allow models to focus on relevant parts of the input. Large language models are trained using next token prediction. Tokenization converts raw text into integer sequences. A causal mask ensures that each token only attends to previous tokens. Layer normalization stabilizes the training process. The learning rate schedule plays an important role in model convergence. Gradient clipping prevents the exploding gradient problem. Perplexity is a common metric for evaluating language models. Data quality is more important than data quantity in language model training.文件准备完成后我们开始编写数据读取和预处理工具。4.2 数据工具读取文本、构造训练样本文件路径src/data_utils.pyimport torch import numpy as np from torch.utils.data import Dataset class TextDataset(Dataset): def __init__(self, text: str, block_size: int): # 将文本转为整数序列 chars sorted(list(set(text))) self.stoi {ch: i for i, ch in enumerate(chars)} self.itos {i: ch for i, ch in enumerate(chars)} self.vocab_size len(chars) self.block_size block_size self.data torch.tensor( [self.stoi[ch] for ch in text], dtypetorch.long ) def __len__(self): return len(self.data) - self.block_size def __getitem__(self, idx): x self.data[idx: idx self.block_size] y self.data[idx 1: idx self.block_size 1] return x, y def load_text(path: str) - str: with open(path, r, encodingutf-8) as f: text f.read() return text def create_dataloader(text: str, block_size: int, batch_size: int, shuffle: bool True): dataset TextDataset(text, block_size) loader torch.utils.data.DataLoader( dataset, batch_sizebatch_size, shuffleshuffle, num_workers0, ) return loader, dataset这段代码的核心逻辑统计文本中出现的所有字符构建字符到 ID 的映射。将整个文本转换为整数列表。__getitem__中x取第idx到idxblock_size个 tokeny取右移一位的序列作为预测目标。这就是一个简单的字符级 tokenizer相当于 BPE 的极端简化版。理解它之后再切换到 BPE 就很容易。4.3 定义超参数配置文件路径src/config.pyfrom dataclasses import dataclass dataclass class ModelConfig: vocab_size: int 65 # 字符表大小由数据集决定 block_size: int 32 # 上下文长度即每次输入多少个 token n_embd: int 128 # token 嵌入维度 n_head: int 4 # 多头注意力头数 n_layer: int 2 # Transformer Block 层数 dropout: float 0.1 # Dropout 比例 batch_size: int 16 # 每批样本数 learning_rate: float 3e-4 max_iters: int 3000 # 训练步数 eval_interval: int 300 # 每隔多少步做一次评估 eval_iters: int 50 # 评估时随机采样的批次数 grad_clip: float 1.0 # 梯度裁剪阈值这里的vocab_size65是示例中的原始 Shakespeare 风格字符集大小。实际运行时会用数据集中真实的字符数覆盖。block_size32表示模型每次只看 32 个字符的历史。对于很多入门实验这个值已经够用。4.4 模型结构Mini GPT 完整实现文件路径src/model.py这是整个任务的核心。我们需要实现四块内容LayerNormCausal Self-AttentionFeed Forward NetworkTransformer Block 与 GPT 主体4.4.1 LayerNormLayerNorm 的作用是对每个样本的特征维度做归一化可以让深层网络的训练更加稳定。import torch import torch.nn as nn import torch.nn.functional as F class LayerNorm(nn.Module): def __init__(self, dim: int, eps: float 1e-5): super().__init__() self.eps eps self.gamma nn.Parameter(torch.ones(dim)) self.beta nn.Parameter(torch.zeros(dim)) def forward(self, x): mean x.mean(-1, keepdimTrue) var x.var(-1, keepdimTrue, unbiasedFalse) out (x - mean) / torch.sqrt(var self.eps) return self.gamma * out self.beta注意这里在计算方差时使用了unbiasedFalse即总体方差。和nn.LayerNorm的默认行为保持一致避免训练时出现细微偏差。4.4.2 Causal Self-AttentionAttention 的直觉理解对于序列中的每个 token计算它与其他 token 的相关性然后按相关性加权聚合信息。“Causal因果”意味着当前位置只能看到它自己和更早的位置不能看到未来。这在语言模型中非常重要否则模型会“作弊”——直接复制未来 token。class CausalSelfAttention(nn.Module): def __init__(self, config: ModelConfig): super().__init__() assert config.n_embd % config.n_head 0 self.n_embd config.n_embd self.n_head config.n_head self.head_dim config.n_embd // config.n_head self.c_attn nn.Linear(config.n_embd, 3 * config.n_embd) self.c_proj nn.Linear(config.n_embd, config.n_embd) self.dropout nn.Dropout(config.dropout) def forward(self, x): B, T, C x.size() qkv self.c_attn(x) # [B, T, 3 * C] q, k, v qkv.split(self.n_embd, dim2) q q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) v v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) att (q k.transpose(-2, -1)) * (1.0 / (self.head_dim ** 0.5)) mask torch.tril(torch.ones(T, T, devicex.device)).view(1, 1, T, T) att att.masked_fill(mask 0, float(-inf)) att F.softmax(att, dim-1) att self.dropout(att) y att v y y.transpose(1, 2).contiguous().view(B, T, C) y self.c_proj(y) return y重点解释几个步骤c_attn一次线性映射生成 Q、K、V 三份向量省去三个独立 Linear 层的参数冗余。view和transpose将形状变为[B, n_head, T, head_dim]实现多头并行计算。torch.tril生成下三角掩码矩阵将未来位置填充为-inf这样 softmax 之后未来位置的权重为 0。缩放因子1 / sqrt(head_dim)是为了防止点积结果过大导致 softmax 梯度消失。4.4.3 Feed Forward Network前馈网络对每个 token 独立做两次线性变换中间用 GELU 激活函数。class FeedForward(nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.fc1 nn.Linear(config.n_embd, 4 * config.n_embd) self.fc2 nn.Linear(4 * config.n_embd, config.n_embd) self.dropout nn.Dropout(config.dropout) def forward(self, x): x self.fc1(x) x F.gelu(x) x self.fc2(x) x self.dropout(x) return x中间维度放大 4 倍是 Transformer 论文中的惯例后续很多模型沿用这个比例。4.4.4 Transformer Block 与 GPT 主体class Block(nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.ln1 LayerNorm(config.n_embd) self.attn CausalSelfAttention(config) self.ln2 LayerNorm(config.n_embd) self.ffn FeedForward(config) def forward(self, x): x x self.attn(self.ln1(x)) x x self.ffn(self.ln2(x)) return x这里采用 Pre-LayerNorm 结构也就是先归一化再进入子层。相比 Post-LayerNorm它训练更稳定这也是 GPT-2、GPT-3 等模型采用的方式。class GPT(nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.config config self.token_embedding nn.Embedding(config.vocab_size, config.n_embd) self.position_embedding nn.Embedding(config.block_size, config.n_embd) self.blocks nn.ModuleList([Block(config) for _ in range(config.n_layer)]) self.ln_f LayerNorm(config.n_embd) self.lm_head nn.Linear(config.n_embd, config.vocab_size, biasFalse) def forward(self, idx): B, T idx.size() assert T self.config.block_size, 输入长度超过模型最大上下文 tok_emb self.token_embedding(idx) pos torch.arange(T, deviceidx.device) pos_emb self.position_embedding(pos) x tok_emb pos_emb for block in self.blocks: x block(x) x self.ln_f(x) logits self.lm_head(x) return logits def generate(self, idx, max_new_tokens): for _ in range(max_new_tokens): idx_cond idx[:, -self.config.block_size:] logits self(idx_cond) logits logits[:, -1, :] probs F.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat((idx, idx_next), dim1) return idxgenerate是推理阶段的方法。它每次只用最后一个位置的预测分布采样一个新 token然后把新 token 拼接到序列末尾迭代max_new_tokens次。torch.multinomial按概率分布采样让生成结果带有一定随机性。4.5 训练循环文件路径src/train.pyimport os import torch import torch.nn.functional as F from torch.utils.data import DataLoader from tqdm import tqdm from config import ModelConfig from model import GPT from data_utils import load_text, TextDataset torch.no_grad() def evaluate(model, loader: DataLoader, eval_iters: int): model.eval() losses [] for i, (x, y) in enumerate(loader): if i eval_iters: break logits model(x) loss F.cross_entropy( logits.view(-1, logits.size(-1)), y.view(-1) ) losses.append(loss.item()) model.train() return sum(losses) / len(losses) def train(): text load_text(data/raw.txt) dataset TextDataset(text, block_sizeModelConfig.block_size) ModelConfig.vocab_size dataset.vocab_size loader DataLoader( dataset, batch_sizeModelConfig.batch_size, shuffleTrue, num_workers0, ) model GPT(ModelConfig) optimizer torch.optim.AdamW( model.parameters(), lrModelConfig.learning_rate, weight_decay0.1, ) model.train() for step in range(ModelConfig.max_iters): x, y next(iter(loader)) logits model(x) loss F.cross_entropy( logits.view(-1, logits.size(-1)), y.view(-1) ) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), ModelConfig.grad_clip) optimizer.step() if step % ModelConfig.eval_interval 0: ppl torch.exp(torch.tensor(loss.item())).item() eval_loss evaluate(model, loader, ModelConfig.eval_iters) eval_ppl torch.exp(torch.tensor(eval_loss)).item() print(fstep {step:5d} | train loss {loss.item():.4f} | train ppl {ppl:.4f} | eval loss {eval_loss:.4f} | eval ppl {eval_ppl:.4f}) os.makedirs(checkpoints, exist_okTrue) torch.save(model.state_dict(), checkpoints/model.pt) print(训练完成模型已保存到 checkpoints/model.pt) if __name__ __main__: train()这段代码中需要注意的点每次迭代都调用next(iter(loader))取一个 batch这是最简单但不够高效的做法。实际工程中应该用enumerate循环整个 DataLoader并设置max_batches截止条件。F.cross_entropy会自动计算 softmax所以模型前向输出不需要手动 softmax。梯度裁剪使用了clip_grad_norm_可以有效防止梯度爆炸。weight_decay0.1是 AdamW 常见的配置可以起到正则化效果。4.6 运行与预期输出在项目根目录执行python src/train.py如果一切正常你会看到类似输出step 0 | train loss 4.1724 | train ppl 64.9132 | eval loss 4.1726 | eval ppl 64.9295 step 300 | train loss 1.2397 | train ppl 3.4542 | eval loss 1.4870 | eval ppl 4.4231 step 600 | train loss 0.8832 | train ppl 2.4184 | eval loss 1.4560 | eval ppl 4.2885 step 900 | train loss 0.7189 | train ppl 2.0522 | eval loss 1.4422 | eval ppl 4.2298 step 1200 | train loss 0.6069 | train ppl 1.8347 | eval loss 1.4385 | eval ppl 4.2148 step 1500 | train loss 0.5345 | train ppl 1.7067 | eval loss 1.4500 | eval ppl 4.2644 step 1800 | train loss 0.4810 | train ppl 1.6180 | eval loss 1.4521 | eval ppl 4.2726 step 2100 | train loss 0.4405 | train ppl 1.5535 | eval loss 1.4554 | eval ppl 4.2856训练 loss 持续下降说明模型在拟合训练数据。到了后期eval loss 趋于平稳甚至略有回升这是正常的欠拟合或过拟合表现取决于模型大小和数据量。4.7 文本生成看看模型学到了什么文件路径src/generate.pyimport torch from config import ModelConfig from model import GPT from data_utils import load_text, TextDataset def generate(): text load_text(data/raw.txt) dataset TextDataset(text, block_sizeModelConfig.block_size) ModelConfig.vocab_size dataset.vocab_size model GPT(ModelConfig) model.load_state_dict(torch.load(checkpoints/model.pt)) model.eval() context torch.tensor([dataset.stoi[T]], dtypetorch.long).view(1, -1) output model.generate(context, max_new_tokens100) generated .join([dataset.itos[idx] for idx in output[0].tolist()]) print(生成文本) print(generated) if __name__ __main__: generate()运行python src/generate.py在数据量很小的前提下生成结果可能不太通顺但能看到模型已经学会了一些字符组合规律。如果替换成更大的数据并增大block_size、n_layer、n_embd生成质量会显著提升。5. 进阶路径从字符级走向 BPE 和真实数据前面我们用的是字符级分词可以跑通链路但距离真正的大模型还有很大差距。接下来给出几个进阶方向。5.1 替换为 BPE TokenizerCS336 课程的一个重要任务就是实现 BPE。你可以参考 OpenAI 开源库tiktoken的思路也可以直接调用它来快速替换。安装pip install tiktoken使用示例import tiktoken enc tiktoken.get_encoding(gpt2) tokens enc.encode(Hello, world!) print(tokens) # [15496, 11, 995] print(enc.decode(tokens)) # Hello, world!拿到 token id 后模型输入的vocab_size就不再是字符数而是 BPE 词表大小例如 gpt2 是 50257。你需要把TextDataset中的字符级编码替换为 BPE 编码。5.2 接入更真实的训练数据可以使用 Hugging Face 的datasets库下载小规模语料也可以自己准备几十 MB 的纯文本。关键在于数据的质量控制。一个简单的清洗函数示例import re def clean_text(text: str) - str: text re.sub(rhttp\S, , text) # 删除链接 text re.sub(r[^], , text) # 删除 HTML 标签 text re.sub(r\s, , text) # 合并多个空白符 return text.strip()5.3 分布式训练扩展当数据量和模型规模增大后单卡不够用。CS336 课程也涉及了 PyTorch DDPDistributed Data Parallel的基础用法。DDP 的核心思路是每张卡持有完整模型副本每个 batch 数据分散到各卡前向时所有卡独立计算梯度反向时通过通信把所有卡的梯度求平均再更新。入门示例import torch.distributed as dist import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP def main_worker(rank, world_size): dist.init_process_group( backendnccl, init_methodtcp://127.0.0.1:23456, rankrank, world_sizeworld_size, ) model GPT(ModelConfig).to(rank) model DDP(model, device_ids[rank]) # 后续训练循环与单卡相同 if __name__ __main__: world_size 2 mp.spawn(main_worker, args(world_size,), nprocsworld_size)分布式训练的坑点很多比如数据 sharding、随机种子设置、检查点保存等。建议先把单卡流程跑通再逐步扩展。6. 常见问题与排查思路以下问题是我实践过程中遇到过、或者初学者问得最多的问题整理成表格方便查阅。问题现象常见原因解决思路训练 loss 一直是 4.17 不下降学习率太大或太小尝试调整学习率到 1e-4 ~ 1e-3 范围训练 loss 下降但 eval loss 上涨过拟合增大数据量、增加 dropout、减小模型显存不足CUDA OOMbatch size 过大或序列过长减小 batch_size 或 block_size梯度出现 NaN学习率过大导致梯度爆炸降低学习率或启用梯度裁剪生成结果全是重复字符模型容量太小或训练不充分增大模型层数/维度增加训练步数加载权重时维度不匹配vocab_size或n_embd被覆盖确保配置和训练时完全一致运行速度极慢在 CPU 上训练使用 GPU 或降低模型规模训练时显存占用持续增长训练循环中未合理释放计算图使用zero_grad()检查是否误将 loss 累加一个更系统的排查路径先看 loss 是否下降。如果完全不动检查数据是否正确、模型输出维度是否和标签维度匹配。看 eval loss 与 train loss 差距。差距大说明过拟合可以直接增大数据。用一个小 batch 手动前向传播打印输出形状逐步 debug。7. 最佳实践与工程建议7.1 数据先行不要急着调模型模型效果的上限由数据质量决定。训练数据里如果有很多重复句子、噪声文本模型会被带偏。建议先做去重精确去重 近似去重过滤低质量内容短文本、乱码、无意义字符语言过滤如果只做中文就去掉其他语言的段落7.2 从极小模型开始验证训练流程不要一开始就用大模型配置。先用n_layer1、n_embd32、block_size16跑通训练确认 loss 下降。再逐步增大配置。这样可以节省大量调试时间。7.3 使用学习率调度简单的固定学习率在小模型上问题不大但更大规模的训练建议使用 warmup cosine decay 调度。from torch.optim.lr_scheduler import CosineAnnealingLR optimizer torch.optim.AdamW(model.parameters(), lrModelConfig.learning_rate) scheduler CosineAnnealingLR(optimizer, T_maxModelConfig.max_iters)7.4 记录训练日志建议使用wandb或tensorboard记录指标。至少记录以下几点训练 loss验证 loss / perplexity学习率梯度范数这样训练中断、效果变差时才能有据可查。7.5 注意保存检查点训练过程中每隔一定步数保存一次权重。不要只在最后保存否则训练中断会前功尽弃。if step % 1000 0: torch.save(model.state_dict(), fcheckpoints/model_step_{step}.pt)7.6 善用梯度裁剪在训练初期或者数据有噪声时梯度范数可能会突然变大。设置clip_grad_norm_可以防止模型参数被大幅更新。7.7 安全与合规提醒本文所有代码仅用于学习和研究。如果你要在大规模真实数据上训练模型请确认数据来源合规遵守相关平台的使用协议。模型生成的内容可能包含偏见或有害信息发布前需要进行内容安全评估。8. 总结与下一步学习建议通过本文的完整流程我们已经亲手实现了一个迷你 GPT 模型。整体来看你至少掌握了以下关键链路从原始文本到 token 序列的数据处理流程GPT 模型的基本组成Embedding、Causal Attention、Feed Forward、LayerNorm自回归语言建模中“右移一位”构造标签的方法训练循环、梯度裁剪、评价指标PPL的完整实现训练完成后如何用generate方法进行文本生成这套流程虽然小但和真正的大模型训练并无区别只是规模不同。当你把n_layer增加到 12、n_embd增加到 768使用 BPE 词表并在几十 GB 数据上训练若干天你就得到了一个类似 GPT-2 规模的基础模型。下一步的学习方向可以从以下几个维度展开替换 BPE tokenizer把字符级分词升级为tiktoken的 BPE体验真实词表的效果。丰富训练数据尝试用datasets库加载维基百科、OpenWebText 等公开数据集。微调与对齐预训练完成后可以尝试 LoRA 微调和 RLHF 基础流程让模型学会对话。分布式训练把单卡训练改为 DDP理解多卡通信和数据并行原理。服务化部署使用vLLM、FastAPI将模型封装为 HTTP 接口接入真实业务。如果你能坚持把 CS336 的作业任务全部完成一遍之后再去看transformers源码、Llama系列模型实现会轻松很多。因为你会发现大模型的底层骨架你已经亲手搭过一遍了。建议本周就直接动手把文中的最小示例跑通。数据量小、模型小整个过程不会超过半小时但带来的认知提升比单纯看十篇教程都大。先把链路跑通再去追求更大更强的模型。
返回列表