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](/course/chapter3/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.

> [!TIP]
> 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

```bash
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:

```python
from unsloth import FastLanguageModel
```

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

```python
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.

> [!TIP]
> No cubriremos los detalles de LoRA en este capítulo, pero puedes aprender más en el [Capítulo 11](/course/chapter11/3).

## 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:

```python
# 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:

```python
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](/course/chapter12/4), 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.

```python
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.

```python
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:

```python
trainer.train()
```

> [!WARNING]
> 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:

```python
model.save_lora("grpo_saved_lora")
```

Ahora probemos el modelo con una pregunta nueva:

```python
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

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

## Subirlo al Hugging Face Hub

```python
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`:

```python
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](https://docs.unsloth.ai/)
- [Discord de Unsloth](https://discord.gg/unsloth)
- [GitHub de Unsloth](https://github.com/unslothai/unsloth)