llmcompressor.modifiers.pruning

Bases: Modifier

Prunes experts from MoE layers using the REAP saliency metric. For each expert j the saliency is

``S_j = mean(g_j * ||f_j||_2)``

averaged over the tokens routed to expert j, where:

• g_j is the router gate weight assigned to expert j (the coefficient that multiplies the expert's output when combining experts), and
• f_j is expert j's output activation for that token, so ||f_j||_2 is its L2 norm.

The lowest-saliency experts are removed per layer. REAP runs during the sequential calibration pipeline: saliency is accumulated via hooks on the MoE experts, the structural pruning for a layer is executed when it completes (SEQUENTIAL_EPOCH_END). The config is updated to reflect the new number of experts in on_finalize.

Parameters:

• sparsity –

  fraction of experts to remove per layer (0, 1).

• ignore –

  module name patterns to skip during MoE layer detection.

  Example recipe::

  REAPPruningModifier:
    sparsity: 0.25
  

Methods:

• on_finalize –

  Finalize the model config to reflect the new number of experts.

• on_sequential_epoch_end –

  Prune any tracked layer whose saliency is

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

208
209
210
211
212
213
214
215
216
217
218
219

def on_finalize(self, state: State, **kwargs) -> bool:
    """Finalize the model config to reflect the new number of experts."""

    model = state.model

    new_num_experts = self._moe_attrs.num_experts - self._n_experts_to_drop
    update_model_config(model, self._moe_attrs, new_num_experts)

    self._saliency_trackers.clear()
    self._norm_buffers.clear()

    return True

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

223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247

def on_sequential_epoch_end(self, state: State, event: Event, **kwargs):
    """Prune any tracked layer whose saliency is
    complete, then release its activation norm buffers."""

    model = state.model

    for layer_name, tracker in list(self._saliency_trackers.items()):
        if tracker.total_count <= 0:
            continue

        retained = tracker.compute_retained_experts(
            self._n_experts_to_drop,
            self._n_experts_to_drop_per_group,
            self._moe_attrs,
        )
        expected = self._moe_attrs.num_experts - self._n_experts_to_drop
        assert (
            len(retained) == expected
        ), f"Expected {expected} retained experts, got {len(retained)}"

        prune_moe_layer(model, layer_name, retained, self._moe_attrs)

        # free this layer's accumulators / buffers now
        del self._saliency_trackers[layer_name]
        self._norm_buffers.pop(layer_name, None)

Bases: SparsityModifierBase

Modifier for applying the one-shot SparseGPT algorithm to a model

Sample yaml:

test_stage:
  obcq_modifiers:
    SparseGPTModifier:
      sparsity: 0.5
      mask_structure: "2:4"
      dampening_frac: 0.001
      block_size: 128
      targets: ['Linear']
      ignore: ['re:.*lm_head']

Lifecycle:

• on_initialize
  • register_hook(module, calibrate_module, "forward")
• on_sequential_batch_end
  • sparsify_weight
• on_finalize
  • remove_hooks()

Parameters:

• sparsity –

  Sparsity to compress model to

• sparsity_profile –

  Can be set to 'owl' to use Outlier Weighed Layerwise Sparsity (OWL), more information can be found in the paper https://arxiv.org/pdf/2310.05175

• mask_structure –

  String to define the structure of the mask to apply. Must be of the form N:M where N, M are integers that define a custom block shape. Defaults to 0:0 which represents an unstructured mask.

• owl_m –

  Number of outliers to use for OWL

• owl_lmbda –

  Lambda value to use for OWL

• 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

• preserve_sparsity_mask –

  Whether or not to preserve the sparsity mask during when applying sparsegpt, this becomes useful when starting from a previously pruned model, defaults to False.

• offload_hessians –

  Set to True for decreased memory usage but increased runtime.

• targets –

  list of layer names to compress during SparseGPT, or 'ALL' to compress every layer in the model

• ignore –

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

Methods:

• calibrate_module –

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

• compress_modules –

  Sparsify modules which have been calibrated

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

 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

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 the first argument is the input
    inp = args[0]

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

    # 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()

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
141
142
143
144
145
146

def compress_modules(self):
    """
    Sparsify modules which have been calibrated
    """
    for module in list(self._num_samples.keys()):
        name = self._module_names[module]
        sparsity = self._module_sparsities[module]
        num_samples = self._num_samples[module]

        logger.info(f"Sparsifying {name} using {num_samples} samples")
        with (
            torch.no_grad(),
            align_module_device(module),
            CompressionLogger(module) as comp_logger,
        ):
            loss, sparsified_weight = sparsify_weight(
                module=module,
                hessians_dict=self._hessians,
                sparsity=sparsity,
                prune_n=self._prune_n,
                prune_m=self._prune_m,
                block_size=self.block_size,
                dampening_frac=self.dampening_frac,
                preserve_sparsity_mask=self.preserve_sparsity_mask,
            )
            comp_logger.set_results(name="SGPT", loss=loss)

        update_offload_parameter(module, "weight", sparsified_weight)

        # self._hessians[module] already deleted by sparsify_weight
        del self._num_samples[module]

Bases: SparsityModifierBase

Modifier for applying the one-shot WANDA algorithm to a model from the paper: https://arxiv.org/abs/2306.11695

Sample yaml:

test_stage:
  sparsity_modifiers:
    WandaPruningModifier:
      sparsity: 0.5
      mask_structure: "2:4"

Lifecycle:

• on_initialize
  • register_hook(module, calibrate_module, "forward")
  • run_sequential / run_basic
    • make_empty_row_scalars
    • accumulate_row_scalars
• on_sequential_batch_end
  • sparsify_weight
• on_finalize
  • remove_hooks()

Parameters:

• sparsity –

  Sparsity to compress model to

• sparsity_profile –

  Can be set to 'owl' to use Outlier Weighed Layerwise Sparsity (OWL), more information can be found in the paper https://arxiv.org/pdf/2310.05175

• mask_structure –

  String to define the structure of the mask to apply. Must be of the form N:M where N, M are integers that define a custom block shape. Defaults to 0:0 which represents an unstructured mask.

• owl_m –

  Number of outliers to use for OWL

• owl_lmbda –

  Lambda value to use for OWL

• targets –

  list of layer names to compress during OBCQ, or 'ALL' to compress every layer in the model

• ignore –

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

Methods:

• calibrate_module –

  Calibration hook used to accumulate the row scalars of the input to the module

• compress_modules –

  Sparsify modules which have been calibrated

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

70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

def calibrate_module(
    self,
    module: torch.nn.Module,
    args: tuple[torch.Tensor, ...],
    _output: torch.Tensor,
):
    """
    Calibration hook used to accumulate the row scalars 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 the first argument is the input
    inp = args[0]

    # Initialize row scalars if not present
    if module not in self._num_samples:
        device = get_execution_device(module)
        self._row_scalars[module] = make_empty_row_scalars(module, device=device)
        self._num_samples[module] = 0

    # Accumulate scalars using data
    self._row_scalars[module], self._num_samples[module] = accumulate_row_scalars(
        inp,
        module,
        self._row_scalars[module],
        self._num_samples[module],
    )

compress_modules()

101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

def compress_modules(self):
    """
    Sparsify modules which have been calibrated
    """
    for module in list(self._num_samples.keys()):
        name = self._module_names[module]
        sparsity = self._module_sparsities[module]
        num_samples = self._num_samples[module]

        logger.info(f"Sparsifying {name} using {num_samples} samples")
        with (
            torch.no_grad(),
            align_module_device(module),
            CompressionLogger(module),
        ):
            sparsified_weight = sparsify_weight(
                module=module,
                row_scalars_dict=self._row_scalars,
                sparsity=sparsity,
                prune_n=self._prune_n,
                prune_m=self._prune_m,
            )

        update_offload_parameter(module, "weight", sparsified_weight)

        # self._row_scalars[module] already deleted by sparsify_weight
        del self._num_samples[module]