Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2024 IBM and the HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Bamba model configuration"""16 17from ...configuration_utils import PretrainedConfig18from ...utils import logging19 20 21logger = logging.get_logger(__name__)22 23 24class BambaConfig(PretrainedConfig):25    r"""26    This is the configuration class to store the configuration of a [`BambaModel`]. It is used to instantiate a27    BambaModel model according to the specified arguments, defining the model architecture. Instantiating a configuration28    with defaults taken from [ibm-fms/Bamba-9.8b-2.2T-hf](https://huggingface.co/ibm-fms/Bamba-9.8b-2.2T-hf).29 30    The BambaModel is a hybrid [mamba2](https://github.com/state-spaces/mamba) architecture with SwiGLU.31    The checkpoints are  jointly trained by IBM, Princeton, and UIUC.32 33    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the34    documentation from [`PretrainedConfig`] for more information.35 36    Args:37        vocab_size (`int`, *optional*, defaults to 128000):38            Vocabulary size of the Bamba model. Defines the number of different tokens that can be represented by the39            `inputs_ids` passed when calling [`BambaModel`]40        tie_word_embeddings (`bool`, *optional*, defaults to `False`):41            Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the42            model has an output word embedding layer.43        hidden_size (`int`, *optional*, defaults to 4096):44            Dimension of the hidden representations.45        intermediate_size (`int`, *optional*, defaults to 14336):46            Dimension of the MLP representations.47        num_hidden_layers (`int`, *optional*, defaults to 32):48            Number of hidden layers in the Transformer encoder.49        num_attention_heads (`int`, *optional*, defaults to 32):50            Number of attention heads for each attention layer in the Transformer encoder.51        num_key_value_heads (`int`, *optional*, defaults to 8):52            This is the number of key_value heads that should be used to implement Grouped Query Attention. If53            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if54            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When55            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed56            by meanpooling all the original heads within that group. For more details, check out [this57            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `8`.58        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):59            The non-linear activation function (function or string) in the decoder.60        initializer_range (`float`, *optional*, defaults to 0.02):61            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.62        rms_norm_eps (`float`, *optional*, defaults to 1e-05):63            The epsilon used by the rms normalization layers.64        use_cache (`bool`, *optional*, defaults to `True`):65            Whether or not the model should return the last key/values attentions (not used by all models). Only66            relevant if `config.is_decoder=True`.67        num_logits_to_keep (`int` or `None`, *optional*, defaults to 1):68            Number of prompt logits to calculate during generation. If `None`, all logits will be calculated. If an69            integer value, only last `num_logits_to_keep` logits will be calculated. Default is 1 because only the70            logits of the last prompt token are needed for generation. For long sequences, the logits for the entire71            sequence may use a lot of memory so, setting `num_logits_to_keep=1` will reduce memory footprint72            significantly.73        pad_token_id (`int`, *optional*, defaults to 0):74            The id of the padding token.75        bos_token_id (`int`, *optional*, defaults to 1):76            The id of the "beginning-of-sequence" token.77        eos_token_id (`int`, *optional*, defaults to 2):78            The id of the "end-of-sequence" token.79        max_position_embeddings (`int`, *optional*, defaults to 262144):80            Max cached sequence length for the model81        attention_dropout (`float`, *optional*, defaults to 0.0):82            The dropout ratio for the attention probabilities.83        attn_layer_indices (`list`, *optional*):84            Specifies the layer indices that will have full attention. Must contain values at most num_hidden_layers.85        mamba_n_heads (`int`, *optional*, defaults to 128):86            The number of mamba heads used in the v2 implementation.87        mamba_d_head (`int`, *optional*, defaults to `"auto"`):88            Head embedding dimension size89        mamba_n_groups (`int`, *optional*, defaults to 1):90            The number of the mamba groups used in the v2 implementation.91        mamba_d_state (`int`, *optional*, defaults to 256):92            The dimension the mamba state space latents93        mamba_d_conv (`int`, *optional*, defaults to 4):94            The size of the mamba convolution kernel95        mamba_expand (`int`, *optional*, defaults to 2):96            Expanding factor (relative to hidden_size) used to determine the mamba intermediate size97        mamba_chunk_size (`int`, *optional*, defaults to 256):98            The chunks in which to break the sequence when doing prefill/training99        mamba_conv_bias (`bool`, *optional*, defaults to `True`):100            Flag indicating whether or not to use bias in the convolution layer of the mamba mixer block.101        mamba_proj_bias (`bool`, *optional*, defaults to `False`):102            Flag indicating whether or not to use bias in the input and output projections (["in_proj", "out_proj"]) of the mamba mixer block103        z_loss_coefficient (`float`, *optional*, defaults to 0.0):104            Coefficient for auxiliary z-loss used to control logit growth during training105 106    """107 108    model_type = "bamba"109    keys_to_ignore_at_inference = ["past_key_values"]110 111    def __init__(112        self,113        vocab_size=128000,114        tie_word_embeddings=False,115        hidden_size=4096,116        intermediate_size=14336,117        num_hidden_layers=32,118        num_attention_heads=32,119        num_key_value_heads=8,120        hidden_act="silu",121        initializer_range=0.02,122        rms_norm_eps=1e-5,123        use_cache=True,124        num_logits_to_keep=1,125        pad_token_id=0,126        bos_token_id=1,127        eos_token_id=2,128        max_position_embeddings=262144,129        attention_dropout=0.0,130        attn_layer_indices=None,131        mamba_n_heads=128,132        mamba_d_head="auto",133        mamba_n_groups=1,134        mamba_d_state=256,135        mamba_d_conv=4,136        mamba_expand=2,137        mamba_chunk_size=256,138        mamba_conv_bias=True,139        mamba_proj_bias=False,140        z_loss_coefficient=0.0,141        **kwargs,142    ):143        self.vocab_size = vocab_size144        self.tie_word_embeddings = tie_word_embeddings145        self.hidden_size = hidden_size146        self.intermediate_size = intermediate_size147        self.num_hidden_layers = num_hidden_layers148        self.num_attention_heads = num_attention_heads149        self.max_position_embeddings = max_position_embeddings150        self.attention_dropout = attention_dropout151        self.attention_bias = False152        self.mlp_bias = False153 154        # for backward compatibility155        if num_key_value_heads is None:156            num_key_value_heads = num_attention_heads157 158        self.num_key_value_heads = num_key_value_heads159        self.hidden_act = hidden_act160        self.initializer_range = initializer_range161        self.rms_norm_eps = rms_norm_eps162 163        self.use_cache = use_cache164        self.num_logits_to_keep = num_logits_to_keep165 166        self.attn_layer_indices = attn_layer_indices167        self.rope_theta = 10000.0168        self.rope_scaling = None169        self.partial_rotary_factor = 0.5170 171        mamba_intermediate = mamba_expand * hidden_size172 173        if mamba_intermediate % mamba_n_heads != 0:174            raise ValueError("mamba_n_heads must divide mamba_expand * hidden_size")175 176        # for the mamba_v2, must satisfy the following177        if mamba_d_head == "auto":178            mamba_d_head = mamba_intermediate // mamba_n_heads179 180        if mamba_d_head * mamba_n_heads != mamba_intermediate:181            raise ValueError("The dimensions for the Mamba head state do not match the model intermediate_size")182 183        self.mamba_n_heads = mamba_n_heads184        self.mamba_d_head = mamba_d_head185        self.mamba_n_groups = mamba_n_groups186        self.mamba_d_state = mamba_d_state187        self.mamba_d_conv = mamba_d_conv188        self.mamba_expand = mamba_expand189        self.mamba_chunk_size = mamba_chunk_size190        self.mamba_conv_bias = mamba_conv_bias191        self.mamba_proj_bias = mamba_proj_bias192        self.z_loss_coefficient = z_loss_coefficient193 194        super().__init__(195            pad_token_id=pad_token_id,196            bos_token_id=bos_token_id,197            eos_token_id=eos_token_id,198            tie_word_embeddings=tie_word_embeddings,199            **kwargs,200        )201 202    @property203    def layers_block_type(self):204        return [205            "attention" if (self.attn_layer_indices and i in self.attn_layer_indices) else "mamba"206            for i in range(self.num_hidden_layers)207        ]208 209 210__all__ = ["BambaConfig"]211