Hugging Face's logo

Transformers documentation

Mamba 2

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Ce modèle a été publié dans HF papers le 2024-05-31 et contribué à Hugging Face Transformers le 2024-08-06.

Mamba 2

Mamba 2 repose sur le cadre de dualité des espaces d'état (SSD), qui relie les modèles d'espace d'état structurés (SSM) aux variantes de l'attention. Il utilise un algorithme SSD plus efficace, 2 à 8 fois plus rapide que Mamba, et modifie l'architecture pour permettre le parallélisme de tenseurs et une structure de tête d'attention à valeurs groupées (GVA).

Vous pouvez trouver tous les points de contrôle originaux de Mamba 2 sous l’organisation State Space Models, mais les exemples présentés ci-dessous utilisent mistralai/Mamba-Codestral-7B-v0.1 car une implémentation Hugging Face n’est pas encore prise en charge pour les points de contrôle originaux.

Les autres architectures basées sur Mamba 2 incluent Bamba, FalconH1 et Zamba2.

Ce modèle a été contribué par ArthurZ. Cliquez sur les modèles Mamba dans la barre latérale droite pour plus d’exemples sur l’application de Mamba à différentes tâches de langage.

L'exemple ci-dessous montre comment générer du texte avec Pipeline, AutoModel et depuis la ligne de commande.

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")

La quantification réduit la charge mémoire des grands modèles en représentant les poids avec une précision inférieure. Reportez-vous à l'aperçu de la Quantization pour plus d'informations sur les backends de quantification disponibles.

L'exemple ci-dessous utilise torchao pour quantiser uniquement les poids en entiers sur 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))

Notes

  • Codestral Mamba possède groups=8, qui sont similaires au nombre de têtes kv dans un modèle basé sur l'attention.

  • Codestral Mamba dispose de deux passes avant différentes, torch_forward ou cuda_kernels_forward, et leurs résultats sont censés être légèrement différents.

    • torch_forward sans compilation est 3 à 4 fois plus rapide que cuda_kernels_forward.
    • cuda_kernels_forward utilise les noyaux CUDA d’origine s’ils sont disponibles dans votre environnement. Il est plus lent lors du préremplissage car il nécessite une « exécution de mise en route » en raison de la surcharge CPU plus élevée (voir ces commentaires pour plus de détails).

  • Ce modèle ne comporte pas d'embeddings positionnels, mais il possède un attention_mask et une logique spécifique pour masquer les états cachés à deux endroits lors de la génération par lots (voir ce commentaire pour plus de détails). Cela (ainsi que l'ajout des noyaux Mamba 2 réimplémentés) entraîne une légère divergence entre la génération par lots et la génération mise en cache.

  • L'algorithme SSM repose fortement sur des contractions de tenseurs, qui ont des équivalents en matmul, mais l'ordre des opérations est légèrement différent. Cela accentue l'écart à des précisions plus faibles.

  • Le padding n'est testé qu'avec un padding à gauche. Le padding à droite met à jour le cache avec les positions de padding et donne des résultats peu fiables. Définissez tokenizer.padding_side="left" pour la génération par lots. Voir Côté de padding.

  • L'exemple ci-dessous montre comment effectuer le fine-tuning de Mamba 2 avec 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 )

Paramètres

  • num_heads (int, facultatif, valeur par défaut 128) — Nombre de têtes d'attention pour chaque couche d'attention du décodeur Transformer.
  • head_dim (int, facultatif, par défaut 64) — La dimension des têtes d'attention. Si None, elle est par défaut égale à hidden_size // num_attention_heads
  • vocab_size (int, facultatif, par défaut 32768) — Taille du vocabulaire du modèle. Définit le nombre de jetons différents qui peuvent être représentés par input_ids.
  • hidden_size (int, facultatif, par défaut 4096) — Dimension des représentations cachées.
  • state_size (int, facultatif, par défaut 128) — Taille de l'état SSM (dimension de l'état latent) dans les couches Mamba.
  • num_hidden_layers (int, facultatif, par défaut 64) — Nombre de couches cachées dans le décodeur Transformer.
  • layer_norm_epsilon (float, facultatif, valeur par défaut 1e-05) — L'épsilon à utiliser dans les couches de normalisation par couche.
  • pad_token_id (int, facultatif, par défaut 1) — Identifiant du jeton utilisé pour le bourrage dans le vocabulaire.
  • bos_token_id (int, facultatif, par défaut 0) — Identifiant du jeton utilisé pour le début de flux dans le vocabulaire.
  • eos_token_id (Union[int, list[int]], facultatif, par défaut 2) — Identifiant du jeton utilisé pour la fin de flux dans le vocabulaire.
  • expand (int, facultatif, par défaut 2) — Facteur d'expansion utilisé pour déterminer la taille intermédiaire.
  • conv_kernel (int, facultatif, par défaut 4) — La taille du noyau de convolution.
  • n_groups (int, facultatif, valeur par défaut 8) — Nombre de groupes pour les matrices d'évolution de mamba 2.
  • use_bias (bool, facultatif, par défaut False) — Indique s'il faut utiliser un biais dans ["in_proj", "out_proj"] du bloc mixer
  • use_conv_bias (bool, facultatif, par défaut True) — Indique s'il faut utiliser un biais dans la couche de convolution du bloc mixer.
  • hidden_act (str, facultatif, par défaut silu) — La fonction d'activation non linéaire (fonction ou chaîne de caractères) du décodeur. Par exemple, "gelu", "relu", "silu", etc.
  • initializer_range (float, facultatif, par défaut 0.1) — L'écart-type de la truncated_normal_initializer utilisé pour initialiser toutes les matrices de poids.
  • residual_in_fp32 (bool, facultatif, valeur par défaut True) — Indique si les résidus doivent être en float32. Si défini sur False, les résidus conservent le même dtype que le reste du modèle
  • time_step_rank (Union[str, int], facultatif, par défaut auto) — Rang de la projection du delta (pas de temps). Peut être "auto" pour le définir automatiquement.
  • time_step_min (float, facultatif, valeur par défaut 0.001) — time_step minimal utilisé pour borner dt_proj.bias.
  • time_step_max (float, facultatif, valeur par défaut 0.1) — time_step maximal utilisé pour borner dt_proj.bias.
  • time_step_floor (float, facultatif, par défaut 0.0001) — Valeur minimale autorisée pour l'incrément discret du temps après l'activation softplus.
  • time_step_limit (Union[list[float], tuple[float, ...]], facultatif, par défaut (0.0, inf)) — Plage de valeurs de pas de temps acceptée pour le clampage.
  • rescale_prenorm_residual (bool, facultatif, par défaut False) — Indique s'il faut ou non rééchelonner les poids de out_proj lors de l'initialisation.
  • use_cache (bool, facultatif, valeur par défaut True) — Indique si le modèle doit retourner les dernières attentions clés/valeurs (non utilisé par tous les modèles). N’est pertinent que si config.is_decoder=True ou si le modèle est un modèle génératif uniquement décodeur.
  • chunk_size (int, facultatif, valeur par défaut 256) — Taille des morceaux qui composeront la séquence.
  • tie_word_embeddings (bool, facultatif, valeur par défaut False) — Indique s'il faut lier les embeddings de poids selon la correspondance tied_weights_keys du modèle.

Ceci est la classe de configuration destinée à stocker la configuration d'un Mamba2Model. Elle est utilisée pour instancier un modèle Mamba2 en fonction des arguments spécifiés, définissant ainsi l'architecture du modèle. L'instanciation d'une configuration avec les valeurs par défaut produira une configuration similaire à celle de state-spaces/mamba2-2.8b

Les objets de configuration héritent de PreTrainedConfig et peuvent être utilisés pour contrôler les sorties du modèle. Consultez la documentation de PreTrainedConfig pour plus d’informations.

Exemple :

>>> 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 )

Paramètres

  • config (Mamba2Config) — Classe de configuration du modèle contenant tous les paramètres du modèle. L’initialisation à partir d’un fichier de configuration ne charge pas les poids associés au modèle, uniquement la configuration. Consultez la méthode from_pretrained() pour charger les poids du modèle.

Le modèle Mamba2 brut, qui produit des états cachés bruts sans aucune tête spécifique au sommet.

Ce modèle hérite de PreTrainedModel. Consultez la documentation de la classe parente pour les méthodes génériques que la bibliothèque implémente pour tous ses modèles (tels que le téléchargement ou l'enregistrement, le redimensionnement des embeddings d'entrée, la taille des têtes, etc.)

Ce modèle est également une sous-classe de torch.nn.Module de PyTorch. Utilisez-le comme un module PyTorch standard et consultez la documentation de PyTorch pour tout ce qui concerne l'utilisation et le comportement généraux.

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)

Paramètres

  • input_ids (torch.LongTensor de forme (batch_size, sequence_length), facultatif) — Indices des jetons de la séquence d'entrée dans le vocabulaire. Le bourrage (padding) sera ignoré par défaut.

    Les indices peuvent être obtenus à l'aide de AutoTokenizer. Consultez PreTrainedTokenizer.encode() et PreTrainedTokenizer.call() pour plus de détails.

    Qu'est-ce que les identifiants d'entrée ?

  • inputs_embeds (torch.LongTensor de forme (batch_size, sequence_length, hidden_size), facultatif) — Facultativement, au lieu de passer input_ids, vous pouvez choisir de passer directement une représentation intégrée. Cela est utile si vous souhaitez un meilleur contrôle sur la manière de convertir les indices input_ids en vecteurs associés que la matrice de recherche d’intégration interne du modèle.
  • cache_params (Cache, facultatif) — Si transmis, le modèle utilise l'état précédent dans tous les blocs (ce qui produira la sortie pour les input_ids fournis comme si le modèle ajoutait state_input_ids + input_ids comme contexte).
  • use_cache (bool, facultatif) — Si défini sur True, cache_params est renvoyé et peut être utilisé pour générer rapidement les logits suivants.
  • output_hidden_states (bool, facultatif) — Indique s'il faut retourner les états cachés de toutes les couches. Voir hidden_states sous les tenseurs retournés pour plus de détails.
  • return_dict (bool, facultatif) — Indique s'il faut retourner un ModelOutput plutôt qu'un simple tuple.
  • attention_mask (torch.Tensor de forme (batch_size, sequence_length), facultatif) — Masque pour éviter d'appliquer l'attention aux indices des jetons de bourrage. Les valeurs du masque sont choisies dans [0, 1]:
    • 1 pour les jetons qui ne sont pas masqués,
    • 0 pour les jetons masqués.

    Qu'est-ce que les masques d'attention ?

Retourne

Mamba2Output ou tuple(torch.Tensor)

Un Mamba2Output ou un tuple de torch.Tensor (si return_dict=False est passé ou lorsque config.return_dict=False) comprenant divers éléments en fonction de la configuration (Mamba2Config) et des entrées.

La méthode forward de Mamba2Model redéfinit la méthode spéciale __call__.

Bien que la recette de la passe avant doive être définie dans cette fonction, il faut ensuite appeler l'instance Module plutôt que celle-ci, car la première gère l'exécution des étapes de prétraitement et de post-traitement, tandis que la seconde les ignore silencieusement.

  • last_hidden_state (torch.Tensor de forme (batch_size, sequence_length, hidden_size), facultatif) — Séquence d'états cachés à la sortie de la dernière couche du modèle.

  • cache_params (~cache_utils.Cache, facultatif) — L'état du modèle à la dernière étape de temps. Peut être utilisé dans une méthode forward avec les input_ids suivants pour éviter de fournir les anciens input_ids.

    Comprend à la fois les matrices d'état du modèle d'espace d'état après le balayage sélectif et les états convolutifs

  • hidden_states (tuple[torch.Tensor], facultatif, renvoyé lorsque output_hidden_states=True est passé ou lorsque config.output_hidden_states=True) — Tuple de torch.Tensor (un pour la sortie des embeddings, si le modèle possède une couche d’embedding, + un pour la sortie de chaque couche) de forme (batch_size, sequence_length, hidden_size).

    États cachés du modèle à la sortie de chaque couche, ainsi que les sorties d'incorporation initiales facultatives.

Mamba2LMHeadModel

classe transformers.Mamba2ForCausalLM

< >

( config: Mamba2Config )

Paramètres

  • config (Mamba2Config) — Classe de configuration du modèle contenant tous les paramètres du modèle. L’initialisation à partir d’un fichier de configuration ne charge pas les poids associés au modèle, uniquement la configuration. Consultez la méthode from_pretrained() pour charger les poids du modèle.

Le transformateur de modèle MAMBA2 avec une tête de modélisation de langage au sommet (couche linéaire dont les poids ne sont pas liés aux embeddings d'entrée).

Ce modèle hérite de PreTrainedModel. Consultez la documentation de la classe parente pour les méthodes génériques que la bibliothèque implémente pour tous ses modèles (tels que le téléchargement ou l'enregistrement, le redimensionnement des embeddings d'entrée, la taille des têtes, etc.)

Ce modèle est également une sous-classe de torch.nn.Module de PyTorch. Utilisez-le comme un module PyTorch standard et consultez la documentation de PyTorch pour tout ce qui concerne l'utilisation et le comportement généraux.

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)

Paramètres

  • input_ids (torch.LongTensor de forme (batch_size, sequence_length), facultatif) — Indices des jetons de la séquence d'entrée dans le vocabulaire. Le bourrage (padding) sera ignoré par défaut.

    Les indices peuvent être obtenus à l'aide de AutoTokenizer. Consultez PreTrainedTokenizer.encode() et PreTrainedTokenizer.call() pour plus de détails.

    Qu'est-ce que les identifiants d'entrée ?

  • inputs_embeds (torch.Tensor de forme (batch_size, sequence_length, hidden_size), facultatif) — Facultativement, au lieu de passer input_ids, vous pouvez choisir de passer directement une représentation intégrée. Cela est utile si vous souhaitez un meilleur contrôle sur la manière de convertir les indices input_ids en vecteurs associés que la matrice de recherche d’intégration interne du modèle.
  • cache_params (Cache, facultatif) — Si transmis, le modèle utilise l'état précédent dans tous les blocs (ce qui produira la sortie pour les input_ids fournis comme si le modèle ajoutait state_input_ids + input_ids comme contexte).
  • labels (torch.LongTensor de forme (batch_size, sequence_length), facultatif) — Étiquettes pour le modélisation du langage. Notez que les étiquettes sont décalées à l'intérieur du modèle, c'est-à-dire que vous pouvez définir labels = input_ids Les indices sont sélectionnés dans [-100, 0, ..., config.vocab_size] Toutes les étiquettes définies sur -100 sont ignorées (masquées), la perte n'est calculée que pour les étiquettes dans [0, ..., config.vocab_size]
  • output_hidden_states (bool, facultatif) — Indique s'il faut retourner les états cachés de toutes les couches. Voir hidden_states sous les tenseurs retournés pour plus de détails.
  • return_dict (bool, facultatif) — Indique s'il faut retourner un ModelOutput plutôt qu'un simple tuple.
  • use_cache (bool, facultatif) — Si défini sur True, cache_params est renvoyé et peut être utilisé pour générer rapidement les logits suivants.
  • attention_mask (torch.Tensor de forme (batch_size, sequence_length), facultatif) — Masque pour éviter d'appliquer l'attention aux indices des jetons de bourrage. Les valeurs du masque sont choisies dans [0, 1]:
    • 1 pour les jetons qui ne sont pas masqués,
    • 0 pour les jetons masqués.

    Qu'est-ce que les masques d'attention ?

  • logits_to_keep (Union[int, torch.Tensor], facultatif, valeur par défaut 0) — Si c'est un int, calcule les logits pour les logits_to_keep derniers jetons. Si 0, calcule les logits pour tous les input_ids (cas particulier). Seuls les logits du dernier jeton sont nécessaires pour la génération, et les calculer uniquement pour ce jeton peut économiser de la mémoire, ce qui devient assez significatif pour les séquences longues ou les tailles de vocabulaire importantes. Si c'est un torch.Tensor, il doit être 1D et correspondre aux indices à conserver dans la dimension de la longueur de la séquence. Cela est utile lors de l'utilisation du format de tenseurs empilés (une seule dimension pour le lot et la longueur de la séquence).

Retourne

Mamba2CausalLMOutput ou tuple(torch.Tensor)

Un Mamba2CausalLMOutput ou un tuple de torch.Tensor (si return_dict=False est passé ou lorsque config.return_dict=False) comprenant divers éléments en fonction de la configuration (Mamba2Config) et des entrées.

La méthode forward de Mamba2ForCausalLM redéfinit la méthode spéciale __call__.

Bien que la recette de la passe avant doive être définie dans cette fonction, il faut ensuite appeler l'instance Module plutôt que celle-ci, car la première gère l'exécution des étapes de prétraitement et de post-traitement, tandis que la seconde les ignore silencieusement.

  • loss (torch.Tensor de forme (1,), facultatif, renvoyé lorsque labels est fourni) — Perte de modélisation du langage (pour la prédiction du jeton suivant).

  • logits (torch.Tensor de forme (batch_size, sequence_length, config.vocab_size)) — Scores de prédiction de la tête de modélisation du langage (scores pour chaque jeton du vocabulaire avant SoftMax).

  • cache_params (~cache_utils.Cache, facultatif) — L'état du modèle à la dernière étape de temps. Peut être utilisé dans une méthode forward avec les input_ids suivants pour éviter de fournir les anciens input_ids.

    Comprend à la fois les matrices d'état du modèle d'espace d'état après le balayage sélectif et les états convolutifs

  • hidden_states (tuple[torch.Tensor], facultatif, renvoyé lorsque output_hidden_states=True est passé ou lorsque config.output_hidden_states=True) — Tuple de torch.Tensor (un pour la sortie des embeddings, si le modèle possède une couche d’embedding, + un pour la sortie de chaque couche) de forme (batch_size, sequence_length, hidden_size).

    États cachés du modèle à la sortie de chaque couche, ainsi que les sorties d'incorporation initiales facultatives.

Exemple :

Update on GitHub

←Mamba MarianMT→