近日,阿里云机器学习平台PAI正式发布自研的基于HLO的全自动分布式深度学习系统 TePDist。它通过在 HLO 上做分布式策略搜索,能够与用户模型构建语言解耦。并且在保持通用性的同时,能够在可接受的策略搜索时间内,追求高性能分布式策略,同时用户无需修改模型主体代码,系统能够全自动地帮助用户做分布式扩展,有效解决了分布式框架长期以来在实际生产场景中手工优化和自动分布式工作存在的诸多性能和效率问题。 TePDist 不仅仅是一个分布式 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,以较为微小策略质量牺牲,换取更快地搜索时间,以此满足落地需求。 性能上,TePDist 通过在 GPT 和 MoE 模型上 SPMD+Pipeline 混合策略的模型扩展性实验,TePDist 能够使 GPT 和 MoE 分别达到峰值能力的 62% 和 58%。同时,在自动化方面的通用性上,TePDist 也通过了 VGG-19,DNABert 和 UNet 等模型实验验证。 一直以来,大模型在模型效果上被证明具有显著优势。而 ChatGPT 的出现,证明了其在工业生产工具方面具有巨大潜力。阿里云机器学习平台PAI也宣布将TePDist开源,与AI开发者共同打造更快更好的自动分布式系统,全面助力 AI 大模型发展! 开源地址:https://github.com/alibaba/TePDist 开源钉群: 
| 
