分布式训练加速TorchAcc概述

更新时间: 2023-06-27 09:36:37

PAI-TorchAcc(Torch Accelerator)是基于PyTorch的训练加速框架,通过GraphCapture技术将PyTorch动态图转换为静态执行图,然后进一步基于计算图完成分布式优化、计算优化,从而提高PyTorch模型训练的效率,使其更加易于使用。

技术简介

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

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

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

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

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

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

  • 训练数据读取优化:

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

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

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

产品架构

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

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

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

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

image..png

TorchAcc的核心逻辑如下:

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

阿里云首页 机器学习平台 PAI 相关技术圈