llmcompressor.modifiers.gptq.base

Bases: Modifier, QuantizationMixin

Implements the GPTQ algorithm from https://arxiv.org/abs/2210.17323. This modifier uses activations to calibrate a hessian matrix, which is then used to determine optimal quantization values and orderings for the model weights.

Sample yaml:

test_stage:
  obcq_modifiers:
    GPTQModifier:
      block_size: 128
      dampening_frac: 0.001
      actorder: static
      config_groups:
        group_0:
          targets:
            - "Linear"
          input_activations: null
          output_activations: null
          weights:
            num_bits: 8
            type: "int"
            symmetric: true
            strategy: group
            group_size: 128

Lifecycle:

• on_initialize
  • apply config to model
• on_calibration_start
  • add activation calibration hooks
  • add gptq weight calibration hooks
• on_sequential_epoch_end
  • quantize_weight
• on_calibration_end
  • remove_hooks()
  • model.apply(freeze_module_quantization)

Parameters:

• block_size –

  Used to determine number of columns to compress in one pass

• dampening_frac –

  Amount of dampening to apply to H, as a fraction of the diagonal norm

• actorder –

  order in which weight columns are quantized. Defaults to "static" activation ordering, which achieves best accuracy recovery with no runtime cost. For more information, see https://github.com/vllm-project/vllm/pull/8135.

• batched_quantization –

  Controls batching of same-shape modules (e.g. linearized MoE experts). "auto" (the default) limits batches to 75% of available CUDA memory. A positive integer sets a maximum batch size without consulting available memory. None disables batching.

• config_groups –

  dictionary specifying quantization schemes to apply to target modules. Modules not matching a scheme target will NOT be quantized.

• targets –

  list of layer names to quantize if a scheme is provided. Defaults to Linear layers

• ignore –

  optional list of module class names or submodule names to not quantize even if they match a target in config_groups. Defaults to empty list.

• scheme –

  a single quantization scheme to apply to the model. This is a dictionary that supports all keys from QuantizationScheme except targets, which will be set to the targets parameter set at the modifier level. Can also be set to a dictionary of the format preset_scheme_name: targets for example: W8A8: ['Linear'] for weight and activation 8-bit.

• kv_cache_scheme –

  optional QuantizationArgs, that specify the quantization of the kv cache. If None, kv cache is not quantized. When applying kv cache quantization to transformer AutoModelForCausalLM, the kv_cache_scheme gets converted into a QuantizationScheme that: - targets the q_proj and k_proj modules of the model. The outputs of those modules are the keys and values that might be cached - quantizes the outputs of the aforementioned layers, so that keys and values are compressed before storing them in the cache There is an explicit assumption that the model contains modules with k_proj and v_proj in their names. If this is not the case and kv_cache_scheme != None, the quantization of kv cache will fail

Methods:

• calibrate_module –

  Calibration hook used to accumulate the hessian of the input to the module

• compress_modules –

  Quantize modules which have been calibrated

• on_calibration_end –

  Finish calibrating by removing observers and calibration hooks

• on_finalize –

  disable the quantization observers used by the OBCQ algorithm

• on_initialize –

  Initialize and run the GPTQ algorithm on the current state

calibrate_module(
    module: Module,
    args: tuple[Tensor, ...],
    _output: Tensor,
)

247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278

def calibrate_module(
    self,
    module: torch.nn.Module,
    args: tuple[torch.Tensor, ...],
    _output: torch.Tensor,
):
    """
    Calibration hook used to accumulate the hessian of the input to the module

    :param module: module being calibrated
    :param args: inputs to the module, the first element of which is the
        canonical input
    :param _output: uncompressed module output, unused
    """
    # Assume that first argument is the input
    inp = args[0]

    # Initialize hessian if not present
    if module not in self._num_samples:
        self._hessians[module] = make_empty_hessian(
            module, device=get_execution_device(module)
        )
        self._num_samples[module] = torch.zeros(
            tuple(), device=get_execution_device(module)
        )

    self._hessians[module], self._num_samples[module] = accumulate_hessian(
        inp,
        module,
        self._hessians[module],
        self._num_samples[module],
    )

compress_modules()

280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306

def compress_modules(self):
    """
    Quantize modules which have been calibrated
    """
    ### Not Distributed
    if not is_distributed():
        self.compress_module_list(list(self._num_samples.keys()))
        return

    ### Distributed
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # Assign modules to ranks
    module_list, rank_to_modules, module_to_rank = greedy_bin_packing(
        list(self._hessians.keys()),
        world_size,
        item_weight_fn=lambda mod: self._hessians[mod].shape[0],
    )

    # send hessians to assigned ranks
    self._reduce_hessian_to_target_rank(module_list, module_to_rank)

    self.compress_module_list(rank_to_modules[rank])

    # broadcast compressed modules to each rank
    broadcast_qparams_and_cleanup(module_list, module_to_rank, _GPTQ_Q_PARAMS)

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

240
241
242
243
244
245

def on_calibration_end(self, state: State, event: Event, **kwargs):
    """
    Finish calibrating by removing observers and calibration hooks
    """
    QuantizationMixin.end_calibration(self, state.model)
    self.remove_hooks()  # remove gptq hooks

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

414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433

def on_finalize(self, state: State, **kwargs) -> bool:
    """
    disable the quantization observers used by the OBCQ algorithm

    :param state: session state storing input model and calibration data
    """
    if not self.ended_:
        self.on_end(state, None)

    self._log_rtn_fallback_summary()

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

    self._hessians = dict()
    self._num_samples = dict()
    self._num_compressed_modules = 0
    self._rtn_fallback_module_names = []

    return True

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

187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205

def on_initialize(self, state: State, **kwargs) -> bool:
    """
    Initialize and run the GPTQ algorithm on the current state

    :param state: session state storing input model and calibration data
    """
    # apply config to model and prepare calibration hooks
    if QuantizationMixin.has_config(self):
        QuantizationMixin.initialize_quantization(self, state.model)

    # prepare module names
    self._module_names = {
        m: name
        for name, m in match_named_modules(
            state.model, self.resolved_targets, self.ignore
        )
    }

    return True