大语言模型（LLM）正无处不在地驱动着聊天机器人、写作工具和代码助手。但从零开始构建一个 LLM 本身是一件截然不同的事情。你不再仅仅是人工智能的使用者，而是成为了它的架构师。本文将结合 Sayak Dutta 的文章，深入探讨使用RLHF (Reinforcement Learning from Human Feedback，人类反馈强化学习) 栈训练和运行 10 亿参数 LLM 的实战经验，以及其中涉及的关键技术，包括模型架构、奖励模型和 PPO (Proximal Policy Optimization，近端策略优化)。
勇攀高峰：从用户到架构师的转变
构建 LLM 不再只是调用 API 或使用现成的库，而是需要深入理解模型的内部机制、训练流程和优化策略。这种转变意味着你必须掌握模型架构设计的精髓、奖励模型的构建方法以及 PPO 算法的应用。Sayak Dutta 的经历正是这种转变的缩影，他放弃了 Hugging Face 提供的便捷工具，选择自主编写训练逻辑、构建奖励模型，并应用 PPO 算法。这种“从零开始”的实践，不仅加深了对模型的理解，也让他体会到了对齐模型的重要性。这好比从一个只会使用计算器的人，变成一个能设计和制造计算器的人，其理解的深度和掌握的技能是完全不同的。
精巧设计：10 亿参数模型架构解析
文章中提到，该模型的配置目标约为 10.3 亿参数，并针对单个 RTX 4090 进行了优化。关键的 ModelConfig 参数如下：
ModelConfig(
vocab_size=50257,
hidden_size=1012,
num_layers=14,
num_attention_heads=22,
intermediate_size=3584,
num_experts=12,
num_activated_experts=2,
moe_layers_start=8,
max_seq_length=512,
dropout=0.1,
bias_update_speed=0.001,
balance_loss_weight=0.0001,
)
这种维度平衡有助于控制内存占用，同时为模型提供足够的表达能力。该模型架构的核心组件包括：
- 旋转位置嵌入（RoPE）： 赋予模型处理序列数据的能力，使其能够理解词语之间的相对位置关系。
- 多头潜在注意力（Multi-head Latent Attention）： 允许模型同时关注输入序列的不同部分，从而捕捉更丰富的语义信息。
- SwiGLU 激活函数： 一种新型激活函数，能够提高模型的表达能力和训练效率。
- 混合专家（MoE）层： 通过动态路由将不同的输入分配给不同的专家网络，从而提高模型的容量和泛化能力。
MoE 的集成是这项工作的亮点。作者实现了具有专家亲和力评分、偏差平衡和 top-k 选择的动态路由。以下是令牌路由的代码片段：
routing_scores = affinity_scores + self.expert_bias.unsqueeze(0)
top_k_scores, top_k_indices = torch.topk(routing_scores, self.num_activated_experts, dim=-1)
top_k_weights = F.softmax(top_k_scores, dim=-1)
对于每个 token，我们计算专家路由分数：
其中 是专家 的质心向量， 是一个可学习的偏差。选择 top-k 个专家，并使用加权输出组合它们。
MoE 的引入使得模型能够在不增加计算量的情况下扩展其容量。MoE 使其能够专业化——一些专家适应对话，另一些适应故事讲述或摘要。这种模块化使得以后的对齐变得更加容易。例如，我们可以想象一个拥有 10 个专家的模型，其中 3 个擅长生成代码，3 个擅长回答问题，4 个擅长创造性写作。通过动态路由，模型可以根据输入选择最合适的专家组合，从而提高性能和效率。
循循善诱：RLHF 流程详解
RLHF 的目标是训练一个能够生成符合人类偏好文本的模型。这个过程通常包括三个步骤：
- 监督微调（SFT）： 使用大量的文本数据对模型进行微调，使其具备基本的语言能力。
- 奖励建模： 训练一个奖励模型，用于评估模型生成的文本的质量。奖励模型通过比较“选择”和“拒绝”的完成情况来学习什么看起来更好。
- PPO： 使用 PPO 算法对模型进行优化，使其能够生成获得更高奖励的文本。
SFT 阶段的关键在于准备高质量的训练数据。作者使用 C4 数据集，并使用 AdamW 优化器和 1e-3 的学习率进行训练。这个阶段让模型学会了语言的流畅性，但缺乏意图性。
奖励模型的训练是 RLHF 的核心。作者使用偏好数据集（例如 Anthropic HH-RLHF），在冻结的基础模型之上训练奖励头。只有最后一层是可训练的。通过比较“选择”和“拒绝”的完成情况，模型学会了什么看起来更好。例如，如果模型生成的文本被人类评价为更有帮助、更无害，那么奖励模型就会给出更高的奖励。
PPO 阶段的目标是在奖励模型的指导下，对模型进行微调，使其能够生成符合人类偏好的文本。PPO 的核心思想是：强化好的完成，但保持新策略接近旧策略。
L(θ) = E_t [min(r_t(θ)A_t, clip(r_t(θ), 1-ε, 1+ε)A_t)]
其中：
- A_t 是优势估计。
- ε 是信任区域裕度（通常为 0.2）。
换句话说：“不要离你过去所说的太远，即使新的东西获得更高的奖励。”
作者使用的关键参数包括：
- KL 惩罚为 0.01（轻轻地束缚策略）。
- 剪辑比率为 0.2。
- 缓慢的 LR 为 5e-7。
- 没有单独的价值头（保持简单）。
结果？模型对齐——微妙但可衡量。它开始喜欢乐于助人。
步步为营：推理过程的经验教训
推理是模型部署的关键环节。作者分享了一些在推理过程中遇到的问题和解决方案：
- 安全反序列化： 使用
torch.serialization.safe_globals()来确保反序列化的安全性，防止恶意代码注入。 - 手动词汇表检查： 检查词汇表以避免形状不匹配。
- 回退解码： 实施回退解码以提高弹性。
这些看似微小的细节，往往会在最后导致整个系统的崩溃。例如，如果模型的词汇表与输入数据的词汇表不一致，就会导致模型无法正确处理输入数据，从而产生错误的结果。
洞见与启示：超越模型本身
训练一个语言模型就像构建一个思想。首先，你赋予它语法（SFT）。然后，你赋予它品味（奖励建模）。最后，你赋予它谨慎（PPO）。
奖励不仅仅是一个有效的模型——而是理解它为什么有效。没有什么可以取代这一点。不要只是使用模型。成为模型。
Sayak Dutta 的经历告诉我们，构建 LLM 不仅仅是技术挑战，更是一种思维方式的转变。我们需要深入理解模型的内部机制、训练流程和优化策略，才能真正掌握 LLM 的力量。 这也如同学习一门外语，不仅仅是记住单词和语法，更重要的是理解其背后的文化和思维方式。只有这样，才能真正流利地使用这门语言，并用它来表达自己的思想。
总结一下，这篇文章的核心在于强调了从零开始构建 LLM 的重要性和挑战性。通过自主设计模型架构，构建奖励模型，并应用 PPO 算法，我们可以更深入地理解 LLM 的工作原理，并更好地控制其行为。 这种“从用户到架构师”的转变，是 LLM 发展的必然趋势，也是我们未来需要努力的方向。 通过理解和掌握这些关键技术，我们可以更好地利用 LLM 来解决实际问题，并创造出更具价值的应用。