精选项目
为什么需要 Zephon？
模型在训练期间看到的数据，是我们影响其性能表现最有力的杠杆之一。改进模型的许多日常工作都发生在数据侧：我们改变数据混合比例，或改变所应用的过滤器，或改变样本的分词与打包方式，然后运行一次消融实验，看看这一改动是否有帮助。这里的第一个问题是，上述每一项改动通常都要求我们对整个训练语料做一次完整遍历，并运行昂贵的离线操作来准备数据、将其存储到磁盘上，之后消融实验才能运行。由于一次典型的训练运行只会读取那份副本恰好一次，而且大多数实验最终并不会产生更好的模型，我们便花费了大量时间和存储空间来创建那些再也不会被读取的数据集。
如今这一点更加重要，因为实验可以由智能体而非人来设计、启动和解读。智能体非常擅长并行运行大量实验，但它们并不特别擅长察觉某个结果实际上取决于它们所做改动之外的其他因素，比如实验恰好运行在多少块 GPU 上（平心而论，人类在这方面也不擅长）。而对于大多数数据加载器来说，GPU 数量 确实会 改变结果，因为它改变了数据顺序。
为说明这一点，我们在 TorchTitan 中用 16 块 GPU 训练一个 1B 参数模型，在第 150 步取一个检查点，然后分别用 8、32 和 64 块 GPU 从该检查点恢复训练，并与原始运行并列对比，一次使用标准数据加载器，一次使用 Zephon（可以把它想象成运行遇到硬件故障，然后在当时恰好空闲的任何一组 GPU 上恢复）。如图 1 所示，使用标准数据加载器时，改变 GPU 数量意味着数据加载器在恢复后交给模型的是不同的样本序列，因此每次运行最终都会得到一条不同的损失曲线。而使用 Zephon 时，各条曲线彼此重合，每个全局批次的 SHA-256 哈希都与基线运行一致。
[图 1：从第 150 步的 16 GPU 检查点分别在 8、32 和 64 块 GPU 上恢复后的损失曲线，标准加载器对比 Zephon，并附有条形码，将每个全局批次的逐步 SHA-256 与原始运行进行对比。]
为体会这一点有多重要，我们使用标准数据加载器，在 8、16、32 和 64 块 GPU 上从头训练同一个 1B 参数模型，训练 20B 个 token。图 2 展示了各 GPU 数量下的得分。模型的 DCLM Core v1 分数会因训练所用 GPU 数量不同而最多相差 0.82 分，其基于 FineWeb 评估套件的分数最多相差 0.64 分。这大约是 DCLM 和 FineWeb 作者在评估数据整理方法时所使用的相关阈值的四倍，这意味着在两种不同资源分配上比较两种数据混合比例的消融实验，很可能实际测量的是 GPU 数量而非数据。而使用 Zephon 时，不同 GPU 数量之间的差异分别仅为 0.05 分和 0.011 分。
[图 2：在 DCLM Core v1 和 FineWeb 上，使用非弹性确定性（标准）数据加载与 Zephon 时，8–64 块 GPU 下的下游评估结果。灰色带标示评估噪声下限。]
请注意，Zephon 在完成这一切的同时还即时进行数据准备工作：它在训练任务内部完成数据的 tokenize、打包、打乱和混合，充分利用模型训练期间闲置的所有 CPU，这意味着无需创建预处理语料副本（也无需在实验失败时将其丢弃）。
Zephon 流水线
Zephon 要求你描述三件事： Datasets 要用于训练的数据集，一个 WorkSource 来决定这些数据集如何混合并排序成训练课程，以及一个 Pipeline 将样本从其存储表示转换为训练循环所消费张量的算子。下面是一个示例，它直接从 Hugging Face 混合两个文本数据集，然后对结果进行分词、打包、打乱和批处理：
from zephon import Pipeline
from zephon.io import Dataset
from zephon.work import StaticMixtureWorkSource, TokenEstimation
source = StaticMixtureWorkSource(
datasets=[
Dataset.from_path("wiki", "hf://Salesforce/wikitext/wikitext-2-raw-v1/train"),
Dataset.from_path("news", "hf://fancyzhx/ag_news/default/train"),
],
mixture={"wiki": 0.5, "news": 0.5}, # a token ratio, not a sample ratio
token_estimation=TokenEstimation(),
seed=42,
)
pipeline = (
Pipeline(source)
.options(io_options={"cache": {"enabled": True}})
.tokenize(tokenizer_id="gpt2", field="text", parallelism=4)
.pack_flat(max_length=2048, algorithm="wrap")
.shuffle(buffer_size=1024, seed=42)
.batch(microbatch_size=8)
)
for batch in pipeline:
train_step(batch.to_training())这个 WorkSource 根据每个数据集的元数据来规划课程，并与实际处理这些样本的算子保持分离，这样你就可以在不触及处理流水线的情况下更改混合比例，反之亦然。为了确保我们在 token 基础上获得正确的混合百分比，我们加入了一个 Dataset实用工具，它会对每个数据集的小样本测量 token/字节比率，并调整计划，以便在样本以原始字符串形式存储的情况下，仍能在 token 层面达到目标混合比例。每个单独的算子都可以拥有自己的并行度；在这里，分词获得四个工作进程，因为它是最昂贵的步骤。并行化预处理算子正是 Zephon 即使在更复杂的流水线上也能让 GPU 保持供给的方式：在一个完整的视觉-语言流水线上，它解耦 TokenEstimation 实用工具，它会对每个数据集的小样本测量 token/字节比率，并调整计划，以便在样本以原始字符串形式存储的情况下，仍能在 token 层面达到目标混合比例。
每个单独的 Pipeline 每个算子都可以拥有自己的并行度；在这里，分词获得四个工作进程，因为它是最昂贵的步骤。并行化预处理算子正是 Zephon 即使在更复杂的流水线上也能让 GPU 保持供给的方式：在一个完整的视觉-语言流水线上，该流水线会解码图像、分词、打包并批处理其样本，Zephon 与 Google Grain 保持同步——后者是我们测试过的最快加载器——同时提供更强的弹性确定性保证。
[图 3：在完整的视觉-语言流水线上，随着并行度增加，Zephon、Grain 以及我们在构建 Zephon 之前内部使用的快速拼凑的内部加载器每 GPU 每秒可用（非填充）token 数。水平线是防止 GPU 停滞所需的吞吐量.]
Zephon 如何保持批次完全一致
在运行过程中即时打包和混合数据，正是弹性确定性难以实现的原因。当数据加载器对文档进行分词、将它们打包成固定长度的序列，并在运行过程中混合多个数据集时，训练数据的顺序就不再是可以通过磁盘上的索引查到的了（我们称之为 有状态 n:m 流水线 ， 我们的论文）。在这种情况下，样本何时以及如何被使用取决于流水线的状态，而在大多数数据加载器中，这种状态无法在不改变输出的情况下转移到不同数量的 GPU 上。现有的加载器只能做到一部分：Mosaic Streaming 和 Google Grain 可以在不同 GPU 数量下保持全局批次一致，但前提是样本顺序可以事先从索引中推算出来，这就排除了即时打包之类的做法。Megatron-Energon 可以在不同数量的 GPU 上复现其批次，但前提是所有 GPU 上的数据加载器工作进程总数保持不变。
Zephon 通过三个层层递进的设计决策，在这些约束之外实现了弹性确定性。
累加器负责决策；工作进程负责执行。 Zephon 的每个算子都拆分为两部分：一个单线程累加器，它按顺序查看每条记录并做出所有有状态决策（例如哪些序列被打包在一起），以及一个无状态函数，它在任意数量的并行工作进程上完成开销较大的工作。由于工作进程可能以任意顺序完成，Zephon 会在分发前为每组记录编号，并在传递结果前将其按顺序还原。与其他数据加载器不同，这使得在 Zephon 中并行度成为一个纯粹的性能调节旋钮：你可以将分词工作进程数量翻倍，而输出不会改变。
用 lane 而非 rank。 Zephon 并没有把数据直接分配给可用的 GPU，而是将其划分为固定数量的逻辑流，称为“lane”。 lane每个 lane 都有自己有序的样本序列，以及算子需要保留的任何状态的独立副本（例如打包算子内部部分填充的 bin），因此无论最终分配给该 lane 的是哪块 GPU，该 lane 产生的记录都完全相同。当数据并行 rank 的数量发生变化时，Zephon 只需在各 rank 之间重新分配 lane，全局批次的结果依然相同。累加器让每个 lane 具有确定性；lane 则让 Zephon 具备弹性确定性。代价（你肯定知道总得有一个）是你需要预先选定 lane 的数量，我们将在下文回到这一点。
采用刷新与重放，而不是保存状态。 Zephon 检查点不包含任何打包 bin 或 shuffle 缓冲区的内容；它只记录在拍摄检查点时输入中哪些块仍在处理中。为了保持该列表简短，Zephon 会定期发送一个 刷新哨兵 沿着流水线，通知每个有状态算子输出其当前持有的内容，并从干净状态重新开始。当我们从检查点恢复时，Zephon 会重放自最近一次刷新以来的输入（这会精确重建算子的状态），并丢弃训练循环已经消费过的任何内容，这使得检查点保持很小（在我们的实验中只有几 MiB），无论训练进行到多深。
接入 TorchTitan
我们的 TorchTitan 集成 用 Zephon 替换 TorchTitan 的默认数据加载器， Pipeline 由一个数据配方构建（即你的数据集、它们的混合权重，以及你想要的任何打包、打乱和并行度设置），而训练运行的其他一切保持不变。Zephon 的状态会与模型一起写入 TorchTitan 的检查点，而且由于该状态覆盖每一条 lane，你可以从同一个检查点以不同的 GPU 数量恢复运行。当运行启动时，集成会检查各 lane 能否在数据并行副本之间均匀划分，以及每个优化器步骤的批次能否在各 lane 之间均匀拆分，因此配置错误的运行会立即失败，而不是在数小时后悄悄改变数据。
该 仓库 还附带一个仅用 CPU 的快速演示：它对数据流做检查点，在不同数量的数据并行副本上恢复，并检查每个全局步是否与不中断的运行一致；此外还有一个 GPU 冒烟测试，用真实（尽管很小）的模型做同样的事。该 README 包含运行上述各项的命令，以及设置说明、完整的配置选项集，还有一些示例配方帮你快速上手。
开始前须知
选择你的 lane 数量。 选择一个数字，它是你预期运行的数据并行规模的公倍数。如果你可能会在 8、16 或 32 块 GPU 上运行，32 条 lane 就是覆盖这三种情况的上限。但也不要过度（例如，当你的最大值是 16 时却选几千），因为每块 GPU 上每多一条 lane 都会增加检查点恢复时间的开销。
恢复训练时，无论使用更多还是更少的 rank，都要保持配方固定不变。 数据、分词器、种子、混合比例、上下文长度和全局批次大小都必须保持不变。当然，你仍然可以调整并行度等性能设置。
Zephon 保证的是数据，而不是算术。 全局批次在不同 GPU 数量下将完全一致，损失也会非常接近，但可能并不完全相同，因为训练循环中的浮点归约顺序仍然取决于有多少块 GPU 分担工作。
展望未来
我们尤其感到兴奋的有几个方向。第一个是无 GIL 的 Python：Zephon 允许用户选择一个 运行器，即如何并行化算子。我们实现的一个运行器是基于线程的运行器。在 Python 3.14t 上，它获取视觉语言样本的速度比 Zephon 的进程运行器在标准 Python 上快约 3.7 倍。我们认为数据加载的发展方向就在于此，而 Zephon 已经为此做好了准备。第二个是动态课程：由于课程与处理流水线保持分离，一个 WorkSource 可以在运行过程中更改混合权重，而无需改动任何下游环节。我们正在开发一个可配置的课程组件，它正是做这件事的，同时保持我们上文所述的确定性保证。
除此之外，我们很想听听你希望把数据流水线的哪些部分迁移到训练任务中，以及我们能如何提供帮助；非常欢迎提交 issue 和 PR！
了解更多： 发布文章 · 论文 · Zephon · GitHub 上的 Zephon · TorchTitan 集成
/row /row