深度解读昇腾CANN动态Shape图调度加速技术
1 引言
在AI模型运行过程中,通常需要CPU和AI专用处理器(如昇腾AI处理器)协同工作。CPU所在位置称为主机端(Host),而NPU所在位置称为设备端(Device)。主机端擅长处理复杂的逻辑计算,而设备端擅长进行高并行计算。通过高效的计算调度机制,实现Host和Device之间的高效协同是提高AI模型性能的关键,能够显著提升异构系统资源的利用率。
2 GE模型调度方式
图引擎(Graph Engine,简称GE)作为昇腾平台计算图编译和运行的控制中心,提供了两种常规的模型调度模式(下沉调度和Host调度),用于实现Host和Device之间的高效协同。
下沉调度模式通常适用于静态Shape模型,由于输入Tensor的Shape固定不变,可以在编译阶段完成内存编排和Tiling计算,因此可以将模型中的算子以整图形式下发到Device上,执行时,只需在Host侧下发一个模型执行的Task即可触发模型在Device上调度执行,从而提升模型调度性能。这种模式的优势在于Host侧的调度开销小,执行过程中Host和Device之间的交互较少。
Host调度模式通常适用于动态Shape模型,由于输入Tensor的Shape不确定,无法直接编译整图,需在执行时逐算子进行InferShape、Tiling和内存分配,然后将Kernel逐一下发至Device执行。这种模式需要Host和Device之间进行更多的协同,确保算子的正确执行。
关于下沉调度的详细介绍可以参考:《深度解读昇腾CANN模型下沉技术,提升模型调度性能》,本文将重点介绍Host调度优化的关键技术,帮助开发者理解和应用。
3 Host Bound和Device Buond
对于动态Shape模型而言,Host和Device之间的下发和执行是异步进行的。Host将算子下发到Device后,会立即继续调度后续算子,而Device则同时执行已下发的计算任务。这样,后一个算子的Host下发可以与前一个算子的Device执行并行进行。
如果一个模型中Host的下发速度快于Device的执行速度,Device不会因为等待Host下发而产生空闲时间。此时,Device的执行时间成为影响端到端性能的主要因素,这种模型被称为Device Bound模型。
相反,如果一个模型中Host的下发速度慢于Device的执行速度,Device在执行完当前任务后会等待Host下发下一个算子,这会导致Device出现较大的空闲时间,从而延长模型的端到端执行时间。这种模型被称为Host Bound模型。对于Host Bound模型,优化Host的调度时间可以提升模型的端到端执行性能。
两种模型的算子在Host和Device上执行的示意图如下:
在Profiling性能分析视图中,Host Bound和Device Bound模型的特征表现通常为:对于Host Bound网络,Host任务下发速度较慢,Device侧需等待Host下发任务,任务接收后会立即执行;而对于Device Bound网络,Host任务下发速度较快,Device侧可能因任务处理不及而产生滞后现象。
4 Host调度优化技术
动态Shape在编译时和执行时所看到的图不同。编译时看到的图称为计算图(ComputeGraph),它表达了用户的逻辑计算语义。执行阶段看到的图称为执行图(ExecuteGraph),在动态Shape模型的执行图中,一个算子的下发被拆分为多个Host执行单元,每个执行单元称为Host Kernel。每个Host Kernel负责完成独立的运算或下发任务,执行图上的节点与Host Kernel一一对应。动态Shape模型在执行时通过依次执行这些Host Kernel来完成整个执行图的下发。
以一个单输入单输出AI Core算子为例,对应的执行图上的表达如下所示:
从上图可以看出,一个AI Core算子的下发会被拆分为InferShape、Tiling、AllocMemHbm、Launch等多个Host Kernel。这些Host Kernel根据执行图的顺序在Host依次执行。
对于Host Bound网络而言,由于瓶颈集中在Host侧,因此提升动态Shape性能的关键在于减少Host Kernel的总下发耗时。
在一些典型的大模型网络中,Host Bound问题较为严重。除了使用高效的数据结构和算法减少Host Kernel的计算复杂度和执行时间外,提升Host调度性能的几个技术优化方向包括:Host缓存技术、Host多核并发技术、小Shape算子计算优化技术等,下面我们将分别进行介绍。
Host缓存技术
Host缓存通过“空间换时间”的思路,将Host Kernel的执行结果进行缓存,避免重复执行耗时较长的计算,从而优化Host Kernel的执行性能。
以Tiling缓存为例,由于算子的Tiling计算结果依赖于特定的Tiling参数,当一个算子的Tiling计算依赖参数完全一致时,其结果可以复用。对于一些复杂算子,Tiling计算时间可能长达几十微秒,通过缓存算子的Tiling参数,可以避免重复执行Tiling函数,从而减少Host侧的计算开销。
Tiling缓存机制仅对多次执行中可能变化的参数进行缓存。Tiling计算依赖的参数分为编译时参数(如输入格式、数据类型等)和执行时参数(如输入Shape等)。编译时参数在执行时固定不变,而执行时参数需要在运行时确定。查询缓存时,系统会根据Tiling依赖的执行时参数动态生成哈希键(HashKey)。通过查询缓存的哈希键,当后续遇到相同计算场景时,可以直接命中缓存的Tiling结果,跳过冗长的Tiling函数计算过程。这一技术显著减少了Host侧的计算开销,提高了整体执行效率。
系统默认启用Host缓存技术。实测表明,启用Tiling缓存后,在pangu、llama2等大语言模型上Host侧Tiling开销可降低50%以上,且模型参数量越大,优化效果越显著。
在端到端性能方面,尽管模型执行仍呈现Host Bound特征(端到端执行开销与Host调度开销相近),但Host调度开销和整体性能均得到显著优化。特别是在pangu38B量化网络中,由于Host侧Tiling计算开销占比高,其性能提升尤为明显。
Host CPU多核并发技术
Host通过CPU多核并发及多级流水技术,提高下发任务的并发度,从而提升调度性能。GE的动态Shape流程将下发过程分为三级流水,分别是普通线程、内存线程和Launch线程。
- 一级流水(普通线程):主要负责执行不涉及Device内存和Launch操作的Host Kernel,如InferShape和Tiling等。
- 二级流水(内存线程):负责与Device内存相关的Host Kernel,如Device内存分配和释放等操作。
- 三级流水(Launch线程):负责Launch相关的Host Kernel,如将AI Core、AI CPU等引擎的算子下发到Device上执行。
理想的调度效果如下:
- 下一级流水依赖上一级流水:下一级流水中的Host Kernel需要等待上一级流水中依赖的Host Kernel执行完成。理想情况下,上一级流水的执行速度比下一级流水快,使得下一级流水的Host Kernel能够依次紧密执行,避免因等待而产生调度间隙。
- 避免反向依赖:上一级流水不会依赖下一级流水,从而避免产生反向依赖导致流水中断。
以普通AI Core算子为例,大致的执行示意图如下(只列举主要节点):
通过多级流水的设计,Host CPU得到充分利用,三级流水之间可以实现并发执行,上一级流水的时间被下一级流水掩盖,从而显著提升了Host调度性能。
在训练与在线推理场景下,针对动态Shape图模式执行的网络,可通过设置环境变量"MAX_RUNTIME_CORE_NUMBER"来开启图执行器(Host侧)的多线程任务调度,建议将该环境变量配置为3,以获得较优性能。
export MAX_RUNTIME_CORE_NUMBER=3
在典型网络上通过以上配置使能多级流水后测试的调度性能如下:
小Shape算子计算优化技术
对于一些小Shape算子,输入Tensor在Host侧,本身计算量很小,下发到Device执行的调度开销往往大于算子的实际计算开销。将这类算子保留在Host侧执行,可以有效减少调度开销带来的性能影响。
上图是一段网络拓扑片段示例,按照一般的调度机制,Gather、Concat算子会下沉到Device侧计算,Shape、Unsqueeze、Reshape算子在Host侧计算。模型E2E执行耗时除了包含算子计算的时间外,还包含Host与Device之间的数据拷贝、算子下沉调度、Stream同步等开销,整体执行E2E耗时在毫秒级别。
而对于小Shape的Gather、Concat,算子本身在Host侧CPU的计算开销上仅微秒级别,与Device侧计算的性能相差无几,此时下发带来的额外开销就显得比较明显。针对上述这种Shape较小且输入Tensor在Host的场景,GE识别将这部分算子保留在Host侧执行,可有效减少调度开销带来的性能影响。
下图是优化前后的E2E耗时对比图:
此外,对于一些值依赖算子,InferShape和Tiling依赖输入Tensor的内容,这些输入Tensor的Shape通常较小。由于Infershape和Tiling在Host上执行,如果对应的输入Tensor依赖Device的计算结果,将产生额外的流同步和Device到Host的数据拷贝,从而打断Host下发的流水,增加额外开销。
采用小Shape算子优化技术后,不仅可以解决小Shape算子本身下发调度耗时长的问题,还能有效缓解值依赖算子引入的Host断流水问题。
系统默认启用小Shape算子计算优化技术,以LLaMA2大语言推理模型为例,启用该技术后,符合上述条件刷新的算子有Pack、Gather、Concat等约650+个,刷新前模型E2E耗时约1.062s,刷新后执行时间优化到了1.009s,吞吐提升5%。
其他技术
除了以上提到的关键技术之外,GE还提供了其他Host调度优化技术,如算子融合等。算子融合是指在计算图中将多个相邻的算子合并成一个更复杂但更高效的算子。通过将多个算子融合,减少算子的总数,从而减少需要下发的Host Kernel的数量,进而提升Host下发的性能。
5 更多学习资源
GE动态Shape图Host调度优化的相关技术的介绍就到这里,欢迎大家关注后续技术分享。欢迎访问昇腾社区GE信息专区,获取海量学习资源:https://www.hiascend.com/graph-engine。
往期技术文章推荐:
- 点赞
- 收藏
- 关注作者
评论(0)