昇腾 CANN 7.0 黑科技:大模型训练性能优化之道
本文分享自华为云社区《昇腾CANN 7.0 黑科技:大模型训练性能优化之道》,作者: 昇腾 CANN 。
目前,大模型凭借超强的学习能力,已经在搜索、推荐、智能交互、AIGC、生产流程变革、产业提效等场景表现出巨大的潜力。大模型经过海量数据的预训练,通常具有良好的通用性和泛化性。用户基于“大模型预训练+微调”开发范式即可在实际业务场景取得理想的效果。从经济成本和时间成本来看,预训练是构筑各种行业模型的主要部分,优化大模型预训练具有非常重要的意义。
Open AI 研究表明:大模型表现好坏强烈依赖于模型规模,弱依赖于架构;模型表现随着计算量、数据量和参数量提升;模型表现随着训练数据量和参数量的提升是可预测的。总体来讲,大参数量、大数据量、大计算量已经成为大模型表现好的主要因素。
来源:OpenAI Scaling laws for neural language models
这样的趋势给大模型训练带来什么挑战呢?
首先是算力问题。1750 亿参数量的 GPT-3 训练 3000 亿 token,需要算力 3.14e11 TFLOPs,千卡 A100 集群训练时长需要 22 天左右(算力利用率约为 50%);1.8 万亿参数的 GPT-4 模型,训练 13 万亿 token,需要算力 2.15e13TFLOPs,2.5 万卡 A100 训练时长达到 90~100 天(算力利用率 32%~36%)。
其次是显存容量问题。1750 亿参数量的 GPT-3 预训练大约需要 3TB 内存,大模型训练在短序列长度时,模型参数是内存占用的主要部分,在长序列长度时,激活内存是主要部分。总体来说,随着序列增长,需要的内存越来越大。
最后是通信开销问题。大模型预训练的高算力和大内存诉求往往通过分布式集群 scale out 的方式应对,但这将带来相当大的通信开销。具体包括模型并行通信开销、数据并行通信拖尾开销、流水并行的通信开销,另外还有流水并行引入的 bubble。通常情况下,通信耗时约占 E2E 耗时的 10%~30%,当存在更大通信域时,通信占比会更大。
因此,大模型预训练需要解决三个核心问题:大模型可部署,在内存受限的 AI 加速卡上放得下;充分发挥算力,提升计算效率,让训练流程更快;提升通信性能,降低通信开销,让计算更专注。
为了释放昇腾硬件算力,昇腾 AI 异构计算架构 CANN 发布更开放、更易用的 CANN 7.0 版本,全面兼容业界的 AI 框架、加速库和主流大模型,同时通过大颗粒算子深度融合、Kernel 调度策略优化、通信并发流水等技术手段,解决大模型训练核心问题,使能大模型性能深度优化。
1 支持分布式切分策略+内存优化策略,让大模型放得下
学术界和工业界发明了很多大模型计算加速策略,包括最为经典的 3D 并行、序列并行、重计算、ZeRO 系列的内存优化策略等等。通过这些并行策略,解决大模型放得下、训的快问题。基于昇腾 AI 处理器的多样算力和丰富的内存资源,异构计算架构 CANN 提供完备的技术栈功能,支持各种并行策略部署,也有很好的扩展能力,能支持新的并行算法和策略。下面来列举下在大模型训练中常用的计算加速策略:
3D 并行(TP/PP/DP)
序列并行
来源:Reducing Activation Recomputation in Large Transformer Models
Checkpoint 完全重计算
选择性重计算
Self-attention 中的 softmax 和 dropout 算子计算结果占用内存大,但计算量相对少,不保留正向计算结果到反向,在反向计算流程可以重新计算 softmax 和 dropout 的结果。
来源:Reducing Activation Recomputation in Large Transformer Models
内存优化:ZeRO1/2/3
来源:ZeRO Memory Optimizations Toward Training Trillion Parameter Models
2 充分发挥昇腾算力,提升计算效率,让模型跑得更快
在大模型计算优化方面,CANN 通过高频算子优化、大颗粒算子融合、构建 Transformer 加速库等提升大模型计算效率。在调度方面,CANN 通过执行器优化、算子 Kernel 重构等有效减少或避免 host bound,提升算子在 Host 侧的下发性能。
1、高频算子优化,充分利用 L1/L2 cache,发挥 cube 大算力
来源:Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
矩阵乘 MatMul 算子是大模型的高频算子,在 GPT-3 模型中,占 60%以上的计算量。CANN 通过算子 tiling 优化合理切分数据,通过算子流水优化使得数据搬运与计算流水并行,同时,通过 cube 与 vector 算子融合,减少内存使用,最终 MatMul 算子 Cube 利用率提升显著,性能达到业界先进水平。
2、Flash/Sparse Attention 大颗粒算子融合,减少显存耗用,提升计算性能
CANN 基于 Ascend C 编程语言开发了 Flash/Sparse Attention 算子,支持 fp16、bf16、fp32 多种数据类型计算,并支持 8K 及以上的长序列,在 bloom、llama、gpt-3 等大模型有很好的应用。相对稠密模型,CANN 还针对 MOE 稀疏模型进行了重要的优化,通过 MOE 融合算子优化内存占用,并实现 All-to-All 通信与计算并发流水,大幅度提升性能。
3、构筑 Transformer 加速库提升核心 Kernel 性能,助力模型加速
基于 Transformer,业界已经诞生了很多大模型,CANN 针对这些模型进行泛化性分析,构筑 Transformer 加速库,主要包括 FlashAttention、MOE FFN、Norm 类等等,加速库为大模型训练保证了高性能下界,客户基于此可以快速改造写出性能更优的算子。
Transformer 高性能库构建策略
4、高性能 Kernel 动态调度和下发,有效减少或避免 host bound
大模型训练不仅要提升 AI Core 上算子的计算性能,也需提升算子在 host 侧的下发性能。若 host 侧下发算子不及时,可能会导致 AI Core 空闲等待,算力利用率降低(即 host bound 现象)。为此,CANN 通过重构 runtime 执行框架,包括优化执行图,提供多种策略调度器,可根据模型结构特点选择最优调度策略,重构算子的 Shape 推导、tiling 策略实现高性能 Kernel,采用用户态直通 device、Host2Device 快速拷贝技术实现高性能 kernel launch,多措并举使能 Tiling 和 Launch 性能倍增,有效减少或避免 host bound。
3 降低通信开销,让计算更专注
大模型训练随着集群规模和通信域的增大,通信耗时占比升高,逐渐成为主要性能矛盾。提升通信性能、降低通信开销显得愈发重要。学术界和工业界发明了很多高性能通信算法,在此不一一介绍。下面介绍两个充分利用通信带宽的优化点:
跨 server all-reduce 场景:SDMA&RDMA 通信流水化
跨 server 实现 all reduce 一般分为四个阶段:server 内 reduce scatter、server 间 reduce scatter、server 间 all-gather,server 内 reduce scatter,在同一时刻,server 间链路和 server 内链路只有一个在工作,由此会造成带宽的浪费。将通信数据切片排布成流水,使 server 内和 server 间的链路并发利用起来,通过小步快跑方式节省通信耗时。
跨 server all-reduce 场景:SDMA&RDMA 通信流水化
单 server 内 all-reduce 场景:TP 通信复用 RDMA 带宽
单 server 内 SDMA 通信(比如 all-reduce)时,server 间的 RDMA 通信链路通常是空闲的。此时将数据按照一定比例切分,分别由 server 内 SDMA 通信、server 间的 RDMA 通信同时传输,从而充分利用 server 内和 server 间的通信链路并发,提升通信性能。
单 server 内 all-reduce 场景:TP 通信复用 RDMA 带宽
4 总结
昇腾 CANN 借助根技术创新对大模型训练过程进行了系统级的优化加速,通过合理的分布式切分策略和内存优化策略,使大模型放得下;通过计算层、调度层和通信域的优化,使大模型跑得快,未来 CANN 也将持续研究、不断突破,以大模型赋能应用,加速千行万业智能化转型。
版权声明: 本文为 InfoQ 作者【华为云开发者联盟】的原创文章。
原文链接:【http://xie.infoq.cn/article/30657f9c8dc51eb1f8f28e64f】。文章转载请联系作者。
评论