vi/3_parameter_efficient_finetuning/notebooks/finetune_sft_peft.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z-6LLOPZouLg"
   },
   "source": [
    "# Cách tinh chỉnh LLM với LoRA Adapters sử dụng Hugging Face TRL\n",
    "\n",
    "Trong Notebook này bạn sẽ được học cách tinh chỉnh hiệu quả các mô hình ngôn ngữ lớn sử dụng LoRA (Low-Rank Adaptation) adapters. LoRA là một kỹ thuật tinh chỉnh tham số hiệu quả:\n",
    "- Đóng băng các trọng số mô hình đã huấn luyện trước\n",
    "- Thêm các ma trận phân rã hạng nhỏ có thể huấn luyện vào các lớp attention  \n",
    "- Thường giảm khoảng 90% tham số có thể huấn luyện\n",
    "- Duy trì hiệu suất mô hình trong khi sử dụng bộ nhớ hiệu quả\n",
    "\n",
    "Chúng ta sẽ tìm hiểu:  \n",
    "1. Cài đặt môi trường phát triển và cấu hình LoRA\n",
    "2. Tạo và chuẩn bị dữ liệu để huấn luyện adapter \n",
    "3. Tinh chỉnh sử dụng `trl` và `SFTTrainer` với LoRA adapters\n",
    "4. Kiểm tra mô hình và gộp adapters (tùy chọn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fXqd9BXgouLi"
   },
   "source": [
    "## 1. Cài đặt môi trường phát triển\n",
    "\n",
    "Bước đầu tiên là cài đặt các thư viện Hugging Face và PyTorch, bao gồm `trl, transformers và datasets`. Nếu bạn chưa nghe nói về `trl`, đừng lo lắng. Đó là một thư viện mới trên nền tảng transformers và datasets, giúp việc tinh chỉnh các mô hình ngôn ngữ lớn (LLM) trở nên dễ dàng hơn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tKvGVxImouLi"
   },
   "outputs": [],
   "source": [
    "# Cài đặt các yêu cầu trong Google Colab\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "# Xác thực với Hugging Face\n",
    "from huggingface_hub import login\n",
    "\n",
    "login()\n",
    "\n",
    "# để thuận tiện bạn có thể tạo một biến môi trường chứa token hub của bạn là HF_TOKEN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XHUzfwpKouLk"
   },
   "source": [
    "## 2. Tải dữ liệu "
   ]
  },
  {
   "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": [
    "# Tải một tập dữ liệu mẫu\n",
    "from datasets import load_dataset\n",
    "\n",
    "# TODO: xác định dataset và config của bạn sử dụng các tham số path và name\n",
    "dataset = load_dataset(path=\"HuggingFaceTB/smoltalk\", name=\"everyday-conversations\")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9TOhJdtsouLk"
   },
   "source": [
    "## 3. Tinh chỉnh LLM sử dụng `trl` và `SFTTrainer` với LoRA\n",
    "\n",
    "[SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) từ `trl` cung cấp tích hợp với LoRA adapters thông qua thư viện [PEFT](https://huggingface.co/docs/peft/en/index). Những lợi thế chính của cài đặt này bao gồm:\n",
    "\n",
    "1. **Hiệu quả bộ nhớ**:\n",
    "   - Chỉ các tham số adapter được lưu trữ trong bộ nhớ GPU \n",
    "   - Trọng số mô hình cơ sở vẫn đóng băng và có thể được tải với độ chính xác thấp hơn\n",
    "   - Cho phép tinh chỉnh các mô hình lớn trên GPU tiêu dùng\n",
    "\n",
    "2. **Tính năng huấn luyện**:\n",
    "   - Tích hợp PEFT/LoRA sẵn có với cài đặt tối thiểu\n",
    "   - Hỗ trợ QLoRA (LoRA lượng tử hóa) cho hiệu quả bộ nhớ tốt hơn\n",
    "\n",
    "3. **Quản lý Adapter**: \n",
    "   - Lưu trọng số adapter trong quá trình checkpoint\n",
    "   - Tính năng gộp adapters trở lại mô hình cơ sở\n",
    "\n",
    "Chúng ta sẽ sử dụng LoRA trong ví dụ của mình, **kết hợp LoRA với lượng tử hóa 4-bit** để giảm thêm việc sử dụng bộ nhớ mà không ảnh hưởng đến hiệu suất. Cài đặt chỉ yêu cầu một vài bước cấu hình:\n",
    "1. Xác định cấu hình LoRA (rank, alpha, dropout)\n",
    "2. Tạo SFTTrainer với cấu hình PEFT \n",
    "3. Huấn luyện và lưu trọng số adapter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Các thư viện cần thiết \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",
    "# Tải mô hình và 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",
    "# Thiết lập định dạng chat\n",
    "model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)\n",
    "\n",
    "# Đặt tên cho bản tinh chỉnh để lưu &/ tải lên\n",
    "finetune_name = \"SmolLM2-FT-LoRA-Adapter\"\n",
    "finetune_tags = [\"smol-course\", \"module_3\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZbuVArTHouLk"
   },
   "source": [
    "`SFTTrainer` hỗ trợ tích hợp sẵn với `peft`, điều này giúp tinh chỉnh hiệu quả các LLM dễ dàng hơn bằng cách sử dụng, ví dụ như LoRA. Chúng ta chỉ cần tạo `LoraConfig` và cung cấp nó cho 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'>Bài tập: Xác định tham số LoRA cho tinh chỉnh</h2>\n",
    "    <p>Lấy một bộ dữ liệu từ Hugging Face hub và tinh chỉnh một mô hình trên nó.</p>\n",
    "    <p><b>Mức độ khó</b></p> \n",
    "    <p>🐢 Sử dụng các tham số chung cho một bản tinh chỉnh tùy ý</p>\n",
    "    <p>🐕 Điều chỉnh các tham số và so sánh trong weights & biases.</p>\n",
    "    <p>🦁 Điều chỉnh các tham số và cho thấy sự thay đổi trong kết quả suy luận.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "blDSs9swouLk"
   },
   "outputs": [],
   "source": [
    "from peft import LoraConfig\n",
    "\n",
    "# TODO: Điều chỉnh cấu hình tham số LoRA\n",
    "# r: chiều rank cho ma trận cập nhật LoRA (nhỏ hơn = nén nhiều hơn)\n",
    "rank_dimension = 4\n",
    "# lora_alpha: hệ số tỷ lệ cho các lớp LoRA (cao hơn = điều chỉnh mạnh hơn) \n",
    "lora_alpha = 8\n",
    "# lora_dropout: xác suất dropout cho các lớp LoRA (giúp tránh overfitting)\n",
    "lora_dropout = 0.05\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    r=rank_dimension,  # Chiều rank - thường từ 4-32\n",
    "    lora_alpha=lora_alpha,  # Hệ số tỷ lệ LoRA - thường gấp 2 lần rank\n",
    "    lora_dropout=lora_dropout,  # Xác suất dropout cho các lớp LoRA \n",
    "    bias=\"none\",  # Loại bias cho LoRA. các bias tương ứng sẽ được cập nhật trong quá trình huấn luyện.\n",
    "    target_modules=\"all-linear\",  # Những module nào áp dụng LoRA\n",
    "    task_type=\"CAUSAL_LM\",  # Loại tác vụ cho kiến trúc mô hình\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l5NUDPcaouLl"
   },
   "source": [
    "Trước khi bắt đầu huấn luyện, chúng ta cần xác định các `siêu tham số (TrainingArguments)` mà chúng ta muốn sử dụng."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NqT28VZlouLl"
   },
   "outputs": [],
   "source": [
    "# Cấu hình huấn luyện  \n",
    "# Siêu tham số dựa trên khuyến nghị từ bài báo QLoRA \n",
    "args = SFTConfig(\n",
    "    # Cài đặt đầu ra\n",
    "    output_dir=finetune_name,  # Thư mục để lưu checkpoint mô hình\n",
    "    # Thời gian huấn luyện\n",
    "    num_train_epochs=1,  # Số epoch huấn luyện\n",
    "    # Cài đặt kích thước batch \n",
    "    per_device_train_batch_size=2,  # Kích thước batch cho mỗi GPU\n",
    "    gradient_accumulation_steps=2,  # Tích lũy gradient cho batch hiệu quả lớn hơn\n",
    "    # Tối ưu bộ nhớ\n",
    "    gradient_checkpointing=True,  # Đánh đổi tính toán để tiết kiệm bộ nhớ\n",
    "    # Cài đặt optimizer\n",
    "    optim=\"adamw_torch_fused\",  # Sử dụng AdamW fusion cho hiệu quả\n",
    "    learning_rate=2e-4,  # Tốc độ học (từ bài báo QLoRA)\n",
    "    max_grad_norm=0.3,  # Ngưỡng cắt gradient\n",
    "    # Lịch trình học\n",
    "    warmup_ratio=0.03,  # Phần bước cho warmup\n",
    "    lr_scheduler_type=\"constant\",  # Giữ tốc độ học không đổi sau warmup\n",
    "    # Ghi log và lưu\n",
    "    logging_steps=10,  # Ghi metrics mỗi N bước\n",
    "    save_strategy=\"epoch\",  # Lưu checkpoint mỗi epoch\n",
    "    # Cài đặt độ chính xác\n",
    "    bf16=True,  # Sử dụng độ chính xác bfloat16\n",
    "    # Cài đặt tích hợp\n",
    "    push_to_hub=False,  # Không đẩy lên HuggingFace Hub\n",
    "    report_to=None,  # Tắt ghi log bên ngoài\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cGhR7uFBouLl"
   },
   "source": [
    "Bây giờ chúng ta đã có mọi thứ cần thiết để tạo `SFTTrainer` và bắt đầu huấn luyện mô hình của mình."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "M00Har2douLl"
   },
   "outputs": [],
   "source": [
    "max_seq_length = 1512  # độ dài chuỗi tối đa cho mô hình và đóng gói (packing) bộ dữ liệu\n",
    "\n",
    "# Tạo SFTTrainer với cấu hình LoRA\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=dataset[\"train\"],\n",
    "    peft_config=peft_config,  # Cấu hình LoRA\n",
    "    max_seq_length=max_seq_length,  # Độ dài chuỗi tối đa \n",
    "    tokenizer=tokenizer,\n",
    "    packing=True,  # Bật đóng gói đầu vào cho hiệu quả \n",
    "    dataset_kwargs={\n",
    "        \"add_special_tokens\": False,  # Token đặc biệt được xử lý bởi template\n",
    "        \"append_concat_token\": False,  # Không cần thêm separator \n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zQ_kRN24ouLl"
   },
   "source": [
    "Bắt đầu huấn luyện mô hình bằng cách gọi phương thức `train()` trên `Trainer` của chúng ta. Việc này sẽ bắt đầu vòng lặp huấn luyện và huấn luyện mô hình của chúng ta trong `3 epochs`. Vì chúng ta đang sử dụng phương pháp PEFT, chúng ta sẽ chỉ lưu phần trọng số của adapter đã điều chỉnh và không lưu toàn bộ mô hình."
   ]
  },
  {
   "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": [
    "# bắt đầu huấn luyện, mô hình sẽ tự động được lưu lên hub và thư mục đầu ra\n",
    "trainer.train()\n",
    "\n",
    "# lưu mô hình\n",
    "trainer.save_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y4HHSYYzouLl"
   },
   "source": [
    "Việc huấn luyện với Flash Attention cho 3 epoch với một dataset 15k mẫu mất `4:14:36` trên một cụm máy `g5.2xlarge` của AWS. Cụm máy này có giá `1.21$/h` và tổng chi phí của lần huấn luyện này chỉ tốn khoảng `5.3$`.\n",
    "\n",
    "**Ghi chú: bạn hoàn toàn có thể sử dụng GPU của Kaggle hoặc Google Colab để huấn luyện**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C309KsXjouLl"
   },
   "source": [
    "### Gộp LoRA Adapter vào Mô hình Gốc\n",
    "\n",
    "Khi sử dụng LoRA, chúng ta chỉ huấn luyện trọng số adapter trong khi giữ nguyên mô hình cơ sở. Trong quá trình huấn luyện, chúng ta chỉ lưu những trọng số adapter nhẹ này (~2-10MB) thay vì một bản sao mô hình đầy đủ. Tuy nhiên, để triển khai, bạn có thể muốn gộp các adapter trở lại mô hình cơ sở để:\n",
    "\n",
    "1. **Đơn giản hóa triển khai**: Một file mô hình thay vì mô hình cơ sở + adapters\n",
    "2. **Tốc độ suy luận**: Không có chi phí tính toán adapter phụ thêm\n",
    "3. **Tương thích Framework**: Tương thích tốt hơn với các framework phục vụ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import AutoPeftModelForCausalLM\n",
    "\n",
    "\n",
    "# Tải mô hình PEFT trên 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",
    "# Gộp LoRA và mô hình cơ sở và lưu\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. Kiểm tra Mô hình\n",
    "\n",
    "Sau khi huấn luyện hoàn tất, chúng ta muốn kiểm tra mô hình của mình. Chúng ta sẽ tải các mẫu khác nhau từ dataset gốc và đánh giá mô hình trên những mẫu đó, sử dụng một vòng lặp đơn giản và độ chính xác làm điểm số đánh giá.\n",
    "\n",
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>Bài tập Bonus: Tải LoRA Adapter</h2>\n",
    "    <p>Sử dụng những gì bạn đã học được từ notebook ví dụ để tải adapter LoRA đã huấn luyện của bạn cho suy luận.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "I5B494OdouLl"
   },
   "outputs": [],
   "source": [
    "# giải phóng bộ nhớ một lần nữa\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",
    "# Tải Mô hình với 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": [
    "Hãy thử một số chỉ định mẫu và xem mô hình hoạt động như thế nào."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "id": "-shSmUbvouLl",
    "outputId": "16d97c61-3b31-4040-c780-3c4de75c3824"
   },
   "outputs": [],
   "source": [
    "prompts = [\n",
    "    \"Thủ đô của Việt Nam thành phố nào? Giải thích tại sao là như vậy và liệu nó có khác trong quá khứ không?\",\n",
    "    \"Viết một hàm Python để tính giai thừa của một số.\",\n",
    "    \"Một khu vườn hình chữ nhật có chiều dài 25 mét và chiều rộng 15 mét. Nếu bạn muốn xây một hàng rào xung quanh toàn bộ khu vườn, bạn sẽ cần bao nhiêu mét hàng rào?\",\n",
    "    \"Sự khác biệt giữa trái cây và rau củ là gì? Đưa ra ví dụ cho mỗi loại.\",\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
}