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

mamba-minimal

• Explicação sobre o mamba-minimal, uma implementação simples do Mamba em um único arquivo escrita em PyTorch.
• Fornece a mesma saída numérica da implementação oficial, tanto na propagação direta quanto na retropropagação.
• O código foi simplificado, é fácil de ler e possui comentários.
• Não inclui recursos como as otimizações de velocidade da implementação oficial.
• Não inclui uma inicialização adequada de parâmetros, mas isso pode ser adicionado sem prejudicar a legibilidade.

Demo

• Arquivo demo.ipynb que mostra um exemplo de conclusão de prompt.
• Fornece um exemplo de geração de texto usando o modelo Mamba e o AutoTokenizer.
• No texto de exemplo gerado, o Mamba é descrito como a cobra venenosa mais longa do mundo.

References

• A arquitetura Mamba foi apresentada no artigo "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", de Albert Gu e Tri Dao.
• A implementação oficial pode ser consultada no GitHub.

Opinião do GN⁺

• O mamba-minimal é um projeto criado para reduzir a complexidade das implementações existentes do Mamba, permitindo que até engenheiros de software iniciantes possam entendê-lo.
• Este projeto contribui para aumentar a legibilidade e a compreensão de código na área de aprendizado de máquina.
• Por meio de exemplos práticos de uso, ele mostra de forma simples como utilizar o modelo Mamba, o que pode ser um material bastante interessante para quem está aprendendo.

[GN⁺]

Opiniões do Hacker News

• Compartilhamento de biblioteca

  • Criou com um colega uma biblioteca para extrair a maior parte do código comum dos modelos. A maioria dos modelos pode ser implementada em cerca de 100 linhas, excluindo o processo de import do Python e os comentários.
  • Cita como exemplos os modelos BERT, Llama 1/2 e MPT, com suporte a TorchScript JIT, PyTorch flash attention etc.

• Compartilhamento de implementação de inferência do Mamba

  • Como o código do Mamba tem muitas partes, como otimizações de velocidade, que dificultam seu uso direto em treinamento, compartilha sua própria implementação de inferência do Mamba para ajudar nisso.

• Pedido de explicação do Mamba para leigos

  • Pede uma explicação acessível sobre o principal insight do Mamba, os modelos de espaço de estados, por que o Mamba teve sucesso e quais outras vantagens ele tem mesmo para quem não se importa com o comprimento de contexto.

• Expectativa sobre o núcleo do algoritmo

  • Esperava que o núcleo do algoritmo Mamba fosse uma varredura paralela de prefixos, mas o código real é implementado de outra forma.

• Piada espirituosa sobre o Mamba

  • Achou engraçada a piada de que a mamba é a cobra venenosa mais longa do mundo. Gostou de poder dar uma olhada no que está por trás da pesquisa consultando o artigo no arXiv.

• Pergunta sobre a dificuldade de treinar o modelo Mamba

  • Pergunta o quão difícil é treinar o modelo Mamba disponível no Huggingface e quantas GPUs e quanto tempo seriam necessários para treinar o maior modelo.

• Compartilhamento de tentativa de interpretar a versão oficial em CUDA

  • Tentou interpretar a versão oficial em CUDA, falhou e não voltou a tentar. Em comparação, considera esta nova implementação muito melhor.

• Admiração por uma implementação PyTorch em arquivo único

  • Admira mais uma implementação em arquivo único em PyTorch e considera esse tipo de trabalho importante para uma pesquisa eficiente em ML. Diz que a velocidade do progresso da pesquisa está relacionada ao inverso do tempo de execução dos experimentos, o que, por sua vez, se relaciona com a complexidade de Kolmogorov do código. Destaca o quanto as ferramentas de pesquisa são importantes para acelerar o processo de descoberta de conhecimento.

• Pedido de discussão do artigo original

  • Pergunta se houve discussão sobre o artigo original e aponta que modelos anteriores como RWKV e RetNet não apresentaram resultados completos para comprimento de contexto de 8k por falta de implementações eficientes, devido a problemas de memória ou exigências computacionais irreais.

• Elogio à simplificação do conteúdo central

  • Gosta de ver coisas complexas sendo simplificadas e apresentadas em seus elementos essenciais.