Implementación de Mamba en un solo archivo de PyTorch

 (github.com/johnma2006)
1 puntos por GN⁺ 2023-12-21 | 1 comentarios | Compartir por WhatsApp
  • mamba-minimal es un proyecto que implementa Mamba de forma simple y mínima en un solo archivo de PyTorch
  • Su objetivo es producir la misma salida numérica que la implementación oficial en el forward/backward pass
  • El código está simplificado y organizado en una forma comentada para que sea fácil de leer
  • No incluye las optimizaciones clave de la implementación oficial, por lo que no ofrece velocidad, ni incluye una inicialización adecuada de parámetros
  • La demo ejecuta un ejemplo de completado de prompt usando state-spaces/mamba-370m y el tokenizador EleutherAI/gpt-neox-20b

Resumen del proyecto

  • mamba-minimal es una implementación simple y mínima de Mamba en un solo archivo de PyTorch
  • El objetivo es mostrar el mismo comportamiento que la implementación oficial con un código más fácil de leer
  • Características principales:
    • Salida numérica equivalente a la implementación oficial en el forward pass y el backward pass
    • Código simplificado
    • Implementación legible y comentada

Lo que no incluye

  • La velocidad no es el objetivo
    • La implementación oficial está fuertemente optimizada
    • Esas optimizaciones forman parte de la contribución central del paper de Mamba
    • Esta implementación mantiene la mayor parte del código simple por legibilidad
  • No incluye una inicialización adecuada de parámetros
    • Se presenta como algo que podría agregarse sin sacrificar legibilidad

Ejemplo de uso de la demo

  • En demo.ipynb se puede ver un ejemplo de completado de prompt
  • El ejemplo usa model.Mamba y AutoTokenizer de Hugging Face transformers
  • Modelo y tokenizador usados:
    • state-spaces/mamba-370m
    • EleutherAI/gpt-neox-20b
  • El prompt de ejemplo es Mamba is the, y el resultado generado incluye una oración que describe a Mamba como una víbora

Material de referencia

Lecturas relacionadas

1 comentarios

 
GN⁺ 2023-12-21
Opiniones de Hacker News

  • Hace un tiempo, con un colega creamos una biblioteca que separaba la mayor parte del código de modelos compartido; al usarla, se pueden implementar muchos modelos en unas 100 líneas, sin contar los imports de Python ni los comentarios
    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/...
    También soporta funciones como TorchScript JIT y PyTorch flash attention

    • Definitivamente voy a revisar esta biblioteca. Me pregunto si también vieron xformers
      xformers aborda un problema similar, pero está más enfocado en ofrecer módulos Transformer de buen rendimiento usando Triton. Sin embargo, no fue fácil usar solo componentes específicos de la biblioteca, y seguían apareciendo errores en tiempo de ejecución, así que lo dejé de lado por ahora. Estoy construyendo algo basado en la arquitectura de BERT, así que lo tomaré como referencia
    • Esta biblioteca me impresiona. Nunca me gustó mucho la implementación de Hugging Face, pero esto parece una API hermosa con el nivel de abstracción justo
      Pienso usarla en mi próximo proyecto

  • El código de Mamba original tiene muchas optimizaciones de velocidad y otros elementos, por lo que no es fácil de entender de inmediato; esta implementación parece útil para aprender
    Al hacer inferencia token por token, todo se vuelve mucho más simple. También tengo una implementación propia de inferencia de Mamba: https://github.com/rbitr/llm.f90/tree/master/ssm

    • Fortran, vaya. Me pregunto por qué usaste Fortran
      Sé que es la base de mucho código de computación científica probado durante años y que a menudo se envuelve con bibliotecas como PyTorch o Numpy, pero no es un lenguaje popular hoy en día. Me da curiosidad saber por qué lo elegiste

  • Hay partes sobre Mamba que me gustaría que explicaran de forma que incluso alguien que no sea investigador de machine learning pueda entenderlas

    1. Cuál es la intuición general detrás de los modelos de espacio de estados más allá de los Transformers
    2. Qué innovaciones incrementales hicieron que Mamba fuera más exitoso o interesante que trabajos anteriores como S4, H3 o Monarch
    3. Qué significado tiene aparte de la escalabilidad subcuadrática con la longitud del contexto. Por ejemplo, si no me interesan longitudes de contexto de más de 100k tokens, me pregunto si Mamba podría tener mejor eficiencia de cómputo de entrenamiento en modelos y datasets de tamaño similar
    • Mi inteligencia está muy por debajo de la de los autores del paper, pero aun así intenté entenderlo. Estudié ciencias de la computación y tengo intuición básica de teoría de control y de sistemas de tiempo discreto a nivel licenciatura, pero parece que para entender bien este paper tendría que estudiar mucho más los modelos de espacio de estados
      La intuición clave de Mamba está en resolver un problema antiguo de los modelos de espacio de estados. Los modelos de espacio de estados son buenos para comprimir el contexto de entrada, pero en el proceso de comprimir la entrada en un estado oculto se borra la información necesaria para aprovechar el contexto con tanta eficacia como un Transformer
      La solución es crear lo que el paper llama un mecanismo de selección. Este mecanismo depende de la entrada, así que cada vez que cambia la entrada, el modelo puede ajustar la salida de cada paso. Para ello, algunas variables del espacio de estados dejan de ser invariantes a la entrada y pasan a depender de ella, y se agregan capas lineales, entre otras cosas, para proyectar la entrada de cada instante a variables de espacio de estados
      Pero hacer que las variables de espacio de estados dependan de la entrada introduce overhead computacional. Esto se resuelve con un algoritmo consciente del hardware que aprovecha al máximo la estructura de memoria de las GPU modernas, evitando en lo posible mover datos hacia dentro y fuera de HBM
      Tri Dao es quien creó Flash Attention, que también era una forma de usar el hardware de manera más eficiente en Transformers. Este tipo de cosas es realmente su especialidad
    • Attention crece cuadráticamente con la longitud del contexto; las redes neuronales recurrentes con gating (LSTM, GRU, etc.) son lineales, y estas nuevas arquitecturas también lo son. Las redes recurrentes iniciales usaban gating para evitar gradientes explosivos, pero los nuevos enfoques usan teoría de sistemas dinámicos que garantiza estabilidad, de modo que el gating no tenga que resolver dos problemas a la vez y pueda concentrarse en la memoria
      Mamba y Based, que salieron justo antes de NeurIPS 2023, incorporaban recuperación asociativa de múltiples consultas (MQAR) y dependencia de los datos en el gating/la selección inspirada en la Attention multi-head. Estos eran los elementos clave que les faltaban a Hyena y a arquitecturas anteriores de espacio de estados, y los nuevos modelos se vuelven tan buenos como Attention en tareas de recuperación asociativa; en tareas que no son de consulta, quizá incluso puedan ser un poco mejores que Attention
      Por supuesto, el gran detalle de Mamba es su implementación eficiente en CUDA. Sin eso, esta arquitectura podría tener menos sentido en tareas para las que los Transformers ya son adecuados
      Aunque no te preocupe demasiado la longitud del contexto, se abren muchas áreas nuevas. El análisis de secuencias de ADN es una tarea lineal con dependencias largas, y también se puede pensar en imágenes, video e información de alta dimensión como streams de tokens. Sería como recorrer píxeles al estilo de los viejos monitores CRT
      Uno de los sueños iniciales de la IA era que la trayectoria única de aprendizaje de un agente que interactúa continuamente con su entorno evolucionara de forma constante, y este tipo de modelos de contexto infinito podría hacer más fácil ese sueño
      Dicho eso, por ahora las aplicaciones downstream de estos modelos en tareas reales importantes están, en general, menos validadas y ajustadas que las aplicaciones maduras basadas en Attention. La analogía con las viejas redes recurrentes ayuda hasta cierto punto, pero durante los últimos 5 años la gente se ha especializado demasiado en Attention y Transformers, así que hay mucha inercia del lado de los Transformers
    • Yo también quisiera saber si Mamba puede entrenarse de forma más eficiente en cómputo con modelos y datasets de tamaño similar
      El paper original explica que, una vez transformados los parámetros, el modelo puede calcularse de dos maneras: como una recurrencia lineal o como una convolución global. En general, durante el entrenamiento, cuando se puede ver de antemano toda la secuencia de entrada, se usa el modo de convolución, que es fácil de paralelizar; y durante la inferencia autorregresiva, donde la entrada se ve un paso a la vez, se cambia al modo recurrente eficiente
      Así que el entrenamiento es paralelizable, como el modo de forward paralelo de RetNet. La inferencia básica se realiza en modo recurrente para obtener un contexto lo más largo posible, y como no hay chunking, es difícil juzgar cuánta RAM y VRAM consumirá durante la inferencia
    • Creo que este video encaja exactamente con lo que estabas buscando
      Explica el paper y, al mismo tiempo, da bastante contexto sobre dónde encaja dentro del panorama general. Es bastante interesante escuchar el desarrollo
      https://youtu.be/ouF-H35atOY?si=y2Ckp9MCFd7ulLL3
    • Hasta donde sé, Mamba es básicamente una continuación de la investigación en modelos de espacio de estados que podrían llamarse convoluciones largas
      En lugar de la Attention cuadrática, que calcula cuánta atención presta cada token a todos los demás tokens, de alguna manera calcula un kernel de convolución largo de la misma longitud que la entrada y luego aplica conv1d
      Según mi comprensión limitada, está algo relacionado con aplicar FFT, hacer una multiplicación de matrices y luego volver con IFFT. Sé que funciona, pero es lento. Hay varias formas de calcular una FFT, y una de ellas son las matrices mariposa. Probablemente sea solo una aproximación, pero parece lo suficientemente buena y muy rápida y eficiente en el hardware actual
      La complejidad cuadrática suena mal, pero en la práctica, por restricciones de hardware, los algoritmos subcuadráticos muchas veces son más lentos. Así que, aunque haya muchas expectativas sobre los modelos de espacio de estados, no es fácil decir que Llama está acabado. Todavía no sabemos si Mamba funcionará bien al escalarlo, y para saberlo habría que gastar realmente millones de dólares en entrenamiento. Aun así, soy optimista
      Otro modelo interesante de la familia subcuadrática es RWKV. Vale la pena revisarlo, aunque creo que probablemente ya lo trataron en el podcast
      Aprendí por mi cuenta y además solo hojeé el paper por encima hace tiempo, así que podría estar equivocado en muchas cosas. Además, Attention normalmente tiene una caché KV, lo que ayuda mucho al rendimiento, y entiendo que en Mamba no se puede hacer eso

  • Me dio risa la frase “Mamba es la serpiente venenosa más larga del mundo, con una longitud estimada de más de 150 m”
    Aun así, es realmente excelente, y agradezco que haga referencia al paper de arXiv, porque permite que personas como yo, que consumimos este tipo de artículos en vez de interpretar directamente el paper, podamos asomarnos un poco al interior

    • El nombre Mamba está bueno. Es [S]elective [S]tructured [S]tate [S]pace [S]equence models, así que queda sSSSS, y suena como una serpiente
    • Yo pensaba que la serpiente venenosa más larga era la cobra real. Una búsqueda rápida en Google también decía eso
      Sería divertido si más adelante tuvieran que publicar una corrección de esa frase

  • Esperaba que el núcleo del algoritmo fuera el escaneo paralelo de prefijos. Creo que ese es el punto 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)

  • Quizá sea una pregunta tonta, pero me da curiosidad qué tan difícil es entrenar el modelo Mamba que está en Hugging Face
    El modelo más grande parece ser de 2.8B; si quisiera entrenarlo con un dataset como The Pile, ¿cuántas GPU harían falta y cuánto tardaría?

    • Es una excelente pregunta que yo también quisiera saber. La respuesta parece ser que es bastante más rápido que un Transformer del mismo tamaño, y que el resultado final también debería obtener mejores puntajes que un Transformer en casi todos los benchmarks
      La inferencia también parece ejecutarse entre 3 y 5 veces más rápido, usando solo la mitad de la RAM

  • Intenté desentrañar la versión oficial en CUDA, pero después de que falló el primer intento terminé dejándola; esta implementación se ve mucho mejor

  • Otra implementación de PyTorch en un solo archivo, de verdad excelente. Espero que hlb-CIFAR10 y proyectos relacionados que hice antes, además de influencias previas como minGPT o DawnBench, hayan ayudado aunque sea un poco a impulsar el formato simple de un solo archivo
    Este tipo de trabajo es importante para la investigación eficiente en machine learning, y quizá sea una de las cosas más importantes que se pueden hacer ahora mismo por el campo
    La investigación avanza a la velocidad de la innovación, y la innovación se acelera de forma inversamente proporcional al tiempo de ejecución de los experimentos; esto claramente está relacionado con la complejidad de Kolmogorov del código para investigación o para hackeos simples
    Nunca está de más enfatizar cuán importantes son estas herramientas para la investigación y cuánto aceleraron, personalmente, mi proceso de descubrimiento de conocimiento. La capacidad de bosquejar ideas rápidamente en cuestión de minutos y recibir de inmediato resultados con alta relación señal-ruido se volvió esencial para avanzar en investigación
    Creo que la destilación de conocimiento y MDL (https://en.wikipedia.org/wiki/Minimum_description_length) son muy importantes para revertir la ornamentación innecesaria, el desorden y la competencia excesivamente densa por temas de bajo valor “para no quedarse atrás” que el proceso actual de envío y revisión de papers parece incentivar
    Últimamente, para evitar este problema e ir hacia una solución de escalado un poco mejor, empecé a publicar el código como “bocetos de código”: gists breves, autocontenidos y de un solo archivo. Reducen el tiempo de desarrollo y permiten entregar directamente a la gente código funcional, tosco y sin pulir, que contiene el concepto. Hasta ahora parece funcionar bastante bien y quiero seguir haciéndolo
    Me gustaría ver más código así. Si los investigadores entrenan datos a gran escala, también deberían ser eficientes en datos en la forma en que difunden información

    • 2023 fue un año interesante simplemente por ver cómo la investigación en IA se desplegaba a una velocidad absurda. Elementos de base como ArXiV, PyTorch, GitHub, Hugging Face y código Python open source conciso están acelerando drásticamente el progreso de este nuevo campo
      Probablemente la humanidad nunca haya desarrollado algo con una complejidad considerable tan rápido
      El único lugar donde se ve una velocidad similar quizá sea SpaceX, que también lanzó este año dos cohetes de vanguardia. Me da curiosidad qué saldrá en 2024
    • Hay una posible pequeña mejora de rendimiento. Como aquí x_proj no tiene bias, parece que se podrían combinar los pesos de x_proj y dt_proj
      Si hay requisitos de ajuste de pesos, quizá se pueda hacer simplemente en runtime, y un solo kernel con bias probablemente termine siendo más rápido. No estoy seguro

  • Me pregunto si hubo una discusión sobre el paper original. Parece que me la perdí, pero es bastante interesante
    No entendí bien la parte que dice: “debido a la falta de una implementación eficiente, lo que causa falta de memoria o requisitos de cómputo poco realistas, faltan los resultados completos con longitud de contexto de 8k para las líneas base RWKV y RetNet, modelos recurrentes fuertes anteriores que también pueden interpretarse como SSM”
    RetNet no usa mucha memoria, y si se usa una implementación de forward por chunks, el uso de VRAM queda limitado por el tamaño del chunk. Esa es la parte clave para probar la longitud de contexto
    Me pregunto si alguien probó el modelo Mamba original. ¿Qué tan rápida será la velocidad de entrenamiento comparada con RetNet en modo de forward paralelo?

  • Siempre se agradece una implementación que reduzca algo complejo a lo esencial