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

Эта модель была опубликована в HF papers 31 мая 2024 года и добавлена в Hugging Face Transformers 6 августа 2024 года.

Mamba 2

Mamba 2 основана на фреймворке двойственности пространств состояний (SSD), который связывает структурированные модели пространств состояний (SSM) и варианты механизма внимания. Она использует более эффективный алгоритм SSD, работающий в 2–8 раз быстрее, чем Mamba, а также модифицированную архитектуру, обеспечивающую тензорный параллелизм и структуру заголовков внимания с группированными значениями (GVA).

Все исходные контрольные точки Mamba 2 можно найти в организации State Space Models, однако в примерах ниже используется mistralai/Mamba-Codestral-7B-v0.1, так как реализация на Hugging Face пока не поддерживает исходные контрольные точки.

Среди других архитектур на основе Mamba 2 — Bamba, FalconH1 и Zamba2.

Эта модель была предоставлена ArthurZ. Нажмите на модели Mamba в боковой панели справа, чтобы увидеть больше примеров применения Mamba к различным языковым задачам.

Приведённый ниже пример показывает, как генерировать текст с помощью Pipeline, AutoModel и из командной строки.

Пайплайн
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")

Квантование снижает нагрузку на память при работе с большими моделями, представляя веса в пониженной точности. Подробнее о доступных бэкендах квантования см. в обзоре Квантование.

Приведённый ниже пример использует torchao для квантования только весов до 4-битных целых чисел.

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

Примечания

  • У Codestral Mamba есть groups=8, которые аналогичны количеству kv-голов в модели, основанной на механизме внимания.

  • У Codestral Mamba есть два различных прохода вперёд: torch_forward и cuda_kernels_forward, и их результаты, как ожидается, будут немного отличаться.

    • torch_forward без компиляции в 3–4 раза быстрее, чем cuda_kernels_forward.
    • cuda_kernels_forward использует исходные CUDA-ядра, если они доступны в вашей среде. На этапе prefill он работает медленнее, так как из-за более высокой нагрузки на CPU требуется «разогрев» (см. эти комментарии для получения подробностей).
  • В этой модели нет позиционных эмбеддингов, но есть attention_mask и специальная логика для маскирования скрытых состояний в двух местах при пакетной генерации (см. этот комментарий для подробностей). Это (а также добавление переосуществленных ядер Mamba 2) приводит к небольшому расхождению между результатами пакетной и кэшированной генерации.

  • Алгоритм SSM в значительной степени опирается на тензорные сокращения, которые имеют эквиваленты в виде операций matmul, однако порядок выполнения операций немного отличается. Это приводит к более заметным различиям при использовании меньших точностей.

  • Проверка выполняется только для выравнивания по левому краю. Выравнивание по правому краю приводит к обновлению кэша на позициях заполнителей и даёт ненадёжные результаты. Для пакетной генерации установите tokenizer.padding_side="left". См. Сторона выравнивания.

  • Приведённый ниже пример демонстрирует, как выполнить дообучение Mamba 2 с помощью 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

класс 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 )

Параметры

  • num_heads (int, необязательный, по умолчанию 128) — Количество головок внимания в каждом слое внимания декодера Transformer.
  • head_dim (int, необязательный, по умолчанию 64) — Размерность головы внимания. Если значение равно None, по умолчанию используется hidden_size // num_attention_heads
  • vocab_size (int, необязательный, по умолчанию 32768) — Размер словаря модели. Определяет количество различных токенов, которые могут быть представлены в input_ids.
  • hidden_size (int, необязательный, по умолчанию 4096) — Размерность скрытых представлений.
  • state_size (int, необязательный, по умолчанию 128) — Размер состояния SSM (размерность латентного состояния) в слоях Mamba.
  • num_hidden_layers (int, необязательный, по умолчанию 64) — Количество скрытых слоев в декодере Transformer.
  • layer_norm_epsilon (float, необязательный, по умолчанию 1e-05) — Значение epsilon, используемое в слоях нормализации.
  • pad_token_id (int, необязательный, по умолчанию 1) — Идентификатор токена, используемого для дополнения (padding) в словаре.
  • bos_token_id (int, необязательный, по умолчанию 0) — Идентификатор токена, используемого для обозначения начала потока в словаре.
  • eos_token_id (Union[int, list[int]], необязательный, по умолчанию 2) — Идентификатор токена, используемого для обозначения конца потока в словаре.
  • expand (int, необязательный, по умолчанию 2) — Коэффициент расширения, используемый для определения промежуточного размера.
  • conv_kernel (int, необязательный, по умолчанию 4) — Размер свёрточного ядра.
  • n_groups (int, необязательный, по умолчанию 8) — Количество групп для эволюционных матриц mamba 2.
  • use_bias (bool, необязательный, по умолчанию False) — Использовать ли смещение в [“in_proj”, “out_proj”] блока миксера
  • use_conv_bias (bool, необязательный, по умолчанию True) — Использовать ли смещение (bias) в свёрточном слое блока миксера.
  • hidden_act (str, необязательный, по умолчанию silu) — Нелинейная функция активации (функция или строка) в декодере. Например, "gelu", "relu", "silu" и т. д.
  • initializer_range (float, необязательный, по умолчанию 0.1) — Стандартное отклонение truncated_normal_initializer, используемое для инициализации всех матриц весов.
  • residual_in_fp32 (bool, необязательный, по умолчанию True) — Определяет, должны ли остаточные соединения храниться в формате float32. Если установлено значение False, остаточные соединения сохраняют тот же dtype, что и остальная часть модели
  • time_step_rank (Union[str, int], необязательный, по умолчанию auto) — Ранг проекции дельты (шага времени). Может быть "auto" для автоматического определения.
  • time_step_min (float, необязательный, по умолчанию 0.001) — Минимальное значение time_step, используемое для ограничения dt_proj.bias.
  • time_step_max (float, необязательный, по умолчанию 0.1) — Максимальное значение time_step, используемое для ограничения dt_proj.bias.
  • time_step_floor (float, необязательный, по умолчанию 0.0001) — Минимально допустимое значение дискретного шага времени delta после активации softplus.
  • time_step_limit (Union[list[float], tuple[float, ...]], необязательный, по умолчанию (0.0, inf)) — Допустимый диапазон значений шага времени для ограничения.
  • rescale_prenorm_residual (bool, необязательный, по умолчанию False) — Указывает, следует ли перемасштабировать веса out_proj при инициализации.
  • use_cache (bool, необязательный, по умолчанию True) — Указывает, должен ли модель возвращать последние ключи/значения внимания (не используется во всех моделях). Актуально только если config.is_decoder=True или если модель является генеративной моделью только с декодером.
  • chunk_size (int, необязательный, по умолчанию 256) — Размер блоков, из которых будет состоять последовательность.
  • tie_word_embeddings (bool, необязательный, по умолчанию False) — Определяет, следует ли связывать веса эмбеддингов согласно маппингу tied_weights_keys модели.

Это класс конфигурации для хранения настроек модели Mamba2Model. Он используется для создания модели Mamba2 в соответствии с указанными аргументами, определяя архитектуру модели. Создание конфигурации с параметрами по умолчанию даст конфигурацию, аналогичную конфигурации state-spaces/mamba2-2.8b

Объекты конфигурации наследуются от PreTrainedConfig и могут использоваться для управления выводами модели. Подробнее читайте в документации PreTrainedConfig.

Пример:

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

класс transformers.Mamba2Model

< >

( config: Mamba2Config )

Параметры

  • config (Mamba2Config) — Класс конфигурации модели, содержащий все параметры модели. Инициализация с помощью файла конфигурации не загружает веса, связанные с моделью, а только конфигурацию. Чтобы загрузить веса модели, используйте метод from_pretrained().

Модель Mamba2 без дополнительных компонентов, выводящая необработанные скрытые состояния без какого-либо специализированного слоя наверху.

Эта модель наследуется от PreTrainedModel. См. документацию родительского класса для общих методов, которые библиотека реализует для всех своих моделей (например, загрузка или сохранение, изменение размера входных эмбеддингов, обрезка голов и т. д.)

Эта модель также является подклассом torch.nn.Module в PyTorch. Используйте её как обычный модуль PyTorch и обращайтесь к документации PyTorch по всем вопросам, связанным с общим использованием и поведением.

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)

Параметры

  • input_ids (torch.LongTensor формы (batch_size, sequence_length), , необязательный) — Индексы токенов входной последовательности в словаре. По умолчанию токены отсечки (padding) игнорируются.

    Индексы можно получить с помощью AutoTokenizer. Подробности см. в PreTrainedTokenizer.encode() и PreTrainedTokenizer.call().

    Что такое входные идентификаторы?

  • inputs_embeds (torch.LongTensor формы (batch_size, sequence_length, hidden_size), необязательный) — При необходимости, вместо передачи input_ids можно напрямую передать встроенное представление. Это полезно, если требуется больший контроль над тем, как индексы input_ids преобразуются в соответствующие векторы, чем это обеспечивает внутренняя матрица поиска эмбеддингов модели.
  • cache_params (Cache, необязательный) — Если передан, модель использует предыдущее состояние во всех блоках (что даст вывод для input_ids, как если бы модель добавила state_input_ids + input_ids в контекст).
  • use_cache (bool, необязательный) — Если установлено в True, возвращается cache_params, которое можно использовать для быстрого генерирования следующих logits.
  • output_hidden_states (bool, необязательный) — Указывает, следует ли возвращать скрытые состояния всех слоев. Подробности см. в описании hidden_states среди возвращаемых тензоров.
  • return_dict (bool, необязательный) — Указывает, следует ли возвращать ModelOutput вместо обычного кортежа.
  • attention_mask (torch.Tensor формы (batch_size, sequence_length), , необязательный) — Маска для исключения токенов-заполнителей из внимания. Значения маски из диапазона [0, 1]:
    • 1 для токенов, которые не замаскированы,
    • 0 для токенов, которые замаскированы.

    Что такое маски внимания?

Возвращает

Mamba2Output или tuple(torch.Tensor)

Mamba2Output или кортеж из torch.Tensor (если передано return_dict=False или если config.return_dict=False), содержащий различные элементы в зависимости от конфигурации (Mamba2Config) и входных данных.

Метод forward класса Mamba2Model переопределяет специальный метод __call__.

Хотя рецепт для прямого прохода должен быть определён внутри этой функции, после этого следует вызывать экземпляр Module, а не эту функцию, поскольку первый вариант обеспечивает выполнение шагов предварительной и последующей обработки, тогда как второй молча их игнорирует.

  • last_hidden_state (torch.Tensor формы (batch_size, sequence_length, hidden_size), необязательный) — Последовательность скрытых состояний на выходе последнего слоя модели.

  • cache_params (~cache_utils.Cache, необязательный) — состояние модели на последнем шаге. Может использоваться в forward-методе вместе с новыми input_ids, чтобы избежать передачи старых input_ids.

    Включает в себя как матрицы состояния модели пространства состояний после селективного сканирования, так и состояния свёрток

  • hidden_states (tuple[torch.Tensor], необязательный, возвращается при передаче output_hidden_states=True или когда config.output_hidden_states=True) — Кортеж из torch.Tensor (один для вывода слоя встраиваний, если модель имеет слой встраиваний, + один для вывода каждого слоя) формы (batch_size, sequence_length, hidden_size).

    Скрытые состояния модели на выходе каждого слоя, а также необязательные выходные значения начального слоя встраивания.

Mamba2LMHeadModel

класс transformers.Mamba2ForCausalLM

< >

( config: Mamba2Config )

Параметры

  • config (Mamba2Config) — Класс конфигурации модели, содержащий все параметры модели. Инициализация с помощью файла конфигурации не загружает веса, связанные с моделью, а только конфигурацию. Чтобы загрузить веса модели, используйте метод from_pretrained().

Трансформер модели MAMBA2 с головной частью для моделирования языка (линейный слой с весами, не связанными с входными эмбеддингами).

Эта модель наследуется от PreTrainedModel. См. документацию родительского класса для общих методов, которые библиотека реализует для всех своих моделей (например, загрузка или сохранение, изменение размера входных эмбеддингов, обрезка голов и т. д.)

Эта модель также является подклассом torch.nn.Module в PyTorch. Используйте её как обычный модуль PyTorch и обращайтесь к документации PyTorch по всем вопросам, связанным с общим использованием и поведением.

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)

Параметры

  • input_ids (torch.LongTensor формы (batch_size, sequence_length), , необязательный) — Индексы токенов входной последовательности в словаре. По умолчанию токены отсечки (padding) игнорируются.

    Индексы можно получить с помощью AutoTokenizer. Подробности см. в PreTrainedTokenizer.encode() и PreTrainedTokenizer.call().

    Что такое входные идентификаторы?

  • inputs_embeds (torch.Tensor формы (batch_size, sequence_length, hidden_size), необязательный) — При необходимости, вместо передачи input_ids можно напрямую передать встроенное представление. Это полезно, если требуется больший контроль над тем, как индексы input_ids преобразуются в соответствующие векторы, чем это обеспечивает внутренняя матрица поиска встраиваний модели.
  • cache_params (Cache, необязательный) — Если передан, модель использует предыдущее состояние во всех блоках (что даст вывод для input_ids, как если бы модель добавила state_input_ids + input_ids в контекст).
  • labels (torch.LongTensor формы (batch_size, sequence_length), необязательный) — Метки для моделирования языка. Обратите внимание, что метки сдвигаются внутри модели, то есть можно задать labels = input_ids Индексы выбираются из диапазона [-100, 0, ..., config.vocab_size] Все метки, установленные в -100, игнорируются (маскируются); потеря вычисляется только для меток из диапазона [0, ..., config.vocab_size]
  • output_hidden_states (bool, необязательный) — Указывает, следует ли возвращать скрытые состояния всех слоев. Подробности см. в описании hidden_states среди возвращаемых тензоров.
  • return_dict (bool, необязательный) — Указывает, следует ли возвращать ModelOutput вместо обычного кортежа.
  • use_cache (bool, необязательный) — Если установлено в True, возвращается cache_params, которое можно использовать для быстрого генерирования следующих logits.
  • attention_mask (torch.Tensor формы (batch_size, sequence_length), , необязательный) — Маска для исключения токенов-заполнителей из внимания. Значения маски из диапазона [0, 1]:
    • 1 для токенов, которые не замаскированы,
    • 0 для токенов, которые замаскированы.

    Что такое маски внимания?

  • logits_to_keep (Union[int, torch.Tensor], необязательный, по умолчанию 0) — Если передано значение типа int, вычисляет logits для последних logits_to_keep токенов. Если указано 0, logits вычисляются для всех input_ids (особый случай). Для генерации нужны logits только последнего токена, и вычисление их только для этого токена позволяет сэкономить память, что становится весьма значимым для длинных последовательностей или большого размера словаря. Если передано значение типа torch.Tensor, оно должно быть одномерным и соответствовать индексам, которые нужно сохранить в измерении длины последовательности. Это полезно при использовании упакованного формата тензоров (один размер для batch и длины последовательности).

Возвращает

Mamba2CausalLMOutput или tuple(torch.Tensor)

Mamba2CausalLMOutput или кортеж из torch.Tensor (если передано return_dict=False или если config.return_dict=False), содержащий различные элементы в зависимости от конфигурации (Mamba2Config) и входных данных.

Метод forward класса Mamba2ForCausalLM переопределяет специальный метод __call__.

Хотя рецепт для прямого прохода должен быть определён внутри этой функции, после этого следует вызывать экземпляр Module, а не эту функцию, поскольку первый вариант обеспечивает выполнение шагов предварительной и последующей обработки, тогда как второй молча их игнорирует.

  • loss (torch.Tensor формы (1,), необязательный, возвращается, если задан labels) — Потери языковой модели (для предсказания следующего токена).

  • logits (torch.Tensor формы (batch_size, sequence_length, config.vocab_size)) — Оценки предсказаний языковой модели (оценки для каждого токена словаря до применения SoftMax).

  • cache_params (~cache_utils.Cache, необязательный) — состояние модели на последнем шаге. Может использоваться в forward-методе вместе с новыми input_ids, чтобы избежать передачи старых input_ids.

    Включает в себя как матрицы состояния модели пространства состояний после селективного сканирования, так и состояния свёрток

  • hidden_states (tuple[torch.Tensor], необязательный, возвращается при передаче output_hidden_states=True или когда config.output_hidden_states=True) — Кортеж из torch.Tensor (один для вывода слоя встраиваний, если модель имеет слой встраиваний, + один для вывода каждого слоя) формы (batch_size, sequence_length, hidden_size).

    Скрытые состояния модели на выходе каждого слоя, а также необязательные выходные значения начального слоя встраивания.

Пример:

Update on GitHub