はじめに

以下の記事で作成した内容をwandbで実験管理しようと思い、wandbのドキュメントを読みながら、アップデートしたので、共有。

想定読者

大規模言語モデルをLoRAでファインチューニングしたい人
wandbを使って実験管理をしたいが、使い方が分からない人

元記事からの更新点含めた共有なので、LoRAについては説明しない。おもにwandbでの実験管理について紹介。
ドキュメントをざっと見て実装したので、もっとよい実験管理手法があるかも...

実行環境

transformers==4.38.2
peft==0.10.0
trl==0.7.10
torch==2.4.1+cu124
学習を動かすのに影響があるバージョンのみ記載。

openaiモジュールについては「Structured Outputs」を使って学習後のLLM評価をしているので、ある程度のバージョンであれば特に問題ないはず。

コード

00_wandb_tutorialのコードは、wandb (Weights & Biases) を利用して、大規模言語モデル（LLM）のファインチューニング実験の追跡、ログ記録、および可視化を行うためのものです。

WandBのワークスペース

実行結果があると分かりやすいかと思うので、作成されるワークスペースを共有します。

コードの説明

以下に、wandbに関連するコード部分とその説明をまとめます。

1. wandb のインポートと初期化:
runのoverviewに登録される内容を主に記載しています。

import wandb
# ...
wandb.login(key=config.WANDB_API_KEY) # wandbへのログイン

# ...
PROJECT = "wandb_tutorial"
NAME = f"{MODEL}_lora_experiment"
NOTES = f"{MODEL}モデルに対してLoRAを適用した実験"
TAGS = [MODEL, "lora", "fine-tuning"]

params = {
     "learning_rate": LEARNING_RATE,
     "architecture": MODEL,
     "dataset": DATASET,
     "train_samples": TRAIN_SAMPLES,
     "eval_samples": EVAL_SAMPLES,
     "batch_size": BATCH_SIZE,
     "epochs": EPOCHS,
     "optimizer": OPTIMIZER,
     "lora_rank": LORA_RANK,
     "lora_alpha": LORA_ALPHA,
     "lora_dropout": LORA_DROPOUT,
     "target_modules": TARGET_MODULES,
     "max_grad_norm": MAX_GRAD_NORM,
     "warmup_ratio": WARMUP_RATIO,
     "weight_decay": WEIGHT_DECAY,
     "lr_scheduler_type": LR_SCHEDULER_TYPE,
 }

 # wandbの初期化
 wandb.init(
     project="wandb_tutorial", # プロジェクト名
     name=NAME, # 実験の名前
     config=params, # 設定パラメータ
     notes=NOTES, # 実験の説明
     tags=TAGS,  # タグ
 )

import wandb: wandb ライブラリをインポートします。
wandb.login(key=config.WANDB_API_KEY): wandb にログインします。config.WANDB_API_KEYには自身のAPIキーが格納されている必要があります。
wandb.init(...): 新しい wandb 実験を開始します。
- project="wandb_tutorial": 実験を所属させるプロジェクト名です。
- name=NAME: 実験の名前です。
- config=params: 実験の設定パラメータを wandb に記録します。
- notes=NOTES: 実験に関する説明を記述します。
- tags=TAGS: 実験にタグを付与し、wandb で検索しやすくします。

2. TrainingArguments での wandb の使用:

 training_args = TrainingArguments(
 # ...
     report_to="wandb"
 )

report_to="wandb": transformers の TrainingArguments を設定する際に、wandb に学習のログを記録するように設定します。これにより、損失やメトリクスなどの情報を wandb で確認できるようになります。

3. wandb を使ったモデルアーティファクトの保存:

  os.environ["WANDB_LOG_MODEL"] = "checkpoint"

os.environ["WANDB_LOG_MODEL"] = "checkpoint": 環境変数 WANDB_LOG_MODEL を "checkpoint" に設定することで、学習中にモデルのチェックポイントを自動的に wandb にアーティファクトとして保存します。

4. wandb へのカスタムデータの記録:

     # ... MTBenchでの評価結果の計算
 
 artifact = wandb.use_artifact("mtbench_score_artifact:latest", type="dataset")
 artifact_dir = artifact.download()
 with open(f"{artifact_dir}/mtbench_score_key.table.json") as f:
     tjs = json.load(f)
 output_table = wandb.Table.from_json(json_obj=tjs, source_artifact=artifact)
 output_df = pd.DataFrame(data=output_table.data, columns=output_table.columns)
 output_df = pd.concat([output_df, df], ignore_index=True)
 artifact = wandb.Artifact("mtbench_score_artifact", type="dataset")
 artifact.add(wandb.Table(dataframe=output_df), "mtbench_score_key")  
 # あるいは df.to_csv() してファイルとして add_file() するなど
 wandb.log_artifact(artifact)

 # ...レーダーチャートを作成
 wandb.log({"MTBench_スコア": fig})

MTBenchで評価した結果をDataFrameに格納し、wandb テーブルとして記録します。
wandb.use_artifact：以前のアーティファクトを読み込みます。
wandb.Artifactと wandb.log_artifact: 新しいアーティファクトを作成して、記録します。
wandb.log({"MTBench_スコア": fig}): 作成したレーダーチャートを wandb にログとして記録します。これにより、wandb のダッシュボード上でグラフを可視化できます。

参考情報

最後に

残りの機能として気になっているのがアラート送信の機能とカスタムコールバックの機能。
LLMバージョンでのカスタムコールバックのコード誰か書いてください...