course/pr_1213/es/chapter12/6.md

Ejercicio práctico: GRPO con Unsloth

En este ejercicio, ajustarás un modelo con GRPO (Group Relative Policy Optimization) usando Unsloth para mejorar sus capacidades de razonamiento. Vimos GRPO en el Capítulo 3.

Unsloth es una librería que acelera el ajuste fino de LLM, lo que permite entrenar modelos más rápido y con menos recursos computacionales. Unsloth se integra con TRL, así que construiremos sobre lo visto en las secciones anteriores y lo adaptaremos a las particularidades de Unsloth.

Este ejercicio puede ejecutarse en una GPU T4 gratuita de Google Colab. Para obtener la mejor experiencia, sigue el notebook enlazado arriba y pruébalo tú mismo.

Instala las dependencias

pip install unsloth vllm
pip install --upgrade pillow

Configura Unsloth

Unsloth proporciona una clase, FastLanguageModel, que integra transformers con las optimizaciones de Unsloth. Importémosla:

from unsloth import FastLanguageModel

Ahora carguemos el modelo Gemma 3 1B Instruct de Google y configúralo para el ajuste:

from unsloth import FastLanguageModel
import torch

max_seq_length = 1024  # Can increase for longer reasoning traces
lora_rank = 32  # Larger rank = smarter, but slower

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="google/gemma-3-1b-it",
    max_seq_length=max_seq_length,
    load_in_4bit=True,  # False for LoRA 16bit
    fast_inference=True,  # Enable vLLM fast inference
    max_lora_rank=lora_rank,
    gpu_memory_utilization=0.6,  # Reduce if out of memory
)

model = FastLanguageModel.get_peft_model(
    model,
    r=lora_rank,
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
    lora_alpha=lora_rank,
    use_gradient_checkpointing="unsloth",
    random_state=3407,
)

Este código carga el modelo con cuantización de 4 bits para ahorrar memoria y aplica LoRA para un ajuste eficiente.

No cubriremos los detalles de LoRA en este capítulo, pero puedes aprender más en el Capítulo 11.

Preparación de datos

Para este ejercicio usaremos el dataset GSM8K, que contiene problemas de matemáticas de primaria. Daremos formato a los datos para animar al modelo a mostrar su razonamiento antes de dar una respuesta.

Primero definiremos el formato de los prompts y las respuestas:

# Define the system prompt that instructs the model to use a specific format
SYSTEM_PROMPT = """
Respond in the following format:

...

...

"""

XML_COT_FORMAT = """\

{reasoning}

{answer}

"""

Ahora preparemos el dataset:

import re
from datasets import load_dataset, Dataset

def extract_xml_answer(text: str) -> str:
    answer = text.split("")[-1]
    answer = answer.split("")[0]
    return answer.strip()

def extract_hash_answer(text: str) -> str | None:
    if "####" not in text:
        return None
    return text.split("####")[1].strip()

def get_gsm8k_questions(split="train") -> Dataset:
    data = load_dataset("openai/gsm8k", "main")[split]
    data = data.map(
        lambda x: {
            "prompt": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": x["question"]},
            ],
            "answer": extract_hash_answer(x["answer"]),
        }
    )
    return data

dataset = get_gsm8k_questions()

Definición de funciones de recompensa

Como comentamos en una página anterior, GRPO puede usar funciones de recompensa para guiar el aprendizaje según criterios verificables como longitud y formato.

En este ejercicio definiremos varias funciones de recompensa que fomentan distintos aspectos de un buen razonamiento.

def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
    responses = [completion[0]["content"] for completion in completions]
    q = prompts[0][-1]["content"]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    print(
        "-" * 20,
        f"Question:\n{q}",
        f"\nAnswer:\n{answer[0]}",
        f"\nResponse:\n{responses[0]}",
        f"\nExtracted:\n{extracted_responses[0]}",
    )
    return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]

def int_reward_func(completions, **kwargs) -> list[float]:
    responses = [completion[0]["content"] for completion in completions]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]

def strict_format_reward_func(completions, **kwargs) -> list[float]:
    pattern = r"^\n.*?\n\n\n.*?\n\n$"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, r) for r in responses]
    return [0.5 if match else 0.0 for match in matches]

def soft_format_reward_func(completions, **kwargs) -> list[float]:
    pattern = r".*?\s*.*?"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, r) for r in responses]
    return [0.5 if match else 0.0 for match in matches]

def count_xml(text) -> float:
    count = 0.0
    if text.count("\n") == 1:
        count += 0.125
    if text.count("\n\n") == 1:
        count += 0.125
    if text.count("\n\n") == 1:
        count += 0.125
        count -= len(text.split("\n\n")[-1]) * 0.001
    if text.count("\n") == 1:
        count += 0.125
        count -= (len(text.split("\n")[-1]) - 1) * 0.001
    return count

def xmlcount_reward_func(completions, **kwargs) -> list[float]:
    contents = [completion[0]["content"] for completion in completions]
    return [count_xml(c) for c in contents]

Estas funciones de recompensa cumplen propósitos distintos:

Reward Function	Purpose
correctness_reward_func	Recompensa al modelo cuando su respuesta coincide con la respuesta correcta
int_reward_func	Recompensa al modelo por dar una respuesta numérica
strict_format_reward_func y soft_format_reward_func	Recompensan al modelo por seguir el formato especificado
xmlcount_reward_func	Recompensa el uso correcto de etiquetas XML y penaliza contenido extra

Entrenamiento con GRPO

Ahora configuraremos GRPOTrainer con nuestro modelo, tokenizador y funciones de recompensa.

from trl import GRPOConfig, GRPOTrainer

max_prompt_length = 256

training_args = GRPOConfig(
    learning_rate=5e-6,
    adam_beta1=0.9,
    adam_beta2=0.99,
    weight_decay=0.1,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine",
    optim="paged_adamw_8bit",
    logging_steps=1,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,
    num_generations=6,
    max_prompt_length=max_prompt_length,
    max_completion_length=max_seq_length - max_prompt_length,
    max_steps=250,
    save_steps=250,
    max_grad_norm=0.1,
    report_to="none",
    output_dir="outputs",
)

trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=[
        xmlcount_reward_func,
        soft_format_reward_func,
        strict_format_reward_func,
        int_reward_func,
        correctness_reward_func,
    ],
    args=training_args,
    train_dataset=dataset,
)

La GRPOConfig establece varios hiperparámetros para el entrenamiento, entre ellos learning_rate, num_generations y max_steps.

Ahora iniciemos el entrenamiento:

trainer.train()

El entrenamiento puede tardar un tiempo. Puede que no veas mejoras inmediatas en las recompensas; a veces hacen falta entre 150 y 200 pasos.

Prueba del modelo

Después del entrenamiento, guardemos primero los pesos de LoRA:

model.save_lora("grpo_saved_lora")

Ahora probemos el modelo con una pregunta nueva:

from vllm import SamplingParams

text = tokenizer.apply_chat_template(
    [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": "Calculate pi."},
    ],
    tokenize=False,
    add_generation_prompt=True,
)

sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=1024,
)
output = (
    model.fast_generate(
        text,
        sampling_params=sampling_params,
        lora_request=model.load_lora("grpo_saved_lora"),
    )[0]
    .outputs[0]
    .text
)

print(output)

Deberías ver que el modelo ahora sigue el formato especificado y muestra su razonamiento antes de dar una respuesta.

Guardar el modelo

model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")

Subirlo al Hugging Face Hub

model.push_to_hub_merged(
    "your-username/model-name", tokenizer, save_method="merged_16bit", token="your-token"
)

Unsloth también permite guardar en formato GGUF para usarlo con llama.cpp:

model.push_to_hub_gguf(
    "your-username/model-name",
    tokenizer,
    quantization_method=["q4_k_m", "q8_0", "q5_k_m"],
    token="your-token",
)

Los archivos GGUF pueden usarse con llama.cpp o con interfaces como Jan u Open WebUI.

Conclusión

En este ejercicio aprendiste a:

1. Configurar Unsloth para un ajuste acelerado.
2. Preparar datos para entrenamiento con GRPO.
3. Definir funciones de recompensa para guiar el aprendizaje.
4. Entrenar un modelo usando GRPO.
5. Probar el modelo ajustado.
6. Guardar el modelo en varios formatos.

GRPO es una técnica potente para alinear modelos de lenguaje con comportamientos específicos, y Unsloth hace que resulte más accesible incluso en hardware limitado.

Para obtener más información y recursos, consulta:

• Documentación de Unsloth
• Discord de Unsloth
• GitHub de Unsloth