mamba

Mamba est une base de code de recherche open source du groupe state-spaces, implémentant une famille d'architectures de modèles à espace d'états (SSM) pour la modélisation de séquences. Elle accompagne les articles « Mamba: Linear-Time Sequence Modeling with Selective State Spaces », « Transformers are SSMs » (Mamba-2) et « Mamba-3: Improved Sequence Modeling using State Space Principles ». Le projet se positionne comme une alternative subquadratique aux Transformers, en particulier pour les données denses en information comme la modélisation du langage, et s'appuie sur les travaux antérieurs sur les espaces d'états structurés (S4) avec une implémentation matériellement consciente dans l'esprit de FlashAttention. Ce que fournit le dépôt : - Couche SSM sélective : l'opération de scan sélectif central, avec source dans ops/selective_scan_interface.py. - Bloc Mamba : le module d'architecture principal encapsulant le SSM sélectif (modules/mamba_simple.py), utilisable comme module PyTorch autonome avec paramètres configurables d_model, d_state, d_conv et expand. - Bloc Mamba-2 : implémenté dans modules/mamba2.py avec une variante plus simple dans modules/mamba2_simple.py, plus un module de référence SSD minimal (dualité espace d'états structuré) dans modules/ssd_minimal.py. - Bloc Mamba-3 : implémenté dans modules/mamba3.py, exposant des options telles que d_state, headdim, mode MIMO avec mimo_rank, chunk_size, une norme de projection de sortie optionnelle, et le dtype bfloat16. - Exemple de modèle de langage : un backbone complet plus une tête LM dans models/mixer_seq_simple.py, utilisé par les scripts de génération. Options d'installation et de compilation : Le paquet s'installe via pip (mamba-ssm) et nécessite Linux, Python 3.10+ et PyTorch 1.12+. Par défaut, l'installation ne compile pas l'extension selective_scan_cuda et ne récupère pas les wheels CUDA en cache. Des extras optionnels et des variables d'environnement modifient ce comportement : l'extra causal-conv1d ajoute cette dépendance ; MAMBA_FORCE_BUILD=TRUE force une compilation locale ; MAMBA_KEEP_CUDA_BUILD=TRUE opte pour l'extension CUDA de scan sélectif, en essayant d'abord une wheel CUDA/HIP précompilée correspondante et en compilant localement si aucune n'est disponible. La combinaison des deux drapeaux force la compilation CUDA locale. Le README note que --no-build-isolation est requis pour les compilations CUDA afin que pip utilise le PyTorch existant compatible CUDA. Les compilations CUDA et l'exécution GPU nécessitent en outre un GPU NVIDIA et CUDA 11.6+. Pour les cartes AMD, le README documente une étape de correctif ROCm 6.0 (non nécessaire à partir de ROCm 6.1). Modèles pré-entraînés : Les poids sont publiés sur Hugging Face sous l'organisation state-spaces, notamment mamba-130m, mamba-370m, mamba-790m, mamba-1.4b, mamba-2.8b, mamba2-130m, mamba2-370m, mamba2-780m, mamba2-1.3b, mamba2-2.7b, transformerpp-2.7b et mamba2attn-2.7b, entraînés sur 300B tokens du Pile, plus mamba-2.8b-slimpj entraîné sur 600B tokens de SlimPajama. Le README les décrit comme des modèles de base sans ajustement par instructions ni autre modification en aval, et indique que les performances devraient être comparables ou supérieures à celles d'autres architectures entraînées sur des données similaires, mais pas à la hauteur de modèles plus grands ou affinés. Les dimensions des modèles suivent un dimensionnement de type GPT-3 (par exemple 130M paramètres avec 24 couches et dimension de modèle 768, jusqu'à 2,8B avec 64 couches et dimension 2560). Outils d'évaluation et d'inférence : L'évaluation zéro-shot se fait via la bibliothèque lm-evaluation-harness (pip install lm-eval==0.4.2), avec des exemples de commandes couvrant des tâches telles que lambada_openai, hellaswag, piqa, arc_easy, arc_challenge, winogrande, openbookqa, boolq, race, truthfulqa_mc2 et mmlu. Le README note que les résultats par tâche peuvent différer des valeurs rapportées de 0,1 à 0,3 en raison du bruit d'évaluation. Un script de benchmark de génération (benchmarks/benchmark_generation_mamba_simple.py) charge automatiquement un modèle depuis le Hugging Face Hub, génère des complétions pour un prompt utilisateur et mesure la vitesse d'inférence, avec des options pour top-p, min-p, top-k, température, pénalité de répétition et taille de lot. Notes de dépannage : Le README met en évidence deux problèmes pratiques. Premièrement, la précision : les modèles ont été entraînés avec PyTorch AMP, qui conserve les paramètres en float32 et les convertit en demi-précision si nécessaire ; comme les SSM sont sensibles à la dynamique récurrente, les frameworks qui stockent les paramètres en float16 (comme DeepSpeed) peuvent provoquer une instabilité, et le stockage des paramètres en fp32 est suggéré comme premier remède. Deuxièmement, l'initialisation : certains composants héritent des initialisations de type S4 (par exemple, la plage ciblée du paramètre delta via le biais de projection linéaire), et les frameworks avec des hooks post-initialisation qui mettent à zéro les biais des nn.Linear peuvent nécessiter une logique personnalisée pour les préserver. Licence et citation : le README fournit des entrées BibTeX pour les articles Mamba, Mamba-2 et Mamba-3 et demande aux utilisateurs de les citer lors de l'utilisation de la base de code.