互联网资讯 / 人工智能 · 2023年12月30日

Transformer推理速度提升4.5倍

近日,一款名为InfinITy的产品被推出,能够以1毫秒的延迟完成Transformer的推理,展现了卓越的性能。

让Transformer的推理速度提高4.5倍,这个trick还能给你省十几万

不过,尽管技术出色,其价格也不菲,年费用至少在十几万人民币左右(约2万美元)。

那么,有没有更为经济的替代方案呢?

答案是肯定的!有一种开源的方法,能够在几乎不费力的情况下,达到InfinITy的一些公共基准。

让Transformer的推理速度提高4.5倍,这个trick还能给你省十几万

目前,通过在此方法上进行一些小调整,可以将Transformer的推理速度提升4.5倍。

让Transformer的推理速度提高4.5倍,这个trick还能给你省十几万

那么,这种“平替”如何能够实现“付费”的效果呢?

提升Transformer推理速度的小技巧

首先,介绍一下这个方法:Transformer-deploy。

让Transformer的推理速度提高4.5倍,这个trick还能给你省十几万

该方法只需一行命令即可优化和部署Hugging Face上的Transformer模型,并且支持大多数基于Transformer编码器的模型,如BERT、RoBERTa、MiniLM、CamemBERT、ALBERT、XLM-R、DistilBERT等。

让Transformer的推理速度提高4.5倍,这个trick还能给你省十几万

Transformer-deploy推理服务器使用的是NVIDIA Triton。

推理引擎包括Microsoft ONNX Runtime(用于CPU和GPU推理)和NVIDIA TensorRT(仅限GPU)。

如果想在GPU上获得最佳性能,NVIDIA Triton与NVIDIA TensorRT的组合无疑是理想选择。

虽然TensorRT使用上略显复杂,但其性能确实比PyTorch快5到10倍。

让Transformer的推理速度提高4.5倍,这个trick还能给你省十几万

在实际性能测试中,使用Batch size为1、Token分别为16和128的输入序列时,Transformer-deploy的推理速度均快于收费的Hugging Face InfinITy:

  • Token为16时,Transformer-deploy需要1.52毫秒,而InfinITy需1.7毫秒;
  • Token为128时,Transformer-deploy为1.99毫秒,InfinITy为2.5毫秒。

让Transformer的推理速度提高4.5倍,这个trick还能给你省十几万

那么,之前提到的小技巧究竟是什么呢?

答案是GPU量化(quantization)。

作者提到,目前尚无任何云服务使用这种方法。

不过,执行GPU量化需要对模型源代码进行修改(在矩阵乘法等高开销操作上添加特定节点),这既容易出错,也比较繁琐,用户需要维护修改后的代码。

因此,作者已为多个基于Transformer的模型手动完成了这一过程。

随后,他们发现仅需修补模型模块的抽象语法树(AST),也能够自动完成。

在用户端,GPU执行模型的基本量化类似于以下形式:

让Transformer的推理速度提高4.5倍,这个trick还能给你省十几万

最终,该方法在RoBERTa-base模型和MNLI数据集(分类任务)上实现了4.53倍的推理速度提升。

让Transformer的推理速度提高4.5倍,这个trick还能给你省十几万

当然,这也牺牲了0.4个点的精度;如果不牺牲精度,速度提升大约在3.2倍左右。

作者表示,与Transformer-deploy的原版本相比,这已经是一个显著的改进,毕竟原版本的加速成本需要超过1个点的精确度。

最终,他们用ALBERT、BERT(包括MiniLM)、DistilBERT、RoBERTa(包括CamemBERT、XLM-R、DistilRoBERTa等)、Electra测试了这一技巧。

结果表明,任何可以导出为ONNX格式的Transformer模型,都能“开箱即用”。