Quantização de modelos

Integração com bitsandbytes

O Accelerate traz a quantização bitsandbytes para o seu modelo. Agora você pode carregar qualquer modelo pytorch em 8 bits ou 4 bits com poucas linhas de código.

Se você deseja usar modelos do Transformers com bitsandbytes, deve seguir esta documentação.

Para saber mais sobre como funciona a quantização do bitsandbytes, confira os posts no blog sobre quantização de 8 bits e quantização de 4 bits.

Pré-requisitos

Você precisará instalar os seguintes requisitos:

• Instale a biblioteca bitsandbytes

pip install bitsandbytes

Para dispositivos que não são cuda, você pode consultar o guia de instalação do bitsandbytes aqui.

• Instale a versão mais recente do accelerate a partir do código-fonte

pip install git+https://github.com/huggingface/accelerate.git

• Instale minGPT e huggingface_hub para executar os exemplos

git clone https://github.com/karpathy/minGPT.git
pip install minGPT/
pip install huggingface_hub

Como funciona

Primeiro, precisamos inicializar nosso modelo. Para economizar memória, podemos inicializar um modelo vazio usando o gerenciador de contexto init_empty_weights().

Vamos pegar o modelo GPT2 da biblioteca minGPT.

from accelerate import init_empty_weights
from mingpt.model import GPT

model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024

with init_empty_weights():
    empty_model = GPT(model_config)

Em seguida, precisamos obter o caminho para os pesos do seu modelo. O caminho pode ser o arquivo state_dict (por exemplo, “pytorch_model.bin”) ou uma pasta contendo os checkpoints fragmentados.

from huggingface_hub import snapshot_download
weights_location = snapshot_download(repo_id="marcsun13/gpt2-xl-linear-sharded")

Por fim, você precisa definir sua configuração de quantização com BnbQuantizationConfig.

Aqui está um exemplo para quantização de 8 bits:

from accelerate.utils import BnbQuantizationConfig
bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, llm_int8_threshold = 6)

Aqui está um exemplo para quantização de 4 bits:

from accelerate.utils import BnbQuantizationConfig
bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")

Para quantizar seu modelo vazio com a configuração selecionada, você precisa usar load_and_quantize_model().

from accelerate.utils import load_and_quantize_model
quantized_model = load_and_quantize_model(empty_model, weights_location=weights_location, bnb_quantization_config=bnb_quantization_config)

Salvando e carregando modelo de 8 bits

Você pode salvar seu modelo de 8 bits com o accelerate usando save_model().

from accelerate import Accelerator
accelerate = Accelerator()
new_weights_location = "path/to/save_directory"
accelerate.save_model(quantized_model, new_weights_location)

quantized_model_from_saved = load_and_quantize_model(empty_model, weights_location=new_weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")

Observe que a serialização de modelos de 4 bits não é suportada no momento.

Deslocar módulos para a CPU e o disco

Você pode descarregar alguns módulos para a CPU/disco se não tiver espaço suficiente na GPU para armazenar o modelo inteiro em suas GPUs. Isso usa inferência de modelos grandes por baixo dos panos. Consulte esta documentação para mais detalhes.

Para a quantização de 8 bits, os módulos selecionados serão convertidos para precisão de 8 bits.

Para a quantização de 4 bits, os módulos selecionados serão mantidos no torch_dtype que o usuário passou em BnbQuantizationConfig. Adicionaremos suporte para converter esses módulos descarregados em 4 bits quando a serialização de 4 bits for possível.

Você só precisa passar um device_map personalizado para descarregar módulos na cpu/disco. Os módulos descarregados serão enviados para a GPU quando necessário. Aqui está um exemplo:

device_map = {
    "transformer.wte": 0,
    "transformer.wpe": 0,
    "transformer.drop": 0,
    "transformer.h": "cpu",
    "transformer.ln_f": "disk",
    "lm_head": "disk",
}

Ajuste fino de um modelo quantizado

Não é possível realizar treinamento puro em 8bit ou 4bit nesses modelos. No entanto, você pode treinar esses modelos aproveitando métodos de ajuste fino eficientes em parâmetros (PEFT) e, por exemplo, treinar adaptadores sobre eles. Consulte a biblioteca peft para mais detalhes.

Atualmente, não é possível adicionar adaptadores sobre qualquer modelo quantizado. No entanto, com o suporte oficial a adaptadores em modelos do Transformers, é possível fazer o ajuste fino de modelos quantizados. Se você deseja fazer o ajuste fino de um modelo do Transformers, siga esta documentação em vez disso. Confira este demo sobre como fazer o ajuste fino de um modelo do Transformers de 4 bits.

Observe que você não precisa passar device_map ao carregar o modelo para treinamento. Ele carregará automaticamente seu modelo na sua GPU. Lembre-se de que device_map=auto deve ser usado apenas para inferência.

Exemplo de demonstração - executando GPT2 1.5b em um Google Colab

Confira a demonstração no Google Colab para executar modelos quantizados em um modelo GPT2. O checkpoint do modelo GPT2-1.5B está em FP32, o que utiliza 6GB de memória. Após a quantização, ele utiliza 1.6GB com módulos de 8 bits e 1.2GB com módulos de 4 bits.