Présentation de FlexModel : cadre révolutionnaire pour dévoiler les secrets des grands modèles d’IA générative
7 décembre 2023
7 décembre 2023
Par Mark Coastworth et Matthew Choi
Le monde de l’apprentissage automatique est témoin de l’essor de réseaux neuronaux gigantesques avec des milliards de paramètres. Ces grands modèles de langage (LLM) ont démontré des capacités incroyables, principalement grâce à leur généralisation et leur capacité d’apprentissage en contexte. Mais cette croissance massive de la taille du modèle apporte un défi important : les exigences matérielles accrues pour leur formation et leur déploiement nécessitent souvent une infrastructure distribuée, divisant le modèle entre plusieurs unités de traitement graphique (GPU) ou même plusieurs nœuds.
Bien que de nombreux outils existent pour la parallélisation des modèles et l’entraînement distribué, des interactions plus profondes avec ces modèles, comme la récupération d’informations intermédiaires ou l’édition, nécessitent une solide maîtrise de l’informatique distribuée. Cela a été un obstacle pour de nombreux chercheurs en apprentissage automatique ayant des connaissances limitées en informatique distribuée. En conséquence, ces grands modèles fonctionnent généralement à l’intérieur d’une boîte noire, ce qui rend difficile de comprendre les raisons d’une sortie donnée d’une manière facilement interprétable pour les humains.
Pour résoudre ce problème, les membres de l’équipe d’ingénierie IA de Vector ont développé FlexModel, un logiciel conçu pour offrir une interface conviviale permettant d’interagir avec des modèles à grande échelle répartis sur des configurations multi-GPU et multi-nœuds.
Introduit dans « FlexModel : A Framework for Interpretability of Distributed Large Language Models », qui a été sélectionné comme article en vedette lors de NeurIPS 2023, Flexmodel y parvient en fournissant une interface commune pour envelopper de grands modèles, peu importe leur répartition (Accelerate, FSDP, DeepSpeed, etc.). Ensuite, il introduit le concept de HookFunctions, qui permet aux utilisateurs d’interagir avec les composants internes distribués du modèle, tant lors des passages vers l’avant que vers l’arrière. Il implémente ces mécanismes via une API simple qui a été publiée sous forme de bibliothèque Python appelée FlexModel. En intégrant cette bibliothèque dans leurs projets, les chercheurs peuvent rapidement et facilement obtenir des informations riches sur les raisons pour lesquelles un modèle se comporte d’une certaine façon.
La bibliothèque FlexModel offre une nouvelle classe comme interface principale pour les interactions utilisateurs. Cette classe FlexModel hérite du PyTorch nn, couramment utilisé. Module class, permettant aux développeurs d’interagir facilement avec le modèle enveloppé via le nn. API du module sans aucun changement de code.
Un exemple simple d’initialisation ressemble à ceci :
model = AutoModelForCausalLM.from_pretrained("model-name")
model = accelerator.prepare(model)
output_dict: Dist[str, Tensor] = {}
model = FlexModel(model, output_dict, data_parallel_size=accelerator.num_processes)
Une fois qu’un FlexModel a été instancié, les utilisateurs peuvent définir un ensemble de HookFunctions : une fonction définie par l’utilisateur pour effectuer des opérations détaillées à chaque couche individuelle d’un réseau de neurones. Le cas d’usage le plus courant ici est d’effectuer la récupération d’activation, c’est-à-dire saisir des informations intermédiaires d’un modèle afin de comprendre comment il prend une décision de sortie. Un autre cas d’utilisation consiste à modifier cette information intermédiaire, pour voir comment différents états internes peuvent mener à différentes sorties.
FlexModel a deux grands objectifs de conception. Ça devrait être intuitif : appliquer le wrapper FlexModel à un nn PyTorch. Le module devrait simplement ajouter des fonctionnalités pour l’inspection du modèle au modèle cible. Déplier le modèle devrait produire le modèle original sans effets secondaires. La fonction d’édition de la HookFunction devrait permettre d’exécuter du code arbitraire sur les activations. Il est aussi conçu pour être évolutif : FlexModel est agnostique quant au nombre de GPU ou de nœuds GPU, à l’architecture du modèle (par exemple LLaMA, Falcon, GPT), à la taille du modèle, ainsi qu’à la stratégie de distribution (par exemple DP, FSDP, TP, PP) ou à sa composition.
FlexModel promet de démocratiser les interactions entre les modèles et de combler le fossé entre les paradigmes des modèles distribués et à dispositif unique. Cela permet aux chercheurs qui ne sont pas forcément experts en informatique distribuée d’interagir avec et de modifier les modèles distribués sans plonger en profondeur dans la complexité des systèmes distribués.
À mesure que les préoccupations concernant les biais et l’équité en IA gagnent en importance, l’interprétabilité peut aider à détecter, comprendre et atténuer les biais cachés dans les décisions du modèle. Démêler comment ces modèles prennent des décisions, comment ils ont appris des comportements spécifiques et comprendre leurs mécanismes internes peut nous donner des perspectives pour bâtir des systèmes d’IA plus robustes, fiables et efficaces.
Déjà, de nombreux secteurs comme la médecine, la finance et le système juridique réglementent que les modèles qui prennent des décisions touchant les humains doivent être interprétables. Cela garantit que les décisions sont prises de manière transparente. Avec des outils comme FlexModel, les chercheurs peuvent désormais mener des recherches sur l’interprétabilité sans être accablés par les complexités techniques de l’informatique distribuée.
Des outils comme FlexModel soulignent l’importance de rendre la recherche avancée en IA inclusive et universellement accessible. En abaissant les barrières à la recherche sur l’interprétabilité dans les LLM, FlexModel nous rapproche d’un pas de plus pour rendre l’apprentissage automatique à la fine pointe de la technologie plus accessible, interprétable, sécuritaire et digne de confiance.