Caricare i dati nel runtime di intelligenza artificiale

Importante

Questa funzionalità è in Anteprima Pubblica.

I dati e le risorse dei modelli sono fondamentali per il deep learning e per i workload di post-addestramento dei modelli linguistici di grandi dimensioni (LLM) e dei modelli vision-language (VLM). Con AI Runtime, tutti i dati e gli asset del modello sono accessibili tramite Unity Catalog:

Volumi del Catalogo Unity: utilizzato principalmente per grandi dataset e file non strutturati, inclusi immagini, audio e testo.
Tabelle del Catalogo Unity: utilizzate per dati strutturati e tabellari, accessibili tramite Spark Connect.

I tuoi volumi e le tue tabelle devono essere registrati nel Catalogo Unity e accessibili al tuo utente o al principale del servizio.

Volume di Unity Catalog per dati non strutturati

I volumi del Catalogo Unity forniscono accesso governato a dati non tabulari in qualsiasi formato, inclusi dati strutturati, semi-strutturati e non strutturati. In AI Runtime, i volumi sono il principale meccanismo per accedere a grandi dataset, testi, asset di modello e checkpoint di modello.

Gli utenti possono elencare, leggere e scrivere file nei volumi del Catalogo Unity utilizzando operazioni familiari del file system, simili al lavoro con file su un disco locale:

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

Analogamente, le operazioni shell funzionano allo stesso modo:

%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

Alcune caratteristiche dei volumi del Catalogo Unity li rendono particolarmente adatti ai carichi di lavoro di machine learning:

Archiviazione distribuita: Unity Catalog è supportato da storage distribuito, permettendo ai carichi di lavoro AI Runtime di leggere e scrivere dati e modellare asset su tutta la piattaforma, sia da notebook che da carichi di lavoro basati su CLI.
Ottimizzato per pattern di accesso ML: Gli storage e i percorsi di accesso sottostanti sono ottimizzati per carichi di lavoro ML comuni, in particolare file grandi con letture e scritture sequenziali. Questo rende Unity Catalog particolarmente adatto per il caricamento dei dati di addestramento, il caricamento degli asset del modello e il salvataggio dei checkpoint del modello.
Accesso simile a un file system: gli utenti possono elencare, leggere e scrivere file nei volumi del Catalogo Unity utilizzando operazioni di file system familiari, simile al lavoro con file su un disco locale.

Grazie ai commit automatici in background, gli utenti possono contare su un accesso continuo ai dati dei volumi di Unity Catalog:

Scritture: AI Runtime esegue automaticamente il commit delle scritture, rendendo visibili le modifiche ad altre applicazioni e carichi di lavoro che accedono allo stesso volume di Unity Catalog.
Lettura: AI Runtime rileva automaticamente le modifiche al volume senza richiedere alcuna operazione esplicita di aggiornamento o sincronizzazione da parte dell'utente.

Esecuzione a volume della melodia

Come detto, i volumi del Catalogo Unity sono supportati da uno storage distribuito e ottimizzati per file di grandi dimensioni con letture e scritture sequenziali.

Alcuni consigli possono aiutarti a ottenere le migliori prestazioni dal tempo di esecuzione dell'IA:

Concatena i dati in file più grandi: Quando possibile, consolida i dati in meno file più grandi, circa da 1 GiB a 10 GiB per file. Questo permette a AI Runtime di pre-raccogliere i dati in modo aggressivo e di raggiungere automaticamente prestazioni di lettura sequenziale quasi ottimali.

Per carichi di lavoro con file piccoli, usa disco locale: se il tuo carico di lavoro coinvolge molti file piccoli, considera di copiare i file sul disco locale usando copie parallele prima di elaborarli. Questo può ridurre il sovraccarico di accedere ripetutamente a molti piccoli file attraverso il 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

Puoi usarlo UCVolumeDataset per i tuoi carichi di lavoro di machine learning. Incorpora le ottimizzazioni sopra descritte per fornire un accesso e un caricamento efficiente ai dati dai volumi del Catalogo Unity. Per altri dettagli, vedere le sezioni seguenti.

Carica dati non strutturati con `UCVolumeDataset`

Per dati non strutturati come immagini, audio e file di testo memorizzati nei volumi del Catalogo Unity, si utilizza UCVolumeDataset dal modulo databricks.air.data. UCVolumeDataset è un PyTorch IterableDataset che copia ogni file dal volume a una cache locale veloce al primo accesso e restituisce il percorso del file locale memorizzato nella cache. Gestisce le prestazioni e gli aspetti legati alla distribuzione che altrimenti dovresti implementare manualmente:

Memorizzazione nella cache locale. I file vengono copiati dal punto di montaggio FUSE in una directory di cache locale al primo accesso e vengono poi letti dalla cache, quindi l'addestramento su più epoche non richiede una nuova lettura del volume.
Partizionamento automatico. Quando torch.distributed viene inizializzato, i file vengono partizionati tra i ranghi e quindi suddivisi ulteriormente tra DataLoader i ruoli di lavoro, quindi ogni (rank, worker) coppia riceve una sezione non sovrapposta senza alcuna configurazione aggiuntiva.

Annotazioni

UCVolumeDataset e databricks.air.data.DataLoader provengono dal databricks-sdk-air pacchetto. Installalo con l'extra data, che include anche un torch compatibile:

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

UCVolumeDataset restituisce percorsi di file locali non elaborati. Per decodificare questi file in tensori, racchiuderli in un secondo IterableDataset che utilizza il flusso di percorsi e applica la logica di parsing. In questo modo, le problematiche di I/O e analisi sono separate.

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)

Il wrapper riceve percorsi locali già memorizzati nella cache, quindi la fase di parsing non accede mai al volume. È possibile concatenare wrapper aggiuntivi per l'aumento, la tokenizzazione o il filtro.

Per ottenere prestazioni ottimali, abbina UCVolumeDataset a databricks.air.data.DataLoader anziché a DataLoader di PyTorch standard. È ottimizzato per l'I/O in tempo di esecuzione AI e recupera e memorizza file in cache contemporaneamente mentre la GPU calcola.

Modelli di checkpoint per i volumi

Per fare checkpoint al tuo modello così da poter riprendere l'addestramento dall'ultimo snapshot o recuperare da un crash, puoi usare i volumi del catalogo Unity proprio come un file system locale.

Databricks raccomanda di utilizzare un checkpoint distribuito (DCP) per migliori prestazioni sia su carichi di lavoro a singola GPU che su più GPU. Vedi l'articolo Addestramento PyTorch rapido e tollerante ai guasti su AI Runtime dal blog di ingegneria di Databricks.

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

Anche l'approccio monolitico torch.save funziona.

Per il checkpoint di modello a GPU singola,

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

Per l'addestramento distribuito lanciato tramite 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"])

Caricare dati tabulari

Usare Spark Connect per caricare dati di Machine Learning tabulari da tabelle Delta.

Per il training a nodo singolo, è possibile convertire i DataFrame Apache Spark in DataFrame pandas utilizzando il metodo PySparktoPandas(), e quindi, facoltativamente, convertire in formato NumPy utilizzando il metodo PySparkto_numpy().

Annotazioni

Spark Connect rinvia l'analisi e la risoluzione dei nomi al tempo di esecuzione, che può modificare il comportamento del codice. Vedi Confronta Spark Connect con Spark Classic.

Spark Connect supporta la maggior parte delle API PySpark, tra cui Spark SQL, API Pandas in Spark, Structured Streaming e MLlib (basata su dataframe). Vedere la documentazione di riferimento dell'API PySpark per le API supportate più recenti.

Per altre limitazioni, vedere Limitazioni di calcolo serverless.

Carica grandi tabelle Delta usando volumi del Catalogo Unity

Per tabelle Delta di grandi dimensioni troppo grandi per la conversione con toPandas(), esportare i dati in un volume di Catalogo Unity e caricarli direttamente usando PyTorch o 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")

Questo approccio evita il sovraccarico di Spark durante il training e funziona bene sia per i flussi di lavoro di training a GPU singola che per i flussi di lavoro di training distribuiti.

Commenti e suggerimenti

Questa pagina è stata utile?

Ultimo aggiornamento: 2026-10-05