跳到内容

llmcompressor.modifiers.awq.base

  • AWQModifier

    实现了 AWQ (Activation-Weighted Quantization) 算法,

AWQModifier

基类: Modifier, QuantizationMixin

实现了 AWQ (Activation-Weighted Quantization) 算法,如 https://arxiv.org/pdf/2306.00978 所述。该算法通过仅保护 1% 最显著的权重通道来显著减少量化误差。

AWQ 不依赖原始权重值,而是通过分析激活模式来识别重要通道,专注于权重张量中最能响应输入的通道。为了减少量化误差,它以一种保留模型原始行为的方式对这些通道进行缩放,使用从激活统计数据中离线计算的缩放因子。

由于此修改器会操作模型的权重,因此它只能用于一次性运行,而不能在训练期间使用。激活范围通过对少量校准数据运行模型来确定。

示例配方

AWQModifier:
  mappings:
    - smooth_layer: "re:.*self_attn_layer_norm"
      balance_layers: ["re:.*q_proj", "re:.*k_proj", "re:.*v_proj"]
    - smooth_layer: "re:.*final_layer_norm"
      balance_layers: ["re:.*fc1"]
  ignore: ["lm_head"]
  config_groups:
    group_0:
      targets:
        - "Linear"
      input_activations: null
      output_activations: null
      weights:
        num_bits: 4
        type: int
        symmetric: false
        strategy: group
        group_size: 128

生命周期

  • on_initialize
    • 解析映射
    • 将前向传播所需的 kwargs 捕获到模块中
  • on_start
    • 设置激活缓存钩子以捕获输入激活以平衡层
  • 在序列 epoch 结束时
    • 对每个平滑层应用平滑
      • 消耗所有批次中的缓存激活
        • 在使用时清除缓存激活
      • 为每个平滑层查找最佳平滑尺度
      • 应用于模型权重
      • 如果存在任何未使用的激活,则引发错误
  • on_end
    • 重新运行顺序 epoch 结束的逻辑(以防是基本管道)
    • 设置尺度和零点
    • 移除激活钩子
  • on_finalize
    • 清除已解析的映射和捕获的激活

参数

  • sequential_targets

    在同一校准过程中要压缩的模块名称列表

  • mappings

    列出要平滑的激活层,以及要缩放输出以实现激活平滑的层。映射列表的每个条目本身都应该是一个列表,其中第一个条目是共享相同输入激活(要平滑的激活)的层列表,第二个条目是输出被缩放以实现平滑的层。如果使用正则表达式,它将匹配模块名称中重叠最大的层。

  • ignore

    要忽略的层列表,即使它们匹配映射中的正则表达式。它应该匹配其输出被缩放以实现平滑的层的名称(映射列表的第二个条目)。

  • offload_device

    将缓存的 args 卸载到此设备,这会减少内存需求,但需要更多时间在 CPU 和执行设备之间移动数据。默认为 None,因此不会卸载缓存的 args。如果遇到 OOM 错误,请考虑将其设置为 torch.device("cpu")

  • duo_scaling

    是否使用 duo scaling,它同时使用输入激活和权重来确定缩放因子。默认为 True。如果为 True,则同时使用激活和权重。如果为 False,则仅使用激活。如果为 "both",则一半的网格搜索使用 duo_scaling=False 进行,另一半使用 duo_scaling=True 进行。

  • n_grid

    在为每个映射执行最佳尺度网格搜索时,此参数指定应使用的网格点数。为了减少运行时间,可能以牺牲稍差的尺度为代价,可以减小此值。默认为 20

方法

  • on_end

    通过设置尺度和零点来完成校准,

  • on_finalize

    通过清除激活和映射数据来清理

  • on_initialize

    在给定状态下初始化 AWQ

  • validate_duo_scaling

    验证 duo_scaling 是否为 True、False 或 'both'(小写)

on_end

on_end(state: State, event: Event, **kwargs)

通过设置尺度和零点、移除观察者和校准钩子来完成校准

源代码在 llmcompressor/modifiers/awq/base.py
def on_end(self, state: State, event: Event, **kwargs):
    """
    Finish calibrating by setting scales and zero-points,
     removing observers and calibration hooks
    """
    self._assert_all_activations_consumed()

    self.ended_ = True

    for _, module in tqdm(
        match_named_modules(state.model, self.resolved_targets, self.ignore),
        desc="Calibrating weights",
    ):
        update_weight_zp_scale(module)

    QuantizationMixin.end_calibration(self, state.model)

    # remove activation hooks
    self.remove_hooks()

on_finalize

on_finalize(state: State, **kwargs) -> bool

通过清除激活和映射数据来清理

参数

  • state

    (State) –

    未使用

返回

  • bool

    True

源代码在 llmcompressor/modifiers/awq/base.py
def on_finalize(self, state: State, **kwargs) -> bool:
    """
    Clean up by clearing the activations and mapping data

    :param state: unused
    :return: True
    """
    if not self.ended_:
        self.on_end(state, None)

    self._parent_args_cache.clear()
    self._smooth_activation_means.clear()
    self._resolved_mappings.clear()

    return True

on_initialize

on_initialize(state: State, **kwargs) -> bool

在给定状态下初始化 AWQ 初始化量化、解析映射、缓存模块 kwargs

参数

  • state

    (State) –

    用于运行 AWQ 的状态

返回

  • bool

    成功运行为 True,否则为 False

源代码在 llmcompressor/modifiers/awq/base.py
def on_initialize(self, state: State, **kwargs) -> bool:
    """
    Initialize AWQ on the given state
    Initialize quantization, resolve mappings, cache module kwargs

    :param state: state to run AWQ on
    :return: True on a successful run, False otherwise
    """

    # apply config to model and prepare calibration hooks
    if QuantizationMixin.has_config(self):
        QuantizationMixin.initialize_quantization(self, state.model)

    # Validate that duo_scaling is only used with per-channel quantization
    if self.duo_scaling is not False:
        for _, module in match_named_modules(
            state.model, self.resolved_targets, self.ignore
        ):
            if (
                hasattr(module, "quantization_scheme")
                and hasattr(module.quantization_scheme, "weights")
                and module.quantization_scheme.weights.strategy
                == QuantizationStrategy.TENSOR
            ):
                raise ValueError(
                    "duo_scaling is only supported with per-channel quantization "
                    "strategies (group or channel), but found TENSOR strategy. "
                    "Please set duo_scaling=False or use a per-channel "
                    "quantization strategy."
                )

    if self.mappings is None:
        logger.info("No AWQModifier.mappings provided, inferring from model...")
        self.mappings = get_layer_mappings_from_architecture(
            architecture=state.model.__class__.__name__
        )

    self._set_resolved_mappings(state.model)

    return True

validate_duo_scaling classmethod

validate_duo_scaling(v)

验证 duo_scaling 是否为 True、False 或 'both'(小写)

源代码在 llmcompressor/modifiers/awq/base.py
@field_validator("duo_scaling")
@classmethod
def validate_duo_scaling(cls, v):
    """Validate that duo_scaling is either True, False, or 'both' (lowercase)"""
    if v not in (True, False, "both"):
        raise ValueError(f"duo_scaling must be True, False, or 'both', got {v!r}")
    return v

get_lowest_common_ancestor_with_avoid

get_lowest_common_ancestor_with_avoid(
    balance_names: Iterator[str],
    model: Module,
    avoid=torch.nn.ModuleList,
)

获取最低的非避免类/类型的祖先。有关大小写处理的详细信息,请参阅 compressed_tensors.utils.get_lowest_common_ancestor_name。

注意:主要用于排除 ModuleList 类型的父级,因为它们与钩子不兼容,因为它们的 forward 方法不会直接为 MoE 模型调用。例如,请参阅 Qwen3MoeSparseMoeBlock,专家根据路由器的输出进行选择,并调用其 forward 方法。 https://github.com/huggingface/transformers/blob/v4.52.4/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py#L233

源代码在 llmcompressor/modifiers/awq/base.py
def get_lowest_common_ancestor_with_avoid(
    balance_names: Iterator[str], model: Module, avoid=torch.nn.ModuleList
):
    """
    Get the lowest ancestor that is not the avoided class/type.
    see compressed_tensors.utils.get_lowest_common_ancestor_name
    for detail on case handling.

    NOTE: primarily used to exclude parents of type ModuleList, which don't play
    nicely with hooks because their forward method is never directly
    called for MoE models. See Qwen3MoeSparseMoeBlock for example, experts
    are selected based on router output and their forward method is called.
    https://github.com/huggingface/transformers/blob/v4.52.4/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py#L233
    """
    ancestor_name = get_lowest_common_ancestor_name(balance_names)

    while True:
        if ancestor_name == "":
            return "", model
        ancestor = model.get_submodule(ancestor_name)
        if not isinstance(ancestor, avoid):
            return ancestor_name, ancestor
        ancestor_name = ".".join(ancestor_name.split(".")[:-1])