Zephon 与 TorchTitan 的集成确保在不同 GPU 数量上恢复训练时，全局批次（global batch）保持完全一致。

在我们最新博客的 1B 参数模型测试中，训练从 16 GPU 的检查点分别在 8、32 和 64 GPU 上恢复，每个全局批次的 SHA-256 哈希值都与原始运行完全匹配。

Zephon 在训练任务内部完成分词、打包、洗牌和混合数据，无需预先生成一份语料的预处理副本。

在恢复训练时，数据、分词器、随机种子、混合比例、上下文长度和全局批次大小必须保持不变。车道数（lane count）必须能被数据并行副本数整除，且每个优化器步骤的批次必须能在各车道之间均匀切分。训练损失仍可能有所不同，因为浮点归约顺序取决于 GPU 数量。

✍️ Josh Wills、Maximilian Böther、Ties Robroek（DatologyAI）

阅读更多：https://pytorch.org/blog/bringing-elastically-deterministic-data-loading-to-torchtitan-with-zephon/…

用 Zephon 为 TorchTitan 带来弹性确定性数据加载 - PyTorch

来源：https://pytorch.org/blog/bringing-elastically-deterministic-data-loading-to-torchtitan-with-zephon/

精选项目

• PyTorch logo (https://pytorch.org/projects/pytorch/)

为什么需要 Zephon？

模型在训练过程中看到的数据，是我们提升其性能最强大的手段之一。改进模型的日常工作大多发生在数据侧：我们调整数据混合比例、修改过滤条件、改变样本的分词与打包方式，然后运行消融实验来验证改动是否有效。问题在于，这些改动通常都需要对训练语料进行完整扫描，并执行昂贵的离线预处理操作来准备数据并存入磁盘，之后消融实验才能开始。而一次典型的训练运行只会完整读取这份数据一次，且大多数实验最终并不能产出更好的模型，因此我们花费大量时间和存储去创建那些将永远不会被再次读取的数据集。

当实验可以由智能体（agent）而非人类来设计、启动和解读时，这一点变得更加重要。智能体擅长并行运行大量实验，但不太擅长察觉某个结果实际上是否依赖于它所做的改动之外的因素——比如实验恰好跑在多少个 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 上恢复后的损失曲线对比，标准加载器 vs. Zephon，附有每个全局批次逐步 SHA-256 值与原始运行的比对条码。]

为了量化这有多重要，我们使用标准数据加载器从头开始在 8、16、32、64 个 GPU 上分别用 20B token 训练同一个 1B 参数模型。如图 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 上 8–64 GPU 的下游评估结果对比，非弹性确定性（标准）数据加载 vs. Zephon。灰色区域标记为评估噪声底线。]

值得注意的是，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}, # token 比例，而非样本比例 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 操作符都可以拥有自己的并行度；在这个例子中，分词获得了 4 个 worker，因为它是开销最大的步骤。并行化预处理操作符正是 Zephon 保持 GPU 持续供数的方式，即使在更复杂的流水线上也是如此：在一个完整的视觉-语言流水线（解码图像、分词、打包和批处理样本）上，Zephon 的吞吐量与我们测试中最快的加载器 Google Grain 相当，同时提供了更强的弹性确定性保证。

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

Zephon 如何保持批次一致

实时打包和混合数据正是弹性确定性难以实现的原因。当数据加载器对文档进行分词、将它们打包成固定长度序列、并在过程中混合多个数据集时，训练数据的顺序就不再是你可以在磁盘上通过索引查找的内容了（我们在论文中将其称为有状态的 n:m 流水线 https://arxiv.org/abs/2610.03087）。在这种情况下，数据顺序取决于样本被使用时流水线的状态，而在大多数数据加载器中，这种状态无法在不改变输出的情况下迁移到不同数量的 GPU 上。现有加载器只能部分做到这一点：Mosaic Streaming 和 Google Grain 可以在不同 GPU 数量之间保持全局批次一致，但前提是样本顺序可以提前通过索引推算出来，这就排除了实时打包等操作。Megatron-Energon 可以在不同 GPU 数量上复现其批次，但前提是所有 GPU 上的数据加载器 worker 总数保持不变。

Zephon 通过三个相互支撑的设计决策实现了无上述约束的弹性确定性。

累加器负责决策，worker 负责执行。 每个 Zephon 操作符都分为两部分：一个单线程的累加器（accumulator），按顺序看到每条记录并做出所有有状态的决策（比如哪些序列被打包在一起）；以及一个无状态函数，在你分配的任意数量并行 worker 上执行耗时的计算。由于 worker 可能以任意顺序完成，Zephon 在分发前给每组记录编号，并在传递结果前将其按顺序排列。与其他数据加载器不同，这使得 Zephon 中的并行度成为纯粹的性能调节参数：你可以将分词 worker 数量翻倍，输出不会改变。

车道（lane）而非排名（rank）。 Zephon 不直接在可用 GPU 之间划分数据，而是将数据分成固定数量的逻辑流，称为车道。每条车道拥有自己有序的样本序列和操作符所需的任何状态的独立副本（比如打包操作符中已部分填充的 bin），因此无论最终被分配到哪个 GPU，每条车道产生的记录完全相同。当数据并行 rank 数量变化时，Zephon 只需将车道重新分配到各 rank 上，全局批次就会保持一致。累加器使每条车道具有确定性；车道使 Zephon 具有弹性确定性。不过有一点限制（你肯定猜到会有一个），就是需要预先选择车道数量，我们稍后再讨论。

刷新与重放，而非保存状态。 Zephon 的检查点不包含任何打包 bin 或洗牌缓冲区的内容；它只记录检查点创建时仍在处理的输入块。为了保持这个列表尽量短，Zephon 会定期向流水线发送一个刷新哨兵（flush sentinel），告诉所有有状态操作符输出当前持有的内容并从干净状态重新开始。当我们从检查点恢复时，Zephon 从最近一次刷新开始重放输入（这会精确地重建操作符的状态），并丢弃训练循环已经消费的数据，这使得检查点始终保持很小（在我们的实验中只有几 MiB），无论训练运行了多久。

接入 TorchTitan

我们的 TorchTitan 集成用基于数据配方（data recipe）构建的 Zephon Pipeline 替换了 TorchTitan 的默认数据加载器（即你的数据集、混合权重以及任何打包、洗牌和并行度设置），训练运行的其他方面保持不变。Zephon 的状态随模型一起存入 TorchTitan 的检查点，由于该状态覆盖了所有车道，你可以用同一个检查点在不同 GPU 数量上恢复训练。当运行启动时，集成会检查车道数能否被数据并行副本数整除，以及每个优化器步骤的批次能否在各车道之间均匀切分，这样配置错误的运行会立即失败，而不是在几个小时后悄悄改变数据。

仓库还附带了一个纯 CPU 的快速演示，它对数据流进行检查点、在不同的数据并行副本数上恢复，并检查每个全局步骤都与不中断的运行一致；另外还有一个 GPU 冒烟测试，使用一个真实（尽管很小的）模型执行相同的操作。README 中有运行这些测试的命令、安装说明、完整的配置选项列表，以及一些示例配方供你快速上手。

开始之前需要了解的事项

选择车道数量。 选一个你预期运行的各数据并行规模的公倍数。如果你可能在 8、16 或 32 个 GPU 上运行，32 条车道是一个能覆盖三种情况的上限。但也不要过度配置（比如最大规模为 16 却设置数千条车道），因为每 GPU 每多一条车道都会增加检查点恢复的开销。

在不同 rank 数量上恢复时保持配方不变。 数据、分词器、随机种子、混合比例、上下文长度和全局批次大小都必须保持相同。当然，你仍然可以调整并行度等性能设置。

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

展望未来

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

除此之外，我们非常希望了解你希望将数据流水线的哪些部分移入训练任务中，以及我们如何提供帮助；欢迎提交 issue 和 PR！

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