PithTrain: A Compact and Agent-Native MoE Training System

发表时间: 2026-05 · arXiv:2605.31463

作者/机构: Ruihang Lai¹, Hao Kang¹, Haozhan Tang¹, Akaash R. Parthasarathy¹, Zichun Yu¹, Junru Shao³, Todd C. Mowry¹, Chenyan Xiong¹,², Tianqi Chen¹,³

¹卡内基梅隆大学, ²Xlue, ³NVIDIA

A1 主要贡献

现代人工智能系统日益依赖于专家混合（MoE）语言模型，其训练需要依赖能够在分布式集群上扩展的系统。现有生产级框架经过多年的工程努力，构建了优化的MoE训练堆栈，但在演进这些堆栈以支持新模型架构和系统优化时，成本高昂且需要深厚的专业知识。AI编码智能体的兴起为自动化部分训练框架开发工作提供了可能，但将它们应用于现有框架会带来隐藏成本，这些成本在当今仅关注吞吐量的评估中被忽略了。

核心问题与研究目标: 本文提出了一个设计问题：我们能否重新设计一个MoE训练框架，以优化“智能体-任务效率”（Agent-Task Efficiency, ATE）？ATE被定义为使用编码智能体来理解、操作和扩展一个框架的成本。为了回答这个问题，本文构建了一个从一开始就为智能体原生设计的端到端MoE训练框架——PithTrain。

创新点与主要贡献:

  1. PithTrain: 一个紧凑、Python原生的端到端MoE训练框架。PithTrain约有11,000行代码，从设计之初就为智能体原生，其训练吞吐量与生产级框架相当。

  2. 四大设计原则: 提出了四个用于构建智能体原生机器学习训练框架的设计原则，并以此指导PithTrain的设计：

    • 代码紧凑性（code compactness）
    • Python原生组件（Python-native components）
    • 无隐式间接调用（no implicit indirection）
    • 智能体技能（agent skills）

  3. 智能体-任务效率（ATE）与ATE-Bench: 提出了ATE作为超越训练吞吐量的评估指标，并开发了一个名为ATE-Bench的基准测试，用于系统性地评估框架设计如何影响真实训练系统任务中的ATE。实证研究表明，PithTrain相比生产级框架具有更高的ATE。

PithTrain实现了双重效率：强大的训练吞吐量和高智能体-任务效率。在NVIDIA H100和B200 GPU上，它在一系列MoE模型和训练设置中达到了与生产级框架相当的吞D吐量。在ATE-Bench上，编码智能体在PithTrain上完成相同的任务，与生产级框架相比，最难的新功能任务的智能体轮次（Agent Turns）减少了高达62%，活跃GPU时间（Active GPU Time）减少了高达64%。

图1：PithTrain概览。智能体发出动作并接收反馈（左）；PithTrain基于四个智能体原生设计原则构建，提供双重效率（中）；ATE-Bench评估跨框架的智能体-任务效率（右）。
图1：PithTrain概览。智能体发出动作并接收反馈（左）；PithTrain基于四个智能体原生设计原则构建，提供双重效率（中）；ATE-Bench评估跨框架的智能体-任务效率（右）。

A3 智能体原生设计原则

本节介绍指导PithTrain系统设计的四个智能体原生设计原则，并对比了当今生产级训练框架与这些原则的符合情况。表1总结了这种比较；不同的生产级框架采纳了这四个原则的不同子集，反映了它们不同的设计优先级。

紧凑的代码库。生产级框架如Megatron-LM和DeepSpeed提供了对模型、训练特性和硬件平台的广泛覆盖，这些都是多年工程努力的积累，其核心代码库超过16万行。更大的代码库增加了定位相关代码、跟踪跨文件依赖以及验证变更完整性的成本。一个紧凑的代码库可以降低这一成本；随着前沿编码智能体的上下文窗口达到20万到100万个token，一个足够紧凑的代码库甚至可以一次性装入单个上下文。我们将紧凑性视为增长的约束：PithTrain可能会随着时间推移而增长，但新增内容应尊重这四个原则。

Python原生的代码库。Python是现代机器学习中的主导语言。一个纯Python框架能让智能体在单一语言环境中浏览整个框架，提供可读的Python回溯信息而非不透明的本地错误，并消除了编译扩展的重建周期。Megatron-LM的核心Transformer层由树外的TransformerEngine【25，NVIDIA. TransformerEngine. 2022. URL: https://github.com/NVIDIA/TransformerEngine】模块组成，而DeepSpeed则捆绑了大量的树内扩展。这些扩展提供了顶尖的内核性能和厂商优化的数值计算，但迫使智能体跨越语言边界，并在每次更改后强制重建 。

无隐式间接调用。生产级框架通过隐式间接调用（存储的可调用对象、插件注册表或字符串键解析）从共享的层骨架中组合出许多模型变体。这种模式实现了跨模型的代码重用，但使得通过局部阅读难以确定在给定调用点运行的代码。图2展示了一个生产级框架中模型构建的实例：TransformerLayer从一个单独文件中的运行时规范（runtime spec）解析其子模块。一种扁平化的代码结构以牺牲跨模型重用为代价，换取了局部可读性，从而减少了智能体为建立端到端理解所付出的努力。

任务特定的智能体技能。智能体技能【2，Anthropic. Equipping agents for the real world with Agent Skills. 2025. URL: https://www.anthropic.com/engineering/equipping-agents-for-the-real-world-with-agent-skills】是编码智能体按需加载的程序化手册。技能编码了智能体仅通过静态阅读无法恢复的程序性知识，因此它能运行一个经过验证的程序。智能体技能是最近的实践，现有的训练框架尚未大规模采用：Megatron-LM提供了六个用于CI/PR审查的技能；TorchTitan提供了一个开发者工作流技能和四个用于引导智能体熟悉代码库的编辑器规则；DeepSpeed提供了两个包含通用规则的markdown文件。这些技能都未针对重复性的训练框架任务 。

表1：对智能体原生设计原则的遵守情况。✓ 满足；◐ 部分满足；✗ 不满足。数字是Python、C++和CUDA的总框架代码行数。
表1：对智能体原生设计原则的遵守情况。✓ 满足；◐ 部分满足；✗ 不满足。数字是Python、C++和CUDA的总框架代码行数。

生产级框架中的模型构建（通过隐式间接调用）

PithTrain

图2：模型构建模式，阐释了“无隐式间接调用”原则。隐式间接调用模式通过运行时规范解析子模块，支持从共享的层骨架中进行模型变体；PithTrain直接实例化层，偏好局部可读性。
图2：模型构建模式，阐释了“无隐式间接调用”原则。隐式间接调用模式通过运行时规范解析子模块，支持从共享的层骨架中进行模型变体；PithTrain直接实例化层，偏好局部可读性。

A2 方法细节

3.2 系统架构与优化

PithTrain的架构与范围。代码库分为三层（应用、引擎、算子），如图3所示。为了实现紧凑的代码库，我们识别了分布式MoE训练框架的必要组件，而PithTrain恰好覆盖了这些组件。生产级框架提供了广泛的开箱即用模型、功能和硬件覆盖，而PithTrain则缩小了范围，以保持代码库的紧凑性，并使其能在单个上下文窗口内被处理。表2将每个原则映射到其实现位置：我们采用扁平的代码结构，没有插件注册表或运行时规范，每个MoE模型都存在于models/下的一个自包含文件中，而不是通过共享的层骨架组合而成。这优先考虑了局部可读性而非跨模型的代码重用。

图3：PithTrain架构及各组件代码行数。
表2：PithTrain中各原则的实现位置。
表2：PithTrain中各原则的实现位置。

支持的特性与目标。PithTrain支持标准的MoE训练，包括流水线并行（PP）、通过FSDP【45，Yanli Zhao et al. Pytorch fsdp: Experiences on scaling fully sharded data parallel. Proc. VLDB Endow., 2023】实现的数据并行（DP）、上下文并行（CP）【20，Hao Liu et al. Ringattention with blockwise transformers for near-infinite context. ICLR, 2024】、专家并行（EP）【17，Dmitry Lepikhin et al. {GS}hard: Scaling giant models with conditional computation and automatic sharding. ICLR, 2021】、FP8训练【22，Paulius Micikevicius et al. Fp8 formats for deep learning, 2022】以及DCP检查点【37，The PyTorch Team. torch.distributed.checkpoint. 2023. URL: https://docs.pytorch.org/docs/stable/distributed.checkpoint.html】，并支 持NVIDIA Hopper和Blackwell GPU。尽管PithTrain的代码库紧凑且为Python原生，但它通过采用标准的MoE优化技术，旨在实现与生产级框架相媲美的训练吞吐量；这些技术本身并非新颖，但对PithTrain的训练吞吐量至关重要，值得特别指出。

DualPipeV流水线调度与计算通信重叠。PithTrain的流水线调度器建立在DeepSeek-V3【11，DeepSeek-AI et al. Deepseek-v3 technical report. 2025】的DualPipe之上。DeepSeek的开源版本提供了一个最小化的流水线编排框架，PithTrain在此基础上实现了实际的计算-通信重叠。每个Transformer层在专家并行（EP）通信边界被分解为五个阶段。EP的all-to-all操作在独立的通信流上运行，调度器将一个微批次（micro-batch）的前向传播与另一个微批次的后向传播重叠。

Torch编译优化。PithTrain对除MoE前向和后向传播之外的所有Transformer计算应用了torch.compile(fullgraph=True)。这种严格模式在编译时拒绝图中断（graph breaks），而不是静默地降低加速效果。我们排除了MoE的前向和后向传播，因为在专家并行下，每个专家的输入形状是数据依赖的。

其他优化技术。PithTrain还实现了：wgrad延迟【29，Penghui Qi et al. Zero bubble (almost) pipeline parallelism. ICLR, 2024】；用于提升吞吐量和减少激活内存占用的融合SwiGLU【34，Noam Shazeer. Glu variants improve transformer, 2020】前向和后向内核；用于降低all-to-all通信量的EP分发去重（EP dispatch deduplication）；跨微批次的FP8权重缓存以避免重新量化；以及用于EP token分散和FP8量化的融合Triton内核。

3.3 智能体技能

Agent Skills的定义和组成。一个技能（skill）为一个重复性的训练框架任务编码了执行过程。PithTrain提供了一套技能，涵盖了几个常见的任务（见表3）。每个技能都是一个自包含的文件夹，包含一个顶级的SKILL.md手册，可选的附加markdown文档，以及可选的辅助脚本。一些技能是纯markdown文件，如add-new-model；另一些则捆绑了脚本以分担确定性的工作。

表3：PithTrain中的代表性技能。
表3：PithTrain中的代表性技能。

Agent Skills的设计属性。PithTrain中的每个技能都围绕三个属性设计：具体范围、明确的前提条件和可验证的成功。图4以validate-correctness技能为例说明了这些属性。描述（description）和触发器（triggers）共同编码了技能的具体范围。前提条件（prerequisites）部分列举了环境、数据和配置的假设，以便在技能开始运行前捕获缺失的状态。执行过程以一个脚本调用结束，该调用返回一个可复现的PASS/FAIL结论，而不是依赖智能体的自我评估。围绕这些属性设计的技能在技术上并不难编写，我们期望随着实践的成熟，其他训练框架也会提供类似的覆盖。

图4：PithTrain中的validate-correctness技能。粉色标签标记了属性。
图4：PithTrain中的validate-correctness技能。粉色标签标记了属性。

A7 补充细节

4 ATE-Bench：一个用于评估智能体-任务效率的基准

ATE-Bench的设计理念。评估一个训练框架的智能体-任务效率需要改变框架，同时保持智能体和任务不变。这与标准的AI编码智能体基准测试【14, 7, 4】相反，后者固定代码库和任务，改变智能体以评估其能力。我们引入了ATE-Bench，这是一个使用固定智能体和精选任务套件在不同框架上运行的基准，使得智能体性能的差异可归因于框架设计。该套件涵盖了研究人员通常在训练框架上执行的各类工作，围绕三个重复出现的模式组织：在不修改框架的情况下理解它、将其作为工具进行插桩和分析、以及用新功能扩展它。

任务类别。任务分布在三个类别中：
* 问答（Q&A，12个任务）：回答答案是代码属性而非运行时测量的问题（例如，“设备网格是如何构建的？”）。
* 操作与分析（Operate and Profile，4个任务）：作为工具运行、插桩和分析框架（例如，捕获Nsight Systems剖析并识别最昂贵的CUDA内核）。
* 新功能（New Feature，4个任务）：根据已发布的参考实现，将新的模型架构端到端地移植到框架中（例如，Mixture of Block Attention (MoBA)【21，Enzhe Lu et al. Moba: Mixture of block attention for long-context llms. arXiv, 2025】）。

各类别Agent循环与评估指标。图5展示了每个类别的智能体循环，智能体的参与程度从问答（只读）到操作与分析（运行、少量插桩）再到新功能（大量修改、测试驱动迭代）逐渐加深。完整的任务描述和各类别正确性检查见附录B。使用ATE-Bench，我们评估了PithTrain和生产级框架，并报告了五个工作量指标：会话时长、活跃GPU时间、智能体轮次、每轮上下文和输出token数。由于没有单一标量指标来衡量智能体-任务效率，我们独立报告每个维度。

ATE-Bench的范围和限制。问答问题被选择为在所有三个框架中都有效，排除了特定于框架的行为。ATE-Bench不涵盖像跨模型传播共享变更这样的任务，在这种任务中，生产级框架的隐式间接调用可能会降低智能体的工作量；我们将其作为未来工作。

图5：各类别智能体循环。不同类别的步骤、工具和输出有所不同。
图5：各类别智能体循环。不同类别的步骤、工具和输出有所不同。

A4 实验环境与结果

实验环境

  • 模型:

    • 训练效率对比: GPT-OSS-20B【26】, Qwen3-30B-A3B【39】, DeepSeek-V2-Lite【10】。
    • ATE-Bench: DeepSeek-V2-Lite【10】作为基础模型。

  • 数据集:

    • ATE-Bench新功能任务: DCLM【18】。
    • 下游任务评估: OpenBookQA【23】, WinoGrande【33】, ARC-Challenge & ARC-Easy【9】, HellaSwag【43】, PIQA【5】。

  • 硬件配置:

    • GPU: NVIDIA H100 和 B200 GPU。
    • ATE-Bench Operate/New-feature任务平台: 单节点，配备8个NVIDIA H100 GPU。

  • 软件配置:

    • 对比框架: PithTrain (23db182), Megatron-LM (3bec9aa), TorchTitan (d84e83d)。DeepSpeed因不支持PP与EP组合的MoE训练而被排除。
    • AI智能体: Claude Code (Opus 4.7 at xhigh effort)。
    • 关键库: PyTorch, FSDP, TransformerEngine, Nsight Systems, vLLM, lm-evaluation-harness。
    • 训练配置: 在并行策略（PP, DP, CP, EP）、序列长度和精度（BF16, FP8）上进行匹配。

实验结果

训练正确性验证 (附录A)

  • 预训练损失: 在Qwen3-30B-A3B模型上，PithTrain和Megatron-LM的预训练损失曲线在整个运行过程中基本对齐，验证了训练过程的数值一致性（图7）。
  • 下游任务准确率: 在六个标准基准测试中，两个框架训练出的模型在所有检查点上都取得了在统计误差范围内的可比准确率（图8）。

训练效率 (§5.1)
* 实验内容: 在H100和B200 GPU上，比较PithTrain、Megatron-LM和TorchTitan在三种MoE模型上的训练吞吐量（tokens/sec）。
* 实验结果: 如表4所示，PithTrain在5个配置中的4个上达到或超过了Megatron-LM的吞吐量，在第5个配置上也仅落后1.4%。
* 分析结论: 结果表明，一个紧凑的、Python原生的代码库通过采用DualPipeV和torch.compile等优化，可以实现与生产级框架相媲美的训练吞吐量。

表4：各框架的训练吞吐量。我们报告聚合的每秒token数作为训练吞吐量。“—”表示不支持；“OOM”表示内存不足。
表4：各框架的训练吞吐量。我们报告聚合的每秒token数作为训练吞吐量。“—”表示不支持；“OOM”表示内存不足。

智能体-任务效率 (§5.2)
* 实验内容: 在ATE-Bench上，使用固定的Claude Code智能体评估PithTrain、Megatron-LM和TorchTitan的ATE。
* Q&A任务结果: PithTrain在回答12个代码理解问题时，智能体轮次比Megatron-LM少高达67%。这归功于其紧凑的代码库和无隐式间接调用设计，缩小了搜索空间（表5）。
* Operate & Profile任务结果: 在操作和分析任务中，PithTrain的智能体轮次比Megatron-LM低70%，比TorchTitan低57%；输出token数也相应减少。这得益于紧凑的代码库和仓库内技能的辅助（表6）。
* New Feature任务结果: 在集成新功能的任务中，PithTrain的活跃GPU时间比Megatron-LM低44%，比TorchTitan低64%，主要因为它需要更少的训练重试次数。Megatron-LM的重试由隐式间接调用和C++段错误引起，TorchTitan则主要因内存压力调试。PithTrain的错误信息是可读的Python回溯，且修复通常在本地文件内完成（表7）。

Agent Skills消融研究 (§5.3)
* 实验内容: 在PithTrain上，对validate-correctness和capture-nsys-profile两个任务，分别在启用和禁用相应技能的情况下进行测试。
* 实验结果: 启用技能后，智能体端的开销指标（轮次、上下文、token数）显著下降，其中智能体轮次分别下降了70%和52%。活跃GPU时间基本不变，因为它由任务本身决定（表8）。
* 分析结论: 任务特定的仓库内技能为智能体提供了固定的执行计划，显著减少了其在设置、启动和解释结果时的推理开销。

案例研究：集成MoBA (§5.4)
* 实验内容: 深入分析在集成MoBA架构时，智能体在不同框架上的行为差异，包括输出token的类别分布（图6a）和每轮的上下文窗口大小（图6b）。
* 实验结果: 在所有框架上，编辑（Editing）操作的token数都占主导。Megatron-LM在探索（Exploring）上花费了更多token。TorchTitan因内存溢出（OOM）调试导致上下文窗口出现峰值。
* 分析结论: PithTrain的紧凑性和无隐式间接调用降低了智能体的探索成本，并将修复限制在局部。Megatron-LM的故障（如命令行标志冲突）需要跨多个文件修复。PithTrain的故障（张量步幅不匹配）则在智能体刚编辑的同一文件中被修复。

图6：在各框架上集成MoBA时的智能体行为。(a) 报告了三次独立尝试中每个动作类别的中位数输出token数；(b) 显示了三次独立尝试中每一轮的输入端上下文。
图6b

A5 结论

本文介绍了PithTrain，一个基于四个设计原则构建的紧凑、智能体原生的MoE训练框架。PithTrain在一系列模型上实现了与生产级框架相当的吞吐量，并且在ATE-Bench上，编码智能体在PithTrain上表现出更高的智能体-任务效率。我们希望PithTrain能成为未来智能体原生训练框架的起点。

A6 附录

附录A 训练正确性

本附录验证了PithTrain在匹配配置下相对于Megatron-LM的训练正确性。我们报告了两个互补的测量结果：预训练损失轨迹（§A.1）和下游任务准确率（§A.2）。

A.1 预训练损失。实验描述。我们使用完全相同的配置，用Megatron-LM和PithTrain预训练Qwen3-30B-A3B。结果。图7报告了完整的训练配置以及两条损失曲线；两条轨迹在整个运行期间保持一致，其中Megatron-LM显示出两次瞬时峰值，但在几个步骤内恢复。

图7：使用Megatron-LM和PithTrain训练Qwen3-30B-A3B的训练配置（左）和预训练损失曲线（右）。
图7：使用Megatron-LM和PithTrain训练Qwen3-30B-A3B的训练配置（左）和预训练损失曲线（右）。

A.2 下游任务准确率。实验描述。我们在六个标准基准上评估下游任务准确率：OpenBookQA【23】和WinoGrande【33】（0-shot设置），以及ARC-Challenge、ARC-Easy【9】、HellaSwag【43】和PIQA【5】（10-shot设置）。结果。图8绘制了每个任务的准确率；在所有六个基准测试的每个检查点，Megatron-LM和PithTrain在统计噪声范围内实现了相当的准确率。

图8：使用Megatron-LM和PithTrain训练的Qwen3-30B-A3B的下游任务准确率。
图8：使用Megatron-LM和PithTrain训练的Qwen3-30B-A3B的下游任务准确率。

附录B 任务套件

本附录扩展了各类别任务的描述以及用于验证每次尝试的正确性检查。所有操作与分析及新功能任务共享一个固定的训练配置：基础模型为DeepSeek-V2-Lite【10】（其训练可在单节点8个NVIDIA H100 GPU内完成），并行网格为PP=4, EP=2, DP=1，序列长度2048，全局批量大小1024，精度为BF16。

B.1 问答（Q&A）任务。任务描述。每个问答任务要求智能体在框架代码库中定位特定行为的位置。智能体接收一个包含通用指令和具体查询文本（Q1-Q12）的单一提示，并拥有只读工具权限（Read, Grep, Glob, Bash）。通用指令要求智能体通过引用代码文件和行号来验证其所有声明。任务列表。Q1-Q12分别询问了进程组/设备网格初始化、配置传播、数据加载与分片、分布式种子管理、注意力核调度、RoPE实现、SwiGLU/MLP块、归一化位置、上下文/序列并行、FSDP/DDP封装、全局梯度裁剪和分布式检查点序列化等具体实现细节。正确性检查。由两名独立的人类评分员追溯智能体答案中引用的代码位置，验证代码是否如其所述。所有108次尝试（12个问题 × 3个框架 × 3次尝试）均被两位评分员判定为满意。

B.2 操作与分析任务。任务描述。这些任务要求智能体端到端地执行一个真实的训练系统工作流。四个任务涵盖了：Getting Started（设置环境并运行冒烟测试）、Train and Evaluate（驱动完整的训练-导出-评估流程）、Collect Routing Trace（插桩以收集MoE路由轨迹）、Report Heavy Kernels（使用Nsight Systems分析并报告最耗时的CUDA内核）。正确性检查。通过自动化脚本和人工检查相结合的方式验证智能体生成的产物。例如，Getting Started任务通过检查训练脚本是否能运行到第5步且损失有限来验证。所有36次尝试均产生了被接受的产物。

B.3 新功能任务。任务描述。每个任务模拟研究员将最近发表的模型架构集成到训练框架中的工作流。智能体根据arXiv论文和参考实现，将新功能集成到基础模型中。四个任务包括：Diff【42】（差分注意力）、DynMoE【13】（动态专家数MoE）、MoBA【21】（块注意力混合）和MoE++【15】（零计算专家）。正确性检查。正确性从两方面验证：首先，交叉熵损失在64步运行中必须下降且保持有限；其次，智能体的修改必须满足三个任务特定的规则，这些规则验证了新功能的关键架构组件是否被正确实现。所有36次尝试均满足其任务的所有三个规则。

附录C 各任务结果

本附录报告了三个任务类别中每次尝试的智能体工作量：表9和表10涵盖12个Q&A任务，表11涵盖4个操作与分析任务，表12涵盖4个新功能任务。每行是一次独立的尝试。会话时长和活跃GPU时间以分钟为单位报告。在所有指标上，数值越低越好。

表9：Q&A任务Q1-Q6的每次尝试智能体工作量（§5.2）。每行是一次独立尝试。会话时长以分钟为单位。
表9

表10：Q&A任务的每次尝试智能体工作量（续，Q7-Q12）。
表10

表11：操作与分析任务的每次尝试智能体工作量（§5.2）。每行是一次独立尝试。会话时长和活跃GPU时间以分钟为单位。
表11

表12：新功能任务的每次尝试智能体工作量（§5.2）。每行是一次独立尝试。会话时长和活跃GPU时间以分钟为单位。
表12

方法细节中的参考文献引用

以下是在《方法细节》、《智能体原生设计原则》和《补充细节》章节中引用的参考文献及其描述：

  • [2] Anthropic. Equipping agents for the real world with Agent Skills. 2025. URL: https://www.anthropic.com/engineering/equipping-agents-for-the-real-world-with-agent-skills

    • 引用段落: 3.1节 “任务特定的智能体技能”
    • 引用描述: 该段落将智能体技能（agent skill）定义为一个编码代理按需加载的程序化手册，并引用[2]作为该定义的来源。

  • [11] DeepSeek-AI, et al. Deepseek-v3 technical report, 2025.

    • 引用段落: 3.2节 “DualPipeV流水线调度与计算通信重叠”
    • 引用描述: PithTrain的流水线调度器建立在DeepSeek-V3的DualPipe之上，并引用了[11]作为其来源。

  • [17] Dmitry Lepikhin, et al. {GS}hard: Scaling giant models with conditional computation and automatic sharding. ICLR, 2021.

    • 引用段落: 3.2节 “支持的特性与目标”
    • 引用描述: 该段落列举了PithTrain支持的并行化技术，其中包括专家并行（EP），并引用了[17]作为EP的参考文献。

  • [20] Hao Liu, et al. Ringattention with blockwise transformers for near-infinite context. ICLR, 2024.

    • 引用段落: 3.2节 “支持的特性与目标”
    • 引用描述: 该段落列举了PithTrain支持的并行化技术，其中包括上下文并行（CP），并引用了[20]作为CP的参考文献。

  • [21] Enzhe Lu, et al. Moba: Mixture of block attention for long-context llms. arXiv preprint arXiv:2502.13189, 2025.

    • 引用段落: 4节 “任务类别”
    • 引用描述: 在介绍ATE-Bench的“新功能”任务类别时，以将MoBA架构移植到框架中为例，并引用了[21]作为MoBA的参考文献。

  • [22] Paulius Micikevicius, et al. Fp8 formats for deep learning, 2022.

    • 引用段落: 3.2节 “支持的特性与目标”
    • 引用描述: 该段落列举了PithTrain支持的功能，其中包括FP8训练，并引用了[22]作为FP8格式的参考文献。

  • [25] NVIDIA. TransformerEngine. 2022. URL: https://github.com/NVIDIA/TransformerEngine

    • 引用段落: 3.1节 “Python原生的代码库”
    • 引用描述: 在讨论生产级框架的非Python原生特性时，指出Megatron-LM的核心Transformer层由树外的TransformerEngine模块组成，并引用了[25]作为TransformerEngine的来源。

  • [29] Penghui Qi, et al. Zero bubble (almost) pipeline parallelism. ICLR, 2024.

    • 引用段落: 3.2节 “其他优化技术”
    • 引用描述: 该段落列举了PithTrain实现的其他优化技术，其中包括wgrad延迟，并引用了[29]作为该技术的参考文献。

  • [34] Noam Shazeer. Glu variants improve transformer, 2020.

    • 引用段落: 3.2节 “其他优化技术”
    • 引用描述: 该段落列举了PithTrain实现的其他优化技术，其中包括融合的SwiGLU前向和后向内核，并引用了[34]作为SwiGLU变体的参考文献。

  • [37] The PyTorch Team. torch.distributed.checkpoint. 2023. URL: https://docs.pytorch.org/docs/stable/distributed.checkpoint.html

    • 引用段落: 3.2节 “支持的特性与目标”
    • 引用描述: 该段落列举了PithTrain支持的功能，其中包括DCP检查点，并引用了[37]作为PyTorch分布式检查点的官方文档。

  • [45] Yanli Zhao, et al. Pytorch fsdp: Experiences on scaling fully sharded data parallel. Proc. VLDB Endow., 2023.

    • 引用段落: 3.2节 “支持的特性与目标”
    • 引用描述: 该段落列举了PithTrain支持的并行化技术，其中包括通过FSDP实现的数据并行（DP），并引用了[45]作为PyTorch FSDP的参考文献。