ブログ

「GemmaとRayを使ったVertex AIでの機械学習入門」

GemmaとRayを使った効率的な機械学習モデルのトレーニング

機械学習の分野は日々進化しており、その中でGoogleが開発した「Gemma」と「Ray」は注目のツールです。この記事では、 GemmaをRayとVertex AI上で活用して、効率的に機械学習モデルをトレーニングする方法について詳しく解説します。

Gemmaとは？

Gemmaは、Geminiモデルで使用されている研究と技術から構築されたオープンモデルのファミリーです。このモデルは、テキスト生成やコードの補完、特定のタスクに対するファインチューニングが可能であり、さまざまなデバイスで実行できます。

Rayとは？

Rayは、AIとPythonアプリケーションのスケーリングのためのオープンソースフレームワークです。分散コンピューティングと並列処理を活用して、機械学習（ML）ワークフローを効率化するためのインフラストラクチャを提供しています。

このチュートリアルの目的

この記事を通じて、読者は「Gemma」をRayを使用してVertex AI上でトレーニングする方法を学びます。最終的には、効果的かつ効率的に機械学習モデルをトレーニングし、提供する方法を理解できます。

前提条件

  • Google Cloudプロジェクトを選択または作成します。
  • プロジェクトに対して請求が有効であることを確認します。
  • APIを有効にします。
  • ローカルでこのチュートリアルを実行するには、Cloud SDKをインストールする必要があります。

必要なリソース

データセット

このチュートリアルでは、Extreme Summarization (XSum)データセットを使用します。このデータセットは、抽象的な単一文書要約システムに関するものです。

Cloud Storageバケット

中間アーティファクト（データセットなど）を保存するために、ストレージバケットを作成する必要があります。次のコマンドを使用して新しいバケットを作成できます。

gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}

Dockerイメージリポジトリ

カスタムクラスタイメージを保存するために、Artifact RegistryにDockerリポジトリを作成します。次のコマンドを使用できます。

gcloud artifacts repositories create your-repo --repository-format=docker --location=your-region --description="Tutorial repository"

Vertex AI TensorBoardインスタンス

チューニングジョブを追跡して監視するためにTensorBoardインスタンスを作成します。以下のコマンドを使用して作成できます。

gcloud ai tensorboards create --display-name your-tensorboard --project your-project --region your-region

RayクラスターをVertex AI上に設定する

カスタムクラスタイメージの構築

RayをVertex AI上で使用するには、カスタムイメージのDockerfileを作成するか、事前に構築されたRay基本イメージを利用します。以下の手順でカスタムクラスタイメージを作成します。

  • 依存関係を含む要件ファイルを準備します。
  • カスタムイメージのDockerfileを作成します。
  • Cloud Buildを使用してRayクラスターのカスタムイメージをビルドします。

Rayクラスターの作成

次のPythonコードを使用して、Vertex AIでRayクラスターを作成できます。

ray_cluster_name = vertex_ray.create_ray_cluster(
    head_node_type=HEAD_NODE_TYPE,
    worker_node_types=WORKER_NODE_TYPES,
    custom_images=CUSTOM_IMAGES,
    cluster_name="your-cluster-name",
)

Gemmaのファインチューニング

Ray Trainを使用して、HuggingFace Transformersを配布してGemmaをファインチューニングできます。以下の手順で進めます。

  • トレーニング関数を定義します。
  • スケーリング構成を定義します。
  • TorchTrainerを起動してファインチューニングを開始します。

Gemmaのバリデーション

トレーニングが成功すれば、ローカルで生成された予測を確認できます。以下の手順を参考にしてください。

  • Rayジョブからすべてのチェックポイントをダウンロードします。
  • ExperimentAnalysisメソッドを用いて最良のチェックポイントを取得します。

モデル評価

ファインチューニングされたモデルは、ベースのGemmaモデルとの比較で評価し、ROUGE指標の改善を算出することができます。以下の手順で評価を行います。

improvements = {}
for rouge_metric, gemma_rouge in gemma_results.items():
    tuned_gemma_rouge = tuned_gemma_results[rouge_metric]
    if gemma_rouge != 0:
        improvement = ((tuned_gemma_rouge - gemma_rouge) / gemma_rouge) * 100
    else:
        improvement = None
    improvements[rouge_metric] = improvement

Gemmaモデルの提供

GemmaモデルをRay Dataを使ってオフライン予測のために提供します。Ray Dataを利用して分散データ読み取りを実施し、バッチ処理を行います。

まとめ

このチュートリアルでは、GemmaとRayを用いた機械学習モデルのトレーニングから、評価、提供まで一通りの流れを学びました。ユーザーはこのプロセスを通じて、データサイエンスや機械学習の実践的なスキルを高めることができるでしょう。

最後に、チュートリアルに使用したGoogle Cloudリソースは適切にクリーンアップすることを忘れないでください。これにより、無駄なコストを防ぐことができます。

ご覧いただきありがとうございました。GemmaとRayを使った素晴らしいプロジェクトに取り組むことを願っています。