pt-br/3_parameter_efficient_finetuning/notebooks/finetune_sft_peft.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z-6LLOPZouLg"
   },
   "source": [
    "# Como Ajustar LLMs com Adaptadores LoRA usando Hugging Face TRL\n",
    "\n",
    "Este notebook demonstra como ajustar eficientemente grandes modelos de linguagem usando adaptadores LoRA (Adaptação de Baixa Classificação). LoRA é uma técnica de ajuste fino eficiente em termos de parâmetros que:\n",
    "- Congela os pesos do modelo pré-treinado\n",
    "- Adiciona pequenas matrizes treináveis de decomposição de classificação às camadas de atenção\n",
    "- Normalmente reduz os parâmetros treináveis em ~90%\n",
    "- Mantém o desempenho do modelo sendo eficiente em termos de memória\n",
    "\n",
    "Vamos abordar:\n",
    "1. Configurar o ambiente de desenvolvimento e a configuração do LoRA\n",
    "2. Criar e preparar o conjunto de dados para o treinamento do adaptador\n",
    "3. Ajustar o modelo usando `trl` e `SFTTrainer` com adaptadores LoRA\n",
    "4. Testar o modelo e mesclar adaptadores (opcional)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fXqd9BXgouLi"
   },
   "source": [
    "## 1. Configurar o ambiente de desenvolvimento\n",
    "\n",
    "Nosso primeiro passo é instalar os módulos do Hugging Face e o PyTorch, incluindo trl, transformers e datasets. Se você ainda não ouviu falar do trl, não se preocupe. É uma nova biblioteca baseada em transformers e datasets, que facilita o ajuste fino, o RLHF (Reinforcement Learning from Human Feedback) e o alinhamento de modelos de linguagem abertos (open LLMs).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tKvGVxImouLi"
   },
   "outputs": [],
   "source": [
    "# Install the requirements in Google Colab\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "# Authenticate to Hugging Face\n",
    "\n",
    "from huggingface_hub import login\n",
    "\n",
    "login()\n",
    "\n",
    "# for convenience you can create an environment variable containing your hub token as HF_TOKEN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XHUzfwpKouLk"
   },
   "source": [
    "## 2. Carregue o conjunto de dados"
   ]
  },
  {
   "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": [
    "# Load a sample dataset\n",
    "from datasets import load_dataset\n",
    "\n",
    "# TODO: define your dataset and config using the path and name parameters\n",
    "dataset = load_dataset(path=\"HuggingFaceTB/smoltalk\", name=\"everyday-conversations\")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9TOhJdtsouLk"
   },
   "source": [
    "## 3. Ajustar uma LLM usando `trl` e o `SFTTrainer` com o LoRA\n",
    "\n",
    "O [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) do `trl` oferece integração com adaptadores LoRA através do módulo [PEFT](https://huggingface.co/docs/peft/en/index). As principais vantagens dessa configuração incluem:\n",
    "\n",
    "1. **Eficiência de Memória**: \n",
    "   - Apenas os parâmetros dos adaptadores são armazenados na memória da GPU\n",
    "   - Os pesos do modelo base permanecem congelados e podem ser carregados em precisão reduzida\n",
    "\n",
    "   - Enables fine-tuning of large models on consumer GPUs\n",
    "\n",
    "2. **Recursos de Treinamento**:\n",
    "   - Integração nativa PEFT/LoRA com configuração mínima\n",
    "   - Suporte ao QLoRA (LoRA Quantizado) para uma eficiência de memória ainda melhor\n",
    "\n",
    "3. **Gerenciamento de Adaptadores**:\n",
    "   - Salvamento de pesos dos adaptadores durante checkpoints\n",
    "   - Recursos para mesclar adaptadores de volta ao modelo base\n",
    "\n",
    "Usaremos LoRA em nosso exemplo, que combina LoRA com quantização de 4 bits para reduzir ainda mais o uso de memória sem sacrificar o desempenho. A configuração requer apenas algumas etapas:\n",
    "1. Definir a configuração do LoRA (rank, alpha, dropout)\n",
    "2. Criar o SFTTrainer com configuração PEFT\n",
    "3. Treinar e salvar os pesos dos adaptadores\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\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",
    "# Load the model and tokenizer\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",
    "# Set up the chat format\n",
    "model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)\n",
    "\n",
    "# Set our name for the finetune to be saved &/ uploaded to\n",
    "finetune_name = \"SmolLM2-FT-MyDataset\"\n",
    "finetune_tags = [\"smol-course\", \"module_1\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZbuVArTHouLk"
   },
   "source": [
    "O `SFTTrainer` oferece uma integração nativa com o `peft`, tornando extremamente fácil ajustar eficientemente LLMs, como com o LoRA, por exemplo. Precisamos apenas criar nosso `LoraConfig` e fornecê-lo ao trainer.\n",
    "\n",
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'> \n",
    "    <h2 style='margin: 0;color:blue'>Exercício: Definir parâmetros LoRA para ajuste fino</h2> \n",
    "    <p>Escolha um conjunto de dados do hub do Hugging Face e ajuste um modelo com ele.</p> \n",
    "    <p><b>Níveis de Dificuldade</b></p> \n",
    "    <p>🐢 Use os parâmetros gerais para um ajuste fino arbitrário.</p> \n",
    "    <p>🐕 Ajuste os parâmetros e revise em pesos & biases.</p> \n",
    "    <p>🦁 Ajuste os parâmetros e mostre a mudança nos resultados de inferência.</p> \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "blDSs9swouLk"
   },
   "outputs": [],
   "source": [
    "from peft import LoraConfig\n",
    "\n",
    "# TODO: Configure LoRA parameters\n",
    "# r: rank dimension for LoRA update matrices (smaller = more compression)\n",
    "rank_dimension = 6\n",
    "# lora_alpha: scaling factor for LoRA layers (higher = stronger adaptation)\n",
    "lora_alpha = 8\n",
    "# lora_dropout: dropout probability for LoRA layers (helps prevent overfitting)\n",
    "lora_dropout = 0.05\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    r=rank_dimension,  # Rank dimension - typically between 4-32\n",
    "    lora_alpha=lora_alpha,  # LoRA scaling factor - typically 2x rank\n",
    "    lora_dropout=lora_dropout,  # Dropout probability for LoRA layers\n",
    "    bias=\"none\",  # Bias type for LoRA. the corresponding biases will be updated during training.\n",
    "    target_modules=\"all-linear\",  # Which modules to apply LoRA to\n",
    "    task_type=\"CAUSAL_LM\",  # Task type for model architecture\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l5NUDPcaouLl"
   },
   "source": [
    "Antes de iniciarmos nosso treinamento, precisamos definir os hiperparâmetros (`TrainingArguments`) que queremos usar.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NqT28VZlouLl"
   },
   "outputs": [],
   "source": [
    "# Training configuration\n",
    "# Hyperparameters based on QLoRA paper recommendations\n",
    "args = SFTConfig(\n",
    "    # Output settings\n",
    "    output_dir=finetune_name,  # Directory to save model checkpoints\n",
    "    # Training duration\n",
    "    num_train_epochs=1,  # Number of training epochs\n",
    "    # Batch size settings\n",
    "    per_device_train_batch_size=2,  # Batch size per GPU\n",
    "    gradient_accumulation_steps=2,  # Accumulate gradients for larger effective batch\n",
    "    # Memory optimization\n",
    "    gradient_checkpointing=True,  # Trade compute for memory savings\n",
    "    # Optimizer settings\n",
    "    optim=\"adamw_torch_fused\",  # Use fused AdamW for efficiency\n",
    "    learning_rate=2e-4,  # Learning rate (QLoRA paper)\n",
    "    max_grad_norm=0.3,  # Gradient clipping threshold\n",
    "    # Learning rate schedule\n",
    "    warmup_ratio=0.03,  # Portion of steps for warmup\n",
    "    lr_scheduler_type=\"constant\",  # Keep learning rate constant after warmup\n",
    "    # Logging and saving\n",
    "    logging_steps=10,  # Log metrics every N steps\n",
    "    save_strategy=\"epoch\",  # Save checkpoint every epoch\n",
    "    # Precision settings\n",
    "    bf16=True,  # Use bfloat16 precision\n",
    "    # Integration settings\n",
    "    push_to_hub=False,  # Don't push to HuggingFace Hub\n",
    "    report_to=None,  # Disable external logging\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cGhR7uFBouLl"
   },
   "source": [
    "Agora temos todos os pedaços que precisamos para criar nosso `SFTTrainer` para então começar a treinar nosso modelo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "M00Har2douLl"
   },
   "outputs": [],
   "source": [
    "max_seq_length = 1512  # max sequence length for model and packing of the dataset\n",
    "\n",
    "# Create SFTTrainer with LoRA configuration\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=dataset[\"train\"],\n",
    "    peft_config=peft_config,  # LoRA configuration\n",
    "    max_seq_length=max_seq_length,  # Maximum sequence length\n",
    "    tokenizer=tokenizer,\n",
    "    packing=True,  # Enable input packing for efficiency\n",
    "    dataset_kwargs={\n",
    "        \"add_special_tokens\": False,  # Special tokens handled by template\n",
    "        \"append_concat_token\": False,  # No additional separator needed\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zQ_kRN24ouLl"
   },
   "source": [
    "Comece treinando nosso modelo chamando a função `train ()` em nossa instância de `Trainer`. Isso iniciará o loop de treinamento e treinará nosso modelo por 3 épocas. Como estamos usando um método PEFT, salvaremos apenas os pesos do modelo adaptado e não o do 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": [
    "# start training, the model will be automatically saved to the hub and the output directory\n",
    "trainer.train()\n",
    "\n",
    "# save model\n",
    "trainer.save_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y4HHSYYzouLl"
   },
   "source": [
    "O treinamento com Flash Attention por 3 épocas com um conjunto de dados de 15 mil amostras levou 4:14:36 ​​em um `g5.2xlarge`. A instância custa `1,21 dólares/h`, o que nos leva a um custo total de apenas ~`5,3 dólares`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C309KsXjouLl"
   },
   "source": [
    "### Mesclar Adaptador LoRA no Modelo Original\n",
    "\n",
    "Ao usar LoRA, treinamos apenas os pesos do adaptador mantendo o modelo base congelado. Durante o treinamento, salvamos apenas esses pesos de adaptador leves (~2-10MB) em vez de uma cópia completa do modelo. No entanto, para implantação, você talvez queira mesclar os adaptadores de volta ao modelo base para:\n",
    "\n",
    "1. **Implantação Simplificada**: Arquivo de modelo único em vez de modelo base + adaptadores\n",
    "2. **Velocidade de Inferência**: Sem sobrecarga de computação de adaptadores\n",
    "3. **Compatibilidade de Framework**: Melhor compatibilidade com serving frameworks \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import AutoPeftModelForCausalLM\n",
    "\n",
    "\n",
    "# Load PEFT model on 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",
    "# Merge LoRA and base model and save\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. Testar Modelo e Executar Inferência\n",
    "\n",
    "Após a conclusão do treinamento, queremos testar nosso modelo. Vamos carregar diferentes amostras do conjunto de dados original e avaliar o modelo baseado nessas amostras, usando um loop simples e a precisão como nossa métrica."
   ]
  },
  {
   "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'>Exercício Bônus: Carregar Adaptador LoRA</h2>\n",
    "    <p>Use o que você aprendeu no caderno de exemplo para carregar seu adaptador LoRA treinado para inferência.</p> \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "I5B494OdouLl"
   },
   "outputs": [],
   "source": [
    "# free the memory again\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",
    "# Load Model with PEFT adapter\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": [
    "Vamos testar uns exemplos de prompts e ver como o modelo reage."
   ]
  },
  {
   "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
}