# 架构设计

## 金箍棒架构

我们分析了一些[业界模型压缩框架实践](https://gitee.com/mindspore/golden-stick/wikis/%E6%9E%B6%E6%9E%84%E8%AE%BE%E8%AE%A1/%E9%87%91%E7%AE%8D%E6%A3%92%E8%AE%BE%E8%AE%A1%E6%96%87%E6%A1%A3#%E4%B8%9A%E7%95%8C%E5%AE%9E%E8%B7%B5%E8%B0%83%E7%A0%94)，总结了一些优秀的特点：

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

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

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

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

- **易用且灵活的API**：多层级的API，兼顾易用性和灵活性，降低使用门槛，同时保留算法定制化的能力；
- **丰富且模块化的算法库**：提供丰富的压缩算法，并且支持灵活组合；
- **高度可扩展的框架架构**：便于集成自定义算法模块，配合灵活的API构建定制化压缩流水线。

![金箍棒架构图](images/arch.png)

### 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](https://gitee.com/mindspore/vllm-mindspore)和[MindSpore Transformers](https://gitee.com/mindspore/mindformers)上完成部署验证。基于Hugging Face格式的标准化设计，理论上支持在其他推理框架上部署应用。