| import mamba_ssm |
| from transformers import PretrainedConfig |
|
|
| mamba_config_defaults = mamba_ssm.models.config_mamba.MambaConfig() |
|
|
| class MambaConfig(PretrainedConfig): |
| model_type = "mamba" |
|
|
| def __init__( |
| self, |
| d_model: int = mamba_config_defaults.d_model, |
| fused_add_norm: bool = mamba_config_defaults.fused_add_norm, |
| n_layer: int = mamba_config_defaults.n_layer, |
| pad_vocab_size_multiple: int = mamba_config_defaults.pad_vocab_size_multiple, |
| residual_in_fp32: bool = mamba_config_defaults.residual_in_fp32, |
| rms_norm: bool = mamba_config_defaults.rms_norm, |
| ssm_cfg: dict = mamba_config_defaults.ssm_cfg, |
| vocab_size: int = mamba_config_defaults.vocab_size, |
| **kwargs, |
| ): |
| self.d_model = d_model |
| self.fused_add_norm = fused_add_norm |
| self.n_layer = n_layer |
| self.pad_vocab_size_multiple = pad_vocab_size_multiple |
| self.residual_in_fp32 = residual_in_fp32 |
| self.rms_norm = rms_norm |
| self.ssm_cfg = ssm_cfg |
| self.vocab_size = vocab_size |
|
|
| super().__init__(**kwargs) |
|
|
