分布式训练加速（TorchAcc）

技术简介

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

• 通过 GraphCapture 将动态图转换为静态图。

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

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

• 通过半精度通信、通信压缩、通信重叠等技术提高通信效率。

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

• 训练数据读取优化：

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

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

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

产品架构

深度学习框架按执行模式可分为两类：

• graph mode：以 TensorFlow 1.* 为代表。系统优化友好、训推一体，但代码不够灵活，开发和调试困难。

• eager mode：以 PyTorch 为代表。灵活性好、易于开发和调试，但系统优化困难。

TorchAcc 在保持 PyTorch 灵活性的同时，对模型训练做系统优化。架构图如下。

TorchAcc 的核心逻辑如下：

LazyTensor 和 HybridDispatcher 将 PyTorch 的 eager execution 转换为 IR 表达式。TorchAcc 对 IR 做计算优化、显存优化和自动并行化，再交给后端进一步优化和生成代码。同时，TorchAcc 提供手动优化的 kernel，针对特定算子提升计算效率。