小朋友做安全教育的网站,青岛网站网站建设,wordpress 摘要,前端开发的工作内容作者#xff1a;沈雯婷、黄奕桐、艾宝乐、王昂、李永
1、简介
PAI-TorchAcc(Torch Accelerator)是阿里云人工智能平台开发的Pytorch上的大模型训练加速框架。
PAI-TorchAcc提供了一套基于Pytorch的简洁、易用的接口#xff0c;无需进行模型转换就可以无缝地接入HuggingFac…作者沈雯婷、黄奕桐、艾宝乐、王昂、李永
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等兼容阿里云上灵骏集群的硬件资源。
与现有框架对比
模型支持分布式训练策略算子优化性能DeepSpeedHuggingFace模型用户自定义模型ZeRO/PP手写Kernel优化⭐️Megatron只有GPT/BERT/T5等少数模型DP/TP/PP手写Kernel优化⭐️⭐️⭐️PAI-TorchAccHuggingFace模型用户自定义模型DP/FSDP/TP/PP/SP自动编译优化手写Kernel优化⭐️⭐️⭐️
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 CallHigh-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的优化。
引用
[1] https://arxiv.org/pdf/2205.01068.pdf
[2] https://huggingface.co/blog/ram-efficient-pytorch-fsdp
[3] https://github.com/NVIDIA/Megatron-LM
[4] https://github.com/microsoft/DeepSpeed
[5] https://github.com/pytorch/xla
