从零构建LLM的完整工程路径：数据管道、模型架构与训练循环

大语言模型的训练通常被视为需要数百张 GPU 和海量资金的工程，但实际上，理解其核心机制并在可控资源下复现，对掌握现代 AI 系统至关重要。本文基于一个开源的完整实现，梳理从零构建 LLM 的工程路径，涵盖数据管道、模型架构、训练循环到文本生成的关键环节。

数据管道：从原始语料到训练张量

训练数据的质量直接决定模型上限。该项目选用 The Pile 数据集，这是一个包含 22 个不同领域数据源、总计 825GB 的多样化语料库，涵盖书籍、论文、代码和网页内容。数据预处理的核心任务是将文本转换为模型可消费的张量格式。

具体流程分为三步：首先下载.jsonl.zst格式的压缩文件，每行是一个 JSON 对象，包含text字段和元数据；然后使用 OpenAI 的 tiktoken 库（r50k_base 编码器）将文本 token 化，该编码器与 GPT-3/GPT-2 兼容，词汇表大小为 50304；最后将 token 序列存储为 HDF5 格式，利用其可扩展数据集特性支持高效随机访问。值得注意的是，每个文本序列末尾需添加<|endoftext|>特殊 token，作为序列边界标记，这对模型学习文本连贯性至关重要。

数据加载器采用迭代器模式，从 HDF5 文件中按批次抽取样本。每个批次包含输入序列xb和目标序列yb，后者是前者右移一位的结果，用于计算下一个 token 的预测损失。通过随机打乱样本索引，确保训练过程的随机性。

模型架构：模块化 Transformer 实现

该项目采用经典的 Decoder-only Transformer 架构，完全基于 PyTorch 从零实现，不依赖 HuggingFace 等高层封装。核心组件按层次递进：

MLP 模块作为前馈网络的基础单元，将输入嵌入维度扩展 4 倍后经 ReLU 激活，再投影回原维度。这种 "扩展 - 压缩" 结构增强了模型的非线性表达能力。

单头注意力实现因果自注意力机制。通过线性层分别计算 Query、Key、Value 投影，利用下三角掩码矩阵确保模型只能关注当前位置及之前的 token，防止信息泄露。注意力权重计算采用缩放点积形式，除法因子为嵌入维度的平方根，用于稳定梯度。

多头注意力并行运行多个单头注意力，每个头的维度为n_embed // n_head，输出沿特征维度拼接。这种设计允许模型同时捕获不同类型的语义关系。

Transformer Block将多头注意力与 MLP 堆叠，配合层归一化和残差连接。层归一化先于子层应用（Pre-LN），有助于稳定深层网络的训练。残差连接则为梯度流动提供捷径，缓解梯度消失问题。

顶层 Transformer 类整合 token 嵌入、位置嵌入、N 个 Transformer Block 堆叠、最终层归一化，以及将隐藏状态映射到词汇表维度的语言模型头。位置嵌入使用可学习的绝对位置编码，最大序列长度由context_length参数决定。

训练循环：参数配置与优化策略

训练配置需权衡模型规模与硬件资源。该项目提供从 1300 万到 20 亿参数的完整参数表，单张消费级 GPU（如 RTX 4090 24GB）可训练约 40 亿参数模型。

对于 20 亿参数模型，推荐配置为：嵌入维度 2048、16 个注意力头、64 个 Transformer Block、上下文长度 512。训练批次大小 32，使用 AdamW 优化器，初始学习率 5e-4，在指定步数衰减至 5e-5。总训练步数 20 万，每 1000 步执行一次验证评估。

训练过程采用滑动窗口平均损失监控训练稳定性，验证集用于检测过拟合。模型状态、优化器状态和训练元数据定期保存为检查点，支持中断恢复。值得注意的是，大模型（20 亿参数）在训练初期损失波动较大，需要更谨慎的学习率调度；而小模型（1300 万参数）收敛更快，适合快速验证架构正确性。

文本生成：采样策略与规模效应

文本生成通过自回归方式实现：给定初始 token 序列，模型预测下一个 token 的概率分布，采样后追加到序列，重复直至达到最大生成长度或遇到结束标记。该项目采用 softmax 归一化配合多项式采样，相比贪婪解码能产生更多样化的输出。

模型规模对生成质量影响显著。1300 万参数模型已能生成语法正确、拼写准确的短句，如 "The villages were directly linked to cities in China" 这类连贯表达。而 20 亿参数模型在长文本生成上表现更稳定，能维持更长的上下文一致性。但实验也表明，若架构设计不当，大模型可能出现过拟合，生成内容反而不如精心训练的小模型连贯。

可落地参数清单

基于上述实现，以下是可直接应用于类似项目的参数配置：

数据预处理

• 分词器：tiktoken r50k_base（词汇表 50304）
• 序列结束标记：<|endoftext|>
• 存储格式：HDF5，dtype=int32

模型架构（20 亿参数级）

• n_embed: 2048
• n_head: 16
• n_blocks: 64
• context_length: 512

训练配置

• 优化器：AdamW
• 初始学习率：5e-4
• 衰减后学习率：5e-5
• 批次大小：32
• 评估频率：每 1000 步

硬件参考

• RTX 4090 24GB：最大约 40 亿参数
• A100 40GB：最大约 60-80 亿参数
• T4 16GB：适合 1300 万 - 20 亿参数模型

从零训练 LLM 的核心价值不在于与工业级模型竞争，而在于深入理解数据流、梯度传播和生成机制的本质。这种端到端的工程实践，是成为 AI 系统工程师的重要阶梯。

资料来源

• GitHub: FareedKhan-dev/train-llm-from-scratch — 完整的从零训练 LLM 实现
• 论文: Attention is All You Need — Transformer 架构原始论文

mlops