ko/3_parameter_efficient_finetuning/notebooks/finetune_sft_peft.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z-6LLOPZouLg"
   },
   "source": [
    "# Hugging Face TRL과 LoRA 어댑터를 활용해 LLM을 미세 조정하는 방법\n",
    "\n",
    "이 노트북에서는 LoRA(Low-Rank Adaptation) 어댑터를 사용하여 대형 언어 모델을 효율적으로 파인튜닝하는 방법을 설명합니다. LoRA는 다음과 같은 특징을 가지는 파라미터 효율적인 미세 조정 기법입니다:\n",
    "- 사전 학습 모델 가중치 고정\n",
    "- 어텐션 레이어에 학습 가능한 저랭크 분해 행렬 추가\n",
    "- 학습 가능한 파라미터 수가 최대 90%까지 감소\n",
    "- 메모리를 효율적으로 사용하면서 모델 성능 유지\n",
    "\n",
    "이 노트북에서 다루는 내용입니다:\n",
    "1. 개발 환경 및 LoRA configuration 설정\n",
    "2. 어댑터 학습을 위한 데이터셋 생성 및 준비\n",
    "3. `trl`과 `SFTTrainer`를 활용한 LoRA 어댑터 기반 미세 조정\n",
    "4. 모델 테스트 및 어댑터 병합(선택 사항)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fXqd9BXgouLi"
   },
   "source": [
    "## 1. 개발 환경 설정\n",
    "\n",
    "첫 번째 단계는 trl, transformers, datasets을 포함한 Hugging Face 라이브러리와 Pytorch를 설치하는 것입니다. trl을 들어본 적이 없더라도 걱정하지 마세요. trl은 transformers와 datasets 위에 구축된 새로운 라이브러리로, 오픈 소스 LLM의 미세 조정과 RLHF, 정렬 작업을 쉽게 만들어줍니다.  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tKvGVxImouLi"
   },
   "outputs": [],
   "source": [
    "# Google Colab에서 requirements 설치\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "# Hugging Face 인증\n",
    "\n",
    "from huggingface_hub import login\n",
    "\n",
    "login()\n",
    "\n",
    "# 허브 토큰을 HF_TOKEN 환경 변수로 설정해두면 편하게 사용할 수 있습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XHUzfwpKouLk"
   },
   "source": [
    "## 2. 데이터셋 불러오기"
   ]
  },
  {
   "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": [
    "# 예시 데이터셋 불러오기\n",
    "from datasets import load_dataset\n",
    "\n",
    "# TODO: path와 name 파라미터를 이용해 원하는 데이터셋 정의하기\n",
    "dataset = load_dataset(path=\"HuggingFaceTB/smoltalk\", name=\"everyday-conversations\")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9TOhJdtsouLk"
   },
   "source": [
    "## 3. LoRA와 함께 `trl`과 `SFTTrainer`를 사용한 LLM 미세 조정\n",
    "\n",
    "`trl`에서 제공하는 [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer)는 [PEFT](https://huggingface.co/docs/peft/en/index) 라이브러리를 통해 LoRA 어댑터와의 통합을 제공합니다. 이 설정의 주요 장점은 다음과 같습니다:\n",
    "\n",
    "1. **메모리 효율성**: \n",
    "   - GPU 메모리에 어댑터 파라미터만 저장\n",
    "   - 기본 모델 가중치는 고정되며 더 낮은 정밀도로 불러오기 가능\n",
    "   - 서버가 아닌 소비자용 GPU에서 대형 모델 미세 조정 가능\n",
    "\n",
    "2. **학습 기능**:\n",
    "   - 최소한의 설정으로 기본 PEFT/LoRA 통합\n",
    "   - 더 나은 메모리 효율을 위한 QLoRA(양자화된 LoRA) 지원\n",
    "\n",
    "3. **어댑터 관리**:\n",
    "   - 체크포인트 저장 시 어댑터 가중치도 저장\n",
    "   - 어댑터를 기본 모델에 병합하는 기능 제공\n",
    "\n",
    "이번 예제에서는 성능 저하 없이 메모리 사용량을 더욱 줄이기 위해 4비트 양자화를 결합한 LoRA를 사용하겠습니다. 설정은 몇 가지 단계만 거치면 됩니다:\n",
    "\n",
    "1. LoRA configuration 설정(랭크, 알파, 드롭아웃)\n",
    "2. PEFT configuration을 추가한 SFTTrainer 생성\n",
    "3. 어댑터 가중치 학습 및 저장\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 주요 라이브러리 불러오기\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",
    "# 모델과 토크나이저 불러오기\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",
    "# 대화 형식 설정\n",
    "model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)\n",
    "\n",
    "# 미세 조정 결과를 저장하고 업로드하기 위한 이름 설정\n",
    "finetune_name = \"SmolLM2-FT-MyDataset\"\n",
    "finetune_tags = [\"smol-course\", \"module_1\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZbuVArTHouLk"
   },
   "source": [
    "`SFTTrainer`는 `peft`와의 기본 통합을 지원하기 때문에 LoRA와 같은 방법으로 LLM 미세 조정 작업을 매우 쉽고 효율적으로 진행할 수 있습니다. 우리는 `LoraConfig`를 생성하고 트레이너에게 제공하기만 하면 됩니다.\n",
    "\n",
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>연습: 미세 조정을 위한 LoRA 파라미터 정의하기</h2>\n",
    "    <p>Hugging Face 허브에서 가져온 데이터셋으로 모델을 미세 조정 해보세요.</p> \n",
    "    <p><b>난이도</b></p>\n",
    "    <p>🐢 일반적인 파라미터 값을 사용하여 임의로 미세 조정 진행하기</p>\n",
    "    <p>🐕 파라미터를 조정하고 가중치와 편향 검토해보기</p>\n",
    "    <p>🦁 파라미터를 조정하고 추론 결과 변화 확인하기</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "blDSs9swouLk"
   },
   "outputs": [],
   "source": [
    "from peft import LoraConfig\n",
    "\n",
    "# TODO: LoRA 파라미터를 구성해보세요.\n",
    "# r: LoRA 업데이트 행렬을 위한 랭크 차원원 (작을수록 압축률이 높아짐)\n",
    "rank_dimension = 6\n",
    "# lora_alpha: LoRA 레이어를 위한 스케일링 계수 (높을수록 모델이 새로운 데이터에 잘 적응함)\n",
    "lora_alpha = 8\n",
    "# lora_dropout: LoRA 레이어의 드롭아웃 확률 (과적합 방지)\n",
    "lora_dropout = 0.05\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    r=rank_dimension,  # 랭크 차원 - 일반적으로 4-32 사이의 값 사용\n",
    "    lora_alpha=lora_alpha,  # LoRA 스케일링 계수 - 일반적으로 랭크 차원의 2배\n",
    "    lora_dropout=lora_dropout,  # LoRA 레이어의의 드롭아웃 확률\n",
    "    bias=\"none\",  # LoRA 편향 유형 - 해당 편향은 학습 중 업데이트됨됨\n",
    "    target_modules=\"all-linear\",  # LoRA를 적용할 모듈듈\n",
    "    task_type=\"CAUSAL_LM\",  # 모델 구조에 맞는 태스크 유형\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l5NUDPcaouLl"
   },
   "source": [
    "학습을 시작하기 전 사용하고자 하는 하이퍼파라미터(`TrainingArguments`) 값을 설정해야 합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NqT28VZlouLl"
   },
   "outputs": [],
   "source": [
    "# 학습을 위한 configuration\n",
    "# QLoRA 논문에 기반한 하이퍼파라미터 설정\n",
    "args = SFTConfig(\n",
    "    # 출력 설정\n",
    "    output_dir=finetune_name,  # 모델 체크포인트 저장 디렉토리\n",
    "    # 학습 기간\n",
    "    num_train_epochs=1,  # 학습 에포크 수\n",
    "    # 배치 크기 설정\n",
    "    per_device_train_batch_size=2,  # GPU당 배치 크기\n",
    "    gradient_accumulation_steps=2,  # 더 큰 배치 크기 효과를 위한 경사 누적\n",
    "    # 메모리 최적화\n",
    "    gradient_checkpointing=True,  # 학습 속도는 느려지지만 메모리 절약 가능\n",
    "    # 옵티마이저 설정\n",
    "    optim=\"adamw_torch_fused\",  # 효율성을 위해 fused AdamW 사용\n",
    "    learning_rate=2e-4,  # 학습률 (QLoRA 논문 참고)\n",
    "    max_grad_norm=0.3,  # 경사 클리핑 임곗값\n",
    "    # 학습률 스케줄링\n",
    "    warmup_ratio=0.03,  # 학습률을 0에서 목표 값까지 선형적으로 증가시키는 웜업을 위한 전체 학습 스텝의 비율\n",
    "    lr_scheduler_type=\"constant\",  # 웜업 후 일정한 학습률 유지\n",
    "    # 로깅 및 저장\n",
    "    logging_steps=10,  # 매 N 스텝마다 지표 기록\n",
    "    save_strategy=\"epoch\",  # 매 에포크마다 체크포인트 저장\n",
    "    # 정밀도 설정\n",
    "    bf16=True,  # bfloat16 정밀도 사용\n",
    "    # 통합 설정\n",
    "    push_to_hub=False,  # Hugging Face 허브로 내보내지 않음\n",
    "    report_to=\"none\",  # 외부 로깅 비활성화\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cGhR7uFBouLl"
   },
   "source": [
    "이제 모델 학습을 위한 `SFTTrainer` 생성에 필요한 요소를 모두 갖췄습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "M00Har2douLl"
   },
   "outputs": [],
   "source": [
    "max_seq_length = 1512  # 모델과 데이터셋 패킹을 위한 최대 시퀀스 길이\n",
    "\n",
    "# LoRA configuration으로 SFTTrainer 생성\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,  # 최대 시퀀스 길이\n",
    "    tokenizer=tokenizer,\n",
    "    packing=True,  # 효율성을 위해 입력 패킹 활성화\n",
    "    dataset_kwargs={\n",
    "        \"add_special_tokens\": False,  # 템플릿에서 추가 토큰 처리\n",
    "        \"append_concat_token\": False,  # 추가 구분자 없음\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zQ_kRN24ouLl"
   },
   "source": [
    "`Trainer` 객체의 `train()` 메서드를 호출해 학습을 시작합니다. 학습 루프가 시작되고 3 에포크 동안 모델이 학습됩니다. 지금은 PEFT 방법을 쓰고 있기 때문에 전체 모델을 저장하지 않고 적응된 모델 가중치만 저장합니다."
   ]
  },
  {
   "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": [
    "# 학습이 시작되면 허브와 출력 디렉토리에 모델 자동 저장\n",
    "trainer.train()\n",
    "\n",
    "# 모델 저장\n",
    "trainer.save_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y4HHSYYzouLl"
   },
   "source": [
    "`g5.2xlarge` 인스턴스에서 플래시 어텐션을 적용해 15,000개의 샘플을 3 에포크 동안 학습하는 데 걸린 시간은 4시간 14분 36초였습니다. 해당 인스턴스 비용은 시간당 1.21달러로, 총 비용은 약 `5.3달러`에 불과합니다."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C309KsXjouLl"
   },
   "source": [
    "### 기존 모델과 LoRA 어댑터 병합\n",
    "\n",
    "LoRA를 사용할 때 기본 모델은 고정된 상태로 두고 어댑터 가중치만 학습합니다. 학습 과정에서는 전체 모델이 아닌 2-10MB 크기의 가벼운 어댑터 가중치만 저장합니다. 하지만 배포할 경우 다음과 같은 이유로 어댑터를 기본 모델과 병합하는 것이 더 나을 수 있습니다:\n",
    "\n",
    "1. **간소화된 배포**: 기본 모델과 어댑터를 따로 관리하지 않고 단일 모델로 처리 가능\n",
    "2. **추론 속도**: 어댑터 계산에 따른 추가적인 오버헤드 없음\n",
    "3. **프레임워크 호환성**: 서빙 프레임워크와의 호환성 향상\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import AutoPeftModelForCausalLM\n",
    "\n",
    "\n",
    "# CPU에서 PEFT 모델 불러오기\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",
    "# LoRA와 기본 모델 병합 후 저장\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. 모델 테스트 및 추론 수행\n",
    "\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'>추가 연습: LoRA 어댑터 불러오기</h2>\n",
    "    <p>예제 노트북에서 배운 내용을 활용하여 학습된 LoRA 어댑터를 불러오고 추론에 사용하세요.</p> \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "I5B494OdouLl"
   },
   "outputs": [],
   "source": [
    "# 메모리 확보\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",
    "# 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": [
    "몇 가지 프롬프트 예제로 테스트해보고 모델이 어떻게 작동하는지 확인해봅시다."
   ]
  },
  {
   "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
}