跳到内容

llmcompressor.modifiers.transform.spinquant

模块

类

  • Event –

    一个用于定义在稀疏化过程中可以触发的事件的类。

  • EventType –

    一个枚举，用于定义可以触发的不同类型的事件

  • Modifier –

    所有修改器继承的基类。

  • NormMapping –

    SpinQuant 需要知道模型中每个 norm 层的位置，

  • SpinQuantMapping –

    SpinQuant 需要知道模型的整个架构，

  • SpinQuantModifier –

    根据“SpinQuant: LLM 量化

  • State –

    状态类保存有关当前压缩状态的信息。

函数

Event dataclass 

Event(
    type_: Optional[EventType] = None,
    steps_per_epoch: Optional[int] = None,
    batches_per_step: Optional[int] = None,
    invocations_per_step: int = 1,
    global_step: int = 0,
    global_batch: int = 0,
)

一个用于定义在稀疏化过程中可以触发的事件的类。

参数

  • 类型_

    (Optional[EventType], default: None ) –

    事件类型。

  • steps_per_epoch

    (Optional[int], 默认值: None ) –

    每个 epoch 的步数。

  • batches_per_step

    (Optional[int], 默认值: None ) –

    每个步骤的批次数量，其中步骤是优化器步骤调用。对于大多数路径，它们是相同的。当它们不同时，请参阅 invocations_per_step 参数以获取更多详细信息。

  • invocations_per_step

    (int, 默认值: 1 ) –

    调用步骤包装器之前 optimizer.step 被调用的次数。通常可以保留为 1（默认值）。对于较旧的 amp 路径，这是在调用包装的优化器步骤函数以处理 fp16 中的累积之前调用缩放器包装器的次数。

  • global_step

    (int, 默认值: 0 ) –

    当前全局步骤。

  • global_batch

    (int, 默认值: 0 ) –

    当前全局批次。

方法

  • new_instance –

    使用提供的关键字参数创建事件的新实例。

  • should_update –

    确定事件是否应触发更新。

属性

  • current_index (float) –

    计算事件的当前索引。

  • epoch (int) –

    计算当前 epoch。

  • epoch_based (bool) –

    确定事件是否基于 epoch。

  • epoch_batch (int) –

    计算当前 epoch 中的当前批次。

  • epoch_full (float) –

    计算当前 epoch 和当前步骤的比例。

  • epoch_step (int) –

    计算当前 epoch 中的当前步骤。

current_index property writable 

current_index: float

计算事件的当前索引。

返回

  • float –

    事件的当前索引，可以是全局步数，也可以是当前步数占 epoch 的比例。

引发

  • ValueError –

    如果事件不是基于 epoch 的，或者每个 epoch 的步数太多。

epoch property 

epoch: int

计算当前 epoch。

返回

  • int –

    当前 epoch。

引发

  • ValueError –

    如果事件不是基于 epoch 的。

epoch_based property 

epoch_based: bool

确定事件是否基于 epoch。

返回

  • bool –

    如果事件基于 epoch，则为 True，否则为 False。

epoch_batch property 

epoch_batch: int

计算当前 epoch 中的当前批次。

返回

  • int –

    当前 epoch 中的当前批次。

引发

  • ValueError –

    如果事件不是基于 epoch 的。

epoch_full property 

epoch_full: float

计算当前 epoch 和当前步骤的比例。

返回

  • float –

    当前 epoch 和当前步骤的比例。

引发

  • ValueError –

    如果事件不是基于 epoch 的。

epoch_step property 

epoch_step: int

计算当前 epoch 中的当前步骤。

返回

  • int –

    当前 epoch 中的当前步骤。

引发

  • ValueError –

    如果事件不是基于 epoch 的。

new_instance

new_instance(**kwargs) -> Event

使用提供的关键字参数创建事件的新实例。

参数

  • kwargs

    –

    要在新实例中设置的关键字参数。

返回

  • Event –

    具有提供的 kwargs 的事件新实例。

源代码位于 llmcompressor/core/events/event.py 
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

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

确定事件是否应触发更新。

参数

  • start

    (Optional[float]) –

    要检查的起始索引，设置为 None 则忽略起始。

  • 结束

    (Optional[float]) –

    要检查的结束索引，设置为 None 则忽略结束。

  • 更新

    (Optional[float]) –

    更新间隔，设置为 None 或 0.0 则始终更新，否则必须大于 0.0，默认为 None。

返回

  • bool –

    如果事件应触发更新，则为 True，否则为 False。

源代码位于 llmcompressor/core/events/event.py 
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

EventType

基类：Enum

一个枚举，用于定义在模型压缩生命周期中可以触发的不同类型的事件。每个 EventType 的目的是在训练或训练后管道中触发相应的修饰符回调。

参数

  • 初始化

    –

    初始化事件类型。

  • 完成

    –

    完成事件类型。

  • BATCH_START

    –

    批次开始事件类型。

  • 损失计算

    –

    损失计算事件类型。

  • BATCH_END

    –

    批次结束事件类型。

  • CALIBRATION_EPOCH_START

    –

    校准 epoch 开始事件类型。

  • SEQUENTIAL_EPOCH_END

    –

    层校准 epoch 结束事件类型，专门用于 src/llmcompressor/pipelines/sequential/pipeline.py

  • CALIBRATION_EPOCH_END

    –

    校准 epoch 结束事件类型。

  • OPTIM_PRE_STEP

    –

    优化前步骤事件类型。

  • OPTIM_POST_STEP

    –

    优化后步骤事件类型。

Modifier

Bases: ModifierInterface, HooksMixin

所有修改器继承的基类。修改器用于修改模型的训练过程。定义了所有修改器可用的基本属性和方法。

生命周期： 1. 初始化 2. on_event -> * 如果 self.start <= event.current_index 则 on_start * 如果 self.end >= event.current_index 则 on_end 5. 最终化

参数

  • index

    –

    修改器在模型修改器列表中的索引

  • group

    –

    修改器的组名

  • start

    –

    修改器的起始步长

  • 结束

    –

    修改器的结束步长

  • 更新

    –

    修改器的更新步长

方法

  • finalize –

    为给定模型和状态最终化修改器。

  • initialize –

    为给定模型和状态初始化修改器。

  • on_end –

    on_end 在修改器结束时调用，必须实现。

  • on_event –

    on_event 在事件触发时调用。

  • on_finalize –

    on_finalize 在修改器最终化时调用。

  • on_initialize –

    on_initialize 在修改器初始化时调用。

  • on_start –

    on_start 在修改器开始时调用。

  • on_update –

    on_update 在模型需要更新时调用。

  • should_end –

    :param event: 检查修改器是否应该结束的事件

  • should_start –

    :param event: 检查修改器是否应该开始的事件

  • update_event –

    根据给定事件更新修改器。进而调用

属性

  • finalized (bool) –

    :return: 如果修饰符已最终确定，则为 True

  • initialized (bool) –

    :return: 如果修饰符已初始化，则为 True

finalized property 

finalized: bool

返回

  • bool –

    如果修饰符已最终确定，则为 True

initialized property 

initialized: bool

返回

  • bool –

    如果修饰符已初始化，则为 True

finalize

finalize(state: State, **kwargs)

为给定模型和状态最终化修改器。

参数

  • state

    (State) –

    模型的当前状态

  • kwargs

    –

    最终化修改器的附加参数

引发

  • RuntimeError –

    如果修改器未初始化

源代码位于 llmcompressor/modifiers/modifier.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
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

initialize(state: State, **kwargs)

为给定模型和状态初始化修改器。

参数

  • state

    (State) –

    模型的当前状态

  • kwargs

    –

    初始化修改器的附加参数

引发

  • RuntimeError –

    如果修改器已最终化

源代码位于 llmcompressor/modifiers/modifier.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
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

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

on_end 在修改器结束时调用，必须由继承的修改器实现。

参数

  • state

    (State) –

    模型的当前状态

  • event

    (Event) –

    触发结束的事件

  • kwargs

    –

    结束修改器的附加参数

源代码位于 llmcompressor/modifiers/modifier.py 
221
222
223
224
225
226
227
228
229
230
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

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

on_event 在事件触发时调用。

参数

  • state

    (State) –

    模型的当前状态

  • event

    (Event) –

    触发更新的事件

  • kwargs

    –

    更新模型的附加参数

源代码位于 llmcompressor/modifiers/modifier.py 
232
233
234
235
236
237
238
239
240
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

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

on_finalize 在修改器最终化时调用，必须由继承的修改器实现。

参数

  • state

    (State) –

    模型的当前状态

  • kwargs

    –

    最终化修改器的附加参数

返回

  • bool –

    如果修改器成功最终化则为 True，否则为 False

源代码位于 llmcompressor/modifiers/modifier.py 
186
187
188
189
190
191
192
193
194
195
196
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 abstractmethod 

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

on_initialize 在修改器初始化时调用，必须由继承的修改器实现。

参数

  • state

    (State) –

    模型的当前状态

  • kwargs

    –

    初始化修改器的附加参数

返回

  • bool –

    如果修改器成功初始化则为 True，否则为 False

源代码位于 llmcompressor/modifiers/modifier.py 
173
174
175
176
177
178
179
180
181
182
183
184
@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

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

on_start 在修改器开始时调用，必须由继承的修改器实现。

参数

  • state

    (State) –

    模型的当前状态

  • event

    (Event) –

    触发开始的事件

  • kwargs

    –

    开始修改器的附加参数

源代码位于 llmcompressor/modifiers/modifier.py 
198
199
200
201
202
203
204
205
206
207
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

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

on_update 在相关模型必须根据传入事件进行更新时调用。必须由继承的修改器实现。

参数

  • state

    (State) –

    模型的当前状态

  • event

    (Event) –

    触发更新的事件

  • kwargs

    –

    更新模型的附加参数

源代码位于 llmcompressor/modifiers/modifier.py 
209
210
211
212
213
214
215
216
217
218
219
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

should_end(event: Event)

参数

  • event

    (Event) –

    检查修改器是否应该结束的事件

返回

  • –

    如果修改器应根据给定事件结束，则为 True

源代码位于 llmcompressor/modifiers/modifier.py 
164
165
166
167
168
169
170
171
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

should_start(event: Event) -> bool

参数

  • event

    (Event) –

    检查修改器是否应该开始的事件

返回

  • bool –

    如果修改器应根据给定事件开始，则为 True

源代码位于 llmcompressor/modifiers/modifier.py 
152
153
154
155
156
157
158
159
160
161
162
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

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

根据给定事件更新修改器。进而根据事件和修改器设置调用 on_start、on_update 和 on_end。如果修改器未初始化则立即返回。

参数

  • state

    (State) –

    稀疏化的当前状态

  • event

    (Event) –

    用于更新修饰符的事件

  • kwargs

    –

    更新修改器的附加参数

引发

  • RuntimeError –

    如果修改器已最终化

源代码位于 llmcompressor/modifiers/modifier.py 
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
150
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 sparsification
    :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)

NormMapping

基类：BaseModel

SpinQuant 需要知道模型中每个 norm 层的位置，以及 norm 传递到的所有后续 Linear 层。这是因为需要对 norm 层的权重进行归一化，然后才能将变换融合到 Linear 层中。

参数

  • norm

    –

    匹配模型中 norm 层的名称或正则表达式

  • linears

    –

    接收来自 norm 输入的 Linear 层的名称或正则表达式列表。

SpinQuantMapping

基类：BaseModel

SpinQuant 需要知道模型的整个架构，因为 R1、R2、R3 和 R4 旋转需要应用于特定层（https://arxiv.org/pdf/2405.16406 图 1）。

参数

  • embedding

    –

    嵌入层的名称或正则表达式

  • attn

    –

    注意力块中 attention 层的名称或正则表达式

  • attn_q

    –

    注意力块中 q_proj 层的名称或正则表达式

  • attn_k

    –

    注意力块中 k_proj 层的名称或正则表达式

  • attn_v

    –

    注意力块中 v_proj 层的名称或正则表达式

  • attn_o

    –

    注意力块中 o_proj 层的名称或正则表达式

  • attn_head_dim

    –

    注意力模块的 head_dim，因为 R2 需要“逐头”应用于 v_proj 和 o_proj

  • mlp_in

    –

    接收 MLP 块输入的 MLP 块的名称或正则表达式列表，通常是 up_proj 和 gate_proj

  • mlp_out

    –

    构成 MLP 块输出的 MLP 块的名称或正则表达式列表，通常是 down_proj

SpinQuantModifier

基类：Modifier

根据“SpinQuant: LLM 量化与学习旋转”实现变换（https://arxiv.org/abs/2405.16406）

变换（旋转）是添加到模型中的额外层，可以减少量化引起的精度损失。这是通过将权重和激活“旋转”到具有较小值动态范围的空间来实现的，从而减小了量化所需的尺度范围。

SpinQuant 作者描述了四种可以应用于模型的不同旋转。R1 和 R2 是“离线”旋转，意味着它们可以融合到现有权重中，因此不会引起运行时成本。R3 和 R4 是“在线”旋转，意味着它们在运行时需要额外的计算。

生命周期

  • on_initialize
    • 推断 SpinQuantMappings 和 NormMappings
    • 根据需要，为 R1、R2、R3 和 R4 创建变换方案
  • on_start
    • 归一化嵌入
    • 将 norm 层融合到后续的 Linear 层
    • 应用 TransformConfig
      • 将变换合并到权重中，以便进行合并变换
      • 添加在线变换的钩子
  • 在序列 epoch 结束时
  • on_end
  • on_finalize

参数

  • 旋转

    –

    一个列表，包含应用于模型的旋转名称。可能的旋转包括 R1、R2、R3 和 R4

  • 变换类型

    –

    应用于模型的变换类型。"hadamard" 性能成本最低，但只支持 2 的幂次幂的尺寸。"random-matrix" 性能成本较高，但支持更多尺寸。"random-matrix" 性能成本最高，但支持任何尺寸

  • 随机化

    –

    如果为 True，则为每次应用创建不同的变换

  • 可学习

    –

    如果为 True，则为变换权重附加梯度以进行训练

  • precision

    –

    所有变换应应用的精度。这适用于权重融合和在线旋转

  • transform_block_size

    –

    用于旋转矩阵的块大小。模型的 hidden_size 和 head_dim 必须能被 transform_block_size 整除。层将通过一个块对角矩阵进行变换，其中每个块是此大小的矩阵。如果提供 None，则 R1、R3 和 R4 将使用模型的 hidden_size，R2 将使用模型的 head_dim

  • mappings

    –

    指定模型中要应用变换的目标层。如果提供 None，则会自动推断映射

  • norm_mappings

    –

    指定模型中要应用 norm 融合的目标层。如果提供 None，则会自动推断映射

  • transform_config

    –

    用于覆盖提供的参数的可选变换配置

State dataclass 

State(
    model: Any = None,
    teacher_model: Any = None,
    optimizer: Any = None,
    optim_wrapped: bool = None,
    loss: Any = None,
    batch_data: Any = None,
    data: Data = Data(),
    hardware: Hardware = Hardware(),
    loggers: Optional[LoggerManager] = None,
    model_log_cadence: Optional[float] = None,
    _last_log_step: Union[float, int, None] = None,
)

状态类保存有关当前压缩状态的信息。

参数

  • model

    (Any, 默认值: None ) –

    用于压缩的模型

  • teacher_model

    (Any, 默认值: None ) –

    用于压缩的教师模型

  • 优化器

    (Any, 默认值: None ) –

    用于训练的优化器

  • optim_wrapped

    (bool, 默认值: None ) –

    优化器是否已被包装

  • 损失

    (Any, 默认值: None ) –

    用于训练的损失函数

  • batch_data

    (Any, 默认值: None ) –

    当前用于压缩的批次数据

  • data

    (Data, default: Data() ) –

    用于训练、验证、测试和/或校准的数据集，包装在 Data 实例中

  • 硬件

    (Hardware, default: Hardware() ) –

    硬件实例，包含有关正在使用的目标硬件的信息

  • loggers

    (Optional[LoggerManager], default: None ) –

    LoggerManager 实例，包含所有要记录的日志记录器

  • model_log_cadence

    (Optional[float], 默认值: None ) –

    记录模型信息相对于 epoch 的节奏。如果为 1，则每个 epoch 记录一次。如果为 2，则每隔一个 epoch 记录一次，依此类推。默认值为 1。

方法

  • update –

    使用给定参数更新状态。

属性

  • compression_ready (bool) –

    检查模型和优化器是否已设置为进行压缩。

compression_ready property 

compression_ready: bool

检查模型和优化器是否已设置为进行压缩。

返回

  • bool –

    如果模型和优化器已设置，则为 True，否则为 False

update

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,
    loggers: Union[
        None, LoggerManager, List[BaseLogger]
    ] = None,
    model_log_cadence: Optional[float] = None,
    **kwargs,
) -> Dict

使用给定参数更新状态。

参数

  • model

    (Any, 默认值: None ) –

    用于更新状态的模型

  • teacher_model

    (Any, 默认值: None ) –

    用于更新状态的教师模型

  • 优化器

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

    用于更新状态的每个 epoch 的步数

  • batches_per_step

    (int, 默认值: None ) –

    用于更新状态的每个步骤的批次数量

  • loggers

    (Union[None, LoggerManager, List[BaseLogger]], default: None ) –

    用于设置记录重要信息和里程碑的指标管理器，也接受 BaseLogger 列表

  • model_log_cadence

    (Optional[float], 默认值: None ) –

    记录模型信息相对于 epoch 的节奏。如果为 1，则每个 epoch 记录一次。如果为 2，则每隔一个 epoch 记录一次，依此类推。默认值为 1。

  • kwargs

    –

    用于更新状态的额外关键字参数

返回

  • Dict –

    更新后的状态字典

源代码位于 llmcompressor/core/state.py 
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
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
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,
    loggers: Union[None, LoggerManager, List[BaseLogger]] = None,
    model_log_cadence: Optional[float] = 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 loggers: The metrics manager to setup logging important info and
        milestones to, also accepts a list of BaseLogger(s)
    :type loggers: Union[None, LoggerManager, List[BaseLogger]]
    :param model_log_cadence: The cadence to log model information w.r.t epochs.
        If 1, logs every epoch. If 2, logs every other epoch, etc. Default is 1.
    :type model_log_cadence: Optional[float]
    :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,
            "loggers": loggers,
            "model_log_cadence": model_log_cadence,
            "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"]

    loggers = loggers or []
    if isinstance(loggers, list):
        loggers = LoggerManager(loggers)
    self.loggers = loggers

    if model_log_cadence is not None:
        self.model_log_cadence = model_log_cadence

    return kwargs

center_embeddings

center_embeddings(embedding: Module)

将每个嵌入的均值设为零

参数

  • embedding

    (Module) –

    包含要居中嵌入的嵌入模块

源代码在 llmcompressor/modeling/fuse.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def center_embeddings(embedding: torch.nn.Module):
    """
    Shift each embedding to have a mean of zero

    :param embedding: embedding module containing embeddings to center
    """
    if not hasattr(embedding, "weight"):
        raise ValueError(f"Cannot fuse norm of type {type(embedding)}")

    with align_module_device(embedding):
        weight_dtype = embedding.weight.dtype
        weight = embedding.weight.to(PRECISION)
        new_weight = weight - weight.mean(dim=-1, keepdim=True)
        new_weight = new_weight.to(weight_dtype)

    update_offload_parameter(embedding, "weight", new_weight)

fuse_norm_linears

fuse_norm_linears(norm: Module, linears: Iterable[Linear])

将归一化层的缩放操作融合到后续的线性层中。这对于确保归一化层和线性层之间的变换不变性很有用。

请注意，幺正变换（旋转）与归一化可交换，但与缩放不可交换

参数

  • norm

    (Module) –

    归一化层，其权重将融合到后续的线性层中

  • linears

    (Iterable[Linear]) –

    紧跟在归一化层之后的线性层

源代码在 llmcompressor/modeling/fuse.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def fuse_norm_linears(norm: torch.nn.Module, linears: Iterable[torch.nn.Linear]):
    """
    Fuse the scaling operation of norm layer into subsequent linear layers.
    This useful for ensuring transform invariance between norm and linear layers.

    Note that unitary transforms (rotation) commute with normalization, but not scaling

    :param norm: norm layer whose weight will be fused into subsequent linears
    :param linears: linear layers which directly follow the norm layer
    """
    if not hasattr(norm, "weight"):
        raise ValueError(f"Cannot fuse norm of type {type(norm)}")

    for linear in linears:
        # NOTE: spinquant does this op in float64
        exec_device = get_execution_device(norm)
        with align_module_device(norm, exec_device), align_module_device(
            linear, exec_device
        ):
            weight_dtype = linear.weight.dtype
            new_weight = linear.weight.to(PRECISION) * norm.weight.to(PRECISION)
            new_weight = new_weight.to(weight_dtype)

        update_offload_parameter(linear, "weight", new_weight)

    new_norm_weight = torch.ones_like(norm.weight, device="cpu")
    update_offload_parameter(norm, "weight", new_norm_weight)

untie_word_embeddings

untie_word_embeddings(model: PreTrainedModel)

如果可能，解绑词嵌入。如果模型定义中找不到嵌入，此函数将发出警告。

模型配置将更新以反映嵌入现已解绑

参数

  • model

    (PreTrainedModel) –

    包含词嵌入的 transformers 模型

源代码位于 llmcompressor/utils/transformers.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def untie_word_embeddings(model: PreTrainedModel):
    """
    Untie word embeddings, if possible. This function raises a warning if
    embeddings cannot be found in the model definition.

    The model config will be updated to reflect that embeddings are now untied

    :param model: transformers model containing word embeddings
    """
    input_embed, output_embed = get_embeddings(model)
    if input_embed is None or output_embed is None:
        logger.warning(
            "Cannot untie embeddings. If this model has word embeddings, please "
            "implement `get_input_embeddings` and `get_output_embeddings`"
        )
        return

    # clone data to untie
    for module in (input_embed, output_embed):
        if not has_offloaded_params(module):
            data = module.weight.data
        else:
            data = module._hf_hook.weights_map["weight"]

        requires_grad = module.weight.requires_grad
        untied_param = Parameter(data.clone(), requires_grad=requires_grad)
        register_offload_parameter(module, "weight", untied_param)

    # modify model config
    if hasattr(model.config, "tie_word_embeddings"):
        model.config.tie_word_embeddings = False