★ FEATURED ARTICLE 简介《Happy-LLM从零开始的大语言模型原理与实践教程》是一份面向NLP初学者与进阶学习者的系统性开源教程专为希望深入理解LLM底层机制并动手训练模型的开发者设计。资源以单个19.06MB PDF文件形式提供内容覆盖NLP基础、Transformer架构详解、语言模型训练范式Encoder-Only/Decoder-Only/Encoder-Decoder、LLaMA2完整实现、LoRA/QLoRA微调、RAG与Agent实战应用等核心模块章节结构清晰理论推导与代码实操深度融合。已有553人学习下载教程由Datawhale社区发起强调“学完即能复现”包含从Tokenizer构建、小规模LLM训练到大模型部署的全流程指导并嵌入大量自注意力机制图解、训练日志分析与典型错误应对提示助力读者真正掌握大模型从原理到落地的关键能力。1. 为什么“从零开始”不是口号而是大语言模型落地的第一道生死线你手头这份《Happy-LLM从零开始的大语言模型原理与实践教程.pdf》名字里带“Happy”但现实中90%的工程师第一次照着它跑通第一个 tokenizer 加载、第一个 embedding 查表、第一个 causal mask 构建时卡在RuntimeError: expected scalar type Float but found Half上超过3小时——不是因为不会写代码而是因为教程没告诉你“从零开始”真正的敌人从来不是公式推导而是数据类型隐式转换、梯度缩放时机错位、以及 PyTorch 默认 device 分配策略在多卡环境下的玄学行为。这份 PDF 不是入门手册它是一份「可执行的逆向工程地图」用最小可运行模块单层 Decoder 手动实现 attention验证你对 KV Cache、RoPE 偏移、LayerNorm 归一化维度的理解是否真实成立用纯 NumPy 实现的 softmax 验证你是否真懂数值稳定性用一行torch.compile(model, fullgraphTrue)对比前后 CUDA kernel launch 次数确认你是否真的关闭了动态图开销。它适合两类人一类是刚读完《Attention Is All You Need》但对着 Hugging Face 源码不敢下断点的算法同学另一类是业务中被迫接手 LLM 微调任务、却连flash_attn和xformers底层 dispatch 逻辑都分不清的后端工程师。别急着下载权重、别急着跑 QLoRA——先让一个 32M 参数的 TinyLLM 在你的笔记本上不依赖任何预编译 wheel纯靠torch.nn.Linear和torch.bmm跑出和 HF 官方实现完全一致的 logits 差异 1e-5。这才是 Happy 的起点。2. 用纯 PyTorch 从头搭起 TinyLLM不碰 transformers只用 nn.Module 和 tensor ops要真正吃透 LLM必须亲手把每个张量的 shape、dtype、device、requires_grad 属性都“摸”一遍。本节目标构建一个仅含 2 层 Decoder、每层 8 个 head、hidden_size256 的极简模型所有组件Embedding、RMSNorm、RoPE、FlashAttention 手动版、MLP全部手写不 import 任何 transformers 或 flash_attn 的核心类。这不是炫技而是为了后续排查时能精准定位到某一行.view()是否破坏了 memory layout。2.1 构建 Token Embedding Positional Encoding 的最小闭环我们不直接用nn.Embedding而是用nn.Parameter显式管理 weight并手动实现 lookupimport torch import torch.nn as nn class TinyEmbedding(nn.Module): def __init__(self, vocab_size: int, hidden_size: int, max_seq_len: int 2048): super().__init__() self.vocab_size vocab_size self.hidden_size hidden_size # 手动初始化避免默认正态分布导致后续 softmax 数值爆炸 self.weight nn.Parameter(torch.empty(vocab_size, hidden_size)) nn.init.normal_(self.weight, mean0.0, std0.02) # RoPE 的 cos/sin 缓存预计算避免 forward 中重复计算 self.register_buffer(cos, torch.empty(max_seq_len, hidden_size // 2)) self.register_buffer(sin, torch.empty(max_seq_len, hidden_size // 2)) self._precompute_rope_freqs(max_seq_len) def _precompute_rope_freqs(self, max_seq_len: int): # RoPE 核心theta_i 10000^(-2i/d), i in [0, d//2) dim self.hidden_size // 2 inv_freq 1.0 / (10000 ** (torch.arange(0, dim, dtypetorch.float32) / dim)) t torch.arange(max_seq_len, dtypetorch.float32) freqs torch.outer(t, inv_freq) # [max_seq_len, dim] self.cos.copy_(torch.cos(freqs)) self.sin.copy_(torch.sin(freqs)) def forward(self, input_ids: torch.LongTensor) - torch.Tensor: # Step 1: token lookup → [B, T, D] x torch.nn.functional.embedding(input_ids, self.weight) # Step 2: RoPE apply — 注意只对最后 D//2 维做旋转前 D//2 维不变 # RoPE 公式[x0, x1, ..., x_{d-1}] → # [x0*cos x1*sin, -x0*sin x1*cos, x2*cos x3*sin, ...] B, T, D x.shape x_half x.view(B, T, D // 2, 2) # [B, T, D//2, 2] x0, x1 x_half[..., 0], x_half[..., 1] # [B, T, D//2] # 截取对应位置的 cos/sin[T, D//2] → [1, T, D//2] cos_t self.cos[:T].unsqueeze(0) sin_t self.sin[:T].unsqueeze(0) # 旋转x0 x0 * cos - x1 * sin; x1 x0 * sin x1 * cos x0_out x0 * cos_t - x1 * sin_t x1_out x0 * sin_t x1 * cos_t # 拼回 [B, T, D] x_rotated torch.stack([x0_out, x1_out], dim-1).view(B, T, D) return x_rotated参数说明max_seq_len2048是为 RoPE 预分配 buffer避免每次 forward 动态计算nn.init.normal_(..., std0.02)是 LLaMA 系列标准初始化比默认std1.0更稳定register_buffer确保 cos/sin 不参与梯度更新且随 model.to(device) 自动迁移。2.2 手写 RMSNorm为什么不能直接用 F.rms_normHugging Face 的F.rms_norm是 CUDA kernel黑匣子。而手写版本让你看清归一化维度、epsilon 作用点、以及torch.rsqrt的数值特性class TinyRMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float 1e-6): super().__init__() self.weight nn.Parameter(torch.ones(hidden_size)) self.eps eps def forward(self, x: torch.Tensor) - torch.Tensor: # RMSNorm: x / sqrt(mean(x^2, dim-1, keepdimTrue) eps) * weight # 注意mean 是对最后一个维度hidden_size求均值不是 batch 维度 variance x.pow(2).mean(dim-1, keepdimTrue) # [B, T, 1] # rsqrt 1/sqrt比 sqrt 1/x 更数值稳定 inv_rms torch.rsqrt(variance self.eps) # [B, T, 1] # weight 是 [D]需广播为 [1, 1, D] return x * inv_rms * self.weight关键逻辑x.pow(2).mean(dim-1, keepdimTrue)—— 这里dim-1是灵魂。若误写为dim1时间维度则整个 batch 内所有 token 共享一个 norm训练必崩。torch.rsqrt比1/torch.sqrt(...)更抗溢出这是 PyTorch 官方推荐写法。2.3 手写 Multi-Head Attention无 Flash理解 causal mask 的物理意义我们不用nn.MultiheadAttention它封装太深也不用flash_attn它绕过你对 mask 的控制而是用最原始的torch.bmmtorch.tril构建因果注意力class TinyAttention(nn.Module): def __init__(self, hidden_size: int, num_heads: int, dropout: float 0.0): super().__init__() self.hidden_size hidden_size self.num_heads num_heads self.head_dim hidden_size // num_heads assert self.head_dim * num_heads hidden_size, hidden_size must be divisible by num_heads # Wq, Wk, Wv, Wo —— 全部独立 Linear self.q_proj nn.Linear(hidden_size, hidden_size, biasFalse) self.k_proj nn.Linear(hidden_size, hidden_size, biasFalse) self.v_proj nn.Linear(hidden_size, hidden_size, biasFalse) self.o_proj nn.Linear(hidden_size, hidden_size, biasFalse) self.dropout nn.Dropout(dropout) # causal mask 缓存[1, 1, T, T]避免每次 forward 重建 self.register_buffer(causal_mask, None) def _get_causal_mask(self, T: int) - torch.Tensor: if self.causal_mask is None or self.causal_mask.size(-1) T: # 构建下三角矩阵tril(ones(T,T)) → [T, T] mask torch.tril(torch.ones(T, T, dtypetorch.bool)) # 扩展为 [1, 1, T, T] 以匹配 bmm 输出 shape self.causal_mask mask.unsqueeze(0).unsqueeze(0) return self.causal_mask[:, :, :T, :T] def forward(self, x: torch.Tensor, cache_k: torch.Tensor None, cache_v: torch.Tensor None) - torch.Tensor: B, T, D x.shape # Step 1: 线性投影 → [B, T, D] q self.q_proj(x) # [B, T, D] k self.k_proj(x) # [B, T, D] v self.v_proj(x) # [B, T, D] # Step 2: reshape for multi-head → [B, num_heads, T, head_dim] q q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) k k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) v v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) # Step 3: 计算 Q K^T → [B, num_heads, T, T] # 注意这里没有除以 sqrt(head_dim)因为后续 softmax 会吸收但为对齐 LLaMA我们显式除 attn_scores torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5) # Step 4: 应用 causal mask —— 关键mask 必须是 bool 类型且值为 True 的位置会被 fill -inf causal_mask self._get_causal_mask(T) attn_scores attn_scores.masked_fill(~causal_mask, float(-inf)) # Step 5: softmax dropout → [B, num_heads, T, T] attn_weights torch.softmax(attn_scores, dim-1) attn_weights self.dropout(attn_weights) # Step 6: 加权求和 → [B, num_heads, T, head_dim] attn_output torch.matmul(attn_weights, v) # Step 7: 拼接 heads → [B, T, D] attn_output attn_output.transpose(1, 2).contiguous().view(B, T, D) # Step 8: 输出投影 return self.o_proj(attn_output)血泪经验attn_scores.masked_fill(~causal_mask, float(-inf))中的~是易错点causal_mask是下三角True矩阵~causal_mask才是上三角True即需要屏蔽的位置。若漏掉~模型会看到未来 token训练 loss 直线下降但推理全错。torch.matmul(q, k.transpose(-2, -1))的-2,-1是安全写法比k.permute(0,1,3,2)更不易出错。3. 把 TinyLLM 接入真实训练流程数据加载、loss 计算与梯度裁剪的硬核细节光有模型结构不够必须让它真正“学起来”。本节聚焦三个常被教程忽略的实操环节如何构造符合 causal LM 的 label、为什么CrossEntropyLoss的ignore_index必须设为-100、以及梯度裁剪为何要在loss.backward()之后、optimizer.step()之前执行。3.1 构造 causal LM 的输入/label 对为什么 input_ids 和 labels 只差一位LLM 训练本质是给定x_0, x_1, ..., x_{t-1}预测x_t。因此labels必须是input_ids右移一位且首 token 无法预测应设为ignore_indexdef prepare_causal_batch(input_ids: torch.LongTensor, ignore_index: int -100) - dict: input_ids: [B, T] —— 原始 token 序列 返回: { input_ids: [B, T], # 用于模型 forward 的输入 labels: [B, T] # 用于 loss 计算的 target其中 labels[:, 0] ignore_index } B, T input_ids.shape # labels [x1, x2, ..., xT, ignore] → 实际就是 input_ids 右移一位 labels torch.full_like(input_ids, ignore_index) labels[:, :-1] input_ids[:, 1:] # 将 x1~x_{T-1} 填入 labels 的位置 0~T-2 # labels[:, -1] 保持 ignore_index因为没有 x_T1 可预测 return {input_ids: input_ids, labels: labels} # 示例 batch torch.tensor([[1, 2, 3, 4, 5]]) # [1, 5] prepared prepare_causal_batch(batch) print(input_ids:, prepared[input_ids]) # [[1, 2, 3, 4, 5]] print(labels: , prepared[labels]) # [[2, 3, 4, 5, -100]]为什么必须用-100torch.nn.CrossEntropyLoss默认ignore_index-100遇到该值会跳过此位置的 loss 计算且不计入梯度。若设为0或其他合法 token id模型会强行学习预测0导致灾难性错误。-100是 PyTorch 的约定俗成不是 magic number。3.2 手写 Trainer Loop暴露 optimizer.step() 前后的所有状态很多教程用Trainer封装一切但当你发现 loss 不降时根本不知道grad_norm是多少、param.grad是否为 NaN、lr是否按 schedule 正确衰减。以下是最小可调试 loopdef train_step( model: nn.Module, batch: dict, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, grad_clip: float 1.0, device: torch.device torch.device(cuda) ) - float: model.train() # Step 1: 数据搬移 input_ids batch[input_ids].to(device) # [B, T] labels batch[labels].to(device) # [B, T] # Step 2: forward logits model(input_ids) # [B, T, vocab_size] # Step 3: loss 计算 —— 注意logits 需 reshape 为 [B*T, vocab_size] loss_fct torch.nn.CrossEntropyLoss(ignore_index-100) loss loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1)) # Step 4: backward optimizer.zero_grad() loss.backward() # Step 5: 梯度裁剪 —— 关键必须在 backward 之后、step 之前 grad_norm torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) # Step 6: step update lr optimizer.step() scheduler.step() return loss.item(), grad_norm.item() # 使用示例 model TinyLLM(vocab_size32000, hidden_size256, num_layers2, num_heads8) model.to(cuda) optimizer torch.optim.AdamW(model.parameters(), lr3e-4, weight_decay0.1) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max1000) for step in range(1000): batch next(data_loader) # 假设 data_loader yield dict with input_ids loss, grad_norm train_step(model, batch, optimizer, scheduler, grad_clip1.0) if step % 10 0: print(fStep {step}: loss{loss:.4f}, grad_norm{grad_norm:.3f})翻车预警logits.view(-1, logits.size(-1))是必须的CrossEntropyLoss要求 input 为[N, C]label 为[N]。若直接传logits[B,T,C]和labels[B,T]PyTorch 会静默失败loss 值异常但不报错。grad_norm torch.nn.utils.clip_grad_norm_(...)返回的是全局 L2 norm若返回inf或nan说明某层梯度爆炸需立即检查weight_init或lr。3.3 验证 logits 正确性用 NumPy 复现 softmax对比 PyTorch 输出当你的 TinyLLM 输出 logits 和 HF 官方模型差异 1e-3问题一定出在实现细节。最可靠验证法用 NumPy 手写 softmax强制使用float64看是否与 PyTorchfloat32结果对齐import numpy as np def numpy_softmax(x: np.ndarray) - np.ndarray: NumPy softmax, stable version x x.astype(np.float64) # 强制高精度 x - np.max(x, axis-1, keepdimsTrue) # 减去最大值防溢出 exp_x np.exp(x) return exp_x / np.sum(exp_x, axis-1, keepdimsTrue) # 获取 PyTorch logits假设已 run 一次 forward with torch.no_grad(): pt_logits model(input_ids).cpu().numpy() # [1, T, V] # 取第一个 token 的 logits 做验证 pt_first pt_logits[0, 0] # [V] np_first numpy_softmax(pt_first) # [V] # 对比 top-5 概率 pt_top5 torch.topk(torch.from_numpy(pt_first), 5) np_top5 np.argsort(pt_first)[-5:][::-1] print(PyTorch top-5 indices:, pt_top5.indices.numpy()) print(NumPy top-5 indices: , np_top5) # 若两者不一致说明你的 RoPE 或 attention mask 逻辑有误玄学提示np.max(x, axis-1, keepdimsTrue)中的keepdimsTrue是关键。若漏掉x - np.max(...)会触发 broadcasting 错误导致 softmax 结果全为 0 或 inf。这个 bug 在低精度float16下更隐蔽。4. 避坑指南TinyLLM 训练中 4 个高频翻车现场与血泪解法别等 loss 爆到inf才查原因。以下是我在多个模拟项目X中反复踩过的坑按发生频率排序每条都附带可复现现象、根因分析和一行修复命令。4.1 现象loss 初始为 nan且grad_norm返回inf原因RMSNorm中variance x.pow(2).mean(...)在x全为 0 时variance0rsqrt(0eps)仍为有限值但若x含极大值如初始化偏差x.pow(2)溢出为infmean(inf)infrsqrt(inf)0最终x * 0 nan。解决在RMSNorm.forward开头加防御# 在 TinyRMSNorm.forward 第一行插入 x torch.clamp(x, min-1e4, max1e4) # 防止输入过大导致 pow(2) 溢出4.2 现象训练 loss 下降极慢1000 步后仍 5.0vocab_size32k 时正常应 3.0原因TinyAttention中attn_scores未除以sqrt(head_dim)导致 softmax 输入值过大概率分布过于尖锐梯度稀疏。解决确认attn_scores torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)这行存在且未被注释。用print(attn_scores.std())检查若 10则大概率没除。4.3 现象input_ids长度为 2048 时 OOM但长度 1024 时正常原因causal_mask在_get_causal_mask中未做缓存复用每次 forward 都新建[T,T]bool 矩阵T2048 时占用2048^2 * 1 byte ≈ 4MB叠加多层 attention显存暴涨。解决确保self.register_buffer(causal_mask, None)和self.causal_mask mask.unsqueeze(0).unsqueeze(0)存在且self.causal_mask在forward中被复用。用nvidia-smi观察若显存随 step 线性增长必是 buffer 泄漏。4.4 现象torch.compile后 loss 反而升高或CUDA error: device-side assert triggered原因torch.compile默认启用fullgraphTrue但手写causal_mask的torch.tril是动态 shapeT 来自 input_ids.shape[1]无法 trace 成静态图。解决禁用 fullgraph 或显式指定 dynamic shape# 方案1关闭 fullgraph保守 model_compiled torch.compile(model, fullgraphFalse) # 方案2声明 dynamic dims推荐 from torch._dynamo import config config.dynamic_shapes True model_compiled torch.compile(model)注意torch.compile不是银弹。在 TinyLLM 这种小模型上fullgraphTrue反而因过度优化引入 bug。我的习惯是先用fullgraphFalse跑通再逐步开启fullgraphTrue并用torch._dynamo.explain(model)查看 graph 是否 clean。5. 进阶技巧用 TinyLLM 做模型行为探针定位 HF 模型中的隐藏逻辑当你能稳定跑通 TinyLLM下一步不是堆参数而是把它变成一把手术刀——用来反向解析 Hugging Face 模型中那些文档没写的隐含行为。本节展示三个真实可用的探针技巧每个都能在 10 行内完成。5.1 探针 1验证 HF 模型是否真的用了 RoPE还是只是套了个壳HF 的LlamaForCausalLM默认启用 RoPE但某些魔改版本可能关掉。用 TinyLLM 的 RoPE 实现作为 reference对比输出from transformers import AutoModelForCausalLM # 加载 HF 模型不加载权重只结构 hf_model AutoModelForCausalLM.from_config( {model_type: llama, hidden_size: 256, num_attention_heads: 8, num_hidden_layers: 2} ) hf_model.eval() # 提取 HF 模型的 embedding 层假设叫 model.embed_tokens hf_embed hf_model.model.embed_tokens # 用 TinyEmbedding 的 RoPE 逻辑处理同一 input_ids input_ids torch.randint(0, 32000, (1, 10)) tiny_embed TinyEmbedding(vocab_size32000, hidden_size256) tiny_out tiny_embed(input_ids) # 提取 HF embedding 输出需 hook hook_out [] def hook_fn(module, input, output): hook_out.append(output.detach().cpu()) hf_embed.register_forward_hook(hook_fn) _ hf_model(input_ids) # 对比tiny_out 和 hook_out[0] 的 cosine similarity sim torch.nn.functional.cosine_similarity( tiny_out.flatten().float(), hook_out[0].flatten().float(), dim0 ) print(fRoPE match score: {sim.item():.4f}) # 0.9999 说明 HF 确实用了 RoPE技巧本质不依赖文档用输出一致性证明内部逻辑。若sim 0.9说明 HF 模型用了别的 positional encoding如 ALiBi或 RoPE 参数base不一致。5.2 探针 2测量 KV Cache 的实际内存占用而非理论值HF 的past_key_values是 tuple of tupletorch.cuda.memory_allocated()测不准。用 TinyLLM 的 cache 结构做精确测量class TinyLLMWithCache(TinyLLM): def forward(self, input_ids: torch.LongTensor, past_k: torch.Tensor None, past_v: torch.Tensor None): # ... 原 forward 逻辑 # 在 attention 后显式返回 new_k, new_v return logits, k, v # k,v shape: [B, num_heads, T, head_dim] # 初始化 cache cache_k torch.zeros(1, 8, 0, 32, devicecuda) # [B, h, 0, d] cache_v torch.zeros(1, 8, 0, 32, devicecuda) # 测单次 forward cache 内存增量 torch.cuda.reset_peak_memory_stats() _, new_k, new_v model(input_ids, cache_k, cache_v) peak_mem torch.cuda.max_memory_allocated() / 1024**2 print(fKV cache memory per token: {peak_mem / input_ids.size(1):.2f} MB/token)价值业务部署时per token memory比total model size更决定你能跑多长序列。若测出 0.5 MB/token说明 cache 未 quantize 或未 paged。5.3 探针 3用 TinyLLM 的梯度流反推 HF 模型哪一层 gradient 最不稳定HF 模型的gradient_checkpointing开关很玄学。用 TinyLLM 的 clean 结构逐层打印param.grad.std()def log_grad_stats(model: nn.Module): for name, param in model.named_parameters(): if param.grad is not None: std param.grad.std().item() print(f{name:30s}: grad_std {std:.6f}) # 在 train_step 的 optimizer.step() 前调用 log_grad_stats(model) # 输出示例 # q_proj.weight : grad_std 0.002145 # k_proj.weight : grad_std 0.000032 ← 这层梯度太小可能 dead # o_proj.weight : grad_std 0.015678行动指南若某层grad_std 1e-5优先检查其weight_init是否 std 过小、activation是否 saturate、或loss scale是否太小。这比盲调gradient_checkpointing有效十倍。我坚持用 TinyLLM 做探针不是因为它多先进而是因为它的每一行代码都在你眼皮底下。当 HF 文档说“RoPE is applied”你可以立刻用torch.allclose()验证当同事说“cache 内存没问题”你可以用torch.cuda.memory_allocated()精确到 KB当 loss 不降你不需要猜“是不是数据有问题”而是直接print(grad.std())看哪一层在躺平。这种掌控感是任何黑盒框架给不了的。希望帮到你。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助？
