Ajustar o Gemma usando o Hugging Face Transformers e o QloRA

Ver em ai.google.dev Executar no Google Colab Executar no Kaggle Abrir na Vertex AI Ver código-fonte no GitHub

Neste guia, mostramos como ajustar a Gemma em um conjunto de dados personalizado de texto para SQL usando as bibliotecas Transformers e TRL do Hugging Face. Você vai aprender:

  • O que é a adaptação de classificação baixa quantizada (QLoRA)?
  • Configurar o ambiente de desenvolvimento
  • Criar e preparar o conjunto de dados de ajuste refinado
  • Ajustar o Gemma usando a TRL e o SFTTrainer
  • Testar a inferência de modelo e gerar consultas SQL

O que é a adaptação de classificação baixa quantizada (QLoRA)?

Este guia demonstra o uso da Adaptação quantizada de classificação baixa (QLoRA), que surgiu como um método popular para ajustar LLMs de maneira eficiente, já que reduz os requisitos de recursos computacionais e mantém o alto desempenho. No QLoRA, o modelo pré-treinado é quantizado para 4 bits e os pesos são congelados. Em seguida, as camadas de adaptadores treináveis (LoRA) são anexadas, e apenas elas são treinadas. Depois, os pesos do adaptador podem ser mesclados com o modelo de base ou mantidos como um adaptador separado.

Configurar o ambiente de desenvolvimento

A primeira etapa é instalar as bibliotecas do Hugging Face, incluindo TRL e conjuntos de dados, para ajustar o modelo aberto, incluindo diferentes técnicas de RLHF e alinhamento.

# Install Pytorch & other libraries
%pip install torch tensorboard

# Install Transformers
%pip install transformers

# Install Hugging Face libraries
%pip install datasets accelerate evaluate bitsandbytes trl peft protobuf sentencepiece

# COMMENT IN: if you are running on a GPU that supports BF16 data type and flash attn, such as NVIDIA L4 or NVIDIA A100
#%pip install flash-attn

Observação: se você estiver usando uma GPU com arquitetura Ampere (como a NVIDIA L4) ou mais recente, poderá usar a atenção rápida. O Flash Attention é um método que acelera significativamente os cálculos e reduz o uso da memória de quadrático para linear no comprimento da sequência, acelerando o treinamento em até 3 vezes. Saiba mais em FlashAttention.

Você precisa de um token do Hugging Face válido para publicar seu modelo. Se você estiver executando em um Google Colab, use seu token do Hugging Face com segurança usando os secrets do Colab. Caso contrário, defina o token diretamente no método login. Verifique se o token também tem acesso de gravação, já que você vai enviar o modelo para o Hub durante o treinamento.

# Login into Hugging Face Hub
from huggingface_hub import login
login()

Criar e preparar o conjunto de dados de ajuste refinado

Ao ajustar LLMs, é importante conhecer seu caso de uso e a tarefa que você quer resolver. Isso ajuda a criar um conjunto de dados para ajustar o modelo. Se você ainda não definiu seu caso de uso, talvez seja necessário voltar ao início.

Por exemplo, este guia se concentra no seguinte caso de uso:

  • Ajuste um modelo de linguagem natural para SQL para integração perfeita em uma ferramenta de análise de dados. O objetivo é reduzir significativamente o tempo e a experiência necessários para gerar consultas SQL, permitindo que até mesmo usuários sem conhecimento técnico extraiam insights significativos dos dados.

O Text-to-SQL pode ser um bom caso de uso para ajuste fino de LLMs, já que é uma tarefa complexa que exige muito conhecimento (interno) sobre os dados e a linguagem SQL.

Depois de determinar que o ajuste fino é a solução certa, você precisa de um conjunto de dados para fazer isso. O conjunto de dados precisa ser diversificado e conter demonstrações das tarefas que você quer resolver. Há várias maneiras de criar um conjunto de dados desse tipo, incluindo:

  • Usando conjuntos de dados de código aberto, como o Spider
  • Usar conjuntos de dados sintéticos criados por LLMs, como o Alpaca
  • Usando conjuntos de dados criados por humanos, como o Dolly.
  • Usando uma combinação dos métodos, como Orca

Cada método tem vantagens e desvantagens e depende do orçamento, do tempo e dos requisitos de qualidade. Por exemplo, usar um conjunto de dados existente é a opção mais fácil, mas pode não ser adequada ao seu caso de uso específico. Já usar especialistas no assunto pode ser a opção mais precisa, mas pode levar tempo e ser caro. Também é possível combinar vários métodos para criar um conjunto de dados de instruções, conforme mostrado em Orca: Progressive Learning from Complex Explanation Traces of GPT-4.

Este guia usa um conjunto de dados já existente (philschmid/gretel-synthetic-text-to-sql), um conjunto de dados sintéticos de alta qualidade de texto para SQL que inclui instruções em linguagem natural, definições de esquema, raciocínio e a consulta SQL correspondente.

A TRL do Hugging Face oferece suporte à criação automática de modelos de formatos de conjuntos de dados de conversas. Isso significa que você só precisa converter o conjunto de dados nos objetos JSON corretos, e o trl cuida da criação de modelos e da formatação adequada.

{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}

O philschmid/gretel-synthetic-text-to-sql (link em inglês) contém mais de 100 mil exemplos. Para manter o guia pequeno, ele é reduzido para usar apenas 10.000 exemplos.

Agora você pode usar a biblioteca de conjuntos de dados do Hugging Face para carregar o conjunto de dados e criar um modelo de solicitação que combine a instrução em linguagem natural, a definição de esquema e uma mensagem do sistema para seu assistente.

from datasets import load_dataset

# System message for the assistant
system_message = """You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA."""

# User prompt that combines the user query and the schema
user_prompt = """Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.

<SCHEMA>
{context}
</SCHEMA>

<USER_QUERY>
{question}
</USER_QUERY>
"""
def create_conversation(sample):
  return {
    "messages": [
      {"role": "system", "content": system_message},
      {"role": "user", "content": user_prompt.format(question=sample["sql_prompt"], context=sample["sql_context"])},
      {"role": "assistant", "content": sample["sql"]}
    ]
  }

# Load dataset from the hub
dataset = load_dataset("philschmid/gretel-synthetic-text-to-sql", split="train")
dataset = dataset.shuffle().select(range(12500))

# Convert dataset to OAI messages
dataset = dataset.map(create_conversation, remove_columns=dataset.features,batched=False)
# split dataset into 80% training samples and 20% test samples
dataset = dataset.train_test_split(test_size=0.2)

# Print formatted user prompt
for item in dataset["train"][0]["messages"]:
  print(item)

README.md:   0%|          | 0.00/737 [00:00<?, ?B/s]
synthetic_text_to_sql_train.snappy.parqu(…):   0%|          | 0.00/32.4M [00:00<?, ?B/s]
synthetic_text_to_sql_test.snappy.parque(…):   0%|          | 0.00/1.90M [00:00<?, ?B/s]
Generating train split:   0%|          | 0/100000 [00:00<?, ? examples/s]
Generating test split:   0%|          | 0/5851 [00:00<?, ? examples/s]
Map:   0%|          | 0/12500 [00:00<?, ? examples/s]
{'content': 'You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA.', 'role': 'system'}
{'content': "Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.\n\n<SCHEMA>\nCREATE TABLE Menu (id INT PRIMARY KEY, name VARCHAR(255), category VARCHAR(255), price DECIMAL(5,2));\n</SCHEMA>\n\n<USER_QUERY>\nCalculate the average price of all menu items in the Vegan category\n</USER_QUERY>\n", 'role': 'user'}
{'content': "SELECT AVG(price) FROM Menu WHERE category = 'Vegan';", 'role': 'assistant'}

Ajustar o Gemma usando a TRL e o SFTTrainer

Agora está tudo pronto para ajustar seu modelo. O SFTTrainer do Hugging Face TRL facilita a supervisão do ajuste de LLMs abertos. O SFTTrainer é uma subclasse do Trainer da biblioteca transformers e oferece suporte a todos os mesmos recursos, incluindo registro, avaliação e checkpoint, mas adiciona outros recursos de qualidade de vida, como:

  • Formatação de conjuntos de dados, incluindo formatos de conversa e instrução
  • Treinamento apenas em conclusões, ignorando comandos
  • Compactar conjuntos de dados para um treinamento mais eficiente
  • Suporte para ajuste fino com eficiência de parâmetros (PEFT), incluindo QloRA
  • Preparar o modelo e o tokenizador para ajuste fino conversacional (como adicionar tokens especiais)

O código a seguir carrega o modelo e o tokenizador do Gemma da Hugging Face e inicializa a configuração de quantização.

import torch
from transformers import AutoTokenizer, AutoModelForImageTextToText, BitsAndBytesConfig

# Hugging Face model id
model_id = "google/gemma-4-E2B" # @param ["google/gemma-4-E2B","google/gemma-4-E4B"] {"allow-input":true}

# Check if GPU benefits from bfloat16
if torch.cuda.get_device_capability()[0] >= 8:
    torch_dtype = torch.bfloat16
else:
    torch_dtype = torch.float16

# Define model init arguments
model_kwargs = dict(
    dtype=torch_dtype,
    device_map="auto", # Let torch decide how to load the model
)

# BitsAndBytesConfig: Enables 4-bit quantization to reduce model size/memory usage
model_kwargs["quantization_config"] = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_compute_dtype=model_kwargs['dtype'],
    bnb_4bit_quant_storage=model_kwargs['dtype'],
)

# Load model and tokenizer
model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-E2B-it") # Load the Instruction Tokenizer to use the official Gemma template

config.json: 0.00B [00:00, ?B/s]
model.safetensors:   0%|          | 0.00/10.2G [00:00<?, ?B/s]
Loading weights:   0%|          | 0/2011 [00:00<?, ?it/s]
generation_config.json:   0%|          | 0.00/181 [00:00<?, ?B/s]
config.json: 0.00B [00:00, ?B/s]
tokenizer_config.json: 0.00B [00:00, ?B/s]
tokenizer.json:   0%|          | 0.00/32.2M [00:00<?, ?B/s]
chat_template.jinja: 0.00B [00:00, ?B/s]

O SFTTrainer oferece suporte a uma integração integrada com o peft, o que facilita o ajuste eficiente de LLMs usando a QLoRA. Você só precisa criar um LoraConfig e fornecer ao instrutor.

from peft import LoraConfig

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.05,
    r=16,
    bias="none",
    target_modules="all-linear",
    task_type="CAUSAL_LM",
    modules_to_save=["lm_head", "embed_tokens"], # make sure to save the lm_head and embed_tokens as you train the special tokens
    ensure_weight_tying=True,
)

Antes de iniciar o treinamento, defina o hiperparâmetro que você quer usar em uma instância SFTConfig.

import torch
from trl import SFTConfig

args = SFTConfig(
    output_dir="gemma-text-to-sql",         # directory to save and repository id
    max_length=512,                         # max length for model and packing of the dataset
    num_train_epochs=3,                     # number of training epochs
    per_device_train_batch_size=1,          # batch size per device during training
    optim="adamw_torch_fused",              # use fused adamw optimizer
    logging_steps=10,                       # log every 10 steps
    save_strategy="epoch",                  # save checkpoint every epoch
    eval_strategy="epoch",                  # evaluate checkpoint every epoch
    learning_rate=5e-5,                     # learning rate
    fp16=True if model.dtype == torch.float16 else False,   # use float16 precision
    bf16=True if model.dtype == torch.bfloat16 else False,   # use bfloat16 precision
    max_grad_norm=0.3,                      # max gradient norm based on QLoRA paper
    lr_scheduler_type="constant",           # use constant learning rate scheduler
    push_to_hub=True,                           # push model to hub
    report_to="tensorboard",                # report metrics to tensorboard
    dataset_kwargs={
        "add_special_tokens": False, # Template with special tokens
        "append_concat_token": True, # Add EOS token as separator token between examples
    }
)

Agora você tem todos os elementos básicos necessários para criar seu SFTTrainer e iniciar o treinamento do modelo.

from trl import SFTTrainer

# Create Trainer object
trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    peft_config=peft_config,
    processing_class=tokenizer,
)

Tokenizing train dataset:   0%|          | 0/10000 [00:00<?, ? examples/s]
Tokenizing eval dataset:   0%|          | 0/2500 [00:00<?, ? examples/s]

Comece o treinamento chamando o método train().

# Start training, the model will be automatically saved to the Hub and the output directory
trainer.train()

# Save the final model again to the Hugging Face Hub
trainer.save_model()

The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'eos_token_id': 1, 'bos_token_id': 2, 'pad_token_id': 0}.

Antes de testar o modelo, libere a memória.

# free the memory again
del model
del trainer
torch.cuda.empty_cache()

Ao usar o QLoRA, você treina apenas adaptadores, e não o modelo completo. Isso significa que, ao salvar o modelo durante o treinamento, você salva apenas os pesos do adaptador e não o modelo completo. Se você quiser salvar o modelo completo, o que facilita o uso com stacks de veiculação como vLLM ou TGI, mescle os pesos do adaptador aos pesos do modelo usando o método merge_and_unload e salve o modelo com o método save_pretrained. Isso salva um modelo padrão, que pode ser usado para inferência.

from peft import PeftModel

# Load Model base model
model = AutoModelForImageTextToText.from_pretrained(model_id, low_cpu_mem_usage=True)

# Merge LoRA and base model and save
peft_model = PeftModel.from_pretrained(model, args.output_dir)
merged_model = peft_model.merge_and_unload()
merged_model.save_pretrained("merged_model", safe_serialization=True, max_shard_size="2GB")

processor = AutoTokenizer.from_pretrained(args.output_dir)
processor.save_pretrained("merged_model")

Loading weights:   0%|          | 0/2011 [00:00<?, ?it/s]
Writing model shards:   0%|          | 0/5 [00:00<?, ?it/s]
('merged_model/tokenizer_config.json',
 'merged_model/chat_template.jinja',
 'merged_model/tokenizer.json')

Testar a inferência de modelo e gerar consultas SQL

Depois do treinamento, avalie e teste o modelo. É possível carregar diferentes amostras do conjunto de dados de teste e avaliar o modelo nelas.

import torch
from transformers import pipeline

model_id = "merged_model"

# Load Model with PEFT adapter
model = AutoModelForImageTextToText.from_pretrained(
  model_id,
  device_map="auto",
  dtype="auto",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

Loading weights:   0%|          | 0/2012 [00:00<?, ?it/s]
The tied weights mapping and config for this model specifies to tie model.language_model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints with different values, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning.

Vamos carregar uma amostra aleatória do conjunto de dados de teste e gerar um comando SQL.

from random import randint
import re
from transformers import pipeline, GenerationConfig

config = GenerationConfig.from_pretrained(model_id)
config.max_new_tokens = 256

# Load the model and tokenizer into the pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Load a random sample from the test dataset
rand_idx = randint(0, len(dataset["test"]))
test_sample = dataset["test"][rand_idx]

# Convert as test example into a prompt with the Gemma template
prompt = pipe.tokenizer.apply_chat_template(test_sample["messages"][:2], tokenize=False, add_generation_prompt=True)
print(prompt)

# Generate our SQL query.
outputs = pipe(prompt, generation_config=config)

# Extract the user query and original answer
print(f"Context:\n", re.search(r'<SCHEMA>\n(.*?)\n</SCHEMA>', test_sample['messages'][1]['content'], re.DOTALL).group(1).strip())
print(f"Query:\n", re.search(r'<USER_QUERY>\n(.*?)\n</USER_QUERY>', test_sample['messages'][1]['content'], re.DOTALL).group(1).strip())
print(f"Original Answer:\n{test_sample['messages'][2]['content']}")
print(f"Generated Answer:\n{outputs[0]['generated_text'][len(prompt):].strip()}")

<bos><|turn>system
You are a text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA.<turn|>
<|turn>user
Given the <USER_QUERY> and the <SCHEMA>, generate the corresponding SQL command to retrieve the desired data, considering the query's syntax, semantics, and schema constraints.

<SCHEMA>
CREATE TABLE broadband_plans (plan_id INT, plan_name VARCHAR(255), download_speed INT, upload_speed INT, price DECIMAL(5,2));
</SCHEMA>

<USER_QUERY>
Delete a broadband plan from the 'broadband_plans' table
</USER_QUERY><turn|>
<|turn>model

Context:
 CREATE TABLE broadband_plans (plan_id INT, plan_name VARCHAR(255), download_speed INT, upload_speed INT, price DECIMAL(5,2));
Query:
 Delete a broadband plan from the 'broadband_plans' table
Original Answer:
DELETE FROM broadband_plans WHERE plan_id = 3001;
Generated Answer:
DELETE FROM broadband_plans
WHERE plan_name = 'Basic';

Resumo e próximas etapas

Neste tutorial, mostramos como ajustar um modelo Gemma usando TRL e QLoRA. Confira os seguintes documentos: