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

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的模块

使用场景:

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

# 对比不同方法
methods = ['mse', 'kl', 'wasserstein', 'percentile']
for method in methods:
    HistogramObserver.reset_scale(model, method)
    evaluate(model)
  1. 混合策略量化:不同层使用不同的校准方法。

# 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对比不同方法后确定最优策略。

Logo

加入社区

更多推荐