Q-bert commited on
Commit
8af6a54
1 Parent(s): ab819c5

Upload 2 files

Browse files
Files changed (2) hide show
  1. configuration_mamba.py +43 -0
  2. modeling_mamba.py +308 -0
configuration_mamba.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ from typing import Optional , Union
3
+
4
+ from transformers import PretrainedConfig
5
+ class MambaConfig(PretrainedConfig):
6
+ model_type = "mamba"
7
+ def __init__(
8
+ self,
9
+ vocab_size=50277,
10
+ d_state=16,
11
+ d_model=2560,
12
+ d_conv=4,
13
+ expand=2,
14
+ conv_bias=True,
15
+ bias=False,
16
+ n_layer=64,
17
+ dt_rank: Union[int, str] = "auto",
18
+ pad_vocab_size_multiple=8,
19
+ initializer_range=0.02,
20
+ **kwargs,
21
+ ):
22
+ self.vocab_size = vocab_size
23
+ self.n_layer= n_layer
24
+ self.conv_bias = conv_bias
25
+ self.expand = expand
26
+ self.pad_vocab_size_multiple = pad_vocab_size_multiple
27
+ self.d_conv = d_conv
28
+ self.d_model = d_model
29
+ self.d_state = d_state
30
+ self.d_inner = int(self.expand * self.d_model)
31
+ self.dt_rank = dt_rank
32
+ self.initializer_range = initializer_range
33
+ self.bias = bias
34
+
35
+ if self.dt_rank == 'auto':
36
+ self.dt_rank = math.ceil(self.d_model / 16)
37
+
38
+ if self.vocab_size % self.pad_vocab_size_multiple != 0:
39
+ self.vocab_size += (self.pad_vocab_size_multiple
40
+ - self.vocab_size % self.pad_vocab_size_multiple)
41
+ super().__init__(
42
+ **kwargs,
43
+ )
modeling_mamba.py ADDED
@@ -0,0 +1,308 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch.nn as nn
2
+ import torch
3
+ from configuration_mamba import MambaConfig
4
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
5
+ from transformers.modeling_utils import PreTrainedModel
6
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
7
+ import math
8
+ import json
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.nn.functional as F
12
+ from dataclasses import dataclass
13
+ from einops import rearrange, repeat, einsum
14
+ from typing import Optional , Union ,Tuple
15
+
16
+ # Dear contributors of the https://github.com/johnma2006/mamba-minimal/tree/master repository, special thanks to Albert Gu and Tri Dao for their articles. (https://arxiv.org/abs/2312.00752)
17
+
18
+
19
+ class MambaRMSNorm(nn.Module):
20
+ def __init__(self,
21
+ d_model: int,
22
+ eps: float = 1e-5):
23
+ super().__init__()
24
+ self.eps = eps
25
+ self.weight = nn.Parameter(torch.ones(d_model))
26
+ def forward(self, x):
27
+ output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
28
+ return output
29
+
30
+
31
+ class MambaBlock(nn.Module):
32
+ def __init__(self, config: MambaConfig):
33
+ """A single Mamba block, as described in Figure 3 in Section 3.4 in the Mamba paper [1]."""
34
+ super().__init__()
35
+ self.config = config
36
+
37
+ self.in_proj = nn.Linear(config.d_model, config.d_inner * 2, bias=config.bias)
38
+
39
+ self.conv1d = nn.Conv1d(
40
+ in_channels=config.d_inner,
41
+ out_channels=config.d_inner,
42
+ bias=config.conv_bias,
43
+ kernel_size=config.d_conv,
44
+ groups=config.d_inner,
45
+ padding=config.d_conv - 1,
46
+ )
47
+
48
+ # x_proj takes in `x` and outputs the input-specific Δ, B, C
49
+ self.x_proj = nn.Linear(config.d_inner, config.dt_rank + config.d_state * 2, bias=False)
50
+
51
+ # dt_proj projects Δ from dt_rank to d_in
52
+ self.dt_proj = nn.Linear(config.dt_rank, config.d_inner, bias=True)
53
+
54
+ A = repeat(torch.arange(1, config.d_state + 1), 'n -> d n', d=config.d_inner)
55
+ self.A_log = nn.Parameter(torch.log(A))
56
+ self.D = nn.Parameter(torch.ones(config.d_inner))
57
+ self.out_proj = nn.Linear(config.d_inner, config.d_model, bias=config.bias)
58
+ self.norm = MambaRMSNorm(config.d_model)
59
+
60
+ def forward(self, x):
61
+ """Mamba block forward. This looks the same as Figure 3 in Section 3.4 in the Mamba paper [1].
62
+
63
+ Args:
64
+ x: shape (b, l, d) (See Glossary at top for definitions of b, l, d_in, n...)
65
+
66
+ Returns:
67
+ output: shape (b, l, d)
68
+
69
+ Official Implementation:
70
+ class Mamba, https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py#L119
71
+ mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311
72
+
73
+ """
74
+
75
+ (b, l, d) = x.shape
76
+ x_copy = x # There was a separate class for residual, I deleted that part and added it here.
77
+ x = self.norm(x)
78
+ x_and_res = self.in_proj(x) # shape (b, l, 2 * d_in)
79
+ (x, res) = x_and_res.split(split_size=[self.config.d_inner, self.config.d_inner], dim=-1)
80
+
81
+ x = rearrange(x, 'b l d_in -> b d_in l')
82
+ x = self.conv1d(x)[:, :, :l]
83
+ x = rearrange(x, 'b d_in l -> b l d_in')
84
+
85
+ x = F.silu(x)
86
+
87
+ y = self.ssm(x)
88
+
89
+ y = y * F.silu(res)
90
+
91
+ output = self.out_proj(y) + x_copy
92
+
93
+ return output
94
+
95
+
96
+ def ssm(self, x):
97
+ """Runs the SSM. See:
98
+ - Algorithm 2 in Section 3.2 in the Mamba paper [1]
99
+ - run_SSM(A, B, C, u) in The Annotated S4 [2]
100
+
101
+ Args:
102
+ x: shape (b, l, d_in) (See Glossary at top for definitions of b, l, d_in, n...)
103
+
104
+ Returns:
105
+ output: shape (b, l, d_in)
106
+
107
+ Official Implementation:
108
+ mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311
109
+
110
+ """
111
+ (d_in, n) = self.A_log.shape
112
+
113
+ # Compute ∆ A B C D, the state space parameters.
114
+ # A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
115
+ # ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
116
+ # and is why Mamba is called **selective** state spaces)
117
+
118
+ A = -torch.exp(self.A_log.float()) # shape (d_in, n)
119
+ D = self.D.float()
120
+
121
+ x_dbl = self.x_proj(x) # (b, l, dt_rank + 2*n)
122
+
123
+ (delta, B, C) = x_dbl.split(split_size=[self.config.dt_rank, n, n], dim=-1) # delta: (b, l, dt_rank). B, C: (b, l, n)
124
+ delta = F.softplus(self.dt_proj(delta)) # (b, l, d_in)
125
+
126
+ y = self.selective_scan(x, delta, A, B, C, D) # This is similar to run_SSM(A, B, C, u) in The Annotated S4 [2]
127
+
128
+ return y
129
+
130
+
131
+ def selective_scan(self, u, delta, A, B, C, D):
132
+ """Does selective scan algorithm. See:
133
+ - Section 2 State Space Models in the Mamba paper [1]
134
+ - Algorithm 2 in Section 3.2 in the Mamba paper [1]
135
+ - run_SSM(A, B, C, u) in The Annotated S4 [2]
136
+
137
+ This is the classic discrete state space formula:
138
+ x(t + 1) = Ax(t) + Bu(t)
139
+ y(t) = Cx(t) + Du(t)
140
+ except B and C (and the step size delta, which is used for discretization) are dependent on the input x(t).
141
+
142
+ Args:
143
+ u: shape (b, l, d_in) (See Glossary at top for definitions of b, l, d_in, n...)
144
+ delta: shape (b, l, d_in)
145
+ A: shape (d_in, n)
146
+ B: shape (b, l, n)
147
+ C: shape (b, l, n)
148
+ D: shape (d_in,)
149
+
150
+ Returns:
151
+ output: shape (b, l, d_in)
152
+
153
+ Official Implementation:
154
+ selective_scan_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L86
155
+ Note: I refactored some parts out of `selective_scan_ref` out, so the functionality doesn't match exactly.
156
+
157
+ """
158
+ (b, l, d_in) = u.shape
159
+ n = A.shape[1]
160
+
161
+ # Discretize continuous parameters (A, B)
162
+ # - A is discretized using zero-order hold (ZOH) discretization (see Section 2 Equation 4 in the Mamba paper [1])
163
+ # - B is discretized using a simplified Euler discretization instead of ZOH. From a discussion with authors:
164
+ # "A is the more important term and the performance doesn't change much with the simplication on B"
165
+ deltaA = torch.exp(einsum(delta, A, 'b l d_in, d_in n -> b d_in l n'))
166
+ deltaB_u = einsum(delta, B, u, 'b l d_in, b l n, b l d_in -> b d_in l n')
167
+
168
+ # Perform selective scan (see scan_SSM() in The Annotated S4 [2])
169
+ x = torch.zeros((b, d_in, n), device=deltaA.device)
170
+ ys = []
171
+ for i in range(l):
172
+ x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
173
+ y = einsum(x, C[:, i, :], 'b d_in n, b n -> b d_in')
174
+ ys.append(y)
175
+ y = torch.stack(ys, dim=1) # shape (b, l, d_in)
176
+
177
+ y = y + u * D
178
+
179
+ return y
180
+
181
+ class MambaPreTrainedModel(PreTrainedModel):
182
+ config_class = MambaConfig
183
+ base_model_prefix = "model"
184
+ supports_gradient_checkpointing = True
185
+ _no_split_modules = ["MambaBlock"]
186
+
187
+ def _init_weights(self, module):
188
+ std = 0.02
189
+ if isinstance(module, (nn.Linear, nn.Conv1d)):
190
+ module.weight.data.normal_(mean=0.0, std=std)
191
+ if module.bias is not None:
192
+ module.bias.data.zero_()
193
+ elif isinstance(module, nn.Embedding):
194
+ module.weight.data.normal_(mean=0.0, std=std)
195
+ if module.padding_idx is not None:
196
+ module.weight.data[module.padding_idx].zero_()
197
+
198
+ class MambaModel(MambaPreTrainedModel):
199
+ def __init__(self, config: MambaConfig):
200
+ """Full Mamba model.
201
+ Mamba model decoder consisting of *config.n_layer* layers. Each layer is a [`MambaBlock`]
202
+
203
+ Args:
204
+ config: MambaConfig
205
+ """
206
+ super().__init__(config)
207
+ self.config = config
208
+
209
+ self.embedding = nn.Embedding(config.vocab_size, config.d_model)
210
+ self.layers = nn.ModuleList([MambaBlock(config) for _ in range(config.n_layer)])
211
+ self.norm_f = MambaRMSNorm(config.d_model)
212
+
213
+ self.gradient_checkpointing = False
214
+ self.post_init()
215
+
216
+ def get_input_embeddings(self):
217
+ return self.embedding
218
+
219
+ def set_input_embeddings(self, value):
220
+ self.embedding = value
221
+
222
+ def forward(self,
223
+ input_ids: torch.LongTensor = None,
224
+ return_dict: Optional[bool] = None,
225
+ )-> Union[Tuple, BaseModelOutputWithPast]:
226
+ x = self.embedding(input_ids)
227
+ all_hidden_states = list()
228
+ for layer in self.layers:
229
+ x = layer(x)
230
+ all_hidden_states.append(x)
231
+
232
+ hidden_states = self.norm_f(x)
233
+
234
+ return BaseModelOutputWithPast(
235
+ last_hidden_state=hidden_states,
236
+ hidden_states=all_hidden_states,
237
+ )
238
+ class MambaForCausalLM(MambaPreTrainedModel):
239
+ _tied_weights_keys = ["lm_head.weight"]
240
+
241
+ def __init__(self, config):
242
+ super().__init__(config)
243
+ self.model = MambaModel(config)
244
+ self.vocab_size = config.vocab_size
245
+ self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
246
+ self.lm_head.weight = self.model.embedding.weight
247
+ self.post_init()
248
+
249
+ def get_input_embeddings(self):
250
+ return self.model.embedding
251
+
252
+ def set_input_embeddings(self, value):
253
+ self.model.embedding = value
254
+
255
+ def get_output_embeddings(self):
256
+ return self.lm_head
257
+
258
+ def set_output_embeddings(self, new_embeddings):
259
+ self.lm_head = new_embeddings
260
+
261
+ def set_decoder(self, decoder):
262
+ self.model = decoder
263
+
264
+ def get_decoder(self):
265
+ return self.model
266
+
267
+ def forward(self,
268
+ input_ids: torch.LongTensor = None,
269
+ labels: Optional[torch.LongTensor] = None,
270
+ output_attentions: Optional[bool] = None,
271
+ output_hidden_states: Optional[bool] = None,
272
+ return_dict: Optional[bool] = None,
273
+ )-> Union[Tuple, CausalLMOutputWithPast]:
274
+ outputs = self.model(
275
+ input_ids=input_ids,
276
+ return_dict=return_dict,
277
+ )
278
+ hidden_states = outputs[0]
279
+ logits = self.lm_head(hidden_states)
280
+ logits = logits.float()
281
+ loss = None
282
+ if labels is not None:
283
+ shift_logits = logits[..., :-1, :].contiguous()
284
+ shift_labels = labels[..., 1:].contiguous()
285
+ loss_fct = CrossEntropyLoss()
286
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
287
+ shift_labels = shift_labels.view(-1)
288
+
289
+ shift_labels = shift_labels.to(shift_logits.device)
290
+ loss = loss_fct(shift_logits, shift_labels)
291
+
292
+ if not return_dict:
293
+ output = (logits,) + outputs[1:]
294
+ return (loss,) + output if loss is not None else output
295
+
296
+ return CausalLMOutputWithPast(
297
+ loss=loss,
298
+ logits=logits,
299
+ hidden_states=outputs.hidden_states,
300
+ )
301
+
302
+ def prepare_inputs_for_generation(
303
+ self, input_ids, **kwargs
304
+ ):
305
+ model_inputs = {"input_ids": input_ids}
306
+ return model_inputs
307
+
308
+