【征程6】校准量化中HistogramObserver解析

声明：本文主要参考开源资料进行学习整理，如有错漏，欢迎评论交流~

1. HistogramObserver的定义与原理

HistogramObserver是horizon_plugin_pytorch中一种基于直方图统计的量化Observer。与MinMaxObserver仅记录最小最大值、MSEObserver在线性搜索中计算MSE不同，HistogramObserver通过构建完整的数值分布直方图来捕获张量的统计特性，并支持多种量化参数计算方法。

MinMaxObserver与MSEObserver解析可参考：【J6】MinMaxObserver与MSEObserver介绍 - 地平线开发者社区

1.1 基本原理

HistogramObserver的工作流程分为两个阶段：

阶段一：直方图收集（forward过程）

对于每个输入张量 x:
1. 计算当前batch的min_val_cur, max_val_cur
2. 初始化或扩展直方图边界:
   - 首次初始化：以max(|min|, |max|)为基准，扩展到最近的2的幂次
   - 后续扩展：当数值超出当前边界时，边界按2倍扩展
3. 将数值映射到直方图桶中，累加到全局直方图

直方图桶（Histogram Bin）：将数值范围划分为若干个等宽的区间，每个区间就是一个"桶"（bin）。落在某个区间内的数值个数就是该桶的计数。

例如，数据范围 -10, 10，设置 bin_num=4：

桶0 [-10, -5) 包含 -8, -6 → 计数 = 2

桶1 [-5, 0) 包含 -3 → 计数 = 1

桶2 [0, 5) 包含 2, 1, 4 → 计数 = 3

桶3 5, 10 包含 7, 9 → 计数 = 2

桶越多（bin_num越大），分布刻画越精细，内存占用越大。

阶段二：量化参数计算（calculate_qparams过程）

HistogramObserver支持多种计算方法，在calculate_qparams时根据method参数选择：

方法	描述
mse	最小化量化误差的积分估计（默认）
kl	最小化KL散度
js	最小化JS散度
cdf_kl	基于累积分布函数的KL散度
cdf_js	基于累积分布函数的JS散度
wasserstein	最小化Wasserstein距离
percentile	百分位截断

重点关注mse 和 percentile 即可。

1.2 核心实现

def forward(self, x_orig):
    x = x_orig.detach().to(self.min_val.dtype)

    # 计算当前batch的min/max
    min_val_cur, max_val_cur = torch.aminmax(x, dim=1)

    # 初始化直方图边界（首次）
    if self.histogram_edges.numel() == 0:
        max_abs = torch.maximum(-min_val_cur, max_val_cur)
        max_abs = 2 ** torch.ceil(torch.log2(max_abs))  # 对齐到2的幂次
        self.histogram_edges = ...

    # 扩展边界（当数值超出范围时）
    if need_update_hist.any():
        multiple = 2 ** torch.ceil(torch.log2(...))
        self.histogram_edges[idx] *= multiple

    # 累加当前batch的直方图
    bin_idx = vmap(self.bucketize_per_sample)(x, self.histogram_edges)
    self.histogram += hist

1.3 直方图扩展机制

当遇到超出当前边界的数值时，边界会按2的幂次扩展，同时通过分桶合并保持bin数量不变：

# 边界扩展倍数
multiple = 2 ** torch.ceil(torch.log2(ratio))

# 重构直方图：将原直方图的bins合并到新范围
new_hist[slice_start:slice_end] = padded_histogram.reshape(-1, multiple).sum(dim=-1)

这种设计确保了：直方图可以处理任意范围的数值、内存占用保持恒定、多卡校准结果一致性。

1.4 特点分析

优点：

• 保留完整分布信息，支持多种量化参数计算方法

• 计算与收集分离，可在校准后灵活切换方法

• 边界对齐机制确保多卡校准结果一致

缺点：

• 内存占用较高（bin_num默认2048）

2. 与MinMaxObserver/MSEObserver的区别

2.1 核心差异对比

特性	MinMaxObserver	MSEObserver	HistogramObserver
统计方式	仅记录min/max	在线搜索最优scale	构建完整直方图
内存占用	极低	中等	较高
计算时机	forward时更新	forward时搜索	forward收集，qparams时计算
分布感知	否	部分（通过MSE）	是（完整分布）
多方法支持	否	否（仅MSE）	是（7种方法）
Per-channel支持	是	是	是

2.2 具体区别分析

与MinMaxObserver的区别：

1. 信息保留：MinMaxObserver只保留极值信息，丢失了分布密度；HistogramObserver保留完整分布信息。
2. 离群值处理：MinMaxObserver对离群值极其敏感；HistogramObserver通过分布感知方法（如percentile、kl）自动处理离群值。
3. 校准稳定性：MinMaxObserver的滑动平均可能导致校准结果不稳定；HistogramObserver的直方图累加更加稳定。

与MSEObserver的区别：

1. 计算时机：MSEObserver在每次forward时进行完整的搜索；HistogramObserver将收集和计算分离。
2. 误差度量：MSEObserver仅支持L2误差；HistogramObserver支持MSE、KL、JS、Wasserstein等多种度量。
3. 灵活性：HistogramObserver支持在计算量化参数时更换方法，无需重新校准。

3. reset_dtype与reset_scale用法

3.1 reset_dtype方法

reset_dtype方法定义在ObserverBase基类中，用于在运行时更改量化数据类型：

def reset_dtype(self, dtype):
    dtype = get_horizon_quant_dtype(dtype)
    if dtype == self.dtype:
        return
    self.dtype = dtype
    self._set_quant_min_max(self.dtype)

使用场景：

1. 混合精度量化：不同层使用不同量化位宽时，动态调整Observer的dtype。
2. 量化位宽探索：尝试INT8、INT4等不同位宽，无需重新校准。

from horizon_plugin_pytorch.dtype import qint8, qint4

observer = HistogramObserver(dtype=qint8)
# 校准完成后，尝试INT16量化
observer.reset_dtype(qint16)
scale, zero_point = observer.calculate_qparams()

3.2 reset_scale类方法

reset_scale是HistogramObserver的类方法，用于基于已收集的直方图重新计算量化参数：

@classmethod
def reset_scale(
    cls,
    model: torch.nn.Module,
    method: str,
    method_kwargs: Dict = None,
    prefix: Tuple[str] = None,
    dtype=None,
):
    for n, m in model.named_modules():
        if isinstance(m, FakeQuantizeBase) and isinstance(
            getattr(m, "activation_post_process", None), cls
        ):
            m.set_qparams(
                m.activation_post_process.calculate_qparams(
                    m.get_dtype(), method, **method_kwargs
                )[0]
            )

参数说明：

参数	说明
model	待处理的模型
method	计算方法：mse/kl/js/cdf_kl/cdf_js/wasserstein/percentile
method_kwargs	方法参数（如steps、percentile）
prefix	仅处理指定前缀的模块
dtype	仅处理指定dtype的模块

使用场景：

方法对比实验： 使用同一份校准数据，对比不同量化参数计算方法的效果。

# 对比不同方法
methods = ['mse', 'kl', 'wasserstein', 'percentile']
for method in methods:
    HistogramObserver.reset_scale(model, method)
    evaluate(model)

混合策略量化： 不同层使用不同的校准方法。

# Attention层使用KL方法
HistogramObserver.reset_scale(model, method='kl', prefix='model.layers.0.self_attn')

# FFN层使用MSE方法
HistogramObserver.reset_scale(model, method='mse', prefix='model.layers.0.mlp')

3.3 reset_scale vs 重新校准

对比项	reset_scale	重新校准
数据需求	无需校准数据	需要校准数据
时间成本	极低	高
适用场景	方法对比、参数调优	分布变化、数据变更

4. 使用场景建议

from horizon_plugin_pytorch.quantization.observer_v2 import HistogramObserver
from horizon_plugin_pytorch.quantization.qconfig import QConfig
from horizon_plugin_pytorch.quantization.fake_quantize import FakeQuantize

# 高精度量化场景
qconfig = QConfig(
    weight=FakeQuantize.with_args(
        observer=HistogramObserver,
        dtype=qint8,
        bin_num=2048,
    ),
    output=FakeQuantize.with_args(
        observer=HistogramObserver,
        dtype=qint8,
        bin_num=2048,
    ),
)

# 校准后尝试不同方法
for method in ['mse', 'kl']:
    HistogramObserver.reset_scale(model, method)
    acc = evaluate(model)

方法选择建议：

数据分布特征	推荐方法
正态分布	mse
长尾分布	percentile 或 kl
多峰分布	wasserstein
存在离群值	percentile(0.9999)

HistogramObserver在实际应用中，适用于高精度量化需求、复杂分布数据、需要灵活调整量化策略的场景。

建议：校准时收集直方图，通过reset_scale对比不同方法后确定最优策略。