Ottimizzare Gemma utilizzando Hugging Face Transformers e QloRA

Questa guida ti mostra come eseguire il fine-tuning di Gemma su un set di dati personalizzato di conversione del testo in SQL utilizzando Transformers e TRL di Hugging Face. Scoprirai:

  • Che cos'è l'adattamento a basso ranking quantizzato (QLoRA)?
  • Configura l'ambiente di sviluppo
  • Crea e prepara il set di dati di fine tuning
  • Ottimizzare Gemma utilizzando TRL e SFTTrainer
  • Testare l'inferenza del modello e generare query SQL

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

Questa guida illustra l'utilizzo di Quantized Low-Rank Adaptation (QLoRA), emersa come metodo popolare per ottimizzare in modo efficiente i LLM, in quanto riduce i requisiti di risorse di calcolo mantenendo prestazioni elevate. In QLoRA, il modello preaddestrato viene quantizzato a 4 bit e i pesi vengono bloccati. Vengono quindi collegati i livelli dell'adattatore addestrabile (LoRA) e vengono addestrati solo questi livelli. Successivamente, 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 "transformers>=4.51.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.21.0" \
  "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 più recente, puoi utilizzare Flash Attention. Flash Attention è un metodo che accelera 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 l'addestramento, 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, è necessario un token Hugging Face valido per accedere al modello. Se esegui l'operazione all'interno di Google Colab, puoi utilizzare in modo sicuro il token Hugging Face utilizzando i segreti di Colab. In alternativa, puoi impostare il token direttamente nel metodo login. Assicurati che il token abbia anche accesso in scrittura, in quanto il modello viene inviato all'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 di fine tuning

Quando esegui il fine-tuning dei 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:

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

Text-to-SQL può essere un buon caso d'uso per il fine-tuning degli LLM, in quanto è un'attività complessa che richiede molte conoscenze (interne) sui dati e sul linguaggio SQL.

Una volta stabilito che l'ottimizzazione è la soluzione giusta, hai bisogno di 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:

  • Utilizzando 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.
  • Utilizzando 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 adatto al tuo caso d'uso specifico, mentre l'utilizzo di esperti del settore potrebbe essere il più accurato, ma può richiedere 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 sintetici Text-to-SQL di alta qualità che include istruzioni in linguaggio naturale, definizioni di schema, ragionamento e la query SQL corrispondente.

Hugging Face TRL supporta la creazione automatica di modelli per i formati dei set di dati delle conversazioni. Ciò significa che devi solo convertire il set di dati negli oggetti JSON corretti e trl si occupa di creare i modelli e di inserirli 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": "..."}]}

Il repository philschmid/gretel-synthetic-text-to-sql contiene oltre 100.000 campioni. Per ridurre le dimensioni della guida, viene eseguito il sottocampionamento 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 puoi eseguire il fine tuning del modello. Hugging Face TRL SFTTrainer semplifica la messa a punto supervisionata 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 conversazionali e delle istruzioni
  • Addestramento solo sui completamenti, ignorando i prompt
  • Compressione dei set di dati per un addestramento più efficiente
  • Supporto della parameter-efficient fine-tuning (PEFT), incluso QLoRA
  • Preparazione del modello e del tokenizer per l'ottimizzazione conversazionale (ad esempio l'aggiunta di token speciali)

Il seguente codice carica il modello Gemma e il tokenizer 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, il che semplifica l'ottimizzazione efficiente degli LLM utilizzando QLoRA. Devi solo creare un LoraConfig e fornirlo al formatore.

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 che vuoi utilizzare in un'istanza SFTConfig.

from trl import SFTConfig

args = SFTConfig(
    output_dir="gemma-text-to-sql",         # directory to save and repository id
    max_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 gli elementi 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
)

Inizia 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, addestri solo gli adattatori e non l'intero modello. Ciò significa che quando salvi il modello durante l'addestramento, salvi solo i pesi dell'adattatore e non l'intero modello. Se vuoi salvare l'intero modello, il che ne facilita l'utilizzo con 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. In questo modo 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")

Testare l'inferenza del modello e generare query SQL

Al termine dell'addestramento, dovrai valutare e testare il modello. Puoi caricare campioni diversi dal set di dati di test e valutare il modello in base a questi campioni.

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)

Carichiamo un campione casuale dal set di dati di test e generiamo 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"])-1)
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 ha illustrato come ottimizzare un modello Gemma utilizzando TRL e QLoRA. Consulta i seguenti documenti: