Ottimizzare Gemma utilizzando Hugging Face Transformers e QloRA

Visualizza su ai.google.dev Esegui in Google Colab Esegui in Kaggle Apri in Vertex AI Visualizza l'origine su GitHub

Questa guida illustra come ottimizzare Gemma su un set di dati personalizzato da testo a SQL utilizzando Hugging Face Transformers e TRL. Scoprirai:

  • Che cos'è l'adattamento a basso ranking quantizzato (QLoRA)
  • Come configurare l'ambiente di sviluppo
  • Come creare e preparare il set di dati di ottimizzazione
  • Come ottimizzare Gemma utilizzando TRL e SFTTrainer
  • Come testare l'inferenza del modello e generare query SQL

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

Questa guida illustra l'utilizzo dell'adattamento a basso ranking quantizzato (QLoRA), che è diventato un metodo diffuso per ottimizzare in modo efficiente gli 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 di adattatore addestrabile (LoRA) e vengono addestrati solo i livelli di adattatore. 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, inclusi TRL e i set di dati, per ottimizzare il modello aperto, incluse le diverse tecniche di RLHF e allineamento.

# 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

Nota: se utilizzi una GPU con architettura Ampere (ad esempio NVIDIA L4) o versioni successive, puoi utilizzare Flash Attention. Flash Attention è un metodo che accelera notevolmente i calcoli e riduce la memoria utilizzata da quadratico a lineare nella lunghezza della sequenza, accelerando l'addestramento fino a 3 volte. Scopri di più su FlashAttention.

Per pubblicare il modello, devi disporre di un token Hugging Face valido. Se esegui l'operazione all'interno di Google Colab, puoi utilizzare in modo sicuro il token Hugging Face utilizzando i secret di Colab. In alternativa, puoi impostare il token direttamente nel metodo login. Assicurati che il token abbia anche accesso in scrittura, poiché il modello viene inviato all'Hub durante l'addestramento.

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

Crea e prepara il set di dati di ottimizzazione

Quando ottimizzi gli LLM, è importante conoscere il caso d'uso e l'attività che vuoi risolvere. In questo modo puoi creare un set di dati per ottimizzare il modello. Se non hai ancora definito il tuo caso d'uso, ti consigliamo di ricominciare da capo.

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

  • Ottimizza un modello da linguaggio naturale a SQL per l'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 esperti di estrarre informazioni significative dai dati.

Da testo a SQL può essere un buon caso d'uso per l'ottimizzazione 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 per l'ottimizzazione. 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:

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

Ognuno dei metodi presenta vantaggi e svantaggi e dipende dal budget, dal tempo e dai requisiti di 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 di dominio potrebbe essere il più accurato, ma può richiedere molto tempo ed essere costoso. È 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 da testo a 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 di formati di set di dati di conversazione. Ciò significa che devi solo convertire il set di dati negli oggetti JSON corretti e trl si occupa della creazione di modelli e della formattazione corretta.

{"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 set di dati philschmid/gretel-synthetic-text-to-sql contiene oltre 100.000 campioni. Per mantenere la guida breve, viene eseguito il downsampling per 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 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'}

Ottimizza Gemma utilizzando TRL e SFTTrainer

Ora puoi ottimizzare il modello. Hugging Face TRL SFTTrainer semplifica l'ottimizzazione supervisionata degli LLM aperti. Il SFTTrainer è una sottoclasse del Trainer della libreria transformers e supporta tutte le stesse funzionalità, tra cui logging, valutazione e checkpoint, ma aggiunge funzionalità di qualità della vita aggiuntive, tra cui:

  • Formattazione del set di dati, inclusi i formati di conversazione e istruzioni
  • Addestramento solo sui completamenti, ignorando i prompt
  • Compressione dei set di dati per un addestramento più efficiente
  • Supporto per l'ottimizzazione efficiente dei parametri (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, 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]

SFTTrainer supporta un'integrazione integrata con peft, che semplifica l'ottimizzazione efficiente degli LLM utilizzando QLoRA. Devi solo creare un LoraConfig e fornirlo al trainer.

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,
)

Prima di poter iniziare l'addestramento, devi definire l'iperparametro che vuoi utilizzare in un'istanza 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
    }
)

Ora hai tutti i componenti di base necessari per creare SFTTrainer per avviare l'addestramento del modello.

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]

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()

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}.

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 il modello completo. Ciò significa che quando salvi il modello durante l'addestramento, salvi solo i pesi dell'adattatore e non il modello completo. Se vuoi salvare il modello completo, che ne semplifica l'utilizzo con gli stack di gestione 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 = 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')

Testa l'inferenza del modello e genera query SQL

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

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.

Carichiamo un campione casuale dal set di dati di test e generiamo un 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';

Riepilogo e passaggi successivi

Questo tutorial ha illustrato come ottimizzare un modello Gemma utilizzando TRL e QLoRA. Consulta i seguenti documenti: