为什么选择 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 算子都可以有自己的并行度；这里，tokenization 分配了四个 worker，因为它是开销最大的步骤。对预处理算子进行并行化，正是 Zephon 能够持续为 GPU 提供数据的方式，即使在更复杂的流水线中也是如此：在一个完整的视觉-语言流水线中，该流水线负责解码图像、tokenize、打包和批处理样本，Zephon 能够跟上 Google Grain（我们测试过的最快加载器）的速度，同时提供更强的弹性确定性保证。

[ 图 3：在完整的视觉-语言流水线上，随着并行度增加，Zephon、Grain 以及我们在构建 Zephon 之前内部使用的快速拼凑的内部加载器每 GPU 每秒可用（非填充）token 数。水平线是防止 GPU 停滞所需的吞吐量。]

Zephon 如何保持批次一致

动态打包和混合数据正是弹性确定性难以实现的原因。当数据加载器对文档进行 tokenize、将它们打包成固定长度的序列，并在过程中混合多个数据集时，训练数据的顺序就不再是可以通过磁盘上的索引来查找的东西（我们在论文中称之为有状态 n:m 流水线）。在这种情况下，它取决于样本何时以及如何被使用时流水线的状态，而在大多数数据加载器中，这种状态无法在不改变输出的情况下转移到不同数量的 GPU 上。现有的加载器只实现了一部分：Mosaic Streaming 和 Google Grain 可以在不同 GPU 数量下保持全局批次相同，但前提是样本顺序可以事先从索引中计算出来，这就排除了动态打包之类的情况。Megatron-Energon 可以在不同数量的 GPU 上重现其批次，但前提是所有 GPU 上的数据加载器 worker 总数保持不变。

Zephon 通过三个相互叠加的设计决策，在没有这些限制的情况下实现了弹性确定性。

累加器做决策；worker 做工作。每个 Zephon 算子都分为两部分：一个单线程累加器，它按顺序查看每条记录并做出所有有状态决策（例如哪些序列被打包在一起），以及一个无状态函数，它在任意数量的并行 worker 上完成开销大的工作。由于 worker 可以以任意顺序完成，Zephon 在分发之前对每组记录进行编号，并在传递之前将结果按顺序放回。与其他数据加载器不同，这使得在 Zephon 中并行度成为一个纯粹的性能调节旋钮：你可以将 tokenization worker 的数量翻倍，而输出不会改变。

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

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

接入 TorchTitan

我们的 TorchTitan 集成用一个由数据配方（即你的数据集、它们的混合权重，以及你想使用的任何打包、shuffle 和并行度设置）构建的 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 可以在运行过程中更改混合权重，而无需触及下游的任何内容。我们正在开发一个可配置的课程组件，它正是这样做的，同时保持我们上面描述的确定性保证。

转载声明：本文为转载整理，原文版权归原作者及原站点所有。

原文来源：PyTorch Blog

原创地址：https://pytorch.org/blog/bringing-elastically-deterministic-data-loading-to-torchtitan-with-zephon/