深度解读昇腾CANN模型下沉技术，提升模型调度性能

news2024/9/22 4:11:36

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执行时间。

模型Host调度

Host CPU将模型中的算子依次下发到Device执行（如下图中的标号①所示），每一个算子在执行流上以1个或多个Task的形式存在，昇腾AI处理器依次拉取执行流上的Task执行（如下图中的标号②所示），我们称这个过程为Host调度。Host调度需要Host和Device频繁交互，在实际的训练或推理场景中，模型会运行多次，每次运行都会触发Host把模型上的所有算子遍历下发一遍。

图1 Host调度

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执行时序分析

单算子模式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双线程调度相比，调度性能持平或略优。

模型下沉调度

图模式调度可以降低Host侧的调度耗时，并一定程度减少模型执行E2E耗时，但如何降低Device上执行时序的空闲时间仍是需要考虑的问题。对于静态shape模型，昇腾支持下沉调度，可大幅优化Host调度性能。

所谓静态shape模型，是指模型每次执行的输入tensor shape是固定不变的，模型中的所有算子，给定输入shape后，都可以确定自己的输出shape，那么该模型为静态shape模型。

静态shape模型在编译时即可确定所有算子的输入输出shape，结合昇腾内存复用算法，可完成模型级内存编排；静态shape模型在编译时还可提前完成所有算子的Tiling计算等Host侧计算。综上，静态shape模型可以在编译时完成所有执行时的Host调度工作，这是静态shape模型下沉调度的理论基础。

下沉调度原理

模型下沉调度分为两个阶段，模型加载和模型执行。

图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，后文有介绍）中会完成模型内算子相关地址的刷新。

下沉效果

模型下沉执行存在如下优势：

减少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%。