ATB是什么?

简介: ATB加速库专为华为Ascend AI处理器设计,针对Transformer模型的训练和推理进行了深度优化。它通过算法、硬件和软件层面的优化,大幅提升模型性能,降低能耗与成本。ATB支持PyTorch、MindSpore等多种框架,提供高效的基础算子及图算子技术,适用于各种应用场景。其软件架构主要包括基础Operation、Plugin机制和Graph Frame三部分,通过优化算子计算和数据传输,实现性能的显著提升。

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加速库

加速库使用指导

相关文章
|
机器学习/深度学习 数据可视化 算法
深度学习之梯度下降参数可视化
深度学习之梯度下降参数可视化
|
11月前
|
容器
vllm+vllm-ascend本地部署QwQ-32B
本指南介绍如何下载、安装和启动基于Ascend的vLLM模型。首先,可通过华为镜像或Hugging Face下载预训练模型;其次,安装vllm-ascend,支持通过基础镜像(如`quay.io/ascend/vllm-ascend:v0.7.3-dev`)或源码编译方式完成;最后,使用OpenAI兼容接口启动模型,例如运行`vllm serve`命令,设置模型路径、并行规模等参数。适用于大模型推理场景,需注意显存需求(如QwQ-32B需70G以上)。
4287 17
|
机器学习/深度学习 存储 缓存
ATB概念之:算子tiling
算子 tiling 是一种优化技术,用于提高大规模张量运算的计算效率。它通过将大任务分解为小块,优化内存使用、支持并行计算,并防止内存溢出。在ATB中,tiling data指kernel的分片参数,用于指导计算。ATB提供了三种 tiling data 搬移策略:整体搬移、多stream搬移及随kernel下发搬移,旨在优化内存拷贝任务,提高计算效率。
|
12月前
|
PyTorch 编译器 算法框架/工具
NPU上如何使能pytorch图模式
本文介绍了PyTorch的`torch.compile`技术和TorchAir的相关内容。`torch.compile`通过将动态图转换为静态图并结合JIT编译,提升模型推理和训练效率。示例代码展示了如何使用`torch.compile`优化模型。TorchAir是昇腾为PyTorch提供的图模式扩展库,支持在昇腾设备上进行高效训练和推理。它基于Dynamo特性,将计算图转换为Ascend IR,并通过图引擎优化执行。文章还提供了TorchAir的使用示例及功能配置方法。
|
机器学习/深度学习 安全 物联网安全
探索未来网络:物联网安全的最佳实践与创新策略
本文旨在深入探讨物联网(IoT)的安全性问题,分析其面临的主要威胁与挑战,并提出一系列创新性的解决策略。通过技术解析、案例研究与前瞻展望,本文不仅揭示了物联网安全的复杂性,还展示了如何通过综合手段提升设备、数据及网络的安全性。我们强调了跨学科合作的重要性,以及在快速发展的技术环境中保持敏捷与适应性的必要性,为业界和研究者提供了宝贵的参考与启示。
|
测试技术 开发者 Docker
llama2 70B mindie推理开箱报错问题
遇到问题,先从基本的检查开始,先检查卡有没有被占用,有的话就kill掉(如果是别人的任务,先知会一下哈)!其次,不要自己随意组合版本,否则会因为经验不足卡在莫名其妙的问题上。
247 3
|
前端开发 数据可视化 API
Python实现智能家居设备的统一控制平台
【10月更文挑战第6天】 Python实现智能家居设备的统一控制平台
864 11
|
机器学习/深度学习 人工智能 搜索推荐
AI与未来医疗:重塑健康的双刃剑
【10月更文挑战第6天】 人工智能作为现代科技的巅峰之作,已经渗透进我们生活的方方面面。从语音助手到自动驾驶,AI不仅改变了我们的日常,更在各个专业领域,特别是医疗行业,扮演着愈发重要的角色。本文探讨了AI在未来医疗中的应用及其潜在影响,揭示了这把“双刃剑”的机遇与挑战。
462 1
|
PyTorch TensorFlow API
大模型中 .safetensors 文件、.ckpt文件、.gguf和.pth以及.bin文件区别、加载和保存以及转换方式
本文讨论了大模型中不同文件格式如`.safetensors`、`.ckpt`、`.gguf`、`.pth`和`.bin`的区别、用途以及如何在TensorFlow、PyTorch和ONNX等框架之间进行加载、保存和转换。
5699 2
|
存储 缓存 前端开发
ATB算子实现原理解读
本文详细介绍了Ascend Transformer Boost(ATB)加速库中三种类型算子的执行流程及其与CANN中其他算子的区别。文章首先概述了ATB算子的实现步骤,接着深入解析了单算子和图算子的执行流程,包括kernel图构建、输入准备、内存计算、tiling数据处理及任务下发等环节。此外,还探讨了ATB在host侧性能优化上的几种机制,如Tiling Cache、Setup与InferShape复用、Runner Pool等,以及ATB中的内存优化与管理策略。最后,介绍了Context类的功能和作用,包括它如何管理ATB内部的各种公共资源。

热门文章

最新文章