架构设计

金箍棒架构

我们分析了一些业界模型压缩框架实践，总结了一些优秀的特点：

• 极简易用：以LLM Compressor为代表，深度集成Transformers，提供极简API设计，学习曲线平缓，易用性极佳。算法高度封装但快速跟进LLM领域前沿算法，满足快速上手、部署、验证的需求，是LLM领域用户的理想选择。

• 精细配置：以DeepSpeed Compression为代表，提供精细化配置API，需要一定的学习成本来理解压缩算法原理和参数。但其丰富且可组合的算法库支持用户像搭积木一样构建压缩流水线，模块化架构便于集成自定义算法，不仅是工具集更是研究平台。

• 极致性能：以TensorRT为例，作为CUDA运行时而非专门的模型压缩工具，通过插件机制（自定义算子、Pass等）实现量化稀疏功能，在NVIDIA硬件上提供极致推理性能，广泛应用于生产环境，但技术门槛较高。

基于上述观察，我们希望金箍棒既是一个容易上手的模型压缩工具，同时也能作为算法研究员的研究平台，我们明确了几个目标：

• 易用且灵活的API：多层级的API，兼顾易用性和灵活性，降低使用门槛，同时保留算法定制化的能力；

• 丰富且模块化的算法库：提供丰富的压缩算法，并且支持灵活组合；

• 高度可扩展的框架架构：便于集成自定义算法模块，配合灵活的API构建定制化压缩流水线。

API设计

金箍棒提供多层次API设计：

• Level 0接口：类似DeepSpeed Compression的配置式API，提供丰富的压缩算法配置选项，支持用户按需定制化配置。

• Level 1接口：类似Transformers的简洁API，支持快速应用内置压缩算法，直接优化Transformers网络和Hugging Face权重。

算法流水线

• 分层架构设计：采用四层架构（Pipeline、AlgoModule、Wrappers、BaseQuantFuncs），支持在任意层级自定义组件，复用底层能力，加速新算法实验。

• 责任链模式：Pipeline和AlgoModule基于责任链设计模式，实现算法模块化，用户可灵活组合不同AlgoModule，构建定制化压缩流水线。

基础设施

兼容Hugging Face权重

本节主要讨论训练后量化技术相关的Hugging Face兼容性实现。

1. Hugging Face权重加载

   • MindOne提供类Transformers接口，支持直接加载Hugging Face格式权重；

   • MindFormers提供专用接口，支持从Hugging Face权重文件构建网络。

2. Hugging Face量化权重导出

   • 将TP并行权重合并为按参数名切分的Hugging Face格式，保存为safetensors文件；

   • 生成safetensors索引文件 index.json；

   • 更新 config.json，添加量化配置信息，供推理框架构建量化方法；

   • 保存量化描述文件，供推理框架构建量化网络；

   • 保留原始权重目录中的其他所有文件。

伪量化评测

本节主要讨论训练后量化技术相关的精度评估方法。

伪量化评测是一种不依赖推理框架的精度评估方法。其核心思想是在浮点网络中插入伪量化算子，通过引入量化误差来模拟真实量化推理过程。这种方法能够独立评估量化算法本身带来的精度损失，排除推理框架或性能优化特性引入的额外误差，便于进行精度分析。

局限性说明：

1. 算法兼容性限制：部分算法无法通过伪量化准确模拟，如FAQuant需要在FlashAttention融合算子的中间计算过程引入量化误差，若网络使用FlashAttention融合算子则无法进行伪量化评估。

2. 性能开销较大：伪量化评测的推理性能显著低于正常推理，特别是对于CoT（Chain-of-Thought）等长序列推理场景，评测耗时较长。

混合策略

为平衡模型压缩率与精度，通常需要根据网络不同层的特性采用差异化的算法策略。

金箍棒提供双层策略配置体系：

• 网络级策略：全局压缩策略，对网络中所有算法敏感层生效；

• 层级策略：通过正则表达式匹配特定层，配置专属压缩策略，优先级高于网络级策略。

昇腾硬件适配层

• 网络融合优化：算子融合是推理优化的基础技术。原始Transformers网络通常未进行融合优化，为确保压缩后网络仍具备融合能力，算法设计需统一考虑融合模块。通过NetworkFuser抽象封装融合能力，配合StatisticMgr避免重复计算，提升整体效率。

• 算子规格适配：昇腾硬件对特定算子有规格要求（如8bit全量化要求左右矩阵均为对称量化，需将offset转换为bias）。金箍棒提供专用适配模块，确保量化权重符合昇腾算子规格要求。

算法与推理解耦

金箍棒采用算法与推理分离的架构设计：

1. 算法独立性：直接对Transformers网络进行压缩，不依赖特定部署框架（如vLLM-MindSpore Plugin）的网络定义；

2. 伪量化评测：支持独立于推理框架的精度评估，便于算法效果验证；

3. 标准化接口：以Hugging Face社区权重格式作为算法与推理框架间的标准接口，实现一次压缩、多处部署。

算法-推理解耦架构提升了压缩权重的通用性，降低了维护成本。

部署验证：金箍棒量化权重已在vLLM-MindSpore Plugin和MindSpore Transformers上完成部署验证。基于Hugging Face格式的标准化设计，理论上支持在其他推理框架上部署应用。