Projects
Created
Nov 9, 2025 01:29 PM
Tags
项目背景
Transformer架构和Scaling Raw规律的碰撞下孕育了ChatGPT3，这个划时代产品的诞生为21世纪带来了新的技术时代。Scaling Raw提到，通过扩大Transformer模型大小可以使得训练后的模型具备极强的泛化能力。随着新的AI模型层出不穷，如何将大语言模型部署到个人手机或者个人PC上成为了新的研究问题。端侧模型推理和常见的云侧模型推理又具备不同的计算特性，端侧模型推理的应用场景中主要集中在个人使用中，大部分情况下Batch Size都不会很大，所以端侧模型推理大部分情况都是Memory Bounded。我们也通过简单的测试验证了这个基本的结论。[setup: Falcon-7B Model Inference deployed on 4090]
对于非长文本推理的场景中，模型推理过程中主要的显存开销来源于模型参数本身。因此如何在保证模型推理精度的同时降低模型推理过程中的显存开销成为了新的研究思路。
Opportunity
前面我们在4090对模型推理过程进行了Breakdown，我们发现围绕FFN层展开的IO操作和计算花费了更多的时间。因此我们将目光放到FFN层的结构上。Falcon-7B模型采用Non-Gated架构，该架构的计算流程可以简述如下：
其中对于现有模型架构来说，W1矩阵的长宽比大部分均为1:4，W2则反之(4:1)。
我们观察到，如果我们可以将激活函数拟合成一个线性函数，那么就可以利用矩阵运算的结合律将W1和W2这两个大矩阵合并成一个长宽比为1:1的小矩阵，减小模型参数量（如果激活函数可以用线性函数完美拟合，则FFN层的模型参数可以节省87.5%）。但是现实是骨感的，现在大部分模型使用的激活函数都是非线性的，比如Gelu或者Silu等等：
直接在整个定义域内对激活函数进行线性拟合会使得整个模型“胡言乱语”，可以说是捡了芝麻丢西瓜了。
Observation
将FFN层参数折叠的思路似乎变得不切实际，一个简单的思路就是缩小线性拟合的区间，对整个定义域进行拟合效果能好才怪嘞。那么我们就需要搞清楚缩小线性拟合空间是否是一个可行的方案呢？针对这个问题，我们分析了不同的数据集，发现了两个规律：
- 激活函数的输入具备Skewness: 大部分激活函数的输入都会落入某个特定的区间内
- 不同Transformer Layer以及同一个Layer的不同Neuron对推理结果的重要性是不同的
第一个观察使得只对一个小区间进行线性拟合成为了可能，第二个观察则告诉我们不同的神经网络层对结果影响不同，对于那些权重较高的神经网络层，我们应该选择相对小的区间，防止细小的误差却对结果产生巨大的影响。
Design
基于前面的观察和思路，我们提出了LLMFolder的架构
LLMFolder接受大语言模型、校准集、以及一个threshold（标识最多有多少输入神经元能够落入到选定的线性区间，用于限制线性拟合区间的大小），LLMFolder的架构分为两个部分：
Offline部分：
- 分析校准集，收集每个Transformer层和神经元的激活函数输入的分布，并根据线性拟合后误差给予不同的权重。
- 基于上一步骤分析的结果，对激活函数进行线性拟合，同时将各个FFN层参数进行Folding。最后输出Constant Folding后的矩阵以及一个用来预测FFN层输入经过W1矩阵处理后是否会落入线性拟合区间的预测器。
Online部分：
在实际Serving过程中，LLMFolder使用Folded Matrix进行FFN层的计算，同时使用Offline阶段得到的Predictor预测神经元经过W1层处理后是否会落在线性拟合区间内，如果存在一些异常值则需要：
1. 先减去使用Folder Matrix所进行的错误计算（减去在这个神经元上的即可）
2. 使用之前的激活函数重新计算该神经元处理后的结果
举一个具体的例子
上图中我们截取了某个模型的FFN层进行介绍（为了演示方便我们对模型进行了简化，只有两个神经元），该FFN层包含两个神经元分别对应图中黄色和绿色的矩阵。Offline阶段做的事情就是通过分析得到不同的神经元的Importance Score，基于这个结果确定不同神经元的线性拟合区间。在上面的例子中，黄色的神经元对应的线性拟合区间就是(-1.5, 0.12)，而绿色的神经元对应的线性拟合区间则是(-3.5, -0.1)。Offline阶段LLMFolder基于这两个信息对这两个区间的激活函数使用最小二乘法进行线性拟合将W1矩阵和W2矩阵合并成图中蓝色部分的矩阵。同时通过训练或者量化等方式得到一个Predictor，该Predictor的输入是Online Inference的输入（aka so called hidden state），输出是Bool变量用于表示该输入是否落入到对应神经元的线性区间内。在Online Serving阶段，如果Predictor预测某个神经元处理结果out of linear range，则我们需要单独减去该神经元的计算结果，结合原来的activation function重新处理该神经元对应的输入。而这就是投机执行过程中投机失败所带来的Penality。
Thoughts
这是我研究生期间和博士一起共事的第一个完整项目，主要负责前期调研和后期Online Serving中性能优化的部分。从工作完成到工作总结却Gap了整整一年的时间…我至今仍然清晰地记得，《总结LLMFolder》这件事情在我的Notion代办上一推再推…最后终于在年底忙完大部分事情之后得以将这个TODO一笔勾销。不禁感叹，不管做什么事情，Never Mind Perfect。近期我看着去年记录在Notion上面的科研记录，望着那些恐怕再无人能懂的记录真的恨不得穿越过去扇自己一巴掌。可是再打开我们最近刚刚竣工的项目记录，我思索了一下觉得还是算了吧，不自欺欺人的话，我也该扇现在的自己一巴掌：）
我，美其名曰是一名系统领域研究者，我自己看来也就只是一个科研民工（能高攀这个称号吗），尽管如此，我仍然从中收获到不少经验教训。传统系统研究需要对每个现象，每个结论做足充分的验证，一步一个脚印。只可惜，在科研过程中我老是出现“That's sufficient.”的想法，这是很常见的现象，当你发现性能已经符合要求同时你脑海中已经有一个mental model可以解释这个现象，那么你还会思考：“性能是最优的吗？还有优化的可能吗？我们的mental model正确吗？需要进一步验证吗？”诸如此类问题。而科研过程中坑绝大部分时候就来源于此。几年前如果有人告诉我“慢就是快，科研工作，做得越慢，做得越快”，我只认为自己又遇到一个串子…现在则不然，或许他只是实现了科研进度的辩证统一呢？🫠