ja/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（低ランク適応）アダプタを使用して大規模言語モデルを効率的に微調整する方法を示します。LoRAは、次のようなパラメータ効率の良い微調整技術です：\n",
    "- 事前学習されたモデルの重みを固定\n",
    "- 注意層に小さな学習可能なランク分解行列を追加\n",
    "- 通常、学習可能なパラメータを約90％削減\n",
    "- メモリ効率を維持しながらモデル性能を維持\n",
    "\n",
    "以下の内容をカバーします：\n",
    "1. 開発環境のセットアップとLoRA設定\n",
    "2. アダプタトレーニング用のデータセットの作成と準備\n",
    "3. `trl`と`SFTTrainer`を使用してLoRAアダプタで微調整\n",
    "4. モデルのテストとアダプタの統合（オプション）\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fXqd9BXgouLi"
   },
   "source": [
    "## 1. 開発環境のセットアップ\n",
    "\n",
    "最初のステップは、Hugging FaceライブラリとPytorchをインストールすることです。`trl`、`transformers`、`datasets`を含みます。`trl`について聞いたことがない場合でも心配ありません。これは`transformers`と`datasets`の上に構築された新しいライブラリで、微調整、RLHF、オープンLLMの調整を容易にします。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tKvGVxImouLi"
   },
   "outputs": [],
   "source": [
    "# Google Colabでの要件のインストール\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "# Hugging Faceへの認証\n",
    "\n",
    "from huggingface_hub import login\n",
    "\n",
    "login()\n",
    "\n",
    "# 便利のため、Hugging Faceのトークンを環境変数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: パスと名前のパラメータを使用してデータセットと設定を定義\n",
    "dataset = load_dataset(path=\"HuggingFaceTB/smoltalk\", name=\"everyday-conversations\")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9TOhJdtsouLk"
   },
   "source": [
    "## 3. `trl`と`SFTTrainer`を使用してLLMをLoRAで微調整\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",
    "例としてLoRAを使用します。これは、LoRAと4ビット量子化を組み合わせて、性能を犠牲にせずにメモリ使用量をさらに削減します。セットアップには次の手順が必要です：\n",
    "1. LoRA設定（ランク、アルファ、ドロップアウト）を定義\n",
    "2. PEFT設定で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": [
    "# トレーニング設定\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\",  # 効率のために融合されたAdamWを使用\n",
    "    learning_rate=2e-4,  # 学習率（QLoRA論文）\n",
    "    max_grad_norm=0.3,  # 勾配クリッピングの閾値\n",
    "    # 学習率スケジュール\n",
    "    warmup_ratio=0.03,  # ウォームアップのステップの割合\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,  # HuggingFace Hubにプッシュしない\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設定でSFTTrainerを作成\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=dataset[\"train\"],\n",
    "    peft_config=peft_config,  # LoRA設定\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": [
    "Flash Attentionを使用して15kサンプルのデータセットで3エポックのトレーニングを行った結果、`g5.2xlarge`で4:14:36かかりました。このインスタンスのコストは`1.21$/h`で、合計コストは約`5.3$`です。\n",
    "\n"
   ]
  },
  {
   "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",
    "トレーニングが完了したら、モデルをテストしたいと思います。元のデータセットから異なるサンプルを読み込み、単純なループと精度をメトリックとして使用してモデルを評価します。\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",
    "    \"ドイツの首都はどこですか？その理由と過去に異なっていたかどうかを説明してください。\",\n",
    "    \"数の階乗を計算するPython関数を書いてください。\",\n",
    "    \"長さ25フィート、幅15フィートの長方形の庭があります。庭全体を囲むフェンスを作りたい場合、何フィートのフェンスが必要ですか？\",\n",
    "    \"果物と野菜の違いは何ですか？それぞれの例を挙げてください。\",\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
}