训练3.8B LLM至0.384 CORE，成本998美元 – Hugo Vergnes

# Hugo Vergnes | 训练一个3.8B参数的LLM：CORE评分0.384，成本998美元 来源：https://hugovergnes.github.io/little-lm-3-8b/ 在“nanoGPT玩具”和“你需要一个研究实验室”之间，存在着一个广阔且鲜被描述的领域——一个人花几千美元就能训练出一个有意义的模型。 我渴望亲眼见证语言和理解能力从随机权重中涌现，并学习那些唯有从零开始才能掌握的知识。这个项目在夜间编写，在5090显卡上调试，最终在租赁的B200上完成。它深受Andrej Karpathy的nanochat（https://github.com/karpathy/nanochat）启发。 最终结果是一个拥有3.8B参数的模型，在**CORE基准测试中获得0.384分**，使用65B tokens训练了43小时，**耗资998美元**。 以下内容将介绍哪些方法有效，哪些无效，以及我仍不明确的部分。 | 模型 | 参数量 | Tokens | 硬件 | 时间 | 成本 | CORE评分 | | :--- | :--- | :--- | :--- | :--- | :--- | :--- | | GPT-2 (OpenAI) | 1.5B | — | — | — | — | 0.2565 | | nanochat d26 | ~561M | 11.2B | 8× H100 | ~3h | — | ~0.258 | | nanochat d32 | ~1B | — | 8× H100 | ~33h | ~$1000 | 0.310 | | little-lm 3.8B (1024 ctx) | 3.848B | 57.3B | 8× B200 | 35.9h | $820 | 0.338 | | **little-lm 3.8B (2048 ctx)** | **3.848B** | **65.3B** | **8× B200** | **43h** | **$998** | **0.384** | 我的模型比nanochat d32更大，并花费了相似的挂钟时间。B200的单位计算成本效益优于H100。但大致花费与nanochat的1000美元配置相同，我的模型取得了显著的领先。这是一个鼓舞人心的数据点，表明在实验室或拥有数百万计算预算的大公司之外所能达到的成就。随着技术前沿的推进，1000美元能带你走得越来越远。 --- ## 设置 我将little-lm构建为一个配置驱动的小型解码器专用LLM训练框架。每次运行完全由YAML文件指定：模型、数据集、优化器、调度器、回调函数。组件通过全局注册表自注册并按名称解析，因此更换优化器或数据集只需一行配置更改。 **优秀的基础设施几乎能立即收回成本。** 普通的软件工程规范（如关注点分离、清晰的接口、可替换的组件）在AI工作中非常重要。起初它花了一点成本，之后又花了几次来修复不良契约或次优性。但这次时间投入在你遇到的第一个收敛问题上就能收回。我发现，最好的基础设施是几乎从不需要手动编辑代码的基础设施。如果你能阅读配置并确切理解发生了什么，没有隐藏机制，那就意味着你做得很好。以下报告的结果得益于能够用三行YAML差异来表示实验，而不是创建一个分支。 最终模型采用Llama风格：RMSNorm、RoPE、GQA（24个查询头，8个KV头）、relu2 MLP、QK-norm、logit softcap、逐层可学习的残差缩放器，以及ResFormer风格的值嵌入。 | 组件 | 参数量 | | :--- | :--- | | Token嵌入 | 154.5M | | LM头（未绑定） | 154.5M | | 28个解码器层 | 2,818.7M | | 值嵌入（14个表） | 721.2M | | **总计** | **3.848B** | 值得注意的是，值嵌入占了参数数量的19%。14个`vocab × kv_dim`的表，每隔一层放置一个。 --- ## 结果 ### 早期实验 在好的运行之前，有很多次失败的尝试。 我在FineWeb-Edu上训练了一个858M参数的Llama模型，使用了16.4B tokens，在单个A100上运行了5.8天。使用AdamW，学习率2.5e-4，余弦衰减至零，5%的预热，通过梯度累积实现批大小256，上下文长度2048。 结果：**PIQA得分60.45%**。GPT-2 124M的得分约为63%。我花了六天的计算资源，构建了一个比2019年的小七倍模型更差的东西。生成内容重复且近乎无意义。 损失曲线说明了问题。 1. **余弦衰减至零**：曲线在约70%的步骤后完全趋于平坦。最后30%的计算预算几乎没有产出，因为学习率可能过低。线性冷却在更长时间内保持了有用的学习率。 2. **峰值学习率过于保守**：2.5e-4对于858M参数来说太低了。对于这种小模型，你可以更激进一些。 3. **所有地方都用AdamW**：在这个规模上，Muon对于矩阵参数每token的表现应该明显更好。事实上，在消融实验中很快就证明了这一点。 4. **数据**：FineWeb-Edu还不错，但并非最佳选择。 那次事后分析得出了五个改变。它们共同构成了上面那次运行和大幅超越GPT-2的模型之间的区别。 **梯形学习率调度**：预热5%。保持平稳，最后从峰值的50%线性冷却至5%。关键在于模型持续学习直到结束，而不是在尾声阶段空转。在3.8B的运行中，评估损失在最后一步*仍在下降*，这正是858M运行未能产生的行为。 **矩阵参数使用Muon，其他地方使用AdamW**：Muon每步更慢（牛顿-舒尔茨正交化不是免费的，在浅层累积基准测试中大约慢25%），但这个成本只在优化器步骤支付一次：在7次梯度累积下，稀释到约4%。与总运行时间相比，整体收敛速度要快得多。 **使用ClimbMix代替FineWeb-Edu**：这带来了收敛速度的巨大飞跃，正如Karpathy所发现的那样。 **FP8 + 词表填充**：通过`torch._scaled_mm`对三个GEMM（通用矩阵乘法）进行动态张量级缩放来实现FP8训练，并将词表从50,257填充到50,304（64的倍数）以使张量核心满意。总共+33%的吞吐量，主要来自fp8。 **1024上下文代替2048**：上下文减半，固定内存下批次大小大约翻倍。每token的吞吐量几乎不变。我们仍然受限于MLP（多层感知机），这是使用硬件效率良好的一个好迹象。下面我们讨论上下文长度对模型的影响。 以下是完整的运行记录： | 步骤 | Tokens | 评估损失 | CORE | | :--- | :--- | :--- | :--- | | 2,500 | 5.7B | 2.3278 | 0.2389 | | 5,000 | 11.5B | 2.2072 | 0.2752 | | 7,500 | 17.2B | 2.1571 | 0.2934 | | 10,000 | 22.9B | 2.1269 | 0.3104 | | 12,500 | 28.7B | 2.1075 | 0.3147 | | 15,000 | 34.4B | 2.0710 | 0.3224 | | 17,500 | 40.1B | 2.0395 | 0.3294 | | 20,000 | 45.9B | 2.0160 | 0.3267 | | 22,500 | 51.6B | 1.9963 | 0.3345 | | **25,000** | **57.3B** | **1.9868** | **0.3384** | 稳态下约480,000 tokens/sec，因此57.3B tokens在33小时内完成。挂钟时间是35.9h。差异来自CORE评估，每次大约15分钟（运行中共进行了10次），占总时间的7%。 以2048 token上下文重新运行相同配方，得分**0.3840**。几乎所有的差距都归因于一些非常依赖上下文的任务。 在GPU上：**92% SM活动率，40% SM占用率。** 高活动率意味着SM几乎从未空闲。没有数据加载器饥饿或网络等待，这是将分片下载到本地而非流式传输的收益，后者会让我们容易受到Hugging Face小规模网络挂起的影响。低占用率是连续大型GEMM的典型表现：矩阵乘法内核有意用占用率换取寄存器分块大小。计算受限且供应充足，这是一个良好的信号，表明我们正在有效利用硬件，能够将每一分钱转化为更好的模型。 这大约相当于每个B200持续**1,047 TFLOP/s**，或者说，相对于Blackwell的密集FP8峰值，实现了约**25%的MFU（模型FLOP利用率）**。（相对于bf16峰值，它达到50%，这个数字更重要一些，因为并非所有线性层都在FP8下运行。） 分布式策略是传统的DistributedDataParallel（分布式数据并行）。在单节点3.8B模型上，梯度通信从来不是瓶颈，分片优化器机制被证明是不必要的。 ### 提高吞吐量 租用GPU并不便宜，在工作中，你经常考虑模型质量而非成本。但当花的是自己的钱时，吞吐量突然变得重要得多。 在租用节点之前，我在单张RTX 5090上做了大量工作。基线858M模型，bf16，编译后：**26,144 tok/s**。最终：**37,621 tok/s**。 * **FP8（+25%）**：三个GEMM（1个前向，2个反向）都使用FP8和动态张量级缩放。需要SM90+，但吞吐量提升非常可观。 * **词表填充（累计+33%）**：填充50,257 → 50,304，增加了47个未使用的嵌入行，解锁了快速张量核心路径。几乎免费。 * **融合线性交叉熵（累计+44%）**：Liger的`FusedLinearCrossEntropyLoss`将lm_head矩阵乘法融合到损失中，并在内部分块，因此永远不会具体化完整的`(B*T, vocab)` logits张量。在相同批大小下，头对头测量，它**慢了6%**： | 配置 | 吞吐量 | 显存 | | :--- | :--- | :--- | | 基线CE，批大小6 | 34,724 tok/s | 27,852 MiB | | 融合CE，批大小6 | 32,952 tok/s | 19,630 MiB | | 融合CE，批大小8 | 35,979 tok/s | 24,028 MiB | | 融合CE，批大小10 | **37,621 tok/s** | 28,872 MiB | 即使每步更慢，它也节省了大量显存（在我的5090上是8 GB），因此增加的微批次大小足以弥补损失的6%。Claude很快拒绝了它，因为慢了6%，但总的来说，这是获得额外吞吐量的好方法。 * **非门控MLP**：在小模型上去掉门投影（SwiGLU → relu2，两个矩阵乘法代替三个）：从183,035 tok/s提升到214,173 tok/s，并减少6 GB显存。消融实验的一个注意点：SwiGLU的2.75中间比例*不*适用于relu2。模型表现明显更差。非门控应使用4×比例。 * **bf16主权重**：将优化器主权重保持为bf16而非fp32，将显存减少了27%，在1.5B配置上将吞吐量从640K提升到1.4M tok/s。这是巨大的加速，2.2倍。质量代价是真实的，但很小：在4,000步时，CORE为0.22 vs 0.23。当你优化每美元的能力时，谨慎的数据类型处理是可用杠杆最高且讨论最少的旋钮之一。 * **硬件**：相同代码，150M模型，FP8：RTX 5090达到184,662 tok/s，B200达到477,440 tok/s。仅硬件就带来2.59倍提升，尚未计算额外显存让你进一步增大批次大小。 ### 无效的方法 * **使用灵活注意力（flex attention）进行文档边界掩码**：将文档打包成一个序列允许token跨越边界关注，因此我正确地修复了它：每token的文档ID，并掩码注意力，使每个token只能关注其当前文档。这很优雅，但我删除了所有这些。Andrej Karpathy也发现，在BOS对齐的打包下，跨文档泄漏不会使情况变得更糟。最佳拟合打包在约10行代码内替代了它，注意力回到了无条件的`F.scaled_dot_product_attention(..., is_causal=True)`。我相信这也取决于数据集和训练文档。 * **Liger RMSNorm和RoPE**：RoPE在微基准测试中快了2.2倍，但端到端吞吐量*没有可测量的变化*。在这个规模上，RoPE并非关键计算瓶颈的一部分。RMSNorm甚至比PyT2.9内置的`F.rms_norm`更慢（0.41ms vs 0.25ms）。两者都恢复了，不值得增加复杂性。 * **Nanochat风格的初始化**：嵌入使用`N(0, 0.8)`，线性权重均匀分布，输出投影零初始化，使残差流从纯恒等变换开始，LM头使用`N(0, 0.001)`。理论上比GPT-2到处使用`N(0, 0.02)`好得多。损失曲线起点稍低，两条曲线在约1,500步后重叠。没有可测量的质量差异。我保留了它，但是为了美观，而非证据。 * **流式数据集**：对于入门很棒，对于真实运行则错误。即使网络看起来健康，本地分片也提供了2-3%的更多吞吐量，偶尔的网络波动代价远高于此。对于运行超过几小时的情况，在训练开始时付费下载一次是值得的。 ### 值嵌入的消融研究 值嵌入对于3.8B模型是721M参数。我使用相同配置但设置`value_embeddings: false`训练了相同的模型，并与原始运行（我已经付费了）进行比较，直到12,500步和29B tokens。 | 参数量 | 12.5K步损失 | 12.5K步CORE | 吞吐量 | | :--- | :--- | :--- | :--- | | 值嵌入开启 | 3.848B | **2.1075** | **0.3147** | 479,445 tok/s | | 值嵌入关闭 | 3.128B | 2.1171 | 0.3047 | 477,908 tok/s | 损失改善0.46%，CORE改善3.2%，参数增加19%。吞吐量相同，因为值嵌入是查找操作。它们消耗内存和优化器状态，但本质上不消耗FLOP。 两个有趣的发现： * **值嵌入购买了相当于约1,200个训练步的效果**：如何定价：在步骤10,000到12,500之间，我的基线损失下降了0.0194，因此2,500步大约购买了这么多。值嵌入的优势是0.0096，大约是一半——我们称之为25,000步中的1,200步。因此，增加19%的参数价值相当于增加约5%的训练量。 * **CORE的移动幅度大约是损失的七倍**（3.2% vs 0.46%），并且在训练过程中差距稳步缩小。如果你使用CORE来做决策，这很重要：它是一个准确度指标，因此决策边界附近的项目会因微小的logit变化而翻转，并且它相对于随机基线居中，在分数仍低时放大了相对差异。 值嵌入对于小模型很有用，且几乎不消耗吞吐量。在这上面花费一点显存，为模型提供了一种偏向于某些概念的形式，这可能对CORE有用。 --- ## 讨论 ### 误导性的微基准测试 我们可能会被诱惑相信1024 token的上下文足以获得高CORE分数。回顾每个任务的日志，对于一些非常依赖上下文的任务，这个结论是错误的。 22个CORE任务中有3个任务的提示词在1024 token中根本无法容纳： | 任务 | 提示词被截断数 | 步骤 2.5K | 步骤 25K | | :--- | :--- | :--- | :--- | | squad | 10570 / 10570 (100%) | 0.1478 | **0.0000** | | boolq | 3265 / 3270 (99.8%) | 0.5798 | 0.5131 | | bigbench_language_id | 9965 / 10000 (99.7%) | 0.2454 | 0.2538 | SQuAD是引人注目的一个。它不是停滞，而是单调衰减至恰好零：0.1478 → 0.0617 → 0.0099 → 0.0007 → 0.0000。模型在这个任务上训练时间越长，表现*越差*，这通常不是模型应有的表现。 两个细节解释了这一点。SQuAD在DCLM bundle中是一个**10-shot**任务，因此每个提示词是10个示例加上真正的那个。在我的评估数据上，中位数是1,998个token。没有一个能放进1024。当提示词过长时，我的测试框架保留*最后*的`max_seq_len` token。 测试段落位于末尾，所以总是保留了下来；一个测试样本只有约169个token。被截断的是那10个演示。模型看到了段落和问题，但几乎从未看到那些教会它预期输出格式的示例。由于SQuAD是基于与黄金答案的精确token匹配来评分的，流畅的散文每次都得零分。 这也解释了下降的原因。早期，一个高熵模型偶尔会生成一些简短且通用的东西，碰巧匹配了。随着它变得更明确，它会致力于格式良好的续写，偶然的匹配就消失了。有趣的是，语言能力变强反而让它更难偶然猜对。 boolq展示了相同模式的更温和版本。它在步骤10,000达到峰值（0.6294）然后下降到0.5131。语言识别任务则根本没有从随机水平上移动。 简而言之，**0.338是在22个任务中3个因与模型质量无关、仅因输入上下文长度大小而得分接近零的情况下测量得到的。** ### 更大上下文的影响 正如我们所看到的，如果我们想要尽可能高的CORE分数，我们需要更大的上下文。但这对训练吞吐量有影响。 上下文长度加倍，微批次大小减半以保持显存恒定，因此每个优化器步骤的token数保持不变。我在约28,000步时停止，以节省最后几小时的租金，因此学习率衰减从未完全完成，下面的数字是一个下限。 **CORE从0.3384提升到0.3840。** 1024和2048上下文运行的评估损失和CORE 在步骤20,000时，两次运行具有相同的评估...