- mamba-minimal est un projet qui propose une implémentation simple et minimale de Mamba dans un seul fichier PyTorch
- Il vise à produire les mêmes résultats numériques que l’implémentation officielle pour les forward/backward pass
- Le code est simplifié et structuré sous une forme commentée pour faciliter la lecture
- Il n’inclut pas les optimisations clés de l’implémentation officielle, n’offre donc pas de vitesse, et ne comprend pas non plus d’initialisation appropriée des paramètres
- La démo exécute un exemple de complétion de prompt avec
state-spaces/mamba-370met le tokenizerEleutherAI/gpt-neox-20b
Présentation du projet
- mamba-minimal est une implémentation minimale et simple de Mamba dans un seul fichier PyTorch
- Son objectif est de montrer le même comportement que l’implémentation officielle dans un code plus lisible
- Principales caractéristiques :
- Résultats numériques équivalents à l’implémentation officielle pour le forward pass et le backward pass
- Code simplifié
- Implémentation lisible et commentée
Ce qui n’est pas inclus
- La vitesse n’est pas un objectif
- L’implémentation officielle est fortement optimisée
- Ces optimisations font partie des contributions clés de l’article Mamba
- Cette implémentation garde la plupart du code simple afin de privilégier la lisibilité
- L’initialisation appropriée des paramètres n’est pas incluse
- Elle est présentée comme un élément qui pourrait être ajouté sans sacrifier la lisibilité
Exemple d’utilisation de la démo
- Un exemple de complétion de prompt est disponible dans
demo.ipynb - L’exemple utilise
model.MambaetAutoTokenizerde Hugging Facetransformers - Modèle et tokenizer utilisés :
state-spaces/mamba-370mEleutherAI/gpt-neox-20b
- Le prompt d’exemple est
Mamba is the, et le résultat généré contient une phrase décrivant le mamba comme un serpent venimeux
Références
- L’architecture Mamba est présentée dans Mamba: Linear-Time Sequence Modeling with Selective State Spaces
- Les auteurs de l’article sont Albert Gu et Tri Dao
- L’implémentation officielle se trouve dans le dépôt
state-spaces/mamba
1 commentaires
Avis sur Hacker News
Il y a quelque temps, avec un collègue, nous avons créé une bibliothèque qui factorise la majeure partie du code de modèle partagé ; avec elle, on peut implémenter beaucoup de modèles en environ 100 lignes, hors imports Python et commentaires.
BERT : https://github.com/explosion/curated-transformers/blob/main/...
Llama 1/2 : https://github.com/explosion/curated-transformers/blob/main/...
MPT : https://github.com/explosion/curated-transformers/blob/main/...
Elle prend aussi en charge des fonctionnalités comme TorchScript JIT et PyTorch flash attention.
xformers traite aussi un problème similaire, mais se concentre davantage sur la fourniture de modules Transformer performants avec Triton. Cela dit, il n’était pas facile d’en utiliser seulement certains composants, et je tombais sans cesse sur des erreurs à l’exécution, donc je l’ai mis de côté pour l’instant. Je travaille sur quelque chose basé sur l’architecture BERT, je vais donc m’en inspirer.
Je compte l’essayer dans mon prochain projet.
Le code Mamba d’origine contient beaucoup d’optimisations de vitesse et d’autres éléments, ce qui le rend difficile à comprendre directement ; cette implémentation devrait être utile pour apprendre.
Quand on fait de l’inférence token par token, tout devient beaucoup plus simple. J’ai aussi une implémentation maison de l’inférence Mamba : https://github.com/rbitr/llm.f90/tree/master/ssm
Je sais que c’est une base éprouvée pour le calcul scientifique, souvent enveloppée par des bibliothèques comme PyTorch ou Numpy, mais ce n’est pas un langage très populaire aujourd’hui. Je serais curieux de connaître la raison de ce choix.
Il y a certains points sur Mamba que j’aimerais voir expliqués d’une façon compréhensible même pour quelqu’un qui n’est pas chercheur en machine learning
L’intuition centrale de Mamba consiste à résoudre un vieux problème des modèles à espace d’états. Ces modèles sont bons pour compresser le contexte d’entrée, mais, dans le processus qui compresse l’entrée en un état caché, ils effacent des informations nécessaires pour exploiter efficacement le contexte comme le font les Transformers
La solution consiste à créer ce que l’article appelle un mécanisme de sélection. Ce mécanisme dépend de l’entrée, de sorte que le modèle peut ajuster la sortie à chaque étape quand l’entrée change. Pour cela, certains paramètres de l’espace d’états ne sont plus invariants à l’entrée mais dépendants de celle-ci, et l’on ajoute notamment des couches linéaires pour projeter l’entrée de chaque instant vers les variables de l’espace d’états
Mais rendre les variables de l’espace d’états dépendantes de l’entrée introduit un surcoût de calcul. Ils le résolvent avec un algorithme conscient du matériel qui exploite au maximum la structure mémoire des GPU modernes, afin d’éviter autant que possible de déplacer des données vers et depuis la HBM
Tri Dao est le créateur de Flash Attention, qui était déjà une manière d’utiliser le matériel plus efficacement avec les Transformers. C’est vraiment son domaine d’expertise
Mamba et Based, sortis juste avant NeurIPS 2023, incluaient le rappel associatif multi-requêtes (MQAR) ainsi qu’une dépendance aux données du gating/de la sélection inspirée de l’Attention multi-têtes. C’étaient les éléments clés qui manquaient à Hyena et aux architectures antérieures à espace d’états, et les nouveaux modèles deviennent aussi bons que l’Attention sur les tâches de rappel associatif, tout en montrant la possibilité d’être légèrement meilleurs que l’Attention sur des tâches autres que la recherche
Bien sûr, le grand détail de Mamba est son implémentation CUDA efficace. Sans elle, cette architecture pourrait perdre de son intérêt sur les tâches pour lesquelles les Transformers sont déjà bien adaptés
Même sans trop se soucier de la longueur de contexte, cela ouvre beaucoup de nouveaux domaines. L’analyse de séquences d’ADN est une tâche linéaire avec de longues dépendances, et l’on peut aussi envisager les images, vidéos et informations de haute dimension comme des flux de tokens. Un peu comme le balayage des pixels sur les anciens moniteurs CRT
L’un des rêves initiaux de l’IA était l’évolution continue d’une trajectoire d’apprentissage unique pour un agent interagissant en permanence avec son environnement, et ces modèles à longueur de contexte infinie pourraient rendre ce rêve plus facile à atteindre
Cela dit, pour l’instant, les applications downstream de ces modèles à des tâches pratiques importantes sont globalement moins validées et moins réglées que les applications plus matures fondées sur l’Attention. L’analogie avec les anciens réseaux récurrents aide dans une certaine mesure, mais ces cinq dernières années, les gens se sont tellement spécialisés sur l’Attention et les Transformers que l’inertie côté Transformers est très forte
L’article original explique qu’une fois les paramètres transformés, le modèle peut être calculé de deux façons : comme une récurrence linéaire ou comme une convolution globale. En général, lors de l’entraînement, où l’on peut voir toute la séquence d’entrée à l’avance, on utilise le mode convolution, plus facile à paralléliser ; en inférence autorégressive, où l’on voit l’entrée pas à pas, on passe au mode récurrent efficace
L’entraînement est donc parallélisable, comme le mode de propagation avant parallèle de RetNet. L’inférence de base s’effectue en mode récurrent pour obtenir le contexte le plus long possible ; comme il n’y a pas de chunking, il est difficile de juger de la quantité de RAM et de VRAM consommée pendant l’inférence
Elle explique l’article tout en donnant beaucoup de contexte sur sa place dans la vue d’ensemble. C’est assez intéressant d’écouter le déroulé
https://youtu.be/ouF-H35atOY?si=y2Ckp9MCFd7ulLL3
Au lieu de l’Attention quadratique, qui calcule à quel point chaque token doit prêter attention à tous les autres tokens, on calcule d’une manière ou d’une autre un long noyau de convolution de la même longueur que l’entrée, puis on applique une conv1d
D’après ma compréhension limitée, cela a un certain rapport avec le fait d’appliquer une FFT, de faire une multiplication matricielle, puis de revenir avec une IFFT. Je sais que cela fonctionne, mais c’est lent. Il existe plusieurs façons de calculer une FFT, et l’une d’elles est la matrice papillon. Ce n’est sans doute qu’une approximation, mais elle semble assez bonne et très rapide et efficace sur le matériel actuel
La complexité quadratique sonne mal, mais en pratique, à cause des contraintes matérielles, les algorithmes sous-quadratiques sont souvent plus lents. Donc, même s’il y a beaucoup d’attentes autour des modèles à espace d’états, il n’est pas facile de dire que Llama est fini. On ne sait pas encore non plus si Mamba fonctionne bien en passant à l’échelle, et pour le savoir il faudra réellement dépenser des millions de dollars en entraînement. Je reste tout de même optimiste
Un autre modèle intéressant de la famille sous-quadratique est RWKV. Ça vaut le coup d’y jeter un œil, même si je crois qu’il a probablement déjà été abordé dans le podcast
J’ai appris tout ça en autodidacte et je n’ai fait que parcourir rapidement l’article il y a quelque temps, donc il peut y avoir beaucoup d’erreurs. Par ailleurs, l’Attention dispose généralement d’un cache KV, ce qui aide beaucoup les performances, mais je ne pense pas que Mamba puisse faire la même chose
La phrase « Mamba est le plus long serpent venimeux du monde, avec une longueur estimée à plus de 150 m » m’a fait rire
Cela dit, c’est vraiment excellent, et j’ai apprécié la référence à l’article arXiv : cela permet à quelqu’un comme moi, qui consomme ce genre d’articles plutôt que d’interpréter directement les papiers, de jeter un petit coup d’œil à l’intérieur
Ce serait amusant s’il fallait publier plus tard une correction pour cette phrase
Je m’attendais à ce que le cœur de l’algorithme soit un scan de préfixe parallèle. Il me semble que c’est un peu le point clé de Mamba
for i in range(l):x = deltaA[:, :, i] * x + deltaB_u[:, :, i]y = einsum(x, C[:, i, :], 'b d_in n , b n -> b d_in')ys.append(y)Question peut-être idiote, mais je me demande à quel point il est difficile d’entraîner le modèle Mamba disponible sur Hugging Face
Le plus gros modèle semble faire 2,8B ; pour l’entraîner sur un jeu de données comme The Pile, combien de GPU faudrait-il et combien de temps cela prendrait-il ?
L’inférence semble également s’exécuter 3 à 5 fois plus vite, en utilisant seulement la moitié de la RAM
Après un premier échec en essayant de décortiquer la version CUDA officielle, j’ai fini par laisser tomber, mais cette implémentation a l’air bien meilleure
Encore une implémentation PyTorch en un seul fichier, c’est vraiment excellent. J’espère que les projets précédents comme hlb-CIFAR10 et les projets associés, ainsi que des influences antérieures comme minGPT ou DawnBench, ont contribué, ne serait-ce qu’un peu, à pousser ce format simple en un seul fichier
Ce type de travail est important pour une recherche en machine learning efficace, et c’est peut-être l’une des choses les plus importantes que l’on puisse faire aujourd’hui pour ce domaine
La recherche progresse au rythme de l’innovation, l’innovation s’accélère avec l’inverse du temps d’exécution des expériences, et cela est clairement lié à la complexité de Kolmogorov du code, que ce soit pour la recherche ou pour du hacking rapide
On n’insistera jamais assez sur l’importance de ces outils pour la recherche, ni sur la vitesse à laquelle ils ont personnellement accéléré mon processus de découverte de connaissances. La capacité à esquisser rapidement des idées en quelques minutes et à recevoir immédiatement des résultats au rapport signal/bruit élevé est devenue indispensable à l’avancement de la recherche
Je pense que la distillation de connaissances et le MDL (https://en.wikipedia.org/wiki/Minimum_description_length) sont très importants pour revenir sur les ornements inutiles, le bazar et la compétition à faible valeur sur des sujets excessivement denses de type « ne pas se faire distancer », que le processus actuel de soumission et de relecture d’articles semble encourager
Ces derniers temps, pour éviter ce problème et aller vers une solution de passage à l’échelle un peu meilleure, j’ai commencé à diffuser mon code sous forme de « croquis de code » : de courts gists autonomes en un seul fichier. Cela réduit le temps de développement et permet de transmettre directement aux gens du code fonctionnel, brut et non poli, qui contient le concept. Jusqu’ici, cela semble plutôt bien fonctionner, et j’aimerais continuer
J’aimerais voir davantage de ce genre de code. Si des chercheurs entraînent des données à grande échelle, ils devraient aussi être efficaces en données dans leur manière de diffuser l’information
L’humanité n’a probablement jamais fait progresser aussi vite quelque chose d’une complexité aussi importante
Le seul endroit où l’on voit une vitesse comparable est peut-être SpaceX, qui a encore lancé cette année deux fusées à la pointe de la technologie. Je me demande ce que 2024 nous réserve
x_projn’a pas de biais, il semble possible de fusionner les poids de x_proj et dt_projS’il y a des exigences d’ajustement des poids, cela pourrait peut-être se faire simplement à l’exécution, et un seul kernel avec un biais finirait probablement par être plus rapide. Je n’en suis pas sûr
Je me demande s’il y a eu une discussion sur l’article original. J’ai dû la manquer, mais c’est assez intéressant
Je n’ai pas bien compris le passage disant qu’« en raison du manque d’implémentations efficaces, qui entraîne des dépassements mémoire ou des exigences de calcul irréalistes, les résultats complets à longueur de contexte 8k manquent pour les baselines RWKV et RetNet, d’anciens modèles récurrents puissants pouvant aussi être interprétés comme des SSM »
RetNet n’utilise pas beaucoup de mémoire, et avec une implémentation de la passe avant par chunks, l’utilisation de VRAM est limitée par la taille du chunk. C’est précisément l’essentiel pour tester la longueur de contexte
Je me demande si quelqu’un a testé le modèle Mamba original. Quelle serait sa vitesse d’entraînement par rapport à RetNet en mode passe avant parallèle ?
https://openreview.net/forum?id=AL1fq05o7H
Une implémentation qui réduit quelque chose de complexe à l’essentiel, c’est toujours appréciable