Optimiza Gemma con Hugging Face Transformers y QloRA

En esta guía, se explica cómo ajustar Gemma en un conjunto de datos personalizado de texto a SQL con Transformers y TRL de Hugging Face. Aprenderás a hacer lo siguiente:

  • ¿Qué es la adaptación de bajo rango cuantificada (QLoRA)?
  • Configura el entorno de desarrollo
  • Crea y prepara el conjunto de datos de ajuste
  • Optimiza Gemma con TRL y SFTTrainer
  • Prueba la inferencia de modelos y genera consultas de SQL

¿Qué es la adaptación de bajo rango cuantificada (QLoRA)?

En esta guía, se muestra el uso de la adaptación de clasificación baja cuantificada (QLoRA), que surgió como un método popular para ajustar con precisión los LLM de manera eficiente, ya que reduce los requisitos de recursos computacionales y, al mismo tiempo, mantiene un alto rendimiento. En QloRA, el modelo previamente entrenado se cuantifica a 4 bits y los pesos se congelan. Luego, se adjuntan las capas de adaptador entrenables (LoRA) y solo se entrenan las capas de adaptador. Luego, las ponderaciones del adaptador se pueden combinar con el modelo de base o mantenerse como un adaptador independiente.

Configura el entorno de desarrollo

El primer paso es instalar las bibliotecas de Hugging Face, incluida la TRL, y los conjuntos de datos para ajustar el modelo abierto, incluidas diferentes técnicas de RLHF y alineación.

# 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: Si usas una GPU con arquitectura Ampere (como NVIDIA L4) o una versión posterior, puedes usar la atención de Flash. La atención de Flash es un método que acelera significativamente los cálculos y reduce el uso de la memoria de cuadrático a lineal en la longitud de la secuencia, lo que permite acelerar el entrenamiento hasta 3 veces. Obtén más información en FlashAttention.

Antes de comenzar a entrenar, debes asegurarte de haber aceptado las condiciones de uso de Gemma. Para aceptar la licencia en Hugging Face, haz clic en el botón Aceptar y acceder al repositorio en la página del modelo: http://huggingface.co/google/gemma-3-1b-pt

Después de aceptar la licencia, necesitas un token de Hugging Face válido para acceder al modelo. Si ejecutas el código en Google Colab, puedes usar de forma segura tu token de Hugging Face con los secretos de Colab. De lo contrario, puedes configurar el token directamente en el método login. Asegúrate de que tu token también tenga acceso de escritura cuando envíes tu modelo al Hub durante el entrenamiento.

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 y prepara el conjunto de datos de ajuste

Cuando ajustas los LLM, es importante conocer tu caso de uso y la tarea que quieres resolver. Esto te ayuda a crear un conjunto de datos para ajustar tu modelo. Si aún no definiste tu caso de uso, te recomendamos que vuelvas al principio.

A modo de ejemplo, esta guía se enfoca en el siguiente caso de uso:

  • Afina un modelo de lenguaje natural a SQL para integrarlo sin problemas en una herramienta de análisis de datos. El objetivo es reducir significativamente el tiempo y la experiencia necesarios para la generación de consultas SQL, lo que permite que incluso los usuarios no técnicos extraigan estadísticas significativas de los datos.

Texto a SQL puede ser un buen caso de uso para ajustar los LLM, ya que es una tarea compleja que requiere mucho conocimiento (interno) sobre los datos y el lenguaje SQL.

Una vez que hayas determinado que el ajuste fino es la solución correcta, necesitarás un conjunto de datos para realizarlo. El conjunto de datos debe ser un conjunto diverso de demostraciones de las tareas que deseas resolver. Existen varias formas de crear un conjunto de datos de este tipo, como las siguientes:

  • Usar conjuntos de datos de código abierto existentes, como Spider
  • Usar conjuntos de datos sintéticos creados por LLM, como Alpaca
  • Usar conjuntos de datos creados por humanos, como Dolly
  • Usar una combinación de los métodos, como Orca

Cada uno de los métodos tiene sus propias ventajas y desventajas, y depende de los requisitos de presupuesto, tiempo y calidad. Por ejemplo, usar un conjunto de datos existente es lo más fácil, pero es posible que no se adapte a tu caso de uso específico, mientras que usar expertos en el dominio puede ser lo más preciso, pero puede ser costoso y llevar mucho tiempo. También es posible combinar varios métodos para crear un conjunto de datos de instrucciones, como se muestra en Orca: Aprendizaje progresivo a partir de seguimientos de explicaciones complejas de GPT-4.

En esta guía, se usa un conjunto de datos ya existente (philschmid/gretel-synthetic-text-to-sql), un conjunto de datos sintético de texto a SQL de alta calidad que incluye instrucciones en lenguaje natural, definiciones de esquemas, razonamientos y la consulta de SQL correspondiente.

Hugging Face TRL admite la creación automática de plantillas de formatos de conjuntos de datos de conversaciones. Esto significa que solo debes convertir tu conjunto de datos en los objetos JSON correctos, y trl se encargará de crear plantillas y colocarlas en el formato correcto.

{"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 más de 100,000 muestras. Para que la guía sea pequeña, se reduce la muestra para usar solo 10,000 muestras.

Ahora puedes usar la biblioteca de conjuntos de datos de Hugging Face para cargar el conjunto de datos y crear una plantilla de instrucción para combinar la instrucción de lenguaje natural, la definición del esquema y agregar un mensaje del sistema para tu asistente.

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"])

Optimiza Gemma con TRL y SFTTrainer

Ya puedes ajustar tu modelo. SFTTrainer de TRL de Hugging Face permite supervisar y ajustar con precisión los LLM abiertos de forma sencilla. SFTTrainer es una subclase de Trainer de la biblioteca transformers y admite las mismas funciones, como el registro, la evaluación y los puntos de control, pero agrega funciones adicionales de calidad de vida, como las siguientes:

  • Aplicación de formato a conjuntos de datos, incluidos los formatos de instrucción y de conversación
  • Entrenamiento solo en las acciones de finalización, sin tener en cuenta las instrucciones
  • Empaquetamiento de conjuntos de datos para un entrenamiento más eficiente
  • Compatibilidad con el ajuste eficiente de parámetros (PEFT), incluido QloRA
  • Prepara el modelo y el analizador para el ajuste fino de la conversación (como agregar tokens especiales)

El siguiente código carga el modelo y el analizador de Gemma desde Hugging Face y, luego, inicializa la configuración de cuantificación.

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 admite una integración nativa con peft, lo que facilita la configuración eficiente de los LLM con QLoRA. Solo debes crear un LoraConfig y proporcionarlo al capacitador.

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
)

Antes de comenzar el entrenamiento, debes definir el hiperparámetro que deseas usar en una instancia de 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
    }
)

Ahora tienes todos los componentes básicos que necesitas para crear tu SFTTrainer y comenzar a entrenar el modelo.

from trl import SFTTrainer

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

Para iniciar el entrenamiento, llama al 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()

Antes de que puedas probar tu modelo, asegúrate de liberar la memoria.

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

Cuando usas QLoRA, solo entrenas adaptadores y no el modelo completo. Esto significa que, cuando guardas el modelo durante el entrenamiento, solo guardas las ponderaciones del adaptador y no el modelo completo. Si deseas guardar el modelo completo, lo que facilita su uso con pilas de entrega como vLLM o TGI, puedes combinar los pesos del adaptador con los pesos del modelo con el método merge_and_unload y, luego, guardar el modelo con el método save_pretrained. Esto guarda un modelo predeterminado, que se puede usar para la inferencia.

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

Prueba la inferencia de modelos y genera consultas de SQL

Una vez que se complete el entrenamiento, deberás evaluar y probar tu modelo. Puedes cargar diferentes muestras del conjunto de datos de prueba y evaluar el modelo en esas muestras.

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)

Carguemos una muestra aleatoria del conjunto de datos de prueba y generemos 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()}")

Resumen y próximos pasos

En este instructivo, se analizó cómo ajustar un modelo de Gemma con TRL y QLoRA. Consulta los siguientes documentos: