TePDist(TEnsorProgramDISTributed)是一个用于DL模型的自动分布式训练系统基础设施，而不仅仅是一种算法。TePDist系统以客户端/服务器模式运行。客户端应该是任何可以生成XLAHLO的前端。服务器负责分布式策略规划和自动分布式任务启动。将客户端和服务器解耦的动机是为了方便未来与不同前端框架的集成。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问题，并使用ILPSolver自动找到通信量最少的切割方案。
评论