企业官网建设流程全解析

1. 模型主体怎么定：第一版60M试水时的“规模边界”计算

如果你正打算从零训练一个大模型，我猜你最开始的步骤可能跟我一样：盯着 GitHub 上那些 7B、13B 的 config 文件发呆，心想先抄哪个好。这个思路不能说错，但很容易让你在第一步就跑偏。因为 7B 的配置是给几百张 GPU 的集群准备的，你在单机四卡上复现它，只会得到一个永远填不满的显存和一张再也跑不动实验的卡。

我们当时在实验室第一次做“从零构建大模型”这个系列实验时，第一位导师给我们的约束其实很朴素：60M 参数左右，序列长度 4096，跑通一个完整的前向和反向，再把 loss 压到能看。你别小看这个 60M，它其实是一个非常好的起点——它足够小，小到单张卡就能训；又足够大，大到所有“大模型”该有的组件——词表、位置编码、RMSNorm、Attention mask、Gradient Accumulation、混合精度——一样都缺不了。

那“规模边界”怎么算？我当时用了一个很实用的经验公式：参数量 (N \approx L \times (12d^2 + 13d) \times \left(1 + \frac{2V}{L d^2}\right))。其中 (L) 是层数，(d) 是 hidden size，(V) 是词表大小。你算一下就会发现，3.4B 词表对参数量的影响在 60M 这个量级是不可忽略的：假如你的词表是 32K，词嵌入矩阵就是 (32768 \times 768 \approx 25M) 参数，占了总体的四成。这里藏着第一个容易踩的坑——如果只看 decoder block 的参数量，你会严重低估你对显存的需求，因为 embedding 的梯度在反向传播里同样要吃显存。所以早期定模型规模时，我建议直接把公式列出来，而不是拍脑袋。

在这个阶段，我们选了这样一个基准配置，后续所有实验都从这里出发：

为什么序列长度一上来就设 4096？不是因为长序列更好，而是因为 RoPE 的位置编码在短序列上测不出外推问题，你迟早要面对长序列的训练稳定性，还不如一开始就把它暴露出来。别以为 60M 模型小就随便弄，架构上的每一个设计，最后都会被放大到 1B、7B 时重新回来折磨你。

2. 前向计算图拆开看：RMSNorm、RoPE、SwiGLU 与 GQA 的实现细节

很多教程讲到 Transformer 就一句话：“self-attention + FFN 堆几层”，但等你真去写代码时，你会发现每一步都有三个选项，而且每个选项都能改变训练曲线的形状。下面我把我们实际用过的四个关键算子逐个拆开，顺便把代码贴出来。

2.1 RMSNorm：为什么不用 LayerNorm

你做分类任务时可能很习惯 LayerNorm，但到生成模型里，RMSNorm 是更常见的选择。原因在于 LayerNorm 的均值减法在残差连接里其实是冗余计算——Transformer 的残差结构本身就在做中心化，你再在每一层强制均值归零，不仅费算力，而且在 bf16 下还会引入额外的数值偏移。RMSNorm 只做缩放，不做平移：

import torch import torch.nn as nn class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): # x: [B, T, D] rms = torch.rsqrt(x.to(torch.float32).pow(2).mean(-1, keepdim=True) + self.eps) return (x * rms).to(x.dtype) * self.weight

这里有一个细节：x必须先转成 fp32 计算方差再转回 bf16。如果你偷懒直接对 bf16 的 x 求pow(2)，当训练到深层的某个位置时，x 的数值范围一旦超过 bf16 的精度上限（大约 (3.39 \times 10^{38})），你会得到一个 inf，然后整个 loss 瞬间飞到 NaN。我们在 1B 模型上就因为这个吃了大亏，后面会细说。

2.2 RoPE 旋转位置编码

RoPE 的核心思想是用旋转矩阵把位置信息“藏”进向量的方向里，而不是简单相加。这样设计的好处是：两个 token 的 attention score 只跟它们的相对位置有关，模型外推能力比绝对位置编码好得多。

它的实现不复杂，关键点在于频率的生成：

import torch def precompute_freqs_cis(dim, max_seq_len, base=10000.0): # 这里 dim 必须是 head_dim，不是 hidden_size freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) t = torch.arange(max_seq_len, device=freqs.device) freqs = torch.outer(t, freqs) return torch.polar(torch.ones_like(freqs), freqs) # 用复数表示旋转 def apply_rotary_pos_emb(q, k, freqs_cis): q = torch.view_as_complex(q.float().reshape(*q.shape[:-1], -1, 2)) k = torch.view_as_complex(k.float().reshape(*k.shape[:-1], -1, 2)) q = (q * freqs_cis.unsqueeze(0).unsqueeze(2)).flatten(3).to(q.dtype) k = (k * freqs_cis.unsqueeze(0).unsqueeze(2)).flatten(3).to(k.dtype) return q, k

你可能注意到我把复数那一套写出来了，虽然torch.polar看起来有点吓人，但它比手动展开cos/sin再拼接要清晰得多。一个小技巧是：freqs_cis 的 shape 是 [max_seq_len, head_dim/2]，你用的时候要广播到 [B, T, heads, head_dim/2]，尤其是在 GQA 之后 key 的 head 数是 query 的 (\frac{1}{4})，广播维度稍不留神就会错位。

2.3 SwiGLU 激活函数：换掉 GELU 的真实收益

FFN 层的激活函数选择，直接决定同样参数量下的表达能力和训练吞吐。如果你只是换一下激活函数，从 GELU 换成 SwiGLU，在相同参数量下，SwiGLU 的训练 loss 大概能低 0.05~0.1 左右，原因是它等于把原来的单线性层拆成两个门控分支，相当于增加了非线性拟合的自由度。

class SwiGLU(nn.Module): def __init__(self, dim, hidden_dim): super().__init__() self.w1 = nn.Linear(dim, hidden_dim, bias=False) self.w2 = nn.Linear(dim, hidden_dim, bias=False) self.w3 = nn.Linear(hidden_dim, dim, bias=False) def forward(self, x): gate = torch.nn.functional.silu(self.w1(x)) hidden = self.w2(x) return self.w3(gate * hidden)

注意这里传统 FFN 通常是hidden_dim = 4 * dim，但 SwiGLU 因为有两个权重矩阵，实际参数是原来的 1.5 倍。如果你想保持参数量不变，hidden_dim要设为 ( \frac{2}{3} \times 4d = \frac{8}{3}d )，取整后一般是dim * 8 // 3。我们第一版直接抄了 LLaMA 的配置用了8/3这个比例，跑出来参数量才没有超预期。

2.4 GQA：为什么 7B 以下也别急着用 MHA

多头注意力（MHA）每个头都有自己的 key 和 value，好处是表达能力强，坏处是显存和算力会随着 head 数线性涨。在小模型阶段，把 MHA 换成 GQA（Grouped Query Attention）并不是为了省显存——影响不大——而是为了让后面的 KV Cache 不会成为推理阶段的瓶颈。哪怕你现在的训练脚本根本不管 KV Cache，你也应该从一开始就把它设计进去，否则模型训完才发现推理速度被内存带宽拖垮，那时候再改架构等于重新训练。

GQA 的实现思路是：query 还是 8 个头，但 key 和 value 只有 4 个头，然后每个 key/value 头被 2 个 query 头共享。

num_query_groups = 4 # 8 query heads -> 4 kv heads q = q.view(B, T, num_heads, head_dim).transpose(1, 2) # [B, 8, T, 64] k = k.view(B, T, num_kv_heads, head_dim).transpose(1, 2) # [B, 4, T, 64] v = v.view(B, T, num_kv_heads, head_dim).transpose(1, 2) k = k.repeat_interleave(num_heads // num_kv_heads, dim=1) # 重复为 [B, 8, T, 64] v = v.repeat_interleave(num_heads // num_kv_heads, dim=1)

这里有个很多人忽略的细节：repeat_interleave会在显存里真正把 k/v 复制一份，而torch.expand只是共享 view。Flash Attention 的实现一般要求 k/v 是 contiguous 的，所以即使你用了 expand，它在 kernel 内部也会 materialize。而在小显存卡上训练时，这一步会影响你的最大 batch size。我们后来改成在attn_mask里直接处理 group 映射，省了一整块中间显存，但这只在纯手工写的 CUDA kernel 里才能这么干，用 PyTorch 原生 API 的话老实 repeat_interleave 就好。

3. 训练稳定性实测：一次 Loss Spike 的完整排查链路

单讲“如何搭模型”你可能会觉得不过瘾，真正让你头皮发麻的是训练到一半突然 loss 从 3.2 跳到 43.7，然后永远回不来。这里我必须把我们在某公司某次 1B 模型训练里排查 loss spike 的完整过程记录下来，因为这个问题如果你没亲眼见过，根本不知道从哪下手。

3.1 现象：grad norm 从 0.8 跳到 12

2000 步之前一切正常，val loss 在平滑下降。第 2047 步时，loss 还在 3.19 附近，第 2048 步之后，loss 直接变成 43.6，grad norm 从 0.8 一路涨到 12。我们在日志里看到这个数字的时候，第一反应是数据出了问题——很可能某个 batch 混入了脏数据，导致 loss 放大。

我们于是把 data loader 改成固定顺序，让每一步吃到的数据是完全可复现的，然后重新训练。结果：loss spike 依旧在第 2048 步左右出现。这说明问题不在数据侧，而是模型在某个临界点发生了数值不稳定。

3.2 排查：AdamW 里的 eps 太小，bf16 的下溢

第二步我们开始怀疑优化器。我们用 AdamW，初始参数是 (\beta_1=0.9, \beta_2=0.95, \epsilon=1e-8)。这个配置在 fp32 下很常见，但在 bf16 下有一个问题：bf16 的精度只有大约 7-8 位十进制有效数字，当你某个参数的梯度过小，bf16 直接把它下溢成 0，而 AdamW 里的 (\epsilon=1e-8) 意味着它会用这个接近 0 的梯度做分母，导致更新幅度被异常放大。

我们在代码里加了 grad norm 的 logging，发现 spike 前 50 步的 grad norm 越来越小，到 spike 那一步突然变成 12。典型的 pre-spike 信号：norm 持续变小，随后单步突然暴涨。懂行的人这时候已经知道问题不是数据，而是某一层某个小梯度被 bf16 下溢，然后 AdamW 用极小的 eps 把更新量放大了。

3.3 根源：RMSNorm 在 bf16 下的灾难链

但为什么是第 2048 步而不是更早？我们把 checkpoint 拿到 CPU 上逐层跑，发现真正的问题在 RMSNorm 和 GQA 的交互上。前面我写的 RMSNorm 里有一个x.to(torch.float32)的步骤——我们最初实现时偷懒，直接对 bf16 的 x 做 rsqrt。这就导致：当某一路残差经过若干层之后，它的 RMS 值变得非常大（超过 bf16 可表示的正常范围），bf16 连这个值本身都不准了。而 2048 这个步数对应的是我们 learning rate warmup 结束的节点，lr 达到峰值，对数值扰动的敏感度正好也达到峰值。

这个问题的完整链条是：

1. RMSNorm 在 bf16 下计算方差不准，导致 output scale 偏大；
2. scale 偏大的 output 进入 attention，和 GQA 的repeat_interleave相互作用，产生很大的内积；
3. 很大的内积过 softmax 后变成接近 one-hot 的分布，反向传播时某个梯度接近 0 但不是 0；
4. bf16 把这个接近 0 的梯度下溢成 0；
5. AdamW 的 (\epsilon) 太小，分母接近 0，于是对这部分参数输出了超大更新；
6. 这一大步把模型参数推出正常区域，loss 暴涨。

3.4 修复：三层防御

我们最后在代码里加了三样东西，缺一不可：

# 防御层 1：RMSNorm 的 fp32 计算 rms = torch.rsqrt(x.to(torch.float32).pow(2).mean(-1, keepdim=True) + self.eps) # 防御层 2：GQA repeat 后做一次 clip k = k.repeat_interleave(num_heads // num_kv_heads, dim=1).clip(-20, 20) # 防御层 3：AdamW 的 eps 调大，并开启 fp32 master weight # optimizer = AdamW(model.parameters(), lr=3e-4, eps=1e-6)

你可能好奇clip(-20, 20)会不会损伤表达能力。实测下来，key 向量经过 RoPE 之后的值域确实很少超过 20，clip 只会在异常时触发，起的是一个保险丝作用。而 AdamW 的 eps 从 1e-8 调到 1e-6，在 fp32 训练里几乎无感，但在 bf16 混合精度里能显著减少下溢放大。修完之后我们再训练到 3000 步，loss 曲线基本平滑，再没有出现单步跳变。

希望这篇模型训练心得能快速帮你定位问题。以后遇到 loss spike，第一件事永远是把 grad norm 和每层输出 norm 打出来，别急着怪数据。90% 的数值问题，都能在梯度统计里找到先兆。

4. 小规模预训练验证：从 60M 到 1B 的参数与数据配比

架构定下来之后，下一个问题就是：训练数据要用多少 token？用什么样的 batch size？这直接决定了你的实验要跑多久、能验证什么结论。

我们说一个简单的规则：训练 token 总数 ≈ 模型参数量 × 20。也就是说 60M 模型配 1.2B token 是最低配，1B 模型配 20B token 才比较充分。这个比例来自所谓“Chinchilla 最优计算量”的结论——在给定算力预算下，模型尺度和数据量要同时增长才划算，而不是只堆参数。我们的实验里，同样 1B 模型，用 10B token 比用 20B token 的最终 loss 差大约 0.15，这个差距在下游任务上会放大很多。

4.1 一套亲测稳定的训练配置

下面是我们跑 300M 模型时用的一套配置，你可以直接复制过去：

你可能会问：为什么 micro batch size 只有 8？因为序列长度 4096 时，单卡显存主要在 attention 的中间激活上，而不是模型参数。我们把模型参数、梯度和 AdamW 状态算了笔账：300M 模型，参数 fp32 是 1.2GB，AdamW 的 fp32 状态再加 2.4GB，bf16 参数 0.6GB，总共约 4.2GB，模型本身占不到多少；但 attention 部分一个 4096×4096 的 score 矩阵就是 64MB，8 个头乘起来轻松上 512MB，再乘 batch 8，就是 4GB。激活内存是隐形的成本，模型越小它越容易成为主要瓶颈。

4.2 用扩展定律验证：60M 模型怎么预测 1B 的表现

这个阶段最有价值的一件事，是用小模型验证扩展定律。我们当时训练了 60M、130M、300M 三个规模的模型，各用相同的数据配比（参数×20 token），记录它们的最终 loss。结果画出来基本是一条直线：log(loss) ≈ −a × log(参数)。意思是，在 log-log 坐标系里，每 10 倍参数量，loss 下降约 0.35 左右。

这个拟合结果有什么用途？它可以帮你提前估算 1B 模型能到多少 loss，而不是等训完 20B token 才知道。我们在 300M 模型训完后，用拟合曲线预测 1B 模型的 loss，准确度在 ±0.05 以内。要注意的是，这个预测有一个前提：小模型和大模型用完全相同的架构、数据配比和训练设置。如果你在 300M 或 1B 之间改了激活函数或位置编码，拟合关系立刻失效。

我建议任何从零构建大模型的实验，至少同时训两个规模（比如 60M 和 300M），不是因为学界要求，而是因为你知道 60M 的曲线长什么样，才能判断 300M 是否发生了偏差。如果 300M 的最终 loss 明显偏离拟合线，说明你的模型在某个规模上出现了未预期的行为——这种模型再往上扩也会有问题。

5. 工程填坑篇：小显存卡上跑大模型的七个低效陷阱

这一章最后，我想记录几个我们在实际训练中反复踩的工程坑。它们不会让训练崩掉，但会让你的卡始终跑在 30% 利用率，而你觉得一切正常。这些坑没有写在任何教程里，你只能靠 profiling 和肌肉记忆去发现。

第一个坑是DataLoader 的异步预取。你训练小模型时，数据处理往往比 GPU 计算快很多，于是你可能会忽略 IO 瓶颈。但当你切到 4096 序列长度后，每个 batch 解压和 tokenize 的开销变得不可忽略。我们用torch.profiler一看，GPU 有 23% 的时间在等 CPU 喂数据。解决方案是 prefetch：把num_workers=4起步，再用pin_memory=True。别小看这个，它能直接让 throughput 提升 15%。

第二个坑是梯度累积的时机。很多人把loss.backward()放在每个 micro batch 里，然后每隔 N 步才optimizer.step()。这在数学上是等价的，但在 bf16 混合精度下，梯度在累积过程中会不断加总，如果中间某个 micro batch 的梯度特别大，fp32 master weight 里的更新可能变化过大。我们的经验是：累积不要超过 16 步，一旦超过，就把 grad clip 的阈值从 1.0 调低到 0.5。

第三个坑是Key 的 repeat_interleave 后没有 contiguous。前面提到过，在 DDP 环境下，带 padding 的 gather 操作会触发隐式的contiguous()拷贝。如果 k/v 没有在构造时直接 contiguous，后面 Flash Attention 里的实际计算会因为内存布局不连续而多一次额外的 copy，这个 copy 在小 batch 下看起来没啥，但在长序列时能占到 20% 的额外显存带宽。

第四个坑是AdamW 里的 weight_decay 混淆了参数组。我们最初图省事，对所有参数做了统一的 weight_decay，结果 bias 和 LayerNorm/RMSNorm 里的 scale 也被 decay 了。这会导致训练后期 loss 降不动。正确做法是像 LLaMA 那样：只有 weight 矩阵进 weight_decay，所有 norm 的 scale 和 bias 不进。实现非常简单，PyTorch 里就是两组参数分开传给 optimizer。

第五个坑是余弦退火步数设错了。如果你 300M 模型训练了 20B token，global step 总数是 (20B / 0.5M = 40000) 步。但如果你在 config 里直接写max_steps=40000，然后在某次实验中改了 batch size，那你实际训练的总 token 数完全变了，而退火的步数却没有跟着改。你要设的不是 step 数，而是 total_tokens，然后在每次跑之前用它算出 step 数。

第六个坑是未开 cuDNN benchmark。这里面有很多 model-specific 的 trick，但 PyTorch 默认的 benchmark 模式是关的，注意力实现的某些分支会被自动 fallback 到较慢的算法。我们在 60M 模型上开了torch.backends.cudnn.benchmark = True之后，吞吐量涨了大约 8%，虽然不是质变，但白捡的算力谁不想要。

第七个坑是用 DDP 时会出现的向量化同步问题。小规模模型 DDP 是够用的，但你会观察到 loss 曲线不如单卡平滑，因为每张卡的梯度在 all-reduce 之前是独立算的。正常现象。如果你想要单卡一样的确定性复现，就把model_parallel和zero1的某些参数设得保守点，否则每次跑出来的微小曲线差异很容易让你误以为是 Agent 代码写错了。

最后，一个长期有用的习惯：在你开始 24 小时长训之前，先用 500 步的 smoke test 把 loss curve、grad norm、lr schedule、throughput 全部打点记录一遍，然后盯着前 500 步看 10 分钟。如果前 500 步是好的，后面 95% 的概率也是好的；反之前 500 步有任何一个数异常，后面几小时都是白白浪费。我们后来所有的大规模训练前，都会刻意跑这一步，比什么监控脚本都管用。

从 60M 走到 1B，过程确实琐碎，但每一步踩的坑都是之后往上扩的资本。希望这一章能帮你省掉几天的排查时间，直接站上一个能稳定复现的平台。下一篇文章我准备写数据配比和 tokenizer 的细节——那又是一个容易被忽略但能拉开最终效果差距的重头戏。