ATB是什么?

1 ATB介绍

Ascend Transformer Boost加速库(下文简称为ATB加速库)是一款高效、可靠的加速库,基于华为Ascend AI处理器,专门为Transformer类模型的训练和推理而设计。

ATB加速库采用了一系列优化策略,包括算法优化、硬件优化和软件优化,能够显著提升Transformer模型的训练和推理速度,同时降低能耗和成本。具体来说,ATB加速库通过优化矩阵乘法等核心算子和注意力机制的实现方式,实现了对Transformer模型的高效加速。此外,ATB加速库还充分利用了Ascend AI处理器的硬件特性,如算力、存储带宽和内存带宽,通过硬件加速和数据重用等技术,进一步提升了性能和效率。ATB加速库目前提供了底层基础的高性能算子以及高效的算子组合技术(Graph图算子),同时上层支持对接多种模型框架如PyTorch、MindSpore、Paddle。

总而言之,ATB加速库中包含了各类Transformer类模型的高度优化模块,在各种应用场景中发挥重要作用,为模型的训练和推理提供了强有力的支持。

2 ATB的软件架构

image.png

ATB架构图
从上述架构看,ATB优化的核心内容是在算子计算方面,通过优化算子计算的方式,比如:增加算计算并行机会,优化内存排布等。涉及到3种算子,如下所述。

2.1 基础Operation(原生算子)

用户可以根据需求使用对应的算子完成计算功能。这类算子为一系列基础算子,提供了如矩阵乘、转置等功能。详细信息请参考atb/infer_op_params.h和atb/train_op_params.h。

2.2 插件(Plugin)机制(插件算子)

自定义插件算子(PluginOperation)是一种为用户实现特定功能提供的机制。如果一些功能通过单算子或图算子无法实现,用户可以通过开发自定义插件算子实现对应的功能。

2.3 Graph Frame(图算子)

提供图算子(Graph)机制,用户根据模型设计对应的图算子,使用加速库提供的原生算子和创建的自定义算子创建图算子,完成相应的计算。

graph也就是图,其实就是将多个算子组合成一个图的形式进行调用。那么组合成图的形式有什么好处呢?

1 算子调用

说明这个问题前,首先来了解下算子的下发过程。举例来说,如下图:


image.png

算子host->device下发过程
一个算子的下发,过程大概要经过host(CPU)-> Device(Ascend, GPU...)的过程。

在这个过程中,Host做的事情大概又可以做如下划分:

1、python侧。模型一般是用python编写,算子首先发起调用是从python侧的接口;

2、c++侧。python侧发起后,一般是通过pybind技术,路由到C++侧,这个过程一般由AI框架完成,如torch或、mindspore等。在c++侧做的事情,一般有算子的infer(mindspore),算子输入输出、workspace内存申请,调用Device侧的接口(驱动层)。

那么明显的,从host->devcie是有个调用过程,这个过程肯定是存在开销的。

那如何能减少这个开销呢?ATB中graph frame的实现,就是将算子组成图的形式,将其整体一次性下发到device上 ,那么host->device的开销是可以极大减少的。这样端到端下的性能也就提升了。

2 ATB 图算子构建

用户需要自行设计并定义图的结构,即图结构中节点(Node)的组合与依赖关系,包括节点对应的单算子、节点的输入Tensor与输出Tensor,并识别这些Tensor为图的输入Tensor、输出Tensor和中间Tensor。

图输入Tensor,为用户使用图算子时需要从外部输入的所有Tensor。

图输出Tensor,为用户使用图算子时不会再进行下一步运算操作的所有Tensor。

图中间Tensor,为图算子运算中产生的临时Tensor。例如同一图算子中有算子A和算子X,算子A的输出Tensor为算子X所用,且该tensor对于图外部模块不可见,则该Tensor为中间tensor。

如下图所示,该示例图算子由两个节点组成,这两个节点均为Elewise_Add算子。a、b、c为图算子的三个输入Tensor,output为图算子的输出Tensor,a_add_b_output为图算子的中间Tensor。其中,节点0的输入为a和b,输出为a_add_b_output,节点1的输入为a_add_b_output和c,输出为output。

用户需要自行设计并定义图的结构,即图结构中节点(Node)的组合与依赖关系,包括节点对应的单算子、节点的输入Tensor与输出Tensor,并识别这些Tensor为图的输入Tensor、输出Tensor和中间Tensor。图输入Tensor为用户使用图算子时需要从外部输入的所有Tensor。图输出Tensor为用户使用图算子时不会再进行下一步运算操作的所有Tensor。图中间Tensor为图算子运算中产生的临时Tensor。例如同一图算子中有算子A和算子X,算子A的输出Tensor为算子X所用,且该tensor对于图外部模块不可见,则该Tensor为中间tensor。

如下图所示,该示例图算子由两个节点组成,这两个节点均为Elewise_Add算子。a、b、c为图算子的三个输入Tensor,output为图算子的输出Tensor,a_add_b_output为图算子的中间Tensor。其中,节点0的输入为a和b,输出为a_add_b_output,节点1的输入为a_add_b_output和c,输出为output。


image.png

3. ATB的算子执行流程

如上所述,ATB存在三种算子,那么每种算子的设计思考是什么,与Ascend上,或者说CANN(CANN:CANN-昇腾社区 (hiascend.com))中的算子区别是什么。由于篇幅问题,请移驾到下述文章。
ATB三种算子的执行区别

本文主要参考:
ATB加速库

加速库使用指导

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 216,001评论 6 498
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 92,210评论 3 392
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 161,874评论 0 351
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 58,001评论 1 291
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 67,022评论 6 388
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 51,005评论 1 295
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 39,929评论 3 416
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 38,742评论 0 271
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 45,193评论 1 309
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,427评论 2 331
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,583评论 1 346
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 35,305评论 5 342
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,911评论 3 325
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,564评论 0 21
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,731评论 1 268
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 47,581评论 2 368
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 44,478评论 2 352

推荐阅读更多精彩内容