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
      offload_hessians: False
      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_start
  • add activation calibration hooks
  • add gptq weight calibration hooks
• on_sequential_epoch_end
  • quantize_weight
• on_finalize
  • 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

• offload_hessians –

  Set to True for decreased memory usage but increased runtime.

• 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_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,
)

255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289

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:
        init_device = (
            "cpu" if self.offload_hessians else get_execution_device(module)
        )
        self._hessians[module] = make_empty_hessian(module, device=init_device)
        self._num_samples[module] = torch.zeros(
            tuple(), device=get_execution_device(module)
        )

    # Accumulate hessian with input with optional offloading
    with self._maybe_onload_hessian(module):
        self._hessians[module], self._num_samples[module] = accumulate_hessian(
            inp,
            module,
            self._hessians[module],
            self._num_samples[module],
        )

compress_modules()

291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317

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
    self._broadcast_quantized_params(module_list, module_to_rank)

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

388
389
390
391
392
393
394

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

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

396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411

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)

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

    self._hessians = dict()
    self._num_samples = dict()

    return True

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

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

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