機械学習基盤としてのKubernetes

By Avinton rk | 技術解説, 技術ブログ＆インタビュー | Comments are Closed | 15 5月, 2024 | 1

最近、急速に進化するAI技術、特に大規模言語モデル（LLM）などを含む人工知能技術が注目を集めています。一般的な報道では、ニューラルネットワーク（NN）のアーキテクチャやパラメータ数に注目されることがよくありますが、これらのモデルがどのような環境で構築されているかについてはあまり言及されません。

今回は、機械学習基盤としてのKubernetesに焦点を当て、大規模学習と大規模モデルがどのような環境で構築されているかを紹介したいと思います。

LLMに限らず、大規模学習を伴うAIモデルを構築するには、それに適した機械学習基盤が必要です。その中でも特に重要なのがGPU（Graphics Processing Unit）です。GPUはもともとは画像処理を得意とするパーツでしたが、その並列処理性能の高さからAI開発にも利用されるようになりました。

数年前までは高性能なGPU数台でAI開発を十分に行うことができましたが、近年の急速な発展に伴ってGPUの必要数/性能がこれまで以上に求められるようになりました。数台のGPUであればマニュアルで管理することもできそうですが、数千台となったら非現実的です。
そこで着目されたのがKubernetesです。Kubernetesはその特性からノードのスケールアップが容易で、GPU等のリソースを断続的に追加するといった拡張を実現できます。これにより、柔軟性と効率性が向上し、AI開発プロジェクトにおけるリソースの最適化が可能となります。また、Kubernetesの自動化機能を活用することで、運用コストや管理の手間も削減されます。AI開発の規模や要求に応じてリソースを調整し、効率的な開発環境を構築することができるため、多様なニーズに柔軟に対応することができます。

用語紹介

Kubernetes

Kubernetesは、コンテナ化されたアプリケーションのデプロイメント、スケーリング、および管理を自動化するためのオープンソースのプラットフォームです。クラウド上やオンプレミスのインフラストラクチャ上で、コンテナ化されたアプリケーションの効率的な運用を可能にします。

主な機能:

1.コンテナオーケストレーション: コンテナを自動的に配置し、スケーリングし、管理します。これにより、アプリケーションの可用性とパフォーマンスが向上します。
2.サービスディスカバリーとロードバランシング: アプリケーションのためのサービスを自動的に構成し、クライアントからのリクエストを適切なバックエンドポッドにルーティングします。
3.自動デプロイメントとローリングアップデート: 新しいバージョンのアプリケーションをデプロイする際に、古いバージョンと新しいバージョンの間でシームレスなトラフィックの切り替えを行います。
4.ストレージオーケストレーション: 永続化されたデータのためのストレージを管理し、コンテナに永続化ボリュームを提供します。
5.自己修復: ノードやコンテナの障害を検出し、自動的に修復します。これにより、アプリケーションの可用性が向上します。

実際に、OpenAIは2018年1月の記事でKubernetesクラスターを2500ノードにスケールアップし、2021年1月の記事では7500ノードに拡張したと報告しています。このような背景から、OpenAIは機械学習基盤の整備を早期に行い、それが現在でも最前線での開発を支える基盤となっていると考えられます。
(補足: 以下表のPending Podsはリソース不足で予約待ち状態のポッド数を意味します)

参照：https://openai.com/index/scaling-kubernetes-to-7500-nodes/

一方で、Kubernetesクラスターのノードを拡張し複数のGPUを組み込んだだけではそれらを十分に活用することができません。複数のGPUを同時に使ってAIモデルトレーニングを行う場合、機械学習フレームワーク(PyTorch、Tensorflowなど)を適切に設定し分散処理を制御させる必要があります。
ここからは一つの例として、KubeflowとPyTorchを用いた分散学習の実装を簡単に紹介したいと思います。

Kubeflow

Kubeflowは、機械学習ワークロードを効率的にデプロイ、管理、スケーリングするためのオープンソースプラットフォームです。Kubernetes上で動作し、機械学習モデルのトレーニング、推論、モデル管理をシームレスに行うためのツールとリソースを提供します。

主な機能:

1.ワークフローのオーケストレーション(Kubeflow Pipeline): 機械学習ワークフローを定義し、Kubernetes上で自動的に実行します。これにより、トレーニング、評価、推論などの各ステップを統合的に管理できます。
2.ノートブックの提供: Jupyterノートブックと統合されており、データの探索、前処理、モデルのトレーニングなどの機械学習タスクを実行できます。
3.ハイパーパラメータの最適化: ハイパーパラメータの自動チューニングをサポートし、モデルのパフォーマンスを最適化します。
4.モデルデプロイメント: トレーニングされたモデルを本番環境にデプロイするためのツールを提供し、モデルの推論をスケーリングします。
5.モデル管理: トレーニングされたモデルとそのバージョンを追跡し、再トレーニングや展開のためのモデルのバージョン管理を支援します。

PyTorch

PyTorchは、ディープラーニングのためのオープンソースのライブラリであり、Pythonをベースとしています。柔軟性と高いパフォーマンスを特徴としており、研究から本番環境まで幅広い用途で使用されています。

主な特徴:

1.動的計算グラフ: PyTorchは、動的計算グラフを採用しています。これにより、モデルの構築や変更が柔軟に行え、デバッグやモデルの理解が容易になります。
2.GPUサポート: PyTorchは、GPUを使用して高速な演算を実行することができます。

核となる機能は、Kubeflow – Training Operator(PyTorchJob) と PyTorch – Distributed Data Parallel(DDP)モジュールです。

複数ノードにまたがったGPUを使用するには、各ノードにGPUが割り当てられたポッドを実行し、それらが独自のネットワークで認識している必要があります。以下は4つのノード（各ノードに1つのGPU）を使用する場合の構成図です。master-0, worker-0, worker-1, worker-2はポッドを表します。

PyTorchJobを使用すれば、1つのマニフェストファイルで全てのポッドを管理することができます。以下はシンプルにしたマニフェストの例です。
例えば、100台のGPUを使用する場合はWorker.replicas:を3から100に変更するだけで実現できます。

apiVersion: kubeflow.org/v1

kind: PyTorchJob

metadata:

spec:

pytorchReplicaSpecs:

Master:

replicas: 1

template:

spec:

containers:

command:

- <your-command>

image: <your-image>

resources:

limits:

nvidia.com/gpu: 1

Worker:

replicas: 3

template:

spec:

containers:

command:

- <your-command>

image: <your-image>

resources:

limits:

nvidia.com/gpu: 1

この時点ですべてのコンテナは同じ処理をする状態です。
また、PyTorchJobがpytorchコンテナに以下のような環境変数が自動で付与されます。これらの環境変数はDDPでの分散処理に必要になります。

# master-0

Environment:

MASTER_PORT: 23456

MASTER_ADDR: localhost

WORLD_SIZE: 4

RANK: 0

# worker-0

Environment:

MASTER_PORT: 23456

MASTER_ADDR: pytorchjob

WORLD_SIZE: 4

RANK: 1

次に、各コンテナ内で実行されるトレーニングプロセスにDDPを適用します。以下は公式ドキュメントのサンプルです。nnモジュールを継承したカスタムモデルクラスをDDPモジュールの引数に入れることで分散処理向けのモデルインスタンスを得ることができます。

# https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#initialize-ddp-with-torch-distributed-run-torchrun

import torch

import torch.distributed as dist

import torch.nn as nn

import torch.optim as optim

from torch.nn.parallel import DistributedDataParallel as DDP

class ToyModel(nn.Module):

def __init__(self):

super(ToyModel, self).__init__()

self.net1 = nn.Linear(10, 10)

self.relu = nn.ReLU()

self.net2 = nn.Linear(10, 5)

def forward(self, x):

return self.net2(self.relu(self.net1(x)))

def demo_basic():

dist.init_process_group("nccl")

rank = dist.get_rank()

print(f"Start running basic DDP example on rank {rank}.")

# create model and move it to GPU with id rank

device_id = rank % torch.cuda.device_count()

model = ToyModel().to(device_id)

ddp_model = DDP(model, device_ids=[device_id])

loss_fn = nn.MSELoss()

optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

optimizer.zero_grad()

outputs = ddp_model(torch.randn(20, 10))

labels = torch.randn(20, 5).to(device_id)

loss_fn(outputs, labels).backward()

optimizer.step()

dist.destroy_process_group()

if __name__ == "__main__":

demo_basic()

実行時は以下のようにコンテナに割り当てられた環境変数を使用して行います。

python -m torch.distributed.run \

--nproc_per_node 1 \

--nnodes $WORLD_SIZE \

--node_rank $RANK \

--master_addr $MASTER_ADDR \

--master_port $MASTER_PORT \

your_training_script.py

以上が、Kubernetesクラスター上でKubeflowとPyTorchを用いた分散学習方法の例です。

まとめ

大規模な学習や大規模なAIモデルの構築は、ハードウェアとソフトウェアの両方に関するスキルが必要であり、高度な領域です。しばしば、ニューラルネットワークのアーキテクチャやパラメータ数に焦点が当てられがちですが、AIの進化を促進するには、機械学習基盤との統合も重要です。

分散学習は、計算リソースの効率的な利用とモデルのトレーニング時間の短縮に役立ちます。Kubernetesクラスターを使用することで、柔軟性と拡張性を持った環境を構築し、複数のノードを活用して大規模なデータセットに対しても効率的に学習を行うことができます。

Kubeflowは、Kubernetes上で機械学習ワークフローを管理するためのオープンソースプロジェクトであり、PyTorchなどの機械学習フレームワークと統合されることができます。これにより、データの前処理からモデルのトレーニング、評価、デプロイまでの一連のプロセスを効率的に管理できます。

また、機械学習基盤の構築においては、データの管理やバージョン管理、モデルの監視やチューニング、セキュリティなどの要素も重要です。これらの側面を総合的に考慮することで、持続的なAIシステムの開発と運用が可能になります。

以上のように、AIの発展を促進するためには、単にモデルの設計やパラメータの最適化だけでなく、機械学習基盤の整備や運用も重要です。KubernetesやKubeflowなどのツールを活用し、ハードウェアとソフトウェアの両面からAIの発展を支援していくことが求められます。

Avintonでは、ITエンジニアリングトレーニングとチームメンバーの継続的な教育に特に重点を置いています。この分野でのスキルの活用、最先端のテクノロジーへの取り組み、国際的なクライアントへのソリューション提供に興味がある場合は、採用情報のページをご覧ください。ご連絡をお待ちしております！

Avintonで働く魅力

機械学習基盤としてのKubernetes