Hugging Face's logo

Transformers documentation

Mamba 2

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Este modelo foi publicado no HF papers em 2024-05-31 e contribuído para o Hugging Face Transformers em 2024-08-06.

Mamba 2

Mamba 2 é baseado no framework de dualidade de espaço de estado (SSD), que conecta modelos de espaço de estado estruturados (SSMs) e variantes de atenção. Ele utiliza um algoritmo SSD mais eficiente, que é de 2 a 8 vezes mais rápido que o Mamba, e modifica a arquitetura para habilitar paralelismo de tensor e uma estrutura de cabeçalho de atenção com valores agrupados (GVA).

Você pode encontrar todos os checkpoints originais do Mamba 2 na organização State Space Models, mas os exemplos mostrados abaixo utilizam mistralai/Mamba-Codestral-7B-v0.1, pois uma implementação no Hugging Face ainda não é suportada para os checkpoints originais.

Outras arquiteturas baseadas em Mamba 2 incluem Bamba, FalconH1 e Zamba2.

Este modelo foi contribuído por ArthurZ. Clique nos modelos Mamba na barra lateral à direita para ver mais exemplos de como aplicar o Mamba a diferentes tarefas de linguagem.

O exemplo abaixo demonstra como gerar texto com Pipeline, AutoModel e pela linha de comando.

Pipeline
AutoModel
from transformers import pipeline


pipeline = pipeline(
    task="text-generation",
    model="mistralai/Mamba-Codestral-7B-v0.1",
    device=0
)
pipeline("Plants create energy through a process known as")

A quantização reduz a carga de memória de modelos grandes ao representar os pesos em uma precisão menor. Consulte a visão geral de Quantização para mais backends de quantização disponíveis.

O exemplo abaixo usa torchao para quantizar apenas os pesos para inteiros de 4 bits.

from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig


quantization_config = TorchAoConfig("int4_weight_only", group_size=128)
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mamba-Codestral-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mamba-Codestral-7B-v0.1", quantization_config=quantization_config, device_map="auto")
input_ids = tokenizer("Plants create energy through a process known as", return_tensors="pt").to(model.device)

output = model.generate(**input_ids)
print(tokenizer.decode(output[0], skip_special_tokens=True))

Observações

  • O Codestral Mamba possui groups=8, que são semelhantes ao número de cabeças de kv em um modelo baseado em atenção.

  • O Codestral Mamba possui dois passes de forward diferentes, torch_forward ou cuda_kernels_forward, e seus resultados são esperados para serem ligeiramente diferentes.

    • torch_forward sem compilação é 3 a 4 vezes mais rápido que cuda_kernels_forward.
    • cuda_kernels_forward usa os kernels CUDA originais, caso estejam disponíveis no seu ambiente. Ele é mais lento durante o prefill, pois exige uma “execução de aquecimento” devido ao maior overhead de CPU (veja esses comentários para mais detalhes).

  • Este modelo não possui embeddings posicionais, mas há um attention_mask e uma lógica específica para mascarar estados ocultos em dois locais no caso de geração em lote (veja este comentário para mais detalhes). Isso (e a adição dos kernels Mamba 2 reimplementados) resulta em uma pequena discrepância entre a geração em lote e a geração com cache.

  • O algoritmo SSM depende fortemente de contrações de tensores, que possuem equivalentes em matmul, mas a ordem das operações é ligeiramente diferente. Isso torna a diferença maior em precisões menores.

  • O padding é testado apenas com padding à esquerda. O padding à direita atualiza o cache com as posições de padding e produz resultados não confiáveis. Defina tokenizer.padding_side="left" para geração em lote. Consulte Lado do padding.

  • O exemplo abaixo demonstra como fazer o ajuste fino do Mamba 2 com PEFT.

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer


model_id = "mistralai/Mamba-Codestral-7B-v0.1"
dataset = load_dataset("Abirate/english_quotes", split="train")
training_args = SFTConfig(dataset_text_field="quote", gradient_checkpointing=True, per_device_train_batch_size=4)
lora_config =  LoraConfig(target_modules=["x_proj", "embeddings", "in_proj", "out_proj"])
trainer = SFTTrainer(
    model=model_id,
    args=training_args,
    train_dataset=dataset,
    peft_config=lora_config,
)
trainer.train()

Mamba2Config

classe transformers.Mamba2Config

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonenum_heads: int = 128head_dim: int = 64vocab_size: int = 32768hidden_size: int = 4096state_size: int = 128num_hidden_layers: int = 64layer_norm_epsilon: float = 1e-05pad_token_id: int | None = 1bos_token_id: int | None = 0eos_token_id: int | list[int] | None = 2expand: int = 2conv_kernel: int = 4n_groups: int = 8use_bias: bool = Falseuse_conv_bias: bool = Truehidden_act: str = 'silu'initializer_range: float = 0.1residual_in_fp32: bool = Truetime_step_rank: str | int = 'auto'time_step_min: float = 0.001time_step_max: float = 0.1time_step_floor: float = 0.0001time_step_limit: list[float] | tuple[float, ...] = (0.0, inf)rescale_prenorm_residual: bool = Falseuse_cache: bool = Truechunk_size: int = 256tie_word_embeddings: bool = False )

Parâmetros

  • num_heads (int, opcional, padrão 128) — Número de cabeças de atenção para cada camada de atenção no decodificador do Transformer.
  • head_dim (int, opcional, padrão 64) — A dimensão da cabeça de atenção. Se for None, o valor padrão será hidden_size // num_attention_heads
  • vocab_size (int, opcional, padrão 32768) — Tamanho do vocabulário do modelo. Define o número de tokens diferentes que podem ser representados pelos input_ids.
  • hidden_size (int, opcional, padrão 4096) — Dimensão das representações ocultas.
  • state_size (int, opcional, padrão 128) — Tamanho do estado do SSM (dimensão do estado latente) nas camadas Mamba.
  • num_hidden_layers (int, opcional, padrão 64) — Número de camadas ocultas no decodificador do Transformer.
  • layer_norm_epsilon (float, opcional, padrão 1e-05) — O epsilon a ser usado nas camadas de normalização de camada.
  • pad_token_id (int, opcional, padrão 1) — ID do token usado para preenchimento no vocabulário.
  • bos_token_id (int, opcional, padrão 0) — ID do token usado para o início do fluxo no vocabulário.
  • eos_token_id (Union[int, list[int]], opcional, padrão 2) — ID do token usado para o fim do fluxo no vocabulário.
  • expand (int, opcional, padrão 2) — Fator de expansão usado para determinar o tamanho intermediário.
  • conv_kernel (int, opcional, padrão 4) — O tamanho do kernel de convolução.
  • n_groups (int, opcional, padrão 8) — Número de grupos para as matrizes de evolução do mamba 2.
  • use_bias (bool, opcional, padrão False) — Indica se deve ou não usar viés em ["in_proj", "out_proj"] do bloco mixer
  • use_conv_bias (bool, opcional, padrão True) — Indica se deve ou não usar viés na camada de convolução do bloco do mixer.
  • hidden_act (str, opcional, padrão silu) — A função de ativação não linear (função ou string) no decodificador. Por exemplo, "gelu", "relu", "silu", etc.
  • initializer_range (float, opcional, padrão 0.1) — O desvio padrão do truncated_normal_initializer para inicializar todas as matrizes de pesos.
  • residual_in_fp32 (bool, opcional, padrão True) — Indica se os resíduos devem estar em float32. Se definido como False, os resíduos manterão o mesmo dtype do restante do modelo
  • time_step_rank (Union[str, int], opcional, padrão auto) — Rank da projeção do delta (time step). Pode ser "auto" para definir automaticamente.
  • time_step_min (float, opcional, padrão 0.001) — Valor mínimo de time_step usado para limitar dt_proj.bias.
  • time_step_max (float, opcional, padrão 0.1) — Valor máximo de time_step usado para limitar dt_proj.bias.
  • time_step_floor (float, opcional, padrão 0.0001) — Valor mínimo permitido para o delta do passo de tempo discreto após a ativação softplus.
  • time_step_limit (Union[list[float], tuple[float, ...]], opcional, padrão (0.0, inf)) — Intervalo aceito de valores de passo de tempo para clamping.
  • rescale_prenorm_residual (bool, opcional, padrão False) — Se as pesos de out_proj devem ser reescalados durante a inicialização.
  • use_cache (bool, opcional, padrão True) — Se o modelo deve ou não retornar as últimas chaves/valores de atenção (não usado por todos os modelos). Relevante apenas se config.is_decoder=True ou quando o modelo for um modelo generativo somente decodificador.
  • chunk_size (int, opcional, padrão 256) — Tamanho dos blocos que comporão a sequência.
  • tie_word_embeddings (bool, opcional, padrão False) — Se deve vincular as embeddings de pesos de acordo com o mapeamento tied_weights_keys do modelo.

Esta é a classe de configuração para armazenar a configuração de um Mamba2Model. Ela é usada para instanciar um modelo Mamba2 de acordo com os argumentos especificados, definindo a arquitetura do modelo. Instanciar uma configuração com os padrões resultará em uma configuração semelhante à do state-spaces/mamba2-2.8b

Os objetos de configuração herdam de PreTrainedConfig e podem ser usados para controlar as saídas do modelo. Consulte a documentação de PreTrainedConfig para mais informações.

Exemplo:

>>> from transformers import Mamba2Config, Mamba2Model

>>> # Initializing a Mamba2 configuration
>>> configuration = Mamba2Config()

>>> # Initializing a model (with random weights) from the configuration
>>> model = Mamba2Model(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

Mamba2Model

classe transformers.Mamba2Model

< >

( config: Mamba2Config )

Parâmetros

  • config (Mamba2Config) — Classe de configuração do modelo com todos os parâmetros do modelo. A inicialização com um arquivo de configuração não carrega os pesos associados ao modelo, apenas a configuração. Consulte o método from_pretrained() para carregar os pesos do modelo.

O modelo Mamba2 puro, que produz estados ocultos brutos sem nenhum cabeçalho específico no topo.

Este modelo herda de PreTrainedModel. Consulte a documentação da classe pai para os métodos genéricos que a biblioteca implementa para todos os seus modelos (como download ou salvamento, redimensionamento dos embeddings de entrada, poda de cabeças, etc.)

Este modelo também é uma subclasse de torch.nn.Module do PyTorch. Use-o como um Módulo PyTorch regular e consulte a documentação do PyTorch para todos os assuntos relacionados ao uso geral e ao comportamento.

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.LongTensor] = Nonecache_params: transformers.cache_utils.Cache | None = Noneuse_cache: bool | None = Noneoutput_hidden_states: bool | None = Nonereturn_dict: bool | None = Noneattention_mask: typing.Optional[torch.Tensor] = None**kwargs ) → Mamba2Output or tuple(torch.Tensor)

Parâmetros

  • input_ids (torch.LongTensor de forma (batch_size, sequence_length), opcional) — Índices dos tokens da sequência de entrada no vocabulário. O padding será ignorado por padrão.

    Os índices podem ser obtidos usando AutoTokenizer. Consulte PreTrainedTokenizer.encode() e PreTrainedTokenizer.call() para mais detalhes.

    O que são IDs de entrada?

  • inputs_embeds (torch.LongTensor de forma (batch_size, sequence_length, hidden_size), opcional) — Opcionalmente, em vez de passar input_ids, você pode escolher passar diretamente uma representação incorporada. Isso é útil se você quiser ter mais controle sobre como converter os índices de input_ids nos vetores associados do que a matriz de busca de incorporação interna do modelo.
  • cache_params (Cache, opcional) — Se for passado, o modelo usa o estado anterior em todos os blocos (o que resultará na saída para os input_ids fornecidos, como se o modelo adicionasse state_input_ids + input_ids como contexto).
  • use_cache (bool, opcional) — Se definido como True, o cache_params é retornado e pode ser usado para gerar rapidamente os próximos logits.
  • output_hidden_states (bool, opcional) — Indica se as hidden states de todas as camadas devem ser retornadas. Consulte hidden_states em tensors retornados para mais detalhes.
  • return_dict (bool, opcional) — Indica se deve retornar um ModelOutput em vez de uma tupla simples.
  • attention_mask (torch.Tensor de forma (batch_size, sequence_length), opcional) — Máscara para evitar a execução de atenção nos índices de tokens de preenchimento. Valores de máscara selecionados em [0, 1]:
    • 1 para tokens que não estão mascarados,
    • 0 para tokens que estão mascarados.

    O que são máscaras de atenção?

Retorna

Mamba2Output ou tuple(torch.Tensor)

Um Mamba2Output ou uma tupla de torch.Tensor (se return_dict=False for passado ou quando config.return_dict=False) contendo vários elementos, dependendo da configuração (Mamba2Config) e das entradas.

O método forward do Mamba2Model substitui o método especial __call__.

Embora a receita para a passagem direta precise ser definida dentro desta função, deve-se chamar a instância de Module em seguida, em vez desta, pois a primeira cuida da execução das etapas de pré e pós-processamento, enquanto a segunda as ignora silenciosamente.

  • last_hidden_state (torch.Tensor de forma (batch_size, sequence_length, hidden_size), opcional) — Sequência de estados ocultos na saída da última camada do modelo.

  • cache_params (~cache_utils.Cache, opcional) — O estado do modelo no último passo de tempo. Pode ser usado em um método forward com o próximo input_ids para evitar fornecer os input_ids antigos.

    Inclui tanto as matrizes de estado do modelo de espaço de estados após o scan seletivo quanto os estados convolucionais

  • hidden_states (tuple[torch.Tensor], opcional, retornado quando output_hidden_states=True é passado ou quando config.output_hidden_states=True) — Tupla de torch.Tensor (um para a saída dos embeddings, se o modelo tiver uma camada de embedding, + um para a saída de cada camada) de forma (batch_size, sequence_length, hidden_size).

    Estados ocultos do modelo na saída de cada camada, além das saídas opcionais do embedding inicial.

Mamba2LMHeadModel

classe transformers.Mamba2ForCausalLM

< >

( config: Mamba2Config )

Parâmetros

  • config (Mamba2Config) — Classe de configuração do modelo com todos os parâmetros do modelo. A inicialização com um arquivo de configuração não carrega os pesos associados ao modelo, apenas a configuração. Consulte o método from_pretrained() para carregar os pesos do modelo.

O transformer do modelo MAMBA2 com uma cabeça de modelagem de linguagem no topo (camada linear com pesos não vinculados aos embeddings de entrada).

Este modelo herda de PreTrainedModel. Consulte a documentação da classe pai para os métodos genéricos que a biblioteca implementa para todos os seus modelos (como download ou salvamento, redimensionamento dos embeddings de entrada, poda de cabeças, etc.)

Este modelo também é uma subclasse de torch.nn.Module do PyTorch. Use-o como um Módulo PyTorch regular e consulte a documentação do PyTorch para todos os assuntos relacionados ao uso geral e ao comportamento.

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonecache_params: transformers.cache_utils.Cache | None = Nonelabels: typing.Optional[torch.LongTensor] = Noneoutput_hidden_states: bool | None = Nonereturn_dict: bool | None = Noneuse_cache: bool | None = Noneattention_mask: typing.Optional[torch.Tensor] = Nonelogits_to_keep: typing.Union[int, torch.Tensor] = 0**kwargs ) → Mamba2CausalLMOutput or tuple(torch.Tensor)

Parâmetros

  • input_ids (torch.LongTensor de forma (batch_size, sequence_length), opcional) — Índices dos tokens da sequência de entrada no vocabulário. O padding será ignorado por padrão.

    Os índices podem ser obtidos usando AutoTokenizer. Consulte PreTrainedTokenizer.encode() e PreTrainedTokenizer.call() para mais detalhes.

    O que são IDs de entrada?

  • inputs_embeds (torch.Tensor de forma (batch_size, sequence_length, hidden_size), opcional) — Opcionalmente, em vez de passar input_ids, você pode escolher passar diretamente uma representação incorporada. Isso é útil se você quiser ter mais controle sobre como converter os índices de input_ids nos vetores associados do que a matriz de busca de incorporação interna do modelo.
  • cache_params (Cache, opcional) — Se for passado, o modelo usa o estado anterior em todos os blocos (o que resultará na saída para os input_ids fornecidos, como se o modelo adicionasse state_input_ids + input_ids como contexto).
  • labels (torch.LongTensor de forma (batch_size, sequence_length), opcional) — Rótulos para modelagem de linguagem. Observe que os rótulos são deslocados dentro do modelo, ou seja, você pode definir labels = input_ids Os índices são selecionados em [-100, 0, ..., config.vocab_size] Todos os rótulos definidos como -100 são ignorados (mascarados); a perda é calculada apenas para rótulos em [0, ..., config.vocab_size]
  • output_hidden_states (bool, opcional) — Indica se as hidden states de todas as camadas devem ser retornadas. Consulte hidden_states em tensors retornados para mais detalhes.
  • return_dict (bool, opcional) — Indica se deve retornar um ModelOutput em vez de uma tupla simples.
  • use_cache (bool, opcional) — Se definido como True, o cache_params é retornado e pode ser usado para gerar rapidamente os próximos logits.
  • attention_mask (torch.Tensor de forma (batch_size, sequence_length), opcional) — Máscara para evitar a execução de atenção nos índices de tokens de preenchimento. Valores de máscara selecionados em [0, 1]:
    • 1 para tokens que não estão mascarados,
    • 0 para tokens que estão mascarados.

    O que são máscaras de atenção?

  • logits_to_keep (Union[int, torch.Tensor], opcional, padrão 0) — Se for um int, calcula os logits para os últimos logits_to_keep tokens. Se for 0, calcula os logits para todos os input_ids (caso especial). Apenas os logits do último token são necessários para a geração, e calculá-los apenas para esse token pode economizar memória, o que se torna bastante significativo para sequências longas ou tamanho de vocabulário grande. Se for um torch.Tensor, deve ser 1D, correspondendo aos índices a serem mantidos na dimensão de comprimento da sequência. Isso é útil ao usar o formato de tensor empacotado (uma única dimensão para lote e comprimento da sequência).

Retorna

Mamba2CausalLMOutput ou tuple(torch.Tensor)

Um Mamba2CausalLMOutput ou uma tupla de torch.Tensor (se return_dict=False for passado ou quando config.return_dict=False) contendo vários elementos, dependendo da configuração (Mamba2Config) e das entradas.

O método forward de Mamba2ForCausalLM substitui o método especial __call__.

Embora a receita para a passagem direta precise ser definida dentro desta função, deve-se chamar a instância de Module em seguida, em vez desta, pois a primeira cuida da execução das etapas de pré e pós-processamento, enquanto a segunda as ignora silenciosamente.

  • loss (torch.Tensor de forma (1,), opcional, retornado quando labels é fornecido) — Perda de modelagem de linguagem (para previsão do próximo token).

  • logits (torch.Tensor de forma (batch_size, sequence_length, config.vocab_size)) — Pontuações de previsão da cabeça de modelagem de linguagem (pontuações para cada token do vocabulário antes do SoftMax).

  • cache_params (~cache_utils.Cache, opcional) — O estado do modelo no último passo de tempo. Pode ser usado em um método forward com o próximo input_ids para evitar fornecer os input_ids antigos.

    Inclui tanto as matrizes de estado do modelo de espaço de estados após o scan seletivo quanto os estados convolucionais

  • hidden_states (tuple[torch.Tensor], opcional, retornado quando output_hidden_states=True é passado ou quando config.output_hidden_states=True) — Tupla de torch.Tensor (um para a saída dos embeddings, se o modelo tiver uma camada de embedding, + um para a saída de cada camada) de forma (batch_size, sequence_length, hidden_size).

    Estados ocultos do modelo na saída de cada camada, além das saídas opcionais do embedding inicial.

Exemplo:

Update on GitHub

←Mamba MarianMT→