Aller au contenu principal
Dernière mise à jour le

Charger des données sur AI Runtime

info

Aperçu

Cette fonctionnalité est en aperçu public.

Les assets de données et de modèles sont essentiels aux charges de travail d’apprentissage profond et de post-entraînement pour les grands modèles de langage (LLM) et les modèles vision-langage (VLM). Avec AI Runtime, l’accès à tous les assets de données et de modèles s’effectue par l’intermédiaire de Unity Catalog :

  • Volumes Unity Catalog : utilisés principalement pour les grands datasets et les fichiers non structurés, notamment les images, l'audio et le texte.
  • Tables Unity Catalog : utilisées pour les données structurées et tabulaires, accessibles via Spark Connect.

Vos volumes et tables doivent être enregistrés dans Unity Catalog et être accessibles à votre utilisateur ou service principal.

Volume Unity Catalog pour les données non structurées​

Les volumes Unity Catalog fournissent un accès gouverné aux données non tabulaires dans n’importe quel format, y compris pour les données structurées, semi-structurées et non structurées. Dans AI Runtime, les volumes constituent le principal mécanisme d’accès aux grands datasets, aux textes, aux assets de modèles et aux points de contrôle de modèles.

Les utilisateurs peuvent lister, lire et écrire des fichiers dans les volumes Unity Catalog à l’aide d’opérations de système de fichiers familières, semblables à la manipulation de fichiers sur un disque local :

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

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

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

Quelques caractéristiques des volumes Unity Catalog les rendent particulièrement bien adaptés aux charges de travail de machine learning :

  • Stockage distribué : Unity Catalog s'appuie sur un stockage distribué, ce qui permet aux charges de travail AI Runtime de lire et d'écrire des données et des assets de modèle sur l'ensemble de la plateforme, depuis des notebooks et des charges de travail basées sur une CLI.
  • Optimisé pour les schémas d'accès au ML : le stockage sous-jacent et les chemins d'accès sont optimisés pour les charges de travail de ML courantes, en particulier les grands fichiers avec des lectures et des écritures séquentielles. Unity Catalog est ainsi parfaitement adapté au chargement des données d'entraînement, au chargement d'asset de modèle et à l'écriture de points de contrôle de modèle.
  • Accès semblable à un système de fichiers : les utilisateurs peuvent lister, lire et écrire des fichiers dans des volumes Unity Catalog à l'aide d'opérations de système de fichiers familières, de la même manière qu'ils travaillent avec des fichiers sur un disque local.

En raison des commits en arrière-plan automatiques, les utilisateurs peuvent compter sur un accès cohérent aux données des volumes Unity Catalog :

  • Écritures : AI Runtime effectue automatiquement les commits des écritures, ce qui rend les modifications visibles par les autres applications et charges de travail accédant au même volume Unity Catalog.
  • Lectures : AI Runtime détecte automatiquement les modifications apportées au volume sans qu’aucune opération de refresh ou de synchronisation explicite ne soit requise de la part de l’utilisateur.

Optimiser les performances des volumes​

Comme indiqué, les volumes Unity Catalog s'appuient sur un stockage distribué et sont optimisés pour les fichiers volumineux soumis à des lectures et écritures séquentielles.

Quelques conseils peuvent vous aider à obtenir les meilleures performances d'AI Runtime :

  • Concaténer les données dans des fichiers plus volumineux : dans la mesure du possible, consolidez les données en un plus petit nombre de fichiers plus grands, d'une taille approximative comprise entre 1 Gio et 10 Gio par fichier. Cela permet à AI Runtime d'effectuer une pré-récupération agressive des données et d'atteindre automatiquement des performances de lecture séquentielle quasi optimales.

  • Pour les charges de travail axées sur les petits fichiers, utilisez le disque local : si votre charge de travail implique de nombreux petits fichiers, envisagez de les copier sur le disque local à l'aide de copies parallèles avant de les traiter. Cela permet de réduire la surcharge liée à l'accès répété à de nombreux petits fichiers via le volume.

    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

  • Vous pouvez utiliser UCVolumeDataset pour vos charges de travail de Machine Learning. Il intègre les optimisations décrites ci-dessus pour assurer un accès et un chargement efficaces des données à partir des volumes Unity Catalog. Consultez les sections suivantes pour plus de détails.

Charger 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 des volumes Unity Catalog, utilisez UCVolumeDataset à partir 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 renvoie le chemin du fichier local mis en cache. Il gère les problèmes de performance et de distribution que vous auriez sinon à mettre en œuvre vous-même :

  • Mise en cache locale. Les fichiers sont copiés du montage FUSE vers un répertoire de cache local lors du premier accès et servis à partir du cache par la suite, de sorte que l'entraînement multi-époque ne relit pas le volume.
  • **Partitionnement automatique.** Lorsque torch.distributed est initialisé, les fichiers sont partitionnés entre les rangs, puis divisés entre DataLoader Workers, de sorte que chaque paire (rank, worker) reçoit une tranche non superposée sans configuration supplémentaire.
remarque

UCVolumeDataset et databricks.air.data.DataLoader proviennent du package databricks-sdk-air. Installez-le avec l'option data, qui intègre également un élément torch compatible :

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

UCVolumeDataset produit des chemins de fichiers locaux bruts. Pour décoder ces fichiers en tenseurs, enveloppez-le dans un deuxième IterableDataset qui consomme le Stream de chemin et applique votre logique d'analyse. Cela permet de séparer les préoccupations liées aux E/S et à l'analyse.

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)

Le wrapper reçoit des chemins locaux déjà mis en cache, de sorte que l'étape d'analyse n'interagit jamais avec 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'au DataLoader PyTorch standard. Il est optimisé pour les E/S d'AI Runtime, et récupère et met en cache des fichiers de manière concurrente pendant que le GPU effectue les calculs.

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

Pour créer un point de contrôle (checkpoint) de votre modèle afin de pouvoir reprendre l'entraînement à partir du dernier instantané ou récupérer après un plantage, vous pouvez utiliser les volumes Unity Catalog exactement comme un système de fichiers local.

Databricks recommande d'utiliser un point de contrôle distribué (DCP) pour de meilleures performances sur les charges de travail à GPU unique ou multiples. Consultez Fast, fault-tolerant PyTorch training on AI Runtime sur le blog de Data Engineering de Databricks.

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

L’approche monolithique torch.save fonctionne également.

  • Pour la création de points de contrôle de modèles à GPU unique,

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

  • Pour l'entraînement distribué lancé via 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"])

Charger des données tabulaires​

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

Pour l'entraînement sur nœud unique, vous pouvez convertir les DataFrames Apache Spark en DataFrames Pandas à l'aide de la méthode PySpark toPandas(), puis, en option, les convertir au format NumPy à l'aide de la méthode PySpark to_numpy().

remarque

Spark Connect diffère l'analyse et la résolution de noms à la durée d'exécution, ce qui peut modifier le comportement de votre code. Voir Comparaison de Spark Connect à Spark Classic.

Spark Connect prend en charge la plupart des API PySpark, y compris Spark SQL, l'API Pandas sur Spark, Structured Streaming et MLlib (basé sur les DataFrames). Consultez la documentation de référence de l'API PySpark pour les dernières APIs prises en charge.

Pour les autres limitations, consultez les limitations du compute Serverless.

Charger de grandes tables Delta à l'aide des volumes Unity Catalog​

Pour les grandes tables Delta qui sont trop grandes pour être converties avec toPandas(), exportez les données vers un volume Unity Catalog et chargez-les directement à l'aide de PyTorch ou de 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")

Cette approche évite le surcoût de Spark pendant l'entraînement et fonctionne bien pour les workflows d'entraînement à GPU unique et distribués.

Sur cette page