Hugging Face's logo

Transformers documentation

Bamba

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

Эта модель была добавлена в Hugging Face Transformers 19.12.2024.

FlashAttention SDPA

Bamba

Bamba — это языковая модель с 9 млрд параметров, использующая только декодер и построенная на архитектуре Mamba-2. Предобучение выполняется в два этапа: сначала модель обучается на 2 трлн токенов из набора данных Dolma v1.7, а затем — на дополнительных 200 млрд токенов из FineWeb и Cosmopedia.

Все исходные контрольные точки Bamba можно найти в коллекции Bamba.

Эта модель была предоставлена ani300 и fabianlim.

Щёлкните по моделям Bamba в боковой панели справа, чтобы увидеть дополнительные примеры применения Bamba к различным задачам генерации текста.

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

Пайплайн
AutoModel
from transformers import pipeline


pipeline = pipeline(
    task="text-generation",
    model="ibm-ai-platform/Bamba-9B-v2",
    device=0
)
pipeline("Plants create energy through a process known as")

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

Приведённый ниже пример использует torchao для квантования только весов до int4.

from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig


quantization_config = TorchAoConfig("int4_weight_only", group_size=128)
tokenizer = AutoTokenizer.from_pretrained("ibm-ai-platform/Bamba-9B-v2")
model = AutoModelForCausalLM.from_pretrained(
   "ibm-ai-platform/Bamba-9B-v2",
   quantization_config=quantization_config,
   device_map="auto",
   attn_implementation="sdpa"
)

inputs = tokenizer("Plants create energy through a process known as", return_tensors="pt").to(model.device)
output = model.generate(**inputs)
print(tokenizer.decode(output[0], skip_special_tokens=True))

Примечания

  • Bamba поддерживает обучение без дополнения (padding), при котором различные примеры обучения конкатенируются, но при этом входные данные обрабатываются как отдельные батчи. Это может значительно ускорить инференс примерно в ~2 раза (в зависимости от модели и распределения данных) и снизить потребление памяти, если в данных присутствуют примеры разной длины, за счёт исключения лишних вычислений и накладных расходов на память, связанных с токенами дополнения.

    Обучение без padding требует пакетов flash-attn, mamba-ssm и causal-conv1d, а также передачи модели следующих аргументов в дополнение к input_ids и labels.

    • position_ids: torch.LongTensor: индекс позиции каждого токена в каждой последовательности.
    • seq_idx: torch.IntTensor: индекс каждой последовательности в батче.
    • Каждый из FlashAttentionKwargs
      • cu_seq_lens_q: torch.LongTensor: накопленные длины последовательностей всех запросов.
      • cu_seq_lens_k: torch.LongTensor: накопленные длины последовательностей всех ключей.
      • max_length_q: int: максимальная длина запроса в батче.
      • max_length_k: int: максимальная длина ключа в батче.

    Входные данные attention_mask не должны передаваться. DataCollatorWithFlattening программно генерирует набор дополнительных аргументов, описанных выше, с использованием return_seq_idx=True и return_flash_attn_kwargs=True. Дополнительную информацию см. в блоге Improving Hugging Face Training Efficiency Through Packing with Flash Attention.

    from transformers import DataCollatorWithFlattening

# Example of using padding-free training
data_collator = DataCollatorWithFlattening(
    tokenizer=tokenizer,
    return_seq_idx=True,
    return_flash_attn_kwargs=True
)

  • Используйте левое выравнивание (left padding) для пакетной генерации. См. Сторона выравнивания.

BambaConfig

класс transformers.BambaConfig

< >

( 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 = Nonevocab_size: int = 128000tie_word_embeddings: bool = Falsehidden_size: int = 4096intermediate_size: int = 14336num_hidden_layers: int = 32num_attention_heads: int = 32num_key_value_heads: int | None = 8hidden_act: str = 'silu'initializer_range: float = 0.02rms_norm_eps: float = 1e-05use_cache: bool = Truenum_logits_to_keep: int | None = 1pad_token_id: int | None = 0bos_token_id: int | None = 1eos_token_id: int | list[int] | None = 2max_position_embeddings: int = 262144attention_dropout: float | int | None = 0.0attn_layer_indices: list[int] | None = Nonemamba_n_heads: int | None = 128mamba_d_head: str | int | None = 'auto'mamba_n_groups: int | None = 1mamba_d_state: int | None = 256mamba_d_conv: int | None = 4mamba_expand: int | None = 2mamba_chunk_size: int | None = 256mamba_conv_bias: bool | None = Truemamba_proj_bias: bool | None = Falsetime_step_limit: list[float] | tuple[float, float] | None = (0.0, inf)z_loss_coefficient: float | None = 0.0rope_parameters: transformers.modeling_rope_utils.RopeParameters | dict | None = Noneattention_bias: bool = Falsemlp_bias: bool = False )

Параметры

  • vocab_size (int, необязательный, по умолчанию 128000) — Размер словаря модели. Определяет количество различных токенов, которые могут быть представлены в input_ids.
  • tie_word_embeddings (bool, необязательный, по умолчанию False) — Определяет, следует ли связывать веса эмбеддингов согласно маппингу tied_weights_keys модели.
  • hidden_size (int, необязательный, по умолчанию 4096) — Размерность скрытых представлений.
  • intermediate_size (int, необязательный, по умолчанию 14336) — Размерность представлений MLP.
  • num_hidden_layers (int, необязательный, по умолчанию 32) — Количество скрытых слоев в декодере Transformer.
  • num_attention_heads (int, необязательный, по умолчанию 32) — Количество головок внимания в каждом слое внимания декодера Transformer.
  • num_key_value_heads (int, необязательный, по умолчанию 8) — Это количество ключевых и value-голов, используемых для реализации Grouped Query Attention. Если num_key_value_heads=num_attention_heads, модель использует Multi Head Attention (MHA); если num_key_value_heads=1 — Multi Query Attention (MQA); в остальных случаях применяется GQA. При конвертации мультиголовного чекпоинта в GQA-чекпоинт каждая групповая ключевая и value-голова должна конструироваться путем усреднения (meanpooling) всех исходных голов в этой группе. Дополнительные подробности — в этой статье. Если значение не указано, по умолчанию используется num_attention_heads.
  • hidden_act (str, необязательный, по умолчанию silu) — Нелинейная функция активации (функция или строка) в декодере. Например, "gelu", "relu", "silu" и т. д.
  • initializer_range (float, необязательный, по умолчанию 0.02) — Стандартное отклонение truncated_normal_initializer, используемое для инициализации всех матриц весов.
  • rms_norm_eps (float, необязательный, по умолчанию 1e-05) — Значение epsilon, используемое слоями нормализации rms.
  • use_cache (bool, необязательный, по умолчанию True) — Указывает, должен ли модель возвращать последние ключи/значения внимания (не используется во всех моделях). Актуально только если config.is_decoder=True или если модель является генеративной моделью только с декодером.
  • num_logits_to_keep (int or None, необязательный, по умолчанию 1) — Количество логитов промпта, которые вычисляются во время генерации. Если None, вычисляются все логиты. Если указано целое число, вычисляются только последние num_logits_to_keep логитов. По умолчанию установлено значение 1, так как для генерации нужны только логиты последнего токена промпта. Для длинных последовательностей логиты для всей последовательности могут занимать много памяти, поэтому установка num_logits_to_keep=1 значительно уменьшает потребление памяти.
  • pad_token_id (int, необязательный, по умолчанию 0) — Идентификатор токена, используемого для дополнения (padding) в словаре.
  • bos_token_id (int, необязательный, по умолчанию 1) — Идентификатор токена, используемого для обозначения начала потока в словаре.
  • eos_token_id (Union[int, list[int]], необязательный, по умолчанию 2) — Идентификатор токена, используемого для обозначения конца потока в словаре.
  • max_position_embeddings (int, необязательный, по умолчанию 262144) — Максимальная длина последовательности, с которой может использоваться эта модель.
  • attention_dropout (Union[float, int], необязательный, по умолчанию 0.0) — Коэффициент dropout для вероятностей внимания.
  • attn_layer_indices (list, необязательный) — Указывает индексы слоёв, в которых будет использоваться полное внимание. Должно содержать значения, не превышающие num_hidden_layers.
  • mamba_n_heads (int, необязательный, по умолчанию 128) — Количество mamba-голов, используемых в реализации v2.
  • mamba_d_head (Union[str, int], необязательный, по умолчанию auto) — Размер размерности встраивания головы
  • mamba_n_groups (int, необязательный, по умолчанию 1) — Количество групп mamba, используемых в реализации v2.
  • mamba_d_state (int, необязательный, по умолчанию 256) — Размер состояния SSM (размерность латентного состояния) в слоях Mamba.
  • mamba_d_conv (int, необязательный, по умолчанию 4) — Размер ядра свёртки mamba
  • mamba_expand (int, необязательный, по умолчанию 2) — Коэффициент расширения (относительно hidden_size), используемый для определения промежуточного размера mamba
  • mamba_chunk_size (int, необязательный, по умолчанию 256) — Размер блоков, на которые разбивается последовательность при выполнении prefill/обучения
  • mamba_conv_bias (bool, необязательный, по умолчанию True) — Флаг, указывающий, следует ли использовать смещение (bias) в свёрточном слое блока mamba mixer.
  • mamba_proj_bias (bool, необязательный, по умолчанию False) — Флаг, указывающий, следует ли использовать смещение (bias) во входных и выходных проекциях ([“in_proj”, “out_proj”]) блока mamba mixer
  • time_step_limit (Union[list[float], tuple[float, float]], необязательный, по умолчанию (0.0, inf)) — Допустимый диапазон значений шага времени для ограничения.
  • z_loss_coefficient (float, необязательный, по умолчанию 0.0) — Коэффициент для вспомогательной z-loss, используемой для контроля роста логитов во время обучения
  • rope_parameters (Union[~modeling_rope_utils.RopeParameters, dict], необязательный) — Словарь, содержащий параметры конфигурации для эмбеддингов RoPE. Словарь должен содержать значение для rope_theta и, при необходимости, параметры, используемые для масштабирования, если вы хотите использовать RoPE с более длинными max_position_embeddings.
  • attention_bias (bool, необязательный, по умолчанию False) — Использовать ли смещение (bias) в слоях проекции query, key, value и output при само-внимании.
  • mlp_bias (bool, необязательный, по умолчанию False) — Использовать ли смещение (bias) в слоях up_proj, down_proj и gate_proj внутри MLP-слоёв.

Модель BambaModel представляет собой гибридную архитектуру mamba2 с SwiGLU. Чекпоинты совместно обучены командами IBM, Princeton и UIUC.

BambaModel

класс transformers.BambaModel

< >

( config: BambaConfig )

Параметры

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

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

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

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

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Nonepast_key_values: transformers.cache_utils.Cache | None = Noneinputs_embeds: typing.Optional[torch.Tensor] = Noneuse_cache: bool | None = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → BaseModelOutputWithPast or tuple(torch.Tensor)

Параметры

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

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

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

  • attention_mask (torch.Tensor формы (batch_size, sequence_length), , необязательный) — Маска для исключения токенов-заполнителей из внимания. Значения маски из диапазона [0, 1]:
    • 1 для токенов, которые не замаскированы,
    • 0 для токенов, которые замаскированы.

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

  • position_ids (torch.LongTensor формы (batch_size, sequence_length), , необязательный) — Индексы позиций токенов каждого входного последовательности в позиционных эмбеддингах. Выбираются из диапазона [0, config.n_positions - 1].

    Что такое идентификаторы позиций?

  • past_key_values (~cache_utils.Cache, необязательный) — Предвычисленные скрытые состояния (ключи и значения в блоках само-внимания и кросс-внимания), которые можно использовать для ускорения последовательного декодирования. Обычно это past_key_values , возвращённые моделью на предыдущем этапе декодирования, когда use_cache=True или config.use_cache=True.

    В качестве входного параметра допускается только экземпляр Cache, см. наше руководство по кэшу KV. Если past_key_values не переданы, по умолчанию будет инициализирован DynamicCache.

    Модель будет выводить тот же формат кэша, который был передан на вход.

    Если используются past_key_values, ожидается, что пользователь передаст только необработанные input_ids (те, для которых в этой модели ещё не заданы состояния прошлых ключей и значений) формы (batch_size, unprocessed_length), а не все input_ids формы (batch_size, sequence_length).

  • inputs_embeds (torch.Tensor формы (batch_size, sequence_length, hidden_size), необязательный) — При необходимости, вместо передачи input_ids можно напрямую передать встроенное представление. Это полезно, если требуется больший контроль над тем, как индексы input_ids преобразуются в соответствующие векторы, чем это обеспечивает внутренняя матрица поиска встраиваний модели.
  • use_cache (bool, необязательный) — Если установлено в True, возвращаются состояния ключей и значений past_key_values, которые можно использовать для ускорения декодирования (см. past_key_values).

Возвращает

BaseModelOutputWithPast или tuple(torch.Tensor)

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

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

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

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

    Если используется past_key_values, выводится только последнее скрытое состояние последовательностей формы (batch_size, 1, hidden_size).

  • past_key_values (Cache, необязательный, возвращается при передаче use_cache=True или когда config.use_cache=True) — это экземпляр Cache. Подробнее см. в нашем руководстве по кэшу kv.

    Содержит предварительно вычисленные скрытые состояния (ключи и значения в блоках само-внимания и, при необходимости, если config.is_encoder_decoder=True, в блоках кросс-внимания), которые можно использовать (см. входной параметр past_key_values) для ускорения последовательного декодирования.

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

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

  • attentions (tuple(torch.Tensor), необязательный, возвращается при передаче output_attentions=True или при config.output_attentions=True) — Кортеж torch.Tensor (по одному на каждый слой) формы (batch_size, num_heads, sequence_length, sequence_length).

    Веса внимания после softmax в механизме внимания, используемые для вычисления взвешенного среднего в головах само-внимания.

BambaForCausalLM

класс transformers.BambaForCausalLM

< >

( config: BambaConfig )

Параметры

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

Модель Bamba для каузального моделирования языка.

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

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

forward

< >

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

Параметры

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

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

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

  • attention_mask (torch.Tensor формы (batch_size, sequence_length), , необязательный) — Маска для исключения токенов-заполнителей из внимания. Значения маски из диапазона [0, 1]:
    • 1 для токенов, которые не замаскированы,
    • 0 для токенов, которые замаскированы.

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

  • position_ids (torch.LongTensor формы (batch_size, sequence_length), , необязательный) — Индексы позиций токенов каждого входного последовательности в позиционных эмбеддингах. Выбираются из диапазона [0, config.n_positions - 1].

    Что такое идентификаторы позиций?

  • past_key_values (~cache_utils.Cache, необязательный) — Предвычисленные скрытые состояния (ключи и значения в блоках само-внимания и кросс-внимания), которые можно использовать для ускорения последовательного декодирования. Обычно это past_key_values , возвращённые моделью на предыдущем этапе декодирования, когда use_cache=True или config.use_cache=True.

    В качестве входного параметра допускается только экземпляр Cache, см. наше руководство по кэшу KV. Если past_key_values не переданы, по умолчанию будет инициализирован DynamicCache.

    Модель будет выводить тот же формат кэша, который был передан на вход.

    Если используются past_key_values, ожидается, что пользователь передаст только необработанные input_ids (те, для которых в этой модели ещё не заданы состояния прошлых ключей и значений) формы (batch_size, unprocessed_length), а не все input_ids формы (batch_size, sequence_length).

  • inputs_embeds (torch.Tensor формы (batch_size, sequence_length, hidden_size), необязательный) — При необходимости, вместо передачи input_ids можно напрямую передать встроенное представление. Это полезно, если требуется больший контроль над тем, как индексы input_ids преобразуются в соответствующие векторы, чем это обеспечивает внутренняя матрица поиска встраиваний модели.
  • labels (torch.LongTensor формы (batch_size, sequence_length), необязательный) — Метки для вычисления потерь маскированного языкового моделирования. Индексы должны находиться в диапазоне [0, ..., config.vocab_size] или быть равны -100 (см. docstring input_ids). Токены с индексами, установленными в -100, игнорируются (маскируются); потери вычисляются только для токенов с метками в диапазоне [0, ..., config.vocab_size].
  • use_cache (bool, необязательный) — Если установлено в True, возвращаются состояния ключей и значений past_key_values, которые можно использовать для ускорения декодирования (см. past_key_values).
  • logits_to_keep (Union[int, torch.Tensor], необязательный, по умолчанию 0) — Если передано значение типа int, вычисляет logits для последних logits_to_keep токенов. Если указано 0, logits вычисляются для всех input_ids (особый случай). Для генерации нужны logits только последнего токена, и вычисление их только для этого токена позволяет сэкономить память, что становится весьма значимым для длинных последовательностей или большого размера словаря. Если передано значение типа torch.Tensor, оно должно быть одномерным и соответствовать индексам, которые нужно сохранить в измерении длины последовательности. Это полезно при использовании упакованного формата тензоров (один размер для batch и длины последовательности).

Возвращает

CausalLMOutputWithPast или tuple(torch.Tensor)

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

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

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

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

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

  • past_key_values (Cache, необязательный, возвращается при передаче use_cache=True или когда config.use_cache=True) — это экземпляр Cache. Подробнее см. в нашем руководстве по кэшу kv.

    Содержит предварительно вычисленные скрытые состояния (ключи и значения в блоках само-внимания), которые можно использовать (см. вход past_key_values) для ускорения последовательного декодирования.

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

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

  • attentions (tuple(torch.Tensor), необязательный, возвращается при передаче output_attentions=True или при config.output_attentions=True) — Кортеж torch.Tensor (по одному на каждый слой) формы (batch_size, num_heads, sequence_length, sequence_length).

    Веса внимания после softmax в механизме внимания, используемые для вычисления взвешенного среднего в головах само-внимания.

Пример:

>>> from transformers import AutoTokenizer, BambaForCausalLM

>>> model = BambaForCausalLM.from_pretrained("...")
>>> tokenizer = AutoTokenizer.from_pretrained("...")

>>> prompt = "Hey, are you conscious? Can you talk to me?"
>>> inputs = tokenizer(prompt, return_tensors="pt")

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
Update on GitHub

←Arcee BART→