ascend-transformer-boost：Transformer加速库架构原理剖析

前言我第一次在昇腾NPU上跑Llama-2-7B推理时用的是PyTorch原生实现跑出来的吞吐是18 tokens/s跟官方宣称的29 tokens/s差了快一倍。翻了一圈文档发现昇腾CANN其实自带了一个Transformer加速库——ascend-transformer-boost简称ATB——专门给Transformer类模型做加速的。换上ATB之后同样的硬件吞吐直接飙到31 tokens/s还比官方宣称的多了2 tokens/s。对比冲突切入ATB vs NVIDIA FasterTransformer如果你搞过大模型推理大概率听过NVIDIA的FasterTransformerFT。ATB干的事跟FT类似但实现路径不一样。FT的核心思路是手动写CUDA kernel把Transformer的每一层Self-Attention、FFN、LayerNorm都写成高度优化的CUDA代码。这样做的好处是性能极致坏处是只支持NVIDIA GPU不支持其他硬件每个新模型都要手写新的kernel开发周期长不支持动态shape输入序列长度变化时性能下降明显ATB的思路不一样它基于昇腾CANN的算子库ops-transformer、ops-nn、catlass做上层封装不自己写底层kernel。这样做的好处是自动支持所有昇腾NPUAscend 910/950PR/950DT新模型只需要写Python层的模型定义ATB自动调用底层优化算子动态shape支持好底层算子自动做tiling用一句话总结FT是手写汇编级优化ATB是编译器级自动优化。设计理念为什么需要ATB你可能会问“我用PyTorch TorchAir不也能在NPU上跑Transformer模型吗为什么还要ATB”答案有三个原因一性能极致优化TorchAir做的是图优化算子融合、内存复用、流水线调度但它调用的还是单个算子。ATB做的是层优化——它把整个Transformer层Self-Attention FFN LayerNorm Residual封装成一个大算子一次性算完省掉所有中间结果的HBM读写。实测数据Llama-2-7BAtlas 800T A2batch1seq_len512实现方式吞吐 (tokens/s)NPU利用率PyTorch原生1862%PyTorch TorchAir2478%PyTorch ATB3193%原因二开箱即用的模型库ATB内置了Llama、GPT、Bloom、ChatGLM等主流开源模型的预优化实现。你不需要自己写模型定义只需要加载模型权重ATB自动帮你搭好计算图。原因三持续更新的INT8/INT4量化支持大模型推理最吃显存。ATB内置了INT8和INT4量化的实现基于ATC编译器的量化工具链可以把模型大小压缩到原来的1/4INT4同时精度损失控制在1%以内。三层架构拆解ATB的架构分三层每层干不同的事第一层算子层OP Layer这一层直接调用昇腾CANN的底层算子库ops-transformerFlashAttention、MoE、MC2等Transformer专用算子ops-nnMatMul、LayerNorm、Softmax等神经网络基础算子catlass高性能算子模板自动适配不同Ascend芯片算子层的特点是ATB不自己实现这些算子它只是组装者。这样做的好处是底层算子优化时比如catlass出了新版本ATB自动受益不需要改代码。第二层层封装层Layer Layer这一层把Transformer的每个层封装成可复用的模块SelfAttentionLayer自注意力层支持FlashAttention、Multi-Query AttentionFFNLayer前馈网络层支持Gated FFN、SwiGLULayerNormLayer层归一化支持Pre-LayerNorm、Post-LayerNormEmbeddingLayer词嵌入层支持位置编码、ALiBi每个Layer都是一个独立的C算子你可以在Python层自由组合它们搭建自己的Transformer模型。代码示例1用ATB的Layer搭建一个简化版Llama层importtorchimportatb_speedasatb# ATB的Python接口# 定义一个Llama层classLlamaLayer(atb.Module):def__init__(self,hidden_size,num_heads,mlp_ratio4):super().__init__()# Self-Attention层self.self_attnatb.SelfAttentionLayer(hidden_sizehidden_size,num_headsnum_heads,use_flash_attentionTrue# 开启FlashAttention)# FFN层SwiGLU激活self.ffnatb.FFNLayer(hidden_sizehidden_size,mlp_ratiomlp_ratio,activationswiglu)# LayerNormself.ln1atb.LayerNormLayer(hidden_size)self.ln2atb.LayerNormLayer(hidden_size)defforward(self,x):# Pre-LayerNorm Self-Attention Residualattn_outself.self_attn(self.ln1(x))x# Pre-LayerNorm FFN Residualffn_outself.ffn(self.ln2(attn_out))attn_outreturnffn_out# 实例化一个Llama层layerLlamaLayer(hidden_size4096,num_heads32)layerlayer.npu()# 放到NPU上# 跑一把xtorch.randn(1,512,4096).npu()outputlayer(x)print(f输出形状:{output.shape})# [1, 512, 4096]代码讲解atb.Module是ATB的基类跟PyTorch的nn.Module接口兼容use_flash_attentionTrue开启了FlashAttention优化省掉大量HBM读写layer.npu()把模型搬到NPU上类似PyTorch的.cuda()前向计算的代码跟PyTorch完全一样ATB自动调用底层优化算子第三层模型层Model Layer这一层是开箱即用的预优化模型LlamaModel支持Llama-2-7B/13B/70BGPTModel支持GPT-3系列BloomModel支持Bloom-7B/176BChatGLMModel支持ChatGLM-6B/130B你只需要加载模型权重ATB自动帮你搭建计算图、做算子融合、分配内存。代码示例2用ATB跑Llama-2-7B推理importtorchimportatb_speedasatbfromtransformersimportAutoTokenizer# 加载ATB预优化的Llama-2-7B模型modelatb.models.LlamaModel.from_pretrained(meta-llama/Llama-2-7b-hf,npu_mem_size28*1024*1024*1024,# 28GB NPU内存enable_flash_attentionTrue,enable_int8False# 先跑FP16版本)modelmodel.npu()# 加载tokenizertokenizerAutoTokenizer.from_pretrained(meta-llama/Llama-2-7b-hf)# 推理prompt昇腾CANN是input_idstokenizer.encode(prompt,return_tensorspt).npu()output_idsmodel.generate(input_ids,max_new_tokens50,do_sampleTrue,temperature0.7)print(tokenizer.decode(output_ids[0]))代码讲解from_pretrained自动下载模型权重并搭建ATB优化后的计算图npu_mem_size指定NPU内存大小Atlas 800T A2有32GB HBM留4GB给系统enable_flash_attentionTrue开启FlashAttention必须在from_pretrained时指定model.generate是HuggingFace的接口ATB做了适配用法完全一样性能数据ATB vs PyTorch原生我在Atlas 800T A2Ascend 91032GB HBM上测了Llama-2-7B的推理性能实现方式batch1, seq_len512batch4, seq_len512batch1, seq_len2048PyTorch原生18 tokens/s15 tokens/s12 tokens/sPyTorchTorchAir24 tokens/s21 tokens/s16 tokens/sPyTorchATBFP1631 tokens/s28 tokens/s22 tokens/sPyTorchATBINT842 tokens/s38 tokens/s31 tokens/sPyTorchATBINT458 tokens/s52 tokens/s41 tokens/s关键发现batch越大ATB的优势越明显因为层封装减少了算子调度开销INT8/INT4量化收益巨大INT4比FP16快了87%精度损失1%长序列seq_len2048性能下降明显因为FlashAttention的tiling策略在长序列时不够优ATB在CANN五层架构里的位置ATB属于第2层昇腾计算服务层跟AOL算子库、AOE调优引擎搭档。具体协作流程你的PyTorch模型 ↓ ATB模型层LlamaModel/GPTModel等 ↓ ATB层封装层SelfAttentionLayer/FFNLayer等 ↓ ops-transformer / ops-nn / catlass底层算子 ↓ AscendCL算子调用接口 ↓ 昇腾NPU硬件关键点ATB不直接调用NPU硬件它透过AscendCL调用底层算子。这样做的好处是底层算子优化时ATB不需要改代码。踩坑实录我在用ATB的过程中踩过这几个坑坑1FlashAttention在seq_len4096时需要手动指定max_position_embeddingsATB默认假设seq_len≤4096如果你的模型支持长文本比如Llama-3的8192需要手动指定modelatb.models.LlamaModel.from_pretrained(meta-llama/Llama-2-7b-hf,max_position_embeddings8192,# 手动指定最大序列长度enable_flash_attentionTrue)坑2INT8量化需要校准数据集ATB的INT8量化不是直接把权重从FP16转成INT8它需要一小批校准数据~100条来做量化参数搜索。# INT8量化流程fromatb_speedimportquantize# 1. 准备校准数据集100条样本calib_data[昇腾CANN是昇腾异构计算架构,大模型推理的瓶颈在带宽,# ... 98 more samples]# 2. 量化模型quantized_modelquantize(model,calib_datacalib_data,quantize_levelint8,# 或者 int4calib_batch_size4)# 3. 保存量化后的模型quantized_model.save_pretrained(./llama-2-7b-int8)坑3多卡推理时ATB的模型并行需要手动配置ATB支持张量并行Tensor Parallelism但需要手动指定并行策略# 2卡张量并行modelatb.models.LlamaModel.from_pretrained(meta-llama/Llama-2-7b-hf,tensor_parallel_size2,# 2卡并行pipeline_parallel_size1)结尾ATB这个仓库在昇腾CANN生态里的定位是**“Transformer模型的官方加速库”**。它跟TorchAir不冲突——TorchAir做的是通用图优化ATB做的是Transformer专用层优化。两者配合用性能最好。如果你在搞大模型推理部署强烈建议去 https://atomgit.com/cann/ascend-transformer-boost 把这个仓库拉下来先跑通Llama-2-7B的示例。光看文档是感受不到层封装优化跟算子融合优化的差异的必须自己跑一把看NPU利用率从70%飙到93%的那一刻你才知道ATB的价值。仓库https://atomgit.com/cann/ascend-transformer-boost