llmcompressor.modifiers.smoothquant

Dataclass for storing the mapping between an activation layer and the following weights that must be balanced during smoothing

Parameters:

• smooth_name

  (str) –

  name of the activation layer

• smooth_layer

  (Module) –

  PyTorch module storing the activation layer

• balance_layers

  (List[Module]) –

  list of PyTorch modules that smooth_layer feeds into, must be balanced to offset the smoothing of smooth_layer

Bases: Modifier

Implements the SmoothQuant algorithm from https://arxiv.org/abs/2211.10438. This modifier performs a channel-wise smoothing of outliers in activations, making them easier to quantize by reducing the dynamic range. The smoothing is offset by applying the inverse operation to the next layer of weights, making the weights slightly more difficult to quantize.

Because this modifier manipulates the weights of the model, it can only be used in in one-shot and not during training. Activation ranges are determined by running a small set of calibration data through the model.

example recipe:

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, intensity of smoothing to perform (0-1 range) :param mappings: list activation layers to smooth, and which layers to scale the output such that activations are smoothed. Each entry of the mapping list should be a list itself, in which the first entry is a list of layers who share the same input activation (the one to be to smoothed) and the second entry is the layer whose output is scaled to achieve the smoothing. If regex is used, it matches layers with the largest overlap in module name. If not supplied the argument will be inferred from the model architecture. :param ignore: list of layers to ignore, even if they match a regex in mappings. It should match the name of layers whose outputs are scaled to achieve smoothing (the second entry of the mappings list). :param num_calibration_steps: number of samples to use for calibration, or None to use the whole dataset

Parameters:

• calibration_function

  –

  optional function to use for the forward pass, or None to use the default tensor_module_forward

Methods:

• on_finalize –

  Clean up by clearing the scale and mapping data

• on_initialize –

  Initialize and run SmoothQuant on the given state

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

164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179

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(state: State, **kwargs) -> bool

112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

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

Dataclass for storing the channel-wise minimum and maximum values for a layer. This is updated each forward pass during calibration

Parameters:

• min_channel_vals

  (Tensor) –

  minimum output value seen so far, per channel

• max_channel_vals

  (Tensor) –

  maximum output value seen so far, per channel