UCDataset を使って Unity Catalog テーブルでの GPU トレーニングを高速化する

通常、GPU はトレーニング実行において最もコストがかかる部分であり、ほとんどのチームは GPU をフル活用しきれていません。ボトルネックは通常モデルではなく、データを保存場所から必要な速度で GPU に取り込むことにあります。チップが待機している1秒ごとに、時間とコストの双方が失われています。

本日、AI Runtime で Unity Catalog Dataset (UCDataset) のベータ版をリリースします。UCDataset は、エクスポートやコピー、カスタムローダーコードなしで、Unity Catalog テーブルをトレーニングジョブ内のすべての GPU に直接ストリーミングするネイティブな PyTorch データローダーです。

シミュレートされた本番ワークロードにおいて、UCDataset は、同じデータのエクスポートされた Parquet コピーの読み込みと比較して、トレーニング時間を 10倍 短縮しました。

課題: GPU がデータの到着を待っている

モデル用のトレーニングデータのほとんどは、すでに Unity Catalog テーブルに存在しています。これを分散トレーニングジョブに取り込むには、3つの回避策のいずれかを選ぶ必要があり、それぞれ時間とコストがかかります。

• 実行前の遅く面倒なエクスポート: 最も一般的なパターンは、実行前にテーブルをボリューム内の Parquet (または他のシャード形式) にエクスポートし、データセットをそのファイルに向けるというものです。2つ目のコピーのストレージ費用が発生し、新しいデータが追加された瞬間に古くなり、Unity Catalog のリネージの対象外になります。
• テーブル全体をメモリに読み込む方法はスケールしない: toPandas() を使ってテーブルを pandas に読み込む方法は、小さなテーブルでは機能しますが、大きなテーブルでは失敗します。また、すべてのデータが1つのノードに集中するため、32個の GPU でトレーニングを行うジョブの役には立ちません。
• ウェアハウス経由のストリーミングは遅い: ウェアハウスは本来 GPU にデータを供給するように構築されておらず、最適化されていない場合、多数の GPU が同時に読み込むとスロットルされます。

さらに、ローダー自体の問題もあります。すべての GPU とワーカーには、データのバランスの取れた分割が必要です。チームはこれを手作業で構築するのに時間を費やしており、わずかなバグによって、誰も気づかないうちに重複またはスキップされた行でモデルがトレーニングされてしまう可能性があります。

結果: 高価な GPU が遊休状態になり、トレーニングに予想以上の時間がかかり、ML エンジニアはモデルの改善ではなくデータの配管の維持に追われることになります。

UCDataset の紹介

UCDataset は、Unity Catalog テーブルを分散 PyTorch トレーニングのファーストクラスのデータソースにします。読み込む対象 (テーブル、必要なカラム、行フィルター) を記述すると、UCDataset がトレーニング用のバッチを各 GPU にストリーミングします。利点は以下の通りです。

• スピード重視の設計: UCDataset はクラウドストレージからテーブルのファイルを直接読み込むため、データと GPU の間に余計なものは一切入りません。一致する行をすべての GPU とデータローダーワーカーに均等に分割します。
• コピーではなくレイクでトレーニング: ライブテーブルを指定します。データの読み込み時にフィルターと削除ベクトルが適用され、次回の実行では新しく追加された行が自動的に取得されます。
• エンドツーエンドのガバナンス: トレーニングジョブは実行者のアイデンティティとしてデータを読み込み、Unity Catalog のアクセス権限が Databricks の他の場所と同様に正確に適用されます。

ベンチマーク: レコメンデーションモデルのワークロードにおけるトレーニングの高速化

300 GB の特徴量テーブルを 8基の NVIDIA H100 GPU にストリーミングしてレコメンデーションモデルをトレーニングするシミュレートされた本番ワークロードで UCDataset をベンチマークしたところ、Parquet から直接データを読み込む場合と比較して、トレーニング時間が平均で 10倍向上しました。

使用方法:

クリックラベル付きのユーザーおよびアイテムの特徴量が main.recsys.ctr_features テーブルに保存されているチームを例にとります。テーブルには毎日新しいインプレッションが追加されます。このテーブルから直接、2つのノード上の 16基の GPU でクリック率モデルをトレーニングしたいとします。

1. 読み込む対象を記述し、ローダーを作成する。 train.py でテーブルを選び、ラベル付きのトレーニング行を保持し、モデルに必要なカラムのみを選択します。フィルターはテーブルのファイルの読み込み時に適用されるため、一致しない行がトレーニングコードに到達することは決してありません。UCDataset はすでに完全なバッチを生成するため、batch_size=None を指定して標準的な PyTorch DataLoader でラップします。

2. トレーニングと 2つノードでの起動。 標準的な PyTorch DDP ループでトレーニングします。UCDataset は 16基の GPU のそれぞれに一致する行の割り当てを自動的に行うため、自分で記述すべきサンプラーやシャーディングコードはありません。workload.yaml にジョブを記述すると、各ノードで torchrun を使用して train.py が開始されます。AI Runtime CLI で実行します:

より多くの GPU でトレーニングするには、workload.yaml の num_accelerators を変更します。トレーニングコードはそのままです。実際に試すためのサンプルテーブルを含む完全な train.py と workload.yaml は、AI Runtime クックブックにあります: Unity Catalog テーブルでのクリック率モデル。

はじめに

PyPI からパッケージをインストールします。ノートブックで %pip install databricks-uc-dataset を実行します。

UCDataset は AI Runtime のベータ版で利用可能です。試してみるには:

• ドキュメントを読む: UCDataset を使用して Unity Catalog テーブルを PyTorch に読み込む
• 動作する例から始める: クックブックのレシピ
• アカウントチームにお気軽にご連絡ください。フィードバックをお待ちしております。

Unity Catalog ボリュームに保存されている非構造化データの場合は、UCVolumeDataset を使用してローカルキャッシュでファイルをストリーミングします。UCDataset が構造化テーブルを処理する一方で、FILE タイプを介したマルチモーダル学習のサポートも拡大しており、テーブルから直接画像、音声、ドキュメントのトレーニングを行えるようになります。ベータ版では Unity Catalog の Delta テーブルを読み込みます。

よくある質問 (FAQ)

UCDataset とは何ですか？

UCDataset (Unity Catalog Dataset) は、Unity Catalog テーブルをトレーニングジョブ内の GPU に直接ストリーミングするネイティブなデータローダーです。Databricks AI Runtime のベータ版で利用できます。テーブルと必要なカラムを選択するだけで、エクスポート手順やカスタムローダーコードなしで、UCDataset がトレーニング用のバッチを提供します。

UCDataset はどのように GPU トレーニングを高速化しますか？

UCDataset は、エクスポートされたコピーを介在させることなく、クラウドストレージからテーブルのデータファイルを直接読み込むことで、GPU を稼働状態に保ちます。一致する行をすべての GPU とデータローダーワーカーに均等に分割するため、各 GPU は自分の分だけを読み込みます。

トレーニングのために Unity Catalog テーブルを Parquet にエクスポートする必要がありますか？

いいえ。UCDataset は Unity Catalog 内のライブテーブルを読み込むため、エクスポートのステップや、保存して同期を維持する必要のある2つ目のコピーはありません。読み込み時にフィルターとテーブルの削除ベクトルが適用され、新しい実行ごとに最後の実行以降に追加された行が取得されます。

UCDataset は Unity Catalog のアクセス権限を適用しますか？

はい。AI Runtime では、トレーニングジョブは実行者のアイデンティティとしてデータを読み込み、テーブルに対する Unity Catalog のアクセス権限が適用されます。ジョブがガバナンスされたテーブルを直接読み込むため、個別に保護すべきエクスポートされたコピーはありません。ベータ版では、行フィルターや列マスクを持つテーブルを UCDataset で読み込むことはできません。

UCDataset はどのテーブルを読み込むことができますか？

UCDataset は、削除ベクトルを持つテーブルを含め、Unity Catalog に登録された Delta テーブルを読み込みます。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事