TVM－MLC LLM 调优方案

本文地址：https://www.cnblogs.com/wanger-sjtu/p/17497249.html

LLM 等GPT大模型大火以后,TVM社区推出了自己的部署方案，支持Llama，Vicuna，Dolly等模型在iOS、Android、GPU、浏览器等平台上部署运行。

https://github.com/mlc-ai/mlc-llm

本文在之前作者介绍的基础上,简要介绍一下mlc的调优部署方案。

pipeline

在正式介绍TVM mlc.ai部署LLM方案之前，首先简要介绍一下当前主流LLM的一个工作流程。

需要说明一点的是，上图中的prefill跟Decode指的的同一个模型，只是输入的shape存在差异。

这里的示意图省略了很多，只是大致描述一下pipeline。

在处理用户输入时，此时长度大小是不能确定的，这时候是完全的是一个完全的动态shape的。但在decode过程中由于是token by token的，这时候网络中的中除了kv cache相关几个部分，其他大多数的操作都是固定shape的，就可以用已有的算法调优了。

MLC.AI 部署调优方案

以下以RedPajama3B模型的tuning跟build过程介绍一下mlc的方案。

pipeline 组成

在已经支持的几个模型里面均有get_model 这个函数，在这个函数里面会创建下面4个IRModel。

• encoding_func
• decoding_func
• create_kv_cache_func
• create_softmax_func
• create_metadata_func

encoding_func

这对应了上图中的prefill过程，在每次用户输入后调用。由于用户输入的不确定性，所以这个过程基本上都是动态shape的，很难确定到底输入是多大，也不适合搜索调优。

decoding_func

这是上图中decode过程的一部分，因为这个过程是token by token的，在计算过程中大部分的计算是固定shape的。

create kv cache func

这里是直接调用的relax.vm中的函数，创建的是kv cache的存储相关。

create softmax func

这个也是解码过程的一部分，确切的说是采样过程中计算的一部分

** create_metadata_func **

模型的meta信息，比如model_name、stop_tokens等

部署优化

构建完以后，就进入到优化的阶段了。下面根据build.py过程描述一下过程。

1. API构图构建了相关的模型，读取权重

2. 量化

3. 优化PASS

   1. FuseTransposeMatmul
   2. FuseDecodeMatmulEwise
   3. DeadCodeElimination
   4. LiftTransformParams
   5. split_transform_deploy_mod

4. Codegen 生成代码

   1. DispatchTIROperatorAdreno/DispatchTIROperator/DefaultGPUSchedule 手动优化的sch
   2. MetaScheduleApplyDatabase搜索的log生成固定shape的sch

Tuning

在MLC-LLM的代码仓里面已经提供了tuning的脚本，有一点需要先做一下，先调用build.py的文件，把静态shape的相关的函数分离出来。就得到了tuning文件中需要的mod_tir_static.py