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工具安装指南》。
原理和实现¶
原理¶
算法核心:
- 使用激活值的均值(mean)计算激活缩放因子:
Act_Mean_Abs = mean(abs(act))。 - 使用实际量化器(LinearQuantizer)评估不同alpha参数下的量化误差。
- 自动搜索或使用配置的alpha参数,beta固定为0。
算法公式:
scales = (Act_Mean_Abs**alpha / Weight_Max_Abs**beta).clamp(min=1e-5)
其中:
Act_Mean_Abs:激活值的均值绝对值(使用均值计算:mean(abs(act)))。Weight_Max_Abs:权重的最大绝对值(取每列的最大值)。alpha:激活缩放的系数,控制激活对缩放因子的影响程度(0-1之间),可通过自动搜索或手动配置。beta:权重缩放的系数,固定为0。
关键特性:
- 实际量化器评估:与Flex Smooth Quant不同,Flex AWQ SSZ使用真实的量化器(LinearQuantizer)来评估不同alpha参数下的量化误差,而不是简单的模拟量化。
- 激活尺度计算:使用激活值的均值(mean)而非最大值(max)来计算激活尺度,更适合低bit量化场景。
- Beta固定为0:算法固定beta为0,简化参数搜索空间,专注于alpha参数的优化。
- 自动参数搜索:如果未提供alpha参数,算法会在[0.0, 1.0]范围内以0.05为步长搜索最优alpha,选择使量化误差(MSE)最小的参数。
Alpha参数搜索流程:
- 初始化:创建FlexAWQSSZAlphaBetaSearcher,使用配置的qconfig。
- 网格搜索:在[0.0, 1.0]范围内以0.05为步长遍历alpha值。
- 量化误差评估:对每个alpha值:
- 计算缩放因子:
scale = max(abs(Act_Mean_Abs)) ** alpha - 应用缩放:
scaled_act = act / scale,scaled_weight = weight * scale - 创建实际量化器(LinearQuantizer)并部署
- 计算量化结果与浮点结果的归一化MSE误差
- 选择最优参数:选择使MSE误差最小的alpha值。
支持的子图类型¶
NormLinearSubgraph(归一化-线性子图)¶
适用于包含归一化层和多个线性层的结构,如:
x = norm(x)
y = torch.cat([linear(x) for linear in linears], dim=-1)
处理方式:
- 计算所有线性层权重的列最大值作为权重缩放因子。
- 使用激活值的均值计算激活缩放因子。
- 对每个线性层应用正向缩放。
- 对归一化层应用反向缩放(1/scales)。
LinearLinearSubgraph(线性-线性子图)¶
适用于两个连续线性层的结构:
y = linear2(linear1(x))
处理方式:
- 基于linear2的权重计算缩放因子。
- 使用linear1输出的激活值均值计算激活缩放因子。
- 对linear2应用正向缩放。
- 对linear1应用反向缩放(1/scales)。
OVSubgraph(注意力输出-值子图)¶
适用于注意力机制中的输出投影和值投影:
- 支持MHA(多头注意力)
- 支持MQA(多查询注意力)
- 支持GQA(分组查询注意力)
处理方式:
- 基于o_proj权重计算缩放因子。
- 使用v_proj输出的激活值均值计算激活缩放因子。
- 对o_proj应用正向缩放。
- 对v_proj应用反向缩放(1/scales)。
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 中实现,处理流程分两阶段。
预处理阶段¶
子图发现与构建:
- 通过
SubgraphProcessor获取全局子图信息,识别四种类型的子图:norm-linear、linear-linear、ov、up-down。 - 根据配置的
include/exclude模式过滤子图。
统计信息收集:
- 为所有子图中的线性模块安装前向钩子(forward hook)。
- 钩子在
[batch, seq, hidden_dim]维度上收集激活值统计信息: - 激活张量数据:收集完整的激活张量,用于后续平滑计算。
- 使用第一个linear的激活统计信息:Flex AWQ SSZ使用子图targets中第一个线性层的激活统计信息。
后处理阶段¶
按优先级处理子图:
- 按默认配置的优先级顺序处理:
up-down(最高)→ov(高)→norm-linear(中)→linear-linear(低)。 - 每种子图类型调用相应的平滑处理方法。
子图平滑处理:
- Norm-Linear子图:对归一化层和后续线性层应用平滑,仅使用前两层线性层进行alpha搜索(如果线性层数量大于3)。
- Linear-Linear子图:对两个线性层应用平滑,调整权重和偏置。
- OV子图:处理注意力机制中的输出投影(Output projection)和值投影(Value projection)之间的连接关系,支持QKV融合模式。
- Up-Down子图:处理MLP门控机制,对上下投影层应用平滑。
资源清理:
- 清理所有安装的统计钩子。
- 释放统计信息内存。
- 恢复模型原始状态。
适用要求¶
- 模型架构要求:模型必须支持
FlexSmoothQuantInterface接口,并正确配置子图映射关系。 - 模块命名要求:模块名称必须与
named_modules()返回的完整路径完全一致。 - 子图类型支持:目前支持四种标准子图类型:
norm-linear、linear-linear、ov、up-down。 - 模块属性要求:目标模块必须存在且具备可写的
weight,其他自定义模块暂不支持。 - 模型结构假设:算法基于标准的Transformer架构设计,对于非标准结构需要谨慎评估适用性。
- 量化配置要求:必须提供qconfig配置,包括激活和权重的量化方式,通常权重使用SSZ方法。
功能介绍¶
YAML配置示例¶
作为Processor使用,YAML配置示例如下:
spec:
process:
- 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配置字段详解¶
整体字段说明:
| 字段名 | 作用 | 说明 |
|---|---|---|
type |
处理器类型标识 | 固定为 "flex_awq_ssz",用于标识这是一个灵活激活感知权重量化平滑处理器。 |
alpha |
激活缩放权重系数 | 取值范围为 [0, 1] 的浮点数,控制激活对缩放因子的影响程度,默认值为 None(自动搜索)。 |
qconfig |
量化配置 | 必填参数,包含激活值(act)和权重(weight)的量化配置,用于实际量化器评估。 |
qconfig.act |
激活值量化配置 | 包含 scope、dtype、symmetric、method 等字段,指定激活值的量化方式。 |
qconfig.weight |
权重量化配置 | 包含 scope、dtype、symmetric、method等字段,指定权重的量化方式,通常结合 SSZ 方法使用。 |
enable_subgraph_type |
开启的子图类型 | 支持的子图类型列表,包括 "norm-linear"、"linear-linear"、"ov"、"up-down"。 |
include |
包含的层 | 通过通配符匹配需要参与处理的模块名。 |
exclude |
排除的层 | 通过通配符匹配需要从处理流程中排除的模块名。 |
qconfig.act(激活值量化配置)¶
作用:配置激活值的量化参数。
| 参数名 | 作用 | 可选值 | 说明 | 默认值 |
|---|---|---|---|---|
scope |
量化范围 | "per_tensor","per_token","pd_mix" |
per_tensor:整个张量使用相同参数per_token:每个 token 独立参数(动态量化)pd_mix:prefilling 时 per_token,decoding 时 per_tensor |
"per_tensor" |
dtype |
量化数据类型 | "int8","int4" |
8 位 / 4 位整数量化 | "int8" |
symmetric |
是否对称量化 | true,false |
true:对称量化,零点为 0false:非对称量化,零点可调整 |
false |
method |
量化方法 | "minmax","histogram" |
minmax:最小最大值量化histogram:直方图量化 |
"minmax" |
qconfig.weight(权重量化配置)¶
作用:配置权重的量化参数。
| 参数名 | 作用 | 可选值 | 说明 | 默认值 |
|---|---|---|---|---|
scope |
量化范围 | "per_tensor","per_channel" |
per_tensor:整个张量使用相同参数per_channel:每个通道独立参数 |
"per_channel" |
dtype |
量化数据类型 | "int8","int4" |
8 位 / 4 位整数量化 | "int8" |
symmetric |
是否对称量化 | true,false |
true:对称量化,零点为 0false:非对称量化,零点可调整 |
true |
method |
量化方法 | "minmax","ssz","gptq" |
minmax:最小最大值量化ssz:SSZ 权重量化gptq:GPTQ 权重量化 |
"minmax" |
模型适配¶
接口与数据结构¶
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。
步骤:
- 继承接口:模型适配器继承
FlexSmoothQuantInterface接口,实现get_adapter_config_for_subgraph()方法。 - 配置子图映射:为每层配置四种类型的子图映射关系:
- Norm-Linear子图:归一化层到后续线性层的映射
- OV子图:注意力机制中V投影到O投影的映射
- Up-Down子图:MLP门控机制中上投影到下投影的映射
- Linear-Linear子图:连续线性层的映射
- 指定模块路径:使用完整的模块路径,如
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¶
模块名不匹配¶
现象: include/exclude 未命中时,日志提示未匹配模式。
解决方案: 核对完整模块名是否与 named_modules() 返回的路径一致。
子图配置错误¶
现象: get_adapter_config_for_subgraph() 返回的配置不正确。
解决方案: 检查配置中的 source 和 targets 字段是否正确。
模块不存在¶
现象: 配置中指定的模块名称在模型中不存在。
解决方案: 通过 model.named_modules() 验证模块是否确实存在。
子图类型不支持¶
现象: 配置的子图类型不被支持。
解决方案: 确保配置的子图类型在支持列表中,当前支持 norm-linear、linear-linear、ov、up-down四种子图类型。
qconfig配置缺失¶
现象: 报错提示qconfig为必填参数。
解决方案: 在YAML配置中添加qconfig字段,包含act和weight的量化配置。
映射关系错误¶
现象: MappingConfig 中的 source 和 targets 指向错误的模块。
解决方案: 检查 MappingConfig 中的 source 和 targets 是否指向正确的模块。