OpenLanguageModel: 面向教育与科研的可读可组合小型语言模型预训练

# 面向教育与研究的可读可组合小语言模型预训练
来源：https://arxiv.org/html/2607.16669
###### 摘要

OpenLanguageModel \(OLM\) 是一个开源 PyTorch 库，用于构建和预训练小语言模型，同时保持其内部机制可见。在 OLM 中，模型代码的呈现方式与架构本身一致：组件是普通的模块，而 `Block`、`Residual`、`Repeat` 和 `Parallel` 则描述了这些组件如何连接。由此产生的模型可以不做任何修改，从教学笔记本直接迁移到完整的预训练运行或研究消融实验。OLM 将这个可读的模型层与分词器、本地及流式数据集、优化、混合精度、回调函数、检查点，以及支持 CPU、单 GPU 和单节点多 GPU 的硬件感知执行相连接。我们通过从示意图到代码追踪 GPT-2，启动一个 FineWeb-Edu 训练脚本，替换一个注意力组件，并让 `AutoTrainer` 自动配置可用机器，来演示完整的流程。该包包含九个熟悉模型家族的 27 个预设，以及从 LM 基础知识到架构研究的文档。验证结果包括：与独立参考实现高度一致、348M 参数规模下四 GPU 弱扩展效率达 90.6%、精简的架构编辑空间以及积极的早期可用性结果。OLM 采用 MIT 许可证，可通过 PyPI、GitHub 及其文档网站获取。

## 1 引言

小语言模型 \(SLMs\) 将完整的语言建模工作流程置于课堂、单台工作站或专注的研究项目中可及的范围。它们规模足够大，能够实践真正的分词、数据流式读取、优化、检查点和分布式训练，同时又足够小，以至于一个人仍然能够理解并修改整个模型。

OpenLanguageModel \(OLM\) 将这两个目标——理解模型和认真运行它——整合到同一个 PyTorch 库中。其模型定义遵循用于解释 Transformer 的示意图。嵌入、注意力、归一化、前馈网络、残差路径、重复层和分支仍然作为普通模块保持可见。然后，同一个对象连接到完整的 SLM 预训练栈，将课程直接带入实验。

这种设计服务于三类用户：希望从头构建 LM 的学生和教育者、需要一个可读预训练方案的实践者，以及希望在保持周边系统稳定的情况下进行局部组件实验的研究人员。用户可以从一个命名的参考模型开始，从组件组装模型，插入手写的 `nn.Module`，或将生成的模型带入自己的 PyTorch 循环中。

OLM 将三件事保持在一起：镜像模型的架构代码、端到端硬件感知的 SLM 预训练，以及能够成长到研究阶段的学习路径。本文直接遵循这个工作流程。我们首先读取一个模型，训练它，改变一个思路，并扩展运行规模；然后我们打开系统设计，并给出关于正确性、定制化、扩展性和可用性的紧凑验证。

## 2 OLM 导览

首次使用的用户可以通过 `pip install openlanguagemodel` 安装 OLM。随后，演示将跟随一个模型从源代码到保存的检查点的完整路径。

#### 读取模型。

图1 (https://arxiv.org/html/2607.16669#S2.F1) 将 GPT-2 示意图与其 OLM 定义放在一起。对应关系是字面意义上的：词元嵌入和位置嵌入输入到重复堆叠的残差注意力和前馈块中，接着是输出头。打开一个 Llama 风格的定义，可以看到相同的结构词汇，只是组件换成了 RMSNorm、分组查询注意力、RoPE 和 SwiGLU。完整的架构在定义中一目了然。

参考图注

图 1：GPT-2 的示意图和 OLM 模型。顺序、嵌套、残差路径和层重复在可执行的 PyTorch 定义中清晰可见。
#### 训练同一个对象。

图2 (https://arxiv.org/html/2607.16669#S2.F2) A 部分展示了一个紧凑的预训练脚本：选择一个分词器，流式读取 FineWeb-Edu，构建一个熟悉的模型，并将普通的 PyTorch 对象传递给 `AutoTrainer`。训练器提供训练循环、混合精度、梯度累积、学习率调度、指标和检查点。本地文本数据集或用户自己编写的循环可以插入到相同的工作流程中，而模型保持不变。

A. 将模型连接到完整的预训练流程

```
tok = HFTokenizer("gpt2")
data = FineWebEduDataset(
    tok, context_length=1024, streaming=True)
loader = DataLoader(data, batch_size=8)

model = GPT2Model(
    vocab_size=tok.vocab_size,
    embed_dim=768, num_layers=12,
    num_heads=12, max_seq_len=1024)

trainer = AutoTrainer(
    model, AdamW, loader,
    device="auto", context_length=1024,
    learning_rate=3e-4, grad_accum_steps=8)
trainer.train(epochs=1, max_steps=1000)
```

B. 进行一次局部架构更改

```
# 本次运行的基线组件：
attention = GroupedQueryAttention(
    d_model, heads, kv_heads, context)

# 对于 ALiBi 消融实验，只需替换
# 上面这一行赋值：
# attention = MultiHeadAttentionwithALiBi(
#     d_model, heads, max_seq_len=context)

layer = Block([
    Residual(Block([
        RMSNorm(d_model), attention])),
    Residual(Block([
        RMSNorm(d_model), SwiGLUFFN(d_model)])),
])
```

图 2：OLM 演示中的两个时刻。左侧面板将可读模型连接到训练栈。右侧面板更改了注意力规则，同时保持了周围的层和训练路径不变。为节省空间省略了导入语句。
#### 改变一个思路。

图2 (https://arxiv.org/html/2607.16669#S2.F2) 的右侧部分将研究工作流程具体化。注意力对象是 Transformer 层中一行代码，其余部分保持不变。用 ALiBi 注意力替换分组查询注意力后，残差连接、归一化、前馈路径、数据管道、优化器和训练器都保持完整。新的注意力规则也可以继承 `AttentionBase` 或作为任何普通的 `nn.Module` 进入块中。

#### 使运行适配机器。

使用 `device="auto"`，OLM 报告检测到的设备、可用内存、估计模型占用空间以及选择的执行策略。同一个调用可以在 CPU 或一个 GPU 上运行，并为单节点多 GPU 配置 DDP 或 FSDP。在训练过程中，用户可以看到损失、困惑度、学习率和吞吐量；回调函数可以添加验证或实验跟踪。演示以保存、重新加载模型并从模型采样结束。

## 3 演示背后的系统

### 3.1 架构作为对象图

OLM 使用 `torch.nn.Module` 对象来表示组件和连接。`Block` 存储一个有序的 `ModuleList`；`Residual` 计算 `x + f(x)`；`Repeat` 从一个工厂函数创建独立参数化的层；`Parallel` 在可选的合并操作之前将同一输入发送到多个分支。由于组合子与注意力或前馈层共享相同的模块入/模块出接口，因此它们可以自然地嵌套。并行残差路径可以写成 `Residual(Parallel([attention, ffn]))`；异构堆栈只是不同层的一个显式列表。

对象图就是模型。它可以被打印、遍历、放置在设备上、优化、保存，或使用标准的 PyTorch 工具嵌入到更大的手写模块中。训练器直接消费这个图，因此源码、检查和执行共享同一个表示。自定义控制流仍然可以使用传统的 `forward` 方法。

### 3.2 可复用组件和参考模型

OLM 将注意力、位置编码、嵌入、归一化、前馈层、输出头和结构组合子作为独立的导入暴露出来。配置化的 GPT-2 (https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/openai/gpt2.py)、Llama 2 (https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/meta/llama2.py)、Llama 3 (https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/meta/llama3.py)、Qwen 2.5 (https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/alibaba/qwen2.py)、Phi-3 (https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/microsoft/phi3.py)、Phi-4 (https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/microsoft/phi4.py)、Gemma 2 (https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/google/gemma2.py)、OLMo (https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/allenai/olmo.py) 和 OPT (https://github.com/openlanguagemodel/openlanguagemodel/blob/main/src/olm/models/facebook/opt.py) 的实现都使用了这些相同的组件，并提供了 27 个命名的预设来对应熟悉的参数化。家族块和模型骨架与预设维度保持分离，使得参考模型既是一个现成的起点，也是一段可读的示例代码。

组件层包含显式的教育实现，以及由 PyTorch 缩放点积注意力支持的高效路径。现代细节如 RoPE、ALiBi、分组查询注意力、RMSNorm、SwiGLU、交替局部注意力和混合专家前馈在进入模型时都保持可见。

### 3.3 连接的预训练栈

创始团队在其研究项目中训练了超过 40 个小于十亿参数的 SLM，这种重复使用塑造了 OLM 的训练路径。该库提供本地和流式数据集、Hugging Face 分词器集成、数据加载、优化器、损失函数、预热和衰减调度、混合精度、梯度累积和裁剪、回调函数、吞吐量日志以及检查点保存/重新加载。

`AutoTrainer` 检查 CPU 和 CUDA 设备、GPU 数量和内存、以及估计的模型占用空间，然后根据平衡、速度或内存导向策略选择基础训练器、DDP 或 FSDP。用户可以直接选择策略，或将模型保留在自己的循环中。架构和训练保持分离，因此组件实验可以继续使用相同经过测试的数据和优化路径。

### 3.4 文档作为系统的一部分

文档位于 https://openlanguagemodel.github.io/openlanguagemodel/，从词元、嵌入、注意力和下一个词元预测开始；继续到一个可运行的小模型，包括训练、生成和保存/重新加载；然后打开块系统、现代架构、分布式训练和 API 参考。图3 (https://arxiv.org/html/2607.16669#S3.F3) 总结了三个入口点。初学者路径中引入的代码与后来研究路径中使用的代码是相同的。

学习 词元 → 嵌入 → 注意力 → 首个 LM 训练 本地 / FineWeb-Edu 数据 → AutoTrainer → 检查点 研究 替换注意力、归一化、FFN 或连接方式 相同的 `torch.nn.Module` 图、数据接口和训练工具

图 3：三个入口点，同一段代码路径。学习、预训练和研究使用相同的可检查模块和连接训练栈。

## 4 紧凑的系统验证

评估验证了导览中展示的能力：模型覆盖范围和参考保真度、局部架构更改、单节点扩展以及早期用户的体验。表1 (https://arxiv.org/html/2607.16669#S4.T1) 给出了简要版本；协议和打包记录见附录A (https://arxiv.org/html/2607.16669#A1)。

表 1：验证总结。
#### 覆盖范围和参考保真度。

所有九个模型家族的简化配置都通过了前向传播、分析参数计数、绑定嵌入和逐位检查点往返检查。所有 27 个命名预设都通过了配置和参数公式检查，本地测试套件通过了 149 个测试。为了进行数值检查，我们将独立的 Hugging Face 实现的权重复制到四层 FP32 的 GPT-2、Llama 3 和 Qwen 2.5 变体中。在每个家族三个种子的情况下，最大绝对 logit 差值为 \(3.576 \times 10^{-7}\)，最大损失差值为 \(4.768 \times 10^{-7}\)，选定的梯度余弦相似度保持在 \(0.9999999999998\) 以上。

#### 局部架构更改。

我们在新的用户脚本中实现了六种下游编辑，同时保持框架、训练和数据文件不变：并行残差、异构块、选定层 MoE、RoPE 到 ALiBi 替换、跨家族块混合以及 N 路学习合并。在 OLM、LitGPT 和 Pico 中完成的四个任务中，新脚本分别包含 117、171 和 195 行代码。在所有六个 OLM–LitGPT 任务中，总行数分别为 167 和 256。这么小的代码量反映了暴露组件和组合子被复用。

#### 单节点扩展。

一个 348,447,744 参数的 Llama-3 风格模型在一个、两个和四个 A100-SXM4-80GB GPU 上进行了弱扩展测试，使用 BF16、AdamW、PyTorch SDPA，每个 GPU 2048 个词元。三次运行的平均吞吐量分别为 17,785、32,426 和 64,427 词元/秒。四 GPU 运行相对于单 GPU 的理想扩展保持了 90.6% 的效率，而峰值分配内存在两个和四个 GPU 上保持在每 GPU 10.47 GB。

#### 早期用户体验。

在一项去标识化的 20 名 OLM 用户调查中，平均系统可用性量表得分为 73.8（95% bootstrap 置信区间 69.4–78.4）。所有 20 名受访者都同意架构可以理解，并且 OLM 与普通 PyTorch 配合自然；19 人同意架构更改是局部化的。在 17 名能够比较工作量的受访者中，14 人认为 OLM 中的架构级别更改更容易。

## 5 定位与可用性

OLM 构建在 PyTorch 的命令式模块系统 [5] 之上，并补充了几个强大的语言模型生态系统。Transformers 提供了对预训练模型的广泛访问 [6]；LitGPT 提供了紧凑的训练配方 [4]；LMFlow 侧重于基础模型的微调和推理 [2]；Pico 支持假设驱动的 SLM 研究 [3]。OLM 的重点是连接这些关注点的路径：可读的架构代码、引导式学习、完整的 SLM 预训练以及同一库中的局部研究编辑。

在撰写本文时，OLM v2.2.0 是最新版本；该项目由七位贡献者积极开发，并在 MIT 许可下分发。该包可从 https://pypi.org/project/openlanguagemodel/ 安装；源码托管在 https://github.com/openlanguagemodel/openlanguagemodel；网站 https://openlanguagemodel.github.io/openlanguagemodel/ 提供初学者课程、首个模型教程、训练指南、架构指南、参考模型源码和生成的 API 文档。这些公开入口使演示的系统可用于教学、研究和扩展。

## 6 结论

OLM 从一个简单的承诺开始：用于解释语言模型的代码，在需要训练和修改该模型时仍然应该有用。学生可以跟随 GPT-2 中的箭头，实践者可以将同一个对象连接到流式数据和硬件感知的训练，而研究人员可以在保持运行其余部分不变的情况下替换注意力行。

该演示展示了从端到端的路径。OLM 结合了可读的 PyTorch 组合、完整的 SLM 预训练套件、熟悉的参考模型以及随用户成长的文档。MIT 许可的包、源码和教程现已可用于教学、实验和新贡献。

## 局限性

验证范围限于演示中显示的系统。数值对等性使用了简化模型，GPU 研究使用了合成词元并在一个 A100 节点上使用 DDP。