mamba

Mamba es una base de código de investigación de código abierto del grupo state-spaces que implementa una familia de arquitecturas de modelo de espacio de estados (SSM) para el modelado de secuencias. Acompaña a los artículos "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", "Transformers are SSMs" (Mamba-2) y "Mamba-3: Improved Sequence Modeling using State Space Principles". El proyecto se posiciona como una alternativa subcuadrática a los Transformers, particularmente para datos densos en información como el modelado del lenguaje, y se basa en trabajos anteriores de espacio de estados estructurados (S4) con una implementación consciente del hardware en el espíritu de FlashAttention. Lo que proporciona el repositorio: - Capa SSM selectiva: la operación principal de escaneo selectivo, con origen en ops/selective_scan_interface.py. - Bloque Mamba: el módulo de arquitectura principal que envuelve el SSM selectivo (modules/mamba_simple.py), utilizable como un módulo PyTorch independiente con parámetros configurables d_model, d_state, d_conv y expand. - Bloque Mamba-2: implementado en modules/mamba2.py con una variante más simple en modules/mamba2_simple.py, además de un módulo de referencia SSD (structured state space duality) mínimo en modules/ssd_minimal.py. - Bloque Mamba-3: implementado en modules/mamba3.py, exponiendo opciones como d_state, headdim, modo MIMO con mimo_rank, chunk_size, una norma de proyección de salida opcional y dtype bfloat16. - Ejemplo de modelo de lenguaje: una columna vertebral completa más una cabeza LM en models/mixer_seq_simple.py, utilizada por los scripts de generación. Opciones de instalación y compilación: El paquete se instala a través de pip (mamba-ssm) y requiere Linux, Python 3.10+ y PyTorch 1.12+. Por defecto, la instalación no compila la extensión selective_scan_cuda y no descarga ruedas CUDA cacheadas. Las opciones adicionales y las banderas del entorno cambian este comportamiento: la opción causal-conv1d agrega esa dependencia; MAMBA_FORCE_BUILD=TRUE fuerza una compilación local; MAMBA_KEEP_CUDA_BUILD=TRUE opta por la extensión de escaneo selectivo CUDA, intentando primero una rueda CUDA/HIP precompilada coincidente y compilando localmente si no hay ninguna disponible. La combinación de ambas banderas fuerza la compilación local de CUDA. El README señala que --no-build-isolation es necesario para las compilaciones de CUDA, de modo que pip utilice PyTorch con capacidad CUDA existente. Las compilaciones de CUDA y la ejecución de GPU requieren además una GPU NVIDIA y CUDA 11.6+. Para tarjetas AMD, el README documenta un paso de parche ROCm 6.0 (no necesario a partir de ROCm 6.1). Modelos preentrenados: Los pesos se publican en Hugging Face bajo la organización state-spaces, incluyendo 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 y mamba2attn-2.7b, entrenados en 300B tokens en The Pile, además de mamba-2.8b-slimpj entrenado en 600B tokens en SlimPajama. El README describe estos como modelos base sin ajuste de instrucciones u otras modificaciones posteriores, y afirma que se espera que el rendimiento sea comparable o mejor que otras arquitecturas entrenadas con datos similares, pero no que iguale a modelos más grandes o ajustados. Las dimensiones del modelo siguen el escalado al estilo GPT-3 (por ejemplo, 130M de parámetros con 24 capas y 768 de dimensión de modelo, hasta 2.8B con 64 capas y 2560 de dimensión). Herramientas de evaluación e inferencia: La evaluación de cero disparos se realiza a través de la biblioteca lm-evaluation-harness (pip install lm-eval==0.4.2), con comandos de ejemplo que cubren tareas como lambada_openai, hellaswag, piqa, arc_easy, arc_challenge, winogrande, openbookqa, boolq, race, truthfulqa_mc2 y mmlu. El README señala que los resultados por tarea pueden diferir de los valores reportados en 0.1-0.3 debido al ruido de evaluación. Un script de benchmark de generación (benchmarks/benchmark_generation_mamba_simple.py) carga automáticamente un modelo del Hugging Face Hub, genera completaciones para una indicación del usuario y mide la velocidad de inferencia, con opciones para top-p, min-p, top-k, temperatura, penalización por repetición y tamaño del lote. Notas de solución de problemas: El README destaca dos problemas prácticos. Primero, precisión: los modelos se entrenaron con PyTorch AMP, que mantiene los parámetros en float32 y los convierte a media precisión según sea necesario; dado que los SSM son sensibles a la dinámica recurrente, los marcos que almacenan parámetros en float16 (como DeepSpeed) pueden causar inestabilidad, y se sugiere el almacenamiento de parámetros en fp32 como primera solución. Segundo, inicialización: algunos componentes heredan inicializaciones al estilo S4 (por ejemplo, el rango objetivo del parámetro delta a través del sesgo de proyección lineal), y los marcos con ganchos de post-inicialización que ponen a cero los sesgos de nn.Linear pueden necesitar lógica personalizada para preservarlos. Licencia y citación: el README proporciona entradas BibTeX para los artículos de Mamba, Mamba-2 y Mamba-3 y pide a los usuarios que los citen al usar la base de código.