Flex AWQ SSZ：灵活激活感知权重量化平滑算法说明¶

简介¶

概述：Flex AWQ SSZ（灵活激活感知权重量化平滑算法）是一种用于大语言模型量化过程中抑制激活离群值的算法。该算法结合了AWQ（Activation-aware Weight Quantization）和SSZ（Smooth Scale Zero）的思想，通过使用实际量化器评估参数有效性，自动搜索最优的alpha参数，在保持模型精度的同时，有效减少量化误差。与传统的平滑算法不同，Flex AWQ SSZ使用真实的量化器进行参数评估，能够更准确地反映量化后的实际效果。
核心思想：Flex AWQ SSZ算法的核心思想是通过实际量化器（LinearQuantizer）来评估不同alpha参数下的量化误差，自动搜索使量化误差最小的最优alpha参数。算法固定beta为0，激活值尺度计算使用均值（mean）而非最大值（max），从而在不同量化场景下获得精度与量化效率的平衡。

使用前准备¶

安装 msModelSlim 工具，详情请参见《msModelSlim工具安装指南》。

原理和实现¶

原理¶

Flex AWQ SSZ算法使用以下公式计算平滑缩放因子：

scales = (A_scale**alpha / W_scale**beta).clamp(min=1e-5)

其中： - A_scale：激活值的缩放因子（使用均值计算：mean(abs(act))）。 - W_scale：权重的缩放因子（取每列的最大值）。 - alpha：激活缩放的系数，控制激活对缩放因子的影响程度（0-1之间），可通过自动搜索或手动配置。 - beta：权重缩放的系数，固定为0。

关键特性： 1. 实际量化器评估：与Flex Smooth Quant不同，Flex AWQ SSZ使用真实的量化器（LinearQuantizer）来评估不同alpha参数下的量化误差，而不是简单的模拟量化。 2. 激活尺度计算：使用激活值的均值（mean）而非最大值（max）来计算激活尺度，更适合低bit量化场景。 3. Beta固定为0：算法固定beta为0，简化参数搜索空间，专注于alpha参数的优化。 4. 自动参数搜索：如果未提供alpha参数，算法会在[0.0, 1.0]范围内以0.05为步长搜索最优alpha，选择使量化误差（MSE）最小的参数。

Alpha参数搜索流程

初始化：创建FlexAWQSSZAlphaBetaSearcher，使用配置的qconfig。
网格搜索：在[0.0, 1.0]范围内以0.05为步长遍历alpha值。
量化误差评估：对每个alpha值：
计算缩放因子：scale = max(abs(act)) ** alpha
应用缩放：scaled_act = act / scale，scaled_weight = weight * scale
创建实际量化器（LinearQuantizer）并部署
计算量化结果与浮点结果的归一化MSE误差
选择最优参数：选择使MSE误差最小的alpha值。

支持的子图类型¶

Flex AWQ SSZ算法支持与Flex Smooth Quant相同的四种标准子图类型：

1. NormLinearSubgraph（归一化-线性子图）¶

适用于包含归一化层和多个线性层的结构，如：

x = norm(x)
y = torch.cat([linear(x) for linear in linears], dim=-1)

处理方式： - 计算所有线性层权重的列最大值作为权重缩放因子。 - 使用激活值的均值计算激活缩放因子。 - 对每个线性层应用正向缩放。 - 对归一化层应用反向缩放（1/scales）。

2. LinearLinearSubgraph（线性-线性子图）¶

适用于两个连续线性层的结构：

y = linear2(linear1(x))

处理方式： - 基于linear2的权重计算缩放因子。 - 使用linear1输出的激活值均值计算激活缩放因子。 - 对linear2应用正向缩放。 - 对linear1应用反向缩放（1/scales）。

3. OVSubgraph（注意力输出-值子图）¶

适用于注意力机制中的输出投影和值投影： - 支持MHA（多头注意力） - 支持MQA（多查询注意力） - 支持GQA（分组查询注意力）

处理方式： - 基于o_proj权重计算缩放因子。 - 使用v_proj输出的激活值均值计算激活缩放因子。 - 对o_proj应用正向缩放。 - 对v_proj应用反向缩放（1/scales）。

4. UpDownSubgraph（上投影-下投影子图）¶

适用于MLP门控机制：

y = down_proj(ReLU(gate_proj(x)) * up_proj(x))

处理方式： - 基于down_proj权重计算缩放因子。 - 使用up_proj输出的激活值均值计算激活缩放因子。 - 对down_proj应用正向缩放。 - 对up_proj应用反向缩放（1/scales）。

实现¶

算法在 msmodelslim/processor/anti_outlier/flex_smooth/processor.py 中实现，处理流程分两阶段：

1) 预处理阶段（preprocess）¶

子图发现与构建： - 通过 SubgraphProcessor 获取全局子图信息，识别四种类型的子图：norm-linear、linear-linear、ov、up-down。 - 根据配置的 include/exclude 模式过滤子图。

统计信息收集： - 为所有子图中的线性模块安装前向钩子（forward hook）。 - 钩子在 [batch, seq, hidden_dim] 维度上收集激活值统计信息： - 激活张量数据：收集完整的激活张量，用于后续平滑计算。 - 使用第一个linear的激活统计信息：Flex AWQ SSZ使用子图targets中第一个线性层的激活统计信息。

2) 后处理阶段（postprocess）¶

按优先级处理子图： - 按默认配置的优先级顺序处理：up-down（最高）→ ov（高）→ norm-linear（中）→ linear-linear（低）。 - 每种子图类型调用相应的平滑处理方法。

子图平滑处理： - Norm-Linear子图：对归一化层和后续线性层应用平滑，仅使用前两层线性层进行alpha搜索（如果线性层数量大于3）。 - Linear-Linear子图：对两个线性层应用平滑，调整权重和偏置。 - OV子图：处理注意力机制中的输出投影（Output projection）和值投影（Value projection）之间的连接关系，支持QKV融合模式。 - Up-Down子图：处理MLP门控机制，对上下投影层应用平滑。

Flex AWQ SSZ算法核心： - 使用激活值的均值（mean）计算激活缩放因子：act_scales = mean(abs(act))。 - 使用实际量化器（LinearQuantizer）评估不同alpha参数下的量化误差。 - 自动搜索或使用配置的alpha参数，beta固定为0。

资源清理： - 清理所有安装的统计钩子。 - 释放统计信息内存。 - 恢复模型原始状态。

适用要求¶

模型架构要求：模型必须支持 FlexSmoothQuantInterface 接口，并正确配置子图映射关系。
模块命名要求：模块名称必须与 named_modules() 返回的完整路径完全一致。
子图类型支持：目前支持四种标准子图类型：norm-linear、linear-linear、ov、up-down。
模块属性要求：目标模块必须存在且具备可写的 weight，其他自定义模块暂不支持。
模型结构假设：算法基于标准的Transformer架构设计，对于非标准结构需要谨慎评估适用性。
量化配置要求：必须提供qconfig配置，包括激活和权重的量化方式，通常权重使用SSZ方法。

功能介绍¶

使用说明¶

作为 Processor 使用

- type: "flex_awq_ssz"              # 固定为 `flex_awq_ssz`，用于指定 Processor。
  alpha: 0.8                       # 激活缩放的系数，取值范围为0-1之间，默认值为None（自动搜索），也支持用户自行配置。
  qconfig:                         # 量化配置，为必填参数。
    act:                           # 激活值量化配置。
      scope: "per_token"           # 量化范围：per_token 或 per_tensor。
      dtype: "int8"                # 量化数据类型：int8。
      symmetric: True              # 是否对称量化：True 或 False。
      method: "minmax"            # 量化方法：minmax 或其他方法。
    weight:                        # 权重量化配置。
      scope: "per_channel"         # 量化范围：per_channel。
      dtype: "int4"                # 量化数据类型：int4 或 int8。
      symmetric: True              # 是否对称量化：True。
      method: "ssz"                # 量化方法：ssz（Smooth Scale Zero）。
      ext:                         # 扩展配置（可选）。
        step: 10                   # SSZ方法的步长参数。
  enable_subgraph_type:            # 字符串列表，指定启用的子图类型，默认启用所有四种类型。
    - 'norm-linear'
    - 'linear-linear'
    - 'ov'
    - 'up-down'
  include:                         # 包含的层，支持通配符。
    - "*"
  exclude:                         # 排除的层，支持通配符。
    - "*self_attn*"

YAML配置示例¶

spec:
  process:
    - type: "flex_awq_ssz"
      alpha: 0.8                          # 激活缩放的系数，取值范围为0-1之间，默认值为None（自动搜索），也支持用户自行配置。
      qconfig:                             # 量化配置，为必填参数。
        act:
          scope: "per_token"
          dtype: "int8"
          symmetric: True
          method: "minmax"
        weight:
          scope: "per_channel"
          dtype: "int4"
          symmetric: True
          method: "ssz"
          ext:
            step: 10
      enable_subgraph_type:                # 开启的子图类型。
        - 'norm-linear'
        - 'linear-linear'
        - 'ov'
        - 'up-down'
      include: ["*"]                      # 包含的层，支持通配符。
      exclude: ["*self_attn*"]            # 排除的层，支持通配符。

YAML配置字段详解¶

字段名	作用	说明
type	处理器类型标识	固定值"flex_awq_ssz"，用于标识这是一个灵活激活感知权重量化平滑处理器。
alpha	激活缩放权重系数	0-1之间的浮点数，控制激活对缩放因子的影响程度，默认None（自动搜索）。
qconfig	量化配置	必填参数，包含激活值（act）和权重（weight）的量化配置，用于实际量化器评估。
qconfig.act	激活值量化配置	包含scope、dtype、symmetric、method等字段，指定激活值的量化方式。
qconfig.weight	权重量化配置	包含scope、dtype、symmetric、method、ext等字段，指定权重的量化方式，通常使用SSZ方法。
enable_subgraph_type	开启的子图类型	支持的子图类型列表，包括"norm-linear"、"linear-linear"、"ov"、"up-down" 。
include	包含的层	支持通配符匹配。
exclude	排除的层	支持通配符匹配。

模型适配¶

接口与数据结构¶

Flex AWQ SSZ使用与Flex Smooth Quant相同的接口：

from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any
from abc import ABC, abstractmethod

@dataclass
class MappingConfig:
    """模块映射关系配置"""
    source: str  # 源模块名称，如 "model.layers.0.input_layernorm"
    targets: List[str]  # 目标模块名称列表，如 ["model.layers.0.self_attn.q_proj", "model.layers.0.self_attn.k_proj"]

@dataclass
class FusionConfig:
    """融合配置，支持QKV融合等高级功能"""
    fusion_type: str = "none"  # 融合类型：none, qkv, custom等
    num_attention_heads: Optional[int] = None  # 注意力头数量
    num_key_value_heads: Optional[int] = None  # 键值头数量
    custom_config: Optional[Dict[str, Any]] = None  # 自定义配置

@dataclass
class AdapterConfig:
    """子图适配器配置"""
    subgraph_type: str  # 子图类型：norm-linear, linear-linear, ov, up-down
    mapping: Optional[MappingConfig] = None  # 模块映射关系
    fusion: FusionConfig = field(default_factory=lambda: FusionConfig())  # 融合配置

# 模型适配Flex AWQ SSZ算法接口（与Flex Smooth Quant相同）
class FlexSmoothQuantInterface(ABC):
    @abstractmethod
    def get_adapter_config_for_subgraph(self) -> List[AdapterConfig]:
        """
        返回模型中所有可进行Flex AWQ SSZ处理的子图配置

        Returns:
            List[AdapterConfig]: 子图配置列表，每个配置包含：
                - subgraph_type: 子图类型
                - mapping: 源模块到目标模块的映射关系
                - fusion: 融合配置（如QKV融合）
        """
        pass

适配步骤¶

前置要求： - 模型需要继承 FlexSmoothQuantInterface 接口（与Flex Smooth Quant相同）。 - 模块名称必须与 named_modules() 返回的完整路径一致。 - 支持的子图类型：norm-linear、linear-linear、ov、up-down。 - 配置中的subgraph_type、mapping 是必要参数。 - 当配置FusionConfig且fusion_type为qkv时，必须给出num_attention_heads和num_key_value_heads。

步骤： 1. 继承接口：模型适配器继承 FlexSmoothQuantInterface 接口，实现 get_adapter_config_for_subgraph() 方法。 2. 配置子图映射：为每层配置四种类型的子图映射关系： - Norm-Linear子图：归一化层到后续线性层的映射 - OV子图：注意力机制中V投影到O投影的映射 - Up-Down子图：MLP门控机制中上投影到下投影的映射 - Linear-Linear子图：连续线性层的映射 3. 指定模块路径：使用完整的模块路径，如 model.layers.{i}.self_attn.q_proj。

参考实现： 可参考 msmodelslim/model/qwen3/model_adapter.py 中的 Qwen3ModelAdapter 实现。

配置示例¶

以下是一个典型的Transformer层配置示例（与Flex Smooth Quant相同）：

def get_adapter_config_for_subgraph(self) -> List[AdapterConfig]:
    adapter_config = []
    for layer_idx in range(self.config.num_hidden_layers):
        # 1. 输入层归一化到QKV投影的Norm-Linear映射
        norm_linear_config1 = AdapterConfig(
            subgraph_type="norm-linear",
            mapping=MappingConfig(
                source=f"model.layers.{layer_idx}.input_layernorm",
                targets=[
                    f"model.layers.{layer_idx}.self_attn.q_proj",
                    f"model.layers.{layer_idx}.self_attn.k_proj",
                    f"model.layers.{layer_idx}.self_attn.v_proj"
                ]
            )
        )

        # 2. 后注意力层归一化到MLP投影的Norm-Linear映射
        norm_linear_config2 = AdapterConfig(
            subgraph_type="norm-linear",
            mapping=MappingConfig(
                source=f"model.layers.{layer_idx}.post_attention_layernorm",
                targets=[
                    f"model.layers.{layer_idx}.mlp.gate_proj",
                    f"model.layers.{layer_idx}.mlp.up_proj"
                ]
            )
        )

        # 3. 注意力机制中的OV映射
        ov_config = AdapterConfig(
            subgraph_type="ov",
            mapping=MappingConfig(
                source=f"model.layers.{layer_idx}.self_attn.v_proj",
                targets=[f"model.layers.{layer_idx}.self_attn.o_proj"]
            )
        )

        # 4. MLP门控机制的Up-Down映射
        up_down_config = AdapterConfig(
            subgraph_type="up-down",
            mapping=MappingConfig(
                source=f"model.layers.{layer_idx}.mlp.up_proj",
                targets=[f"model.layers.{layer_idx}.mlp.down_proj"]
            )
        )

        adapter_config.extend([norm_linear_config1, norm_linear_config2, ov_config, up_down_config])

    return adapter_config

FAQ¶

1. 模块名不匹配¶

现象: include/exclude 未命中时，日志提示未匹配模式。
解决方案: 核对完整模块名是否与 named_modules() 返回的路径一致。

2. 子图配置错误¶

现象: get_adapter_config_for_subgraph() 返回的配置不正确。
解决方案: 检查配置中的 source 和 targets 字段是否正确。

3. 模块不存在¶

现象: 配置中指定的模块名称在模型中不存在。
解决方案: 通过 model.named_modules() 验证模块是否确实存在。

4. 子图类型不支持¶

现象: 配置的子图类型不被支持。
解决方案: 确保配置的子图类型在 ENABLE_SUBGRAPH_TYPES 列表中。

5. qconfig配置缺失¶

现象: 报错提示qconfig为必填参数。
解决方案: 在YAML配置中添加qconfig字段，包含act和weight的量化配置。

6. 映射关系错误¶

现象: MappingConfig 中的 source 和 targets 指向错误的模块。
解决方案: 检查 MappingConfig 中的 source 和 targets 是否指向正确的模块。