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

Mô hình này được công bố trên HF papers vào ngày 31-05-2024 và được đóng góp cho Hugging Face Transformers vào ngày 06-08-2024.

Mamba 2

Mamba 2 dựa trên khung state space duality (SSD), kết nối các mô hình state space có cấu trúc (SSM) và các biến thể của attention. Nó sử dụng thuật toán SSD hiệu quả hơn, nhanh hơn 2-8 lần so với Mamba, và sửa đổi kiến trúc để cho phép tensor parallelism cũng như cấu trúc đầu grouped-value attention (GVA).

Bạn có thể tìm thấy tất cả các checkpoint gốc của Mamba 2 dưới tổ chức State Space Models, nhưng các ví dụ được hiển thị bên dưới sử dụng mistralai/Mamba-Codestral-7B-v0.1 vì một triển khai trên Hugging Face chưa được hỗ trợ cho các checkpoint gốc.

Các kiến trúc khác dựa trên Mamba 2 bao gồm Bamba, FalconH1 và Zamba2.

Mô hình này được đóng góp bởi ArthurZ. Nhấp vào các mô hình Mamba ở thanh bên phải để xem thêm các ví dụ về cách áp dụng Mamba cho các tác vụ ngôn ngữ khác nhau.

Ví dụ dưới đây minh họa cách tạo văn bản bằng Pipeline, AutoModel và từ dòng lệnh.

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

Quantization giảm gánh nặng bộ nhớ của các mô hình lớn bằng cách biểu diễn trọng số ở độ chính xác thấp hơn. Tham khảo tổng quan về Quantization để biết thêm các backend quantization khả dụng.

Ví dụ bên dưới sử dụng torchao để chỉ lượng tử hóa trọng số thành số nguyên 4-bit.

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

Ghi chú

  • Codestral Mamba có groups=8, tương tự như số lượng kv heads trong một mô hình dựa trên attention.

  • Codestral Mamba có hai đường truyền tiến (forward pass) khác nhau, torch_forward hoặc cuda_kernels_forward, và kết quả của chúng được kỳ vọng sẽ khác nhau đôi chút.

    • torch_forward không biên dịch nhanh hơn cuda_kernels_forward từ 3 đến 4 lần.
    • cuda_kernels_forward sử dụng các CUDA kernel gốc nếu chúng có sẵn trong môi trường của bạn. Nó chạy chậm hơn trong giai đoạn prefill vì yêu cầu một “lần chạy làm nóng” do chi phí CPU cao hơn (xem những bình luận này để biết thêm chi tiết).
  • Mô hình này không có positional embeddings, nhưng có attention_mask và một logic cụ thể để mask các hidden states ở hai vị trí trong trường hợp batched generation (xem comment này để biết thêm chi tiết). Điều này (cùng với việc thêm các Mamba 2 kernels được viết lại) dẫn đến một sự khác biệt nhỏ giữa batched generation và cached generation.

  • Thuật toán SSM phụ thuộc nhiều vào các phép co tensor, vốn có các tương đương với phép nhân ma trận nhưng thứ tự thực hiện phép toán hơi khác. Điều này khiến sự khác biệt trở nên lớn hơn ở các độ chính xác thấp hơn.

  • Padding chỉ được kiểm tra với left padding. Right padding cập nhật cache với các vị trí pad và cho kết quả không đáng tin cậy. Đặt tokenizer.padding_side="left" cho batched generation. Xem Padding side.

  • Ví dụ dưới đây minh họa cách tinh chỉnh Mamba 2 bằng 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

lớp 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 )

Tham số

  • num_heads (int, tùy chọn, mặc định là 128) — Số lượng đầu chú ý cho mỗi lớp chú ý trong bộ giải mã Transformer.
  • head_dim (int, tùy chọn, mặc định là 64) — Kích thước chiều của đầu chú ý. Nếu là None, giá trị mặc định sẽ là hidden_size // num_attention_heads
  • vocab_size (int, tuỳ chọn, mặc định là 32768) — Kích thước từ vựng của mô hình. Xác định số lượng token khác nhau có thể được biểu diễn bởi input_ids.
  • hidden_size (int, tùy chọn, mặc định là 4096) — Kích thước của các biểu diễn ẩn.
  • state_size (int, tuỳ chọn, mặc định là 128) — Kích thước của trạng thái SSM (chiều của trạng thái ẩn) trong các lớp Mamba.
  • num_hidden_layers (int, tùy chọn, mặc định là 64) — Số lượng lớp ẩn trong bộ giải mã Transformer.
  • layer_norm_epsilon (float, tùy chọn, mặc định là 1e-05) — Giá trị epsilon được sử dụng trong các lớp chuẩn hóa lớp (layer normalization).
  • pad_token_id (int, tùy chọn, mặc định là 1) — Token id được sử dụng để đệm (padding) trong từ vựng.
  • bos_token_id (int, tùy chọn, mặc định là 0) — Token id được sử dụng cho đầu luồng trong bảng từ vựng.
  • eos_token_id (Union[int, list[int]], tùy chọn, mặc định là 2) — Token id được sử dụng cho end-of-stream trong bảng từ vựng.
  • expand (int, tuỳ chọn, mặc định là 2) — Hệ số mở rộng được sử dụng để xác định kích thước trung gian.
  • conv_kernel (int, tuỳ chọn, mặc định là 4) — Kích thước của nhân tích chập.
  • n_groups (int, tùy chọn, mặc định là 8) — Số lượng nhóm cho các ma trận tiến hóa của mamba 2.
  • use_bias (bool, tùy chọn, mặc định là False) — Có sử dụng bias trong [“in_proj”, “out_proj”] của khối mixer hay không
  • use_conv_bias (bool, tùy chọn, mặc định là True) — Có sử dụng bias trong lớp convolution của khối mixer hay không.
  • hidden_act (str, tùy chọn, mặc định là silu) — Hàm kích hoạt phi tuyến (hàm hoặc chuỗi) trong bộ giải mã. Ví dụ, "gelu", "relu", "silu", v.v.
  • initializer_range (float, tùy chọn, mặc định là 0.1) — Độ lệch chuẩn của truncated_normal_initializer để khởi tạo tất cả các ma trận trọng số.
  • residual_in_fp32 (bool, tùy chọn, mặc định là True) — Xác định xem các phần dư (residuals) có nên ở dạng float32 hay không. Nếu đặt thành False, các phần dư sẽ giữ nguyên dtype giống như phần còn lại của mô hình
  • time_step_rank (Union[str, int], tuỳ chọn, mặc định là auto) — Rank của phép chiếu delta (time step). Có thể đặt là "auto" để thiết lập tự động.
  • time_step_min (float, tùy chọn, mặc định là 0.001) — time_step tối thiểu được dùng để giới hạn dt_proj.bias.
  • time_step_max (float, tùy chọn, mặc định là 0.1) — time_step tối đa được dùng để giới hạn dt_proj.bias.
  • time_step_floor (float, tùy chọn, mặc định là 0.0001) — Giá trị tối thiểu được phép cho delta của bước thời gian rời rạc sau kích hoạt softplus.
  • time_step_limit (Union[list[float], tuple[float, ...]], tùy chọn, mặc định là (0.0, inf)) — Khoảng giá trị bước thời gian được chấp nhận để kẹp (clamp).
  • rescale_prenorm_residual (bool, tùy chọn, mặc định là False) — Có nên điều chỉnh lại trọng số của out_proj khi khởi tạo hay không.
  • use_cache (bool, tùy chọn, mặc định là True) — Xác định xem mô hình có nên trả về các chú ý key/values cuối cùng hay không (không được sử dụng bởi tất cả các mô hình). Chỉ liên quan nếu config.is_decoder=True hoặc khi mô hình là mô hình sinh chỉ có decoder.
  • chunk_size (int, tùy chọn, mặc định là 256) — Kích thước của các khối cấu thành chuỗi.
  • tie_word_embeddings (bool, tùy chọn, mặc định là False) — Có nên gắn kết các trọng số nhúng theo ánh xạ tied_weights_keys của mô hình hay không.

Đây là lớp cấu hình để lưu trữ cấu hình của một Mamba2Model. Nó được sử dụng để khởi tạo một mô hình Mamba2 theo các đối số được chỉ định, xác định kiến trúc của mô hình. Khởi tạo một cấu hình với các giá trị mặc định sẽ tạo ra một cấu hình tương tự như cấu hình của state-spaces/mamba2-2.8b

Các đối tượng cấu hình kế thừa từ PreTrainedConfig và có thể được sử dụng để kiểm soát đầu ra của mô hình. Đọc tài liệu từ PreTrainedConfig để biết thêm thông tin.

Ví dụ:

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

lớp transformers.Mamba2Model

< >

( config: Mamba2Config )

Tham số

  • config (Mamba2Config) — Lớp cấu hình mô hình chứa tất cả các tham số của mô hình. Khởi tạo bằng tệp cấu hình không tải trọng số liên quan đến mô hình, chỉ tải cấu hình. Xem phương thức from_pretrained() để tải trọng số của mô hình.

Mô hình Mamba2 cơ bản, xuất ra các trạng thái ẩn thô mà không có bất kỳ đầu (head) cụ thể nào ở phía trên.

Mô hình này kế thừa từ PreTrainedModel. Xem tài liệu của lớp cha để biết các phương thức chung mà thư viện triển khai cho tất cả các mô hình của nó (chẳng hạn như tải xuống hoặc lưu trữ, thay đổi kích thước các nhúng đầu vào, cắt tỉa các đầu, v.v.)

Mô hình này cũng là một lớp con của torch.nn.Module trong PyTorch. Hãy sử dụng nó như một Module PyTorch thông thường và tham khảo tài liệu PyTorch cho mọi vấn đề liên quan đến cách sử dụng và hành vi chung.

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

Tham số

  • input_ids (torch.LongTensor có dạng (batch_size, sequence_length), tùy chọn) — Chỉ số của các token trong chuỗi đầu vào theo từ điển. Mặc định, padding sẽ bị bỏ qua.

    Các chỉ số có thể được lấy bằng cách sử dụng AutoTokenizer. Xem PreTrainedTokenizer.encode() và PreTrainedTokenizer.call() để biết thêm chi tiết.

    ID đầu vào là gì?

  • inputs_embeds (torch.LongTensor có hình dạng (batch_size, sequence_length, hidden_size), tùy chọn) — Tùy chọn, thay vì truyền input_ids, bạn có thể chọn truyền trực tiếp một biểu diễn đã nhúng. Điều này hữu ích nếu bạn muốn kiểm soát nhiều hơn cách chuyển đổi các chỉ số input_ids thành các vector tương ứng so với ma trận tra cứu nhúng nội bộ của mô hình.
  • cache_params (Cache, tùy chọn) — Nếu được truyền vào, mô hình sẽ sử dụng trạng thái trước đó trong tất cả các khối (điều này sẽ tạo ra đầu ra cho input_ids được cung cấp như thể mô hình đã thêm state_input_ids + input_ids làm ngữ cảnh).
  • use_cache (bool, tùy chọn) — Nếu được đặt thành True, cache_params sẽ được trả về và có thể được sử dụng để nhanh chóng tạo ra các logits tiếp theo.
  • output_hidden_states (bool, tùy chọn) — Có trả về trạng thái ẩn (hidden states) của tất cả các lớp hay không. Xem hidden_states trong phần tensor trả về để biết thêm chi tiết.
  • return_dict (bool, tùy chọn) — Có nên trả về một ModelOutput thay vì một tuple thông thường hay không.
  • attention_mask (torch.Tensor có hình dạng (batch_size, sequence_length), tùy chọn) — Mặt nạ để tránh thực hiện chú ý trên các chỉ số token đệm. Các giá trị mặt nạ được chọn trong [0, 1]:
    • 1 cho các token không bị che,
    • 0 cho các token bị che.

    Mặt nạ chú ý là gì?

Trả về

Mamba2Output hoặc tuple(torch.FloatTensor)

Một Mamba2Output hoặc một tuple của torch.FloatTensor (nếu return_dict=False được truyền hoặc khi config.return_dict=False) bao gồm các thành phần khác nhau tùy thuộc vào cấu hình (Mamba2Config) và đầu vào.

Phương thức forward của Mamba2Model ghi đè lên phương thức đặc biệt __call__.

Mặc dù công thức cho bước truyền tiến cần được định nghĩa bên trong hàm này, người dùng nên gọi instance của Module sau đó thay vì gọi hàm này, vì instance Module sẽ đảm nhận việc thực hiện các bước tiền xử lý và hậu xử lý, trong khi hàm này sẽ lặng lẽ bỏ qua chúng.

  • last_hidden_state (torch.FloatTensor có hình dạng (batch_size, sequence_length, hidden_size), tùy chọn) — Chuỗi các hidden-state tại đầu ra của lớp cuối cùng của mô hình.

  • cache_params (~cache_utils.Cache, tùy chọn) — Trạng thái của mô hình tại bước thời gian cuối cùng. Có thể được sử dụng trong phương thức forward với input_ids tiếp theo để tránh việc cung cấp lại input_ids cũ.

    Bao gồm cả các ma trận trạng thái của mô hình không gian trạng thái sau khi quét chọn lọc, và các trạng thái của lớp tích chập

  • hidden_states (tuple[torch.FloatTensor], tuỳ chọn, được trả về khi output_hidden_states=True được truyền vào hoặc khi config.output_hidden_states=True) — Bộ tuple các torch.FloatTensor (một cho đầu ra của embeddings, nếu mô hình có lớp embedding, + một cho đầu ra của mỗi lớp) có hình dạng (batch_size, sequence_length, hidden_size).

    Trạng thái ẩn của mô hình tại đầu ra của mỗi lớp cùng với các đầu ra embedding ban đầu tùy chọn.

Mamba2LMHeadModel

lớp transformers.Mamba2ForCausalLM

< >

( config: Mamba2Config )

Tham số

  • config (Mamba2Config) — Lớp cấu hình mô hình chứa tất cả các tham số của mô hình. Khởi tạo bằng tệp cấu hình không tải trọng số liên quan đến mô hình, chỉ tải cấu hình. Xem phương thức from_pretrained() để tải trọng số của mô hình.

Model transformer MAMBA2 với đầu mô hình hóa ngôn ngữ ở trên (tầng tuyến tính với trọng số không liên kết với các nhúng đầu vào).

Mô hình này kế thừa từ PreTrainedModel. Xem tài liệu của lớp cha để biết các phương thức chung mà thư viện triển khai cho tất cả các mô hình của nó (chẳng hạn như tải xuống hoặc lưu trữ, thay đổi kích thước các nhúng đầu vào, cắt tỉa các đầu, v.v.)

Mô hình này cũng là một lớp con của torch.nn.Module trong PyTorch. Hãy sử dụng nó như một Module PyTorch thông thường và tham khảo tài liệu PyTorch cho mọi vấn đề liên quan đến cách sử dụng và hành vi chung.

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.FloatTensor] = 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.FloatTensor)

Tham số

  • input_ids (torch.LongTensor có dạng (batch_size, sequence_length), tùy chọn) — Chỉ số của các token trong chuỗi đầu vào theo từ điển. Mặc định, padding sẽ bị bỏ qua.

    Các chỉ số có thể được lấy bằng cách sử dụng AutoTokenizer. Xem PreTrainedTokenizer.encode() và PreTrainedTokenizer.call() để biết thêm chi tiết.

    ID đầu vào là gì?

  • inputs_embeds (torch.FloatTensor có hình dạng (batch_size, sequence_length, hidden_size), tùy chọn) — Tùy chọn, thay vì truyền input_ids, bạn có thể chọn truyền trực tiếp một biểu diễn nhúng. Điều này hữu ích nếu bạn muốn kiểm soát nhiều hơn cách chuyển đổi các chỉ số input_ids thành các vector liên quan so với ma trận tra cứu nhúng nội bộ của mô hình.
  • cache_params (Cache, tùy chọn) — Nếu được truyền vào, mô hình sẽ sử dụng trạng thái trước đó trong tất cả các khối (điều này sẽ tạo ra đầu ra cho input_ids được cung cấp như thể mô hình đã thêm state_input_ids + input_ids làm ngữ cảnh).
  • labels (torch.LongTensor có shape (batch_size, sequence_length), tùy chọn) — Nhãn cho mô hình hóa ngôn ngữ. Lưu ý rằng các nhãn được dịch chuyển bên trong mô hình, tức là bạn có thể đặt labels = input_ids Các chỉ số được chọn trong [-100, 0, ..., config.vocab_size] Tất cả các nhãn được đặt thành -100 sẽ bị bỏ qua (che), loss chỉ được tính cho các nhãn trong [0, ..., config.vocab_size]
  • output_hidden_states (bool, tùy chọn) — Có trả về trạng thái ẩn (hidden states) của tất cả các lớp hay không. Xem hidden_states trong phần tensor trả về để biết thêm chi tiết.
  • return_dict (bool, tùy chọn) — Có nên trả về một ModelOutput thay vì một tuple thông thường hay không.
  • use_cache (bool, tùy chọn) — Nếu được đặt thành True, cache_params sẽ được trả về và có thể được sử dụng để nhanh chóng tạo ra các logits tiếp theo.
  • attention_mask (torch.Tensor có hình dạng (batch_size, sequence_length), tùy chọn) — Mặt nạ để tránh thực hiện chú ý trên các chỉ số token đệm. Các giá trị mặt nạ được chọn trong [0, 1]:
    • 1 cho các token không bị che,
    • 0 cho các token bị che.

    Mặt nạ chú ý là gì?

  • logits_to_keep (Union[int, torch.Tensor], tùy chọn, mặc định là 0) — Nếu là int, tính logits cho logits_to_keep token cuối. Nếu là 0, tính logits cho tất cả input_ids (trường hợp đặc biệt). Chỉ cần logits của token cuối cho quá trình sinh, và việc chỉ tính cho token đó có thể tiết kiệm bộ nhớ, điều này trở nên khá đáng kể với các chuỗi dài hoặc kích thước từ vựng lớn. Nếu là torch.Tensor, phải là 1 chiều tương ứng với các chỉ số cần giữ trong chiều độ dài chuỗi. Điều này hữu ích khi sử dụng định dạng tensor đóng gói (một chiều duy nhất cho batch và độ dài chuỗi).

Trả về

Mamba2CausalLMOutput hoặc tuple(torch.FloatTensor)

Một Mamba2CausalLMOutput hoặc một tuple của torch.FloatTensor (nếu return_dict=False được truyền hoặc khi config.return_dict=False) bao gồm các thành phần khác nhau tùy thuộc vào cấu hình (Mamba2Config) và đầu vào.

Phương thức forward của Mamba2ForCausalLM ghi đè lên phương thức đặc biệt __call__.

Mặc dù công thức cho bước truyền tiến cần được định nghĩa bên trong hàm này, người dùng nên gọi instance của Module sau đó thay vì gọi hàm này, vì instance Module sẽ đảm nhận việc thực hiện các bước tiền xử lý và hậu xử lý, trong khi hàm này sẽ lặng lẽ bỏ qua chúng.

  • loss (torch.FloatTensor có shape (1,), optional, được trả về khi labels được cung cấp) — Loss của language modeling (cho next-token prediction).

  • logits (torch.FloatTensor có hình dạng (batch_size, sequence_length, config.vocab_size)) — Điểm dự đoán của đầu mô hình hóa ngôn ngữ (điểm cho mỗi token từ vựng trước SoftMax).

  • cache_params (~cache_utils.Cache, tùy chọn) — Trạng thái của mô hình tại bước thời gian cuối cùng. Có thể được sử dụng trong phương thức forward với input_ids tiếp theo để tránh việc cung cấp lại input_ids cũ.

    Bao gồm cả các ma trận trạng thái của mô hình không gian trạng thái sau khi quét chọn lọc, và các trạng thái của lớp tích chập

  • hidden_states (tuple[torch.FloatTensor], tuỳ chọn, được trả về khi output_hidden_states=True được truyền vào hoặc khi config.output_hidden_states=True) — Bộ tuple các torch.FloatTensor (một cho đầu ra của embeddings, nếu mô hình có lớp embedding, + một cho đầu ra của mỗi lớp) có hình dạng (batch_size, sequence_length, hidden_size).

    Trạng thái ẩn của mô hình tại đầu ra của mỗi lớp cùng với các đầu ra embedding ban đầu tùy chọn.

Ví dụ:

Update on GitHub