分布式训练加速（TorchAcc）

技术简介

TorchAcc是动静一体的分布式训练加速框架，主要功能如下：

• 通过GraphCapture技术将动态图转化为静态图。

• 通过编译优化手段提升训练性能。

• 通过显存优化降低资源开销。

• 通过半精度通信、通信压缩、通信overlap等通信优化技术来提高通信效率。

• 提供自动和半自动分布式策略，支持大模型高效训练。

• 训练数据读取优化：

  • Prefetcher：进行数据预取，让数据预处理和训练能够同时进行，从而减少数据处理的等待时间，提高训练效率。

  • Packed dataset：通过高效的数据打包方式，减少无效计算，并提高数据读取效率。

  • Preprocess Cache：缓存预处理后的数据，减少数据预处理开销。

产品架构

深度学习框架按照执行模式可分为两个大的类别：

• graph mode：以TensorFlow 1.*为代表的框架采用graph mode的方式运行。其优点是系统优化友好、面向生产、训推一体，而缺点是面向用户不够友好、代码撰写不够灵活、开发和Debug困难。

• eager mode：以Pytorch为代表的框架采用eager mode的方式运行。其优点是灵活性好、容易开发和Debug。但是对于框架优化来说不友好，系统优化困难。

针对上述问题，TorchAcc的目标是在保持Pytorch灵活性的基础上，为模型训练提供系统的优化处理。TorchAcc的架构图如下所示。

TorchAcc的核心逻辑如下：

通过LazyTensor+HybridDispatcher将PyTorch中的eager execution转换为IR表达式，然后进行计算优化、显存优化以及自动并行化等一系列优化处理后，再将优化处理后的IR交给后端进行进一步的优化和CodeGen。同时，TorchAcc还提供了许多手动算子优化的kernel实现，可以针对特定算子进行优化，来提高算子的计算效率和性能。