Charger des données sur AI Runtime

Les données et les ressources de modèles sont essentielles pour l’apprentissage profond et les charges post-entraînement pour les grands modèles de langage (LLM) et les modèles vision-langage (VLM). Avec l’exécution IA, toutes les données et les ressources du modèle sont accessibles via le catalogue Unity :

Volumes du catalogue Unity : principalement utilisé pour de grands ensembles de données et fichiers non structurés, y compris images, audio et texte.
Tables du catalogue Unity : utilisées pour les données structurées et tabulaires, accessibles via Spark Connect.

Vos volumes et tables doivent être enregistrés dans le catalogue Unity et accessibles à votre utilisateur ou responsable de service.

Volume du catalogue Unity pour les données non structurées

Les volumes du catalogue Unity offrent un accès régi aux données non tabulaires dans n’importe quel format, y compris les données structurées, semi-structurées et non structurées. En temps d’exécution IA, les volumes sont le principal mécanisme d’accès à de grands ensembles de données, textes, ressources de modèles et points de contrôle de modèles.

Les utilisateurs peuvent lister, lire et écrire des fichiers dans des volumes du catalogue Unity en utilisant des opérations familières du système de fichiers, un peu comme travailler avec des fichiers sur un disque local :

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

De même, les opérations shell fonctionnent de la même manière :

%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

Quelques caractéristiques des volumes du catalogue Unity les rendent bien adaptés aux charges de travail d’apprentissage automatique :

Stockage distribué : Unity Catalog est soutenu par un stockage distribué, permettant aux charges de travail en temps exécutable IA de lire et d’écrire des données et de modéliser les actifs sur toute la plateforme, à partir de notebooks et de charges de travail basées sur CLI.
Optimisé pour les patrons d’accès ML : Les chemins de stockage et d’accès sous-jacents sont optimisés pour les charges de travail ML courantes, en particulier les gros fichiers avec lectures et écritures séquentielles. Cela rend Unity Catalog très adapté au chargement de données d’entraînement, au chargement d’actifs de modèles et à l’écriture de points de contrôle de modèles.
Accès similaire à un système de fichiers : Les utilisateurs peuvent lister, lire et écrire des fichiers dans des volumes du catalogue Unity en utilisant des opérations familières du système de fichiers, similaire au travail avec des fichiers sur un disque local.

Grâce aux commits automatiques en arrière-plan, les utilisateurs peuvent s’attendre à un accès régulier aux données volumineuses du catalogue Unity :

Écritures : AI Runtime valide automatiquement les écritures, ce qui rend les modifications visibles aux autres applications et charges de travail accédant au même volume Unity Catalog.
Lecture : l’environnement d’exécution de l’IA détecte automatiquement les modifications du volume sans nécessiter d’actualisation ni de synchronisation explicites de la part de l’utilisateur.

Régler les performances sonores

Comme mentionné, les volumes du catalogue Unity sont soutenus par un stockage distribué et optimisés pour les gros fichiers avec lectures et écritures séquentielles.

Quelques conseils peuvent vous aider à obtenir les meilleures performances de l’IA Runtime :

Concaténouez les données en fichiers plus volumineux : Lorsque possible, consolidez les données en moins de fichiers plus grands, soit environ 1 Go à 10 Gi par fichier. Cela permet à l’exécution IA de précharger agressivement les données et d’atteindre automatiquement des performances de lecture séquentielle quasi optimales.

Pour les charges de travail de petits fichiers, utilisez le disque local : si votre charge de travail comprend de nombreux petits fichiers, envisagez de copier les fichiers sur le disque local en utilisant des copies parallèles avant de les traiter. Cela peut réduire la surcharge liée à l’accès répété à de nombreux petits fichiers via le 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

Vous pouvez l’utiliser UCVolumeDataset pour vos charges de travail en apprentissage automatique. Il intègre les optimisations décrites ci-dessus afin de fournir un accès et un chargement efficaces des données depuis les volumes du catalogue Unity. Pour plus d’informations, consultez les sections suivantes.

Chargez des données non structurées avec `UCVolumeDataset`

Pour les données non structurées telles que les images, les fichiers audio et texte stockés dans les volumes du catalogue Unity, utilisez UCVolumeDataset du module databricks.air.data. UCVolumeDataset est un PyTorch IterableDataset qui copie chaque fichier du volume vers un cache local rapide lors du premier accès et génère le chemin d’accès au fichier local mis en cache. Il gère les aspects liés aux performances et à la distribution que vous devriez autrement implémenter à la main :

Mise en cache locale. Les fichiers sont copiés à partir du montage FUSE vers un répertoire de cache local sur le premier accès et distribués à partir du cache par la suite, de sorte que l’apprentissage multi-époque ne réécrit pas le volume.
Partitionnement automatique. Lorsque torch.distributed est initialisé, les fichiers sont répartis entre les rangs, puis subdivisés entre les processus de travail DataLoader, de sorte que chaque paire (rank, worker) reçoit une portion distincte sans chevauchement, sans configuration supplémentaire.

Note

UCVolumeDataset et databricks.air.data.DataLoader proviennent du databricks-sdk-air paquet. Installez-le avec l’extra data, qui installe aussi une dépendance torch compatible :

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

UCVolumeDataset génère des chemins d’accès de fichiers locaux bruts. Pour décoder ces fichiers en tenseurs, encapsulez le tout dans un second IterableDataset qui consomme le flux de chemins d’accès et applique votre logique d'analyse. Cela maintient les problèmes d’E/S et d’analyse distincts.

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)

Le wrapper reçoit des chemins locaux déjà mis en cache, donc l’étape d’analyse ne touche jamais le volume. Vous pouvez chaîner des wrappers supplémentaires pour l’augmentation, la tokenisation ou le filtrage.

Pour des performances optimales, associez UCVolumeDataset à databricks.air.data.DataLoader plutôt qu’à la version standard de PyTorch DataLoader. Il est réglé pour l’E/S d’exécution IA et collecte et met en cache les fichiers simultanément pendant que le GPU calcule.

Modèles de points de contrôle sur les volumes

Pour enregistrer des points de contrôle de votre modèle afin de pouvoir reprendre l’entraînement à partir du dernier instantané enregistré ou récupérer après un plantage, vous pouvez utiliser les volumes Unity Catalog comme un système de fichiers local.

Databricks recommande d’utiliser un point de contrôle distribué (DCP) pour de meilleures performances aussi bien sur des charges à GPU unique qu’à plusieurs GPU. Voir l’entraînement PyTorch rapide et tolérant aux défaillances sur AI Runtime du blog d’ingénierie de 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"],
)

L’approche monolithique torch.save fonctionne également.

Pour le checkpoint de modèle GPU unique,

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

Pour la formation distribuée lancée via 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"])

Charger les données tabulaires

Utilisez Spark Connect pour charger des données De Machine Learning tabulaires à partir de tables Delta.

Pour l’entraînement à nœud unique, vous pouvez convertir des DataFrames Apache Spark en DataFrames pandas à l’aide de la méthode toPandas(), puis éventuellement convertir en format NumPy à l’aide de la méthodeto_numpy() PySpark.

Note

Spark Connect reporte l’analyse et la résolution de noms au temps d’exécution, ce qui peut modifier le comportement de votre code. Consultez Comparer Spark Connect à Spark Classic.

Spark Connect prend en charge la plupart des API PySpark, notamment Spark SQL, l’API Pandas sur Spark, Structured Streaming et MLlib (dataFrame). Consultez la documentation de référence de l’API PySpark pour connaître les dernières API prises en charge.

Pour d’autres limitations, consultez Limitations du calcul serverless.

Chargez de grandes tables Delta à l’aide de volumes du catalogue Unity

Pour les grandes tables Delta qui sont trop grandes pour les convertir avec toPandas(), exportez les données vers un volume du Unity Catalog et chargez-les directement en utilisant PyTorch ou 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")

Cette approche évite la surcharge Spark pendant l’entraînement et fonctionne correctement pour les flux de travail d’entraînement à GPU unique et distribué.

Commentaires

Cette page a-t-elle été utile ?

Dernière mise à jour : 2026-10-22