llmcompressor.modifiers.transform.spinquant

A class for defining an event that can be triggered during sparsification.

Parameters:

• type_ (Optional[EventType], default: None ) –

  The type of event.

• steps_per_epoch (Optional[int], default: None ) –

  The number of steps per epoch.

• batches_per_step (Optional[int], default: None ) –

  The number of batches per step where step is an optimizer step invocation. For most pathways, these are the same. See the invocations_per_step parameter for more details when they are not.

• invocations_per_step (int, default: 1 ) –

  The number of invocations of the step wrapper before optimizer.step was called. Generally can be left as 1 (default). For older amp pathways, this is the number of times the scaler wrapper was invoked before the wrapped optimizer step function was called to handle accumulation in fp16.

• global_step (int, default: 0 ) –

  The current global step.

• global_batch (int, default: 0 ) –

  The current global batch.

Methods:

• new_instance –

  Creates a new instance of the event with the provided keyword arguments.

• should_update –

  Determines if the event should trigger an update.

Attributes:

• current_index (float) –

  Calculates the current index of the event.

• epoch (int) –

  Calculates the current epoch.

• epoch_based (bool) –

  Determines if the event is based on epochs.

• epoch_batch (int) –

  Calculates the current batch within the current epoch.

• epoch_full (float) –

  Calculates the current epoch with the fraction of the current step.

• epoch_step (int) –

  Calculates the current step within the current epoch.

current_index: float

epoch: int

epoch_based: bool

epoch_batch: int

epoch_full: float

epoch_step: int

new_instance(**kwargs) -> Event

239
240
241
242
243
244
245
246
247
248
249
250
251

def new_instance(self, **kwargs) -> "Event":
    """
    Creates a new instance of the event with the provided keyword arguments.

    :param kwargs: Keyword arguments to set in the new instance.
    :return: A new instance of the event with the provided kwargs.
    :rtype: Event
    """
    logger.debug("Creating new instance of event with kwargs: {}", kwargs)
    instance = deepcopy(self)
    for key, value in kwargs.items():
        setattr(instance, key, value)
    return instance

should_update(
    start: Optional[float],
    end: Optional[float],
    update: Optional[float],
) -> bool

208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237

def should_update(
    self, start: Optional[float], end: Optional[float], update: Optional[float]
) -> bool:
    """
    Determines if the event should trigger an update.

    :param start: The start index to check against, set to None to ignore start.
    :type start: Optional[float]
    :param end: The end index to check against, set to None to ignore end.
    :type end: Optional[float]
    :param update: The update interval, set to None or 0.0 to always update,
        otherwise must be greater than 0.0, defaults to None.
    :type update: Optional[float]
    :return: True if the event should trigger an update, False otherwise.
    :rtype: bool
    """
    current = self.current_index
    logger.debug(
        "Checking if event should update: "
        "current_index={}, start={}, end={}, update={}",
        current,
        start,
        end,
        update,
    )
    if start is not None and current < start:
        return False
    if end is not None and current > end:
        return False
    return update is None or update <= 0.0 or current % update < 1e-10

Bases: Enum

An Enum for defining the different types of events that can be triggered during model compression lifecycles. The purpose of each EventType is to trigger the corresponding modifier callback during training or post training pipelines.

Parameters:

• INITIALIZE –

  Event type for initialization.

• FINALIZE –

  Event type for finalization.

• BATCH_START –

  Event type for the start of a batch.

• LOSS_CALCULATED –

  Event type for when loss is calculated.

• BATCH_END –

  Event type for the end of a batch.

• CALIBRATION_EPOCH_START –

  Event type for the start of a calibration epoch.

• SEQUENTIAL_EPOCH_END –

  Event type for the end of a layer calibration epoch, specifically used by src/llmcompressor/pipelines/sequential/pipeline.py

• CALIBRATION_EPOCH_END –

  Event type for the end of a calibration epoch.

• OPTIM_PRE_STEP –

  Event type for pre-optimization step.

• OPTIM_POST_STEP –

  Event type for post-optimization step.

Bases: ModifierInterface, HooksMixin

A base class for all modifiers to inherit from. Modifiers are used to modify the training process for a model. Defines base attributes and methods available to all modifiers

Lifecycle: 1. initialize 2. on_event -> * on_start if self.start <= event.current_index * on_end if self.end >= event.current_index 5. finalize

Parameters:

• index –

  The index of the modifier in the list of modifiers for the model

• group –

  The group name for the modifier

• start –

  The start step for the modifier

• end –

  The end step for the modifier

• update –

  The update step for the modifier

Methods:

• finalize –

  Finalize the modifier for the given model and state.

• initialize –

  Initialize the modifier for the given model and state.

• on_end –

  on_end is called when the modifier ends and must be implemented

• on_event –

  on_event is called whenever an event is triggered

• on_finalize –

  on_finalize is called on modifier finalization and

• on_initialize –

  on_initialize is called on modifier initialization and

• on_start –

  on_start is called when the modifier starts and

• on_update –

  on_update is called when the model in question must be

• should_end –

  :param event: The event to check if the modifier should end

• should_start –

  :param event: The event to check if the modifier should start

• update_event –

  Update modifier based on the given event. In turn calls

Attributes:

• finalized (bool) –

  :return: True if the modifier has been finalized

• initialized (bool) –

  :return: True if the modifier has been initialized

finalized: bool

initialized: bool

finalize(state: State, **kwargs)

 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102

def finalize(self, state: State, **kwargs):
    """
    Finalize the modifier for the given model and state.

    :raises RuntimeError: if the modifier has not been initialized
    :param state: The current state of the model
    :param kwargs: Additional arguments for finalizing the modifier
    """
    if self.finalized_:
        raise RuntimeError("cannot finalize a modifier twice")

    if not self.initialized_:
        raise RuntimeError("cannot finalize an uninitialized modifier")

    # TODO: all finalization should succeed
    self.finalized_ = self.on_finalize(state=state, **kwargs)

initialize(state: State, **kwargs)

61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

def initialize(self, state: State, **kwargs):
    """
    Initialize the modifier for the given model and state.

    :raises RuntimeError: if the modifier has already been finalized
    :param state: The current state of the model
    :param kwargs: Additional arguments for initializing the modifier
    """
    if self.initialized_:
        raise RuntimeError(
            "Cannot initialize a modifier that has already been initialized"
        )

    if self.finalized_:
        raise RuntimeError(
            "Cannot initialize a modifier that has already been finalized"
        )

    self.initialized_ = self.on_initialize(state=state, **kwargs)

    # trigger starts
    fake_start_event = Event(type_=EventType.BATCH_START, global_step=0)
    if self.should_start(fake_start_event):
        self.on_start(state, fake_start_event, **kwargs)
        self.started_ = True

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

220
221
222
223
224
225
226
227
228
229

def on_end(self, state: State, event: Event, **kwargs):
    """
    on_end is called when the modifier ends and must be implemented
    by the inheriting modifier.

    :param state: The current state of the model
    :param event: The event that triggered the end
    :param kwargs: Additional arguments for ending the modifier
    """
    pass

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

231
232
233
234
235
236
237
238
239

def on_event(self, state: State, event: Event, **kwargs):
    """
    on_event is called whenever an event is triggered

    :param state: The current state of the model
    :param event: The event that triggered the update
    :param kwargs: Additional arguments for updating the model
    """
    pass

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

185
186
187
188
189
190
191
192
193
194
195

def on_finalize(self, state: State, **kwargs) -> bool:
    """
    on_finalize is called on modifier finalization and
    must be implemented by the inheriting modifier.

    :param state: The current state of the model
    :param kwargs: Additional arguments for finalizing the modifier
    :return: True if the modifier was finalized successfully,
        False otherwise
    """
    return True

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

172
173
174
175
176
177
178
179
180
181
182
183

@abstractmethod
def on_initialize(self, state: State, **kwargs) -> bool:
    """
    on_initialize is called on modifier initialization and
    must be implemented by the inheriting modifier.

    :param state: The current state of the model
    :param kwargs: Additional arguments for initializing the modifier
    :return: True if the modifier was initialized successfully,
        False otherwise
    """
    raise NotImplementedError()

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

197
198
199
200
201
202
203
204
205
206

def on_start(self, state: State, event: Event, **kwargs):
    """
    on_start is called when the modifier starts and
    must be implemented by the inheriting modifier.

    :param state: The current state of the model
    :param event: The event that triggered the start
    :param kwargs: Additional arguments for starting the modifier
    """
    pass

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

208
209
210
211
212
213
214
215
216
217
218

def on_update(self, state: State, event: Event, **kwargs):
    """
    on_update is called when the model in question must be
    updated based on passed in event. Must be implemented by the
    inheriting modifier.

    :param state: The current state of the model
    :param event: The event that triggered the update
    :param kwargs: Additional arguments for updating the model
    """
    pass

should_end(event: Event)

163
164
165
166
167
168
169
170

def should_end(self, event: Event):
    """
    :param event: The event to check if the modifier should end
    :return: True if the modifier should end based on the given event
    """
    current = event.current_index

    return self.end is not None and current >= self.end

should_start(event: Event) -> bool

151
152
153
154
155
156
157
158
159
160
161

def should_start(self, event: Event) -> bool:
    """
    :param event: The event to check if the modifier should start
    :return: True if the modifier should start based on the given event
    """
    if self.start is None:
        return False

    current = event.current_index

    return self.start <= current and (self.end is None or current < self.end)

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

104
105
106
107
108
109
110
111
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
141
142
143
144
145
146
147
148
149

def update_event(self, state: State, event: Event, **kwargs):
    """
    Update modifier based on the given event. In turn calls
    on_start, on_update, and on_end based on the event and
    modifier settings. Returns immediately if the modifier is
    not initialized

    :raises RuntimeError: if the modifier has been finalized
    :param state: The current state of compression
    :param event: The event to update the modifier with
    :param kwargs: Additional arguments for updating the modifier
    """
    if not self.initialized_:
        raise RuntimeError("Cannot update an uninitialized modifier")

    if self.finalized_:
        raise RuntimeError("Cannot update a finalized modifier")

    self.on_event(state, event, **kwargs)

    # handle starting the modifier if needed
    if (
        event.type_ == EventType.BATCH_START
        and not self.started_
        and self.should_start(event)
    ):
        self.on_start(state, event, **kwargs)
        self.started_ = True
        self.on_update(state, event, **kwargs)

        return

    # handle ending the modifier if needed
    if (
        event.type_ == EventType.BATCH_END
        and not self.ended_
        and self.should_end(event)
    ):
        self.on_end(state, event, **kwargs)
        self.ended_ = True
        self.on_update(state, event, **kwargs)

        return

    if self.started_ and not self.ended_:
        self.on_update(state, event, **kwargs)

Bases: BaseModel

SpinQuant needs to know where every norm layer exists in the model, as well as all the subsequent Linear layers the norm passes into. This is because the norm layer weights need to normalized before transforms can be fused into Linear layers.

Parameters:

• norm –

  name or regex that matches norm layer in model

• linears –

  list of names or regexes of Linear layers that receive input from norm.

Bases: BaseModel

SpinQuant needs to know the entire architecture of the model, as R1, R2, R3, and R4 rotations need to be applied to specific layers (https://arxiv.org/pdf/2405.16406 Fig. 1).

Parameters:

• embedding –

  name or regex of embedding layer

• attn –

  name or regex of attention block in decoder layer

• attn_q –

  name or regex of q_proj layer in attention block

• attn_k –

  name or regex of k_proj layer in attention block

• attn_v –

  name or regex of v_proj layer in attention block

• attn_o –

  name or regex of o_proj layer in attention block

• attn_head_dim –

  head_dim of the attention module, needed because R2 needs to be applied "head-wisely" to v_proj and o_proj

• mlp_in –

  list of names or regexes for the mlp blocks that receive the input to the MLP block, usually up_proj and gate_proj

• mlp_out –

  list of names or regexes for the mlp blocks that constitute the output of the MLP block, usually down_proj

Bases: Modifier

Implements the transforms according to "SpinQuant: LLM quantization with learned rotations" (https://arxiv.org/abs/2405.16406)

Transforms (rotations) are extra layers added to a model which reduce the accuracy loss induced by quantization. This is achieved through "rotating" weights and activations into a space with a smaller dynamic range of values, thus decreasing the range of scales required for quantization.

The SpinQuant authors describe four different rotations which can be applied to a model. R1 and R2 are "offline" rotations, meaning that they can be fused into existing weights and therefore do not induce runtime cost. R3 and R4 are "online" rotations, meaning that they require additional computation at runtime.

Lifecycle:

• on_initialize
  • infer SpinQuantMappings & NormMappings
  • as needed, create transform schemes for R1, R2, R3, & R4
• on_start
  • normalize embeddings
  • fuse norm layers into subsequent Linear layers
  • apply TransformConfig
    • fuse transforms into weights for mergeable transforms
    • add hooks for online transforms
• on sequential epoch end
• on_end
• on_finalize

Parameters:

• rotations –

  A list containing the names of rotations to apply to the model. Possible rotations include R1, R2, R3, and R4

• transform_type –

  The type of transform to apply to the model. "hadamard" has the least performance cost but only supports sizes which are powers of power of two. "random-matrix" has more performance cost, but supports a much larger set of sizes. "random-matrix" has the greatest performance cost, but supports any size

• randomize –

  if True, create distinct transforms for each application

• learnable –

  if True, attach gradients to transform weights for training

• precision –

  Precision at which all transforms should be applied. This applies to both weight fusing and online rotations

• transform_block_size –

  Block size to use for rotation matrices. The model's hidden_size and head_dim must be evenly divisible by transform_block_size. Layers will be transformed by a block-diagonal matrix where each block is a matrix of this size. If None is provided, model's hidden_size will be used for R1, R3, and R4 and model's head_dim will be used for R2

• mappings –

  Specifies layers within a model to target for transforms. A mapping will be inferred if None is provided

• norm_mappings –

  Specifies layers within a model to target for norm fusing. A mapping will be inferred if None is provided

• transform_config –

  Optional transform config for overriding provided arguments

State class holds information about the current compression state.

Parameters:

• model (Any, default: None ) –

  The model being used for compression

• teacher_model (Any, default: None ) –

  The teacher model being used for compression

• optimizer (Any, default: None ) –

  The optimizer being used for training

• optim_wrapped (bool, default: None ) –

  Whether or not the optimizer has been wrapped

• loss (Any, default: None ) –

  The loss function being used for training

• batch_data (Any, default: None ) –

  The current batch of data being used for compression

• data (Data, default: Data() ) –

  The data sets being used for training, validation, testing, and/or calibration, wrapped in a Data instance

• hardware (Hardware, default: Hardware() ) –

  Hardware instance holding info about the target hardware being used

Methods:

• update –

  Update the state with the given parameters.

Attributes:

• compression_ready (bool) –

  Check if the model and optimizer are set for compression.

compression_ready: bool

update(
    model: Any = None,
    teacher_model: Any = None,
    optimizer: Any = None,
    attach_optim_callbacks: bool = True,
    train_data: Any = None,
    val_data: Any = None,
    test_data: Any = None,
    calib_data: Any = None,
    copy_data: bool = True,
    start: float = None,
    steps_per_epoch: int = None,
    batches_per_step: int = None,
    **kwargs,
) -> dict

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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206

def update(
    self,
    model: Any = None,
    teacher_model: Any = None,
    optimizer: Any = None,
    attach_optim_callbacks: bool = True,
    train_data: Any = None,
    val_data: Any = None,
    test_data: Any = None,
    calib_data: Any = None,
    copy_data: bool = True,
    start: float = None,
    steps_per_epoch: int = None,
    batches_per_step: int = None,
    **kwargs,
) -> dict:
    """
    Update the state with the given parameters.

    :param model: The model to update the state with
    :type model: Any
    :param teacher_model: The teacher model to update the state with
    :type teacher_model: Any
    :param optimizer: The optimizer to update the state with
    :type optimizer: Any
    :param attach_optim_callbacks: Whether or not to attach optimizer callbacks
    :type attach_optim_callbacks: bool
    :param train_data: The training data to update the state with
    :type train_data: Any
    :param val_data: The validation data to update the state with
    :type val_data: Any
    :param test_data: The testing data to update the state with
    :type test_data: Any
    :param calib_data: The calibration data to update the state with
    :type calib_data: Any
    :param copy_data: Whether or not to copy the data
    :type copy_data: bool
    :param start: The start index to update the state with
    :type start: float
    :param steps_per_epoch: The steps per epoch to update the state with
    :type steps_per_epoch: int
    :param batches_per_step: The batches per step to update the state with
    :type batches_per_step: int
    :param kwargs: Additional keyword arguments to update the state with
    :return: The updated state as a dictionary
    :rtype: Dict
    """
    logger.debug(
        "Updating state with provided parameters: {}",
        {
            "model": model,
            "teacher_model": teacher_model,
            "optimizer": optimizer,
            "attach_optim_callbacks": attach_optim_callbacks,
            "train_data": train_data,
            "val_data": val_data,
            "test_data": test_data,
            "calib_data": calib_data,
            "copy_data": copy_data,
            "start": start,
            "steps_per_epoch": steps_per_epoch,
            "batches_per_step": batches_per_step,
            "kwargs": kwargs,
        },
    )

    if model is not None:
        self.model = model
    if teacher_model is not None:
        self.teacher_model = teacher_model
    if optimizer is not None:
        self.optim_wrapped = attach_optim_callbacks
        self.optimizer = optimizer
    if train_data is not None:
        self.data.train = train_data if not copy_data else deepcopy(train_data)
    if val_data is not None:
        self.data.val = val_data if not copy_data else deepcopy(val_data)
    if test_data is not None:
        self.data.test = test_data if not copy_data else deepcopy(test_data)
    if calib_data is not None:
        self.data.calib = calib_data if not copy_data else deepcopy(calib_data)

    if "device" in kwargs:
        self.hardware.device = kwargs["device"]

    return kwargs