本文分享自华为云社区《深度解读昇腾CANN模型下沉技术,提升模型调度性能》,作者:昇腾CANN。
AI模型的运行通常情况下需要CPU和NPU(昇腾AI处理器)等AI专用处理器协同工作,CPU所在位置称为主机端(Host),而NPU所在位置称为设备端(Device)。对于采用Host调度的AI模型来说,Host下发Task的时序和Device执行Task的时序是异步的,如果Device执行Task的速度比Host下发Task的速度快,则Device会处于空闲状态。比如,大模型场景的增量推理或训练的FineTune阶段,都是计算量较小的场景,此时很容易出现单个算子的Host下发时间比Device上的算子执行时间还长,从而导致Device间歇处于空闲状态。这种现象通常称为Host Bound,这种模型也称为Host Bound模型。
如何减少Host Bound模型的Device空闲时间,从而优化模型执行性能显得尤其重要,GE(Graph Engine)图引擎通过图模式的Host调度和模型下沉调度的方式,可提升模型调度性能,缩短模型E2E执行时间。
1 模型Host调度
Host CPU将模型中的算子依次下发到Device执行(如下图中的标号①所示),每一个算子在执行流上以1个或多个Task的形式存在,昇腾AI处理器依次拉取执行流上的Task执行(如下图中的标号②所示),我们称这个过程为Host调度。Host调度需要Host和Device频繁交互,在实际的训练或推理场景中,模型会运行多次,每次运行都会触发Host把模型上的所有算子遍历下发一遍。
图1 Host调度:
1.1 Device Bound与Host Bound
如果Device上Task执行速度比Host下发慢,则Device上Task调度开销和算子执行时间成为瓶颈,这种模型通常称为Device Bound模型,如图2所示。Device Bound的模型,Task的下发开销会被已下发Task的执行时间掩盖起来。
图2 Host调度场景Device Bound执行时序分析:
如果Device上Task执行速度比Host下发快,则Host调度开销成为瓶颈,这种模型通常被称为Host Bound模型,如图3所示。Host Bound的模型,Task的下发开销不能完全被已下发Task的执行时间掩盖,Device执行时序上存在间歇的空闲状态。
图3 Host调度场景Host Bound执行时序分析:
1.2 单算子模式Host调度与图模式Host调度
在前面几期介绍中,我们知道当前业界主流的深度学习框架提供了Eager执行(Eager Execution,即时执行)模式与图执行模式。Eager模式也叫单算子模式,它是一种Host调度模式,由Host CPU逐个下发算子,一个算子的下发流程包含Python处理、Python到C++数据结构转换、Tiling计算,申请算子的Workspace内存和输出内存,Launch等Host操作。为了加速单算子Host调度,在PyTorch中,昇腾适配层采用了生产者---消费者双线程模式加速,生产者线程主要负责Launch之前的处理动作,消费者线程主要负责Launch算子。
相比于单算子模式,图模式的Host调度可以避免总是返回Python调用栈,避免冗余流程与数据结构转换,并且可以直接使用图编译阶段完成的Infer Shape与Tiling计算结果。因此,图模式Host单线程调度与单算子模式Host双线程调度相比,调度性能持平或略优。
2 模型下沉调度
图模式调度可以降低Host侧的调度耗时,并一定程度减少模型执行E2E耗时,但如何降低Device上执行时序的空闲时间仍是需要考虑的问题。对于静态shape模型,昇腾支持下沉调度,可大幅优化Host调度性能。
所谓静态shape模型,是指模型每次执行的输入tensor shape是固定不变的,模型中的所有算子,给定输入shape后,都可以确定自己的输出shape,那么该模型为静态shape模型。
静态shape模型在编译时即可确定所有算子的输入输出shape,结合昇腾内存复用算法,可完成模型级内存编排;静态shape模型在编译时还可提前完成所有算子的Tiling计算等Host侧计算。综上,静态shape模型可以在编译时完成所有执行时的Host调度工作,这是静态shape模型下沉调度的理论基础。
2.1 下沉调度原理
模型下沉调度分为两个阶段,模型加载和模型执行。
图4 模型下沉示意图:
- 模型加载
模型加载的具体动作和Host调度类似,即遍历图中的所有算子并将其整体下发至Device流上,区别在于下发到流上不立即执行。模型加载是一次性的动作,在首次模型执行时完成模型加载,如图4中的过程①所示。
- 模型执行
模型加载完成之后,可以像下发单算子Task一样,向执行流下发一个模型执行Task,昇腾AI处理器调度到该Task时(如图4 "执行流"中的E),执行模型中所有Task(如图4中的过程③)。如果需要多次运行模型,仅需多次下发模型执行Task,如图4中的过程②所示。
Host Bound调度和模型下沉调度的时序比较如图5所示,可以看出模型下沉执行的开始有一个模型下发的头开销,模型下沉执行E2E会有一个相对于Host调度的收益,模型的下沉头开销越小,收益将越大。
图5 模型下沉调度Host/Device时序分析:
每次模型下发时,支持更新模型的Feature Map内存地址和输入输出内存地址,如果模型的Feature Map内存和模型输入输出内存发生了更新,则在模型下沉头开销(即上图中的m_l_t,后文有介绍)中会完成模型内算子相关地址的刷新。
2.2 下沉效果
模型下沉执行存在如下优势:
- 减少CPU的负载
对于模型下沉执行的方式,执行时不需要Host和Device频繁交互,调度的开销仅包含模型下沉的头开销和Task从流上调度到加速器上的开销。总的来说,模型下沉执行的方式,Host CPU的负载降低了,在一定程度上整机、集群的功耗也会降低。
- 减少Rank之间的通信抖动
对于Host调度执行方式,集群场景下由于不同卡上的Host下发时序不一定完全同步,卡间集合通信可能存在一定程度上的抖动。而下沉执行方式则不存在该问题,因为Task已提前排布到Device上。
- 提升E2E的收益
由于Task已提前下沉到了Device,对于Host Bound的模型,E2E会有性能提升。以Host Bound的LLaMA-7B模型的推理场景为例,图6是Host调度的Profiling性能数据,图7是模型下沉执行的Profiling性能数据。通过Profiling结果可以看出该模型是一个Host Bound模型,采用Host调度方式,Device上计算时间和空隙的时间比例接近1:1,采用模型下沉执行方式,Device空隙大幅减少,并且端到端有18ms的收益,吞吐量整体提升37%。
图6 LLaMA-7B Decoding Host调度Profiling:
图7 LLaMA-7B Decoding模型下沉Profiling:
2.3 下沉头开销
下沉执行头开销包括以下几部分:
1、模型输入输出Tensor到内部InputData/OutputData数据结构转换。如图8中的stage1_t。
2、如果模型的Feature Map内存和模型输入输出内存有变更,刷新相关联算子的相关地址。如图8中的stage2_t。
3、如果模型的输入不支持零拷贝,(如模型的输入为Host内存),则下发异步拷贝Task。如图8中的stage3_t。
4、下发模型执行Task。如图8中的stage4_t。
图8 模型下沉头开销拆分:
以盘古71B增量推理模型为例(模型输入输出总个数约1600个,模型内节点总数约6300个,模型运行时,Feature Map内存地址不变更,10个输入输出内存地址变更),当前模型下沉的头开销约2ms,关于头开销的性能优化在持续进行中。
3 更多介绍
GE模型下沉技术的相关介绍就到这里,欢迎大家关注后续技术分享。如需获取更多学习资源请登录昇腾社区。