从零训练语言模型:小模型跑通全流程再说大的 很多人上来就想训7B、13B模型结果环境配不好、OOM、loss不收敛到处踩坑。我的建议是先用小模型跑通整个训练流程确认每一步都理解了、没问题了再放大。这篇从零开始手写一个完整的语言模型训练代码。不依赖HuggingFace Trainer不用任何封装每一步都看得见摸得着。模型一个极简的语言模型先写一个最小的语言模型Embedding 4层Transformer LM Head。这里不展开Transformer的实现细节下篇专门讲先用PyTorch内置的nn.TransformerEncoder。import torchimport torch.nn as nnimport mathclass SmallLM(nn.Module): 一个小型因果语言模型 def __init__( self, vocab_size: int 32000, d_model: int 512, n_heads: int 8, n_layers: int 4, max_seq_len: int 2048, dropout: float 0.1, ): super().__init__() self.d_model d_model # Token embedding 位置编码 self.token_embedding nn.Embedding(vocab_size, d_model) self.pos_embedding nn.Embedding(max_seq_len, d_model) self.dropout nn.Dropout(dropout) # Transformer编码器 encoder_layer nn.TransformerEncoderLayer( d_modeld_model, nheadn_heads, dim_feedforwardd_model * 4, dropoutdropout, activationgelu, batch_firstTrue, # 输入格式 (batch, seq, dim) norm_firstTrue, # Pre-LN大模型标配 ) self.encoder nn.TransformerEncoder( encoder_layer, num_layersn_layers, normnn.LayerNorm(d_model), # 最后加一个LayerNorm ) # 语言模型头映射回词表空间 self.lm_head nn.Linear(d_model, vocab_size, biasFalse) # 权重共享embedding和lm_head用同一套权重 self.lm_head.weight self.token_embedding.weight # 因果掩码下三角矩阵防止看到未来的token self.register_buffer( causal_mask, torch.tril(torch.ones(max_seq_len, max_seq_len)).bool() ) # 初始化权重 self.apply(self._init_weights) def _init_weights(self, module: nn.Module): GPT风格的权重初始化 if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean0.0, std0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean0.0, std0.02) def forward( self, input_ids: torch.Tensor, labels: torch.Tensor | None None, ) - dict: B, T input_ids.shape # Embedding positions torch.arange(T, deviceinput_ids.device).unsqueeze(0) # (1, T) x self.token_embedding(input_ids) self.pos_embedding(positions) x self.dropout(x) # Transformer 因果掩码 # causal_mask: (T, T)True允许注意力False遮蔽 mask self.causal_mask[:T, :T] x self.encoder(x, maskmask, is_causalTrue) # LM Head logits self.lm_head(x) # (B, T, vocab_size) # 计算loss loss None if labels is not None: # CrossEntropyLoss自带softmax不需要手动算 loss nn.functional.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index-100, # padding位置不算loss ) return {logits: logits, loss: loss}几个设计决策的解释**为什么要norm_firstTrue**Pre-LN先归一化再注意力比Post-LN训练更稳定。GPT-2之后的主流做法。Post-LN在深层网络里容易梯度爆炸得用learning rate warmup慢慢拉起来。Pre-LN没这问题。为什么要权重共享token_embedding和lm_head是同一个矩阵的两面——一个把token id映射到向量一个把向量映射回token概率。共享权重减少参数量还有正则化效果。GPT-2就是这么做的。**为什么用ignore_index-100**变长序列padding的部分不应该影响loss。标记为-100的位置会被cross_entropy跳过。数据一个玩具级但真实的Dataset为了跑通流程我写一个生成随机数据的Dataset。实际训练时替换成上篇讲的tokenized数据就行。from torch.utils.data import Dataset, DataLoaderclass ToyLMDataset(Dataset): 玩具数据集生成随机token序列 def __init__( self, vocab_size: int 32000, seq_len: int 256, num_samples: int 10000, ): self.vocab_size vocab_size self.seq_len seq_len self.num_samples num_samples # 预生成所有数据确保确定性 self.data torch.randint(0, vocab_size, (num_samples, seq_len 1)) def __len__(self) - int: return self.num_samples def __getitem__(self, idx: int) - dict: tokens self.data[idx] return { input_ids: tokens[:-1], # 除了最后一个 labels: tokens[1:], # 除了第一个错一位 }def create_dataloader( vocab_size: int 32000, seq_len: int 256, batch_size: int 16, num_workers: int 2,) - DataLoader: dataset ToyLMDataset(vocab_sizevocab_size, seq_lenseq_len) return DataLoader( dataset, batch_sizebatch_size, shuffleTrue, num_workersnum_workers, pin_memoryTrue, drop_lastTrue, )训练循环最核心的代码这是整个系列最重要的代码段。训练循环写不好模型训不动、训不对、训不快。import timefrom pathlib import Pathclass Trainer: 极简训练器 def __init__( self, model: nn.Module, train_loader: DataLoader, learning_rate: float 3e-4, weight_decay: float 0.1, warmup_steps: int 100, max_steps: int 5000, grad_clip: float 1.0, log_interval: int 100, save_interval: int 1000, save_dir: str checkpoints, device: str auto, ): if device auto: self.device torch.device(cuda if torch.cuda.is_available() else cpu) else: self.device torch.device(device) self.model model.to(self.device) self.train_loader train_loader self.max_steps max_steps self.grad_clip grad_clip self.log_interval log_interval self.save_interval save_interval self.save_dir Path(save_dir) self.save_dir.mkdir(exist_okTrue) # 分离decay和no-decay参数 # bias和LayerNorm不施加weight_decay decay_params [] no_decay_params [] for name, param in model.named_parameters(): if not param.requires_grad: continue if param.dim() 2 or norm in name or bias in name: no_decay_params.append(param) else: decay_params.append(param) optim_groups [ {params: decay_params, weight_decay: weight_decay}, {params: no_decay_params, weight_decay: 0.0}, ] # AdamW优化器 self.optimizer torch.optim.AdamW( optim_groups, lrlearning_rate, betas(0.9, 0.95), eps1e-8 ) # 学习率调度cosine with warmup self.scheduler self._create_scheduler( learning_rate, warmup_steps, max_steps ) # 训练状态 self.global_step 0 self.tokens_seen 0 def _create_scheduler(self, lr: float, warmup_steps: int, max_steps: int): Cosine学习率调度 warmup def lr_lambda(step: int) - float: if step warmup_steps: return step / max(warmup_steps, 1) progress (step - warmup_steps) / max(max_steps - warmup_steps, 1) return max(0.1, 0.5 * (1.0 math.cos(math.pi * progress))) return torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda) def train(self): self.model.train() data_iter iter(self.train_loader) total_tokens 0 start_time time.time() for step in range(self.max_steps): # 取数据处理迭代器耗尽 try: batch next(data_iter) except StopIteration: data_iter iter(self.train_loader) batch next(data_iter) # 数据搬到GPU input_ids batch[input_ids].to(self.device) labels batch[labels].to(self.device) # 前向传播 outputs self.model(input_idsinput_ids, labelslabels) loss outputs[loss] # 反向传播 self.optimizer.zero_grad(set_to_noneTrue) # 比zero_grad()快 loss.backward() # 梯度裁剪 if self.grad_clip 0: torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.grad_clip ) # 参数更新 self.optimizer.step() self.scheduler.step() # 统计 self.global_step 1 batch_tokens input_ids.numel() self.tokens_seen batch_tokens total_tokens batch_tokens # 日志 if self.global_step % self.log_interval 0: elapsed time.time() - start_time tokens_per_sec total_tokens / elapsed current_lr self.scheduler.get_last_lr()[0] print( fStep {self.global_step} | fLoss {loss.item():.4f} | fLR {current_lr:.2e} | fTokens/s {tokens_per_sec:.0f} | fTokens {self.tokens_seen / 1e6:.1f}M ) # 保存检查点 if self.global_step % self.save_interval 0: self._save_checkpoint() # 训练结束保存 self._save_checkpoint() def _save_checkpoint(self): path self.save_dir / fstep_{self.global_step}.pt torch.save( { model_state_dict: self.model.state_dict(), optimizer_state_dict: self.optimizer.state_dict(), scheduler_state_dict: self.scheduler.state_dict(), global_step: self.global_step, tokens_seen: self.tokens_seen, }, path, ) print(fCheckpoint saved: {path})训练循环里的关键设计1.zero_grad(set_to_noneTrue)比zero_grad()快因为它把.grad设成None而不是填零。None在backward时可以直接分配新内存而不用先清零再填充。2. 分离weight_decay的参数bias和LayerNorm的参数不应该施加weight_decay。weight_decay本质是L2正则化对bias做L2正则化没有意义——bias是用来偏移的不应该被拉向0。3. cosine学习率调度 warmup大模型训练的标配。warmup阶段学习率从0线性增长到目标值防止训练初期梯度太大把参数震飞。之后按余弦衰减让训练后期学得更精细。4. 梯度裁剪clip_grad_norm_把梯度的L2范数限制在max_norm以内。这是训练稳定性的保险措施——万一某个batch梯度爆炸裁剪一下不会崩。跑起来def main(): # 超参数 vocab_size 32000 seq_len 256 d_model 512 n_heads 8 n_layers 4 batch_size 16 learning_rate 3e-4 max_steps 5000 # 创建模型 model SmallLM( vocab_sizevocab_size, d_modeld_model, n_headsn_heads, n_layersn_layers, max_seq_lenseq_len, ) # 打印参数量 total_params sum(p.numel() for p in model.parameters()) trainable_params sum(p.numel() for p in model.parameters() if p.requires_grad) print(fTotal params: {total_params / 1e6:.1f}M) print(fTrainable params: {trainable_params / 1e6:.1f}M) # 创建数据 loader create_dataloader( vocab_sizevocab_size, seq_lenseq_len, batch_sizebatch_size, ) # 创建训练器 trainer Trainer( modelmodel, train_loaderloader, learning_ratelearning_rate, max_stepsmax_steps, ) # 开始训练 trainer.train()if __name__ __main__: main()预期输出随机数据loss不会降到很低但应该能看到在下降Total params: 27.8MTrainable params: 27.8MStep 100 | Loss 10.3245 | LR 3.00e-04 | Tokens/s 125000 | Tokens 0.4MStep 200 | Loss 9.8761 | LR 3.00e-04 | Tokens/s 128000 | Tokens 0.8M...从小模型到大模型什么要改上面跑通的是个27M的小模型。要训大模型主要改这些维度小模型(27M)中模型(350M)大模型(7B)d_model51210244096n_heads81632n_layers42432batch_size1664256~1024seq_len25610242048~8192显存需求2GB~8GB~30GB(单卡fp16)数据量玩具~10B tokens~1T tokens但代码结构不用改上面的训练循环、Trainer类、模型定义的逻辑7B模型也是这么写。区别只在于混合精度训练加autocast和GradScaler分布式训练用FSDP分片到多卡后面专门讲checkpointing用梯度检查点换显存用时间换空间数据换成真实语料的tokenized数据混合精度训练的修改很小加几行就行from torch.cuda.amp import autocast, GradScalerclass AMPTrainer(Trainer): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.scaler GradScaler() def train(self): # ...同上... for step in range(self.max_steps): # ... # 混合精度前向传播 with autocast(dtypetorch.float16): outputs self.model(input_idsinput_ids, labelslabels) loss outputs[loss] # 缩放loss后反向传播 self.optimizer.zero_grad(set_to_noneTrue) self.scaler.scale(loss).backward() # 梯度裁剪要先unscale self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.grad_clip ) # 参数更新 self.scaler.step(self.optimizer) self.scaler.update() self.scheduler.step()训练loss不收敛的排查清单如果loss不降或者降了又弹回来按这个顺序排查1. 学习率太大或太小太大loss振荡或NaN太小loss几乎不动试1e-4、3e-4、1e-3三个值2. 梯度爆炸现象loss突然变成NaN解法加梯度裁剪clip_grad_norm_1.0检查print(torch.nn.utils.clip_grad_norm_(model.parameters(), 1e10))看梯度范数3. 数据问题labels和input_ids没对齐padding的label不是-100数据太少模型记住了4. 模型bug因果掩码写反了看未来信息了embedding的scale没做大d_model时需要× sqrt(d_model)weight_decay施加到了不该施加的参数上5. 初始化问题残差连接的输出需要缩放× 1/sqrt(2*n_layers)否则深层梯度会爆这篇跑通了从模型定义到训练循环的完整流程。小模型先跑通大模型只是改参数加工程。训练循环里的每个设计决策——weight_decay分离、cosine调度、梯度裁剪、set_to_none——都是大模型训练的标配理解了小模型上的效果放大就不慌。学AI大模型的正确顺序千万不要搞错了2026年AI风口已来各行各业的AI渗透肉眼可见超多公司要么转型做AI相关产品要么高薪挖AI技术人才机遇直接摆在眼前有往AI方向发展或者本身有后端编程基础的朋友直接冲AI大模型应用开发转岗超合适就算暂时不打算转岗了解大模型、RAG、Prompt、Agent这些热门概念能上手做简单项目也绝对是求职加分王给大家整理了超全最新的AI大模型应用开发学习清单和资料手把手帮你快速入门学习路线:✅大模型基础认知—大模型核心原理、发展历程、主流模型GPT、文心一言等特点解析✅核心技术模块—RAG检索增强生成、Prompt工程实战、Agent智能体开发逻辑✅开发基础能力—Python进阶、API接口调用、大模型开发框架LangChain等实操✅应用场景开发—智能问答系统、企业知识库、AIGC内容生成工具、行业定制化大模型应用✅项目落地流程—需求拆解、技术选型、模型调优、测试上线、运维迭代✅面试求职冲刺—岗位JD解析、简历AI项目包装、高频面试题汇总、模拟面经以上6大模块看似清晰好上手实则每个部分都有扎实的核心内容需要吃透我把大模型的学习全流程已经整理好了抓住AI时代风口轻松解锁职业新可能希望大家都能把握机遇实现薪资/职业跃迁这份完整版的大模型 AI 学习资料已经上传CSDN朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】