PyTorch 为 TorchTitan 引入 Zephon 弹性确定性数据加载

为什么选择 Zephon？

模型在训练期间所看到的数据，是我们影响其表现好坏的最有力杠杆之一。改进模型的许多日常工作都发生在数据侧：我们改变数据混合比例，或改变所应用的过滤器，或改变样本的分词和打包方式，然后运行一次消融实验，看看这一改动是否有帮助。这样做首先有个问题：这些改动中的每一项，往往都需要我们对整个训练语料库做一次完整遍历，并运行昂贵的离线操作来准备数据、将其存储到磁盘上，之后消融实验才能运行。由于一次典型的训练运行只会恰好读取那份副本一次，而且大多数实验最终并不会产生更好的模型，我们便花费了大量时间和存储空间来创建那些再也不会被读取的数据集。

如今实验可以由智能体而非人来设计、启动和解读，这一点就更加重要了。智能体非常擅长并行运行大量实验，但它们并不特别擅长注意到某个结果实际上取决于它们所做改动之外的其他因素，比如实验恰好运行在多少块 GPU 上（平心而论，人类也不擅长注意到这一点）。而且对于大多数数据加载器来说，GPU 的数量确实会改变结果，因为它改变了数据顺序。

为说明这一点，我们在 TorchTitan 中用 16 块 GPU 训练一个 1B 参数模型，在第 150 步取一个检查点，然后从该检查点在 8、32 和 64 块 GPU 上恢复训练，并与原始运行并列，分别使用标准数据加载器和 Zephon 各做一次（可以把它想象成运行遇到硬件故障，然后在当时恰好空闲的那组 GPU 上恢复）。如图 1 所示，使用标准数据加载器时，改变 GPU 数量意味着数据加载器在恢复后交给模型的是不同的样本序列，因此每一次最终都会得到一条不同的损失曲线。使用 Zephon 时，这些曲线彼此重合，而且每个全局批次的 SHA-256 哈希都与基线运行一致。

[图 1：从 16 块 GPU 在第 150 步的检查点出发，在 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：在 8–64 块 GPU 上，使用非弹性确定性（标准）数据加载和 Zephon 时，在 DCLM Core v1 和 FineWeb 上的下游评估。灰色带标记评估噪声底限。]

请注意，Zephon 在完成所有这些工作的同时，还能即时进行数据准备工作：它在训练任务内部对数据进行分词、打包、打乱和混合，充分利用那些在模型训练时闲置的 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 根据每个 Dataset 的元数据规划课程，并与实际处理这些样本的算子保持分离，这样你就可以在不触及处理流水线的情况下更改混合方式，反之亦然。为了确保我们在 token 基础上获得正确的混合百分比，我们包含了一个 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 中并行度成为一个纯粹的性能旋钮：你可以将分词工作器的数量翻倍，而输出不会改变。

用通道代替秩。 Zephon 不直接将数据划分给可用的 GPU，而是将其划分为固定数量的逻辑流，称为通道。每个通道都有自己的有序样本序列，以及算子需要保留的任何状态的自己的副本（比如打包算子内部部分填充的箱子），因此无论最终分配给哪个 GPU，一个通道都会产生完全相同的记录。当数据并行秩的数量发生变化时，Zephon 只需在秩之间重新分配通道，全局批次结果保持不变。累加器使每个通道具有确定性；通道使 Zephon 具有弹性确定性。问题（你知道肯定有一个）是你需要预先选择通道的数量，我们将在下面回到这一点。

刷新和重放，而不是保存状态。 Zephon 检查点不包含任何打包箱或洗牌缓冲区的内容；它只记录在拍摄检查点时仍在处理哪些输入块。为了保持该列表简短，Zephon 会定期通过流水线发送一个刷新哨兵，告诉每个有状态算子发出它持有的任何内容，并从干净的状态重新开始。当我们从检查点恢复时，Zephon 会从最近一次刷新开始重放输入（这会精确重建算子的状态），并丢弃训练循环已经消耗的任何内容，这使得检查点无论运行进行到多远都保持很小（在我们的实验中为几 MiB）。

接入 TorchTitan

我们的 TorchTitan 集成 用由数据配方（即你的数据集、它们的混合权重，以及你想要的任何打包、洗牌和并行设置）构建的 Zephon Pipeline 替换 TorchTitan 的默认数据加载器，训练运行的其他一切保持不变。Zephon 的状态与模型一起进入 TorchTitan 的检查点，并且由于该状态覆盖每个通道，你可以从同一个检查点在不同数量的 GPU 上恢复运行。当运行启动时，集成会检查通道是否在数据并行副本之间均匀划分，以及每个优化器步骤的批次是否在通道之间均匀拆分，因此配置错误的运行会立即失败，而不是在几个小时后悄悄改变数据。 

该 repo 还附带一个仅用 CPU 的快速演示，它会为一个数据流创建检查点，在不同数量的数据并行副本上恢复它，并检查每个全局步是否与不间断运行一致，此外还有一个 GPU 冒烟测试，用真实（尽管很小）的模型做同样的事情。README 中有运行这些的命令，以及设置说明、完整的配置选项集和一些入门示例配方。

开始前需要了解的事项

选择你的 lane 数量。 选择一个你预期运行的数据并行大小的公倍数。如果你可能在 8、16 或 32 个 GPU 上运行，32 个 lane 是覆盖这三者的上限。但不要过度（例如，当你的最大值是 16 时却用几千个），因为每个 GPU 上每多一个 lane 都会增加检查点恢复时间的开销。

在更多/更少 rank 上恢复时保持配方固定。 数据、分词器、种子、混合、上下文长度和全局批大小都需要保持不变。当然，你仍然可以调整并行度等性能设置。

Zephon 保证数据，而不是算术。 全局批次在不同 GPU 数量下将完全相同，损失也会非常接近，但可能不完全相同，因为训练循环中的浮点归约顺序仍然取决于有多少 GPU 分担工作。

展望未来

有几个方向让我们特别兴奋。第一个是无 GIL 的 Python：Zephon 允许用户选择一个 runner，即如何并行化算子。我们实现的一个 runner 是基于线程的 runner。在 Python 3.14t 上，它获取视觉-语言样本的速度比 Zephon 的进程 runner 在标准 Python 上快约 3.7 倍。我们认为这就是数据加载的发展方向，而 Zephon 已经为此做好了准备。 第二个是动态课程：因为课程与处理流水线保持分离，一个 WorkSource 可以在运行过程中改变混合权重，而无需触及下游的任何内容。我们正在开发一个可配置的课程组件，它正是这样做的，同时保持我们上面描述的确定性保证。

除此之外，我们很想听听你想把数据流水线的哪些部分移入训练作业，以及我们如何提供帮助；非常欢迎提交 issue 和 PR！

了解更多：发布文章 · 论文 · Zephon · GitHub 上的 Zephon · TorchTitan 集成

文章 Bringing Elastically Deterministic Data Loading to TorchTitan with Zephon 首次发表于 PyTorch。