MSCTools 设计特点

MSCTools 设计特点#

MSCTools是MSC架构独有的核心设计之一,也是MSC区别于传统AI编译器或压缩工具链的最重要的点。MSCTools的设计也经历了多次迭代,从一开始的耦合于torch的工具链,到依附于TensorRT的特定功能模块,再到最终的完全从训练和部署框架解耦变成独立体系,可以说一路走来都是教训不断推动着我对模型压缩逻辑进行拆分解耦。

一方面,压缩算法和训练框架耦合太深会存在对硬件信息失去感知的情况。例如随着硬件配套的生态逐步迭代,图优化和算子融合的策略也会持续更新,这会影响到量化策略,这样在训练框架中开发量化就不得不持续去反推硬件行为,随着硬件的选择越来越多,这种方式的维护成本会迅速增加。

另一方面,压缩工具和部署框架耦合太深则放弃了训练能力,基本只能选各种PTQ的方式,对于压缩效果更好但需要训练的稀疏化、剪枝和蒸馏等技术基本只能放弃。而比较好的压缩技术往往都需要配合训练,所以耦合部署框架的压缩工具上限并不高。

经历了各种坑之后,我在设计MSC时选择将压缩工具抽象出来,这样:1.方便新的压缩工具的开发,以及向新训练/部署框架集成;2.统一管理调配各种压缩工具,从而实现不同压缩算法的配合,如剪枝+量化+蒸馏。但相应的开发MSCTool有一些代价,即压缩算法开发过程使用的基础数据变成了MSCGraph这一层抽象,而不是torch.Module、tvm.IRModule这类具体的框架计算图。

MSCTool流程#

MSCTool和MSCRunner共同作用对模型进行压缩,基本流程如下:

make plan#

如果MSCTool配置中指定的plan文件不存在,则首先通过MSCRunner的build方法构建用于生成plan的runnable对象,build过程中MSCTool会根据strategys的配置在构建model时插入埋点,例如数据收集节点(量化)和weight channel分析节点(剪枝、稀疏化)。之后MSCRunner加载数据集进行forward,MSCTool插入的节点完成数据的分析生成plan。

apply plan#

如果MSCTool找到了plan文件或者通过2.1生成了plan文件,则在MSCRunner构建runnable对象时根据plan对tensor进行改造,例如对tensorr进行量化/反量化操作(量化),或者对weight进行修剪(剪枝)。

MSCTool将压缩过程分成两个阶段主要是考虑和模型交付系统配合,一般make plan阶段对资源的消耗较大,且可能依赖蒸馏、训练等费时手段,所以对此阶段结束时使用plan保存结果,起到checkpoint作用。交付系统可以直接加载plan从而跳过make plan这一高消耗阶段。