Le parcours de KT pour réduire le temps de formation d’un modèle de transformateurs de vision en utilisant Amazon SageMaker

Le chemin de KT vers une réduction du temps de formation d'un modèle de transformateurs de vision en utilisant Amazon SageMaker

KT Corporation est l’un des plus grands fournisseurs de télécommunications en Corée du Sud, proposant une large gamme de services, notamment la téléphonie fixe, la communication mobile et internet, et des services d’IA. Le système de balisage alimentaire basé sur l’IA de KT est une solution de gestion diététique basée sur l’IA qui identifie le type et la teneur nutritionnelle des aliments sur des photos à l’aide d’un modèle de vision artificielle. Ce modèle développé par KT repose sur un modèle pré-entraîné comportant une grande quantité de données d’images non étiquetées pour analyser la teneur nutritionnelle et les informations caloriques de divers aliments. Le système de balisage alimentaire basé sur l’IA peut aider les patients atteints de maladies chroniques telles que le diabète à gérer leur alimentation. KT a utilisé AWS et Amazon SageMaker pour former ce modèle de balisage alimentaire basé sur l’IA 29 fois plus rapidement qu’auparavant et l’optimiser pour le déploiement en production grâce à une technique de distillation de modèle. Dans cet article, nous décrivons le parcours de développement du modèle de KT et son succès avec SageMaker.

Présentation du projet KT et définition du problème

Le modèle de balisage alimentaire basé sur l’IA pré-entraîné par KT est basé sur l’architecture des transformateurs de vision (ViT) et comporte plus de paramètres de modèle que leur précédent modèle de vision pour améliorer la précision. Pour réduire la taille du modèle pour la production, KT utilise une technique de distillation des connaissances (KD) pour réduire le nombre de paramètres de modèle sans impact significatif sur la précision. Avec la distillation des connaissances, le modèle pré-entraîné est appelé modèle enseignant et un modèle de sortie léger est entraîné en tant que modèle étudiant, comme illustré dans la figure suivante. Le modèle étudiant léger a moins de paramètres de modèle que le modèle enseignant, ce qui réduit les besoins en mémoire et permet un déploiement sur des instances plus petites et moins coûteuses. L’étudiant maintient une précision acceptable même s’il est plus petit en apprenant à partir des sorties du modèle enseignant.

Le modèle enseignant reste inchangé pendant la KD, mais le modèle étudiant est entraîné à l’aide des logits de sortie du modèle enseignant en tant qu’étiquettes pour calculer la perte. Avec ce paradigme de KD, le modèle enseignant et le modèle étudiant doivent être sur une seule mémoire GPU pour la formation. KT utilisait initialement deux GPU (A100 80 GB) dans leur environnement interne sur site pour former le modèle étudiant, mais le processus prenait environ 40 jours pour couvrir 300 époques. Pour accélérer la formation et générer un modèle étudiant en moins de temps, KT s’est associé à AWS. Ensemble, les équipes ont considérablement réduit le temps de formation du modèle. Cet article décrit comment l’équipe a utilisé Amazon SageMaker Training, la SageMaker Data Parallelism Library, Amazon SageMaker Debugger et Amazon SageMaker Profiler pour développer avec succès un modèle léger de balisage alimentaire basé sur l’IA.

• Des chercheurs de l’Université Stanford présentent FlashFFTConv un nouveau système d’intelligence artificielle pour optimiser les convolutions FFT pour les longues séquences.
• Le Tencent AI Lab présente Chain-of-Noting (CoN) pour améliorer la robustesse et la fiabilité des modèles de langage améliorés par la recherche.
• Construire un trieur Lego Technic avec une reconnaissance avancée en temps réel des objets

Création d’un environnement de formation distribué avec SageMaker

SageMaker Training est un environnement de formation en apprentissage automatique (ML) géré sur AWS qui offre une suite de fonctionnalités et d’outils pour simplifier l’expérience de formation et peut être utile dans le calcul distribué, comme illustré dans le diagramme suivant.

Les clients de SageMaker peuvent également accéder à des images Docker intégrées avec divers frameworks d’apprentissage en profondeur préinstallés et les packages Linux, NCCL et Python nécessaires à la formation du modèle. Les scientifiques des données ou ingénieurs en ML qui souhaitent effectuer une formation de modèle peuvent le faire sans la contrainte de configurer l’infrastructure de formation ou de gérer Docker et la compatibilité de différentes bibliothèques.

Lors d’un atelier d’un jour, nous avons pu mettre en place une configuration de formation distribuée basée sur SageMaker dans le compte AWS de KT, accélérer les scripts de formation de KT en utilisant la bibliothèque SageMaker Distributed Data Parallel (DDP), et même tester un travail de formation en utilisant deux instances ml.p4d.24xlarge. Dans cette section, nous décrivons l’expérience de KT en travaillant avec l’équipe AWS et en utilisant SageMaker pour développer leur modèle.

Dans la preuve de concept, nous voulions accélérer un travail de formation en utilisant la bibliothèque SageMaker DDP, qui est optimisée pour l’infrastructure AWS lors de la formation distribuée. Pour passer de PyTorch DDP à SageMaker DDP, il vous suffit de déclarer le package torch_smddp et de changer le backend en smddp, comme indiqué dans le code suivant:

import smdistributed.dataparallel.torch.torch_smddpdist.init_process_group(backend='smddp',rank=args.rank,world_size=args.world_size)

Pour en savoir plus sur la bibliothèque SageMaker DDP, référez-vous à la bibliothèque de parallélisme des données de SageMaker.

Analyser les causes de la lenteur de la vitesse de formation avec le débogueur et le profileur de SageMaker

La première étape de l’optimisation et de l’accélération d’une charge de travail de formation consiste à comprendre et à diagnostiquer les goulots d’étranglement. Pour le travail de formation de KT, nous avons mesuré le temps de formation par itération du chargeur de données, de la passe avant et de la passe arrière:

Temps d’une itération – chargeur de données : 0,00053 seconde, passe avant : 7,77474 secondes, passe arrière : 1,58002 seconde

Temps de deux itérations – chargeur de données : 0,00063 seconde, passe avant : 0,67429 seconde, passe arrière : 24,74539 secondes

Temps de trois itérations – chargeur de données : 0,00061 seconde, passe avant : 0,90976 seconde, passe arrière : 8,31253 secondes

Temps de quatre itérations – chargeur de données : 0,00060 seconde, passe avant : 0,60958 seconde, passe arrière : 30,93830 secondes

Temps de cinq itérations – chargeur de données : 0,00080 seconde, passe avant : 0,83237 seconde, passe arrière : 8,41030 secondes

Temps de six itérations – chargeur de données : 0,00067 seconde, passe avant : 0,75715 seconde, passe arrière : 29,88415 secondes

En examinant le temps affiché pour chaque itération, nous avons constaté que le temps d’exécution de la passe arrière variait considérablement d’une itération à l’autre. Cette variation est inhabituelle et peut avoir un impact sur le temps total de formation. Pour trouver la cause de cette vitesse de formation incohérente, nous avons d’abord essayé d’identifier les goulots d’étranglement des ressources en utilisant le Moniteur système (interface utilisateur SageMaker Debugger), qui vous permet de déboguer les travaux de formation sur SageMaker Training et de visualiser l’état des ressources telles que le CPU, le GPU, le réseau et l’E/S de la plateforme de formation gérée pendant un certain nombre de secondes.

L’interface utilisateur SageMaker Debugger fournit des données détaillées essentielles qui peuvent aider à identifier et diagnostiquer les goulots d’étranglement d’un travail de formation. Plus précisément, le graphique linéaire d’utilisation du CPU et les tableaux de chaleur d’utilisation du CPU/GPU par instance ont retenu notre attention.

Dans le graphique linéaire d’utilisation du CPU, nous avons constaté que certains CPU étaient utilisés à 100%.

Dans la carte thermique (où les couleurs plus sombres indiquent une utilisation plus élevée), nous avons constaté que quelques cœurs de CPU avaient une utilisation élevée tout au long de la formation, tandis que l’utilisation du GPU n’était pas constamment élevée au fil du temps.

À partir de là, nous avons commencé à suspecter qu’une des raisons de la lenteur de la vitesse d’entraînement était un goulôt d’étranglement du CPU. Nous avons revu le code du script d’entraînement pour voir si quelque chose causait ce goulôt d’étranglement du CPU. La partie la plus suspecte était la grande valeur de num_workers dans le chargeur de données, alors nous avons changé cette valeur à 0 ou 1 pour réduire l’utilisation du CPU. Nous avons ensuite relancé le travail d’entraînement et vérifié les résultats.

Les captures d’écran suivantes montrent le graphique de l’utilisation du CPU, l’utilisation du GPU et la carte thermique après atténuation du goulôt d’étranglement du CPU.

En modifiant simplement num_workers, nous avons constaté une diminution significative de l’utilisation du CPU et une augmentation générale de l’utilisation du GPU. Il s’agissait d’un changement important qui a considérablement amélioré la vitesse d’entraînement. Cependant, nous voulions voir où nous pourrions optimiser l’utilisation du GPU. Pour cela, nous avons utilisé SageMaker Profiler.

SageMaker Profiler aide à identifier des indices d’optimisation en fournissant une visibilité sur l’utilisation par opérations, notamment en suivant les métriques d’utilisation du GPU et du CPU ainsi que la consommation du kernel du GPU/CPU dans les scripts d’entraînement. Il aide les utilisateurs à comprendre quelles opérations consomment des ressources. Tout d’abord, pour utiliser SageMaker Profiler, vous devez ajouter ProfilerConfig à la fonction qui invoque le travail d’entraînement en utilisant le SDK SageMaker, comme indiqué dans le code suivant :

from sagemaker import ProfilerConfig, Profilerfrom sagemaker.debugger import (ProfilerRule, rule_configs)rules=[ProfilerRule.sagemaker(rule_configs.ProfilerReport())]profiler_config = ProfilerConfig(profile_params = Profiler(cpu_profiling_duration=3600))from sagemaker.pytorch import PyTorchregion_name = 'us-west-2'image_uri=f'763104351884.dkr.ecr.{region_name}.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker'estimator = PyTorch(entry_point='train.py',source_dir='src',role=role,image_uri=image_uri,instance_count=4,instance_type='ml.p4d.24xlarge',distribution={'smdistributed': {'dataparallel': {'enabled': True}}},profiler_config=profiler_config,hyperparameters=hyperparameters,sagemaker_session=sagemaker_session,)

Dans le SDK Python de SageMaker, vous avez la flexibilité d’ajouter les fonctions annotate pour SageMaker Profiler afin de sélectionner le code ou les étapes du script d’entraînement qui nécessitent un profilage. Voici un exemple de code à déclarer pour SageMaker Profiler dans les scripts d’entraînement :

import smppySMProf = smppy.SMProfiler.instance()config = smppy.Config()config.profiler = {"EnableCuda": "1",}SMProf.configure(config)SMProf.start_profiling()…with smppy.annotate("Forward"):student_out = student_model(inp)with smppy.annotate("Backward"):loss.backward()…SMProf.stop_profiling()

Après avoir ajouté le code précédent, si vous exécutez un travail d’entraînement en utilisant les scripts d’entraînement, vous pouvez obtenir des informations sur les opérations consommées par le kernel du GPU (comme indiqué dans la figure suivante) après que l’entraînement s’est déroulé pendant une certaine période. Dans le cas des scripts d’entraînement de KT, nous l’avons exécuté pendant une époque et avons obtenu les résultats suivants.

Lorsque nous avons vérifié les cinq premiers temps de consommation des opérations par les noyaux GPU parmi les résultats de SageMaker Profiler, nous avons constaté que pour le script d’entraînement KT, le temps le plus long est consommé par l’opération de produit de matrices, qui est une opération de multiplication générale de matrices (GEMM) sur les GPU. Avec cette information importante provenant de SageMaker Profiler, nous avons commencé à enquêter sur les moyens d’accélérer ces opérations et d’améliorer l’utilisation des GPU.

Accélérer le temps d’entraînement

Nous avons examiné différentes façons de réduire le temps de calcul de la multiplication de matrices et appliqué deux fonctions de PyTorch.

Partager les états de l’optimiseur avec ZeroRedundancyOptimizer

Si vous regardez le zero Redundancy Optimizer (ZeRO), la technique DeepSpeed/ZeRO permet de former efficacement un grand modèle avec une meilleure vitesse de formation en éliminant les redondances dans la mémoire utilisée par le modèle. Le ZeroRedundancyOptimizer de PyTorch utilise la technique de partage de l’état de l’optimiseur pour réduire l’utilisation de mémoire par processus dans Distributed Data Parallel (DDP). DDP utilise des gradients synchronisés dans la passe en arrière de sorte que toutes les répliques de l’optimiseur itèrent sur les mêmes paramètres et valeurs de gradient, mais au lieu d’avoir tous les paramètres du modèle, chaque état de l’optimiseur est maintenu en le partageant uniquement pour les différents processus DDP afin de réduire l’utilisation de mémoire.

Pour l’utiliser, vous pouvez laisser votre Optimiseur existant dans optimizer_class et déclarer un ZeroRedundancyOptimizer avec le reste des paramètres du modèle et le taux d’apprentissage comme paramètres.

student_optimizer = ZeroRedundancyOptimizer(student_model.parameters(),optimizer_class=torch.optim.AdamW,lr=initial_lr)

Précision mixte automatique

La précision mixte automatique (AMP) utilise le type de données torch.float32 pour certaines opérations et torch.bfloat16 ou torch.float16 pour d’autres, pour la commodité de calcul rapide et une utilisation réduite de la mémoire. En particulier, les modèles d’apprentissage profond sont généralement plus sensibles aux bits d’exposant qu’aux bits de fraction dans leurs calculs, torch.bfloat16 est équivalent aux bits d’exposant de torch.float32, ce qui leur permet d’apprendre rapidement avec une perte minimale. torch.bfloat16 ne fonctionne que sur les instances avec une architecture NVIDIA A100 (Ampere) ou supérieure, comme ml.p4d.24xlarge, ml.p4de.24xlarge et ml.p5.48xlarge.

Pour appliquer l’AMP, vous pouvez déclarer torch.cuda.amp.autocast dans les scripts d’entraînement comme indiqué dans le code ci-dessus et déclarer dtype comme torch.bfloat16.

with torch.cuda.amp.autocast(dtype="torch.bfloat16"):teacher = teacher_model(input_data)student = student_model(input_data)loss = loss(teacher, student, target)loss.requires_grad_(True)loss.backward()student_optimizer.step()student_optimizer.zero_grad(set_to_none=True)

Résultats dans SageMaker Profiler

Après avoir appliqué les deux fonctions aux scripts d’entraînement et avoir exécuté un travail d’entraînement pour un seul cycle, nous avons vérifié les cinq premiers temps de consommation des opérations pour le noyau GPU dans SageMaker Profiler. La figure suivante montre nos résultats.

Nous pouvons voir que l’opération GEMM, qui était en tête de liste avant d’appliquer les deux fonctions Torch, a disparu des cinq premières opérations, remplacée par l’opération ReduceScatter, qui se produit généralement dans l’entraînement distribué.

Résultats de vitesse d’entraînement du modèle KT distillé

Nous avons augmenté la taille du lot d’entraînement de 128 de plus pour tenir compte des économies de mémoire résultant de l’application des deux fonctions Torch, ce qui donne une taille de lot finale de 1152 au lieu de 1024. L’entraînement du modèle étudiant final a pu exécuter 210 époques par jour; le temps d’entraînement et l’accélération entre l’environnement d’entraînement interne de KT et SageMaker sont résumés dans le tableau suivant.

Environnement d’entraînement	Spécification du GPU d’entraînement	Nombre de GPU	Temps d’entraînement (heures)	Époque	Heures par époque	Taux de réduction
Environnement d’entraînement interne de KT	A100 (80GB)	2	960	300	3.20	29
Amazon SageMaker	A100 (40GB)	32	24	210	0.11	1

La scalabilité d’AWS nous a permis de terminer le travail d’entraînement 29 fois plus rapidement qu’auparavant en utilisant 32 GPU au lieu de 2 sur site. Par conséquent, l’utilisation de plus de GPU sur SageMaker aurait considérablement réduit le temps d’entraînement sans différence de coût d’entraînement global.

Conclusion

Park Sang-min (Chef d’équipe de la technologie Vision AI Serving) du laboratoire AI2XL du Centre de technologie de convergence de KT a commenté la collaboration avec AWS pour développer le modèle AI Food Tag:

“Récemment, avec la multiplication des modèles basés sur des transformateurs dans le domaine de la vision, les paramètres du modèle et la mémoire GPU requise augmentent. Nous utilisons une technologie légère pour résoudre ce problème, et cela prend beaucoup de temps, environ un mois pour apprendre une fois. Grâce à ce PoC avec AWS, nous avons pu identifier les goulots d’étranglement des ressources avec l’aide de SageMaker Profiler et Debugger, les résoudre, puis utiliser la bibliothèque de parallélisme de données de SageMaker pour terminer la formation en environ un jour avec du code de modèle optimisé sur quatre instances ml.p4d.24xlarge.”

SageMaker a permis à l’équipe de Sang-min de gagner des semaines de temps dans la formation et le développement du modèle.

Sur la base de cette collaboration sur le modèle de vision, AWS et l’équipe de SageMaker continueront à collaborer avec KT sur divers projets de recherche en IA/ML afin d’améliorer le développement de modèles et la productivité des services en appliquant les fonctionnalités de SageMaker.

Pour en savoir plus sur les fonctionnalités connexes de SageMaker, consultez les liens suivants:

• Former des modèles d’apprentissage automatique
• Bibliothèque de parallélisme de données de SageMaker
• Utiliser Amazon SageMaker Debugger pour déboguer et améliorer les performances du modèle
• Utiliser Amazon SageMaker Profiler pour profiler les activités sur les ressources de calcul AWS

We will continue to update IPGirl; if you have any questions or suggestions, please contact us!