Implementación de Mamba en un solo archivo de PyTorch

mamba-minimal

• Explicación de mamba-minimal, una implementación simple de Mamba en un solo archivo escrita en PyTorch.
• Proporciona la misma salida numérica que la implementación oficial, tanto en la propagación hacia adelante como en la retropropagación.
• El código está simplificado, es fácil de leer y está comentado.
• No incluye funciones como las optimizaciones de velocidad de la implementación oficial.
• No incluye una inicialización de parámetros adecuada, pero se puede agregar sin afectar la legibilidad.

Demo

• Archivo demo.ipynb que muestra un ejemplo de finalización de prompt.
• Ofrece un ejemplo de generación de texto usando el modelo Mamba y AutoTokenizer.
• En el texto de ejemplo generado, Mamba es descrita como la serpiente venenosa más larga del mundo.

References

• La arquitectura Mamba fue presentada en el artículo "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" escrito por Albert Gu y Tri Dao.
• La implementación oficial se puede consultar en GitHub.

Opinión de GN⁺

• mamba-minimal es un proyecto creado para reducir la complejidad de las implementaciones existentes de Mamba y hacerlas comprensibles incluso para ingenieros de software principiantes.
• Este proyecto contribuye a mejorar la legibilidad y la comprensión del código en el campo del aprendizaje automático.
• A través de ejemplos de uso reales, muestra de forma sencilla cómo aprovechar el modelo Mamba, lo que puede convertirlo en un material muy interesante para quienes están aprendiendo.

[GN⁺]

Opiniones de Hacker News

• Compartiendo una biblioteca

  • Creó una biblioteca junto con un colega para extraer la mayor parte del código de modelo en común. La mayoría de los modelos pueden implementarse en unas 100 líneas, excluyendo el proceso de import de Python y los comentarios.
  • Menciona como ejemplos los modelos BERT, Llama 1/2 y MPT, y que soporta TorchScript JIT, PyTorch flash attention, entre otros.

• Compartiendo una implementación de inferencia de Mamba

  • El código de Mamba tiene muchas partes, como optimizaciones de velocidad, que dificultan aplicarlo directamente al entrenamiento, por lo que comparte su propia implementación de inferencia de Mamba que puede ayudar con eso.

• Solicitud de una explicación para no especialistas sobre Mamba

  • Pide una explicación para no especialistas sobre la idea principal de Mamba y los modelos de espacio de estados, por qué Mamba ha tenido éxito y qué otras ventajas tiene Mamba incluso si no te interesa la longitud de contexto.

• Expectativa clave sobre el algoritmo

  • Esperaba que el núcleo del algoritmo de Mamba fuera un escaneo de prefijos paralelo, pero el código real está implementado de otra manera.

• Broma ingeniosa sobre Mamba

  • Le causó gracia la broma de que Mamba es la serpiente venenosa más larga del mundo. Le gustó poder echar un vistazo a lo que hay detrás de la investigación consultando el artículo de arXiv.

• Pregunta sobre la dificultad de entrenar modelos Mamba

  • Pregunta qué tan difícil es entrenar un modelo Mamba disponible en Huggingface y cuántas GPU y cuánto tiempo se necesitan para entrenar el modelo más grande.

• Compartiendo un intento de interpretar la versión oficial en CUDA

  • Intentó interpretar la versión oficial en CUDA, pero fracasó y no volvió a intentarlo. En comparación, considera que esta nueva implementación es mucho mejor.

• Admiración por una implementación de PyTorch en un solo archivo

  • Expresa admiración por otra implementación de PyTorch en un solo archivo y cree que este tipo de trabajo es importante para una investigación de ML eficiente. La velocidad del progreso en la investigación está relacionada con el inverso del tiempo de ejecución de los experimentos, lo que a su vez está relacionado con la complejidad de Kolmogorov del código. Destaca lo importantes que son las herramientas de investigación para acelerar el proceso de descubrimiento de conocimiento.

• Solicitud de discusión del artículo original

  • Pregunta si hubo una discusión sobre el artículo original y señala que modelos anteriores como RWKV y RetNet carecían de implementaciones eficientes, por lo que faltaban resultados completos para longitudes de contexto de 8k debido a problemas de memoria o a requisitos de cómputo poco realistas.

• Elogio por simplificar el contenido clave

  • Le gusta que las cosas complejas se hayan simplificado y expresado en sus partes esenciales.