mamba

Mamba é um código-fonte de pesquisa de código aberto do grupo state-spaces que implementa uma família de arquiteturas de modelo de espaço de estados (SSM) para modelagem de sequências. Ele acompanha os artigos "Mamba: Linear-Time Sequence Modeling with Selective State Spaces", "Transformers are SSMs" (Mamba-2) e "Mamba-3: Improved Sequence Modeling using State Space Principles". O projeto se posiciona como uma alternativa subquadrática aos Transformers, particularmente para dados densos em informação, como modelagem de linguagem, e baseia-se em trabalhos anteriores de espaço de estados estruturado (S4) com uma implementação ciente de hardware no espírito do FlashAttention. O que o repositório fornece: - Camada SSM seletiva: a operação central de varredura seletiva, com código-fonte em ops/selective_scan_interface.py. - Bloco Mamba: o principal módulo de arquitetura que envolve o SSM seletivo (modules/mamba_simple.py), utilizável como um módulo PyTorch substituto direto com parâmetros configuráveis d_model, d_state, d_conv e expand. - Bloco Mamba-2: implementado em modules/mamba2.py com uma variante mais simples em modules/mamba2_simple.py, além de um módulo de referência SSD (dualidade de espaço de estados estruturado) mínimo em modules/ssd_minimal.py. - Bloco Mamba-3: implementado em modules/mamba3.py, expondo opções como d_state, headdim, modo MIMO com mimo_rank, chunk_size, uma norma opcional de projeção de saída e dtype bfloat16. - Exemplo de modelo de linguagem: um backbone completo mais cabeça de LM em models/mixer_seq_simple.py, usado pelos scripts de geração. Opções de instalação e build: O pacote é instalado via pip (mamba-ssm) e requer Linux, Python 3.10+ e PyTorch 1.12+. Por padrão, a instalação não compila a extensão selective_scan_cuda e não busca wheels CUDA em cache. Extras opcionais e flags de ambiente mudam esse comportamento: o extra causal-conv1d adiciona essa dependência; MAMBA_FORCE_BUILD=TRUE força um build local; MAMBA_KEEP_CUDA_BUILD=TRUE opta pela extensão CUDA de varredura seletiva, primeiro tentando um wheel CUDA/HIP pré-construído correspondente e compilando localmente se nenhum estiver disponível. Combinar ambas as flags força a compilação CUDA local. O README observa que --no-build-isolation é necessário para builds CUDA, para que o pip use o PyTorch existente habilitado para CUDA. Builds CUDA e execução em GPU também exigem uma GPU NVIDIA e CUDA 11.6+. Para placas AMD, o README documenta uma etapa de patch ROCm 6.0 (não necessária a partir do ROCm 6.1). Modelos pré-treinados: Os pesos são publicados no Hugging Face sob a organização state-spaces, incluindo 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 e mamba2attn-2.7b, treinados em 300B tokens no Pile, além de mamba-2.8b-slimpj treinado em 600B tokens no SlimPajama. O README descreve esses modelos como modelos base, sem ajuste de instruções ou outra modificação downstream, e afirma que o desempenho é esperado como comparável ou melhor que outras arquiteturas treinadas em dados semelhantes, mas não deve corresponder a modelos maiores ou ajustados. As dimensões dos modelos seguem o escalonamento estilo GPT-3 (por exemplo, 130M parâmetros com 24 camadas e dimensão de modelo 768, até 2.8B com 64 camadas e dimensão 2560). Ferramentas de avaliação e inferência: A avaliação zero-shot é feita por meio da biblioteca lm-evaluation-harness (pip install lm-eval==0.4.2), com comandos de exemplo cobrindo tarefas como lambada_openai, hellaswag, piqa, arc_easy, arc_challenge, winogrande, openbookqa, boolq, race, truthfulqa_mc2 e mmlu. O README observa que os resultados por tarefa podem diferir dos valores relatados em 0,1-0,3 devido ao ruído de avaliação. Um script de benchmark de geração (benchmarks/benchmark_generation_mamba_simple.py) carrega automaticamente um modelo do Hugging Face Hub, gera continuações para um prompt do usuário e mede a velocidade de inferência, com opções para top-p, min-p, top-k, temperatura, penalidade de repetição e tamanho do lote. Notas de solução de problemas: O README destaca dois problemas práticos. Primeiro, precisão: os modelos foram treinados com AMP do PyTorch, que mantém os parâmetros em float32 e converte para meia precisão conforme necessário; como SSMs são sensíveis à dinâmica recorrente, frameworks que armazenam parâmetros em float16 (como DeepSpeed) podem causar instabilidade, e o armazenamento de parâmetros em fp32 é sugerido como primeiro remédio. Segundo, inicialização: alguns componentes herdam inicializações estilo S4 (por exemplo, o intervalo alvo do parâmetro delta via viés de projeção linear), e frameworks com hooks pós-inicialização que zeram os vieses de nn.Linear podem precisar de lógica personalizada para preservá-los. Licenciamento e citação: o README fornece entradas BibTeX para os artigos Mamba, Mamba-2 e Mamba-3 e pede que os usuários os citem ao usar o código.