大模型推理迁移的核心内容是基于加速库提供的基础算子和组图接口,构建与原始模型计算图等价的加速库图,整体流程可以归结为以下几个部分:
- 基础环境搭建。
- 模型分析:基于原始在线推理代码梳理模型计算图,主要包含模型拓扑关系、算子列表、输入输出Shape信息等。
- 模型开发:开发流程如图1所示。
- 自定义Operation开发:基于模型分析得到的计算图,如果有加速库原生不支持的Operation,用户需要自己开发自定义的PluginOperation。
- Layer开发:Layer为Operation和PluginOperation 的组合,与具体模型结构、输入输出定义强相关,如典型的Transformer Block。对于大模型而言,模型结构由N个Transformer Block组成,可以重复调用。Layer的开发需要基于加速库原生的Operation和自定义Operation开发中的PluginOperation进行Layer的组图,同时实现OutShape的推导。
- Model开发:Model为前两者的组合,通常包含了基本完整的模型结构,如Prepare_inputs/TransformerBlock(n个)/lm_head等,能够实现极致的性能优化。开发过程需要基于Layer/Operation组成完整Model计算图,同时完成输入Tensor/输出Tensor/中间Tensor的传递过程。
- 框架接口调用:实现Layer/Model层后,可以在PyTorch侧通过torch.classes的方式实现对应类的调用,替换原始的推理接口。
图1 模型开发流程