見出し画像

WandB LaunchのSageMaker連携によるLLMの自動評価

Weights & Biases Japan
Weights & Biases のNoteをフォローしてください

こんにちは、W&Bの山本です。今回は、WandB Launchを使用してAmazon SageMakerでllm-jp-evalの評価を自動化し、LLM勉強会(llm-jp)にご提供したユースケースを紹介したいと思います。
WandB LaunchのSageMaker連携では例えばモデルIDを入力して１クリックするだけ、あるいはレジストリから自動的に評価ジョブをSageMakerに投げることで、LLMモデル評価をインスタンスのプロビジョニングも含めて自動実行することができます。では、順に見ていきましょう。


WandB Launchとは

WandB Launchは、MLエンジニアの生産性向上のために開発されたツールで、Docker、SageMaker、VertexAI、Kubernetesなどの様々なGPUインフラをシームレスに抽象化します。これにより、MLエンジニアはインフラストラクチャの詳細を気にすることなく、機械学習ワークフローに集中することができます。

WandB Launchの仕組みの概念図

Launchは、ジョブのソースコード、依存関係、入力データを含む再現可能なRunのブループリントを作成します。これらのジョブは、W&Bのアーティファクトとして保存され、様々なコンピューティングリソースに送信できます。Launchは、Dockerイメージ、Gitリポジトリ、またはコードアーティファクトからジョブを作成することができます。

ジョブを実行するには、Launchキューを作成する必要があります。キューは、ジョブを実行するためのコンピューティングリソースを定義します。現在、Launchは、Docker、Kubernetes、AWS SageMaker、GCP Vertex AIをサポートしています。キューを作成したら、W&B UIまたはCLIを使用してジョブをキューに送信できます。

このように多様な計算リソース上に任意の機械学習ジョブをスケーラブルに展開可能なWandB Launchですが、今回は冒頭に述べた通り、SageMaker上にキューを構成する例をご紹介いたします。

SageMaker用のLaunchキューの設定

LaunchをSageMakerと組み合わせて使用する場合、キューの設定が重要になります。SageMakerのキューを作成するには、以下の手順を実行します。

  1. W&B Appにアクセスし、Launchページに移動

  2. "Create Queue"ボタンをクリック

  3. キューを作成したいEntityを選択

  4. キューの名前を入力

  5. ResourceとしてSageMakerを選択

  6. Configurationフィールドに、SageMakerジョブの設定を入力（YAML or JSON）

Configurationでは、SageMakerのCreateTrainingJobリクエストで使用するパラメータを指定します。以下の項目などを指定しますが、詳細は公式ドキュメントをご参照ください。

  • RoleArn: SageMakerがジョブを実行するためのIAMロールのARN

  • ResourceConfig: インスタンスタイプ、インスタンス数、ストレージサイズなどのリソース設定

  • OutputDataConfig.S3OutputPath: ジョブの出力を保存するS3のパス

  • StoppingCondition: ジョブの最大実行時間などの停止条件

これらの設定は、各ジョブの要件に応じて適切な値を指定してください。

LaunchをSageMakerと組み合わせて使用する大きな利点の一つは、インスタンスのプロビジョニングと終了が自動的に行われることです。ジョブが開始されると、SageMakerがキューの設定に基づいてインスタンスを起動し、ジョブの実行環境を準備します。ジョブが完了すると、インスタンスは自動的に終了されます。これにより、ユーザーはインフラストラクチャの管理に煩わされることなく、機械学習のタスクに専念できます。

Launch Agentの設定と起動

ジョブをSageMakerのキューに投入し、実行状況をモニタリングするには、Launch Agentを起動する必要があります。Agentは、比較的軽量なコンピューティングリソース上で動作させることができます。例えば、小型のCPUインスタンス（t2.microなど）やローカルマシンで十分です。

Agentの設定は、設定ファイル（launch-config.yaml）で行います。主な設定項目は以下の通りです。

environment:
  type: aws
  region: <aws-region>  # 例: us-west-2
registry:
  type: ecr
  repository: <ecr-repo-name>

environmentセクションでは、AWSの認証情報とリージョンを指定します。registryセクションでは、ジョブのDockerイメージを保存するECRリポジトリを指定します。

設定ファイルを用意したら、以下のコマンドでAgentを起動します。

wandb launch-agent -e <entity> -q <queue-name> -c <path-to-config>

llm-jp-evalの評価自動化

今回のユースケースでは、SageMakerにキューを作成し、llm-jp-evalの評価を自動化しました。llm-jp-evalは、日本語言語モデルの評価ベンチマークで、LLM勉強会(llm-jp)のコミュニティによって開発されています。

手順の概要

以下にWandB LaunchのSageMaker連携について設定の手順を示します。概要としては下図のようにJob, キュー、エージェントを作成して連携を行います。

Launch on SageMakerの概念図

  • SageMaker用のDockerイメージベースのJobを作成 最初のステップは、SageMaker用のDockerイメージベースのJobを作成することです。Jobは、評価に必要なソースコード、依存関係、設定を含むDockerイメージから作成します。Dockerfileを作成し、必要なライブラリとllm-jp-evalのソースコードをインストールします。そして、DockerイメージをビルドしてAmazon ECRにプッシュします。例えば、本ユースケースでは以下のようなDockerfileを使用しました。なお、WandB Launchのジョブとして使用するためにはconfigの情報がコード内に適切に渡されるように注意しましょう。

# ベースイメージの選択
FROM nvidia/cuda:12.0.1-cudnn8-runtime-ubuntu22.04

# Pythonとpipのインストール
RUN apt update && \
    apt install -y bash \
                   build-essential \
                   git \
	                   curl \
                   ca-certificates \
                   python3 \
                   python3-pip && \
    rm -rf /var/lib/apt/lists

RUN python3 -m pip install --no-cache-dir --upgrade pip

# 作業ディレクトリの設定
WORKDIR /app

# リポジトリのコピー＆インストール
COPY . /app
RUN pip install .

# データセットの前処理スクリプトの実行
RUN python3 scripts/preprocess_dataset.py --dataset-name all --output-dir eval_datasets --version-name 1.2.0

# コンテナ起動時に実行されるコマンド
ENTRYPOINT ["python3", "scripts/evaluate_llm.py", "-cn", "config.yaml", "dataset_dir=eval_datasets/1.2.0/evaluation/test/"]

  • SageMakerのIAMロールとインスタンスの設定 次に、SageMakerがジョブを実行するためのIAMロールとインスタンスを設定します。IAMロールには、Amazon ECR、Amazon S3、CloudWatch、SageMakerへのアクセス権限を付与します。インスタンスタイプ、インスタンス数、ストレージサイズなどのリソース設定を定義します。

  • Launch Agentの設定と起動 Launch Agentは、SageMakerへのジョブの投入とモニタリングを担当します。Agentの設定ファイル(launch-config.yaml)を作成し、AWS認証情報、リージョン、ECRリポジトリなどを指定します。そして、wandb launch-agentコマンドでAgentを起動します。

  • Launch UIまたはCLIからジョブをキューに投入 ジョブを作成したら、W&B UIまたはCLIを使用してSageMakerのキューにジョブを投入します。UIでは、ジョブの設定をGUIで変更でき、"Launch now"ボタンをクリックするだけでジョブが開始されます。CLIでは、wandb launchコマンドにジョブ名、キュー名、設定ファイルのパスを指定してジョブを投入します。

W&B UIからSageMakerのキューにジョブを投入する。

  • ジョブの実行とモニタリング ジョブがキューに投入されると、Launch AgentがSageMakerにトレーニングジョブを作成します。SageMakerは、キューの設定に基づいてインスタンスを起動し、ジョブを実行します。ジョブの実行状況は、W&BダッシュボードまたはAWS SageMakerコンソールでモニタリングできます。ジョブが完了すると、評価結果がW&Bプロジェクトにログされ、簡単に確認できます。また、インスタンスは自動的に終了されるため、リソースの無駄を防ぐことができます。

WandBダッシュボード上でLLM評価の結果をインタラクティブに確認できる。
ジョブの実行履歴はCloudWatchからも確認することができ、連携ができていることがわかる。

今回のユースケースでは、AWS Japan様よりご提供頂いた計算リソース上に弊社がLaunchキューを作成し、llm-jp-evalの評価ジョブを実行できるように構築いたしました。これによって、モデルの学習やチューニングを行う研究者やエンジニアは、評価に用いるインスタンスの用意や評価ツールの準備の手間を気にすることなく、最新の学習済みモデルをmodel_idを指定するだけで簡単にllm-jp-evalで評価することができるようになりました。評価プロセスが自動化されたことで、研究開発のイテレーションを速く回すことができ、生産性が大幅に向上しました。

まとめ

WandB Launchは、機械学習ワークフローを自動化し、スケーラブルな実行環境を提供する強力なツールです。今回は、LLM勉強会(llm-jp)でのllm-jp-evalの評価自動化を例に、LaunchとSageMakerを組み合わせた使い方を紹介しました。

WandB Launchは、その他にも例えばOptunaとの連携によって高度でスケーラブルな最適化の実行や、モデルレジストリと連携してモデルの保存などの事前に指定したトリガーによって自動的に起動するなど、活用の幅と製品機能の両面において更なる可能性を有しています。機械学習ワークフローの自動化にご興味のある方は、ぜひお気軽にお問い合わせください。一緒にLaunchの可能性を探求しましょう！

この記事が気に入ったらサポートをしてみませんか？