数据截至 2025-11-24(上游 commit 090253dac668)

OLMo — 全开源预训练项目的完整生命周期

30 秒导读： OLMo 是 AI2（Allen Institute for AI）的「彻底开放」大模型项目：别人开放权重，它连训练代码、预训练数据（Dolma）、每 1000 步一个的中间 checkpoint、W&B 训练日志、数据顺序文件全部公开。这个仓库就是当年真正用来训练 OLMo-1 / OLMo-2（7B/13B）的训练代码 + 真实训练配置。想搞清楚「一次认真的工业级预训练到底由哪些部件组成、挂了怎么恢复、数据怎么混、训到一半怎么评测」，这是最完整的公开样本。

1. 这是什么（零基础也能懂）

一句话定义

OLMo 是一个 LLM 预训练框架 + 一整套真实训练配方：它包含一个配置驱动的 transformer 训练栈（模型 / 数据 / 训练器 / 评测 / 生成），以及训练 OLMo-2 系列模型实际使用的全部 YAML 配置。

它解决谁的什么问题

设想你想回答这些预训练研究者天天面对的问题：

一个 7B 模型在 4 万亿 token 上训练，数据按什么顺序喂？挂了重启后怎么保证不重不漏？
训练到一半 loss 突然尖峰（spike），当时的数据和超参到底是什么？
别人的论文写「我们用 cosine 调度、4T token」——但真实的配置里每一个字段是什么值？

闭源项目里这些问题无处可查。OLMo 的答案是把一切都摆出来：这个仓库装代码和配置，Hugging Face 上放着 Dolma 数据集和几千个中间 checkpoint，W&B 上放着训练曲线，连每个 rank 每一步读了哪些数据都写成了 TSV 文件（scripts/train.py:243 的 data-indices/rank{N}.tsv.gz）。

它能做什么

能力	具体支持	位置
模型定义	配置驱动的 decoder-only transformer（RoPE/ALiBi、RMSNorm、QK-norm、SwiGLU、MQA/GQA）	`olmo/model.py`
数据管线	内存映射 token 流、确定性全局 shuffle、实例级重复过滤、文档内注意力掩码	`olmo/data/`
训练	FSDP / DDP / 单机三种策略，micro-batch 梯度累积，z-loss，自适应梯度裁剪	`olmo/train.py`、`olmo/optim.py`
Checkpoint	sharded / unsharded / ephemeral 三种，带 RNG 和数据位置的完整恢复	`olmo/checkpoint.py`、`olmo/train.py`
评测	训练中途按域算困惑度 + 20 多个下游多选任务（in-loop）	`olmo/eval/`
生成	自带 beam search（采样器 / 长度惩罚 / n-gram 阻断约束），HF 桥接层	`olmo/beam_search.py`、`hf_olmo/`
配方	OLMo-2 7B/13B 两阶段训练的真实 YAML 配置 + 数据 provenance CSV	`configs/official-1124/`

用起来什么样

最小例子是在 Mac 上跑一个 20M 玩具模型（README.md 给的官方示例）：

python scripts/train.py configs/tiny/OLMo-20M.yaml --save_overwrite

真实形态是 8 卡 torchrun + 一份官方配置（摘自 README.md）：

torchrun --nproc_per_node=8 scripts/train.py configs/official-1124/OLMo2-7B-stage1.yaml

打开 configs/official-1124/OLMo2-7B-stage1.yaml 你会看到 OLMo-2 7B 第一阶段的全部真相：d_model: 4096、n_layers: 32、rope_theta: 500000、max_sequence_length: 4096、AdamW learning_rate: 3.0e-4、cosine 调度按 token 计（t_max: 5e12，即 5 万亿 token）、以及一千多行逐文件列出的预训练数据 URL。

一句话直觉

把 OLMo 想成「一个预训练项目的完整 git 仓库 + 飞行记录仪」。 一般开源模型只给你看飞机（权重）；OLMo 把设计图纸（代码）、装配手册（配置）、每一段航程的黑匣子数据（checkpoint + 日志 + 数据顺序）全部摆在了货架上。

2. 顶层全景（它大概怎么转）

2.1 顶层图

一条训练命令从启动到产出模型，流经这些部件：

                  configs/official-1124/OLMo2-7B-stage1.yaml
                    (模型结构 + 数据清单 + 优化器 + 调度 + 评测清单)
                                   │
                                   ▼
   torchrun ──► scripts/train.py  main()                    ① 载入/校验配置
                                   │                           建数据加载器
                                   ▼                           建模型/优化器
        ┌──────────────────────────────────────┐
        │   Trainer.fit()   (olmo/train.py)    │  ② 主循环：逐 batch 训练
        │                                      │     定期 checkpoint / 评测
        └───┬──────────┬───────────┬───────────┘
            │          │           │
            ▼          ▼           ▼
   ③ 数据管线     ④ 模型前反向    ⑤ 优化器步进
   IterableDataset  OLMo blocks   AdamW + cosine
   → MemMapDataset  (FSDP 分片)   + 梯度裁剪
   → 磁盘/HTTP .npy
                                   │
            ┌──────────────────────┼──────────────────┐
            ▼                      ▼                  ▼
   ⑥ checkpoint 体系      ⑦ in-loop 评测        ⑧ 日志与监控
   sharded/unsharded      域困惑度 + 下游任务   W&B + 控制台 + 数据索引
   (checkpoint.py)        (eval/)               (train.py)
                                   │
                                   ▼
                    step1000/ ... step928646-unsharded/
                    → 转换上传 Hugging Face（hf_olmo/）

怎么读这张图： 一切从一份 YAML 配置出发（①）；Trainer.fit() 是唯一的中央循环（②），它左手拽着数据管线（③）右手驱动模型（④⑤），外围三件「后勤装备」——checkpoint（⑥）、评测（⑦）、日志（⑧）——保证万亿 token 级别的长跑可恢复、可观测。

2.2 部件职责

部件	干什么	在哪个文件
`TrainConfig` 族	全部配置的 dataclass + OmegaConf 加载/合并/覆盖	`olmo/config.py:947`
`main`	进程入口：设设备、存配置、建 dataloader/模型/优化器、组装 Trainer	`scripts/train.py:54`
`OLMo`	模型本体：`wte` 嵌入 + N 个 block + `ln_f` + 输出头	`olmo/model.py:1070`
`OLMoSequentialBlock` / `OLMoLlamaBlock`	两种 transformer block 实现（融合 QKV / 仿 Llama 分离投影）	`olmo/model.py:678`、`:826`
`MemMapDataset`	把一堆 `.npy` token 文件当成一个连续数组按 chunk 切实例	`olmo/data/memmap_dataset.py:20`
`IterableDataset`	全局确定性 shuffle + rank/worker 切片 + 多线程预取	`olmo/data/iterable_dataset.py:20`
`DataCollator`	把变长实例 pad 成 batch（右 pad、attention mask、doc lens）	`olmo/data/collator.py:14`
`Trainer`	中央循环：micro-batch、loss、裁剪、调度、checkpoint、评测、取消检测	`olmo/train.py:208`
`Optimizer` / `AdamW` / `LionW`	优化器 + 逐参数梯度指标收集 + 两种梯度裁剪	`olmo/optim.py:38`
`Scheduler` 族	LR 调度：cosine/linear/inverse-sqrt/max/bolt-on warmup 等	`olmo/optim.py:651`
`Checkpointer` 族	`FullCheckpointer`（单文件）与四种 sharded checkpointer	`olmo/checkpoint.py:502`
`Evaluator` / `ICLMetric`	in-loop 评测：LM 困惑度与下游多选任务	`olmo/eval/evaluator.py:14`、`olmo/eval/downstream.py:28`
`BeamSearch`	生成：通用 beam search + 采样器 + 约束	`olmo/beam_search.py:649`
`OLMoForCausalLM`	Hugging Face 兼容封装	`hf_olmo/modeling_olmo.py:41`

2.3 主线走一遍（一个训练 step）

对应 Trainer.fit()（olmo/train.py:1109）主循环体内的一次迭代：

① 取 batch     train_loader 吐出一个 device batch
               (IterableDataset 已按 seed+epoch 决定好全局顺序)
② 记账         global_step += 1，累计 tokens seen；校验 batch 形状
               （train.py:1206-1208 硬断言 seq_len 与 batch size）
③ 前反向       train_step() → split_batch() 切 micro-batch
               → 逐 micro-batch autocast 前向 + loss.backward()
④ 指标与裁剪   聚合各 rank loss；clip_grads_and_collect_metrics()
               按组做固定/自适应梯度裁剪，顺便收集逐参数范数
⑤ 调 LR        scheduler.get_lr() 更新每个 param group
⑥ 优化器步进   optim.step()；NaN 检查（train.py:891）
⑦ 后勤         定期：写日志 → 检测取消信号 → 存 checkpoint → 跑评测

这条线上最容易被忽略、但最体现「严肃项目」气质的是 ②⑦：每个 batch 的形状被硬断言（这样 token 计数才可信），每个 rank 每个 step 读了哪些实例被写进 TSV（olmo/train.py:836-838）——数据顺序本身成了可审计的公共产物。

3. 阅读地图（建议顺序）

五章由浅入深。时间有限的话，读 01 → 02 → 03 就能拿到这个仓库 80% 的教学价值。

顺序	章节	讲什么	适合谁
1	01-model.md	配置驱动的模型骨架；OLMo-2 相对 GPT-2 默认架构的每处改动为什么存在	想读懂现代 LLM 架构细节的人
2	02-data-pipeline.md	memmap token 流、全局 shuffle、三级切片、真实数据混合配方	关心数据工程与可复现性的人
3	03-trainer.md	训练循环、micro-batch、z-loss、梯度裁剪、checkpoint 与精确恢复	要写自己训练循环的人（核心章）
4	04-eval.md	in-loop 评测怎么设计：域困惑度 + 下游 loglikelihood 多选	关心训练监控与评测的人
5	05-generation.md	beam search 全家桶、HF 桥接、checkpoint 加载	要用/要转换模型的人

4. 巧妙之处（可借鉴的技术）

数据顺序是公共产物。 训练时每步把 batch 里的实例 index 写进 data-indices/rank{N}.tsv.gz（olmo/train.py:836-838），配合 IterableDataset 的确定性全局 shuffle（olmo/data/iterable_dataset.py:91-114），任何人都能精确复现「第 928646 步模型看到了什么」。
恢复 = 模型 + 优化器 + RNG + 数据位置四件套。 trainer_state_dict（olmo/train.py:322-339）连 Python/NumPy/PyTorch/CUDA 四套 RNG 状态都存；恢复时用 start_index 直接跳回数据流中部（olmo/train.py:403-406），还能用 fast_forward_batches 跳过引发 loss 尖峰的数据（olmo/config.py:999）。
「先存再读验证」的 pre-train checkpoint。 正式开训前先存一个 checkpoint 再立刻读回来验证（scripts/train.py:334-344）——把「磁盘满了/权限错了」这类灾难消灭在第 0 步。
NaN 哨兵恢复法。 FullCheckpointer.restore_checkpoint 加载前先把所有参数填成 NaN，加载后断言无 NaN 残留（olmo/checkpoint.py:697-739）——任何一个参数没被恢复都会当场爆炸，而不是悄悄带着随机值继续训。
BufferCache 而非 buffer。 注意力偏置、RoPE 频率这些「通常做成 PyTorch buffer」的东西全部放进一个普通 dict（olmo/model.py:108-116），因为 FSDP 会跨进程同步 buffer，而这些缓存含 -inf，同步时可能变成 NaN——注释里写明了这是踩过的坑。
Vocab padding。 embedding_size 默认取 vocab 向上对齐 128 的倍数（olmo/config.py:410-416），OLMo-2 把 100278 的词表 pad 到 100352（configs/official-1124/OLMo2-7B-stage1.yaml:29-30），纯为吞吐量。
按 token 而不是按 step 思考。 max_duration 可以直接写 "5e12T"（5 万亿 token）或 "1ep"，由 Trainer.max_steps 换算（olmo/train.py:262-279）；调度器也可以用 token 作单位（SchedulerUnits.tokens）。规模实验语言直接进入了配置语法。
W&B 标签当急停按钮。 训练循环定期检查 run 是否被打了 cancel 标签，有则存 checkpoint 优雅退出（olmo/train.py:1075-1089）——万卡长跑的运维动作简化成了网页上打个标签。

5. 边界与局限

本仓库已冻结。 README 顶部明确声明：此仓库不再活跃，最新版本（OLMo-2 32B、OLMo-3）在 allenai/OLMo-core。本仓库覆盖到 OLMo-2 13B 为止（README.md:25 的 NOTICE）。
只支持纯 decoder-only 稠密 LM。 没有 MoE、没有多模态、没有张量并行/流水并行——分布式只有 FSDP / DDP / 单机三档（olmo/config.py:741 的 DistributedStrategy）。想学大模型并行切分请去 megatron-lm。
Llama block 是兼容层。 OLMoLlamaBlock 不支持文档内注意力掩码（olmo/model.py:899-902 直接 NotImplementedError），也不支持 flash-attn varlen；官方配置全部用 sequential block。
数据混合是「文件列表级」的。 混合比例靠配置文件里逐条列出 .npy 路径（想上采样就重复列，如 olmo/data/named_data_mixes.py 里 wiki 出现两次）；没有在线按比率采样。改配比 = 改文件列表。
in-loop 下游评测只支持 zero-shot loglikelihood 式任务。 ICLMultiChoiceTaskDataset 类注释写明「Only supports zero-shot for now」（olmo/eval/downstream.py:169）；严肃的少样本/生成式评测在仓库外的 OLMo-eval / olmes 项目。
源卡片提到的 docs/「训练不稳定复盘」在本 commit 里其实不存在——docs/ 只有 Checkpoints.md、RELEASE_PROCESS.md、Safetensors.md 三篇运维笔记。loss spike 的分析要以 OLMo 论文（arXiv:2402.00838）和 W&B 公开报告为准，代码里只剩 fast_forward_batches、early_stopping_factor 这些应对机制的开关。

6. 横向对比

维度	OLMo（本库）	nanogpt	megatron-lm
定位	可复现的完整预训练项目（代码+数据+checkpoint+日志）	教学用最小 GPT 实现	工业级大规模并行训练框架
模型定义	配置驱动单文件，支持多种 norm/位置编码/block 变体	单文件极简 GPT-2	高度模块化 + 张量/流水/专家并行
数据	memmap token 流 + 确定性全局 shuffle + 实例过滤	单个 `.bin` 文件按 offset 读	复杂 indexed dataset + 混合配比
并行	仅 FSDP/DDP/单机	DDP	TP/PP/EP/DP 全家桶
评测	训练循环内嵌域困惑度 + 下游任务	无	外挂
独特价值	公开了真实 run 的全部配置与数据顺序	代码最短	并行策略最深

数据侧的兄弟库：datatrove（HF 的 PB 级数据处理管线）、dolma（OLMo 预训练数据 Dolma 的构建工具包）、hf-datasets（通用数据集抽象）。OLMo 本仓库消费的是 dolma 产出的、已 token 化的 .npy——数据「制造」在 dolma/datatrove，数据「投喂」在本仓库。

7. 代码地图（导航索引）

主题	文件路径	符号名
命令行入口	`scripts/train.py`	`main`
配置体系	`olmo/config.py`	`TrainConfig`、`ModelConfig`、`DataConfig`、`BaseConfig.load`
模型本体	`olmo/model.py`	`OLMo`、`OLMoSequentialBlock`、`OLMoLlamaBlock`、`RotaryEmbedding`、`BufferCache`
真实模型配置	`configs/official-1124/OLMo2-7B-stage1.yaml`	（YAML：模型/优化器/数据/评测全量字段）
数据：memmap	`olmo/data/memmap_dataset.py`	`MemMapDataset`、`_read_chunk_from_memmap`
数据：顺序与切片	`olmo/data/iterable_dataset.py`	`IterableDataset`、`_build_global_indices`
数据：混合配方	`olmo/data/named_data_mixes.py`、`configs/official-1124/provenance.csv`	`DATA_PATHS`
训练循环	`olmo/train.py`	`Trainer.fit`、`train_step`、`train_batch`、`get_labels`
优化器与调度	`olmo/optim.py`	`AdamW`、`LionW`、`CosWithWarmup`、`BoltOnWarmupScheduler`、`clip_grads_and_collect_metrics`
Checkpoint	`olmo/checkpoint.py`	`FullCheckpointer`、`TorchLegacyShardedCheckpointer`、`OlmoCoreCheckpointer`、`build_sharded_checkpointer`
评测	`olmo/eval/downstream.py`、`olmo/eval/evaluator.py`	`ICLMetric`、`ICLMultiChoiceTaskDataset`、`label_to_task_map`、`Evaluator`
生成	`olmo/beam_search.py`、`olmo/model.py:1603`	`BeamSearch`、`TopKSampler`、`TopPSampler`、`OLMo.generate`
HF 桥	`hf_olmo/modeling_olmo.py`、`hf_olmo/convert_olmo_to_hf.py`	`OLMoForCausalLM`、`write_model`
数据制备脚本	`scripts/prepare_memmap_dataset.py`	（jsonl.gz → memmap `.npy`）

1. 这是什么（零基础也能懂）​

一句话定义​

它解决谁的什么问题​

它能做什么​

用起来什么样​

一句话直觉​

2. 顶层全景（它大概怎么转）​

2.1 顶层图​

2.2 部件职责​

2.3 主线走一遍（一个训练 step）​

3. 阅读地图（建议顺序）​

4. 巧妙之处（可借鉴的技术）​

5. 边界与局限​

6. 横向对比​

7. 代码地图（导航索引）​