easydel.modules.mamba2.mamba2_configuration

easydel.modules.mamba2.mamba2_configuration#

class easydel.modules.mamba2.mamba2_configuration.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)[source]#

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)[source]#

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'#
Contents