大语言模型的训练过程常被封装为黑盒,开发者通过调用 API 或加载预训练权重即可使用。然而,要真正理解注意力机制如何捕捉长程依赖、残差连接如何稳定深层网络训练、以及分词器如何将文本映射为模型可处理的离散单元,从零实现并训练一个完整模型仍是不可替代的学习路径。
Fareed Khan 开源的train-llm-from-scratch项目提供了一个端到端的解决方案,支持从百万到十亿参数规模的模型训练。该项目基于 PyTorch 实现完整的 Transformer 架构,使用 The Pile 数据集进行训练,并提供了从数据下载到文本生成的全链路脚本。本文将解析这一训练管道的核心环节,为希望深入理解 LLM 内部机制的开发者提供可落地的工程参考。
数据工程:从原始语料到训练样本
训练数据的质量直接决定模型的上限。该项目选用 The Pile 数据集 —— 一个包含 22 个不同来源、总计 825GB 的多样化语料库,涵盖书籍、学术论文、代码仓库和网页内容。数据以jsonl.zst格式存储,即 JSON Lines 结合 Zstandard 压缩,每条记录包含text字段和元数据。
数据预处理流程包含三个关键步骤。首先是下载与解压,项目提供脚本从 HuggingFace 拉取数据分片,单个训练文件约 11GB。其次是分词编码,使用 OpenAI 开源的tiktoken库(r50k_base编码器,与 GPT-3/GPT-2 同款)将文本转换为整数序列,并在每个样本末尾添加<|endoftext|>标记以标识序列边界。最后是 HDF5 持久化,将编码后的 token 序列存储为 HDF5 格式,利用其支持动态扩展和快速随机访问的特性,为后续训练迭代提供高效的数据加载。
在内存受限的环境下,开发者可通过max_data参数限制处理的样本数量,或使用较小的验证集分片进行超参数调试。对于 1300 万参数模型的训练,约 1-5GB 的文本数据已足够产生可观测的语法学习能力。
架构实现:模块化构建 Transformer
该项目的核心价值在于提供了清晰、模块化的 Transformer 实现,每个组件均可独立理解与测试。
** 多层感知机(MLP)** 作为前馈网络,负责在注意力层之后引入非线性变换。实现上采用 "扩展 - 激活 - 压缩" 结构:先将输入从n_embed维度扩展至4 * n_embed,应用 ReLU 激活函数,再投影回原始维度。这种设计增加了模型的表达能力,同时保持参数量可控。
** 单头注意力(Single Head Attention)** 是理解自注意力机制的关键。每个头维护独立的 Query、Key、Value 投影矩阵(无偏置),通过计算 Query 与 Key 的点积得到注意力权重,经缩放(除以sqrt(head_size))和 Softmax 归一化后,与 Value 相乘得到输出。项目通过注册tril缓冲区实现因果掩码,确保模型仅关注当前位置及之前的 token,满足语言生成的自回归特性。
** 多头注意力(Multi-Head Attention)** 并行运行多个单头注意力,每个头处理n_embed // n_head维度的子空间。这种多头设计允许模型同时关注不同类型的语义关系 —— 例如语法结构、指代消解或主题连贯性。所有头的输出在最后一维拼接,恢复至n_embed维度。
Transformer 块将上述组件组合,并引入层归一化(LayerNorm)和残差连接。每个块先对输入进行层归一化,经多头注意力处理后与输入相加(残差连接),再经层归一化送入 MLP,最后再次残差相加。这种 "先归一化后计算" 的 Pre-Norm 结构配合残差连接,有效缓解了深层网络的梯度消失问题。
最终模型整合 token 嵌入、位置嵌入、堆叠的 Transformer 块、输出层归一化和语言模型头。位置嵌入与 token 嵌入相加,为模型提供序列顺序信息;语言模型头将最终表示投影至词汇表维度,输出每个位置的下一个 token 预测概率。
训练策略:参数配置与硬件适配
项目提供了从 1300 万到 20 亿参数的完整配置参考,开发者可根据硬件条件选择合适规模。
对于1300 万参数模型(T4 GPU 可训练),推荐配置为:上下文长度 128、嵌入维度 128、8 个注意力头、1 个 Transformer 块。该配置下模型约含 1300 万参数,训练约 20 万步即可收敛。学习率采用分段衰减策略:初始值 5e-4,在 5 万步时降至 5e-5,有助于模型先快速探索解空间,再精细调整参数。
对于20 亿参数模型(需 A100 或 RTX 4090),配置扩展为:上下文长度 512、嵌入维度 2048、16 个注意力头、64 个 Transformer 块。训练 batch size 设为 32,每 1000 步在训练集和验证集上评估平均损失,监控过拟合迹象。
硬件适配方面,项目提供了详细的 GPU 显存与可训练模型规模对照表。16GB 显存(如 T4、V100)可支持约 15 亿参数模型训练,24GB 显存(RTX 3090/4090)可扩展至 35-40 亿参数,而 40GB A100 可训练 60-80 亿参数模型。对于显存受限的情况,可通过减小context_length或batch_size换取更大的模型宽度或深度。
训练循环采用 AdamW 优化器,配合 HDF5 数据迭代器实现高效的批量采样。数据迭代器支持随机打乱和无限循环,每个 epoch 自动重新随机化样本顺序。损失跟踪使用滑动窗口平均(默认 64 步),平滑训练曲线的波动。
生成与评估:规模与质量的权衡
模型训练完成后,可通过generate方法进行自回归文本生成:将输入序列截断至context_length,取最后一个位置的 logits 经 Softmax 采样得到下一个 token,循环直至达到指定生成长度。
项目对比了 1300 万参数和 20 亿参数模型的生成效果,揭示了一个重要发现:并非参数越多越好。1300 万参数模型在短上下文下能生成语法正确、拼写准确的句子,如 "The villages were directly linked to cities in China"。而 20 亿参数模型在未经充分调优的情况下,反而产生语义断裂的长句,如 "blacksmith, musician and boutique hospitality and inspire the strain delivered Canadians"。
这一现象说明,大模型需要更深层的架构设计和更精细的训练策略,否则容易过拟合训练数据或陷入局部最优。对于特定应用场景(如邮件撰写、代码补全),1300 万至 5 亿参数的模型经过领域数据微调,往往比通用大模型更具实用价值,且可在消费级 GPU 上本地部署。
实践建议与局限
对于希望复现该项目的开发者,建议遵循渐进式路径:首先训练 1300 万参数模型验证流程正确性,然后逐步扩大规模。数据预处理阶段应仔细检查 token 序列的边界标记和 HDF5 文件完整性,避免因数据损坏导致训练异常。
该项目的局限在于未实现分布式训练、梯度累积或混合精度训练等工业级优化技术,因此不适合直接用于生产级大模型训练。然而,正是这种简洁性使其成为理解 Transformer 原理的理想教学工具 —— 每个模块的输入输出维度清晰可追踪,损失曲线与生成质量的对应关系直观可见。
从零训练 LLM 的过程揭示了现代大模型背后的工程权衡:数据质量与规模的平衡、模型容量与计算成本的取舍、通用能力与领域专精的抉择。掌握这些底层原理,有助于开发者在使用预训练模型时做出更明智的架构选择和调优决策。
参考来源
- Fareed Khan, "train-llm-from-scratch", GitHub, 2025. https://github.com/FareedKhan-dev/train-llm-from-scratch
内容声明:本文无广告投放、无付费植入。
如有事实性问题,欢迎发送勘误至 i@hotdrydog.com。