양자화 방법 기여하기

양자화 방법 기여하기 (Contribute)

Transformers는 QLoRA, GPTQ, LLM.int8, AWQ 같은 많은 양자화 방법을 지원합니다. 하지만 아직 통합되지 않은 양자화 접근 방식도 아직 훨씬 더 많아요. 이 양자화 방법들을 Transformers에 더 쉽게 추가하고 사용하게 하려면 HfQuantizer 클래스를 사용하세요. HfQuantizer는 모든 PyTorch 모듈에 적용하는 것 대신 양자화 방법을 추가하기 위한 내부 헬퍼 클래스로 설계되었습니다.

출처: 문서

본문

이 가이드는 HfQuantizer로 새 양자화 방법을 통합하는 방법을 보여줍니다.

요구 사항 (Requirements)

새 양자화 방법을 Transformers에 통합하기 전에 그 방법이 다음 요구 사항을 충족하는지 확인하세요. PyTorch 모듈로 실행할 수 있는 양자화 방법만 지원됩니다.

  • 양자화 방법이 pip으로 설치 가능한 Python 패키지로 제공되어야 합니다(패키지를 소스에서만 설치할 수 있어도 괜찮아요). 이상적으로는 사전 컴파일된 커널이 pip 패키지에 포함되어 있어야 합니다.

  • 방법이 흔히 사용되는 하드웨어(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에 올릴 수 있어야 해요.

  • 양자화 커널/프리미티브를 담은 패키지가 안정적인지(빈번한 breaking change가 없는지) 확인하세요.

일부 양자화 방법은 데이터 보정(calibration)을 통한 모델 "사전 양자화"(AWQ)가 필요할 수 있습니다. 이 경우 Transformers에서는 추론만 지원하고, ML 커뮤니티가 유지관리하는 서드파티 라이브러리가 모델 양자화 자체를 처리하게 두는 것을 선호합니다.

새 HFQuantizer 클래스 만들기 (Create new HFQuantizer class)

  1. 가장 좋은 출발점은 Finegrained Fp8 같은 다른 양자화 방법을 살펴보는 것입니다. 총 세 개의 파일을 업데이트하거나 만들어야 해요: config 파일, integration 파일, quantizer 파일입니다.

  2. src/transformers/utils/quantization_config.py 안에 새 양자화 config 클래스를 만듭니다. 새 양자화 config를 Transformers의 src/transformers/init.py 파일 안의 _import_structure에 추가하세요.

  3. src/transformers/quantizers/ 안에 quantizer_your_method.py라는 새 파일을 만들고, [`~quantizers.HfQuantizer]를 상속하게 하세요. src/transformers/quantizers/auto.py의 양자화 자동 매핑에 새 quantizer와 양자화 config를 꼭 추가하세요.

  4. 양자화 방법에 대해 다음 클래스 속성과 프로퍼티 메서드를 정의합니다.

    • requires_calibration: 양자화 방법이 데이터 보정 과정을 요구하는지. True로 설정하면 추론(양자화된 가중치로)만 지원하고 추론과 양자화를 모두 지원하지는 못합니다.
    • is_serializable: 방법이 직렬화 가능한지 결정하는 프로퍼티 메서드.
    • is_trainable: 양자화 방법 위에서(PEFT 접근 방식 포함) 모델을 미세 조정할 수 있는지 결정하는 프로퍼티 메서드.
  5. validate_environment와 update_dtype 메서드를 작성합니다. 이 메서드들은 양자화된 모델을 만들기 전에 호출되어 사용자가 올바른 설정을 쓰도록 보장합니다. 어떻게 구현하는지에 대한 예시는 다른 quantizer를 참고하세요.

  6. _process_model_before_weight_loading 메서드를 작성합니다. Transformers에서 양자화된 모델은 가중치를 로드하기 전에 먼저 "meta" 기기에서 초기화됩니다. 즉 _process_model_before_weight_loading 메서드가 모델 스켈레톤을 조작해 일부 모듈(nn.Linear)을 대상 모듈(양자화 모듈)로 교체하는 역할을 합니다.

transformers/src/integrations/에 새 파일을 만들고 그 폴더의 __init__.py 파일에서 관련 메서드를 노출하면, 모듈 교체 로직이나 다른 유틸리티 메서드를 정의할 수 있어요.

  1. 양자화가 그 자리에서 양자화를 지원한다면 quantizer 클래스에 get_quantize_ops 메서드를 추가합니다. transformers에서는 각 텐서를 구체화(materialize)하고 그 위에 일련의 서로 다른 연산을 적용합니다. 우리의 경우 양자화 연산은 마지막에 일어나요. ConversionOps의 하위 클래스인 XXXQuantize를 만들고 convert 메서드를 추가해야 합니다. convert 메서드에서 가중치를 양자화하고 양자화된 매개변수의 딕셔너리를 반환해야 해요.

  2. 양자화가 사전 양자화된 가중치 로딩을 지원한다면 quantizer 클래스에 get_weight_conversions 메서드를 추가합니다. transformers에서는 여러 텐서를 모아 그 위에 연산을 적용할 수 있습니다. 체크포인트에 양자화된 텐서를 다시 만들기 위해 재그룹화가 필요한 텐서가 있을 때 특히 유용합니다.

  3. 필요하다면 _process_model_after_weight_loading 메서드를 작성합니다. 이 메서드는 가중치를 로드한 뒤 모델을 조작해야 하는 추가 기능을 구현할 수 있게 해줍니다.

  4. 모든 것을 문서화하세요! docs/source/en/quantization 아래에 새 파일을 추가해 양자화 방법이 문서화되었는지 확인하세요.

  5. 테스트를 추가해야 합니다. docker/transformers-quantization-latest-gpu 안의 나이틀리 Dockerfile에 패키지를 추가하고, tests/quantization/xxx에 새 테스트 파일을 추가하세요. 기존 양자화 방법이 어떻게 구현되었는지 자유롭게 살펴보세요.

파일 개요 (Files overview)

파일 용도
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에 quantizer와 config를 등록
docs/source/en/quantization/your_method.md 사용자를 위한 사용법 문서
tests/quantization/your_method/ 통합 테스트 추가

get_quantize_ops와 get_weight_conversions 이해하기 (Understanding get_quantize_ops vs get_weight_conversions)

이 두 메서드는 가중치 로딩의 서로 다른 시나리오를 다룹니다. 각각 언제 쓰는지 이해하는 것이 필수적이에요.

get_quantize_ops — 그 자리에서 양자화 (Quantize on the fly)

양자화되지 않은 체크포인트(예: 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 — 사전 양자화된 체크포인트 로드 (Load pre-quantized checkpoints)

사전 양자화된 체크포인트를 로드할 때 사용합니다. 양자화된 가중치가 (데이터, 스케일, 제로 포인트 같은) 여러 개의 분리된 구성요소로 저장되어 있고, 로딩 중에 하나의 텐서로 결합되어야 할 때 사용해요. 모든 양자화 방법이 이 재구성 단계를 요구하는 것은 아닙니다. 예를 들어 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}

더 알아보기 (Learn more)