wandbを使ったLLM実験管理の練習

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

https://qiita.com/t-hashiguchi/items/9f3b394ca0ae1c7e4d02



## 想定読者
- 大規模言語モデルを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評価をしているので、ある程度のバージョンであれば特に問題ないはず。

https://platform.openai.com/docs/guides/structured-outputs

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

https://github.com/t-hashiguchi1995/wandb_tutorial


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

https://wandb.ai/tomoya1995/wandb_tutorial/workspace?nw=nwusertomoya1995

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

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

   ```python
   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` の使用:**
   ```python
    training_args = TrainingArguments(
    # ...
        report_to="wandb"
    )
   ```
   - `report_to="wandb"`: `transformers` の `TrainingArguments` を設定する際に、`wandb` に学習のログを記録するように設定します。これにより、損失やメトリクスなどの情報を `wandb` で確認できるようになります。

**3. `wandb` を使ったモデルアーティファクトの保存:**
  
   ```python
     os.environ["WANDB_LOG_MODEL"] = "checkpoint"
   ```

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

**4. `wandb` へのカスタムデータの記録:**
   ```python
        # ... 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` のダッシュボード上でグラフを可視化できます。


# 参考情報
https://docs.wandb.ai/ja/guides/integrations/huggingface/

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