7 décembre 2023
Par Mark Coastworth et Matthew Choi
Le monde de l'apprentissage automatique est témoin de la montée en puissance de réseaux neuronaux gigantesques comportant des milliards de paramètres. Ces grands modèles de langage (LLM) ont démontré des capacités incroyables, principalement en raison de leurs capacités de généralisation et d'apprentissage en contexte. Mais cette croissance massive de la taille des modèles s'accompagne d'un défi de taille : les exigences matérielles accrues pour leur formation et leur déploiement nécessitent souvent une infrastructure distribuée, répartissant le modèle sur plusieurs unités de traitement graphique (GPU) ou même plusieurs nœuds.
Bien qu'il existe de nombreux outils pour la parallélisation des modèles et l'apprentissage distribué, les interactions plus profondes avec ces modèles, telles que la récupération d'informations intermédiaires ou l'édition, nécessitent une connaissance approfondie de l'informatique distribuée. Cela a été un obstacle pour de nombreux chercheurs en apprentissage automatique ayant une connaissance limitée de l'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 la compréhension des raisons d'un résultat donné d'une manière facilement interprétable par les humains.
Pour résoudre ce problème, les membres de l'équipe AI Engineering de Vector ont développé FlexModel, un logiciel conçu pour fournir une interface conviviale permettant d'interagir avec des modèles à grande échelle répartis sur des configurations multi-GPU et multi-nœuds.
Présenté dans "FlexModel : A Framework for Interpretability of Distributed Large Language Models", qui a été sélectionné comme article phare à NeurIPS 2023, Flexmodel accomplit ceci en fournissant une interface commune pour envelopper de grands modèles indépendamment de la façon dont ils ont été distribués (Accelerate, FSDP, DeepSpeed, etc.). Ensuite, il introduit le concept de HookFunctions qui permet aux utilisateurs d'interagir avec les internes des modèles distribués, à la fois pendant les passes avant et arrière. Il implémente ces mécanismes via une API simple qui a été publiée sous la forme d'une bibliothèque Python appelée FlexModel. En implémentant cette bibliothèque dans leurs projets, les chercheurs peuvent rapidement et facilement obtenir de riches informations sur les raisons pour lesquelles un modèle se comporte d'une certaine manière.
La bibliothèque FlexModel fournit une nouvelle classe comme interface principale pour les interactions avec l'utilisateur. Cette classe FlexModel hérite de la classe PyTorch nn.Module couramment utilisée, ce qui permet aux développeurs d'interagir facilement avec le modèle enveloppé via l'API nn.Module sans modifier le 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 une collection de HookFunctions : une fonction définie par l'utilisateur pour effectuer des opérations fines à chaque couche individuelle d'un réseau neuronal. Le cas d'utilisation le plus courant est la récupération d'activation, c'est-à-dire l'extraction d'informations intermédiaires d'un modèle afin de comprendre comment il aboutit à une décision de sortie. Un autre cas d'utilisation consiste à modifier ces informations intermédiaires, afin de voir comment différents états internes peuvent conduire à différentes sorties.
FlexModel a deux objectifs majeurs en matière de conception. Il doit être intuitif : l'application du wrapper FlexModel à un PyTorch nn.Module doit simplement ajouter des fonctionnalités d'inspection du modèle au modèle cible. Le déballage du modèle doit produire le modèle original sans effets secondaires. La fonction d'édition de la HookFunction doit permettre d'exécuter un code arbitraire sur les activations. FlexModel est également conçu pour être évolutif : FlexModel est indépendant du nombre de GPU ou de nœuds GPU, de l'architecture du modèle (par exemple LLaMA, Falcon, GPT), de la taille du modèle et de la stratégie de distribution (par exemple DP, FSDP, TP, PP) ou de leur 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 des modèles à appareil unique. Cela permet aux chercheurs qui ne sont pas nécessairement des experts en informatique distribuée d'interagir avec des modèles distribués et de les modifier sans avoir à se plonger dans les complexités des systèmes distribués.
Les préoccupations concernant les biais et l'équité dans l'IA ayant pris de l'importance, l'interprétabilité peut aider à détecter, comprendre et atténuer les biais cachés dans les décisions des modèles. L'élucidation de la manière dont ces modèles parviennent à prendre des décisions, de la manière dont ils ont appris des comportements spécifiques et la compréhension de leurs mécanismes internes peuvent nous aider à construire des systèmes d'IA plus robustes, plus fiables et plus efficaces.
Déjà, de nombreux secteurs tels que la médecine, la finance et le système juridique exigent que les modèles qui prennent des décisions ayant un impact sur les êtres humains soient interprétables, ce qui garantit la transparence des décisions. Grâce à des outils tels que FlexModel, les chercheurs peuvent désormais s'engager dans la recherche 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 sur l'IA inclusive et accessible à tous. En abaissant les barrières à la recherche sur l'interprétabilité dans les LLM, FlexModel nous rapproche de l'objectif de rendre l'apprentissage automatique de pointe plus accessible, plus interprétable, plus sûr et plus fiable.
