Aller au contenu principal
Dernière mise à jour le

Charger des données sur AI Runtime

Les données et les modèles sont cruciaux pour les 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, toutes les données et les modèles sont accessibles via Unity Catalog :

  • Volumes Unity Catalog : utilisés principalement pour les grands jeux de données 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 accessibles à votre utilisateur ou service principal.

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

Les volumes Unity Catalog 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. Dans AI Runtime, les volumes sont le principal mécanisme d'accès aux grands datasets, au texte, aux assets de modèle et aux checkpoints de modèle.

Les utilisateurs peuvent répertorier, lire et écrire des fichiers dans des volumes Unity Catalog à l’aide d’opérations système familières, similaires à celles utilisées pour travailler avec des 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 actifs de modèle sur la plateforme, à partir de notebooks comme de charges de travail basées sur une CLI.
  • Optimisé pour les modèles d'accès ML : le stockage sous-jacent et les chemins d'accès sont optimisés pour les charges de travail ML courantes, en particulier les fichiers volumineux avec des lectures et des écritures séquentielles. C'est pourquoi Unity Catalog est particulièrement adapté au chargement des données d'entraînement, au chargement des assets de modèle et à l'écriture des points de contrôle des modèles.
  • Accès similaire à un système de fichiers : les utilisateurs peuvent répertorier, lire et écrire des fichiers dans les volumes Unity Catalog en utilisant des opérations de système de fichiers familières, semblables à l’utilisation de fichiers sur un disque local.

En raison des commit en arrière-plan automatiques, les utilisateurs peuvent bénéficier d'un accès cohérent aux données de volume Unity Catalog :

  • Écritures : AI Runtime commit automatiquement les écritures, rendant les modifications visibles aux 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 nécessiter d’opération de refresh ou de synchronisation explicite de la part de l’utilisateur.

Optimiser les performances des volumes​

Comme mentionné, les volumes Unity Catalog sont pris en charge par un stockage distribué et optimisés pour les gros fichiers avec des lectures et des écritures séquentielles.

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

  • Concaténer les données dans des fichiers plus volumineux : dans la mesure du possible, consolidez les données dans un nombre limité de fichiers plus volumineux, d'environ 1 Gio à 10 Gio par fichier. Cela permet à l'IA Runtime de préextraire activement les données et d'atteindre automatiquement des performances de lecture séquentielle quasi optimales.

  • Pour les charges de travail contenant de petits fichiers, utilisez le disque local : si votre charge de travail implique 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 les frais généraux liés à 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 en machine learning. Elle intègre les optimisations décrites ci-dessus pour assurer un accès efficace aux données et leur chargement depuis des volumes Unity Catalog. Consultez les sections suivantes pour en savoir plus.

Charger des données non structurées avec UCVolumeDataset​

Pour les données non structurées comme les images, l’audio et les fichiers texte stockés dans des volumes Unity Catalog, utilisez UCVolumeDataset du module databricks.air.data. UCVolumeDataset est un IterableDataset PyTorch qui copie chaque fichier du volume dans un cache local rapide lors du premier accès et donne le chemin du fichier local mis en cache. Il gère les problèmes de performance et de distribution que vous devriez autrement mettre en œuvre manuellement :

  • 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 sont servis à partir du cache par la suite, de sorte que l’entraînement multi-époques ne relit pas le volume.
  • Partitionnement automatique. Lorsque torch.distributed est initialisé, les fichiers sont partitionnés entre les rangs puis divisés entre les DataLoader Worker, de sorte que chaque paire (rank, worker) reçoit une tranche sans chevauchement sans configuration supplémentaire.
remarque

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

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

UCVolumeDataset donne des chemins de fichiers locaux bruts. Pour décoder ces fichiers en tenseurs, enveloppez-les dans un second IterableDataset qui consomme le flux de chemin Stream et applique votre logique d'analyse. Cela permet de séparer les problèmes d'E/S et d'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à en cache, de sorte que l’étape d’analyse ne touche jamais le volume. Vous pouvez enchaî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 PyTorch DataLoader standard. Il est optimisé pour les opérations d’E/S d’AI Runtime et récupère et met en cache les fichiers simultanément pendant que le GPU compute.

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

Pour créer un checkpoint de votre modèle afin de pouvoir reprendre l’entraînement à partir du dernier snapshot 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 sur les charges de travail GPU simple et multi-GPU. Voir Entraînement PyTorch rapide et tolérant aux pannes sur AI Runtime dans le blog de Databricks Data Engineering.

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 le checkpointing de modèles sur 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 un 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 à nœud unique, vous pouvez convertir les DataFrames Apache Spark en DataFrames pandas à l'aide de la méthode PySpark toPandas(), puis éventuellement convertir au format NumPy à l'aide de la méthode PySpark to_numpy().

remarque

Spark Connect diffère l’analyse et la résolution des noms au moment de l’exécution, ce qui peut modifier le comportement de votre code. Voir Comparaison de Spark Connect et Spark Classic.

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

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

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

Pour les tables Delta volumineuses trop importantes pour être converties avec toPandas(), exportez les données vers un volume Unity Catalog et chargez-les directement avec PyTorch ou 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 la surcharge Spark pendant l'entraînement et fonctionne bien pour les workflows d'entraînement à GPU unique et distribué.

Sur cette page