Implementação do Mamba em um único arquivo PyTorch

• mamba-minimal é um projeto que implementa o Mamba de forma simples e mínima em um único arquivo PyTorch
• O objetivo é produzir a mesma saída numérica da implementação oficial no forward/backward pass
• O código foi simplificado e estruturado em um formato comentado para facilitar a leitura
• Não inclui as otimizações centrais da implementação oficial, portanto não oferece velocidade, nem inclui inicialização adequada de parâmetros
• A demo executa um exemplo de conclusão de prompt usando state-spaces/mamba-370m e o tokenizador EleutherAI/gpt-neox-20b

---

Visão geral do projeto

• mamba-minimal é uma implementação simples e mínima do Mamba em um único arquivo PyTorch
• O objetivo é mostrar o mesmo comportamento da implementação oficial com um código mais fácil de ler
• Principais características:
  • Saída numérica equivalente à implementação oficial no forward pass e no backward pass
  • Código simplificado
  • Implementação comentada e fácil de ler

O que não está incluído

• Velocidade não é o objetivo
  • A implementação oficial é fortemente otimizada
  • Essas otimizações fazem parte da principal contribuição do artigo do Mamba
  • Esta implementação mantém a maior parte do código simplificada em favor da legibilidade
• Inicialização adequada de parâmetros não está incluída
  • Isso é apresentado como algo que pode ser adicionado sem sacrificar a legibilidade

Exemplo de uso da demo

• É possível ver um exemplo de conclusão de prompt em demo.ipynb
• O exemplo usa model.Mamba e o AutoTokenizer do Hugging Face transformers
• Modelo e tokenizador usados:
  • state-spaces/mamba-370m
  • EleutherAI/gpt-neox-20b
• O prompt de exemplo é Mamba is the, e o resultado gerado inclui uma frase descrevendo a mamba como uma cobra venenosa

Referências

• A arquitetura Mamba é apresentada em Mamba: Linear-Time Sequence Modeling with Selective State Spaces
• Os autores do artigo são Albert Gu e Tri Dao
• A implementação oficial está no repositório state-spaces/mamba

[GN⁺]

Opiniões do Hacker News

• Há algum tempo, com um colega, criei uma biblioteca que separava a maior parte do código de modelos compartilhado; usando-a, é possível implementar muitos modelos em cerca de 100 linhas, excluindo imports em Python e comentários
  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/...
  Também oferece suporte a recursos como TorchScript JIT e PyTorch flash attention

  • Com certeza vou dar uma olhada nessa biblioteca. Fico curioso se você também viu o xformers
    O xformers também aborda problemas parecidos, mas é mais focado em fornecer módulos Transformer de alto desempenho usando Triton. Porém, não foi fácil usar apenas componentes específicos da biblioteca, e erros em tempo de execução continuavam aparecendo, então deixei de lado por enquanto. Como estou criando algo baseado na arquitetura BERT, vou usá-la como referência
  • Fiquei impressionado com essa biblioteca. Nunca gostei muito da implementação da Hugging Face, mas esta parece uma API bonita, com o nível de abstração na medida certa
    Pretendo usá-la no meu próximo projeto

• O código original do Mamba tem muitas otimizações de velocidade e outros elementos, então não é fácil entendê-lo de imediato; esta implementação deve ajudar no aprendizado
  Ao fazer inferência token por token, tudo fica muito mais simples. Também tenho uma implementação própria de inferência do Mamba: https://github.com/rbitr/llm.f90/tree/master/ssm

  • Fortran, hein. Fico curioso por que você usou Fortran
    Sei que ele é a base de código de computação científica há muito validado e que muitas vezes é usado por trás de bibliotecas como PyTorch ou Numpy, mas não é uma linguagem popular hoje em dia. Gostaria de saber o motivo da escolha

• Há partes sobre o Mamba que eu gostaria que fossem explicadas de um jeito que até quem não é pesquisador de machine learning consiga entender:

  1. Qual é a visão geral dos modelos de espaço de estados além dos Transformers
  2. Qual foi a inovação incremental que tornou o Mamba mais bem-sucedido ou interessante do que antecessores como S4, H3 e Monarch
  3. Que significado ele tem além da escalabilidade subquadrática em relação ao comprimento do contexto. Por exemplo, se eu não me interesso por comprimentos de contexto acima de 100k tokens, fico curioso se o Mamba poderia ter melhor eficiência computacional de treinamento em modelos e datasets de tamanho parecido
  • Minha inteligência está muito abaixo da dos autores do artigo, mas ainda assim tentei entender. Estudei ciência da computação e tenho noções básicas de teoria de controle e uma intuição de sistemas em tempo discreto em nível de graduação, mas parece que, para entender bem este artigo, vou precisar estudar muito mais modelos de espaço de estados. O insight central do Mamba está em resolver um problema antigo dos modelos de espaço de estados. Modelos de espaço de estados são bons para comprimir o contexto de entrada, mas, no processo de comprimir a entrada em um estado oculto, informações necessárias para usar o contexto de forma eficaz como em um Transformer acabam sendo apagadas. A solução é criar o que o artigo chama de mecanismo de seleção. Esse mecanismo depende da entrada, de modo que, sempre que a entrada muda, o modelo consegue ajustar a saída de cada etapa. Para isso, algumas variáveis do espaço de estados deixam de ser invariantes à entrada e passam a depender dela, e são anexadas camadas lineares etc. para projetar a entrada de cada instante nas variáveis do espaço de estados. Mas tornar as variáveis do espaço de estados dependentes da entrada cria overhead computacional. Isso é resolvido com um algoritmo ciente do hardware que aproveita ao máximo a estrutura de memória das GPUs modernas, evitando tanto quanto possível mover dados para dentro e para fora da HBM. Tri Dao é a pessoa que criou o Flash Attention, que também era uma forma de usar o hardware de maneira mais eficiente em Transformers. Esse tipo de coisa é realmente a especialidade dele.
  • Attention cresce quadraticamente com o comprimento do contexto, redes neurais recorrentes com gating (LSTM, GRU etc.) são lineares, e essas novas arquiteturas também são lineares. As redes recorrentes iniciais usavam gating para evitar gradientes explosivos, mas as novas abordagens usam teoria de sistemas dinâmicos para garantir estabilidade, permitindo que o gating se concentre na memória em vez de resolver dois problemas ao mesmo tempo. Mamba e Based, que saíram pouco antes da NeurIPS 2023, incluíam recuperação associativa de múltiplas consultas (MQAR) e a dependência dos dados no gating/seleção inspirada em Attention multi-head. Esses eram os elementos essenciais que faltavam no Hyena e em arquiteturas anteriores de espaço de estados, e os novos modelos ficaram tão bons quanto Attention em tarefas de recuperação associativa, além de mostrarem a possibilidade de serem talvez um pouco melhores do que Attention em tarefas que não são de lookup. Claro, um grande detalhe do Mamba é a implementação CUDA eficiente. Sem isso, em trabalhos para os quais Transformers já são adequados, o valor dessa arquitetura pode diminuir. Mesmo sem se preocupar muito com comprimento de contexto, muitas áreas novas se abrem. Análise de sequências de DNA é uma tarefa linear com dependências longas, e também dá para pensar em imagens, vídeos e informações de alta dimensão como fluxos de tokens. É como varrer pixels, como em um monitor CRT antigo. Um dos sonhos iniciais da IA era uma única trajetória de aprendizado de um agente interagindo continuamente com o ambiente evoluir de forma contínua, e esses modelos de contexto infinito talvez tornem esse sonho mais fácil. Por enquanto, porém, as aplicações downstream desses modelos em tarefas reais importantes em geral foram menos validadas e ajustadas do que as aplicações maduras baseadas em Attention. A analogia com redes recorrentes antigas ajuda em certa medida, mas nos últimos 5 anos as pessoas se especializaram demais em Attention e Transformers, então a inércia do lado dos Transformers é grande.
  • Eu também gostaria de saber se, em modelos e datasets de tamanho parecido, o Mamba pode ser treinado com mais eficiência computacional. O artigo original explica que, depois que os parâmetros são transformados, o modelo pode ser calculado de duas formas: como uma recorrência linear ou como uma convolução global. Em geral, no treinamento, quando é possível ver a sequência de entrada inteira de antemão, usa-se o modo de convolução, que é fácil de paralelizar; já na inferência autorregressiva, em que a entrada é vista um passo de cada vez, alterna-se para o modo recorrente eficiente. Portanto, o treinamento é paralelizável, como o modo de forward pass paralelo do RetNet. A inferência básica é feita no modo recorrente para obter o contexto mais longo possível, e, como não há chunking, é difícil julgar quanta RAM e VRAM a inferência consumirá.
  • Este vídeo parece ser exatamente o que você estava procurando. Ele explica o artigo, mas também dá bastante contexto sobre onde ele se encaixa no panorama geral. É bem interessante acompanhar o desenvolvimento. https://youtu.be/ouF-H35atOY?si=y2Ckp9MCFd7ulLL3
  • Pelo que sei, o Mamba é basicamente uma continuação da pesquisa em modelos de espaço de estados que podem ser chamados de convoluções longas. Em vez de Attention quadrático, que calcula quanto cada token deve prestar atenção a todos os outros tokens, calcula-se de alguma forma um kernel de convolução longo do mesmo comprimento que a entrada e depois aplica-se conv1d. Pelo meu entendimento limitado, isso tem alguma relação com aplicar FFT, fazer multiplicação de matrizes e depois voltar com IFFT. Sei que funciona, mas é lento. Há várias maneiras de calcular uma FFT, e uma delas é a matriz borboleta. Talvez seja apenas uma aproximação, mas parece boa o suficiente e muito rápida e eficiente no hardware atual. Complexidade quadrática soa ruim, mas, na prática, por causa de restrições de hardware, muitas vezes algoritmos subquadráticos acabam sendo mais lentos. Então, mesmo com grandes expectativas para modelos de espaço de estados, não é fácil dizer que o Llama acabou. Também ainda não sabemos se o Mamba continua funcionando bem em escala, e, para saber isso, é preciso de fato gastar milhões de dólares em treinamento. Ainda assim, estou otimista. Outro modelo interessante da família subquadrática é o RWKV. Vale a pena dar uma olhada, mas acho que provavelmente já foi abordado no podcast. Aprendi por conta própria e só dei uma olhada superficial no artigo há um tempo, então posso estar bastante errado. Além disso, Attention normalmente tem cache KV, o que ajuda muito no desempenho, mas acho que não dá para fazer isso no Mamba.

• Ri da frase “a mamba é a cobra venenosa mais longa do mundo, com comprimento estimado acima de 150 m”
  Ainda assim, é realmente excelente, e achei bom que tenha referenciado o artigo no arXiv, porque pessoas como eu, que consomem textos assim em vez de interpretar o artigo diretamente, conseguem dar uma espiada por dentro

  • O nome Mamba é bom. Como é [S]elective [S]tructured [S]tate [S]pace [S]equence models, vira sSSSS, e parece som de cobra
  • Eu achava que a cobra venenosa mais longa era a cobra-real. Uma busca rápida no Google também mostrou isso
    Seria engraçado se depois tivessem que publicar uma correção para essa frase

• Eu esperava que o núcleo do algoritmo fosse o parallel prefix scan. Acho que esse é o ponto do 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)

• Pode ser uma pergunta boba, mas fico curioso sobre a dificuldade de treinar o modelo Mamba disponível no Hugging Face
  O maior modelo parece ser de 2,8B; quantas GPUs seriam necessárias e quanto tempo levaria para treiná-lo em um dataset como The Pile?

  • Também acho uma ótima pergunta e queria saber. A resposta parece ser que ele é consideravelmente mais rápido do que um Transformer do mesmo tamanho, e o resultado final também deve pontuar melhor que Transformers em praticamente todos os benchmarks
    A inferência também parece rodar 3 a 5 vezes mais rápido, usando apenas metade da RAM

• Tentei destrinchar a versão CUDA oficial, mas depois que a primeira tentativa falhou acabei não mexendo mais; esta implementação parece muito melhor

• Mais uma implementação em PyTorch em arquivo único, realmente excelente. Espero que trabalhos anteriores como hlb-CIFAR10 e projetos relacionados, além de influências anteriores como minGPT e DawnBench, tenham ajudado a impulsionar nem que seja um pouco esse formato simples de arquivo único
  Esse tipo de trabalho é importante para pesquisa eficiente em machine learning, e talvez seja uma das coisas mais importantes que se pode fazer agora pela área
  A pesquisa avança no ritmo da inovação, a inovação acelera na proporção inversa ao tempo de execução dos experimentos, e isso claramente tem relação com a complexidade de Kolmogorov do código para fins de pesquisa ou hacking simples
  Não dá para enfatizar o bastante o quanto ferramentas assim são importantes para a pesquisa e o quanto, pessoalmente, elas aceleraram meu processo de descoberta de conhecimento. A capacidade de rascunhar ideias rapidamente em minutos e receber resultados imediatos com alta relação sinal-ruído se tornou essencial para o progresso da pesquisa
  Vejo destilação de conhecimento e MDL(https://en.wikipedia.org/wiki/Minimum_description_length) como muito importantes para reverter os enfeites desnecessários, a bagunça e a competição excessivamente densa por temas de baixo valor para “não ficar para trás” que o processo atual de submissão e revisão de artigos parece incentivar
  Recentemente, para evitar esse problema e caminhar para uma solução de escala um pouco melhor, comecei a distribuir código como “esboços de código”: gists curtos, autocontidos, de um único arquivo. Isso reduz o tempo de desenvolvimento e permite entregar diretamente às pessoas código funcional, bruto e sem polimento, que contém o conceito. Até agora parece estar funcionando bem, e quero continuar
  Gostaria de ver mais código assim. Se são pesquisadores que treinam dados em larga escala, também deveriam ser eficientes em dados na forma como disseminam informação

  • 2023 foi um ano interessante só de ver a pesquisa em IA se desenrolar em uma velocidade absurda. Elementos de base como ArXiV, PyTorch, GitHub, Hugging Face e código Python open source conciso estão acelerando dramaticamente o avanço desse novo campo
    Provavelmente a humanidade nunca desenvolveu algo com complexidade considerável tão rápido assim
    Um lugar onde se vê uma velocidade parecida talvez seja a SpaceX, que também lançou dois foguetes de ponta este ano. Fico curioso para saber o que virá em 2024
  • Há uma pequena possibilidade de melhoria de desempenho. Como aqui x_proj não tem bias, parece que seria possível combinar os pesos de x_proj e dt_proj
    Se houver exigências de ajuste de pesos, talvez isso possa ser feito simplesmente em tempo de execução, e no fim um único kernel com bias provavelmente será mais rápido. Não tenho certeza

• Fico curioso se houve uma discussão sobre o artigo original. Acho que perdi, mas é bem interessante
  Não entendi bem a parte que diz: “por falta de uma implementação eficiente, há falta de memória ou exigências computacionais irrealistas, então faltam os resultados completos para comprimento de contexto de 8k dos baselines RWKV e RetNet, modelos recorrentes fortes anteriores que também podem ser interpretados como SSMs”
  RetNet não usa muita memória e, usando uma implementação de forward em chunks, o uso de VRAM fica limitado pelo tamanho do chunk. Essa é a essência de testar comprimento de contexto
  Fico curioso se alguém testou o modelo Mamba original. Qual seria a velocidade de treinamento em comparação com o RetNet no modo de forward paralelo?

  • https://news.ycombinator.com/item?id=38522428
    https://openreview.net/forum?id=AL1fq05o7H
  • O treinamento é mais rápido, a inferência é muito mais rápida, e o uso de VRAM durante a inferência é aproximadamente metade

• Implementações que reduzem algo complexo ao essencial são sempre boas