Laden von Daten auf AI-Runtime

Von Bedeutung

Dieses Feature befindet sich in der Public Preview.

Daten- und Modellressourcen sind für Deep Learning und Post-Training-Workloads großer Sprachmodelle (LLMs) und Vision-Language-Modelle (VLMs) von entscheidender Bedeutung. Mit AI Runtime werden alle Daten und Modellressourcen über den Unity Catalog abgerufen:

Unity-Katalog-Bände: werden hauptsächlich für große Datensätze und unstrukturierte Dateien verwendet, einschließlich Bildern, Audio und Text.
Unity-Katalogtabellen: Verwendet für strukturierte und tabellarische Daten, zugänglich über Spark Connect.

Ihre Volumes und Tabellen müssen im Unity-Katalog registriert und für Ihren Benutzer oder Service-Hauptverantwortlichen zugänglich sein.

Unity-Katalog-Volume für unstrukturierte Daten

Unity-Katalog-Volumes bieten kontrollierten Zugriff auf nicht-tabellarische Daten in jedem Format, einschließlich strukturierter, halbstrukturierter und unstrukturierter Daten. In AI Runtime sind Volumes der primäre Mechanismus zum Zugriff auf große Datensätze, Texte, Modellassets und Modell-Checkpoints.

Benutzer können Dateien in Unity-Catalog-Volumes mit vertrauten Dateisystemoperationen auflisten, lesen und schreiben, ähnlich wie mit Dateien auf einer lokalen Festplatte:

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!")

Ähnlich funktionieren Shell-Operationen auf die gleiche Weise:

%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

Einige Eigenschaften von Unity-Katalog-Volumes machen sie gut geeignet für maschinelles Lernen:

Verteilter Speicher: Der Unity Catalog wird von verteiltem Speicher unterstützt, der es KI-Laufzeit-Workloads ermöglicht, Daten und Modellassets über die gesamte Plattform zu lesen und zu schreiben, sowohl aus Notebooks als auch aus CLI-basierten Workloads.
Optimiert für ML-Zugriffsmuster: Die zugrunde liegenden Speicher- und Zugriffspfade sind für gängige ML-Workloads optimiert, insbesondere große Dateien mit sequentiellen Lese- und Schreibvorgängen. Dadurch eignet sich Unity Catalog gut zum Laden von Trainingsdaten, zum Laden von Modell-Assets und zum Schreiben von Modell-Checkpoints.
Dateisystem-ähnlicher Zugriff: Benutzer können Dateien in Unity-Catalog-Volumes mit vertrauten Dateisystemoperationen listen, lesen und schreiben, ähnlich wie mit Dateien auf einer lokalen Festplatte.

Aufgrund automatischer Hintergrundcommits können Benutzer einen konsistenten Zugriff auf die Unity Catalog-Volumendaten erwarten:

Schreibvorgänge: AI Runtime führt Schreibvorgänge automatisch aus, sodass Änderungen für andere Anwendungen und Workloads, die auf dieses Unity Catalog-Volume zugreifen, sichtbar werden.
Liest: KI-Laufzeit erkennt automatisch Änderungen am Volumen, ohne dass der Benutzer eine explizite Aktualisierung oder Synchronisation benötigt.

Lautstärke optimieren

Wie erwähnt, werden Unity-Katalog-Volumes durch verteilten Speicher unterstützt und für große Dateien mit sequentiellen Lese- und Schreibvorgängen optimiert.

Ein paar Tipps können Ihnen helfen, die beste Leistung aus AI Runtime zu erzielen:

Daten zu größeren Dateien verketten: Wenn möglich, konsolidieren Sie die Daten zu wenigeren, größeren Dateien, etwa 1 GiB bis 10 GiB pro Datei. Dies ermöglicht es AI Runtime, Daten aggressiv im Voraus zu laden und automatisch eine nahezu optimale sequentielle Leseleistung zu erreichen.

Für Workloads mit kleinen Dateien verwenden Sie lokale Festplatte: Wenn Ihre Arbeitslast viele kleine Dateien umfasst, sollten Sie die Dateien vor der Verarbeitung parallel auf die lokale Festplatte kopieren. Dies kann den Mehraufwand beim wiederholten Zugriff auf viele kleine Dateien über das Volume verringern.

# 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

Sie können UCVolumeDataset für Ihre Machine-Learning-Workloads verwenden. Sie integriert die oben beschriebenen Optimierungen, um einen effizienten Datenzugriff und das Laden von Unity-Katalog-Volumes zu ermöglichen. Weitere Informationen finden Sie in den folgenden Abschnitten.

Laden Sie unstrukturierte Daten mit `UCVolumeDataset`

Für unstrukturierte Daten wie Bilder, Audio- und Textdateien, die in Unity Catalog-Volumes gespeichert sind, verwenden Sie das UCVolumeDataset aus dem databricks.air.data-Modul. UCVolumeDataset ist ein PyTorch IterableDataset, der beim ersten Zugriff jede Datei aus dem Volume in einen schnellen lokalen Cache kopiert und den zwischengespeicherten lokalen Dateipfad liefert. Es behandelt die Leistungs- und Verteilungsbedenken, die Sie sonst von Hand implementieren würden:

Lokales Zwischenspeichern. Dateien werden beim ersten Zugriff vom FUSE-Mount in ein lokales Cacheverzeichnis kopiert und anschließend aus dem Cache bereitgestellt, sodass beim Training über mehrere Epochen nicht erneut aus dem Volume gelesen wird.
Automatische Partitionierung. Bei der Initialisierung von torch.distributed werden Dateien auf die Ränge aufgeteilt und anschließend weiter auf die DataLoader-Worker verteilt, sodass jedes (rank, worker)-Paar einen nicht überlappenden Abschnitt ohne zusätzliche Konfiguration erhält.

Hinweis

UCVolumeDataset und databricks.air.data.DataLoader kommen aus dem Paket databricks-sdk-air Installiere es mit dem Extra data, das auch eine kompatible torch mitinstalliert:

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

UCVolumeDataset liefert rohe lokale Dateipfade. Um diese Dateien in Tensoren zu decodieren, umschließen Sie sie mit einem zweiten IterableDataset, das den Pfadstream nutzt und Ihre Parsinglogik anwendet. Dadurch bleiben I/O- und Parsing-Aspekte getrennt.

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)

Der Wrapper empfängt bereits zwischengespeicherte lokale Pfade, sodass der Parsing-Schritt das Volumen nie berührt. Sie können zusätzliche Wrapper zur Erweiterung, Tokenisierung oder Filterung verketten.

Um eine optimale Leistung zu erzielen, kombinieren Sie UCVolumeDataset mit databricks.air.data.DataLoader anstelle des standardmäßigen PyTorch-DataLoader. Er ist auf KI-Laufzeit-I/O optimiert und ruft Dateien gleichzeitig ab und cachet, während die GPU berechnet.

Checkpoint-Modelle auf den Volumina

Um einen Checkpoint deines Modells zu erstellen, damit du das Training mit dem neuesten Snapshot fortsetzen oder den Trainingslauf nach einem Absturz wiederherstellen kannst, kannst du Unity Catalog-Volumes wie ein lokales Dateisystem verwenden.

Databricks empfiehlt die Verwendung eines verteilten Checkpoints (DCP) für bessere Leistung sowohl bei Single-GPU- als auch Multi-GPU-Workloads. Siehe schnelles, fehlertolerantes PyTorch-Training zur KI-Laufzeit im Databricks-Engineering-Blog.

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"],
)

Der monolithische Torch.save-Ansatz funktioniert ebenfalls.

Für das Speichern von Prüfpunkten eines Einzel-GPU-Modells,

# 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"])

Für verteiltes Training, das über torchrun gestartet wurde,

# 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"])

Laden von Tabellendaten

Verwenden Sie Spark Connect, um tabellarische Machine Learning-Daten aus Delta-Tabellen zu laden.

Für ein einzelnes Knoten-Training können Sie Apache Spark DataFrames mithilfe der PySpark-MethodetoPandas() in Pandas DataFrames konvertieren und dann optional mithilfe der PySpark-Methodeto_numpy() in das NumPy-Format konvertieren.

Hinweis

Spark Connect verschiebt die Analyse und Namensauflösung auf die Ausführungszeit, wodurch sich das Verhalten Ihres Codes ändern kann. Siehe "Vergleichen von Spark Connect" mit Spark Classic.

Spark Connect unterstützt die meisten PySpark-APIs, einschließlich Spark SQL, Pandas-API für Spark, Structured Streaming und MLlib (DataFrame-basiert). In der PySpark-API-Referenzdokumentation finden Sie die neuesten unterstützten APIs.

Weitere Einschränkungen finden Sie unter Serverless Compute Limitations.

Laden Sie große Delta-Tabellen mit Unity-Catalog-Volumes

Exportieren Sie die Daten von großen Delta-Tabellen, die zu groß sind, um mit toPandas() konvertiert zu werden, in ein Unity-Katalog-Volume und laden Sie sie direkt mithilfe von PyTorch oder 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")

Dieser Ansatz vermeidet Spark-Overhead während der Schulung und eignet sich gut für Einzel-GPU- und verteilte Schulungsworkflows.This approach avoid Spark overhead during training and works well for both single-GPU and distributed training workflows.

Feedback

War diese Seite hilfreich?

Zuletzt aktualisiert um 2026-10-05