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

データおよびモデル資産は、大規模言語モデル (LLM) およびビジョン言語モデル (VLM) のディープラーニングおよびポストトレーニングワークロードにとって極めて重要です。AI ランタイムでは、すべてのデータとモデル資産に Unity Catalog を介してアクセスします。

Unity Catalog ボリューム：主に、画像、オーディオ、テキストなどの大規模なデータセットと非構造化ファイルに使用されます。
Unity Catalog テーブル: 構造化および表形式データに使用され、 Spark Connect を通じてアクセスされます。

ボリュームとテーブルは Unity Catalogに登録されており、ユーザーまたはService Principalからアクセス可能である必要があります。

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

Unity Catalogボリュームは、構造化データ、半構造化データ、非構造化データなど、任意の形式の非表形式データに対するガバナンスされたアクセスを提供します。AI ランタイムでは、ボリュームは大規模なデータセット、テキスト、モデルアセット、モデルチェックポイントにアクセスするための主要なメカニズムです。

ユーザーは、ローカルディスク上のファイルの操作と同様の使い慣れたファイルシステム操作を使用して、Unity Catalogボリューム内のファイルを一覧表示、読み取り、および書き込みできます。

Python
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
%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 ボリュームには、Machine Learning ワークロードに非常に適したいくつかの特性があります。

分散ストレージ : Unity Catalog は分散ストレージによって支えられており、AI ランタイムワークロードはノートブックと CLIベースのワークロードの両方からプラットフォーム全体でデータとモデル資産を読み書きできます。
機械学習アクセスパターン向けに最適化 : 基盤となるストレージとアクセスパスは、一般的な機械学習ワークロード、特にシーケンシャル読み取り/書き込みを伴う大きなファイルに対して最適化されています。これにより、 Unity Catalog はトレーニングデータの読み込み、モデル資産の読み込み、およびモデルのチェックポイントの書き込みに最適化されます。
ファイルシステムのようなアクセス: ユーザーは、ローカルディスク上のファイルを操作するのと同様に、使い慣れたファイルシステム操作を使用して Unity Catalog ボリューム内のファイルをリスト、読み取り、書き込みできます。

自動バックグラウンド commit のため、ユーザーは Unity Catalog ボリュームデータに一貫してアクセスできます。

書き込み : AIランタイムは自動的に書き込みをcommitし、同じUnity Catalogボリュームにアクセスする他のアプリケーションやワークロードから変更を参照可能にします。
読み取り: AI ランタイムは、ユーザーによる明示的な更新または同期操作を必要とせずに、ボリュームへの変更を自動的に取得します。

ボリュームパフォーマンスの調整

前述のように、 Unity Catalog ボリュームは分散ストレージによってバックアップされ、逐次読み取りおよび書き込みを伴う大きなファイルに対して最適化されています。

AI ランタイムのパフォーマンスを最大限に引き出すためのヒントをいくつか紹介します。

データを大きなファイルに連結：可能な限り、データをより少ない大きなファイル（ファイルあたり約1 GiBから10 GiB）に統合してください。これにより、AI ランタイムは積極的にデータをプリフェッチし、ほぼ最適なシーケンシャル読み取りパフォーマンスを自動的に達成できます。

小規模なファイルワークロードの場合は、ローカルディスクを使用します。ワークロードに多くの小さなファイルが含まれている場合は、処理する前に並列コピーを使用してファイルをローカルディスクにコピーすることを検討してください。これにより、ボリュームを介して多くの小さなファイルに繰り返しアクセスするオーバーヘッドを削減できます。

sh
# 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

Machine LearningワークロードにUCVolumeDatasetを使用できます。これには上記の最適化が組み込まれており、Unity Catalog ボリュームからの効率的なデータアクセスとロードを提供します。詳細については、以下のセクションを参照してください。

非構造化データの読み込み `UCVolumeDataset`

Unity Catalogボリュームに保存された画像、音声、テキストファイルなどの非構造化データには、databricks.air.dataモジュールのUCVolumeDatasetを使用します。UCVolumeDatasetは、初回アクセス時にボリュームから各ファイルを高速なローカルキャッシュにコピーし、キャッシュされたローカルファイルパスを返すPyTorch IterableDatasetです。本来なら手動で実装する必要があるパフォーマンスや分散処理に関する対応を処理します：

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

注記

UCVolumeDataset および databricks.air.data.DataLoader は databricks-sdk-air パッケージから取得されます。data extraを使用してインストールします。これにより、互換性のあるtorchもプルされます。

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

UCVolumeDataset 未加工のローカルファイルパスを生成します。これらのファイルをテンソルにデコードするには、パスストリームを消費して解析ロジックを適用する2番目の IterableDataset でラップします。これにより、I/O と解析の懸念が分離された状態に保たれます。

Python
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 DataLoader ではなく、 UCVolumeDataset と databricks.air.data.DataLoader をペアリングしてください。AI ランタイムの I/O に調整されており、GPU のコンピュート中にファイルを同時に取得およびキャッシュします。

ボリューム上のチェックポイントモデル

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

Databricks では、シングルGPUとマルチGPUの両方のワークロードでパフォーマンスを向上させるために、分散チェックポイント (DCP) を使用することを推奨しています。Databricks エンジニアリングブログの「 AI ランタイムでの高速で耐障害性のあるPyTorchトレーニング」を参照してください。

Python
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モデルのチェックポイント作成について：

Python
# 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経由で起動される分散トレーニングの場合、

Python
# 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テーブルから表形式のMachine Learningデータを読み込みます。

単一ノードトレーニングでは、Apache Spark DataFramesをPySparkメソッドtoPandas()を使用してPandas DataFramesに変換し、オプションで PySparkメソッドto_numpy() を使用してNumPy形式に変換できます。

注記

Spark Connect では、解析と名前の解決が実行時まで遅延されるため、コードの動作が変更される可能性があります。「Spark Connect と Spark Classic の比較」を参照してください。

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

その他の制限については、「 Serverless コンピュートの制限事項」を参照してください。

Unity Catalogボリュームを使用して大規模なDeltaテーブルを読み込む

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

Python
# 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)

Python
# 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と分散トレーニングの両方のワークフローでうまく機能します。

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

ボリューム パフォーマンスの調整​

非構造化データの読み込み UCVolumeDataset​

ボリューム上のチェックポイントモデル​

表形式データを読み込む​

Unity Catalogボリュームを使用して大規模なDeltaテーブルを読み込む​