Re Reference AI

インフラ

Rayとは

Pythonの並列・分散処理を簡単に書けるようにする、機械学習ワークロード向けのオープンソース分散実行フレームワーク

分散処理スケーリング

ひとことで言うと

Pythonのプログラムを、複数のマシンやGPUに分散させて動かしやすくするための土台。

概要

Rayは、通常のPython関数やクラスに簡単なデコレータを付けるだけで、それらを複数のCPU・GPU・マシンにまたがって並列・分散実行できるようにするオープンソースのフレームワークで、分散システムの複雑な実装を意識せずにスケールするコードを書けるようにすることを目的としている。 コア機能である汎用的な分散実行エンジンの上に、ハイパーパラメータチューニングのRay Tune、強化学習のRay RLlib、分散データ処理のRay Data、モデルサービングのRay Serveなど、機械学習の各工程に特化したライブラリ群(Ray AI Libraries)が構築されている。 LLMの分散学習や大規模なバッチ推論、モデルサービングのスケーリングなど、単一のGPU・マシンでは処理しきれない規模の機械学習ワークロードを扱う場面で広く利用されている。 vLLMなど一部のLLM推論エンジンも、複数GPUにまたがる分散推論の基盤としてRayを内部で利用している。

背景

機械学習ワークロードは、データ量やモデルサイズの増加に伴い、単一マシンの計算資源では処理しきれない規模になることが多い。 Rayは、Pythonの関数・クラスを最小限のコード変更で分散実行できるようにすることで、分散システムの専門知識がなくてもスケールする機械学習パイプラインを構築できるようにする目的で開発された。

歴史

2017年: UC BerkeleyのRISELabがRayを発表。 2019年: RayをベースとしたスタートアップAnyscaleが設立され、開発が本格化。 2020年代: LLMの分散学習・分散推論基盤として、vLLMHugging Faceのエコシステムとの統合が進む。

アーキテクチャ

Rayはヘッドノードとワーカーノードから成るクラスタ上で動作し、Python関数を「タスク」、Pythonクラスを状態を持つ「アクター」として分散実行する。 スケジューラがクラスタ全体のリソース(CPU・GPU・メモリ)の空き状況を把握し、タスクやアクターを空きリソースのあるノードへ割り当てる。 オブジェクトストアが、ノード間で共有されるデータをメモリ上に保持し、ワーカー間のデータ受け渡しのオーバーヘッドを抑える。

コード例

@ray.remoteによる並列実行例

import ray

ray.init()

@ray.remote
def process(x):
    return x * x

futures = [process.remote(i) for i in range(4)]
results = ray.get(futures)

利点

  • 既存のPythonコードに最小限の変更(デコレータの追加など)を加えるだけで分散実行に対応できる
  • 学習・チューニング・データ処理・サービングなど、機械学習パイプライン全体を一貫したAPIでスケールできる
  • クラスタの規模拡大に応じて、ノードの追加だけで処理能力を柔軟に拡張できる

欠点

  • 分散システム特有の障害(ノード障害やネットワーク遅延など)を考慮した設計・運用の知識が必要になる
  • 小規模な処理や単一マシンで完結するタスクでは、Rayを導入する分のオーバーヘッドが目立ちやすい
  • クラスタ構成やリソース割り当ての最適化には、ある程度の試行錯誤が必要になる

比較

  • vLLMvLLMは一部の分散推論機能の実装にRayを利用しており、両者は組み合わせて使われることがある
  • GPURayは複数のGPU・マシンにまたがるリソースを統一的に管理し、分散学習・推論のスケジューリングを担う
  • バッチ推論Ray Dataを使うと、大規模なバッチ推論を複数ノードに分散させて高速に処理できる

関連用語

vLLMGPUバッチ推論

よくある質問

Rayは何のためのツール?

Pythonのコードを複数のCPU・GPU・マシンにまたがって並列・分散実行するためのフレームワーク。機械学習の学習・チューニング・データ処理・サービングなど、様々な工程のスケーリングに使われる。

RayとKubernetesの違いは?

Kubernetesはコンテナのオーケストレーション基盤であるのに対し、RayはPythonコード自体を分散実行するためのプログラミングフレームワーク。RayのクラスタをKubernetes上で運用する組み合わせもよく使われる。

参考文献

関連Zenn記事