Ray Data、GPU ネイティブ演算子を導入し処理効率を向上

Ray Data は、マルチモーダルデータの処理やトレーニングデータロードといった AI ワークロードに特化して構築されたデータエンジンです。AI ワークロードへの注力により、GPU が Ray Data の中核機能となっています。Ray Data では GPU をスケジューリングリソースとして扱えるため、データ処理タスクを GPU 間で容易にスケーリングできます。

しかしながら、Ray Data はこれまで、PyTorch や JAX などの GPU 対応ライブラリを呼び出すユーザー定義関数（UDF）に依存し、GPU 上で直接実行されるネイティブ機能は限定的でした。

最近、Anyscale の Ray Data チームと NVIDIA の cuDF チームが協力し、Ray Data への GPU ベース演算子の統合を強化しました。特に注目すべき新機能として、以下の 2 つを発表します。

• Ray Data の map_batches 用のバックエンドとしての cuDF
• RapidsMPF に基づく GPU ベースのシャッフル機能

特定のデータキュレーションワークロードにおいて、これらの強化により CPU ベースのソリューションと比較して TCO（総所有コスト）が最大 3 倍改善されました。

Ray を使用して今日から試すには、pip install -U ray==2.58 でインストールしてください。

LinkBackground

Ray Data は、IO からシャッフル、UDF までを含む一連の標準的なデータ処理操作を提供します。Ray Data の差別化機能の一つは、UDF に対して GPU をファーストクラスとしてサポートしている点です。例えば、GPU を活用したカスタムコードを以下のようにスケジューリングできます。

import ray

ds = ray.data.read_parquet("s3://my-bucket/documents/")

def compute_gpu(batch: pd.DataFrame):
    ...

ds = ds.map_batches(compute_gpu, num_gpus=1, batch_size=4096, batch_format="pandas")

上記のコードを実行すると、Ray クラスター内の異なる GPU に接続された複数のプロセス上で、それぞれが別々の GPU で動作する compute_gpu 呼び出しが自動的に起動されます。

Ray Data をデータ重複排除タスクに活用する際、私たちは改善の余地を2点見出しました。

• パイプライン内には計算集約的なデータフレーム操作が含まれており、cuDF を利用すれば高速化が可能ですが、ユーザーは手動でデータ形式を変換する必要がありました。
• 特定のタスクでは、計算負荷の高いハッシュ集計ステップが存在し、GPU アクセラレーションの強化によって恩恵を受けられる可能性がありました。

これらの課題に対応するため、Ray Data に2つの重要な機能を追加しました。1 つは cuDF バッチフォーマットへの対応、もう 1 つは RapidsMPF を基盤とした GPU ネイティブなシャッフルバックエンドです。

cuDF をファーストクラスのバッチフォーマットとしてリンクする

Ray Data には、Pandas データフレーム、NumPy アレイ、PyArrow テーブルの3つのネイティブなバッチフォーマットが用意されています。バッチフォーマットの役割は2つあります。1 つはステージ間のデータ交換形式としての機能、もう 1 つはユーザーがカスタム処理関数を書く際に親しみやすいインターフェースを提供することです。

Ray 2.58 以降、map_batches は batch_format="cudf" を受け付けるようになりました。これにより、ユーザー定義関数（UDF）に渡されるデータ型が pandas DataFrame や Arrow テーブルから、GPU 上で動作する cudf.DataFrame に変更されます。

cuDFは、NVIDIA が提供するオープンソースの GPU アクセラレーション対応 DataFrame ライブラリです。これはデータ処理用の NVIDIA CUDA-XTM ライブラリの一つであり、Pandas と似た API を提供しながら、GPU 上で高速な演算を可能にします。計算負荷の高い処理においては、従来の CPU ベースの同等機能と比較して、速度が桁違いに速く、コストも大幅に削減できます。

import ray

ds = ray.data.read_parquet("s3://my-bucket/documents/")

def add_minhash(batch):  # batch is a cudf.DataFrame
    batch["minhash"] = batch["text"].str.minhash(seed=42, ...)
    return batch

ds = ds.map_batches(add_minhash, batch_format="cudf", num_gpus=1, batch_size=4096)

GPU ネイティブなシャッフルを LinkGPU で実現

ユーザーは、RapidsMPF を基盤とした新しい GPU 実行パスを通じて、GPU アクセラレーションされた repartition やハッシュ集計を利用できるようになりました。ランク間では UCXX を介して直接通信が行われ、シャッフルされるデータは RDMA や NVLink といった高性能な転送経路を活用します。これにより、CPU や Ray オブジェクトストアを経由せずに処理を進めることが可能になります。

cuDF と同様に、計算負荷の高い特定のオペレーションにおいては、RapidsMPF ベースのシャッフルが従来の CPU 版に比べて最大で桁違いに高速かつ低コストで動作します。

from ray.data.context import DataContext, ShuffleStrategy

ctx = DataContext.get_current()
ctx.shuffle_strategy = ShuffleStrategy.GPU_SHUFFLE

ds.repartition(num_blocks=256, keys=["doc_id"])
ds.groupby("band_id").count()

LinkBenchmarks: 曖昧な重複排除

当社の共同研究で焦点を当てたのは、文書の曖昧な重複排除（fuzzy document deduplication）です。これは基盤モデルのトレーニングにおけるデータキュレーションのための標準的なワークロードとなっています。

ベンチマークには、以下の手順で説明される標準的な MinHash-LSH 手法 を採用しました。

図 1: デュプリケーション（重複除去）ワークロードにおける各工程

*Fig 1 Stages in the deduplication workload*

• すべてのドキュメントに対して MinHash シグネチャを生成する。
• そのシグネチャから LSH バンドを作成する。
• バンドごとにグループ化し、候補となる重複ペアを見つける。
• 候補エッジの重複除去を行う。
• グラフを走査して連結成分を計算する。
• 重複データをフィルタリングする。
• 残ったデータを永続化する。

1 工程目は各 n-gram のハッシュ値を生成する処理ですが、ハッシュ演算は計算集約型のオペレーションです。

3 工程目と 5 工程目はシャッフル（データ再配置）負荷が最も高い部分です。重複ペアを形成するためにグループ化操作を使用しており、連結成分アルゴリズムも一連のグループ化操作で構成されています。

以下の実験では、GPU の統合と加速によりこれらの工程のパフォーマンス向上を確認できました。

実験には FineWeb 10BT データセットを使用し、クラスタサイズは 2 ノードから 8 ノードの間で設定しました。CPU 環境では m8id.4xlarge インスタンスタイプを採用しました。これは 16 vCPU と 64 GiB メモリを備え、時間あたりのコストは 1.04416 ドルです。

GPU 評価には g6.4xlarge インスタンスタイプを使用しました。これも 16 vCPU と 64 GiB メモリを搭載しており、時間あたりのコストは 1.3232 ドルです。

実験コードは Ray 2.58 を使用して こちら で公開されています。

LinkGenerating MinHash 署名

このフェーズでは、各ドキュメントを n-gram のセットに分割します（n は 5 に設定）。その後、各 n-gram から 128 個のハッシュ値を生成します。

ハッシング処理は計算集約型のため、Numpy ベースの従来手法と、cuDF に組み込まれたハッシュ実装を比較しました。

上記の結果から、GPU によるアクセラレーションにより、この処理段階のワークロードが同規模の CPU クラスターと比較して最大で15 倍高速化されることがわかります。

LinkThe grouping stages（グループ化ステージ）

ステップ 3 と 5 では、局所感度ハッシュ（locality-sensitive hashing）を実行し、ハッシュ値が衝突する文書をグループ化します。また、反復的な連結成分検出アルゴリズム を実装して、推移的な衝突を持つ文書群を特定しています。

具体的には、「文書 A が文書 B と類似しており、かつ文書 C も文書 B と類似している場合、文書 A と文書 C も同様に類似しているとみなす」といった関係性を検出します。

GPU 数を増やすと実行時間は短縮されますが、大規模化するとその効果は頭打ちになり、コスト削減のメリットも減少します。

これらの処理段階はシャッフル（データ再配置）を多用するため、GPU を追加すると通信オーバーヘッドが増加します。今回の特定のデータセットとワークロードでは 2 GPU が最適であり、8 GPU の性能を活かしきるにはデータ量が不足していたと考えられます。

リンク：エンドツーエンドパイプライン

固定されたハードウェア構成におけるフルパイプラインの結果から、GPU 加速ソリューションの優位性が明確になりました。

• 最安コストの CPU 比較対象（2 ノード）と比較して、処理速度が 4 倍向上し、TCO（総所有コスト）は 3.1 倍改善されています。
• すべてのノードスケールにおいて、TCO と実行時間の両面で改善効果が得られています。

次のステップ

上記の結果は有望なものですが、まだ性能の上限にはほど遠い状況です。特に、cuDF と RapidsMPF を使用する連続した演算子間の融合（operator fusion）を実現することで、パフォーマンスに大きな飛躍をもたらすことが期待されます。これにより、Ray Data はステージ間での不要なホストからデバイスへのデータ転送を回避し、データを GPU メモリ内に保持したまま処理できるようになります。

エンドユーザーとの協働を通じて、次のステップとしてデータ前処理における GPU シャッフル機能の提供を検討しています。これにより、Ray Data のプリプロセッサ API をベースに構築された特徴量エンジニアリングパイプラインも、GPU 経路を活用できるようになります。

AI データワークロード向けに最良のデータ処理パフォーマンスを提供するため、この分野への投資を継続していく予定です。なお、Ray Data の GPU シャッフル機能は現在実験段階であり、Ray 2.58 から利用可能です。実際に試していただいた際は、ぜひフィードバックをお聞かせください！