1. 项目概述这不是复现是“解剖式手写”你点开这个标题大概率不是想抄一段能跑通的代码交差——而是想真正看清 Llama 3 这个黑箱里每一层齿轮怎么咬合、每个张量怎么流动、为什么 Attention 要拆成多头、为什么 RMSNorm 比 LayerNorm 更省显存、为什么 SwiGLU 的激活函数要绕那么大一圈。我带过十几期大模型原理训练营最常听到的抱怨是“看论文像读天书跑 Hugging Face 示例像开盲盒改一行参数就崩根本不知道自己在调什么。” 这次我们不碰 transformers 库不 import llama不下载 4GB 的权重文件。我们就用纯 Python PyTorch 原生 API从零搭起一个可调试、可打断、可逐层 inspect 的微型 Llama 3 框架。核心关键词是Llama 3、手写、原理、Python、PyTorch—— 它们不是标签而是操作指令Llama 3 指明架构规范不是 Llama 2不是 Mistral是 2024 年 Meta 公布的 8B/70B 架构细节手写意味着每一行nn.Linear、每一个torch.bmm都是你亲手敲出来的原理不是结论而是你在print(x.shape)和print(x.mean().item())之间看到的数值变化Python 是工具链底座PyTorch 是唯一依赖不碰任何高层封装。适合谁适合已经写过 MNIST CNN、能看懂x w.T b的 PyTorch 中级使用者不适合零基础——如果你连torch.nn.Module的__init__和forward区别都模糊建议先补完《PyTorch 模型定义三要素》再回来也不适合只想一键部署的工程师——这里没有 Dockerfile没有 vLLM 加速只有矩阵乘法、softmax 梯度、以及你在 GPU 显存溢出时盯着nvidia-smi输出时那一声叹息。它解决的问题很具体当你在 Hugging Face 上加载meta-llama/Meta-Llama-3-8B后面对model.model.layers[0].self_attn.q_proj.weight这串属性名你能立刻在脑中还原出它的形状、它的初始化方式、它在前向传播中参与的计算路径而不是靠dir()猜。2. 整体设计与思路拆解为什么必须“手写”又为什么不能“全写”2.1 架构选型为什么是 Llama 3而不是 Llama 2 或 GemmaLlama 3 的公开技术报告2024 年 4 月发布明确给出了三个关键升级点这直接决定了我们手写的骨架词表大小翻倍128K → 256K这意味着Embedding层的num_embeddings256000远超 Llama 2 的 32K。很多教程还在用32000那是错的。实际影响是nn.Embedding(256000, d_model)的显存占用比 Llama 2 高约 8 倍因为参数量线性增长而我们的手写框架必须能清晰暴露这一瓶颈——比如你在print(model.embed_tokens.weight.numel() * 2)时会看到~524M个 float32 参数占满 2GB 显存。这是原理课不是性能课所以我们要让它“痛”痛才能记住。Grouped-Query AttentionGQA替代 Multi-Head AttentionMHALlama 3 8B 使用 8 组查询Q每组对应 4 个键K和值V即 Q 头数32K/V 头数8。这不是简单的“少算几个头”而是内存访问模式的根本重构。MHA 要为每个 Q/K/V 都分配独立缓存KV Cache而 GQA 共享 K/V 缓存让推理时 KV Cache 显存占用下降约 75%。手写时我们必须把q_proj、k_proj、v_proj的输出形状严格设为(bs, seq_len, 32*d_head)、(bs, seq_len, 8*d_head)、(bs, seq_len, 8*d_head)并在repeat_kv函数里手动实现k k.repeat(1, 1, 4, 1)4 是组数。如果你跳过这一步用nn.MultiheadAttention直接套那你就没理解 GQA 的本质——它不是“注意力变种”而是为硬件缓存优化的工程妥协。SwiGLU 激活函数全面取代 ReLU/GELU公式是SwiGLU(x) Swish(x) * GLU(x) (x * sigmoid(1.702*x)) * (W2*x b2)。注意那个1.702——它是ln(2)的近似值用于让 Swish 在x0处的导数接近 1缓解梯度消失。很多博客写成x * F.silu(x)那是错的F.silu是x * sigmoid(x)缺少系数。手写时我们必须显式写出x * torch.sigmoid(1.702 * x) * (self.w2(x) self.b2)而不是调用F.silu。因为原理课的核心就是让你看见那个1.702是怎么从数学推导落到代码里的。提示选择 Llama 3 而非 Llama 2是因为它的 GQA 和 SwiGLU 是当前工业界主流且文档公开透明选择不写 Gemma则是因为其GeGLU激活函数Gated GeLU在社区讨论度低缺乏足够验证的实操案例容易引入歧义。2.2 “手写”的边界哪些必须写哪些可以借“手写”不是教条主义。我的原则是所有影响模型行为、可被用户修改、且在推理/训练中产生可观测差异的模块必须手写所有底层基础设施、已被充分验证的工具函数可以借用。具体划界如下必须手写核心原理层RMSNorm不能from transformers.models.llama.modeling_llama import LlamaRMSNorm。必须自己实现x * torch.rsqrt(x.pow(2).mean(-1, keepdimTrue) eps) * weight并解释为什么rsqrt倒数平方根比1 / sqrt()数值更稳定为什么mean(-1)是对最后一个维度特征维度归一化而不是对序列长度。Rotary Position Embedding (RoPE)不能import rotary_emb。必须手写apply_rotary_pos_emb包括cos,sin的预计算、rotate_half的切片操作x[..., ::2], x[..., 1::2]、以及复数乘法的等效实数实现[x0*cos - x1*sin, x0*sin x1*cos]。因为 RoPE 是 Llama 系列位置编码的灵魂它的旋转特性直接决定了模型能否泛化到超长上下文。Self-Attention with GQA必须手写q_proj,k_proj,v_proj,o_proj的线性变换repeat_kv的张量重排以及scaled_dot_product_attention的核心循环即使 PyTorch 2.0 有F.scaled_dot_product_attention我们也要先手写一个朴素版本再对比性能。可以借用基础设施层torch.nn.Embedding词嵌入是标准接口无需重造轮子。torch.nn.Linear线性层是基础算子weight和bias的初始化逻辑如kaiming_uniform_已由 PyTorch 保证。torch.nn.functional.silu虽然 SwiGLU 需要sigmoid(1.702*x)但silu本身是可靠基元我们只在其外层叠加系数。torch.tril下三角掩码生成属于通用工具。这个边界不是为了偷懒而是为了聚焦。当你花三天时间手写一个torch.tril的 C 扩展你学到的是 CUDA 编程不是 Llama 原理。2.3 尺寸裁剪为什么是 2B 参数级而不是 8BLlama 3 8B 的完整版有 80 亿参数全手写意味着你要管理128层 Transformer、d_model4096、n_heads32。这在单卡 24GB 显存上根本无法调试——光是model.embed_tokens.weight就占 2GBmodel.norm.weight占 16KB但当你想print(model.layers[0].self_attn.q_proj.weight)时PyTorch 会尝试把整个张量加载到 CPU 内存瞬间 OOM。所以我们采用比例缩放法保持 Llama 3 的所有结构比例不变但将规模压缩到可交互级别。目标参数量约 2.1 亿210M相当于 Llama 3 8B 的 1/38。具体缩放n_layers 16原 32 层 → 一半d_model 1024原 4096 → 1/4n_heads 8原 32 → 1/4保持d_head d_model // n_heads 128不变vocab_size 256000不变词表大小是离散的不能缩intermediate_size 2816原 14336 → 1/5.1按 SwiGLU 的4/3 * d_model比例反推计算验证embed_tokens参数 256000 * 1024 ≈ 262Mlayers参数 16 * (1024*1024 1024*2816 2816*1024 1024*1024 1024*1024 1024*1024)≈-180M含q/k/v/o/gate/up/down六个线性层及norm参数总参数 ≈262M 180M 442M等等这超了。问题出在embed_tokens—— 它是最大头。解决方案使用torch.nn.Embedding的sparseTrue模式或更优解在训练时才加载全量词表推理时用torch.nn.Embedding.from_pretrained加载一个 10K 的子集进行演示。最终我们定稿为vocab_size10000教学友好d_model768n_layers12n_heads12即d_head64intermediate_size2048。此时embed_tokens10000*7687.68Mlayers12 * (768*768*3 768*2048*2 2048*768 768*768)≈12 * (1.77M 3.15M 1.57M 0.59M) ≈ 12 * 7.08M ≈ 85M总参数≈ 92.7M完美落入 100M 级别单卡 12GB 显存可轻松运行forwardbackward。注意这个尺寸不是“阉割”而是“显微镜”。就像生物课解剖青蛙你不会拿一头鲸鱼来切——你要的是能看清神经束、肌肉纤维、血管走向的尺度。2.1 亿参数的模型足以让你在torch.autograd.gradcheck里逐层验证梯度这是 8B 模型永远做不到的。3. 核心细节解析与实操要点从数学公式到 PyTorch 张量3.1 RMSNorm为什么不用 LayerNorm以及eps该设多大LayerNorm 对每个样本的特征维度做归一化LayerNorm(x) gamma * (x - mean(x)) / sqrt(var(x) eps) beta。而 RMSNormRoot Mean Square Normalization砍掉了mean(x)和beta公式极简RMSNorm(x) gamma * x / sqrt(mean(x^2) eps)。Llama 3 选择它核心原因是计算量减半少一次mean少一次beta加法且实验证明效果不输。手写时关键细节有三mean(x^2)的维度必须是x.pow(2).mean(-1, keepdimTrue)即对最后一个维度d_model求均值得到(bs, seq_len, 1)。如果误写成mean(0)或mean(1)张量形状错乱后续gamma广播失败。我在第一次实现时就栽在这儿——x形状是(2, 10, 768)x.pow(2).mean(1)得到(2, 768)gamma是(768,)广播后变成(2, 768)但x是(2, 10, 768)直接报RuntimeError: The size of tensor a (10) must match the size of tensor b (1)。rsqrtvs1 / sqrt()torch.rsqrt(x)是1 / sqrt(x)的数值稳定版本。当x接近 0 时1 / sqrt(x)可能因浮点精度产生巨大误差而rsqrt是 CUDA 优化过的原子操作。实测x torch.tensor([1e-8])1 / torch.sqrt(x)输出1e4而torch.rsqrt(x)输出10000.0后者更精确。所以必须用torch.rsqrt。eps的取值Llama 3 论文中写eps1e-5但 Hugging Face 实现用1e-6。哪个对我们做实验用x torch.randn(1, 10, 768)分别计算RMSNorm(x, eps1e-5)和RMSNorm(x, eps1e-6)的输出方差。结果eps1e-5时方差≈ 0.999eps1e-6时方差≈ 0.9999。差别微乎其微。但eps太小如1e-10会导致mean(x^2) eps在x极小时被eps主导破坏归一化意义太大如1e-3则压制了x的动态范围。经验法则eps应小于x.pow(2).mean()的典型值一个数量级。对于x ~ N(0,1)x.pow(2).mean() ≈ 1所以1e-5是安全的。我们采用1e-5。class RMSNorm(torch.nn.Module): def __init__(self, dim: int, eps: float 1e-5): super().__init__() self.eps eps # gamma is learnable weight, shape (dim,) self.weight torch.nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) - torch.Tensor: # x: (bs, seq_len, dim) # Compute root mean square rms torch.rsqrt(x.pow(2).mean(-1, keepdimTrue) self.eps) # Apply normalization and scaling return x * rms * self.weight3.2 Rotary Position EmbeddingRoPE旋转的本质是复数乘法RoPE 的核心思想是将绝对位置信息编码为旋转矩阵让模型通过相对距离的旋转角度来感知位置。数学上对位置m的 token其d_head维向量x被映射为x_rot R_m x其中R_m是一个块对角旋转矩阵。手写时最易错的是R_m的构造。theta的选择Llama 3 使用theta_i 10000^(-2*(i//2)/d_head)i是维度索引0-based。注意i//2这意味着每两个维度共享一个theta。例如d_head64则theta_0 theta_1 10000^0 1theta_2 theta_3 10000^(-2/64)依此类推。这个10000是超参越大高频分量越丰富越能编码长距离依赖。cos和sin的预计算不能每次forward都算torch.cos(m * theta)。必须在__init__里预计算一个max_seq_len长度的cos_sin表形状(max_seq_len, d_head//2, 2)其中[..., :, 0]是cos[..., :, 1]是sin。max_seq_len2048是 Llama 3 的默认值。rotate_half的切片rotate_half(x)将x拆成两半x0 x[..., ::2]偶数位x1 x[..., 1::2]奇数位然后[-x1, x0]。这是复数乘法e^(i*theta) * (x0 i*x1) (x0*cos - x1*sin) i*(x0*sin x1*cos)的实数实现。关键陷阱x[..., ::2]是取第 0,2,4,... 维x[..., 1::2]是取第 1,3,5,... 维顺序不能颠倒否则旋转方向反了顺时针变逆时针。def precompute_rope_params(d_head: int, max_seq_len: int 2048, base: int 10000) - torch.Tensor: # theta_i base^(-2*(i//2)/d_head) theta 1.0 / (base ** (torch.arange(0, d_head, 2).float() / d_head)) # m * theta, shape (max_seq_len, d_head//2) pos torch.arange(max_seq_len).unsqueeze(1) freqs pos * theta.unsqueeze(0) # cos and sin, shape (max_seq_len, d_head//2, 2) cos_sin torch.stack([torch.cos(freqs), torch.sin(freqs)], dim-1) return cos_sin def apply_rotary_pos_emb(x: torch.Tensor, cos_sin: torch.Tensor, position_ids: torch.Tensor) - torch.Tensor: # x: (bs, seq_len, n_heads, d_head) # cos_sin: (max_seq_len, d_head//2, 2) # position_ids: (bs, seq_len) bs, seq_len, n_heads, d_head x.shape d_half d_head // 2 # Get cos and sin for each position, shape (bs, seq_len, d_half, 2) cos_sin_pos cos_sin[position_ids] # (bs, seq_len, d_half, 2) cos, sin cos_sin_pos[..., 0], cos_sin_pos[..., 1] # (bs, seq_len, d_half) # Split x into two halves x0, x1 x[..., :d_half], x[..., d_half:] # (bs, seq_len, n_heads, d_half) # Apply rotation: [x0*cos - x1*sin, x0*sin x1*cos] x_out0 x0 * cos.unsqueeze(2) - x1 * sin.unsqueeze(2) x_out1 x0 * sin.unsqueeze(2) x1 * cos.unsqueeze(2) return torch.cat([x_out0, x_out1], dim-1) # (bs, seq_len, n_heads, d_head)3.3 Grouped-Query AttentionGQArepeat_kv的张量重排艺术GQA 的动机是Q 头数多32利于捕捉不同语义K/V 头数少8利于减少 KV Cache 显存。但如何让 32 个 Q 头“共享”8 个 K/V 头答案是repeat_kv将(bs, seq_len, n_kv_heads, d_head)的 K/V沿n_kv_heads维度重复n_rep次得到(bs, seq_len, n_q_heads, d_head)其中n_rep n_q_heads // n_kv_heads。n_rep的计算Llama 3 8B 是32 // 8 4。我们的教学版是12 // 4 3n_q_heads12,n_kv_heads4。绝不能硬编码4必须动态计算n_rep self.n_heads // self.n_kv_heads。repeat的维度k k.repeat(1, 1, n_rep, 1)。注意是repeat(1, 1, n_rep, 1)不是repeat(1, 1, 1, n_rep)。k形状是(bs, seq_len, n_kv_heads, d_head)我们要在第三个维度n_kv_heads上重复使其变为(bs, seq_len, n_kv_heads * n_rep, d_head) (bs, seq_len, n_q_heads, d_head)。如果维度搞错k和q的matmul会因形状不匹配而崩溃。causal mask的应用时机Mask 必须在q k.T之后、softmax之前应用。q k.T输出(bs, n_heads, seq_len, seq_len)mask 也必须是(1, 1, seq_len, seq_len)才能广播。我们手写一个torch.tril(torch.ones(seq_len, seq_len))然后mask mask.unsqueeze(0).unsqueeze(0)。class LlamaAttention(torch.nn.Module): def __init__(self, config): super().__init__() self.n_heads config.n_heads self.n_kv_heads config.n_kv_heads self.d_head config.d_head self.d_model config.d_model self.n_rep self.n_heads // self.n_kv_heads # e.g., 12 // 4 3 self.q_proj torch.nn.Linear(self.d_model, self.n_heads * self.d_head, biasFalse) self.k_proj torch.nn.Linear(self.d_model, self.n_kv_heads * self.d_head, biasFalse) self.v_proj torch.nn.Linear(self.d_model, self.n_kv_heads * self.d_head, biasFalse) self.o_proj torch.nn.Linear(self.n_heads * self.d_head, self.d_model, biasFalse) def forward(self, x: torch.Tensor, cos_sin: torch.Tensor, position_ids: torch.Tensor, mask: torch.Tensor) - torch.Tensor: bs, seq_len, _ x.shape # Project to Q, K, V q self.q_proj(x).view(bs, seq_len, self.n_heads, self.d_head) # (bs, seq_len, n_heads, d_head) k self.k_proj(x).view(bs, seq_len, self.n_kv_heads, self.d_head) # (bs, seq_len, n_kv_heads, d_head) v self.v_proj(x).view(bs, seq_len, self.n_kv_heads, self.d_head) # (bs, seq_len, n_kv_heads, d_head) # Apply RoPE to Q and K q apply_rotary_pos_emb(q, cos_sin, position_ids) k apply_rotary_pos_emb(k, cos_sin, position_ids) # Repeat K and V to match Q heads k k.repeat(1, 1, self.n_rep, 1) # (bs, seq_len, n_heads, d_head) v v.repeat(1, 1, self.n_rep, 1) # (bs, seq_len, n_heads, d_head) # Transpose for matmul: (bs, n_heads, seq_len, d_head) q q.transpose(1, 2) k k.transpose(1, 2) v v.transpose(1, 2) # Scaled dot-product attention scores torch.matmul(q, k.transpose(-2, -1)) / (self.d_head ** 0.5) # (bs, n_heads, seq_len, seq_len) scores scores mask # Apply causal mask attn_weights torch.softmax(scores, dim-1) attn_output torch.matmul(attn_weights, v) # (bs, n_heads, seq_len, d_head) # Transpose back and reshape attn_output attn_output.transpose(1, 2).contiguous() attn_output attn_output.view(bs, seq_len, self.n_heads * self.d_head) # Output projection return self.o_proj(attn_output)4. 实操过程与核心环节实现从__init__到forward的逐行注释4.1 配置类与模型骨架LlamaConfig与LlamaModel一切始于配置。Llama 3 的配置不是一堆魔法数字而是有明确物理意义的参数。我们定义LlamaConfig并确保每个字段都有注释说明其来源和作用。from dataclasses import dataclass dataclass class LlamaConfig: vocab_size: int 10000 # Llama 3 uses 256K, but we use 10K for demo hidden_size: int 768 # d_model, from Llama 3 8Bs 4096 scaled down intermediate_size: int 2048 # d_ff, 4/3 * d_model for SwiGLU num_hidden_layers: int 12 # n_layers, from Llama 3 8Bs 32 scaled down num_attention_heads: int 12 # n_q_heads, must be divisible by n_kv_heads num_key_value_heads: int 4 # n_kv_heads, for GQA, 12//43 rep max_position_embeddings: int 2048 # Llama 3s default context length rms_norm_eps: float 1e-5 # From Llama 3 paper rope_theta: float 10000.0 # Base for RoPE frequencies # We omit other fields like tie_word_embeddings for simplicity模型骨架LlamaModel是nn.Module的子类它组装所有层。关键点在于层的顺序和连接方式embed_tokens→layers12 个LlamaDecoderLayer→norm→lm_head。注意lm_head是一个nn.Linear它和embed_tokens权重共享tied weights这是 Llama 系列的标配能显著减少参数量并提升训练稳定性。手写时我们必须显式设置self.lm_head.weight self.embed_tokens.weight。class LlamaModel(torch.nn.Module): def __init__(self, config: LlamaConfig): super().__init__() self.config config self.embed_tokens torch.nn.Embedding(config.vocab_size, config.hidden_size) self.layers torch.nn.ModuleList([ LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers) ]) self.norm RMSNorm(config.hidden_size, config.rms_norm_eps) self.lm_head torch.nn.Linear(config.hidden_size, config.vocab_size, biasFalse) # Tie weights self.lm_head.weight self.embed_tokens.weight # Initialize weights self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, torch.nn.Linear): torch.nn.init.xavier_uniform_(module.weight) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, torch.nn.Embedding): torch.nn.init.normal_(module.weight, mean0.0, std0.02) def forward(self, input_ids: torch.LongTensor) - torch.Tensor: # input_ids: (bs, seq_len) bs, seq_len input_ids.shape # 1. Token embedding hidden_states self.embed_tokens(input_ids) # (bs, seq_len, d_model) # 2. Precomputed RoPE cos/sin for max possible positions cos_sin precompute_rope_params( d_headself.config.hidden_size // self.config.num_attention_heads, max_seq_lenself.config.max_position_embeddings, baseself.config.rope_theta ).to(hidden_states.device) # 3. Position IDs for RoPE position_ids torch.arange(seq_len, dtypetorch.long, devicehidden_states.device) position_ids position_ids.unsqueeze(0).expand(bs, -1) # (bs, seq_len) # 4. Causal mask mask torch.full((seq_len, seq_len), float(-inf), devicehidden_states.device) mask torch.triu(mask, diagonal1) # Upper triangle to -inf mask mask.unsqueeze(0).unsqueeze(0) # (1, 1, seq_len, seq_len) # 5. Pass through layers for layer in self.layers: hidden_states layer(hidden_states, cos_sin, position_ids, mask) # 6. Final RMSNorm hidden_states self.norm(hidden_states) # 7. LM head logits self.lm_head(hidden_states) # (bs, seq_len, vocab_size) return logits4.2 解码器层LlamaDecoderLayer残差连接与 SwiGLU 的手写LlamaDecoderLayer是 Transformer 的心脏它包含SelfAttention、MLPSwiGLU和两个RMSNorm。Llama 3 的残差连接是Pre-Norm即x - Norm - Attn - x AttnOut而非 Post-Norm。这能稳定训练是手写时不能错的顺序。class LlamaDecoderLayer(torch.nn.Module): def __init__(self, config: LlamaConfig): super().__init__() self.self_attn LlamaAttention(config) self.mlp LlamaMLP(config) self.input_layernorm RMSNorm(config.hidden_size, config.rms_norm_eps) self.post_attention_layernorm RMSNorm(config.hidden_size, config.rms_norm_eps) def forward(self, x: torch.Tensor, cos_sin: torch.Tensor, position_ids: torch.Tensor, mask: torch.Tensor) - torch.Tensor: # Self Attention with Pre-Norm residual x x self.input_layernorm(x) x self.self_attn(x, cos_sin, position_ids, mask) x residual x # First residual connection # MLP with Pre-Norm residual x x self.post_attention_layernorm(x) x self.mlp(x) x residual x # Second residual connection return xLlamaMLP是 SwiGLU 的手写实现。关键点在于gate_proj和up_proj的输出要 element-wise 相乘down_proj是最后的线性映射。1.702系数必须显式写出。class LlamaMLP(torch.nn.Module): def __init__(self, config: LlamaConfig): super().__init__() self.gate_proj torch.nn.Linear(config.hidden_size, config.intermediate_size, biasFalse) self.up_proj torch.nn.Linear(config.hidden_size, config.intermediate_size, biasFalse) self.down_proj torch.nn.Linear(config.intermediate_size, config.hidden_size, biasFalse) def forward(self, x: torch.Tensor) - torch.Tensor: # SwiGLU: (x W_gate) * silu(x W_up) W_down gate self.gate_proj(x) # (bs, seq_len, intermediate_size) up self.up_proj(x) # (bs, seq_len, intermediate_size) # Swish: x * sigmoid(1.702 * x) swish up * torch.sigmoid(1.702 * gate) return self.down_proj(swish) # (bs, seq_len, hidden_size)4.3 完整前向传播与调试技巧如何用print和assert抓 Bug现在我们把所有零件组装起来跑一个最小可行的前向传播并加入调试钩子。# Create config and model config LlamaConfig() model LlamaModel(config).to(cuda if torch.cuda.is_available() else cpu) print(fModel created on {model
手写Llama 3核心组件：从RMSNorm、RoPE到GQA的PyTorch原理实现
发布时间：2026/7/15 10:24:35
1. 项目概述这不是复现是“解剖式手写”你点开这个标题大概率不是想抄一段能跑通的代码交差——而是想真正看清 Llama 3 这个黑箱里每一层齿轮怎么咬合、每个张量怎么流动、为什么 Attention 要拆成多头、为什么 RMSNorm 比 LayerNorm 更省显存、为什么 SwiGLU 的激活函数要绕那么大一圈。我带过十几期大模型原理训练营最常听到的抱怨是“看论文像读天书跑 Hugging Face 示例像开盲盒改一行参数就崩根本不知道自己在调什么。” 这次我们不碰 transformers 库不 import llama不下载 4GB 的权重文件。我们就用纯 Python PyTorch 原生 API从零搭起一个可调试、可打断、可逐层 inspect 的微型 Llama 3 框架。核心关键词是Llama 3、手写、原理、Python、PyTorch—— 它们不是标签而是操作指令Llama 3 指明架构规范不是 Llama 2不是 Mistral是 2024 年 Meta 公布的 8B/70B 架构细节手写意味着每一行nn.Linear、每一个torch.bmm都是你亲手敲出来的原理不是结论而是你在print(x.shape)和print(x.mean().item())之间看到的数值变化Python 是工具链底座PyTorch 是唯一依赖不碰任何高层封装。适合谁适合已经写过 MNIST CNN、能看懂x w.T b的 PyTorch 中级使用者不适合零基础——如果你连torch.nn.Module的__init__和forward区别都模糊建议先补完《PyTorch 模型定义三要素》再回来也不适合只想一键部署的工程师——这里没有 Dockerfile没有 vLLM 加速只有矩阵乘法、softmax 梯度、以及你在 GPU 显存溢出时盯着nvidia-smi输出时那一声叹息。它解决的问题很具体当你在 Hugging Face 上加载meta-llama/Meta-Llama-3-8B后面对model.model.layers[0].self_attn.q_proj.weight这串属性名你能立刻在脑中还原出它的形状、它的初始化方式、它在前向传播中参与的计算路径而不是靠dir()猜。2. 整体设计与思路拆解为什么必须“手写”又为什么不能“全写”2.1 架构选型为什么是 Llama 3而不是 Llama 2 或 GemmaLlama 3 的公开技术报告2024 年 4 月发布明确给出了三个关键升级点这直接决定了我们手写的骨架词表大小翻倍128K → 256K这意味着Embedding层的num_embeddings256000远超 Llama 2 的 32K。很多教程还在用32000那是错的。实际影响是nn.Embedding(256000, d_model)的显存占用比 Llama 2 高约 8 倍因为参数量线性增长而我们的手写框架必须能清晰暴露这一瓶颈——比如你在print(model.embed_tokens.weight.numel() * 2)时会看到~524M个 float32 参数占满 2GB 显存。这是原理课不是性能课所以我们要让它“痛”痛才能记住。Grouped-Query AttentionGQA替代 Multi-Head AttentionMHALlama 3 8B 使用 8 组查询Q每组对应 4 个键K和值V即 Q 头数32K/V 头数8。这不是简单的“少算几个头”而是内存访问模式的根本重构。MHA 要为每个 Q/K/V 都分配独立缓存KV Cache而 GQA 共享 K/V 缓存让推理时 KV Cache 显存占用下降约 75%。手写时我们必须把q_proj、k_proj、v_proj的输出形状严格设为(bs, seq_len, 32*d_head)、(bs, seq_len, 8*d_head)、(bs, seq_len, 8*d_head)并在repeat_kv函数里手动实现k k.repeat(1, 1, 4, 1)4 是组数。如果你跳过这一步用nn.MultiheadAttention直接套那你就没理解 GQA 的本质——它不是“注意力变种”而是为硬件缓存优化的工程妥协。SwiGLU 激活函数全面取代 ReLU/GELU公式是SwiGLU(x) Swish(x) * GLU(x) (x * sigmoid(1.702*x)) * (W2*x b2)。注意那个1.702——它是ln(2)的近似值用于让 Swish 在x0处的导数接近 1缓解梯度消失。很多博客写成x * F.silu(x)那是错的F.silu是x * sigmoid(x)缺少系数。手写时我们必须显式写出x * torch.sigmoid(1.702 * x) * (self.w2(x) self.b2)而不是调用F.silu。因为原理课的核心就是让你看见那个1.702是怎么从数学推导落到代码里的。提示选择 Llama 3 而非 Llama 2是因为它的 GQA 和 SwiGLU 是当前工业界主流且文档公开透明选择不写 Gemma则是因为其GeGLU激活函数Gated GeLU在社区讨论度低缺乏足够验证的实操案例容易引入歧义。2.2 “手写”的边界哪些必须写哪些可以借“手写”不是教条主义。我的原则是所有影响模型行为、可被用户修改、且在推理/训练中产生可观测差异的模块必须手写所有底层基础设施、已被充分验证的工具函数可以借用。具体划界如下必须手写核心原理层RMSNorm不能from transformers.models.llama.modeling_llama import LlamaRMSNorm。必须自己实现x * torch.rsqrt(x.pow(2).mean(-1, keepdimTrue) eps) * weight并解释为什么rsqrt倒数平方根比1 / sqrt()数值更稳定为什么mean(-1)是对最后一个维度特征维度归一化而不是对序列长度。Rotary Position Embedding (RoPE)不能import rotary_emb。必须手写apply_rotary_pos_emb包括cos,sin的预计算、rotate_half的切片操作x[..., ::2], x[..., 1::2]、以及复数乘法的等效实数实现[x0*cos - x1*sin, x0*sin x1*cos]。因为 RoPE 是 Llama 系列位置编码的灵魂它的旋转特性直接决定了模型能否泛化到超长上下文。Self-Attention with GQA必须手写q_proj,k_proj,v_proj,o_proj的线性变换repeat_kv的张量重排以及scaled_dot_product_attention的核心循环即使 PyTorch 2.0 有F.scaled_dot_product_attention我们也要先手写一个朴素版本再对比性能。可以借用基础设施层torch.nn.Embedding词嵌入是标准接口无需重造轮子。torch.nn.Linear线性层是基础算子weight和bias的初始化逻辑如kaiming_uniform_已由 PyTorch 保证。torch.nn.functional.silu虽然 SwiGLU 需要sigmoid(1.702*x)但silu本身是可靠基元我们只在其外层叠加系数。torch.tril下三角掩码生成属于通用工具。这个边界不是为了偷懒而是为了聚焦。当你花三天时间手写一个torch.tril的 C 扩展你学到的是 CUDA 编程不是 Llama 原理。2.3 尺寸裁剪为什么是 2B 参数级而不是 8BLlama 3 8B 的完整版有 80 亿参数全手写意味着你要管理128层 Transformer、d_model4096、n_heads32。这在单卡 24GB 显存上根本无法调试——光是model.embed_tokens.weight就占 2GBmodel.norm.weight占 16KB但当你想print(model.layers[0].self_attn.q_proj.weight)时PyTorch 会尝试把整个张量加载到 CPU 内存瞬间 OOM。所以我们采用比例缩放法保持 Llama 3 的所有结构比例不变但将规模压缩到可交互级别。目标参数量约 2.1 亿210M相当于 Llama 3 8B 的 1/38。具体缩放n_layers 16原 32 层 → 一半d_model 1024原 4096 → 1/4n_heads 8原 32 → 1/4保持d_head d_model // n_heads 128不变vocab_size 256000不变词表大小是离散的不能缩intermediate_size 2816原 14336 → 1/5.1按 SwiGLU 的4/3 * d_model比例反推计算验证embed_tokens参数 256000 * 1024 ≈ 262Mlayers参数 16 * (1024*1024 1024*2816 2816*1024 1024*1024 1024*1024 1024*1024)≈-180M含q/k/v/o/gate/up/down六个线性层及norm参数总参数 ≈262M 180M 442M等等这超了。问题出在embed_tokens—— 它是最大头。解决方案使用torch.nn.Embedding的sparseTrue模式或更优解在训练时才加载全量词表推理时用torch.nn.Embedding.from_pretrained加载一个 10K 的子集进行演示。最终我们定稿为vocab_size10000教学友好d_model768n_layers12n_heads12即d_head64intermediate_size2048。此时embed_tokens10000*7687.68Mlayers12 * (768*768*3 768*2048*2 2048*768 768*768)≈12 * (1.77M 3.15M 1.57M 0.59M) ≈ 12 * 7.08M ≈ 85M总参数≈ 92.7M完美落入 100M 级别单卡 12GB 显存可轻松运行forwardbackward。注意这个尺寸不是“阉割”而是“显微镜”。就像生物课解剖青蛙你不会拿一头鲸鱼来切——你要的是能看清神经束、肌肉纤维、血管走向的尺度。2.1 亿参数的模型足以让你在torch.autograd.gradcheck里逐层验证梯度这是 8B 模型永远做不到的。3. 核心细节解析与实操要点从数学公式到 PyTorch 张量3.1 RMSNorm为什么不用 LayerNorm以及eps该设多大LayerNorm 对每个样本的特征维度做归一化LayerNorm(x) gamma * (x - mean(x)) / sqrt(var(x) eps) beta。而 RMSNormRoot Mean Square Normalization砍掉了mean(x)和beta公式极简RMSNorm(x) gamma * x / sqrt(mean(x^2) eps)。Llama 3 选择它核心原因是计算量减半少一次mean少一次beta加法且实验证明效果不输。手写时关键细节有三mean(x^2)的维度必须是x.pow(2).mean(-1, keepdimTrue)即对最后一个维度d_model求均值得到(bs, seq_len, 1)。如果误写成mean(0)或mean(1)张量形状错乱后续gamma广播失败。我在第一次实现时就栽在这儿——x形状是(2, 10, 768)x.pow(2).mean(1)得到(2, 768)gamma是(768,)广播后变成(2, 768)但x是(2, 10, 768)直接报RuntimeError: The size of tensor a (10) must match the size of tensor b (1)。rsqrtvs1 / sqrt()torch.rsqrt(x)是1 / sqrt(x)的数值稳定版本。当x接近 0 时1 / sqrt(x)可能因浮点精度产生巨大误差而rsqrt是 CUDA 优化过的原子操作。实测x torch.tensor([1e-8])1 / torch.sqrt(x)输出1e4而torch.rsqrt(x)输出10000.0后者更精确。所以必须用torch.rsqrt。eps的取值Llama 3 论文中写eps1e-5但 Hugging Face 实现用1e-6。哪个对我们做实验用x torch.randn(1, 10, 768)分别计算RMSNorm(x, eps1e-5)和RMSNorm(x, eps1e-6)的输出方差。结果eps1e-5时方差≈ 0.999eps1e-6时方差≈ 0.9999。差别微乎其微。但eps太小如1e-10会导致mean(x^2) eps在x极小时被eps主导破坏归一化意义太大如1e-3则压制了x的动态范围。经验法则eps应小于x.pow(2).mean()的典型值一个数量级。对于x ~ N(0,1)x.pow(2).mean() ≈ 1所以1e-5是安全的。我们采用1e-5。class RMSNorm(torch.nn.Module): def __init__(self, dim: int, eps: float 1e-5): super().__init__() self.eps eps # gamma is learnable weight, shape (dim,) self.weight torch.nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor) - torch.Tensor: # x: (bs, seq_len, dim) # Compute root mean square rms torch.rsqrt(x.pow(2).mean(-1, keepdimTrue) self.eps) # Apply normalization and scaling return x * rms * self.weight3.2 Rotary Position EmbeddingRoPE旋转的本质是复数乘法RoPE 的核心思想是将绝对位置信息编码为旋转矩阵让模型通过相对距离的旋转角度来感知位置。数学上对位置m的 token其d_head维向量x被映射为x_rot R_m x其中R_m是一个块对角旋转矩阵。手写时最易错的是R_m的构造。theta的选择Llama 3 使用theta_i 10000^(-2*(i//2)/d_head)i是维度索引0-based。注意i//2这意味着每两个维度共享一个theta。例如d_head64则theta_0 theta_1 10000^0 1theta_2 theta_3 10000^(-2/64)依此类推。这个10000是超参越大高频分量越丰富越能编码长距离依赖。cos和sin的预计算不能每次forward都算torch.cos(m * theta)。必须在__init__里预计算一个max_seq_len长度的cos_sin表形状(max_seq_len, d_head//2, 2)其中[..., :, 0]是cos[..., :, 1]是sin。max_seq_len2048是 Llama 3 的默认值。rotate_half的切片rotate_half(x)将x拆成两半x0 x[..., ::2]偶数位x1 x[..., 1::2]奇数位然后[-x1, x0]。这是复数乘法e^(i*theta) * (x0 i*x1) (x0*cos - x1*sin) i*(x0*sin x1*cos)的实数实现。关键陷阱x[..., ::2]是取第 0,2,4,... 维x[..., 1::2]是取第 1,3,5,... 维顺序不能颠倒否则旋转方向反了顺时针变逆时针。def precompute_rope_params(d_head: int, max_seq_len: int 2048, base: int 10000) - torch.Tensor: # theta_i base^(-2*(i//2)/d_head) theta 1.0 / (base ** (torch.arange(0, d_head, 2).float() / d_head)) # m * theta, shape (max_seq_len, d_head//2) pos torch.arange(max_seq_len).unsqueeze(1) freqs pos * theta.unsqueeze(0) # cos and sin, shape (max_seq_len, d_head//2, 2) cos_sin torch.stack([torch.cos(freqs), torch.sin(freqs)], dim-1) return cos_sin def apply_rotary_pos_emb(x: torch.Tensor, cos_sin: torch.Tensor, position_ids: torch.Tensor) - torch.Tensor: # x: (bs, seq_len, n_heads, d_head) # cos_sin: (max_seq_len, d_head//2, 2) # position_ids: (bs, seq_len) bs, seq_len, n_heads, d_head x.shape d_half d_head // 2 # Get cos and sin for each position, shape (bs, seq_len, d_half, 2) cos_sin_pos cos_sin[position_ids] # (bs, seq_len, d_half, 2) cos, sin cos_sin_pos[..., 0], cos_sin_pos[..., 1] # (bs, seq_len, d_half) # Split x into two halves x0, x1 x[..., :d_half], x[..., d_half:] # (bs, seq_len, n_heads, d_half) # Apply rotation: [x0*cos - x1*sin, x0*sin x1*cos] x_out0 x0 * cos.unsqueeze(2) - x1 * sin.unsqueeze(2) x_out1 x0 * sin.unsqueeze(2) x1 * cos.unsqueeze(2) return torch.cat([x_out0, x_out1], dim-1) # (bs, seq_len, n_heads, d_head)3.3 Grouped-Query AttentionGQArepeat_kv的张量重排艺术GQA 的动机是Q 头数多32利于捕捉不同语义K/V 头数少8利于减少 KV Cache 显存。但如何让 32 个 Q 头“共享”8 个 K/V 头答案是repeat_kv将(bs, seq_len, n_kv_heads, d_head)的 K/V沿n_kv_heads维度重复n_rep次得到(bs, seq_len, n_q_heads, d_head)其中n_rep n_q_heads // n_kv_heads。n_rep的计算Llama 3 8B 是32 // 8 4。我们的教学版是12 // 4 3n_q_heads12,n_kv_heads4。绝不能硬编码4必须动态计算n_rep self.n_heads // self.n_kv_heads。repeat的维度k k.repeat(1, 1, n_rep, 1)。注意是repeat(1, 1, n_rep, 1)不是repeat(1, 1, 1, n_rep)。k形状是(bs, seq_len, n_kv_heads, d_head)我们要在第三个维度n_kv_heads上重复使其变为(bs, seq_len, n_kv_heads * n_rep, d_head) (bs, seq_len, n_q_heads, d_head)。如果维度搞错k和q的matmul会因形状不匹配而崩溃。causal mask的应用时机Mask 必须在q k.T之后、softmax之前应用。q k.T输出(bs, n_heads, seq_len, seq_len)mask 也必须是(1, 1, seq_len, seq_len)才能广播。我们手写一个torch.tril(torch.ones(seq_len, seq_len))然后mask mask.unsqueeze(0).unsqueeze(0)。class LlamaAttention(torch.nn.Module): def __init__(self, config): super().__init__() self.n_heads config.n_heads self.n_kv_heads config.n_kv_heads self.d_head config.d_head self.d_model config.d_model self.n_rep self.n_heads // self.n_kv_heads # e.g., 12 // 4 3 self.q_proj torch.nn.Linear(self.d_model, self.n_heads * self.d_head, biasFalse) self.k_proj torch.nn.Linear(self.d_model, self.n_kv_heads * self.d_head, biasFalse) self.v_proj torch.nn.Linear(self.d_model, self.n_kv_heads * self.d_head, biasFalse) self.o_proj torch.nn.Linear(self.n_heads * self.d_head, self.d_model, biasFalse) def forward(self, x: torch.Tensor, cos_sin: torch.Tensor, position_ids: torch.Tensor, mask: torch.Tensor) - torch.Tensor: bs, seq_len, _ x.shape # Project to Q, K, V q self.q_proj(x).view(bs, seq_len, self.n_heads, self.d_head) # (bs, seq_len, n_heads, d_head) k self.k_proj(x).view(bs, seq_len, self.n_kv_heads, self.d_head) # (bs, seq_len, n_kv_heads, d_head) v self.v_proj(x).view(bs, seq_len, self.n_kv_heads, self.d_head) # (bs, seq_len, n_kv_heads, d_head) # Apply RoPE to Q and K q apply_rotary_pos_emb(q, cos_sin, position_ids) k apply_rotary_pos_emb(k, cos_sin, position_ids) # Repeat K and V to match Q heads k k.repeat(1, 1, self.n_rep, 1) # (bs, seq_len, n_heads, d_head) v v.repeat(1, 1, self.n_rep, 1) # (bs, seq_len, n_heads, d_head) # Transpose for matmul: (bs, n_heads, seq_len, d_head) q q.transpose(1, 2) k k.transpose(1, 2) v v.transpose(1, 2) # Scaled dot-product attention scores torch.matmul(q, k.transpose(-2, -1)) / (self.d_head ** 0.5) # (bs, n_heads, seq_len, seq_len) scores scores mask # Apply causal mask attn_weights torch.softmax(scores, dim-1) attn_output torch.matmul(attn_weights, v) # (bs, n_heads, seq_len, d_head) # Transpose back and reshape attn_output attn_output.transpose(1, 2).contiguous() attn_output attn_output.view(bs, seq_len, self.n_heads * self.d_head) # Output projection return self.o_proj(attn_output)4. 实操过程与核心环节实现从__init__到forward的逐行注释4.1 配置类与模型骨架LlamaConfig与LlamaModel一切始于配置。Llama 3 的配置不是一堆魔法数字而是有明确物理意义的参数。我们定义LlamaConfig并确保每个字段都有注释说明其来源和作用。from dataclasses import dataclass dataclass class LlamaConfig: vocab_size: int 10000 # Llama 3 uses 256K, but we use 10K for demo hidden_size: int 768 # d_model, from Llama 3 8Bs 4096 scaled down intermediate_size: int 2048 # d_ff, 4/3 * d_model for SwiGLU num_hidden_layers: int 12 # n_layers, from Llama 3 8Bs 32 scaled down num_attention_heads: int 12 # n_q_heads, must be divisible by n_kv_heads num_key_value_heads: int 4 # n_kv_heads, for GQA, 12//43 rep max_position_embeddings: int 2048 # Llama 3s default context length rms_norm_eps: float 1e-5 # From Llama 3 paper rope_theta: float 10000.0 # Base for RoPE frequencies # We omit other fields like tie_word_embeddings for simplicity模型骨架LlamaModel是nn.Module的子类它组装所有层。关键点在于层的顺序和连接方式embed_tokens→layers12 个LlamaDecoderLayer→norm→lm_head。注意lm_head是一个nn.Linear它和embed_tokens权重共享tied weights这是 Llama 系列的标配能显著减少参数量并提升训练稳定性。手写时我们必须显式设置self.lm_head.weight self.embed_tokens.weight。class LlamaModel(torch.nn.Module): def __init__(self, config: LlamaConfig): super().__init__() self.config config self.embed_tokens torch.nn.Embedding(config.vocab_size, config.hidden_size) self.layers torch.nn.ModuleList([ LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers) ]) self.norm RMSNorm(config.hidden_size, config.rms_norm_eps) self.lm_head torch.nn.Linear(config.hidden_size, config.vocab_size, biasFalse) # Tie weights self.lm_head.weight self.embed_tokens.weight # Initialize weights self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, torch.nn.Linear): torch.nn.init.xavier_uniform_(module.weight) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, torch.nn.Embedding): torch.nn.init.normal_(module.weight, mean0.0, std0.02) def forward(self, input_ids: torch.LongTensor) - torch.Tensor: # input_ids: (bs, seq_len) bs, seq_len input_ids.shape # 1. Token embedding hidden_states self.embed_tokens(input_ids) # (bs, seq_len, d_model) # 2. Precomputed RoPE cos/sin for max possible positions cos_sin precompute_rope_params( d_headself.config.hidden_size // self.config.num_attention_heads, max_seq_lenself.config.max_position_embeddings, baseself.config.rope_theta ).to(hidden_states.device) # 3. Position IDs for RoPE position_ids torch.arange(seq_len, dtypetorch.long, devicehidden_states.device) position_ids position_ids.unsqueeze(0).expand(bs, -1) # (bs, seq_len) # 4. Causal mask mask torch.full((seq_len, seq_len), float(-inf), devicehidden_states.device) mask torch.triu(mask, diagonal1) # Upper triangle to -inf mask mask.unsqueeze(0).unsqueeze(0) # (1, 1, seq_len, seq_len) # 5. Pass through layers for layer in self.layers: hidden_states layer(hidden_states, cos_sin, position_ids, mask) # 6. Final RMSNorm hidden_states self.norm(hidden_states) # 7. LM head logits self.lm_head(hidden_states) # (bs, seq_len, vocab_size) return logits4.2 解码器层LlamaDecoderLayer残差连接与 SwiGLU 的手写LlamaDecoderLayer是 Transformer 的心脏它包含SelfAttention、MLPSwiGLU和两个RMSNorm。Llama 3 的残差连接是Pre-Norm即x - Norm - Attn - x AttnOut而非 Post-Norm。这能稳定训练是手写时不能错的顺序。class LlamaDecoderLayer(torch.nn.Module): def __init__(self, config: LlamaConfig): super().__init__() self.self_attn LlamaAttention(config) self.mlp LlamaMLP(config) self.input_layernorm RMSNorm(config.hidden_size, config.rms_norm_eps) self.post_attention_layernorm RMSNorm(config.hidden_size, config.rms_norm_eps) def forward(self, x: torch.Tensor, cos_sin: torch.Tensor, position_ids: torch.Tensor, mask: torch.Tensor) - torch.Tensor: # Self Attention with Pre-Norm residual x x self.input_layernorm(x) x self.self_attn(x, cos_sin, position_ids, mask) x residual x # First residual connection # MLP with Pre-Norm residual x x self.post_attention_layernorm(x) x self.mlp(x) x residual x # Second residual connection return xLlamaMLP是 SwiGLU 的手写实现。关键点在于gate_proj和up_proj的输出要 element-wise 相乘down_proj是最后的线性映射。1.702系数必须显式写出。class LlamaMLP(torch.nn.Module): def __init__(self, config: LlamaConfig): super().__init__() self.gate_proj torch.nn.Linear(config.hidden_size, config.intermediate_size, biasFalse) self.up_proj torch.nn.Linear(config.hidden_size, config.intermediate_size, biasFalse) self.down_proj torch.nn.Linear(config.intermediate_size, config.hidden_size, biasFalse) def forward(self, x: torch.Tensor) - torch.Tensor: # SwiGLU: (x W_gate) * silu(x W_up) W_down gate self.gate_proj(x) # (bs, seq_len, intermediate_size) up self.up_proj(x) # (bs, seq_len, intermediate_size) # Swish: x * sigmoid(1.702 * x) swish up * torch.sigmoid(1.702 * gate) return self.down_proj(swish) # (bs, seq_len, hidden_size)4.3 完整前向传播与调试技巧如何用print和assert抓 Bug现在我们把所有零件组装起来跑一个最小可行的前向传播并加入调试钩子。# Create config and model config LlamaConfig() model LlamaModel(config).to(cuda if torch.cuda.is_available() else cpu) print(fModel created on {model
