es/3_parameter_efficient_finetuning/notebooks/finetune_sft_peft.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z-6LLOPZouLg"
   },
   "source": [
    "# Cómo afinar los modelos de lenguaje grandes con adaptadores LoRA usando Hugging Face TRL\n",
    "\n",
    "Este notebook demuestra cómo afinar con precisión de manera eficiente modelos de lenguaje grandes usando adaptadores LoRA (adaptación de bajo rango). LoRA es una técnica de fine-tuning con eficiencia de parámetros que:\n",
    "- Congela los pesos del modelo pre-entrenado.\n",
    "- Agrega pequeñas matrices de descomposición de rangos entrenables a las capas de atención\n",
    "- Generalmente reduce los parámetros entrenables en aproximadamente un 90 %\n",
    "- Mantiene el rendimiento del modelo mientras es eficiente en el uso de la memoria\n",
    "\n",
    "Cubriremos:\n",
    "1. Configuración del entorno de desarrollo y configuración de LoRA\n",
    "2. Creación y preparación del conjunto de datos para el entrenamiento del adaptador\n",
    "3. _Fine-tuning_ utilizando `trl` y `SFTTrainer` con adaptadores LoRA\n",
    "4. Prueba del modelo y fusión de adaptadores (opcional)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fXqd9BXgouLi"
   },
   "source": [
    "## 1. Configura el entorno de desarrollo\n",
    "\n",
    "Nuestro primer paso es instalar las librerías de Hugging Face y PyTorch, incluidas `trl`, `transformers` y `datasets`. Si todavía no sabes qué es `trl`, no te preocupes. Se trata de una librería nueva integrada con `transformers`y `datasets`, que facilita el fine-tuning, aprendizaje por refuerzo o el alineamiento de modelos de lenguaje.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tKvGVxImouLi"
   },
   "outputs": [],
   "source": [
    "# Instala los requisitos en Google Colab\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "# Autenticación en Hugging Face\n",
    "\n",
    "from huggingface_hub import login\n",
    "\n",
    "login()\n",
    "\n",
    "# Por comodidad, puedes crear una variable de entorno que contenga tu token de Hugging Face como HF_TOKEN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XHUzfwpKouLk"
   },
   "source": [
    "## 2. Carga el conjunto de datos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "z4p6Bvo7ouLk"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['full_topic', 'messages'],\n",
       "        num_rows: 2260\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['full_topic', 'messages'],\n",
       "        num_rows: 119\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Carga el conjunto de datos\n",
    "from datasets import load_dataset\n",
    "\n",
    "# TODO: define tu conjunto de datos y configuración utilizando los parámetros de ruta y nombre.\n",
    "dataset = load_dataset(path=\"HuggingFaceTB/smoltalk\", name=\"everyday-conversations\")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9TOhJdtsouLk"
   },
   "source": [
    "## 3. Fine-tuning del LLM usando `trl` y el `SFTTrainer` con LoRA\n",
    "\n",
    "El [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) de `trl` está integrado con los adaptadores LoRA mediante la biblioteca [PEFT](https://huggingface.co/docs/peft/en/index). Las ventajas clave de esta configuración incluyen:\n",
    "\n",
    "1. **Eficiencia de memoria**: \n",
    "   - Solo los parámetros del adaptador se almacenan en la memoria de la GPU\n",
    "   - Los pesos del modelo base permanecen congelados y pueden cargarse en menor precisión\n",
    "   - Permite el fine-tuning de modelos grandes en GPUs de consumo\n",
    "\n",
    "2. **Características de entrenamiento**:\n",
    "   - Integración nativa de PEFT/LoRA con configuración mínima\n",
    "   - Soporte para QLoRA (LoRA Cuantizado) para una eficiencia de memoria aún mejor\n",
    "\n",
    "3. **Gestión de los adaptadores**:\n",
    "   - Guardado de pesos del adaptador durante puntos de control\n",
    "   - Funciones para fusionar adaptadores de vuelta al modelo base\n",
    "\n",
    "Usaremos LoRA en nuestro ejemplo, que combina LoRA con cuantización de 4 bits para reducir aún más el uso de memoria sin sacrificar el rendimiento. La configuración requiere solo unos pocos pasos:\n",
    "1. Definir la configuración de LoRA (rank, alpha, dropout)\n",
    "2. Crear el SFTTrainer con la configuración PEFT\n",
    "3. Entrenar y guardar los pesos del adaptador"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importa las librerias necesarias\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "from trl import SFTConfig, SFTTrainer, setup_chat_format\n",
    "import torch\n",
    "\n",
    "device = (\n",
    "    \"cuda\"\n",
    "    if torch.cuda.is_available()\n",
    "    else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    ")\n",
    "\n",
    "# Carga el modelo y el tokenizador\n",
    "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    pretrained_model_name_or_path=model_name\n",
    ").to(device)\n",
    "tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)\n",
    "\n",
    "# Establece el formato de chat\n",
    "model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)\n",
    "\n",
    "# Establece el nombre para nuestro fine-tuning que se guardará y/o subirá\n",
    "finetune_name = \"SmolLM2-FT-MyDataset\"\n",
    "finetune_tags = [\"smol-course\", \"module_1\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZbuVArTHouLk"
   },
   "source": [
    "El SFTTrainer cuenta con una integración nativa con peft, lo que hace que sea muy sencillo ajustar eficientemente los LLMs utilizando, por ejemplo, LoRA. Solo necesitamos crear nuestra LoraConfig y proporcionarla al entrenador. \n",
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'> \n",
    "    <h2 style='margin: 0;color:blue'>Ejercicio: Define los parámetros de LoRA para fine-tuning</h2> \n",
    "    <p>Escoge un conjunto de datos de Hugging Face y realiza fine-tuning a un modelo con él.</p> \n",
    "    <p><b>Niveles de dificultad</b></p> \n",
    "    <p>🐢 Usa los parámetros generales para un fine-tuning arbitrario</p> \n",
    "    <p>🐕 Ajusta los parámetros y revisalos en weights & biases.</p> \n",
    "    <p>🦁 Ajusta los parámetros y muestra el cambio en los resultados de inferencia.\n",
    "    </p> \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "blDSs9swouLk"
   },
   "outputs": [],
   "source": [
    "from peft import LoraConfig\n",
    "\n",
    "# TODO: Configura los parámetros de LoRA\n",
    "# r: dimensión de rango para las matrices de actualización de LoRA (más pequeño = más compresión)\n",
    "rank_dimension = 6\n",
    "# lora_alpha: factor de escala para las capas LoRA (más alto = adaptación más fuerte)\n",
    "lora_alpha = 8\n",
    "# lora_dropout: probabilidad de dropout para las capas LoRA (ayuda a prevenir el sobreajuste)\n",
    "lora_dropout = 0.05\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    r=rank_dimension,  # Dimensión de rango - normalmente entre 4-32\n",
    "    lora_alpha=lora_alpha,  # Factor de escala LoRA - normalmente 2x el rango\n",
    "    lora_dropout=lora_dropout,  # Probabilidad de dropout para las capas LoRA\n",
    "    bias=\"none\",  # Tipo de sesgo para LoRA. los sesgos correspondientes se actualizarán durante el entrenamiento.\n",
    "    target_modules=\"all-linear\",  # A qué módulos aplicar LoRA\n",
    "    task_type=\"CAUSAL_LM\",  # Tipo de tarea para la arquitectura del modelo\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l5NUDPcaouLl"
   },
   "source": [
    "Antes de comenzar nuestro entrenamiento, necesitamos definir los hiperparámetros (TrainingArguments) que queremos utilizar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NqT28VZlouLl"
   },
   "outputs": [],
   "source": [
    "# Configuración del entrenamiento\n",
    "# Hiperparámetros basados en las recomendaciones del artículo sobre QLoRA\n",
    "args = SFTConfig(\n",
    "    # Configuración de salida\n",
    "    output_dir=finetune_name,  # Directorio para guardar los checkpoints del modelo\n",
    "    # Duración del entrenamiento\n",
    "    num_train_epochs=1,  # Número de épocas de entrenamiento\n",
    "    # Configuración del tamaño de lote\n",
    "    per_device_train_batch_size=2,  # Tamaño de lote por GPU\n",
    "    gradient_accumulation_steps=2,  # Acumular gradientes para un lote efectivo más grande\n",
    "    # Optimización de memoria\n",
    "    gradient_checkpointing=True,  # Intercambiar cómputo por ahorro de memoria\n",
    "    # Configuración del optimizador\n",
    "    optim=\"adamw_torch_fused\",  # Usar AdamW fusionado para mayor eficiencia\n",
    "    learning_rate=2e-4,  # Tasa de aprendizaje (artículo sobre QLoRA)\n",
    "    max_grad_norm=0.3,  # Umbral de recorte de gradiente\n",
    "    # Programación de la tasa de aprendizaje\n",
    "    warmup_ratio=0.03,  # Porción de pasos para el calentamiento\n",
    "    lr_scheduler_type=\"constant\",  # Mantener la tasa de aprendizaje constante después del calentamiento\n",
    "    # Registro y guardado\n",
    "    logging_steps=10,  # Registrar métricas cada N pasos\n",
    "    save_strategy=\"epoch\",  # Guardar checkpoint cada época\n",
    "    # Configuración de precisión\n",
    "    bf16=True,  # Usar precisión bfloat16\n",
    "    # Configuración de integración\n",
    "    push_to_hub=False,  # No subir a HuggingFace\n",
    "    report_to=\"none\",  # Desactivar registro externo\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cGhR7uFBouLl"
   },
   "source": [
    "Ahora tenemos todos los componentes básicos que necesitamos para crear nuestro `SFTTrainer` y comenzar a entrenar nuestro modelo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "M00Har2douLl"
   },
   "outputs": [],
   "source": [
    "max_seq_length = 1512  # longitud máxima de secuencia para el modelo y el empaquetado del conjunto de datos\n",
    "\n",
    "# Crea SFTTrainer con configuración LoRA\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=dataset[\"train\"],\n",
    "    peft_config=peft_config,  # Configuración LoRA\n",
    "    max_seq_length=max_seq_length,  # Longitud máxima de la secuencia\n",
    "    tokenizer=tokenizer,\n",
    "    packing=True,  # Habilita el empaquetado de entrada para mayor eficiencia\n",
    "    dataset_kwargs={\n",
    "        \"add_special_tokens\": False,  # Los tokens especiales son manejados por la plantilla\n",
    "        \"append_concat_token\": False,  # No se necesita un separador adicional\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zQ_kRN24ouLl"
   },
   "source": [
    "Comienza a entrenar nuestro modelo llamando al método `train()` en nuestra instancia de `Trainer`. Esto iniciará el bucle de entrenamiento y entrenará nuestro modelo durante 3 épocas. Dado que estamos utilizando un método PEFT, solo guardaremos los pesos del modelo adaptado y no el modelo completo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Tq4nIYqKouLl"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "300e5dfbb4b54750b77324345c7591f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/72 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=72, training_loss=1.6402628521124523, metrics={'train_runtime': 195.2398, 'train_samples_per_second': 1.485, 'train_steps_per_second': 0.369, 'total_flos': 282267289092096.0, 'train_loss': 1.6402628521124523, 'epoch': 0.993103448275862})"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Inicia el entrenamiento, el modelo se guardará automáticamente en Hugging Face y en el directorio de salida\n",
    "trainer.train()\n",
    "\n",
    "# Guarda el modelo\n",
    "trainer.save_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y4HHSYYzouLl"
   },
   "source": [
    "El entrenamiento con Flash Attention durante 3 épocas con un conjunto de datos de 15k muestras tardó 4:14:36 en un `g5.2xlarge`. La instancia cuesta `1.21$/h`, lo que nos lleva a un coste total de solo ~`5.3$`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C309KsXjouLl"
   },
   "source": [
    "### Fusiona el Adaptador LoRA con el Modelo Original\n",
    "\n",
    "Cuando se usa LoRA, solo entrenamos los pesos del adaptador mientras mantenemos el modelo base congelado. Durante el entrenamiento, guardamos solo estos pesos de adaptador ligeros (~2-10MB) en lugar de una copia completa del modelo. Sin embargo, para la implementación, es posible que desee volver a fusionar los adaptadores con el modelo base para:\n",
    "\n",
    "1.  **Implementación Simplificada**: Un solo archivo de modelo en lugar de modelo base + adaptadores.\n",
    "2.  **Velocidad de Inferencia**: Sin sobrecarga de computación del adaptador.\n",
    "3.  **Compatibilidad con Frameworks**: Mejor compatibilidad con frameworks de servicio.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import AutoPeftModelForCausalLM\n",
    "\n",
    "\n",
    "# Carga el modelo PEFT en CPU\n",
    "model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "    pretrained_model_name_or_path=args.output_dir,\n",
    "    torch_dtype=torch.float16,\n",
    "    low_cpu_mem_usage=True,\n",
    ")\n",
    "\n",
    "# Fusiona LoRA y el modelo base, y guarda\n",
    "merged_model = model.merge_and_unload()\n",
    "merged_model.save_pretrained(\n",
    "    args.output_dir, safe_serialization=True, max_shard_size=\"2GB\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-yO6E9quouLl"
   },
   "source": [
    "## 3. Prueba el Modelo y Ejecuta la Inferencia\n",
    "\n",
    "Una vez que finaliza el entrenamiento, queremos probar nuestro modelo. Cargaremos diferentes muestras del conjunto de datos original y evaluaremos el modelo en esas muestras, utilizando un bucle simple y la precisión como nuestra métrica.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>Ejercicio Adicional: Carga el Adaptador LoRA</h2>\n",
    "    <p>Utiliza lo que aprendiste del notebook de ejemplo para cargar tu adaptador LoRA entrenado para la inferencia.</p> \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "I5B494OdouLl"
   },
   "outputs": [],
   "source": [
    "# Libera memoria otra vez\n",
    "del model\n",
    "del trainer\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P1UhohVdouLl"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from peft import AutoPeftModelForCausalLM\n",
    "from transformers import AutoTokenizer, pipeline\n",
    "\n",
    "# Carga el modelo con el adaptador PEFT\n",
    "tokenizer = AutoTokenizer.from_pretrained(finetune_name)\n",
    "model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "    finetune_name, device_map=\"auto\", torch_dtype=torch.float16\n",
    ")\n",
    "pipe = pipeline(\n",
    "    \"text-generation\", model=merged_model, tokenizer=tokenizer, device=device\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "99uFDAuuouLl"
   },
   "source": [
    "Probemos algunas muestras de *prompts* y veamos cómo se desempeña el modelo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "id": "-shSmUbvouLl",
    "outputId": "16d97c61-3b31-4040-c780-3c4de75c3824"
   },
   "outputs": [],
   "source": [
    "prompts = [\n",
    "    \"What is the capital of Germany? Explain why thats the case and if it was different in the past?\",\n",
    "    \"Write a Python function to calculate the factorial of a number.\",\n",
    "    \"A rectangular garden has a length of 25 feet and a width of 15 feet. If you want to build a fence around the entire garden, how many feet of fencing will you need?\",\n",
    "    \"What is the difference between a fruit and a vegetable? Give examples of each.\",\n",
    "]\n",
    "\n",
    "\n",
    "def test_inference(prompt):\n",
    "    prompt = pipe.tokenizer.apply_chat_template(\n",
    "        [{\"role\": \"user\", \"content\": prompt}],\n",
    "        tokenize=False,\n",
    "        add_generation_prompt=True,\n",
    "    )\n",
    "    outputs = pipe(\n",
    "        prompt,\n",
    "    )\n",
    "    return outputs[0][\"generated_text\"][len(prompt) :].strip()\n",
    "\n",
    "\n",
    "for prompt in prompts:\n",
    "    print(f\"    prompt:\\n{prompt}\")\n",
    "    print(f\"    response:\\n{test_inference(prompt)}\")\n",
    "    print(\"-\" * 50)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 0
}