跳到内容

llmcompressor.modifiers.pruning.sparsegpt.sgpt_base

类

SparsityModifierBase

Bases: Modifier

实现与一次性稀疏性相关功能的抽象基类。继承者必须实现 calibrate_module 和 compress_modules。

方法

  • on_initialize –

    初始化并在当前状态下运行 SparseGPT 算法。

on_initialize

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

初始化并在当前状态下运行 SparseGPT 算法。

参数

  • state

    (State) –

    存储输入模型和校准数据的会话状态

源代码位于 llmcompressor/modifiers/pruning/sparsegpt/sgpt_base.py。 
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
def on_initialize(self, state: "State", **kwargs) -> bool:
    """
    Initialize and run the SparseGPT algorithm on the current state

    :param state: session state storing input model and calibration data
    """
    model: torch.nn.Module = state.model
    dataloader: torch.utils.data.DataLoader = state.data.calib

    # infer module and sequential targets
    self.sequential_targets = self._infer_sequential_targets(model)
    layers = get_layers(self.sequential_targets, model)
    self._target_layers = get_layers(
        self.targets, model
    )  # layers containing targets

    # infer layer sparsities
    if self.sparsity_profile == "owl":
        logger.info(
            "Using OWL to infer target layer-wise sparsities from "
            f"{len(dataloader) if dataloader else 0} calibration samples..."
        )
        self.sparsity = self._infer_owl_layer_sparsity(model, layers, dataloader)

    # get layers and validate sparsity
    if isinstance(self.sparsity, (list, dict)) and len(self._target_layers) != len(
        self.sparsity
    ):
        raise ValueError(
            f"{self.__repr_name__} was initialized with {len(self.sparsity)} "
            f"sparsities values, but model has {len(layers)} target layers"
        )

    return True