OpenShift AIで分散学習を試してみる

こんにちは、Red Hatでソリューションアーキテクトをしている石川です。

過去に何度かこのブログの中でOpenShift Data Scienceについて取り上げてきましたが、 直近のリリースバージョンであるv2.4から機械学習における分散学習を実現する機能がTech Previwとして追加されました。 access.redhat.com

今回はこちらの機能をデプロイしてみてどういった仕組みで分散学習を実現しているのか見ていきたいと思います。 なおTech Preview機能についての制約についてはこちらを参照下さい。

またプロダクトの正式名称がRed Hat OpenShift Data Science(RHODS)からRed Hat OpenShift AI(RHOAI)に変更となりました。
ドキュメントなどまだ直っていない部分もありますが今後はこちらの名称を使っていきたいと思います。

そもそもRHOAIがどういった製品なのか、等については過去のブログにまとめてあるためこちらも合わせて参照下さい。

rheb.hatenablog.com
rheb.hatenablog.com
rheb.hatenablog.com

CodeFlare Projectについて

分散学習に関する機能はCodeFlareと呼ばれるOSSプロジェクトの中で開発が進められています。 CodeFlareプロジェクトには複数のコンポーネントで構成されます。
・CodeFlare SDK: Code Flareの各ツールをPythonで利用するためのSDK。
・MCAD: 学習のバッチ処理を行うためジョブのキューイングとディスパッチ。
・InstaScale: ジョブ要求に応じてオンデマンドにクラスタにGPUノードを追加。
・KubeRay: K8s上に分散処理を行うためのRayクラスタを構築。

以下では順番を追って、インストール方法から各コンポーネントがどのように動作するのか見ていきたいと思います。

Operatorのインストール

まずOpenShiftコンソールのOperatorHubより最新のOperatorをインストールします。

CodeFlare関連の機能は最新のv2.4よりTPの対象となるため正しいバージョンをインストールできているか確認しましょう。

Operatorのインストールが完了したら必要なコンポーネントをインストールするためDataScienceCluster(DSC) CRを作成します。 Operatorのバージョンがv1.xの時はKFDef CRを作成し、インストール対象のコンポーネントを選択していましたが、v2からはこのDSCにて設定するよう変更となりました。

デフォルトの設定では.spec.components.codeflare.managementStateがRemovedになっているため、画像の通りManagedに変更します。同じく.spec.components.ray.managementStateについてもManagedとしましょう。

設定したDSCを作成すると必要なコンポーネントがOpenShiftにインストールされます。

Ray Clusterの定義と立ち上げ

CodeFalreでは分散学習用の基盤としてOSSのRayをOpenShift環境で立ち上げます。Rayは汎用的な分散処理を実現するOSSのフレームワークであり、複数のサーバーリソースをまとめることで、機械学習の各タスク（前処理、ハイパーパラメーターサーチ、学習、推論、等）における処理を効率化することが可能となります。

Ray自体はコンテナ環境以外でも利用することが可能ですが、KubeRay Operatorを利用することでKubernetes環境でも利用することができます。 CodeFlareを利用する場合、このKubeRay Operatorも自動でインストールされます。

まずユーザーが行うことはCodeFlare SDKを使用し、必要なRay Clusterのリソースを定義することです。
SDKについてはpip install codeflare-sdkでインストールできます。

PythonスクリプトとしてRay Clusterの情報を定義し、その情報をMCADに対して送信します。 こちらのリポジトリにCodeFlareを試すためのサンプルNotebookファイルがあります。 以下はそこから一部抜粋、改変したコードとなります。

# CodeFlare SDKのインポート
from codeflare_sdk.cluster.cluster import Cluster, ClusterConfiguration
from codeflare_sdk.cluster.auth import TokenAuthentication

# OpenShiftへのログイン情報定義とログイン
auth = TokenAuthentication(
    token = "SHA256...", # oc whoami -t の実行結果 
    server = "https://api.sample.openshiftapps.com:6443", # OpenShift APIサーバーのURL
    skip_tls=False
)
auth.login()

# Ray Clusterの定義
cluster = Cluster(ClusterConfiguration(
    name='raytest',
    namespace='default', # Ray Clusterが立ち上がるNameSpace
    num_workers=2, # Ray Clusterを構成するWorker Podの数
    # Rayの各Workerに割り当てるリソース (CPU/Memory/GPU)
    min_cpus=1,
    max_cpus=1,
    min_memory=4,
    max_memory=4,
    num_gpus=0,
    image="quay.io/project-codeflare/ray:latest-py39-cu118", # Worker Podのイメージ
    instascale=False # InstaScaleの利用有無
))

# Cluster情報のMCADへの送信
cluster.up()

cluster.wait_ready()

# 作成されたRay Clusterの表示
cluster.details()

SDKを通じてOpenShiftにログインし、clusterオブジェクトの中で必要なリソースの定義を行っています。 SDKの詳細な利用方法については以下のページを参照下さい。
project-codeflare.github.io

ここからRay Clusterの立ち上げが開始します。 OpenShift内に先ほど定義したRay Clusterを立ち上げるのに十分なリソースがある場合、すぐに作成が始まりますが、そうでない場合Clusterの立ち上げリクエストはMCADの中でキューイングされた状態となります。 この時に活躍するのがCodeFlareのコンポーネントの一つであるInstaScaleです。 上記のコードの中でclusterオブジェクトの定義の中でInstaScaleを有効にすると不足するリソース分のNodeをOpenShiftクラスターに自動で追加することができます。

Ray Clusterを作成するのに必要なWorker Nodeを動的にOpenShiftに追加し、ジョブ実行が完了したら追加したNodeを自動で削除することが可能となります。オンデマンドに追加Nodeを調達できるパブリッククラウド環境で有効な機能と言えるでしょう。

RHOAIのv2.4ではInstaScaleがTPの対象機能となっていないため、ここでは概念の紹介のみに留めたいと思います。 access.redhat.com

Ray Clusterが作成されるとcluster.details()からUIへのアクセスが可能となります。

分散学習の実行

作成したRay Clusterの上で学習を実行するには大きく2つの方法があります。
a. CodeFlare SDKを通じてDDPJobを定義し実行する
b. Rayライブラリを使用し学習を実行する

今回は一つ目のCode Flare SDKが提供するDDPJobDefinitionクラスでジョブの定義を行います。 こちらのNotebookファイルからジョブの定義方法を見てみましょう。

from codeflare_sdk.job.jobs import DDPJobDefinition

# ジョブの定義
jobdef = DDPJobDefinition(
    name="mnisttest",
    script="mnist.py", # 学習実行用のスクリプト
    scheduler_args={"requirements": "requirements.txt"} # 学習時の引数
)

# ジョブの実行
job = jobdef.submit(cluster)

DDPJobDefinitionの中で学習用のスクリプトであるmnist.pyと、Rayの各Worker Podで追加でインストールするライブラリをrequirements.txtとして渡しています。 これにより必要な依存関係のインストールを行なった上でジョブを実行します。

DDPJobDefinitionではTorchXが提供するDDP(Distributed Data Parallel)の仕組みを活用しており、Ray Clusterの各Worker上にデータをシャードするデータパラレルな分散学習を実現します。
pytorch.org

ジョブの実行状況についてはRayのダッシュボードや、job.status()により確認できます。 ジョブが完了したらcluster.down()にて作成したRay Clsuterを削除します。 これでCodeFalreを利用した分散学習の一連のプロセスが完了しました。

もう一つの方法であるRayの提供するライブラリを利用する方法については、 こちらを参考にしてみて下さい。

まとめ

今回はOpenShift AIでTech PreviewとなったCodeFlareによる分散学習についてご紹介しました。LLMをはじめとした大規模モデルの活用や、そのファインチューニングのニーズが増えるにつれ、こうした機能へのニーズはより高まってくると思われます。今後も機能のアップデートなど紹介していきたいと思います。