Accélérez l'entraînement sur GPU avec les tables Unity Catalog et UCDataset

Les GPU représentent généralement la partie la plus coûteuse de tout entraînement, et la plupart des équipes n'arrivent pas à les faire tourner à plein régime. Le goulet d'étranglement ne vient généralement pas du modèle. Il s'agit d'acheminer les données depuis leur emplacement vers le GPU assez rapidement, et chaque seconde où ces puces restent inactives, vous gaspillez à la fois du temps et de l'argent.

Nous lançons aujourd'hui Unity Catalog Dataset (UCDataset) en version bêta sur AI Runtime. UCDataset est un chargeur de données PyTorch natif qui diffuse les tables Unity Catalog directement vers chaque GPU de votre tâche d'entraînement, sans exportation, sans copie et sans code de chargement personnalisé.

Sur des charges de travail de production simulées, UCDataset a réduit le temps d'entraînement par 10 par rapport à la lecture d'une copie Parquet exportée des mêmes données.

Le problème : vos GPU attendent vos données

La plupart des données d'entraînement pour vos modèles résident déjà dans des tables Unity Catalog. Les intégrer dans une tâche d'entraînement distribué implique de choisir l'une de ces trois solutions de contournement, et chacune d'elles vous coûte du temps et de l'argent :

• Des exportations lentes et fastidieuses avant chaque exécution : Le modèle le plus courant consiste à exporter la table au format Parquet (ou dans un autre format fragmenté) dans un volume avant chaque exécution, puis à pointer un jeu de données vers les fichiers. Vous payez pour stocker une seconde copie, elle devient obsolète dès que de nouvelles données arrivent, et elle échappe à la traçabilité de Unity Catalog.
• Charger toute la table en mémoire ne passe pas à l'échelle : Lire la table dans pandas avec toPandas() fonctionne pour les petits jeux de données et échoue pour les grands. Cela charge également tout sur un seul nœud, ce qui n'aide pas une tâche s'entraînant sur 32 GPU.
• Le streaming via un entrepôt de données est lent : Les entrepôts n'ont pas été conçus à l'origine pour alimenter des GPU et, sans optimisation, ils seront bridés lorsque de nombreux GPU liront en même temps.

Vient ensuite le chargeur lui-même. Chaque GPU et chaque worker a besoin de sa propre portion équilibrée de données. Les équipes passent du temps à construire cela manuellement, et un léger bug peut faire en sorte que le modèle s'entraîne sur des lignes dupliquées ou ignorées sans que personne ne s'en aperçoive.

Le résultat : des GPU coûteux restent inactifs, l'entraînement prend plus de temps que nécessaire et les ingénieurs ML s'occupent de la plomberie des données au lieu d'améliorer les modèles.

Présentation de UCDataset

UCDataset fait des tables Unity Catalog une source de données de premier ordre pour l'entraînement PyTorch distribué. Vous décrivez ce qu'il faut lire (une table, les colonnes dont vous avez besoin et un filtre de lignes) et UCDataset diffuse des lots prêts pour l'entraînement vers chaque GPU. Les avantages incluent :

• Conçu pour la vitesse : UCDataset lit les fichiers de la table directement depuis le stockage cloud, de sorte que rien ne s'intercale entre vos données et vos GPU. Il répartit uniformément les lignes correspondantes entre chaque GPU et chaque worker de chargement de données.
• Entraînez-vous sur le lake, pas sur une copie : Pointez vers la table en direct. Les filtres et les vecteurs de suppression s'appliquent au fur et à mesure que les données sont lues, et la prochaine exécution récupère automatiquement les nouvelles lignes ajoutées.
• Gouvernance de bout en bout : Les tâches d'entraînement lisent les données avec l'identité de l'utilisateur qui les exécute, et les autorisations Unity Catalog s'appliquent exactement comme partout ailleurs sur Databricks.

Benchmark : un entraînement plus rapide sur une charge de travail de modèle de recommandation

Nous avons évalué UCDataset sur une charge de travail de production simulée, en diffusant une table de caractéristiques de 300 Go sur 8 GPU NVIDIA H100 pour entraîner un modèle de recommandation, et nous avons observé en moyenne une amélioration de 10x du temps d'entraînement par rapport au chargement direct des données depuis Parquet.

Comment l'utiliser :

Imaginons qu'une équipe conserve des caractéristiques d'utilisateurs et d'éléments avec un libellé de clic stocké dans une table main.recsys.ctr_features. De nouvelles impressions arrivent chaque jour dans la table. Ils souhaitent entraîner un modèle de taux de clics sur 16 GPU répartis sur deux nœuds, directement à partir de la table.

1. Décrivez ce qu'il faut lire et créez le chargeur. Dans train.py, sélectionnez la table, conservez les lignes d'entraînement étiquetées et choisissez uniquement les colonnes dont le modèle a besoin. Le filtre est appliqué lors de la lecture des fichiers de la table, de sorte que les lignes ne correspondant pas n'atteignent jamais votre code d'entraînement. UCDataset génère déjà des lots entiers, vous pouvez donc l'envelopper dans un DataLoader PyTorch standard avec batch_size=None.

2. Entraînez et lancez sur deux nœuds. Entraînez avec une boucle DDP PyTorch standard. UCDataset attribue à chacun des 16 GPU sa propre part des lignes correspondantes, il n'y a donc aucun code d'échantillonnage ou de partitionnement à écrire. Décrivez la tâche dans workload.yaml, qui lance train.py avec torchrun sur chaque nœud, et exécutez-la avec l'interface CLI d'AI Runtime :

Pour vous entraîner sur un plus grand nombre de GPU, modifiez num_accelerators dans workload.yaml. Le code d'entraînement reste le même. Le code complet de train.py et workload.yaml, avec une table d'exemple pour les tester, se trouve dans le livre de recettes AI Runtime : Modèle de taux de clics sur une table Unity Catalog.

Pour commencer

Installez le package depuis PyPI. Dans un notebook, exécutez %pip install databricks-uc-dataset.

UCDataset est disponible en version bêta sur AI Runtime. Pour l'essayer :

• Consultez la documentation : Charger des tables Unity Catalog dans PyTorch avec UCDataset
• Partez d'un exemple fonctionnel : recette du livre de recettes
• Contactez votre équipe de compte, nous serions ravis d'entendre vos commentaires.

Pour les données non structurées stockées dans les volumes Unity Catalog, utilisez UCVolumeDataset pour diffuser des fichiers avec mise en cache locale. Bien qu'UCDataset gère les tables structurées, nous l'étendons également pour prendre en charge l'entraînement multimodal via le type FILE, vous permettant de vous entraîner sur des images, des fichiers audio et des documents directement à partir d'une table. La version bêta lit les tables Delta dans Unity Catalog.

FAQ

Qu'est-ce qu'UCDataset ?

UCDataset (Unity Catalog Dataset) est un chargeur de données natif qui diffuse les tables Unity Catalog directement vers les GPU d'une tâche d'entraînement. Il est disponible en version bêta sur Databricks AI Runtime. Vous choisissez une table, les colonnes dont vous avez besoin, et UCDataset fournit des lots prêts pour l'entraînement sans étape d'exportation ni code de chargement personnalisé.

Comment UCDataset accélère-t-il l'entraînement sur GPU ?

UCDataset maintient les GPU actifs en lisant les fichiers de données d'une table directement depuis le stockage cloud, sans copie exportée intermédiaire. Il répartit uniformément les lignes correspondantes entre chaque GPU et chaque worker de chargement de données, de sorte que chaque GPU ne lit que sa propre part.

Dois-je exporter les tables Unity Catalog au format Parquet pour les entraîner ?

Non. UCDataset lit la table en direct dans Unity Catalog, il n'y a donc pas d'étape d'exportation ni de seconde copie à stocker ou à maintenir synchronisée. Il applique votre filtre et les vecteurs de suppression de la table au fur et à mesure de la lecture, et chaque nouvelle exécution récupère les lignes ajoutées depuis la précédente.

UCDataset applique-t-il les autorisations Unity Catalog ?

Oui. Sur AI Runtime, les tâches d'entraînement lisent les données avec l'identité de l'utilisateur qui les exécute, et les autorisations Unity Catalog sur la table s'appliquent. Comme la tâche lit directement la table gouvernée, il n'y a pas de copie exportée à sécuriser séparément. Dans la version bêta, UCDataset ne peut pas lire les tables qui comportent des filtres de lignes ou des masques de colonnes.

Quelles tables UCDataset peut-il lire ?

UCDataset lit les tables Delta enregistrées dans Unity Catalog, y compris les tables avec des vecteurs de suppression.

(Cet article de blog a été traduit à l'aide d'outils basés sur l'intelligence artificielle) Article original