lgq12697 commited on
Commit
463ecb5
1 Parent(s): 452ddae

update model

Browse files
Files changed (3) hide show
  1. config.json +15 -1
  2. configuration_mamba.py +157 -0
  3. modeling_mamba.py +973 -0
config.json CHANGED
@@ -3,6 +3,12 @@
3
  "architectures": [
4
  "MambaForSequenceClassification"
5
  ],
 
 
 
 
 
 
6
  "bos_token_id": 0,
7
  "conv_kernel": 4,
8
  "d_inner": 1536,
@@ -12,8 +18,16 @@
12
  "fused_add_norm": true,
13
  "hidden_act": "silu",
14
  "hidden_size": 768,
 
 
 
 
15
  "initializer_range": 0.1,
16
  "intermediate_size": 1536,
 
 
 
 
17
  "layer_norm_epsilon": 1e-05,
18
  "model_type": "mamba",
19
  "n_layer": 24,
@@ -37,6 +51,6 @@
37
  "use_bias": false,
38
  "use_cache": false,
39
  "use_conv_bias": true,
40
- "use_mambapy": false,
41
  "vocab_size": 27
42
  }
 
3
  "architectures": [
4
  "MambaForSequenceClassification"
5
  ],
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_mamba.MambaConfig",
8
+ "AutoModel": "modeling_mamba.MambaModel",
9
+ "AutoModelForCausalLM": "modeling_mamba.MambaForCausalLM",
10
+ "AutoModelForSequenceClassification": "modeling_mamba.MambaForSequenceClassification"
11
+ },
12
  "bos_token_id": 0,
13
  "conv_kernel": 4,
14
  "d_inner": 1536,
 
18
  "fused_add_norm": true,
19
  "hidden_act": "silu",
20
  "hidden_size": 768,
21
+ "id2label": {
22
+ "0": "Not conserved",
23
+ "1": "Conserved"
24
+ },
25
  "initializer_range": 0.1,
26
  "intermediate_size": 1536,
27
+ "label2id": {
28
+ "Not conserved": 0,
29
+ "Conserved": 1
30
+ },
31
  "layer_norm_epsilon": 1e-05,
32
  "model_type": "mamba",
33
  "n_layer": 24,
 
51
  "use_bias": false,
52
  "use_cache": false,
53
  "use_conv_bias": true,
54
+ "use_mambapy": true,
55
  "vocab_size": 27
56
  }
configuration_mamba.py ADDED
@@ -0,0 +1,157 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2024 The HuggingFace Inc. team.
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 at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # 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 and
14
+ # limitations under the License.
15
+ """MAMBA configuration"""
16
+
17
+ import math
18
+
19
+ from transformers.configuration_utils import PretrainedConfig
20
+ from transformers.utils import logging
21
+
22
+
23
+ logger = logging.get_logger(__name__)
24
+
25
+
26
+ class MambaConfig(PretrainedConfig):
27
+ """
28
+ This is the configuration class to store the configuration of a [`MambaModel`]. It is used to instantiate a MAMBA
29
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
30
+ defaults will yield a similar configuration to that of the MAMBA
31
+ [state-spaces/mamba-2.8b](https://huggingface.co/state-spaces/mamba-2.8b) architecture.
32
+
33
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
34
+ documentation from [`PretrainedConfig`] for more information.
35
+
36
+
37
+ Args:
38
+ vocab_size (`int`, *optional*, defaults to 50280):
39
+ Vocabulary size of the MAMBA model. Defines the number of different tokens that can be represented by the
40
+ `inputs_ids` passed when calling [`MambaModel`].
41
+ hidden_size (`int`, *optional*, defaults to 768):
42
+ Dimensionality of the embeddings and hidden states.
43
+ state_size (`int`, *optional*, defaults to 16): shape of the state space latents.
44
+ num_hidden_layers (`int`, *optional*, defaults to 32):
45
+ Number of hidden layers in the model.
46
+ layer_norm_epsilon (`float`, *optional*, defaults to 1e-05):
47
+ The epsilon to use in the layer normalization layers.
48
+ pad_token_id (`int`, *optional*, defaults to 0):
49
+ Padding token id.
50
+ bos_token_id (`int`, *optional*, defaults to 0):
51
+ The id of the beginning of sentence token in the vocabulary.
52
+ eos_token_id (`int`, *optional*, defaults to 0):
53
+ The id of the end of sentence token in the vocabulary.
54
+ expand (`int`, *optional*, defaults to 2): Expanding factor used to determine the intermediate size.
55
+ conv_kernel (`int`, *optional*, defaults to 4): Size of the convolution kernel.
56
+ use_bias (`bool`, *optional*, defaults to `False`):
57
+ Whether or not to use bias in ["in_proj", "out_proj"] of the mixer block
58
+ use_conv_bias (`bool`, *optional*, defaults to `True`):
59
+ Whether or not to use bias in the convolution layer of the mixer block.
60
+ hidden_act (`str`, *optional*, defaults to `"silu"`):
61
+ The non-linear activation function (function or string) in the decoder.
62
+ initializer_range (`float`, *optional*, defaults to 0.1):
63
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
64
+ residual_in_fp32 (`bool`, *optional*, defaults to `True`):
65
+ Whether or not residuals should be in `float32`. If set to `False` residuals will keep the same `dtype` as the rest of the model
66
+ time_step_rank (`Union[int,str]`, *optional*, defaults to `"auto"`):
67
+ Rank of the discretization projection matrix. `"auto"` means that it will default to `math.ceil(self.hidden_size / 16)`
68
+ time_step_scale (`float`, *optional*, defaults to 1.0):
69
+ Scale used used to scale `dt_proj.bias`.
70
+ time_step_min (`float`, *optional*, defaults to 0.001):
71
+ Minimum `time_step` used to bound `dt_proj.bias`.
72
+ time_step_max (`float`, *optional*, defaults to 0.1):
73
+ Maximum `time_step` used to bound `dt_proj.bias`.
74
+ time_step_init_scheme (`float`, *optional*, defaults to `"random"`):
75
+ Init scheme used for `dt_proj.weight`. Should be one of `["random","uniform"]`
76
+ time_step_floor (`float`, *optional*, defaults to 0.0001):
77
+ Minimum clamping value of the `dt_proj.bias` layer initialization.
78
+ rescale_prenorm_residual (`bool`, *optional*, defaults to `False`):
79
+ Whether or not to rescale `out_proj` weights when initializing.
80
+ use_cache (`bool`, *optional*, defaults to `True`):
81
+ Whether or not the cache should be used.
82
+ use_mambapy (`bool`, *optional*, defaults to `False`):
83
+ Determines the fallback strategy during training if the CUDA-based official implementation of Mamba is not avaiable. If `True`, the mamba.py implementation is used. If `False`, the naive and slower implementation is used. Consider switching to the naive version if memory is limited.
84
+
85
+
86
+ Example:
87
+
88
+ ```python
89
+ >>> from transformers import MambaConfig, MambaModel
90
+
91
+ >>> # Initializing a Mamba configuration
92
+ >>> configuration = MambaConfig()
93
+
94
+ >>> # Initializing a model (with random weights) from the configuration
95
+ >>> model = MambaModel(configuration)
96
+
97
+ >>> # Accessing the model configuration
98
+ >>> configuration = model.config
99
+ ```"""
100
+
101
+ model_type = "mamba"
102
+
103
+ def __init__(
104
+ self,
105
+ vocab_size=50280,
106
+ hidden_size=768,
107
+ state_size=16,
108
+ num_hidden_layers=32,
109
+ layer_norm_epsilon=1e-5,
110
+ pad_token_id=0,
111
+ bos_token_id=0,
112
+ eos_token_id=0,
113
+ expand=2,
114
+ conv_kernel=4,
115
+ use_bias=False,
116
+ use_conv_bias=True,
117
+ hidden_act="silu",
118
+ initializer_range=0.1,
119
+ residual_in_fp32=True,
120
+ time_step_rank="auto",
121
+ time_step_scale=1.0,
122
+ time_step_min=0.001,
123
+ time_step_max=0.1,
124
+ time_step_init_scheme="random",
125
+ time_step_floor=1e-4,
126
+ rescale_prenorm_residual=False,
127
+ use_cache=True,
128
+ use_mambapy=False,
129
+ **kwargs,
130
+ ):
131
+ self.vocab_size = vocab_size
132
+ self.hidden_size = hidden_size
133
+ self.state_size = state_size
134
+ self.num_hidden_layers = num_hidden_layers
135
+ self.layer_norm_epsilon = layer_norm_epsilon
136
+ self.conv_kernel = conv_kernel
137
+ self.expand = expand
138
+ self.intermediate_size = int(expand * self.hidden_size)
139
+ self.bos_token_id = bos_token_id
140
+ self.eos_token_id = eos_token_id
141
+ self.pad_token_id = pad_token_id
142
+ self.use_bias = use_bias
143
+ self.use_conv_bias = use_conv_bias
144
+ self.hidden_act = hidden_act
145
+ self.initializer_range = initializer_range
146
+ self.time_step_rank = math.ceil(self.hidden_size / 16) if time_step_rank == "auto" else time_step_rank
147
+ self.time_step_scale = time_step_scale
148
+ self.time_step_min = time_step_min
149
+ self.time_step_max = time_step_max
150
+ self.time_step_init_scheme = time_step_init_scheme
151
+ self.time_step_floor = time_step_floor
152
+ self.rescale_prenorm_residual = rescale_prenorm_residual
153
+ self.residual_in_fp32 = residual_in_fp32
154
+ self.use_cache = use_cache
155
+ self.use_mambapy = use_mambapy
156
+
157
+ super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id, **kwargs)
modeling_mamba.py ADDED
@@ -0,0 +1,973 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2024 state-spaces/mamba org and HuggingFace Inc. team.
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 at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # 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 and
14
+ # limitations under the License.
15
+ """PyTorch MAMBA model."""
16
+
17
+ import math
18
+ from dataclasses import dataclass
19
+ from typing import Any, Dict, Optional, Tuple, Union
20
+
21
+ import torch
22
+ import torch.utils.checkpoint
23
+ from torch import nn
24
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
25
+
26
+ from transformers.activations import ACT2FN
27
+ from transformers.cache_utils import MambaCache
28
+ from transformers.modeling_utils import PreTrainedModel
29
+ from transformers.utils import (
30
+ ModelOutput,
31
+ add_code_sample_docstrings,
32
+ add_start_docstrings,
33
+ add_start_docstrings_to_model_forward,
34
+ logging,
35
+ replace_return_docstrings,
36
+ )
37
+ from transformers.utils.import_utils import is_causal_conv1d_available, is_mamba_ssm_available, is_mambapy_available
38
+ from .configuration_mamba import MambaConfig
39
+
40
+
41
+ logger = logging.get_logger(__name__)
42
+
43
+ # Check if we can use the fast path
44
+ if is_mambapy_available():
45
+ try:
46
+ from mambapy.pscan import pscan
47
+ except ImportError:
48
+ pscan = None
49
+ else:
50
+ pscan = None
51
+
52
+ if is_mamba_ssm_available():
53
+ try:
54
+ from mamba_ssm.ops.selective_scan_interface import mamba_inner_fn, selective_scan_fn
55
+ from mamba_ssm.ops.triton.selective_state_update import selective_state_update
56
+ except ImportError:
57
+ selective_state_update, selective_scan_fn, mamba_inner_fn = None, None, None
58
+ else:
59
+ selective_state_update, selective_scan_fn, mamba_inner_fn = None, None, None
60
+
61
+ if is_causal_conv1d_available():
62
+ try:
63
+ from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
64
+ except ImportError:
65
+ causal_conv1d_update, causal_conv1d_fn = None, None
66
+ else:
67
+ causal_conv1d_update, causal_conv1d_fn = None, None
68
+
69
+ is_fast_path_available = all(
70
+ (selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)
71
+ )
72
+
73
+
74
+ _CHECKPOINT_FOR_DOC = "state-spaces/mamba-130m-hf"
75
+ _CONFIG_FOR_DOC = "MambaConfig"
76
+
77
+
78
+ class MambaMixer(nn.Module):
79
+ """
80
+ Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
81
+ A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
82
+ ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
83
+ and is why Mamba is called **selective** state spaces)
84
+ """
85
+
86
+ def __init__(self, config: MambaConfig, layer_idx: int):
87
+ super().__init__()
88
+ self.config = config
89
+ self.hidden_size = config.hidden_size
90
+ self.ssm_state_size = config.state_size
91
+ self.conv_kernel_size = config.conv_kernel
92
+ self.intermediate_size = config.intermediate_size
93
+ self.time_step_rank = int(config.time_step_rank)
94
+ self.layer_idx = layer_idx
95
+ self.use_conv_bias = config.use_conv_bias
96
+ self.conv1d = nn.Conv1d(
97
+ in_channels=self.intermediate_size,
98
+ out_channels=self.intermediate_size,
99
+ bias=config.use_conv_bias,
100
+ kernel_size=config.conv_kernel,
101
+ groups=self.intermediate_size,
102
+ padding=config.conv_kernel - 1,
103
+ )
104
+
105
+ self.activation = config.hidden_act
106
+ self.act = ACT2FN[config.hidden_act]
107
+
108
+ self.use_mambapy = config.use_mambapy
109
+
110
+ # projection of the input hidden states
111
+ self.in_proj = nn.Linear(self.hidden_size, self.intermediate_size * 2, bias=config.use_bias)
112
+ # selective projection used to make dt, B and C input dependant
113
+ self.x_proj = nn.Linear(self.intermediate_size, self.time_step_rank + self.ssm_state_size * 2, bias=False)
114
+ # time step projection (discretization)
115
+ self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True)
116
+
117
+ # S4D real initialization. These are not discretized!
118
+ # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded
119
+ A = torch.arange(1, self.ssm_state_size + 1, dtype=torch.float32)[None, :]
120
+ A = A.expand(self.intermediate_size, -1).contiguous()
121
+
122
+ self.A_log = nn.Parameter(torch.log(A))
123
+ self.D = nn.Parameter(torch.ones(self.intermediate_size))
124
+ self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
125
+ self.use_bias = config.use_bias
126
+
127
+ if not is_fast_path_available:
128
+ if self.use_mambapy:
129
+ if is_mambapy_available():
130
+ logger.warning_once(
131
+ "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
132
+ " is None. Falling back to the mamba.py backend. To install follow https://github.com/state-spaces/mamba/#installation and"
133
+ " https://github.com/Dao-AILab/causal-conv1d"
134
+ )
135
+ else:
136
+ raise ImportError(
137
+ "use_mambapy is set to True but the mambapy package is not installed. To install it follow https://github.com/alxndrTL/mamba.py."
138
+ )
139
+ else:
140
+ logger.warning_once(
141
+ "The fast path is not available because one of `(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)`"
142
+ " is None. Falling back to the sequential implementation of Mamba, as use_mambapy is set to False. To install follow https://github.com/state-spaces/mamba/#installation and"
143
+ " https://github.com/Dao-AILab/causal-conv1d. For the mamba.py backend, follow https://github.com/alxndrTL/mamba.py."
144
+ )
145
+
146
+ def cuda_kernels_forward(
147
+ self,
148
+ hidden_states: torch.Tensor,
149
+ cache_params: Optional[MambaCache] = None,
150
+ cache_position: Optional[torch.LongTensor] = None,
151
+ attention_mask: Optional[torch.LongTensor] = None,
152
+ ):
153
+ # 1. Gated MLP's linear projection
154
+ projected_states = self.in_proj(hidden_states).transpose(1, 2)
155
+
156
+ if self.training and cache_params is None: # Doesn't support outputting the states -> used for training
157
+ contextualized_states = mamba_inner_fn(
158
+ projected_states,
159
+ self.conv1d.weight,
160
+ self.conv1d.bias if self.use_conv_bias else None,
161
+ self.x_proj.weight,
162
+ self.dt_proj.weight,
163
+ self.out_proj.weight,
164
+ self.out_proj.bias.float() if self.use_bias else None,
165
+ -torch.exp(self.A_log.float()),
166
+ None, # input-dependent B
167
+ None, # input-dependent C
168
+ self.D.float(),
169
+ delta_bias=self.dt_proj.bias.float(),
170
+ delta_softplus=True,
171
+ )
172
+
173
+ else:
174
+ hidden_states, gate = projected_states.chunk(2, dim=1)
175
+
176
+ if attention_mask is not None:
177
+ hidden_states = hidden_states * attention_mask.unsqueeze(1)
178
+
179
+ # 2. Convolution sequence transformation
180
+ conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0), self.conv1d.weight.size(2))
181
+ if cache_params is not None and cache_position[0] > 0:
182
+ hidden_states = causal_conv1d_update(
183
+ hidden_states.squeeze(-1),
184
+ cache_params.conv_states[self.layer_idx],
185
+ conv_weights,
186
+ self.conv1d.bias,
187
+ self.activation,
188
+ )
189
+ hidden_states = hidden_states.unsqueeze(-1)
190
+ else:
191
+ if cache_params is not None:
192
+ conv_states = nn.functional.pad(
193
+ hidden_states, (self.conv_kernel_size - hidden_states.shape[-1], 0)
194
+ )
195
+ cache_params.update_conv_state(self.layer_idx, conv_states, cache_position)
196
+ hidden_states = causal_conv1d_fn(
197
+ hidden_states, conv_weights, self.conv1d.bias, activation=self.activation
198
+ )
199
+
200
+ if attention_mask is not None:
201
+ hidden_states = hidden_states * attention_mask.unsqueeze(1)
202
+
203
+ # 3. State Space Model sequence transformation
204
+ # 3.a. input varying initialization of time_step, B and C
205
+ ssm_parameters = self.x_proj(hidden_states.transpose(1, 2))
206
+ time_step, B, C = torch.split(
207
+ ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1
208
+ )
209
+ discrete_time_step = self.dt_proj.weight @ time_step.transpose(1, 2)
210
+
211
+ A = -torch.exp(self.A_log.float())
212
+ # 3.c perform the recurrence y ← SSM(A, B, C)(x)
213
+ time_proj_bias = self.dt_proj.bias.float() if hasattr(self.dt_proj, "bias") else None
214
+ if cache_params is not None and cache_position[0] > 0:
215
+ scan_outputs = selective_state_update(
216
+ cache_params.ssm_states[self.layer_idx],
217
+ hidden_states[..., 0],
218
+ discrete_time_step[..., 0],
219
+ A,
220
+ B[:, 0],
221
+ C[:, 0],
222
+ self.D,
223
+ gate[..., 0],
224
+ time_proj_bias,
225
+ dt_softplus=True,
226
+ ).unsqueeze(-1)
227
+ else:
228
+ scan_outputs, ssm_state = selective_scan_fn(
229
+ hidden_states,
230
+ discrete_time_step,
231
+ A,
232
+ B.transpose(1, 2),
233
+ C.transpose(1, 2),
234
+ self.D.float(),
235
+ gate,
236
+ time_proj_bias,
237
+ delta_softplus=True,
238
+ return_last_state=True,
239
+ )
240
+ if ssm_state is not None and cache_params is not None:
241
+ cache_params.update_ssm_state(self.layer_idx, ssm_state)
242
+
243
+ # 4. Final linear projection
244
+ contextualized_states = self.out_proj(scan_outputs.transpose(1, 2))
245
+ return contextualized_states
246
+
247
+ # fmt: off
248
+ def slow_forward(self, input_states, cache_params: Optional[MambaCache]=None, cache_position:Optional[torch.LongTensor]=None, attention_mask: Optional[torch.LongTensor] = None):
249
+ batch_size, seq_len, _ = input_states.shape
250
+ dtype = input_states.dtype
251
+ # 1. Gated MLP's linear projection
252
+ projected_states = self.in_proj(input_states).transpose(1, 2) # [batch, 2 * intermediate_size, seq_len]
253
+ hidden_states, gate = projected_states.chunk(2, dim=1)
254
+
255
+ if attention_mask is not None:
256
+ hidden_states = hidden_states * attention_mask.unsqueeze(1)
257
+
258
+ # 2. Convolution sequence transformation
259
+ if cache_params is not None:
260
+ ssm_state = cache_params.ssm_states[self.layer_idx].clone()
261
+ ssm_state = ssm_state.to(hidden_states.device)
262
+ # use `cache_position.shape[0]` to check whether we are in prefill
263
+ # stage, it's equivalent to check `cache_position[0] == 0`, which
264
+ # breaks dynamo fullgraph constraints
265
+ if cache_position.shape[0] == self.conv_kernel_size:
266
+ conv_state = nn.functional.pad(
267
+ hidden_states,
268
+ (self.conv_kernel_size - hidden_states.shape[-1], 0)
269
+ )
270
+
271
+ cache_params.update_conv_state(self.layer_idx, conv_state, cache_position)
272
+ hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len]) # [batch, intermediate_size, seq_len]
273
+ else:
274
+ conv_state = cache_params.update_conv_state(self.layer_idx, hidden_states, cache_position)
275
+ hidden_states = torch.sum(conv_state * self.conv1d.weight[:, 0, :], dim=-1)
276
+ if self.use_conv_bias:
277
+ hidden_states += self.conv1d.bias
278
+ hidden_states = self.act(hidden_states).to(dtype).unsqueeze(-1) # [batch, intermediate_size, 1] : decoding
279
+ else:
280
+ ssm_state = torch.zeros(
281
+ (batch_size, self.intermediate_size, self.ssm_state_size),
282
+ device=hidden_states.device, dtype=dtype
283
+ )
284
+ hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len]) # [batch, intermediate_size, seq_len]
285
+
286
+ if attention_mask is not None:
287
+ hidden_states = hidden_states * attention_mask.unsqueeze(1)
288
+
289
+ # 3. State Space Model sequence transformation
290
+ # 3.a. Selection: [batch, seq_len, self.time_step_rank + self.ssm_state_size * 2]
291
+ ssm_parameters = self.x_proj(hidden_states.transpose(1, 2))
292
+ time_step, B, C = torch.split(
293
+ ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1
294
+ )
295
+ discrete_time_step = self.dt_proj(time_step) # [batch, seq_len, intermediate_size]
296
+ discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(1, 2) # [batch, intermediate_size, seq_len]
297
+
298
+ # 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM)
299
+ A = -torch.exp(self.A_log.float()) # [intermediate_size, ssm_state_size]
300
+ discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None]) # [batch, intermediate_size, seq_len, ssm_state_size]
301
+ discrete_B = discrete_time_step[:, :, :, None] * B[:, None, :, :].float() # [batch, intermediate_size, seq_len, ssm_state_size]
302
+ deltaB_u = discrete_B * hidden_states[:, :, :, None].float()
303
+
304
+ # 3.c perform the recurrence y ← SSM(A, B, C)(x)
305
+ if self.use_mambapy and self.training and cache_params is None:
306
+ hs = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)) # [batch, seq_len, intermediate_size, ssm_state_size]
307
+
308
+ scan_output = (hs @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) # [batch, intermediate_size, seq_len]
309
+ scan_output = scan_output + hidden_states * self.D[None, :, None]
310
+ scan_output = scan_output * self.act(gate)
311
+ else:
312
+ scan_outputs = []
313
+ for i in range(seq_len):
314
+ ssm_state = discrete_A[:, :, i, :] * ssm_state + deltaB_u[:, :, i, :] # [batch, intermediade_size, ssm_state]
315
+ scan_output = torch.matmul(ssm_state.to(dtype), C[:, i, :].unsqueeze(-1)) # [batch, intermediade_size, 1]
316
+ scan_outputs.append(scan_output[:, :, 0])
317
+ scan_output = torch.stack(scan_outputs, dim=-1) # [batch, seq_len, intermediade_size]
318
+ scan_output = scan_output + (hidden_states * self.D[None, :, None])
319
+ scan_output = (scan_output * self.act(gate))
320
+
321
+ if cache_params is not None:
322
+ cache_params.ssm_states[self.layer_idx].copy_(ssm_state)
323
+
324
+ # 4. Final linear projection
325
+ contextualized_states = self.out_proj(scan_output.transpose(1, 2)) # [batch, seq_len, hidden_size]
326
+ return contextualized_states
327
+ # fmt: on
328
+
329
+ def forward(
330
+ self,
331
+ hidden_states,
332
+ cache_params: Optional[MambaCache] = None,
333
+ cache_position: Optional[torch.LongTensor] = None,
334
+ attention_mask: Optional[torch.LongTensor] = None,
335
+ ):
336
+ if is_fast_path_available and "cuda" in self.x_proj.weight.device.type and not torch._dynamo.is_compiling():
337
+ return self.cuda_kernels_forward(hidden_states, cache_params, cache_position, attention_mask)
338
+ return self.slow_forward(hidden_states, cache_params, cache_position, attention_mask)
339
+
340
+
341
+ class MambaRMSNorm(nn.Module):
342
+ def __init__(self, hidden_size, eps=1e-6):
343
+ """
344
+ MambaRMSNorm is equivalent to T5LayerNorm and LlamaRMSNorm
345
+ """
346
+ super().__init__()
347
+ self.weight = nn.Parameter(torch.ones(hidden_size))
348
+ self.variance_epsilon = eps
349
+
350
+ def forward(self, hidden_states):
351
+ input_dtype = hidden_states.dtype
352
+ hidden_states = hidden_states.to(torch.float32)
353
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
354
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
355
+ return self.weight * hidden_states.to(input_dtype)
356
+
357
+ def extra_repr(self):
358
+ return f"{self.weight.shape[0]}, eps={self.variance_epsilon}"
359
+
360
+
361
+ class MambaBlock(nn.Module):
362
+ def __init__(self, config, layer_idx):
363
+ super().__init__()
364
+ self.config = config
365
+ self.layer_idx = layer_idx
366
+ self.residual_in_fp32 = config.residual_in_fp32
367
+ self.norm = MambaRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
368
+ self.mixer = MambaMixer(config, layer_idx=layer_idx)
369
+
370
+ def forward(
371
+ self,
372
+ hidden_states,
373
+ cache_params: Optional[MambaCache] = None,
374
+ cache_position: Optional[torch.LongTensor] = None,
375
+ attention_mask: Optional[torch.LongTensor] = None,
376
+ ):
377
+ residual = hidden_states
378
+ hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
379
+ if self.residual_in_fp32:
380
+ residual = residual.to(torch.float32)
381
+
382
+ hidden_states = self.mixer(
383
+ hidden_states, cache_params=cache_params, cache_position=cache_position, attention_mask=attention_mask
384
+ )
385
+ hidden_states = residual + hidden_states
386
+ return hidden_states
387
+
388
+
389
+ class MambaPreTrainedModel(PreTrainedModel):
390
+ """
391
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
392
+ models.
393
+ """
394
+
395
+ config_class = MambaConfig
396
+ base_model_prefix = "backbone"
397
+ _no_split_modules = ["MambaBlock", "MambaMixer"]
398
+ supports_gradient_checkpointing = True
399
+ _is_stateful = True
400
+
401
+ def _init_weights(self, module):
402
+ """Initialize the weights."""
403
+ if isinstance(module, MambaMixer):
404
+ module.A_log._no_weight_decay = True
405
+ module.D._no_weight_decay = True
406
+
407
+ dt_init_std = self.config.time_step_rank**-0.5 * self.config.time_step_scale
408
+ if self.config.time_step_init_scheme == "constant":
409
+ nn.init.constant_(module.dt_proj.weight, dt_init_std)
410
+ elif self.config.time_step_init_scheme == "random":
411
+ nn.init.uniform_(module.dt_proj.weight, -dt_init_std, dt_init_std)
412
+
413
+ dt = torch.exp(
414
+ torch.rand(self.config.intermediate_size)
415
+ * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min))
416
+ + math.log(self.config.time_step_min)
417
+ ).clamp(min=self.config.time_step_floor)
418
+ # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
419
+ inv_dt = dt + torch.log(-torch.expm1(-dt))
420
+ with torch.no_grad():
421
+ module.dt_proj.bias.copy_(inv_dt)
422
+ module.dt_proj.bias._no_reinit = True
423
+
424
+ if isinstance(module, nn.Linear):
425
+ if module.bias is not None:
426
+ if not getattr(module.bias, "_no_reinit", False):
427
+ nn.init.zeros_(module.bias)
428
+ elif isinstance(module, nn.Embedding):
429
+ nn.init.normal_(module.weight, std=self.config.initializer_range)
430
+
431
+ if self.config.rescale_prenorm_residual:
432
+ # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
433
+ # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
434
+ # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
435
+ # > -- GPT-2 :: https://openai.com/blog/better-language-models/
436
+ #
437
+ # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
438
+ for name, p in module.named_parameters():
439
+ if name in ["out_proj.weight"]:
440
+ # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
441
+ # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
442
+ # We need to reinit p since this code could be called multiple times
443
+ # Having just p *= scale would repeatedly scale it down
444
+ nn.init.kaiming_uniform_(p, a=math.sqrt(5))
445
+ with torch.no_grad():
446
+ p /= math.sqrt(self.config.num_hidden_layers)
447
+
448
+
449
+ @dataclass
450
+ class MambaOutput(ModelOutput):
451
+ """
452
+ Class for the MAMBA model outputs.
453
+
454
+ Args:
455
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
456
+ Sequence of hidden-states at the output of the last layer of the model.
457
+ cache_params (`MambaCache`):
458
+ The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
459
+ avoid providing the old `input_ids`.
460
+
461
+ Includes both the State space model state matrices after the selective scan, and the Convolutional states
462
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
463
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
464
+ one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
465
+
466
+ Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
467
+ """
468
+
469
+ last_hidden_state: Optional[torch.FloatTensor] = None
470
+ cache_params: Optional[MambaCache] = None
471
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
472
+
473
+
474
+ @dataclass
475
+ class MambaSequenceClassifierOutput(ModelOutput):
476
+ """
477
+ Base class for outputs of sentence classification models.
478
+
479
+ Args:
480
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
481
+ Classification (or regression if config.num_labels==1) loss.
482
+ logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
483
+ Classification (or regression if config.num_labels==1) scores (before SoftMax).
484
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
485
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
486
+ one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
487
+ Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
488
+ cache_params (`MambaCache`):
489
+ The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
490
+ avoid providing the old `input_ids`.
491
+ """
492
+
493
+ loss: Optional[torch.FloatTensor] = None
494
+ logits: torch.FloatTensor = None
495
+ cache_params: Optional[MambaCache] = None
496
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
497
+
498
+
499
+ @dataclass
500
+ class MambaCausalLMOutput(ModelOutput):
501
+ """
502
+ Base class for causal language model (or autoregressive) outputs.
503
+
504
+ Args:
505
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
506
+ Language modeling loss (for next-token prediction).
507
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
508
+ Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
509
+ cache_params (`MambaCache`):
510
+ The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
511
+ avoid providing the old `input_ids`.
512
+
513
+ Includes both the State space model state matrices after the selective scan, and the Convolutional states
514
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
515
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
516
+ one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
517
+
518
+ Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
519
+ """
520
+
521
+ loss: Optional[torch.FloatTensor] = None
522
+ logits: Optional[torch.FloatTensor] = None
523
+ cache_params: Optional[MambaCache] = None
524
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
525
+
526
+
527
+ MAMBA_START_DOCSTRING = r"""
528
+
529
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
530
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
531
+ etc.)
532
+
533
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
534
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
535
+ and behavior.
536
+
537
+ Parameters:
538
+ config ([`MambaConfig`]): Model configuration class with all the parameters of the model.
539
+ Initializing with a config file does not load the weights associated with the model, only the
540
+ configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
541
+ """
542
+
543
+ MAMBA_INPUTS_DOCSTRING = r"""
544
+ Args:
545
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
546
+ Indices of input sequence tokens in the vocabulary.
547
+
548
+ If `cache_params.seqlen_offset>0`, only `input_ids` that do not have their past calculated should be passed as
549
+ `input_ids`.
550
+
551
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
552
+ [`PreTrainedTokenizer.__call__`] for details.
553
+
554
+ [What are input IDs?](../glossary#input-ids)
555
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
556
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
557
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
558
+ model's internal embedding lookup matrix.
559
+ cache_params (`MambaCache`, *optional*):
560
+ If passed along, the model uses the previous state in all the blocks (which will give the output for the
561
+ `input_ids` provided as if the model add `state_input_ids + input_ids` as context).
562
+ use_cache (`bool`, *optional*):
563
+ If set to `True`, the `cache_params` is returned and can be used to quickly generate the next logits.
564
+ output_hidden_states (`bool`, *optional*):
565
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
566
+ more detail.
567
+ return_dict (`bool`, *optional*):
568
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
569
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
570
+ Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
571
+ this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
572
+ the complete sequence length.
573
+ """
574
+
575
+
576
+ @add_start_docstrings(
577
+ "The bare MAMBA Model transformer outputting raw hidden-states without any specific head on top.",
578
+ MAMBA_START_DOCSTRING,
579
+ )
580
+ class MambaModel(MambaPreTrainedModel):
581
+ def __init__(self, config):
582
+ super().__init__(config)
583
+
584
+ self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
585
+ self.layers = nn.ModuleList([MambaBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)])
586
+
587
+ self.gradient_checkpointing = False
588
+ self.norm_f = MambaRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
589
+ # Initialize weights and apply final processing
590
+ self._register_load_state_dict_pre_hook(self.load_hook)
591
+ self.post_init()
592
+
593
+ def load_hook(self, state_dict, prefix, *args):
594
+ for k in state_dict:
595
+ if "embedding." in k:
596
+ state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k)
597
+ break
598
+
599
+ def get_input_embeddings(self):
600
+ return self.embeddings
601
+
602
+ def set_input_embeddings(self, new_embeddings):
603
+ self.embeddings = new_embeddings
604
+
605
+ @add_start_docstrings_to_model_forward(MAMBA_INPUTS_DOCSTRING)
606
+ @add_code_sample_docstrings(
607
+ checkpoint=_CHECKPOINT_FOR_DOC,
608
+ output_type=MambaOutput,
609
+ config_class=_CONFIG_FOR_DOC,
610
+ )
611
+ def forward(
612
+ self,
613
+ input_ids: Optional[torch.LongTensor] = None,
614
+ inputs_embeds: Optional[torch.LongTensor] = None,
615
+ cache_params: Optional[MambaCache] = None,
616
+ use_cache: Optional[bool] = None,
617
+ output_hidden_states: Optional[bool] = None,
618
+ return_dict: Optional[bool] = None,
619
+ cache_position: Optional[torch.LongTensor] = None,
620
+ attention_mask: Optional[torch.LongTensor] = None,
621
+ ) -> Union[Tuple, MambaOutput]:
622
+ output_hidden_states = (
623
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
624
+ )
625
+ use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
626
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
627
+
628
+ if (input_ids is None) ^ (inputs_embeds is not None): # ^ is python for xor
629
+ raise ValueError(
630
+ "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
631
+ )
632
+
633
+ if inputs_embeds is None:
634
+ inputs_embeds = self.embeddings(input_ids)
635
+
636
+ if self.gradient_checkpointing and self.training and use_cache:
637
+ use_cache = False
638
+
639
+ if use_cache:
640
+ if cache_params is None:
641
+ cache_params = MambaCache(
642
+ self.config, inputs_embeds.size(0), device=inputs_embeds.device, dtype=inputs_embeds.dtype
643
+ )
644
+ cache_position = torch.arange(0, self.config.conv_kernel, device=inputs_embeds.device)
645
+ elif cache_position is None:
646
+ # cases when we do manual forward instead of using `model.generate` which will initiate
647
+ # `cache_position` and makes sure it is not None, throw error here instead of doing some
648
+ # hack to conjecture the current cache position
649
+ raise ValueError(
650
+ "You have to specify the `cache_position` manually when `use_cache=True` and `cache_params` is passed, "
651
+ "you don't have to pass a `cache_params` if you are in prefilling stage because in that case it will "
652
+ "be initialized for you automatically"
653
+ )
654
+ else:
655
+ cache_params = None
656
+
657
+ hidden_states = inputs_embeds
658
+ all_hidden_states = () if output_hidden_states else None
659
+ for mixer_block in self.layers:
660
+ if self.gradient_checkpointing and self.training:
661
+ hidden_states = self._gradient_checkpointing_func(
662
+ mixer_block.__call__, hidden_states, cache_params, cache_position, attention_mask
663
+ )
664
+ else:
665
+ hidden_states = mixer_block(
666
+ hidden_states,
667
+ cache_params=cache_params,
668
+ cache_position=cache_position,
669
+ attention_mask=attention_mask,
670
+ )
671
+
672
+ if output_hidden_states:
673
+ all_hidden_states = all_hidden_states + (hidden_states,)
674
+
675
+ hidden_states = self.norm_f(hidden_states)
676
+
677
+ if output_hidden_states:
678
+ all_hidden_states = all_hidden_states + (hidden_states,)
679
+
680
+ if not return_dict:
681
+ return tuple(v for v in [hidden_states, cache_params, all_hidden_states] if v is not None)
682
+
683
+ return MambaOutput(
684
+ last_hidden_state=hidden_states,
685
+ cache_params=cache_params if use_cache else None,
686
+ hidden_states=all_hidden_states,
687
+ )
688
+
689
+
690
+ @add_start_docstrings(
691
+ """
692
+ The MAMBA Model transformer with a language modeling head on top (linear layer with weights tied to the input
693
+ embeddings).
694
+ """,
695
+ MAMBA_START_DOCSTRING,
696
+ )
697
+ class MambaForCausalLM(MambaPreTrainedModel):
698
+ _tied_weights_keys = ["lm_head.weight"]
699
+
700
+ def __init__(self, config):
701
+ super().__init__(config)
702
+ self.backbone = MambaModel(config)
703
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
704
+ # Initialize weights and apply final processing
705
+ self.post_init()
706
+
707
+ def get_output_embeddings(self):
708
+ return self.lm_head
709
+
710
+ def set_output_embeddings(self, new_embeddings):
711
+ self.lm_head = new_embeddings
712
+
713
+ def get_input_embeddings(self):
714
+ return self.backbone.get_input_embeddings()
715
+
716
+ def set_input_embeddings(self, new_embeddings):
717
+ return self.backbone.set_input_embeddings(new_embeddings)
718
+
719
+ def _update_model_kwargs_for_generation(
720
+ self, outputs: ModelOutput, model_kwargs: Dict[str, Any], num_new_tokens: int = 1, **kwargs
721
+ ) -> Dict[str, Any]:
722
+ model_kwargs["cache_params"] = outputs.get("cache_params", None)
723
+ if (
724
+ model_kwargs.get("use_cache", True)
725
+ and "cache_position" in model_kwargs
726
+ and model_kwargs["cache_position"] is not None
727
+ ):
728
+ model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
729
+
730
+ if "attention_mask" in model_kwargs:
731
+ attention_mask = model_kwargs["attention_mask"]
732
+ model_kwargs["attention_mask"] = torch.cat(
733
+ [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
734
+ )
735
+
736
+ return model_kwargs
737
+
738
+ def prepare_inputs_for_generation(
739
+ self,
740
+ input_ids,
741
+ inputs_embeds=None,
742
+ use_cache=None,
743
+ cache_params: Optional[MambaCache] = None,
744
+ cache_position: Optional[torch.LongTensor] = None,
745
+ attention_mask: Optional[torch.LongTensor] = None,
746
+ **kwargs,
747
+ ):
748
+ if use_cache:
749
+ # `cache_position` should have been initialized in `generate`
750
+ if cache_position is None:
751
+ raise ValueError(
752
+ "`cache_position` should not be None as it should have been initialized in "
753
+ "`model.generate`, you are responsible for passing in a valid `cache_position` if "
754
+ "you are calling `prepare_inputs_for_generation` directly with `use_cache=True`"
755
+ )
756
+ if cache_position[0] > 0:
757
+ input_ids = input_ids[:, -1].unsqueeze(-1)
758
+
759
+ if attention_mask is not None:
760
+ attention_mask = None
761
+
762
+ else:
763
+ # we initialize the `cache_position` to full size of `conv_states` at prefill stage
764
+ # considering padding will be applied when input length is shorter, and truncation
765
+ # will be applied when it is longer, so it will be equivalent to always have it match
766
+ # the length of `cache_params.conv_states`, which is `config.conv_kernel`
767
+ cache_position = torch.arange(0, self.config.conv_kernel, device=input_ids.device)
768
+
769
+ if inputs_embeds is not None and cache_params is None:
770
+ model_inputs = {"inputs_embeds": inputs_embeds}
771
+ else:
772
+ model_inputs = {"input_ids": input_ids.contiguous()}
773
+
774
+ model_inputs.update(
775
+ {
776
+ "cache_params": cache_params,
777
+ "use_cache": use_cache,
778
+ "cache_position": cache_position,
779
+ "attention_mask": attention_mask,
780
+ }
781
+ )
782
+ return model_inputs
783
+
784
+ @add_start_docstrings_to_model_forward(MAMBA_INPUTS_DOCSTRING)
785
+ @add_code_sample_docstrings(
786
+ checkpoint=_CHECKPOINT_FOR_DOC,
787
+ output_type=MambaCausalLMOutput,
788
+ config_class=_CONFIG_FOR_DOC,
789
+ )
790
+ def forward(
791
+ self,
792
+ input_ids: Optional[torch.LongTensor] = None,
793
+ attention_mask: Optional[torch.LongTensor] = None,
794
+ inputs_embeds: Optional[torch.FloatTensor] = None,
795
+ cache_params: Optional[MambaCache] = None,
796
+ labels: Optional[torch.LongTensor] = None,
797
+ output_hidden_states: Optional[bool] = None,
798
+ return_dict: Optional[bool] = None,
799
+ use_cache: Optional[bool] = None,
800
+ cache_position: Optional[torch.Tensor] = None,
801
+ **kwargs, # for now we need this for generation
802
+ ) -> Union[Tuple, MambaCausalLMOutput]:
803
+ r"""
804
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
805
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
806
+ `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
807
+ are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
808
+ """
809
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
810
+
811
+ mamba_outputs = self.backbone(
812
+ input_ids,
813
+ cache_params=cache_params,
814
+ inputs_embeds=inputs_embeds,
815
+ output_hidden_states=output_hidden_states,
816
+ return_dict=return_dict,
817
+ use_cache=use_cache,
818
+ cache_position=cache_position,
819
+ attention_mask=attention_mask,
820
+ )
821
+ hidden_states = mamba_outputs[0]
822
+
823
+ logits = self.lm_head(hidden_states.to(self.lm_head.weight.dtype)).float()
824
+
825
+ loss = None
826
+ if labels is not None:
827
+ # move labels to correct device to enable model parallelism
828
+ labels = labels.to(logits.device)
829
+ # Shift so that tokens < n predict n
830
+ shift_logits = logits[..., :-1, :].contiguous()
831
+ shift_labels = labels[..., 1:].contiguous()
832
+ # Flatten the tokens
833
+ loss_fct = CrossEntropyLoss()
834
+ loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
835
+
836
+ if not return_dict:
837
+ output = (logits,) + mamba_outputs[1:]
838
+ return ((loss,) + output) if loss is not None else output
839
+
840
+ return MambaCausalLMOutput(
841
+ loss=loss,
842
+ logits=logits,
843
+ cache_params=mamba_outputs.cache_params,
844
+ hidden_states=mamba_outputs.hidden_states,
845
+ )
846
+
847
+
848
+ @add_start_docstrings(
849
+ """
850
+ Mamba Model backbone with a sequence classification/regression head on top
851
+ (a linear layer on top of the pooled output) e.g. for GLUE tasks.
852
+
853
+ [`MambaForSequenceClassification`] uses the last token in order to do the classification, as other causal models
854
+ (e.g. GPT-2) do.
855
+
856
+ Since it does classification on the last token, it requires to know the position of the last token.
857
+ If a `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row.
858
+ If no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
859
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
860
+ each row of the batch).
861
+ """,
862
+ MAMBA_START_DOCSTRING,
863
+ )
864
+ class MambaForSequenceClassification(MambaPreTrainedModel):
865
+ def __init__(self, config):
866
+ super().__init__(config)
867
+ self.num_labels = config.num_labels
868
+ self.config = config
869
+ self.backbone = MambaModel(config)
870
+ self.classifier = nn.Linear(config.hidden_size, self.num_labels, bias=True)
871
+
872
+ # Initialize weights and apply final processing
873
+ self.post_init()
874
+
875
+ @add_start_docstrings_to_model_forward(MAMBA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
876
+ @replace_return_docstrings(output_type=MambaSequenceClassifierOutput, config_class=_CONFIG_FOR_DOC)
877
+ @add_code_sample_docstrings(
878
+ checkpoint=_CHECKPOINT_FOR_DOC,
879
+ output_type=MambaSequenceClassifierOutput,
880
+ config_class=_CONFIG_FOR_DOC,
881
+ )
882
+ def forward(
883
+ self,
884
+ input_ids: Optional[torch.LongTensor] = None,
885
+ inputs_embeds: Optional[torch.FloatTensor] = None,
886
+ cache_params: Optional[MambaCache] = None,
887
+ labels: Optional[torch.LongTensor] = None,
888
+ output_hidden_states: Optional[bool] = None,
889
+ return_dict: Optional[bool] = None,
890
+ use_cache: Optional[bool] = None,
891
+ **kwargs,
892
+ ) -> Union[MambaSequenceClassifierOutput, Tuple[torch.FloatTensor]]:
893
+ r"""
894
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
895
+ Labels for computing the sequence classification/regression loss.
896
+ Indices should be in `[0, ..., config.num_labels - 1]`.
897
+ If `config.num_labels == 1` a regression loss is computed (Mean-Square loss),
898
+ If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
899
+ """
900
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
901
+
902
+ mamba_outputs = self.backbone(
903
+ input_ids,
904
+ cache_params=cache_params,
905
+ inputs_embeds=inputs_embeds,
906
+ output_hidden_states=output_hidden_states,
907
+ return_dict=return_dict,
908
+ use_cache=use_cache,
909
+ )
910
+
911
+ last_hidden_states = mamba_outputs[0]
912
+
913
+ if input_ids is not None:
914
+ batch_size, _ = input_ids.shape[:2]
915
+ else:
916
+ batch_size, _ = inputs_embeds.shape[:2]
917
+
918
+ if self.config.pad_token_id is None and batch_size > 1:
919
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
920
+
921
+ if self.config.pad_token_id is None:
922
+ sequence_lengths = -1
923
+ else:
924
+ if input_ids is not None:
925
+ # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
926
+ sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
927
+ sequence_lengths = sequence_lengths % input_ids.shape[-1]
928
+ sequence_lengths = sequence_lengths.to(last_hidden_states.device)
929
+ else:
930
+ sequence_lengths = -1
931
+ logger.warning(
932
+ f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
933
+ "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
934
+ )
935
+
936
+ pooled_last_hidden_states = last_hidden_states[
937
+ torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths
938
+ ]
939
+ pooled_logits = self.classifier(pooled_last_hidden_states)
940
+
941
+ loss = None
942
+ if labels is not None:
943
+ if self.config.problem_type is None:
944
+ if self.num_labels == 1:
945
+ self.config.problem_type = "regression"
946
+ elif self.num_labels > 1 and (labels.dtype in [torch.long, torch.int]):
947
+ self.config.problem_type = "single_label_classification"
948
+ else:
949
+ self.config.problem_type = "multi_label_classification"
950
+
951
+ if self.config.problem_type == "regression":
952
+ loss_fct = MSELoss()
953
+ if self.num_labels == 1:
954
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
955
+ else:
956
+ loss = loss_fct(pooled_logits, labels)
957
+ elif self.config.problem_type == "single_label_classification":
958
+ loss_fct = CrossEntropyLoss()
959
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
960
+ elif self.config.problem_type == "multi_label_classification":
961
+ loss_fct = BCEWithLogitsLoss()
962
+ loss = loss_fct(pooled_logits, labels)
963
+
964
+ if not return_dict:
965
+ output = (pooled_logits,) + mamba_outputs[1:]
966
+ return ((loss,) + output) if loss is not None else output
967
+
968
+ return MambaSequenceClassifierOutput(
969
+ loss=loss,
970
+ logits=pooled_logits,
971
+ cache_params=mamba_outputs.cache_params,
972
+ hidden_states=mamba_outputs.hidden_states,
973
+ )