kagglehub を使った大規模言語モデル gemma のファインチューニングとモデル共有

はじめに
事前準備
モデルのダウンロード
ファインチューニングの準備
データの準備
ファインチューニングの実行
モデルの保存
モデルのアップロード
まとめ
ノートブック
参考サイト
1. 関連

はじめに

こんにちは。この記事では、Kaggle の新機能である Kaggle Models を使って、大規模言語モデル gemma をファインチューニングし、コミュニティで共有する方法を初心者向けに解説します。

Kaggle Models では、誰でも自分のモデルを公開し、他の人と共有することができます。今回は、gemma という大規模言語モデルを数学問題の解答生成タスクに適用する例を通して、モデルの共有方法を学んでいきましょう。

こちらの記事もおすすめ

Meta社の新しいオープンソースLLMモデル「Llama 3」がリリースされました

Meta社が、オープンソースの大規模言語モデル（LLM）ファミリー「Llama」の最新版である「Llama 3」を2024年4月18日にリリースしました。Llama 3モデルはHugging Faceからアクセス可能で、Hugging Fa...

LLama 3のSFTTrainer＋Weights & Biasesでファインチューニング

はじめにLLama 3は、Meta社が開発した大規模言語モデルです。高性能でありながら、一般的なGPUでも扱えるサイズのモデルが提供されています。このモデルをファインチューニングすることで、様々なタスクに適用できます。本記事では、Huggi...

事前準備

まずは、必要なライブラリをインストールしましょう。ここでは、モデルの管理に便利な kagglehub というライブラリを使用します。

!pip install kagglehub

モデルのダウンロード

Kaggle からモデルをダウンロードするには、以下のようなコードを実行します。

import kagglehub

# 認証を行う（他の認証方法も利用可能です）
kagglehub.login() 

# モデルの最新バージョンをダウンロード
path = kagglehub.model_download("google/gemma/pyTorch/2b")

# 特定のバージョン（ここではバージョン1）をダウンロード  
path = kagglehub.model_download("google/gemma/pyTorch/2b/1")

print("モデルファイルのパス:", path)

ファインチューニングの準備

次に、ファインチューニングに必要なライブラリをインポートし、学習済みモデルの保存先を指定します。

import os
import keras
import keras_nlp

FINETUNED_WEIGHTS_PATH = "gemma_finetune_demo.weights.h5"

gemma モデルを読み込みます。ここでは JAX バックエンドを使用します。

os.environ["KERAS_BACKEND"] = "jax"
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"]="1.00"

gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma_instruct_2b_en")
gemma_lm.summary()

データの準備

ファインチューニングには、数学問題とその解答のデータセットを使用します。ここでは、Kaggle の "AI Mathematical Olympiad Prize" コンペのデータを使います。

import pandas as pd

df1 = pd.read_csv("/kaggle/input/ai-mathematical-olympiad-prize/train.csv")
Math_data = []

for index, row in df1.iterrows():
    question, answer = row['problem'], row['answer']
    template = (f"""
    Context: You are an intelligent math tutor tasked with solving mathematical problems and explaining the solution steps in a clear and concise manner.

Problem: ${question}

Instructions:
- Carefully analyze the given problem and identify the key information, known values, and the unknown quantity to be found.
- Break down the problem into smaller steps if necessary, and apply the appropriate mathematical concepts, formulas, and operations to solve the problem.
- Show your step-by-step working, explicitly stating the reasoning behind each step.
- If relevant, provide additional explanations, examples, or visualizations to aid understanding.
- Finally, state the final answer to the problem clearly and concisely.

Solution: ${answer}

    """)
    Math_data.append(template)

ファインチューニングの実行

gemma モデルに LoRA (Low-Rank Adaptation) を適用し、軽量なファインチューニングを行います。

gemma_lm.backbone.enable_lora(rank=64)

gemma_lm.preprocessor.sequence_length = 512

optimizer = keras.optimizers.AdamW(
    learning_rate=5e-5,
    weight_decay=0.01,  
)
optimizer.exclude_from_weight_decay(var_names=["bias", "scale"])

gemma_lm.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=optimizer,
    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()], 
)

gemma_lm.fit(Math_data, epochs=24, batch_size=1)

モデルの保存

ファインチューニングしたモデルの重みを保存します。

gemma_lm.save_weights(FINETUNED_WEIGHTS_PATH)

モデルのアップロード

最後に、ファインチューニングしたモデルを Kaggle にアップロードして共有します。

handle = 'makimakiai/gemma/jax/gemma_demo0'

kagglehub.model_upload(handle, FINETUNED_WEIGHTS_PATH, license_name='Apache 2.0', version_notes='demo v0')

model_upload の第1引数にはモデルの識別子、第2引数には重みファイルのパス、license_name にはライセンス、version_notes にはバージョン情報を指定します。