贡献新的量化方法
Transformers 支持并集成了许多量化方法,如 QLoRA、GPTQ、LLM.int8 和 AWQ。然而,还有一些量化方法尚未集成。为了使这些量化方法更容易地与 Transformers 模型一起添加和使用,您应该使用 HfQuantizer
类。HfQuantizer
被设计为一个内部辅助类,用于添加量化方法,而不是应用于每个 PyTorch 模块的东西。
本指南将向您展示如何将新的量化方法与HfQuantizer
类集成。
需求
在将新的量化方法集成到Transformers之前,请确保您尝试添加的方法满足以下先决条件。目前仅支持可以与PyTorch模块一起运行的量化方法。
- 量化方法可通过一个Python包获得,任何人都可以通过pip安装(如果只能从源代码安装包也可以)。理想情况下,预编译的内核包含在pip包中。
- 该方法可以在常用的硬件(CPU、GPU等)上运行。
- 该方法被封装在一个
nn.Module
中(例如,Linear8bitLt
,Linear4bit
),量化线性层应具有以下定义:
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类
在src/transformers/utils/quantization_config.py中创建一个新的量化配置类,并确保通过将其添加到
_import_structure
对象中,将新的量化配置暴露在Transformers主init
中。在src/transformers/quantizers/目录下创建一个名为
quantizer_your_method.py
的新文件,并使其继承自src/transformers/quantizers/base.py::HfQuantizer。确保在src/transformers/quantizers/auto.py中的量化自动映射中添加新的量化器和量化配置。为您的量化方法定义以下类属性/属性方法:
requires_calibration
: 量化方法是否需要数据校准过程。如果设置为True
,则只能支持推理(使用量化权重),而不能同时支持推理和量化。required_packages
: 使用量化权重所需的包列表。你可能需要定义一些新的实用方法,例如在transformers/src/utils/import_utils.py中的is_auto_awq_available
。requires_parameters_quantization
: 仅当您的量化方法需要对底层的nn.Parameter
对象进行额外关注时才需要。例如,bitsandbytes使用Params4bit
和Int8Param
,在量化模型时需要一些额外的关注。大多数最近的量化方法将int2/int4权重打包在torch.uint8
权重中,因此这个标志通常不需要(默认设置为False
)。is_serializable
: 一个属性方法,用于确定该方法是否可序列化。is_trainable
: 一个属性方法,用于确定是否可以在量化方法的基础上微调模型(使用或不使用PEFT方法)。
编写
validate_environment
和update_torch_dtype
方法。这些方法在创建量化模型之前调用,以确保用户使用正确的配置。您可以参考其他量化器中是如何实现这一点的。编写
_process_model_before_weight_loading
方法。在Transformers中,量化模型首先在"meta"
设备上初始化,然后再加载权重。这意味着_process_model_before_weight_loading
方法负责操作模型骨架,以替换一些模块(例如nn.Linear
)为目标模块(量化模块)。您可以通过在transformers/src/integrations/中创建一个新文件并在该文件夹的__init__.py
文件中暴露相关方法来定义模块替换逻辑或任何其他实用方法。最好的起点是查看其他量化方法,例如quantizer_awq.py。编写
_process_model_after_weight_loading
方法。此方法允许实现加载权重后需要操作模型的附加功能。记录一切!确保你的量化方法被记录,通过在
docs/source/en/quantization
下添加一个新文件,并在docs/source/en/quantization/overview.md
的表格中添加一个新行。添加测试!你应该首先在我们的夜间Dockerfile中的
docker/transformers-quantization-latest-gpu
中添加包,然后在tests/quantization/xxx
中添加一个新的测试文件。随时查看其他量化方法的实现方式。