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

PyTorch：Blog· Josh Wills、Maximilian Böther、Ties Robroek 及 Zephon 团队·2026-10-17 00:30· 52 分钟前AI 评分48

借助 Zephon 为 TorchTitan 带来弹性确定性数据加载

AI 导读

PyTorch 为 TorchTitan 引入 Zephon，一种弹性确定性数据加载方案，在训练任务内即时完成 tokenize、pack、shuffle 和 mix，无需预处理语料副本。

正文

为什么需要 Zephon？

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

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

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

[图 1：从 16 块 GPU 在第 150 步的检查点出发，在 8、32 和 64 块 GPU 上恢复后的损失曲线，标准加载器对比 Zephon，并附条形码，将每个全局 batch 的逐步 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 在完成所有这些工作的同时，还即时进行数据准备工作：它在训练任务内部对数据进行分词、打包、打乱和混合，利用模型训练时闲置的所有 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 中并行度成为一个纯粹的性能调节旋钮：你可以将分词工作进程的数量翻倍，而输出不会改变。

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

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

接入 TorchTitan

我们的 TorchTitan 集成用基于数据配方（即你的数据集、它们的混合权重，以及你想要的任何打包、洗牌和并行设置）构建的 Zephon Pipeline 替换 TorchTitan 的默认数据加载器，训练运行的其他一切保持不变。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 允许用户选择一种 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。

来源：PyTorch：Blog · pytorch.org