Transformers documentation

贡献

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

贡献

Transformers 支持多种量化方法,例如 QLoRA、GPTQ、LLM.int8 和 AWQ。不过,目前仍有许多量化方案尚未集成。为了更便捷地在 Transformers 中添加和使用这些量化方法,请使用 HfQuantizer 类。HfQuantizer 被设计为用于添加量化方法的内部辅助类,而不是应用于每个 PyTorch 模块的组件。

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

要求

在将新的量化方法集成到 Transformers 之前,请确保该方法满足以下要求。仅支持与 PyTorch 模块兼容的量化方法。

  • 量化方法可通过一个支持 pip 安装的 Python 包提供(如果只能从源码安装该包,也可以)。理想情况下,pip 包中应包含预编译的 kernel。

  • 该方法可以在常用硬件(CPU、GPU 等)上运行。

  • 该方法被封装在一个 nn.Module(~bitsandbytes.nn.Linear8bitLt、~bitsandbytes.nn.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 类

  1. 最好的起点是参考另一种量化方法,例如 Finegrained Fp8。你总共需要更新或创建三个文件:配置文件、集成文件和 量化器文件。

  2. 在 src/transformers/utils/quantization_config.py 中创建一个新的量化配置类。将新的量化配置添加到 Transformers 的 src/transformers/__init__.py 文件中的 _import_structure 里。

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

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

    • requires_calibration:量化方法是否需要数据校准过程。如果设置为 True,则仅支持推理(使用量化后的权重),而不支持推理和量化。
    • is_serializable:一个属性方法,用于判断该方法是否可序列化。
    • is_trainable:一个属性方法,用于确定是否可以在量化方法之上微调模型(无论是否使用 PEFT 方法)。
  5. 编写 validate_environment 和 update_dtype 方法。这些方法在创建量化模型之前被调用,以确保用户使用了正确的配置。请参考其他量化器了解其实现方式。

  6. 编写 _process_model_before_weight_loading 方法。在 Transformers 中,量化模型会先在 "meta" 设备上初始化,然后再加载权重。这意味着 _process_model_before_weight_loading 方法负责操作模型骨架,将某些模块(nn.Linear)替换为目标模块(量化模块)。

你可以通过在 transformers/src/integrations/ 中创建新文件,并在该文件夹的 __init__.py 文件中暴露相关方法,来定义模块替换逻辑或其他任何实用方法。

  1. 如果量化支持即时量化,请向量化器类添加 get_quantize_ops 方法。在 transformers 中,我们会实例化每个张量,并对其进行一系列不同的操作。在我们的场景中,量化操作发生在最后。你需要创建一个 XXXQuantize,即 ConversionOps 的子类,并添加一个 convert 方法。在 convert 方法中,你需要对权重进行量化,并返回一个包含量化后参数的字典。

  2. 如果量化支持加载预量化权重,请在量化器类中添加 get_weight_conversions 方法。在 transformers 中,我们可以收集多个张量并对它们执行操作。当检查点中的张量需要重新分组以重建量化张量时,这一点尤其有用。

  3. 如有需要,请编写 _process_model_after_weight_loading 方法。此方法可用于实现需要在加载权重后对模型进行额外操作的功能。

  4. 记录一切!请通过在 docs/source/en/quantization 下添加新文件,确保你的量化方法已被记录。

  5. 您应在 docker/transformers-quantization-latest-gpu 中的 nightly Dockerfile 里添加该包,然后在 tests/quantization/xxx 中添加一个新的测试文件。您可以参考现有的量化方法,了解其实现方式。

文件概览

文件目的
utils/quantization_config.py定义继承自 QuantizationConfigMixin 的 YourMethodConfig
quantizers/quantizer_your_method.py实现继承自 HfQuantizer 的 YourMethodHfQuantizer
integrations/your_method.py实现 ConversionOps 的子类和辅助函数
quantizers/auto.py在 AUTO_QUANTIZER_MAPPING 和 AUTO_QUANTIZATION_CONFIG_MAPPING 中注册量化器和配置
docs/source/en/quantization/your_method.md为用户记录使用情况
tests/quantization/your_method/添加集成测试

理解 get_quantize_ops 与 get_weight_conversions

这两种方法用于处理不同的权重加载场景。了解何时使用每种方法至关重要。

get_quantize_ops — 即时量化

在加载未量化的检查点(例如 float16/bfloat16 权重)并在加载过程中进行量化时,请使用此方法。

Checkpoint: model.safetensors (float16 weights for example)
     ↓
get_quantize_ops → YourQuantize.convert()
     ↓
Result: Quantized weights in memory

convert 方法每次接收一个张量,对其进行量化,并可以返回一个包含量化后参数的字典,例如:

class YourQuantize(ConversionOps):
    def convert(self, input_dict, model, full_layer_name, missing_keys, **kwargs):
        # input_dict = {"layer.weight": <float16 tensor>}
        value = list(input_dict.values())[0]
        module, tensor_name = get_module_from_name(model, full_layer_name)

        # Quantize and assign
        quantized, scale, zero_point = your_quantize_fn(value)
        return {full_layer_name: quantized, full_layer_name + ".scale": scale, full_layer_name + ".zero_point": zero_point}

get_weight_conversions — 加载预量化检查点

在加载预量化检查点时使用此功能,其中量化权重以多个独立组件(例如 data、scale 和 zero point)的形式保存,并且需要在加载过程中将这些组件合并为一个张量。并非所有量化方法都需要此重建步骤:例如,FP8 等某些方法会直接加载权重和缩放因子,而无需进行合并。其他方法(如 torchao)则需要从多个已保存的组件中重新组装量化张量。

Checkpoint: model.safetensors (quantized components)
  - layer._weight_qdata
  - layer._weight_scale
  - layer._weight_zero_point
     ↓
get_weight_conversions → WeightConverter + YourDeserialize.convert()
     ↓
Result: Reconstructed quantized tensor → layer.weight

WeightConverter 会根据 source_patterns 收集相关的张量,然后将它们传递给你的 convert 方法:

def get_weight_conversions(self):
    if self.pre_quantized:
        return [
            WeightConverter(
                source_patterns=["_weight_qdata", "_weight_scale", "_weight_zero_point"],
                target_patterns="weight",
                operations=[YourDeserialize(self)],
            ),
        ]
    return []


class YourDeserialize(ConversionOps):
    def convert(self, input_dict, model, full_layer_name, **kwargs):
        # input_dict contains all collected tensors
        # Reconstruct the quantized tensor from components
        reconstructed_tensor = reconstruct_from_components(input_dict)
        return {full_layer_name: reconstructed_tensor}
Update on GitHub