很多人上来就想训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_model=d_model, nhead=n_heads, dim_feedforward=d_model * 4, dropout=dropout, activation="gelu", batch_first=True, # 输入格式 (batch, seq, dim) norm_first=True, # Pre-LN，大模型标配 ) self.encoder = nn.TransformerEncoder( encoder_layer, num_layers=n_layers, norm=nn.LayerNorm(d_model), # 最后加一个LayerNorm ) # 语言模型头：映射回词表空间 self.lm_head = nn.Linear(d_model, vocab_size, bias=False) # 权重共享：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, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward( self, input_ids: torch.Tensor, labels: torch.Tensor | None = None, ) -> dict: B, T = input_ids.shape # Embedding positions = torch.arange(T, device=input_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, mask=mask, is_causal=True) # 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_first=True**：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_size=vocab_size, seq_len=seq_len) return DataLoader( dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, drop_last=True, )

训练循环：最核心的代码

这是整个系列最重要的代码段。训练循环写不好，模型训不动、训不对、训不快。

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_ok=True) # 分离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, lr=learning_rate, betas=(0.9, 0.95), eps=1e-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_ids=input_ids, labels=labels) loss = outputs["loss"] # 反向传播 self.optimizer.zero_grad(set_to_none=True) # 比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( f"Step {self.global_step} | " f"Loss {loss.item():.4f} | " f"LR {current_lr:.2e} | " f"Tokens/s {tokens_per_sec:.0f} | " f"Tokens {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 / f"step_{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(f"Checkpoint saved: {path}")

训练循环里的关键设计

1.zero_grad(set_to_none=True)

比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_size=vocab_size, d_model=d_model, n_heads=n_heads, n_layers=n_layers, max_seq_len=seq_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(f"Total params: {total_params / 1e6:.1f}M") print(f"Trainable params: {trainable_params / 1e6:.1f}M") # 创建数据 loader = create_dataloader( vocab_size=vocab_size, seq_len=seq_len, batch_size=batch_size, ) # 创建训练器 trainer = Trainer( model=model, train_loader=loader, learning_rate=learning_rate, max_steps=max_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的小模型。要训大模型，主要改这些：

但代码结构不用改！上面的训练循环、Trainer类、模型定义的逻辑，7B模型也是这么写。区别只在于：

1. 混合精度训练：加autocast和GradScaler
2. 分布式训练：用FSDP分片到多卡（后面专门讲）
3. checkpointing：用梯度检查点换显存（用时间换空间）
4. 数据：换成真实语料的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(dtype=torch.float16): outputs = self.model(input_ids=input_ids, labels=labels) loss = outputs["loss"] # 缩放loss后反向传播 self.optimizer.zero_grad(set_to_none=True) 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%免费】