阿里云机器学习PAI发布基于HLO的全自动分布式系统 TePDist，并宣布开源！

导读

一直以来，大模型在模型效果上被证明具有显著优势。而ChatGPT的出现，证明了其在工业生产工具方面具有巨大潜力。毫无疑问，大模型的训练需要巨大的算力，这对分布式并行框架是一项考验。现如今，业界越来越多的团队纷纷转向分布式框架的研究与开发之中，既有针对特定场景特定模型的极致手工优化，也包括面向通用模型通用场景的自动分布式工作。然而，它们在实际生产场景仍然有一些挑战。一方面，对于手工优化的框架，虽然其特定情况下的性能较为极致，但通用性存在限制。另一方面，现有的自动分布式框架可能在通用性和性能方面表现较好，但通常需要忍受较长的策略搜索时间，进而对业务落地产生一定的影响。为了应对以上问题，PAI 团队推出并开源了TePDist（Tensor Program Distributed System），它通过在HLO上做分布式策略搜索，实现与用户模型语言解耦。并且在保持通用性的同时，在可接受的策略搜索时间内，追求高性能分布式策略。

TePDist背后的技术框架如何设计？未来有哪些规划？今天一起来深入了解。

TePDist是什么？

TePDist是阿里云PAI团队自研的基于HLO IR层的全自动分布式深度学习系统，它不仅仅是一个分布式Compiler，还拥有自己的分布式Runtime，有效地解决了深度学习模型并行策略的自动搜索与分布式策略实施问题。

在架构方面，TePDist采用Client/Server模式，实现分布式策略与模型描述的解耦。Server端是TePDist最重要部分，它以HLO IR作为输入，自动探索并实施分布式并行策略；Client端以用户描述的模型为输入，将其转换成HLO IR。因此，任何具有转换HLO IR能力的Client，都可经过适配后接入Server端。

在功能方面，TePDist分为两个部分。一是在HLO IR上进行SPMD（Data Parallel和Sharding）和Pipeline并行的策略搜索。并以此构建编译基于Task Graph的执行计划。二是高效运行执行计划的分布式执行引擎。

在使用方面，TePDist提供了不同优化级别，高优化级别更加追求分布式策略质量，低优化级别会额外采取一些Heuristic，以较为微小策略质量牺牲，换取更快地搜索时间，以此满足落地需求。

项目开源地址：https://github.com/alibaba/TePDist

TePDist的主要特性如下：

• 一套完整的系统：采用Client/Server模式。Client可以是任何能够生成XLA HLO的前端。Server负责进行分布式策略规划，以及自动分布式任务拉起。Client与Server解耦的动机是期望将来更方便地对接不同的前端框架。
• 以HLO IR作为Server端输入：HLO的粒度刚刚好，目前看到的超大模型基本上有上万条HLO指令。在HLO这个层级做分布式策略规划，可以避免依赖灵活多变的模型高层次抽象，以不变应万变。
• 全自动探索分布式策略：TePDist支持在没有任何annotation的情况下做策略搜索。当然，用户也可以选择通过annotation进行一定程度的干预。
• 定义不同优化级别：用户可以选择