AI 加速引擎 PAI-TorchAcc：整体介绍与性能概述

2024-03-04
浙江
本文字数：3894 字
阅读完需：约 13 分钟

作者作者：沈雯婷、黄奕桐、艾宝乐、王昂、李永

1、简介

PAI-TorchAcc(Torch Accelerator)是阿里云人工智能平台开发的 Pytorch 上的大模型训练加速框架。

PAI-TorchAcc 提供了一套基于 Pytorch 的简洁、易用的接口，无需进行模型转换就可以无缝地接入 HuggingFace 上的模型，并用多种分布式策略进行训练加速。

PAI-TorchAcc 借助社区 PyTorch/XLA，通过 LazyTensor 技术将 Pytorch 代码转换为静态执行图，基于计算图，结合阿里云上的计算资源情况，进行了大量的 GPU 硬件上模型训练的针对性分布式优化、计算优化。

得益于简单的模型接入方式、基于计算图的优化，PAI-TorchAcc 能够灵活地支持各种大模型的多种规模，兼容不同的硬件。PAI-TorchAcc 支持常见大模型 1B-175B 的训练，训练吞吐相对 PyTorch 原生、Megatron-LM 均有提升，如 LLaMA 系列模型，相比 PyTorch 原生提升了 140%，相比 Megatron-LM 提升了 5%，在 A100 上 MFU 达到 70%，8 卡到 128 卡线性加速比达到 15.6X。

2、背景和需求

2.1 背景

大模型训练

近年来，大语言模型、视频生成类模型迅速发展，它们基于庞大的文本、图片、视频等数据集进行训练，执行多种自然语言处理、图像生成、视频生成等任务，具备强大的理解和生成能力。随着计算资源和技术的不断进步，大模型的参数量已增长到数亿甚至数万亿级别，例如 LLaMA、GPT-3、通义千问、Sora 等，这些模型在许多基准测试上表现出了前所未有的性能。

然而，训练大模型需要极高的成本。比如使用 Megatron-LM 预训练一个 OPT-175B 模型需要上千张 A100 训练 2 个月[1]，硬件利用率 MFU 约 47%，期间因为硬件故障经历了几十次 checkpoint 的加载和续训练。使用 PyTorch FSDP 进行 LLaMA-2-70B 的微调也需要 16 张 A100 运行约 13.5 小时[2]。NVIDIA A100、H100 等硬件资源价格高昂且不易获取，市面上也逐渐出现了其他性价比更高的硬件资源。

加速不同的大模型的预训练、续训练、微调，充分利用不同的硬件资源，提升资源利用率，是降低大模型训练成本的一个有效途径。

Megatron-LM

NVIDIA Megatron-LM[3]是一个基于 PyTorch 的分布式训练框架，用来训练基于 Transformer 的大模型。Megatron-LM 综合应用了数据并行、模型并行、流水并行来实现 GPT-3 等特定模型的训练。然而，不同的大模型、训练数据集接入 Megatron-LM 十分不灵活，需要将 checkpoint 和数据格式进行转换。同时，Megatron-LM 虽然对一些模型算子做了手动的优化，在面对不同模型的不同计算模式时，难以自动地应用这种手动的优化。

DeepSpeed

DeepSpeed[4]是微软开源的一个 PyTorch 上的大模型分布式训练框架，支持 ZeRO 和流水并行，并且可以结合 Megatron-LM 运行 3D 并行。DeepSpeed 已经成为 HuggingFace transformers 库中一个训练组件。然而 DeepSpeed 性能表现较差，并且和 Megatron-LM 同样存在面对不同计算模式时无法灵活优化的限制。

PyTorch/XLA

PyTorch/XLA[5]将 PyTorch 和 OpenXLA 相结合，使用 LazyTenor 技术，将 PyTorch 代码转换为静态执行图，在静态图上进行计算图优化和后端编译优化。Pytorch/XLA 主要是针对 TPU 场景进行优化，在 GPU 上还存在一定问题和优化空间，如不支持 Transformers 模型常用的 FlashAttention 加速算子、不支持 torchrun 拉起、计算通信 Overlap 差、显存开销大等问题。

2.2 需求

基于以上背景，我们需要一个大模型分布式训练引擎，能够方便接入多变的 PyTorch 模型，尤其是 Transformer 类模型，兼容多种硬件。在不同模型变化的计算模式下，在不同硬件变化的硬件架构和计算、访存能力下，能够自动地对计算进行优化，尤其在阿里云的硬件上能够表现较高的性能。同时，大模型导致单卡内存和显存无法完全放下，不同的模型需要结合不同的分布式策略，合理通信，完成多卡训练并提升线性加速比。

3、PAI-TorchAcc 核心技术特性

灵活的模型接入

支持 LLaMA 系列、Qwen、BaiChuan、ChatGLM、OLMo、Bloom 等常见的大模型 1B-175B 的训练；
无缝对接 HuggingFace 中的模型；
一键接入和加速 Pytorch 模型。

千亿级模型参数量

已经支持 1B 到 175B 大模型训练；

全面的训练模式

支持混合精度训练，包括 Float32、Float16、BFloat16 等；
支持 Pytorch 模型的预训练、微调和续训练。

组合的分布式策略

支持 Data Parallel、Tensor Parallel、Sequence Parallel、Fully Sharded Data Parallel、Pipeline 等分布式策略及其组合。

自动计算优化和显存优化

使用手动的 Gradient Checkpoint 和自动的 Rematerialization 降低峰值显存；
自动进行显存规划和管理，降低峰值显存和减少显存碎片化；
自动对 Kernel 进行编译优化，提高计算效率；
自动接入 SOTA 的高性能 Kernel。

兼容多种硬件

兼容 NVIDIA A100/800, H100/800, V100 等；
兼容阿里云上灵骏集群的硬件资源。

与现有框架对比

4、PAI-TorchAcc 架构

4.1 总体架构

PAI-TorchAcc 的架构自顶向下分为以下几层：

模型层：支持计算机视觉、自然语言处理、语音合成等深度学习模型训练的加速；
算法库：支持 HuggingFace Transfomers、PAI-EasyNLP、TIMM 等算法库构建的模型；
前端：支持以 PyTorch 为前端语言的模型训练；
Lowering：使用 LazyTensor、Symbolic Trace 等技术将前端代码转换为静态执行图;
IR：使用多层中间表达，包含 High-Level 的设备无关的 IR 和 Low-Level 的设备相关的 IR，基于两层 IR 上分别做计算图优化和后端编译优化。
编译优化引擎：TorchAcc 的编译优化引擎包括计算图优化引擎 TorchAcc Compiler 和多种后端编译优化引擎 BladeDISC 和 OpenXLA。基于两层 IR，进行分布式优化、显存优化、通信优化、计算优化以及算子调度和显存管理等优化，生成优化的设备码。
硬件：最终产生硬件相关的设备码在不同算力、带宽和显存的硬件设备上执行。

4.2 接口

PAI-TorchAcc 抽取了一套简洁的接口，灵活接入并加速任意的 Pytorch 模型，而不需要改动原有的模型代码。

通过 PAI-TorchAcc 加速模型训练一般需要三步：

定义 torchacc.Config，并指定加速选项。
调用 torchacc.accelerate，并传入 model 和 config，完成加速训练的准备。
通过 torchacc.AsyncLoader 对 torch dataset_loader 进行封装，加速数据加载。

model = ...  dataloader = ...
+ # 一行代码加速模型，也可传入Config配置更丰富的加速功能，如分布式策略、编译优化选项等+ model = torchacc.accelerate(model)
+ # 异步加速数据加载+ dataloader = torchacc.AsyncLoader(dataloader, model.device)
  model.train()  for source, labels in dataloader:      ...

复制代码

4.3 编译优化

PAI-TorchAcc 通过 LazyTensor、Symbolic Trace 等技术将前端 Pytorch 代码转换为静态执行图，并在静态图上进行自动优化，在分布式的硬件设备上高效运行。

4.4 计算图优化

在 Tensor Graph 上进行优化，这层优化基于 High-Level IR——StableHLO 进行。

分布式：通过分图和通信算子插入，完成流水并行、SPMD 等。
显存优化：通过算子级别的显存 Live range 和复用分析、静态调度策略、自动重算、显存管理优化等来减少显存的峰值和碎片化。
计算优化：通过 CSE 等简化计算，通过算子大粒度融合来优化访存密集型算子，减少 kernel launch，减少访存，提升计算效率；通过自动的计算图匹配重写的方式接入 Flash Attention 等高性能 Kernel。
通信优化：通过通信算子的合并、拆分、异步化以及算子的调度来提升通信效率，提高计算和通信的 overlap。

4.5 后端编译优化

在 Buffer Graph 上进行优化，这层优化基于 Low-Level 的 IR，包括 LHLO、LLVM IR 和多种 MLIR 的 dialect。

多后端：支持 OpenXLA 和阿里自研的 BladeDISC 两种编译后端；
Lowering 和 Codegen：将上层的 StableHLO Lowering 成 LHLO 和多种 MLIR 的 dialect，并在各级 Lowering 过程中进行优化，最终表达为 LLVM IR，通过 LLVM 生成针对硬件的优化代码；
Custom Call：High-Level IR 自动 Pattern rewrite 的优化 kernel，通过 custom call 调用。

5、实践案例和性能

PAI-TorchAcc 在 A100 上能够达到 70%的 MFU，并且在多卡下几乎线性扩展（8 卡到 128 卡加速比 15.6X），在灵活支持各种模型的基础上，性能能够高于 Megatron-LM。我们在常见的开源大模型上做了性能测试，使用相同的硬件资源，PAI-TorchAcc 的训练吞吐相对 PyTorch 原生、Megatron 均有提升，如 LLaMA 系列模型相对 PyTorch 原生提升了 140%，相对 Megatron 提升了 5%。

我们将在后续的系列文章中提供一个具体的实践案例：PAI-TorchAcc 在 OLMo 模型训练上的接入示例和加速效果，并且给出加速的来源分析。

6、总结和未来展望

PAI-TorchAcc 可以灵活接入 Pytorch 模型，并通过并行化策略、显存优化、计算优化和调度优化等方法来加速大模型以及视觉类、语音类模型的训练。PAI-TorchAcc 已经在常见大模型上如 LLaMA、LLaMA-2、BaiChuan、ChatGLM、QWen、OLMo、Bloom 取得了不错的效果。未来我们将从以下方向继续深入优化，以支持更多的场景，取得更好的加速效果。

Graph Capture 优化和子图编译：在生成计算图的过程中遇到无法识别的算子将导致编译失败，我们将进一步优化 Graph Capture，并支持子图的编译优化。
自动分布式：PAI-TorchAcc 提供了多种分布式策略，然而在不同的模型和硬件上，使用哪种组合的分布式策略、如何进行分图能够取得最优的性能，仍然需要根据经验手动配置。PAI-TorchAcc 将借助静态计算图和模型、硬件特性，做自动的分布式。
AutoGC：借助静态计算图和模型、硬件特性，自动进行 checkpoint 选点。
动态 Shape 性能优化：动态 Shape 导致重编译引起的性能下降，当前我们通过分桶的方式减少了重编译的次数，仍然存在大量的 padding，如何做更高性能的动态 Shape 支持，是一个深入优化的方向。
自研编译优化引擎 BladeDISC 的优化。