easydel.modules.mamba2.__init__

class easydel.modules.mamba2.__init__.Mamba2Config(num_heads=128, head_dim=64, vocab_size=32768, hidden_size=4096, state_size=128, num_hidden_layers=64, layer_norm_epsilon=1e-05, pad_token_id=1, bos_token_id=0, eos_token_id=2, expand=2, conv_kernel=4, n_groups=8, use_bias=False, use_conv_bias=True, hidden_act='silu', initializer_range=0.1, residual_in_fp32=True, time_step_rank='auto', time_step_min=0.001, time_step_max=0.1, time_step_floor=0.0001, time_step_limit=(0.0, inf), rescale_prenorm_residual=False, use_cache=True, norm_before_gate=True, rms_norm=True, chunk_size=256, tie_word_embeddings=False, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, **kwargs)

Bases: EasyDeLBaseConfig

Configuration objects inherit from [EasyDeLBaseConfig] and can be used to control the model outputs. Read the documentation from [EasyDeLBaseConfig] for more information.

Parameters

• vocab_size (int, optional, defaults to 32768) – Vocabulary size of the Mamba2 model. Defines the number of different tokens that can be represented by the inputs_ids passed to the forward method.

• hidden_size (int, optional, defaults to 4096) – Dimensionality of the encoder layers and the pooler layer.

• state_size (int, optional, defaults to 128) – State size of the Mamba2 model.

• num_hidden_layers (int, optional, defaults to 64) – Number of hidden layers in the Mamba2 encoder.

• num_heads (int, optional, defaults to 128) – Number of attention heads for the grouped selective scan.

• head_dim (int, optional, defaults to 64) – Dimension of each attention head.

• n_groups (int, optional, defaults to 8) – Number of groups for the grouped selective scan.

• layer_norm_epsilon (float, optional, defaults to 1e-5) – The epsilon used by the layer normalization layers.

• pad_token_id (int, optional, defaults to 1) – The index of the padding token in the vocabulary.

• bos_token_id (int, optional, defaults to 0) – The id of the beginning-of-sequence token.

• eos_token_id (int, optional, defaults to 2) – The id of the end-of-sequence token.

• expand (int, optional, defaults to 2) – Expansion factor for the intermediate size.

• conv_kernel (int, optional, defaults to 4) – Kernel size of the convolution layer.

• use_bias (bool, optional, defaults to False) – Whether to use bias in the linear layers.

• use_conv_bias (bool, optional, defaults to True) – Whether to use bias in the convolution layer.

• hidden_act (str or function, optional, defaults to “silu”) – The non-linear activation function (function or string) to use in the encoder and pooler. If string, “gelu”, “relu”, “swish” and “gelu_new” are supported.

• initializer_range (float, optional, defaults to 0.1) – The standard deviation of the truncated_normal_initializer for initializing all weight matrices.

• residual_in_fp32 (bool, optional, defaults to True) – Whether to compute the residual connection in float32.

• time_step_rank (str or int, optional, defaults to “auto”) – The rank of the time step embedding. If set to “auto”, the rank is calculated as math.ceil(self.hidden_size / 16).

• time_step_min (float, optional, defaults to 0.001) – The minimum value for the time step embedding.

• time_step_max (float, optional, defaults to 0.1) – The maximum value for the time step embedding.

• time_step_floor (float, optional, defaults to 1e-4) – The floor value for the time step embedding.

• time_step_limit (tuple, optional, defaults to (0.0, float(“inf”))) – The minimum and maximum limits for the time step.

• rescale_prenorm_residual (bool, optional, defaults to False) – Whether to rescale the pre-norm residual.

• use_cache (bool, optional, defaults to True) – Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if config.is_decoder=True.

• norm_before_gate (bool, optional, defaults to True) – Whether to apply normalization before the gate activation.

• rms_norm (bool, optional, defaults to True) – Whether to use root mean square normalization.

• chunk_size (int, optional, defaults to 256) – Size of chunks for processing long sequences.

• tie_word_embeddings (bool, optional, defaults to False) – Whether to tie the word embedding weights with the output projection weights.

• gradient_checkpointing (str, optional, defaults to “nothing_saveable”) – The gradient checkpointing configuration.

get_partition_rules(*args, **kwargs)

Get the partition rules for distributing the Mamba2 model parameters across multiple devices.

These rules define how parameters should be partitioned when using techniques like Fully Sharded Data Parallelism (FSDP), Sharded Parallelism (SP), and Tensor Parallelism (TP). Each rule consists of a regex pattern matching parameter names and a corresponding PartitionSpec.

Returns

A tuple of tuples where each inner tuple contains:

• A regex pattern matching parameter names

• A PartitionSpec object specifying how to partition matching parameters

Return type

tuple

model_type: str = 'mamba2'

class easydel.modules.mamba2.__init__.Mamba2ForCausalLM(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

init_cache(batch_size: int, max_length: int, starts: int | None = None, shardings: dict | None = None, pad_token_id: int | None = None)

Initializes and returns a standard (non-paged) Key-Value cache.

This method first creates the necessary metadata using create_cache_metadata and then calls TransformerCache.init_cache to allocate and initialize the cache tensors based on the model’s configuration, dtype, sharding, quantization settings, and provided batch size and maximum length.

Parameters

• batch_size (int) – The batch size for the cache.

• max_length (int) – The maximum sequence length the cache needs to support.

• starts (int | None) – Optional starting positions for the cache sequences. If provided, influences the initial state. Defaults to None (usually 0).

• shardings (dict | None) – Optional dictionary specifying sharding configurations. (Note: This argument appears unused in the current implementation shown).

• pad_token_id (int | None) – The ID of the padding token. If None, it’s inferred.

Returns

An initialized standard TransformerCache object.

Return type

TransformerCache

prepare_inputs_for_generation(input_ids, inputs_embeds=None, cache_params: Optional[Mamba2Cache] = None, cache_position: Optional[Union[Array, ndarray, bool, number]] = None, attention_mask: Optional[Union[Array, ndarray, bool, number]] = None, **kwargs)

Sets up the initial inputs required for starting autoregressive generation.

This function initializes the Key-Value cache (past_key_values) using init_cache, calculates the initial position_ids based on the input attention_mask (or assumes a contiguous range if no mask is provided), and prepares an extended attention_mask suitable for caching. It ensures inputs are placed on the correct devices/shards.

Parameters

• input_ids (chex.Array) – The initial sequence of token IDs. Shape (batch_size, seq_length).

• max_length (int) – The maximum sequence length that the KV cache should support.

• pad_token_id (int) – The ID used for padding tokens. Used to calculate starts if not provided.

• starts (int | None) – Optional pre-calculated starting positions (number of leading pads). If None, calculated using compute_prefill_length.

• shardings (dict | None) – Optional sharding configuration passed to init_cache.

• attention_mask (tp.Optional[chex.Array]) – An optional mask indicating which tokens should be attended to. Shape (batch_size, seq_length).

• token_type_ids (tp.Optional[chex.Array]) – Optional segment IDs for models that use them.

Returns

A dictionary containing the prepared inputs, typically including:

• ”past_key_values”: The initialized KV cache.

• ”attention_mask”: The extended attention mask for generation.

• ”position_ids”: The calculated initial position IDs.

• ”token_type_ids”: (Optional) Prepared token type IDs.

This dictionary is then passed through prepare_inputs_for_call.

Return type

dict

update_inputs_for_generation(model_outputs, model_kwargs)

Updates the keyword arguments for the next generation step.

Specifically, it takes the past_key_values from the model_outputs of the current step and updates the model_kwargs with them. It also increments the position_ids by one for the next token prediction.

Parameters

• model_outputs – The output object from the model’s forward pass in the previous step (should contain a past_key_values attribute).

• model_kwargs (dict) – The dictionary of keyword arguments used for the model call. This dictionary will be modified in-place or a new one returned.

Returns

The updated model_kwargs dictionary ready for the next generation step.

Return type

dict

class easydel.modules.mamba2.__init__.Mamba2Model(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule