Gegevens laden in AI Runtime

Belangrijk

Deze functie bevindt zich in openbare preview-versie.

Data- en modelassets zijn cruciaal voor deep learning en post-training workloads voor grote taalmodellen (LLM's) en vision-language modellen (VLM's). Met AI Runtime worden alle data en modelassets toegankelijk via Unity Catalog:

  • Unity Catalogus-volumes: voornamelijk gebruikt voor grote datasets en ongestructureerde bestanden, waaronder afbeeldingen, audio en tekst.
  • Unity Catalog-tabellen: gebruikt voor gestructureerde en tabulaire data, toegankelijk via Spark Connect.

Je volumes en tabellen moeten geregistreerd zijn in de Unity Catalog en toegankelijk zijn voor je gebruiker of diensthoofd.

Unity Catalog volume voor ongestructureerde data

Unity Catalogus-volumes bieden gereguleerde toegang tot niet-tabulaire data in elk formaat, inclusief gestructureerde, semi-gestructureerde en ongestructureerde data. In AI Runtime zijn volumes het primaire mechanisme voor toegang tot grote datasets, tekst, modelassets en modelcontrolepunten.

Gebruikers kunnen bestanden vermelden, lezen en schrijven in Unity Catalog-volumes met behulp van bekende bestandssysteembewerkingen, vergelijkbaar met het werken met bestanden op een lokale schijf:

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

Evenzo werken shell-operaties op dezelfde wijze:

%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

Enkele kenmerken van Unity Catalog-volumes maken ze goed geschikt voor machine learning-workloads:

  • Gedistribueerde opslag: Unity Catalog wordt ondersteund door gedistribueerde opslag, waardoor AI-runtime-workloads data en modelassets over het platform kunnen lezen en schrijven, zowel vanaf notebooks als CLI-gebaseerde workloads.
  • Geoptimaliseerd voor ML-toegangspatronen: De onderliggende opslag- en toegangspaden zijn geoptimaliseerd voor veelvoorkomende ML-workloads, met name grote bestanden met sequentiële lees- en schrijfopdrachten. Dit maakt Unity Catalog zeer geschikt voor het trainen van dataloading, het laden van modelassets en het schrijven van modelcheckpoints.
  • Bestandssysteem-achtige toegang: Gebruikers kunnen bestanden vermelden, lezen en schrijven in Unity Catalog-volumes met behulp van bekende bestandssysteembewerkingen, vergelijkbaar met het werken met bestanden op een lokale schijf.

Door automatische achtergrondcommits kunnen gebruikers consistente toegang tot volumegegevens van de Unity Catalog verwachten:

  • Schrijfbewerkingen: AI Runtime legt schrijfbewerkingen automatisch vast, waardoor wijzigingen zichtbaar worden voor andere applicaties en workloads die toegang hebben tot hetzelfde Unity Catalog-volume.
  • Lezen: AI Runtime detecteert automatisch wijzigingen op het volume zonder dat de gebruiker expliciet een vernieuwings- of synchronisatiebewerking hoeft uit te voeren.

Optimaliseer de prestaties van volumes

Zoals vermeld, worden Unity Catalog-volumes ondersteund door gedistribueerde opslag en geoptimaliseerd voor grote bestanden met sequentiële lees- en schrijfopdrachten.

Een paar tips kunnen je helpen de beste prestaties uit AI Runtime te halen:

  • Verbind data in grotere bestanden: Konsolideer indien mogelijk gegevens tot minder, grotere bestanden, ongeveer 1 GiB tot 10 GiB per bestand. Dit stelt AI Runtime in staat om data agressief vooraf te prefeten en automatisch bijna optimale sequentiële leesprestaties te bereiken.

  • Voor kleine bestandsworkloads gebruik lokale schijf: Als je workload veel kleine bestanden bevat, overweeg dan om de bestanden naar de lokale schijf te kopiëren met parallelle kopieën voordat je ze verwerkt. Dit kan de overhead verminderen van het herhaaldelijk openen van veel kleine bestanden via het volume.

    # 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

  • Je kunt het gebruiken UCVolumeDataset voor je machine learning-workloads. Het bevat de hierboven beschreven optimalisaties om efficiënte gegevenstoegang en laden vanuit Unity Catalog-volumes te bieden. Zie de volgende secties voor meer informatie.

Laad ongestructureerde data met UCVolumeDataset

Gebruik voor ongestructureerde data zoals afbeeldingen, audio en tekstbestanden die zijn opgeslagen in Unity Catalog-volumes UCVolumeDataset uit de module databricks.air.data. UCVolumeDataset is een PyTorch IterableDataset waarmee elk bestand van het volume naar een snelle lokale cache wordt gekopieerd bij eerste toegang en het lokale bestandspad in de cache wordt gegenereerd. Hiermee worden de prestatie- en distributieproblemen afgehandeld die u anders handmatig zou implementeren:

  • Lokale cache. Bestanden worden gekopieerd van de FUSE-koppeling naar een lokale cachemap bij eerste toegang en worden daarna vanuit de cache geleverd, zodat training met meerdere tijdsperioden het volume niet opnieuw leest.
  • Automatisch partitioneren. Wanneer torch.distributed wordt geïnitialiseerd, worden bestanden verdeeld over ranks en vervolgens verder verdeeld over DataLoader workers, zodat elk (rank, worker)-paar een niet-overlappend deel ontvangt zonder extra configuratie.

Opmerking

UCVolumeDataset en databricks.air.data.DataLoader komen uit het databricks-sdk-air pakket. Installeer het met de data extra, die ook een compatibele torch installeert:

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

UCVolumeDataset levert onbewerkte lokale bestandspaden op. Als u deze bestanden in tensors wilt decoderen, verpakt u deze in een seconde IterableDataset die de padstroom verbruikt en past u uw parseerlogica toe. Hierdoor blijven I/O- en parseringsproblemen gescheiden.

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)

De wrapper ontvangt reeds gecachte lokale paden, dus de parsingstap raakt het volume nooit. U kunt extra wrappers koppelen voor uitbreiding, tokenisatie of filteren.

Voor optimale prestaties combineert u UCVolumeDataset met databricks.air.data.DataLoader in plaats van de standaard-PyTorch DataLoader. Het is afgestemd op AI Runtime I/O en haalt en cachet tegelijkertijd bestanden op terwijl de GPU berekent.

Checkpointmodellen op volumes

Om je model te checkpointen zodat je de training kunt hervatten vanaf de nieuwste snapshot of kunt herstellen van een crash, kun je Unity Catalog-volumes gebruiken zoals een lokaal bestandssysteem.

Databricks raadt aan om een gedistribueerd checkpoint (DCP) te gebruiken voor betere prestaties op zowel single-GPU- als multi-GPU-workloads. Bekijk snelle, fouttolerante PyTorch-training op AI Runtime van de Databricks engineeringblog.

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

De monolithische torch.save-aanpak werkt ook.

  • Voor checkpointing met een enkel GPU-model,

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

  • Voor gedistribueerde training die met torchrun wordt gestart,

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

Tabulaire gegevens laden

Gebruik Spark Connect om machine learning-gegevens uit Delta-tabellen in tabelvorm te laden.

Voor training met één knooppunt kunt u Apache Spark DataFrames converteren naar pandas DataFrames met behulp van de PySpark-methodetoPandas() en optioneel converteren naar NumPy-indeling met behulp van de PySpark-methodeto_numpy().

Opmerking

Spark Connect zorgt ervoor dat analyse en naamomzetting worden uitgesteld tot uitvoeringstijd, waardoor het gedrag van uw code kan worden gewijzigd. Zie Spark Connect vergelijken met Spark Classic.

Spark Connect ondersteunt de meeste PySpark-API's, waaronder Spark SQL, Pandas-API op Spark, Structured Streaming en MLlib (op basis van DataFrame). Raadpleeg de referentiedocumentatie voor de PySpark-API voor de meest recente ondersteunde API's.

Zie Serverloze rekenbeperkingen voor andere beperkingen.

Laad grote Delta-tabellen met behulp van Unity Catalog-volumes

Voor grote Delta-tabellen die te groot zijn om mee te converteren toPandas(), exporteert u de gegevens naar een Unity Catalog-volume en laadt u deze rechtstreeks met behulp van PyTorch of 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")

Deze aanpak voorkomt spark-overhead tijdens de training en werkt goed voor zowel één GPU als gedistribueerde trainingswerkstromen.

  • Laatst bijgewerkt om