Le chargeur de données de recherche IBM améliore la formation aux modèles d'IA pour la communauté open source

Comment surmontez-vous les blocages lorsque vous entraînez des modèles IA sur des quantités massives de données ? Lors de la conférence PyTorch de cette année, IBM Recherche a présenté un chargeur de données révolutionnaire pour la formation à grande échelle de LLM. Cet outil, désormais disponible pour les utilisateurs de PyTorch, vise à simplifier la formation à grande échelle pour un public aussi large que possible.

L’idée du chargeur de données à haut débit découle des problèmes pratiques observés par les chercheurs lors de l’entraînement des modèles, car leur travail nécessitait un outil capable de traiter de grandes quantités de données sur plusieurs appareils, tout en permettant de gérer des GPU à efficacité progressive. Comme le note IBM Research dans son blog consacré à cette publication, « tout cela est dû à une équipe de chercheurs qui construisaient simplement les outils dont ils avaient besoin pour accomplir leur travail ».

Davis Wertheimer d'IBM recherche explique certains des défis qui peuvent survenir lors d'une formation à grande échelle : « Il existe une sorte de règle des 80/20 en matière de formation à grande échelle. Quatre-vingts pour cent de la littérature publiée examine les compromis algorithmiques entre la mémoire GPU, la communication et le calcul. Mais lorsque vous essayez de construire quelque chose, dans 80 % des cas, vous pouvez dépendre d’un très grand nombre de problèmes pratiques, car le pipeline fonctionne à la vitesse du goulot d’étranglement le plus étroit. »

Au fur et à mesure que l'équipe IBM développait sa plateforme de formation, elle continuait à rencontrer des problèmes. « À mesure que nous améliorons l’utilisation de nos GPU, le goulot d’étranglement est de plus en plus souvent le chargeur de données », observe M. Wertheimer.

Cette prise de conscience a conduit à un double processus de développement. « Il y a eu un parcours parallèle, d'une part, pour faire évoluer notre plateforme d'entraînement, et, d'autre part, pour faire évoluer en permanence notre chargeur de données pour s'adapter aux exigences de vitesse de notre plateforme et éviter les goulots d'étranglement », explique-t-il.

Linsong Chu d'IBM Research présente les fonctionnalités essentielles du chargeur de données :

Avec état et point de contrôle : « Chaque fois que vous enregistrez un modèle, l'état de votre chargeur de données est également enregistré, et chaque fois que vous récupérez à partir d'un point de contrôle, l'état du modèle et les états du chargeur de données doivent être récupérés en même temps », explique M. Chu.

Réajustement automatique des points de contrôle : le chargeur de données s’adapte automatiquement aux variations de Workload lors des sessions d’entraînement prolongées. « La formation peut facilement prendre des semaines, voire des mois, et il existe une multitude de raisons pour lesquelles vous pourriez avoir à réajuster votre Workload en cours de route », souligne M. Chu.

Diffusion en continu de données efficace : le système prend en charge la diffusion en continu de données sans frais généraux de construction pour le brouillage.

Opération distribuée asynchrone : « Nous voulons que le chargeur de données ne soit pas bloquant », explique M. Chu. « Tout en économisant l’état du chargeur de données, nous voulons que l’économie soit répartie sous une forme où aucune communication n’est impliquée. »

Mélange dynamique des données : Le chargeur de données peut s'adapter à différents ratios de mélange de données, ce qui est utile pour répondre à l'évolution des besoins de formation.

Brassage global efficace : L'outil adresse les goulots d'étranglement de mémoire lors du traitement de grands jeux de données, rendant le brassage efficace même lorsque les données augmentent.

PyTorch natif, modulaire et étendu : conçu pour l'adaptabilité et l'évolutivité, le chargeur de données est prêt pour la croissance future. « Et si l’année prochaine nous devions gérer 30 billions, 50 billions ou 100 billions de tokens ? » demande M. Chu. « Le monde évolue rapidement. Nous devons donc créer le chargeur de données pour qu'il puisse non seulement survivre aujourd'hui, mais aussi demain. »

L’équipe de recherche d’IBM a testé rigoureusement son chargeur de données pendant plusieurs mois, en exécutant des centaines de petites et grandes tâches. Ils ont observé des numéros de code stables et fluides. De plus, l’ensemble du chargeur de données fonctionne de manière asynchrone et n’est pas bloquant.

« Nous avons tiré parti de nombreuses capacités intégrées de PyTorch pour y parvenir », explique Wertheimer. « C'est pourquoi nous contribuons, nous y contribuons en retour. »