Ottimizzare Gemma utilizzando Hugging Face Transformers e QloRA

Questa guida illustra come perfezionare Gemma su un set di dati text-to-sql personalizzato utilizzando Transformers di Hugging Face e TRL. Scoprirai:

  • Che cos'è l'adattamento a basso ranking quantizzato (QLoRA)
  • Configura l'ambiente di sviluppo
  • Crea e prepara il set di dati per la messa a punto fine
  • Ottimizzare Gemma utilizzando TRL e SFTTrainer
  • Testa l'inferenza del modello e genera query SQL

Che cos'è l'adattamento a basso ranking quantizzato (QLoRA)

Questa guida illustra l'utilizzo dell'adattamento a rango basso quantizzato (QLoRA), che si è affermato come metodo popolare per ottimizzare in modo efficiente gli LLM in quanto riduce i requisiti delle risorse di calcolo mantenendo un'elevata efficienza. In QloRA, il modello preaddestrato viene quantizzato a 4 bit e i pesi vengono bloccati. Vengono quindi collegati i livelli di adattamento addestrabili (LoRA) e vengono addestrati solo questi livelli. In seguito, i pesi dell'adattatore possono essere uniti al modello di base o mantenuti come adattatore separato.

Configura l'ambiente di sviluppo

Il primo passaggio consiste nell'installare le librerie Hugging Face, tra cui TRL, e i set di dati per ottimizzare il modello aperto, incluse diverse tecniche di RLHF e allineamento.

# Install Pytorch & other libraries
%pip install "torch>=2.4.0" tensorboard

# Install Gemma release branch from Hugging Face
%pip install git+https://github.com/huggingface/transformers@v4.49.0-Gemma-3

# Install Hugging Face libraries
%pip install  --upgrade \
  "datasets==3.3.2" \
  "accelerate==1.4.0" \
  "evaluate==0.4.3" \
  "bitsandbytes==0.45.3" \
  "trl==0.15.2" \
  "peft==0.14.0" \
  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

Nota: se utilizzi una GPU con architettura Ampere (ad esempio NVIDIA L4) o successiva, puoi utilizzare la funzionalità Flash di attenzione. Flash Attention è un metodo che velocizza notevolmente i calcoli e riduce l'utilizzo della memoria da quadratico a lineare in base alla lunghezza della sequenza, accelerando l'addestramento fino a tre volte. Scopri di più su FlashAttention.

Prima di poter iniziare la formazione, devi assicurarti di aver accettato i termini di utilizzo di Gemma. Puoi accettare la licenza su Hugging Face facendo clic sul pulsante Accetta e accedi al repository nella pagina del modello all'indirizzo: http://huggingface.co/google/gemma-3-1b-pt

Dopo aver accettato la licenza, devi disporre di un token Hugging Face valido per accedere al modello. Se esegui il codice in Google Colab, puoi utilizzare in sicurezza il token Hugging Face utilizzando i secret di Colab, altrimenti puoi impostare il token direttamente nel metodo login. Assicurati che il token abbia anche l'accesso in scrittura, quando carichi il modello nell'hub durante l'addestramento.

from google.colab import userdata
from huggingface_hub import login

# Login into Hugging Face Hub
hf_token = userdata.get('HF_TOKEN') # If you are running inside a Google Colab 
login(hf_token)

Crea e prepara il set di dati per la messa a punto fine

Quando perfezioni gli LLM, è importante conoscere il caso d'uso e l'attività che vuoi risolvere. In questo modo puoi creare un set di dati per perfezionare il modello. Se non hai ancora definito il tuo caso d'uso, ti consigliamo di tornare alla fase di progettazione.

Ad esempio, questa guida si concentra sul seguente caso d'uso:

  • Ottimizza un modello di linguaggio naturale per SQL per un'integrazione perfetta in uno strumento di analisi dei dati. L'obiettivo è ridurre in modo significativo il tempo e le competenze necessarie per la generazione di query SQL, consentendo anche agli utenti non tecnici di estrarre informazioni significative dai dati.

La conversione di testo in SQL può essere un buon caso d'uso per l'ottimizzazione dei modelli LLM, in quanto si tratta di un'attività complessa che richiede molte conoscenze (interne) sui dati e sul linguaggio SQL.

Una volta stabilito che l'ottimizzazione fine è la soluzione giusta, devi avere un set di dati da ottimizzare. Il set di dati deve essere un insieme diversificato di dimostrazioni delle attività che vuoi risolvere. Esistono diversi modi per creare un set di dati di questo tipo, tra cui:

  • Utilizzo di set di dati open source esistenti, come Spider
  • Utilizzo di set di dati sintetici creati da LLM, come Alpaca
  • Utilizzo di set di dati creati da persone, come Dolly.
  • Utilizzare una combinazione di metodi, ad esempio Orca

Ciascun metodo presenta vantaggi e svantaggi e dipende dai requisiti di budget, tempo e qualità. Ad esempio, l'utilizzo di un set di dati esistente è il più semplice, ma potrebbe non essere personalizzato per il tuo caso d'uso specifico, mentre l'utilizzo di esperti del dominio potrebbe essere il più accurato, ma può essere dispendioso in termini di tempo e denaro. È anche possibile combinare diversi metodi per creare un set di dati di istruzioni, come mostrato in Orca: Progressive Learning from Complex Explanation Traces of GPT-4.

Questa guida utilizza un set di dati già esistente (philschmid/gretel-synthetic-text-to-sql), un set di dati Text-to-SQL sintetico di alta qualità che include istruzioni in linguaggio naturale, definizioni dello schema, ragionamento e la query SQL corrispondente.

Hugging Face TRL supporta la creazione automatica di modelli per i formati dei set di dati di conversazione. Ciò significa che devi solo convertire il set di dati negli oggetti JSON corretti e trl si occupa di creare il modello e di impostarlo nel formato corretto.

{"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": "..."}]}

philschmid/gretel-synthetic-text-to-sql contiene oltre 100.000 esempi. Per mantenere la guida di piccole dimensioni, viene sottocampionata in modo da utilizzare solo 10.000 campioni.

Ora puoi utilizzare la libreria Hugging Face Datasets per caricare il set di dati e creare un modello di prompt per combinare l'istruzione in linguaggio naturale, la definizione dello schema e aggiungere un messaggio di sistema per l'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 10,000 training samples and 2,500 test samples
dataset = dataset.train_test_split(test_size=2500/12500)

# Print formatted user prompt
print(dataset["train"][345]["messages"][1]["content"])

Ottimizzare Gemma utilizzando TRL e SFTTrainer

Ora è tutto pronto per ottimizzare il modello. Hugging Face TRL SFTTrainer semplifica la supervisione della messa a punto di LLM aperti. SFTTrainer è una sottoclasse di Trainer della libreria transformers e supporta tutte le stesse funzionalità, tra cui logging, valutazione e checkpoint, ma aggiunge ulteriori funzionalità di qualità della vita, tra cui:

  • Formattazione del set di dati, inclusi i formati di conversazione e di istruzioni
  • Addestramento solo sui completamenti, ignorando i prompt
  • Imballaggio dei set di dati per un addestramento più efficiente
  • Supporto dell'ottimizzazione fine efficiente dei parametri (PEFT), incluso QloRA
  • Preparazione del modello e del tokenizzatore per l'ottimizzazione di conversazione (ad esempio l'aggiunta di token speciali)

Il seguente codice carica il modello e lo tokenizer Gemma da Hugging Face e inizializza la configurazione della quantizzazione.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForImageTextToText, BitsAndBytesConfig

# Hugging Face model id
model_id = "google/gemma-3-1b-pt" # or `google/gemma-3-4b-pt`, `google/gemma-3-12b-pt`, `google/gemma-3-27b-pt`

# Select model class based on id
if model_id == "google/gemma-3-1b-pt":
    model_class = AutoModelForCausalLM
else:
    model_class = AutoModelForImageTextToText

# 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(
    attn_implementation="eager", # Use "flash_attention_2" when running on Ampere or newer GPU
    torch_dtype=torch_dtype, # What torch dtype to use, defaults to auto
    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['torch_dtype'],
    bnb_4bit_quant_storage=model_kwargs['torch_dtype'],
)

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

SFTTrainer supporta un'integrazione nativa con peft, che semplifica la messa a punto efficiente degli LLM utilizzando QLoRA. Devi solo creare un LoraConfig e fornirlo all'istruttore.

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
)

Prima di poter iniziare l'addestramento, devi definire l'iperparametro da utilizzare in un'istanza SFTConfig.

from trl import SFTConfig

args = SFTConfig(
    output_dir="gemma-text-to-sql",         # directory to save and repository id
    max_seq_length=512,                     # max sequence length for model and packing of the dataset
    packing=True,                           # Groups multiple samples in the dataset into a single sequence
    num_train_epochs=3,                     # number of training epochs
    per_device_train_batch_size=1,          # batch size per device during training
    gradient_accumulation_steps=4,          # number of steps before performing a backward/update pass
    gradient_checkpointing=True,            # use gradient checkpointing to save memory
    optim="adamw_torch_fused",              # use fused adamw optimizer
    logging_steps=10,                       # log every 10 steps
    save_strategy="epoch",                  # save checkpoint every epoch
    learning_rate=2e-4,                     # learning rate, based on QLoRA paper
    fp16=True if torch_dtype == torch.float16 else False,   # use float16 precision
    bf16=True if torch_dtype == torch.bfloat16 else False,   # use bfloat16 precision
    max_grad_norm=0.3,                      # max gradient norm based on QLoRA paper
    warmup_ratio=0.03,                      # warmup ratio 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, # We template with special tokens
        "append_concat_token": True, # Add EOS token as separator token between examples
    }
)

Ora hai tutti i componenti di base necessari per creare il tuo SFTTrainer e iniziare l'addestramento del modello.

from trl import SFTTrainer

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

Avvia l'addestramento chiamando il metodo 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()

Prima di poter testare il modello, assicurati di liberare la memoria.

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

Quando utilizzi QLoRA, addestrini solo gli adattatori e non il modello completo. Ciò significa che, quando salvi il modello durante l'addestramento, vengono salvati solo i pesi dell'adattatore e non il modello completo. Se vuoi salvare il modello completo, che è più facile da utilizzare con gli stack di pubblicazione come vLLM o TGI, puoi unire i pesi dell'adattatore ai pesi del modello utilizzando il metodo merge_and_unload e poi salvare il modello con il metodo save_pretrained. Viene salvato un modello predefinito, che può essere utilizzato per l'inferenza.

from peft import PeftModel

# Load Model base model
model = model_class.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")

Testa l'inferenza del modello e genera query SQL

Al termine dell'addestramento, valuta e testa il modello. Puoi caricare diversi esempi dal set di dati di test e valutare il modello su questi esempi.

import torch
from transformers import pipeline

model_id = "gemma-text-to-sql"

# Load Model with PEFT adapter
model = model_class.from_pretrained(
  model_id,
  device_map="auto",
  torch_dtype=torch_dtype,
  attn_implementation="eager",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

Carica un campione casuale dal set di dati di test e genera un comando SQL.

from random import randint
import re

# 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
stop_token_ids = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<end_of_turn>")]
prompt = pipe.tokenizer.apply_chat_template(test_sample["messages"][:2], tokenize=False, add_generation_prompt=True)

# Generate our SQL query.
outputs = pipe(prompt, max_new_tokens=256, do_sample=False, temperature=0.1, top_k=50, top_p=0.1, eos_token_id=stop_token_ids, disable_compile=True)

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

Riepilogo e passaggi successivi

Questo tutorial spiega come perfezionare un modello Gemma utilizzando TRL e QLoRA. Consulta poi la seguente documentazione: