AI Platform(GCP)でGPU 100個同時に使いテンションあがった

12052 ワード
MLEngine gcp AWS SageMaker AIPlatform AWS テキストリンク

ひゃっはー

概要

機械学習基盤系のサービスとしてAWSのSageMakerが有名かと思いますが
本投稿ではGCPの類似サービスAI Platformを使ってGPU 100並列同時に動かした話をします。

AI Platformでできること

SageMakerと似た感じで、AI Platformでは以下のようにうたっています。
（公式ドキュメントより引用）

AI Platform を使用して、ML ワークフローの次の段階を管理できます。

  • データで ML モデルをトレーニングする。
    • モデルをトレーニング
    • モデルの精度を評価する
    • ハイパーパラメータを調整する
  • トレーニング済みモデルをデプロイする。
  • モデルに予測リクエストを送信する。
    • オンライン予測
    • バッチ予測（TensorFlow のみ）
  • 継続的に予測をモニタリングする。
  • モデルとモデル バージョンを管理する。

いわゆる機械学習基盤として使えますよ便利ですね、ってサービスなんですが
単に時間のかかるバッチジョブ（MLの場合、特に学習ジョブ）をサーバレスで走行する環境として使うだけでも割とその恩恵を受けられます。
（ というかマジメに機械学習基盤として使うにはちょっとまだ機能連携が足りていないような... ）

GPUをGCEなどのIaaSで利用するとインスタンスの起動・停止がめんどいので
学習が走行されている時間だけ課金されるマネージドなジョブ実行ができるのはありがたいです。

バッチを走行するだけなら他にも方法あるやん、って思われるかもしれないですが
意外と「GPUが使える」かつ「時間制限がない」ってサービスはあまりなくて
他の選択肢だとGKEでk8s jobくらいなのかな、と思います（他にもあったらごめんなさい...）

なので、GKEよりは手軽にML学習のジョブを走らせられるのは、それだけでもメリットと感じます。

GPUの同時利用上限

AI Platformには同時に利用できるGPU数に上限が設けられています。

  • Tesla K80 GPUs: 30
  • Tesla P100 GPUs: 30
  • Tesla V100 GPUs: 8

SageMakerのデフォルト上限は下記のようです。（最新版は公式ドキュメントを参照してください）

インスタンスタイプ GPU 同時利用インスタンス数上限
ml.p2.xlarge K80 * 1 1
ml.p2.8xlarge K80 * 8 1
ml.p2.16xlarge K80 * 16 1
ml.p3.2xlarge V100 * 1 2
ml.p3.8xlarge V100 * 4 2
ml.p3.16xlarge V100 * 8 2

合計のGPU数はさほどAI Platformと変わらないように思えますが
1つのインスタンスにGPUが複数搭載されていても、GPU分散で学習させないと意味がないです。
そういう意味だと、1GPUのジョブを何個も同時に実行したい場合にはSageMakerのデフォルトではちょっと足りないです。
（もちろんSageMakerでも上限を増やす申請はできます）

上限を増やす

通常、GCPの上限を増やす際の申請は、GCPコンソール上の「IAMと管理」=>「割り当て」から行うのですが
AI PlatformのGPU上限はなぜか別管理されており、ここのGoogleフォームから行います。
（公式ドキュメント参照）

で、私はここで "70 P100s", "30 V100s" と申請し、無事通りました。
（申請から承認まで大体1週間くらいかかりました）

これでGPUめっちゃ同時に使う環境が手に入りました。

ユースケース

ところでGPUを同時にたくさん利用したいケースってどんなんでしょうか。

  • モデル実装、パラメータ、データ、等のバリエーションを短期間で試したいとき
  • GPU分散で学習させたいとき
  • 上記のようなタスクが複数人同時に発生するとき
  • Qiitaで釣りっぽいタイトルを書きたいとき

とかですかね。用法・用量を守って正しく使いましょう。

ちなみに私はデータのバリエーションがたくさんあって、一気に複数のモデルを作って比較したい
というちゃんとした理由は一応ありました。

学習ジョブの実行

AI Platformでの学習ジョブ実行手順はざっくり以下です。

  1. 「pythonのコマンドとして実行できる形式でパッケージ」 or 「Dockerイメージを作成」
  2. 「パッケージをGCSにアップロード」 or 「DockerイメージをGCRにpush」
  3. 学習データをGCSにアップロード
  4. gcloudコマンドでAI Platformジョブを開始

以前はDockerコンテナでの実行はベータ版だったのですが
今はGAになっていますし、Dockerがやりやすいでしょう。

以下ではDockerコンテナを用いた手順を記載
と思ったのですが、 時間がないので 公式チュートリアルが充実しているので
そちらにリンクを貼っておき、「？」となりそうなところだけ補足します。（みんな大好きMNISTです）
基本的にチュートリアルのコピペで動くはずです。

また、単にGPUを使ったジョブ実行を体験したいだけなら、↑のチュートリアルすべてを実施する必要はありません。

具体的には

だけで可能です。

Dockerfile

チュートリアルなんで素直にそのまま使えばいいのですが、ちょっとだけ補足。

FROM nvidia/cuda:9.0-cudnn7-runtime

古い。けどまぁ無視していいかと。
例えばTensorflow 1.12以上を使いたいときとかはCUDA 10のイメージを指定しましょう。

RUN wget -nv \
    https://dl.google.com/dl/cloudsdk/release/google-cloud-sdk.tar.gz && \
    mkdir /root/tools && \
    tar xvzf google-cloud-sdk.tar.gz -C /root/tools && \
    rm google-cloud-sdk.tar.gz && \
    /root/tools/google-cloud-sdk/install.sh --usage-reporting=false \
        --path-update=false --bash-completion=false \
        --disable-installation-options && \
    rm -rf /root/.config/* && \
    ln -s /root/.config /config && \
    # Remove the backup directory that gcloud creates
    rm -rf /root/tools/google-cloud-sdk/.install/.backup

# Path configuration
ENV PATH $PATH:/root/tools/google-cloud-sdk/bin
# Make sure gsutil will use the default service account
RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg

Dockerコンテナではなく、Pythonパッケージをアップロードする方法でやる場合は
すでに色々とインストールされているインスタンスで実行されるのですが
Dockerコンテナ内では当然必要なツール等は自分でインストールしておく必要があります。

AI PlatformはGCSを入出力ファイル置き場として使うので、
ここではそれに必要なgcloud SDKのインストール、およびgsutilコマンドの認証設定を行っています。

ジョブの実行

DockerイメージをGCRにpushしたら、あとはgcloudコマンドでAI Platformのジョブを実行するだけです。かんたんですね。

$ export MODEL_DIR=pytorch_model_gpu_$(date +%Y%m%d_%H%M%S)
$ export REGION=us-central1
$ export JOB_NAME=custom_container_job_gpu_$(date +%Y%m%d_%H%M%S)

と変数を定義しておいて、下記を実行。


$ gcloud ai-platform jobs submit training $JOB_NAME \
  --scale-tier BASIC_GPU \
  --region $REGION \
  --master-image-uri $IMAGE_URI \
  -- \
  --epochs=5 \
  --model-dir=gs://$BUCKET_NAME/$MODEL_DIR

  • --scale-tier
    ジョブを実行するインスタンスのスペックを指定します。（スケール階層と呼ばれてます）
    BASIC_GPU はk80が1個アサインされたインスタンスが使われます。
    指定できるスケール階層は こちら。

  • --region
    ジョブを実行するリージョンを選びます。
    アクセスするGCSと同一リージョンを指定しなければならない、という謎制約があります。（ファイル転送の速度を考慮してくれてる、というポジティブ解釈）
    また、リージョンによっては使えるGPUの種類が違ったりします。

  • -- 以降の引数は、gcloudコマンドの引数ではなく、実行するpythonコマンドの引数として与えられます。
    チュートリアルのソースコードを見ればわかるのですが
    --epochs と --model-dir という引数がコマンドに渡されて実行されることになります。
    自作コマンドを作るときも、こんな感じで可変なパラメータを引数化します。

上記のgcloudコマンドを実行して

jobId: custom_container_job_gpu_20191219_142126
state: QUEUED

と出力されれば、ジョブ発行は成功です。

チュートリアルでは

$ gcloud ai-platform jobs describe $JOB_NAME
$ gcloud ai-platform jobs stream-logs $JOB_NAME

でジョブの進捗を確認しよう、となっていますが
GCPコンソール上で AI Platformのジョブ一覧を見たほうが簡単かもです。

あとはGPUの数だけジョブを発行しまくれば、ひゃっはーできます。

最後に

タイトルから釣りっぽい投稿になってしまい、反省しています。

そもそもGPUを100個同時に使うケースはあまりないにしても、
GPUを使いたいとき使いたいだけ利用できるサーバレスタイプの環境として、
という観点だけでもAI Platformを使ってみる価値はあるかもです。

また、今回取り上げませんでしたがハイパーパラメータの自動調整機能を使うと、複数GPUを一気に使えるメリットもあるってもんです。
MNISTを100並列で動かしてどうすんねん、という感じなので、もっと有効に使っていきましょう。
（一応私ももっとマジメに使ってます）