Transformers 文档

贡献新的量化方法

贡献新的量化方法

Transformers 支持并集成了许多量化方法,如 QLoRA、GPTQ、LLM.int8 和 AWQ。然而,还有一些量化方法尚未集成。为了使这些量化方法更容易地与 Transformers 模型一起添加和使用,您应该使用 HfQuantizer 类。HfQuantizer 被设计为一个内部辅助类,用于添加量化方法,而不是应用于每个 PyTorch 模块的东西。

本指南将向您展示如何将新的量化方法与HfQuantizer类集成。

需求

在将新的量化方法集成到Transformers之前,请确保您尝试添加的方法满足以下先决条件。目前仅支持可以与PyTorch模块一起运行的量化方法。

  • 量化方法可通过一个Python包获得,任何人都可以通过pip安装(如果只能从源代码安装包也可以)。理想情况下,预编译的内核包含在pip包中。
  • 该方法可以在常用的硬件(CPU、GPU等)上运行。
  • 该方法被封装在一个nn.Module中(例如,Linear8bitLtLinear4bit),量化线性层应具有以下定义:
class Linear4bit(nn.Module):
    def __init__(self, ...):
        ...
    
    def forward(self, x):
        return my_4bit_kernel(x, self.weight, self.bias)

这样,通过用目标类替换一些nn.Linear的实例,可以轻松地对Transformers模型进行量化。

  • 量化方法应该是可序列化的。您可以将量化后的权重保存在本地或推送到Hub。
  • 确保包含量化内核/原语的包是稳定的(没有频繁的重大更改)。

对于一些量化方法,它们可能需要通过数据校准进行“预量化”(例如,AWQ)。在这种情况下,我们更倾向于仅在Transformers中支持推理,并让由机器学习社区维护的第三方库处理模型量化本身。

构建一个新的HFQuantizer类

  1. src/transformers/utils/quantization_config.py中创建一个新的量化配置类,并确保通过将其添加到_import_structure对象中,将新的量化配置暴露在Transformers主init中。

  2. src/transformers/quantizers/目录下创建一个名为quantizer_your_method.py的新文件,并使其继承自src/transformers/quantizers/base.py::HfQuantizer。确保在src/transformers/quantizers/auto.py中的量化自动映射中添加新的量化器和量化配置。

  3. 为您的量化方法定义以下类属性/属性方法:

  • requires_calibration: 量化方法是否需要数据校准过程。如果设置为True,则只能支持推理(使用量化权重),而不能同时支持推理和量化。
  • required_packages: 使用量化权重所需的包列表。你可能需要定义一些新的实用方法,例如在transformers/src/utils/import_utils.py中的is_auto_awq_available
  • requires_parameters_quantization: 仅当您的量化方法需要对底层的nn.Parameter对象进行额外关注时才需要。例如,bitsandbytes使用Params4bitInt8Param,在量化模型时需要一些额外的关注。大多数最近的量化方法将int2/int4权重打包在torch.uint8权重中,因此这个标志通常不需要(默认设置为False)。
  • is_serializable: 一个属性方法,用于确定该方法是否可序列化。
  • is_trainable: 一个属性方法,用于确定是否可以在量化方法的基础上微调模型(使用或不使用PEFT方法)。
  1. 编写validate_environmentupdate_torch_dtype方法。这些方法在创建量化模型之前调用,以确保用户使用正确的配置。您可以参考其他量化器中是如何实现这一点的。

  2. 编写_process_model_before_weight_loading方法。在Transformers中,量化模型首先在"meta"设备上初始化,然后再加载权重。这意味着_process_model_before_weight_loading方法负责操作模型骨架,以替换一些模块(例如nn.Linear)为目标模块(量化模块)。您可以通过在transformers/src/integrations/中创建一个新文件并在该文件夹的__init__.py文件中暴露相关方法来定义模块替换逻辑或任何其他实用方法。最好的起点是查看其他量化方法,例如quantizer_awq.py

  3. 编写_process_model_after_weight_loading方法。此方法允许实现加载权重后需要操作模型的附加功能。

  4. 记录一切!确保你的量化方法被记录,通过在docs/source/en/quantization下添加一个新文件,并在docs/source/en/quantization/overview.md的表格中添加一个新行。

  5. 添加测试!你应该首先在我们的夜间Dockerfile中的docker/transformers-quantization-latest-gpu中添加包,然后在tests/quantization/xxx中添加一个新的测试文件。随时查看其他量化方法的实现方式。

< > Update on GitHub