はじめに
以下の記事で作成した内容を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バージョンでのカスタムコールバックのコード誰か書いてください...