跳到内容

llmcompressor.modifiers.smoothquant

模块

SmoothQuantMapping dataclass

SmoothQuantMapping(
    smooth_name: str,
    smooth_layer: Module,
    balance_layers: list[Module],
)

用于存储激活层与在平滑过程中必须进行缩放的后续权重之间的映射关系的类

参数

  • smooth_name

    (str) –

    激活层的名称

  • smooth_layer

    (Module) –

    存储激活层的 PyTorch 模块

  • balance_layers

    (list[Module]) –

    smooth_layer 输入到的 PyTorch 模块列表,必须进行平衡以抵消 smooth_layer 的平滑

SmoothQuantModifier

继承自 Modifier

实现了来自 https://arxiv.org/abs/2211.10438 的 SmoothQuant 算法。此修饰符对激活中的离群值进行通道级平滑处理,通过减小动态范围使其更容易量化。平滑操作通过将逆操作应用于下一层权重来抵消,从而使权重更难量化。

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

示例配方

SmoothQuantModifier:
  smoothing_strength: 0.5
  mappings: [
    [["re:.*q_proj", "re:.*k_proj", "re:.*v_proj"], "re:.*self_attn_layer_norm"],
    [["re:.*fc1"], "re:.*final_layer_norm"]
  ]
  ignore: ["model.decoder.final_layer_norm"]

:param smoothing_strength: alpha,平滑强度(0-1 范围) :param mappings: 要平滑的激活层列表,以及用于缩放输出以实现平滑的层。映射列表中的每个条目本身都应该是一个列表,第一个条目是共享相同输入激活(需要平滑的那个)的层列表,第二个条目是用于实现平滑的输出缩放层。如果使用正则表达式,它将匹配模块名称重叠最大的层。如果未提供,则参数将从模型架构中推断。 :param ignore: 要忽略的层列表,即使它们匹配 mappings 中的正则表达式。它应该匹配用于实现平滑的输出缩放层(mappings 条目中的第二个条目)的名称。 :param num_calibration_steps: 用于校准的样本数量,如果为 None 则使用整个数据集

参数

  • calibration_function

    用于前向传播的可选函数,或 None 表示使用默认的 tensor_module_forward

方法

  • on_finalize

    通过清除尺度和映射数据来进行清理

  • on_initialize

    初始化并在给定状态上运行 SmoothQuant

on_finalize

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

通过清除尺度和映射数据来进行清理

源代码位于 llmcompressor/modifiers/smoothquant/base.py
def on_finalize(self, state: State, **kwargs) -> bool:
    """
    Clean up by clearing the scale and mapping data
    """
    if not self.ended_:
        self.on_end(state, None)

    if len(self.scales_) > 0:
        raise ValueError(f"Failed to compress {len(self.scales_)} modules")

    if self.scales_ is not None:
        self.scales_.clear()
    if self.resolved_mappings_ is not None:
        self.resolved_mappings_.clear()

    return True

on_initialize

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

初始化并在给定状态上运行 SmoothQuant

参数

  • state

    (State) –

    要运行 SmoothQuant 的状态

返回

  • bool

    成功运行为 True,否则为 False

源代码位于 llmcompressor/modifiers/smoothquant/base.py
def on_initialize(self, state: State, **kwargs) -> bool:
    """
    Initialize and run SmoothQuant on the given state

    :param state: state to run SmoothQuant on
    :return: True on a successful run, False otherwise
    """
    if self.end and self.end != -1:
        raise ValueError(
            f"{self.__class__.__name__} can only be applied during one-shot. "
            f" Expected end to be None or -1, got {self.end}"
        )
    if self.start and self.start != -1:
        raise ValueError(
            f"{self.__class__.__name__} can only be applied during one-shot. "
            f"Expected start to be None or -1, got {self.end}"
        )

    if not hasattr(state, "data") or state.data.calib is None:
        raise ValueError(
            f"{self.__class__.__name__} requires a calibration dataset to be "
            "provided"
        )
    self.ignore = [] if not self.ignore else self.ignore
    self.mappings = self._infer_mappings_from_model(state.model)
    self.resolved_mappings_ = self._resolve_mappings(state.model)
    self.scales_ = {}

    return True

SmoothQuantScale dataclass

SmoothQuantScale(
    min_channel_vals: Tensor, max_channel_vals: Tensor
)

用于存储层通道的最小值和最大值的类。在校准期间的每次前向传播中都会更新此值

参数

  • min_channel_vals

    (Tensor) –

    截至目前看到的每通道的最小输出值

  • max_channel_vals

    (Tensor) –

    截至目前看到的每通道的最大输出值