TePDist自动分布式训练系统基础设施

TePDist (TEnsor Program DISTributed) 是一个用于 DL 模型的自动分布式训练系统基础设施，而不仅仅是一种算法。

TePDist 系统以客户端/服务器模式运行。客户端应该是任何可以生成 XLA HLO 的前端。服务器负责分布式策略规划和自动分布式任务启动。将客户端和服务器解耦的动机是为了方便未来与不同前端框架的集成。TePDist 有自己的运行时图和任务调度器，用于分布式运行。

TePDist 系统现在是在以前版本的社区 TensorFlow 的基础上开发的。此存储库中设置了链接到原始代码版本的子模块以供参考。我们计划很快将代码迁移到更新的社区版本。

特征

TePDist 选择 HLO 作为分布式策略规划的输入 IR。我们见过的最大模型包含数万条 HLO 指令。我们的系统可以轻松处理这种规模。在 HLO 级别，指令之间的连接是稀疏的。大多数指令只读取一两个其他指令，尽管指令可能超过数千条。对于 SPMD 策略探索，分布式通信的成本来自于这些指令之间的连接。连接的稀疏性为 TePDist 提供了探索 HLO 策略的机会。

TePDist 的分布式策略探索是完全自动化的。TePDist 的自动计划策略可以涵盖目前已知的各种并行方案，例如Data 并行（包括token 并行）、Model 并行（例如sharding 或Zero）和Pipeline 并行。当然，TePDist也允许用户通过标注的方式介入策略探索。

TePDist 对策略探索问题进行了合理的分解。TePDist 使用多种方法将策略探索问题分解为优化子问题，并使用多种算法分别求解，有效地管理了问题的复杂性。总之，TePDist 根据关键节点将整个图划分为子图（有关更多详细信息，请参阅论文）。在子图中，锥体被进一步划分。子图之间使用动态规划算法，而子图内的锥体之间使用ILP算法。

流水线阶段自动分区很有特色。在划分阶段之前，无需将 DAG 排列成拓扑线性序列。TePDist 将阶段划分建模为 ILP 问题，并使用 ILP Solver 自动找到通信量最少的切割方案。