Amazon Web Services ブログ

TorchElastic Controller for Kubernetes を使った耐障害性のある機械学習分散トレーニング



はじめに

機械学習チームが Kubernetes を使用すると、Amazon EC2 P3 などの強力な GPU インスタンスのフリート全体に分散したトレーニングジョブを実行でき、トレーニング時間を数日から数時間に短縮することができます。しかし、Kubernetes によく関連付けられている従来のマイクロサービスベースのアプリケーションと比較すると、分散トレーニングには制限があります。分散トレーニングジョブは耐障害性がありませんし、ノードの障害または回復によってトレーニングが中断された場合、ジョブを続行できません。それに加え、要求されたリソースをすべて取得しないとジョブを開始できず、また、再起動せずに拡張や縮小することもできません。このように、回復性と柔軟性が欠如しているため、トレーニングの時間とコストが増大します。

この投稿では、AWS の Kubernetes チームと Facebook の PyTorch チーム間の新しいオープンソースコラボレーションである TorchElastic Controller for Kubernetesについて説明します。これで、上記のような制限に対処したり、PyTorch ビルドモデルや Kubernetes 分散トレーニングの新機能を利用できます。たとえば、EC2 スポットインスタンスでのトレーニング、ハードウェア障害に耐性のあるジョブの実行、クラスターリソースの可用性に基づいた動的なジョブのスケーリングなどが可能となります。

PyTorch Elastic と Kubernetes の統合

PyTorch Elastic は大規模な深層学習モデルをトレーニングするためのライブラリで、動的な可用性に基づいたコンピューティングリソースのスケーリングに不可欠です。PyTorch ジョブを作成するためのプリミティブとインターフェースで、複数のマシンで伸縮性を持ちつつ実行できるようになります。最小限のワーカーが集まるとすぐに分散ジョブを開始でき、停止または再起動を実行しなくても最大数のワーカーまで拡張できます。伸縮性とは、トレーニングジョブ中にフレームワークがノード数を拡張または縮小する機能です。耐障害性とは、ジョブを再起動せずに、フレームワークがトレーニングジョブ中に失敗したノードを検出して置き換えることができる機能です。

PyTorch Elastic には Rendezvous と呼ばれるコンポーネントが含まれています。このコンポーネントはトレーニングジョブの参加者 (ワーカー) を集め、そこで、全員が参加者と全員のロールとの同じリストに同意し、トレーニングをいつ開始/再開できるかについて一貫した集団での決定を行います。詳細については、GitHub のオープンソースプロジェクトにアクセスしてください。

TorchElastic Controller for Kubernetes (TECK) は PyTorch Elastic インターフェースのネイティブ Kubernetes 実装で、TorchElastic のトレーニングに必要な Kubernetes ポッドとサービスのライフサイクルを自動で管理します。シンプルなカスタムリソース定義を使用して、Elastic と互換性のあるトレーニングイメージを、希望するワーカーの最小数と最大数とともに指定します。これで、リクエストしたコンピューティングリソースの一部を使用して、Kubernetes クラスターでミッションクリティカルな分散トレーニングジョブを開始できます。また、リソースがさらに利用可能になると、ジョブを停止および再起動しなくても動的にスケーリングできるようになります。加えて、ノードの障害やレクラメーションにより置き換えられたノードから、ジョブを復旧できます。TorchElastic Controller for Kubernetes を使用することで、アイドル状態のクラスターリソースと Amazon EC2 スポットインスタンスでのトレーニングを制限して、分散トレーニングの時間とコストを削減できます。

しくみ

TorchElastic Controller for Kubernetes は特定のカスタムリソースタイプの ElasticJob を監視します。この ElasticJob は TorchElastic ジョブの送信の一部として作成されます。このリソースで関連イベントが発生すると、突合わせループ (reconcile loop) が開始して、現在の状態が目的の状態と一致することを確認します。

ElasticJob ワーカーはすべて同等で、コントローラーはポッド間の通信のためのヘッドレスサービスを作成して、必要な引数を PyTorch Elastic トレーニングランチャーに渡します。希望するワーカーサイズを [最小、最大] の範囲で動的に変更でき、コントローラーはこの変更に応じてジョブを中断することなくワーカー数を拡張したり縮小します。一連のポッドのメンバーシップが変更されるたびに、PyTorch Elastic は rendezvous を実行し、トレーニングを続行します。

チュートリアル

TECK は Kubernetes クラスターのバージョン 1.12 以降で動作します。この例では、Amazon EKS を使用しています。Kubernetes のスケーラビリティと AWS で利用可能な強力な高速インスタンスタイプが組み合わさっており、機械学習ワークロードには最適な選択肢となっています。さらに、EKS に最適化した GPU AMI により、複雑な依存関係を管理する必要がなく、機械学習ベースのアプリケーションを簡単に実行できます。

開始するには、exctl を使って GPU インスタンスのクラスターを作成します。すでに eksctl がインストールされている場合は、こちらの手順に従ってください。

eksctl create cluster \
    --name=torchelastic \
    --region=us-west-2 \
    --ssh-access \
    --ssh-public-key=~/.ssh/id_rsa.pub \
    --node-type=p3.2xlarge \
    --nodes=2 \
    --nodes-min=1 \
    --nodes-max=3

次に、NVIDIA デバイスプラグインをクラスターにインストールして、Kubernetes スケジューラがワーカーノードの GPU デバイスを把握できるようにします。

kubectl apply -f https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/1.0.0-beta4/nvidia-device-plugin.yml

PyTorch Elastic リポジトリのクローンを作成し、TorchElastic Controller for Kubernetes をクラスターにインストールします。

git clone https://github.com/pytorch/elastic.git
cd elastic/kubernetes

kubectl apply -k config/default

以下のような出力が表示されます。

namespace/elastic-job created
customresourcedefinition.apiextensions.k8s.io/elasticjobs.elastic.pytorch.org created
role.rbac.authorization.k8s.io/leader-election-role created
clusterrole.rbac.authorization.k8s.io/elastic-job-k8s-controller-role created
rolebinding.rbac.authorization.k8s.io/leader-election-rolebinding created
clusterrolebinding.rbac.authorization.k8s.io/elastic-job-k8s-controller-rolebinding created
deployment.apps/elastic-job-k8s-controller created

次のコマンドを実行して、コントローラーの準備ができていることを確認します。

kubectl get pods -n elastic-job

NAME                                          READY   STATUS    RESTARTS   AGE
elastic-job-k8s-controller-845ff64bc7-6lcqj   1/1     Running   0          3m45s

etcd に基づいた Rendezvous コンポーネントの実装はすでに PyTorch Elastic で提供されており、ほとんどのユーザーに推奨されています。ここでは、クラスターに単一メンバーの etcd インスタンスをインストールします。本番環境で使用する場合は、可用性の高い etcd クラスターを使用することをお勧めします。

kubectl apply -f config/samples/etcd.yaml

この例では、ImageNet データセットでトレーニングしています。ElasticJob 仕様の config/samples/imagenet.yaml の rdzvEndpoint 値を etcd サービス DNS 名「etcd-service:2379」に更新します。 このジョブ定義の他のパラメータに注意してください。

apiVersion: elastic.pytorch.org/v1alpha1
kind: ElasticJob
metadata:
  name: imagenet
  namespace: elastic-job
spec:
  rdzvEndpoint: "etcd-service:2379"
  minReplicas: 1
  maxReplicas: 3
  replicaSpecs:
    Worker:
      replicas: 2
      restartPolicy: ExitCode
      template:
        apiVersion: v1
        kind: Pod
        spec:
          containers:
            - name: elasticjob-worker
              image: torchelastic/examples:0.2.0rc1
              imagePullPolicy: Always
              args:
                - "--nproc_per_node=1" # set nproc_per_node = # Num GPUs on instance type
                - "/workspace/examples/imagenet/main.py"
                - "--arch=resnet18"
                - "--epochs=20"
                - "--batch-size=32"
                # data loader workers, not trainers.
                # zero means load the data on the same process as the trainer
                - "--workers=0"
                - "/workspace/data/tiny-imagenet-200"
              resources:
                limits:
                  nvidia.com/gpu: 1

現在、クラスターには 2 つの Amazon P3 インスタンスがあり、それぞれに 1 つの GPU があります。ではまず、トレーニングジョブで希望する数のレプリカを設定して、利用できる限り最大のリソースを使用しましょう。ワーカーの最小数は 1 に、最大数は 3 に設定されています。次に、トレーニングジョブをクラスターに送信します。

kubectl apply -f config/samples/imagenet.yaml

ワーカーが実行していることを確認してください。この時点で、PyTorch トレーニングジョブは大変順調に進んでいます。

kubectl get pods -n elastic-job

NAME                                          READY   STATUS    RESTARTS   AGE
elastic-job-k8s-controller-845ff64bc7-6lcqj   1/1     Running   0          3h32m
imagenet-worker-0                             1/1     Running   0          2m43s
imagenet-worker-1                             1/1     Running   0          2m43s

AWS ではオンデマンドインスタンスでワーカーノードを実行していますが、TECK は EC2 スポットインスタンスでのトレーニングもサポートしており、次のコマンドを使ってノードグループをスケールダウンすることで、スポットインスタンスの中断をシミュレーションできます。

eksctl get nodegroup --cluster torchelastic

# record the name of your node group and use in the command below
eksctl scale nodegroup --cluster=torchelastic --nodes=1 --name=ng-a345f4e1

ワーカーノードがクラスターから削除され、ポッドを再度確認すると、残りのポッドが 1 つだけ表示されます。

kubectl get pods -n elastic-job

NAME                                          READY   STATUS        RESTARTS   AGE
elastic-job-k8s-controller-845ff64bc7-vldbf   1/1     Running       0          13s
imagenet-worker-0                             1/1     Running       0          2m16s

典型的な分散トレーニングのシナリオでは、これはジョブを再起動する必要があることを意味しますが、トレーニングジョブについて説明したように、ジョブがノードの削除後も存続し、まだ実行していることがわかります。

kubectl describe elasticjob imagenet -n elastic-job

Message:               ElasticJob imagenet is running.
    Reason:                ElasticJobRunning
    Status:                True
    Type:                  Running
  Replica Statuses:
    Worker:
      Active:  1

では、さらに一歩進んで、クラスターのキャパシティを増やしましょう。実際のシナリオでは、これが別のトレーニングジョブが完了したときに発生する可能性があるので、GPU インスタンスを開放します。

eksctl scale nodegroup --cluster=torchelastic --nodes=3 --name=ng-a345f4e1

現在、クラスターには 3 つのワーカーノードがありますが、トレーニングジョブではそのうちの 1 つだけを使用しています。伸縮性のない分散トレーニングジョブでは、これらの余分な GPU はアイドル状態のままになります。しかし、PyTorch Elastic と TECK を使用することでジョブの設定が簡単になり、動的にスケーリングして新しく利用可能なキャパシティを利用できるようにできます。

kubectl edit elasticjob imagenet -n elastic-job

# set .spec.replicaSpecs[Worker].replicas to 3, and save.

最後にもう一度、ジョブを説明しましょう。

kubectl describe elasticjob imagenet -n elastic-job

Message:               ElasticJob imagenet is running.
    Reason:                ElasticJobRunning
    Status:                True
    Type:                  Running
  Replica Statuses:
    Worker:
      Active:  3

確かに、トレーニングジョブは 3 人のワーカーに動的にスケーリングされています。トレーニングジョブの完了までの時間が短縮されて、データサイエンティストは満足しますし、クラスター内の GPU を常に有効に活用していることで経理部も喜びます。

まとめ

このブログ投稿では、AWS と Facebook が開発した新しいオープンソースプロジェクト、TorchElastic Controller for Kubernetes をご紹介しました。これを使ってアイドル状態のクラスターリソースを制限し、障害からジョブを回復できるようにすることで、分散トレーニングの時間とコストを削減できるようになります。コントローラー内部の詳細について解説し、Amazon EKS クラスターでシンプルかつ耐障害性のある分散トレーニングジョブを実行する方法の一例をご紹介しました。TorchElastic Controller for Kubernetes を使用して柔軟で耐障害性のあるトレーニングを行うことで、ML モデルをより迅速に本番環境へ導入できます。また、アーキテクチャのサイズや複雑さが増大し続ける中でも、最新のアプローチを使ってモデルを探索できます。

今後の改善に向けて、ワーカーの自動サイズ変更、ジョブの優先度や先取りの概念の追加といった提案もあります。Kubernetes コミュニティにこのエキサイティングな新しい機械学習が加わったことについて、皆さまのご感想や使用方法をお知らせください。また、pytorch/elastic/kubernetes のオープンソースプロジェクトにもぜひご参加ください。