AI ランタイムにデータを読み込む

Important

この機能は パブリック プレビュー段階です。

データおよびモデル資産は、大規模言語モデル(LLM)やビジョン言語モデル(VLM)のディープラーニングおよびポストトレーニングワークロードにおいて不可欠です。 AI Runtimeでは、すべてのデータおよびモデルアセットにUnity Catalog経由でアクセスできます:

  • Unity Catalogボリューム:主に大規模なデータセットや非構造化ファイル(画像、音声、テキストを含む)に使用されます。
  • Unity Catalogのテーブル:構造化および表形式データに使用され、Spark Connect経由でアクセス可能です。

ボリュームとテーブルはUnity Catalogに登録され、ユーザーまたはサービス主体がアクセスできるようになっています。

非構造化データ向けの Unity Catalog ボリューム

Unity Catalogのボリュームは、構造化データ、半構造化データ、非構造化データを含むあらゆる形式の非表形式データに対するガバナンス制御されたアクセスを提供します。 AI Runtimeにおいて、ボリュームは大規模なデータセット、テキスト、モデルアセット、モデルチェックポイントにアクセスするための主要な手段です。

ユーザーは、ローカルディスク上のファイルを扱うのと同様に、馴染みのあるファイルシステム操作を使ってUnity Catalogボリューム内のファイルをリストアップ、読み書き、書き込みできます。

import os

dir_path = "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir"
file_path = os.path.join(dir_path, "test_file")

os.makedirs(dir_path, exist_ok=True)

# Write to the file
with open(file_path, "w") as file:
    file.write("Hello, World!")

同様に、シェル演算も同様に動作します:

%sh ls -l /Volumes/<catalog-name>/<schema-name>/<volume-name>
%sh mkdir -p /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir
%sh touch /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/test_file

Unity Catalogボリュームのいくつかの特徴は、機械学習ワークロードに適しています。

  • 分散ストレージ:Unity Catalogは分散ストレージに支えられており、AIランタイムワークロードがノートブックやCLIベースのワークロードからプラットフォーム上でデータの読み書きや資産のモデル化が可能です。
  • MLアクセスパターンに最適化:基盤となるストレージおよびアクセスパスは、特に順次読み書きを行う大規模なMLワークロード向けに最適化されています。 これにより、Unity Catalogはトレーニングデータロード、モデルアセットロード、モデルチェックポイント作成に適しています。
  • ファイルシステムのようなアクセス:ユーザーは、ローカルディスク上のファイルを扱うのと似た、慣れたファイルシステム操作を用いてUnityカタログボリューム内のファイルをリスト、読み書きできます。

自動バックグラウンドコミットにより、ユーザーはUnity Catalogのボリュームデータへの一貫したアクセスを期待できます:

  • 書き込み:AI Runtimeは書き込みを自動的にコミットし、同じUnity Catalogボリュームにアクセスする他のアプリケーションやワークロードにも変更を表示します。
  • 読み上げ:AIランタイムは、ユーザーが明示的にリフレッシュや同期操作を行うことなく、ボリュームの変更を自動的に拾います。

音量性能を調整

前述の通り、Unity Catalogのボリュームは分散ストレージでバックアップされており、連続的な読み書きを行う大容量ファイルに最適化されています。

AIランタイムで最高のパフォーマンスを得るためのいくつかのヒントを紹介します:

  • データを大きなファイルに連結する:可能であれば、1ファイルあたりおおよそ1 GiBから10 GiBの小さな大きなファイルにデータを統合しましょう。 これにより、AIランタイムは積極的にデータをプリフェッチし、ほぼ最適な逐次読み取り性能を自動的に達成できます。

  • 小ファイルワークロードにはローカルディスクを使いましょう:多数の小さなファイルを含むワークロードは、処理前に並列コピーを使ってローカルディスクにコピーすることを検討してください。 これにより、ボリュームを通じて多くの小さなファイルに繰り返しアクセスする負担を軽減できます。

    # Recommended using parallel copy (256 concurrency in this example, you can tune)
#
# This takes only 22 seconds to copy 15,375 150KiB small image files.
%sh cd /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ && find . -type f -print0 | xargs -0 -P 256 -I {} cp --parents "{}" /tmp/


# !!! Avoid doing this !!!
#
# Because the files are copied in serial, this copies the same 15,375 150KiB small image files much more slowly.
# %sh cp -r /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/* /tmp

  • 機械学習のワークロードには UCVolumeDataset を使うことができます。 上記の最適化を取り入れ、Unity Catalogボリュームからの効率的なデータアクセスと読み込みを実現します。 詳細については、次のセクションを参照してください。

UCVolumeDataset で非構造化データを読み込む

Unityカタログのボリュームに保存されている画像、音声、テキストファイルなどの非構造化データについては、UCVolumeDataset モジュールの databricks.air.data を使用します。 UCVolumeDataset は、最初のアクセス時に各ファイルをボリュームから高速ローカル キャッシュにコピーし、キャッシュされたローカル ファイル パスを生成する PyTorch IterableDataset です。 これは、手動で実装するパフォーマンスとディストリビューションの問題を処理します。

  • ローカル キャッシュ。 ファイルは最初のアクセス時に FUSE マウントからローカル キャッシュ ディレクトリにコピーされ、その後キャッシュから提供されるため、マルチエポック トレーニングではボリュームが再読み取りされません。
  • パーティションの自動分割。 torch.distributedが初期化されると、ファイルはランク間でパーティション分割され、さらにDataLoaderワーカー間で分割されるため、各(rank, worker)ペアは、追加のセットアップなしで重複しないスライスを受け取ります。

注

UCVolumeDataset そして databricks.air.data.DataLoader は databricks-sdk-air パッケージから来ています。 data extra をインストールすると、互換性のあるtorchも一緒にインストールされます:

%pip install "databricks-sdk-air[data]"

UCVolumeDataset では、生のローカル ファイル パスが生成されます。 これらのファイルをテンソルに変換するには、パス ストリームを受け取り、指定した解析ロジックを適用する 2 つ目の IterableDataset でラップします。 これにより、I/O と解析の問題が分離されます。

from databricks.air.data import UCVolumeDataset
from torch.utils.data import IterableDataset
from PIL import Image
import torchvision.transforms.functional as TF

class ImageDataset(IterableDataset):
    """Decodes each cached file path from UCVolumeDataset into a tensor."""

    def __init__(self, path_dataset: UCVolumeDataset):
        self._path_dataset = path_dataset

    def __iter__(self):
        for local_path in self._path_dataset:
            image = Image.open(local_path).convert("RGB")
            yield TF.to_tensor(image)

path_dataset = UCVolumeDataset("/Volumes/catalog/schema/my_volume/images")
dataset = ImageDataset(path_dataset)

ラッパーはすでにキャッシュされたローカルパスを受け取るため、解析ステップはボリュームに触れません。 拡張、トークン化、またはフィルター処理用の追加ラッパーを連結できます。

最適なパフォーマンスを得られるように、ストック PyTorch UCVolumeDatasetではなく、databricks.air.data.DataLoaderとDataLoaderをペアにしてください。 AIランタイムI/Oに対応しており、GPUが計算している間にファイルを同時にフェッチとキャッシュします。

体積上のチェックポイントモデル

最新のスナップショットからトレーニングを再開したり、クラッシュ後に復旧したりできるようにモデルのチェックポイントを保存するには、Unity Catalog ボリュームをローカル ファイル システムと同じように使用できます。

Databricksは、シングルGPUおよびマルチGPUワークロードの両方でより良いパフォーマンスを得るために分散チェックポイント(DCP)の使用を推奨しています。 DatabricksのエンジニアリングブログにあるAI Runtime での高速かつ耐障害性のある PyTorch トレーニングをご覧ください。

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

import databricks.air.data

checkpoint_path = "/Volumes/my-catalog/my-schema/my-volume/checkpoints/step_1000"

# Save
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd, "step": 1000}
dcp.async_save(
    state_dict,
    storage_writer=databricks.air.data.UCVolumeWriter(checkpoint_path))

# Load
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd}
dcp.load(
    state_dict,
    storage_reader=databricks.air.data.UCVolumeReader(checkpoint_path))

set_state_dict(
    model,
    optimizer,
    model_state_dict=state_dict["model"],
    optim_state_dict=state_dict["optim"],
)

一体型の torch.save アプローチも機能します。

  • 単一GPUモデルのチェックポイントの場合、

    # The monolithic torch.save approach for single GPU chip

# Save
torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
           "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")


# Load
ckpt = torch.load(
    "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
    weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["opt"])

  • Torchrunで実施された分散型トレーニングの場合、

    # The monolithic torch.save approach for multi-GPU distributed training.
# This snippet assumes your launcher has already called
# dist.init_process_group(...).

import os
import torch.distributed as dist

# Save only on rank 0.
if dist.get_rank() == 0:
    torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
               "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")

# Wait for rank 0 to finish writing before any rank reads.
dist.barrier()

# Load on ALL ranks (map to current rank's local GPU).
local_rank = int(os.environ["LOCAL_RANK"])
ckpt = torch.load(
    "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
    map_location=f"cuda:{local_rank}",
    weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["opt"])

表形式データを読み込む

Spark Connect を使用して 、Delta テーブルから表形式の機械学習データを読み込みます。

単一ノード トレーニングでは、PySpark メソッドを使用して Apache Spark DataFrames を pandas DataFrames に変換しtoPandas()必要に応じて PySpark メソッドを使用して NumPy 形式に変換できますto_numpy()。

注

Spark Connect は、分析と名前解決を実行時間に遅延させ、コードの動作を変更する可能性があります。 Spark Connect と Spark クラシックの比較を参照してください。

Spark Connect では、Spark SQL、Spark 上の Pandas API、構造化ストリーミング、MLlib (DataFrame ベース) など、ほとんどの PySpark API がサポートされています。 サポートされている最新の API については、 PySpark API リファレンス ドキュメント を参照してください。

その他の制限事項については、「 サーバーレス コンピューティングの制限事項」を参照してください。

Unityカタログボリュームを使って大きなDeltaテーブルを読み込む

toPandas()で変換するには大きすぎる大きな差分テーブルの場合は、データを Unity カタログ ボリュームにエクスポートし、PyTorch または Hugging Face を使用して直接読み込みます。

# Step 1: Export the Delta table to Parquet files in a UC volume
output_path = "/Volumes/catalog/schema/my_volume/training_data"
spark.table("catalog.schema.my_table").write.mode("overwrite").parquet(output_path)

# Step 2: Load the exported data directly using Hugging Face datasets
from datasets import load_dataset

dataset = load_dataset("parquet", data_files="/Volumes/catalog/schema/my_volume/training_data/*.parquet")

この方法では、トレーニング中の Spark のオーバーヘッドを回避し、単一 GPU と分散トレーニングの両方のワークフローに適しています。

  • 最終更新: