Files changed (1) hide show
  1. modeling_mpt.py +76 -13
modeling_mpt.py CHANGED
@@ -1,5 +1,4 @@
1
  """A simple, flexible implementation of a GPT model.
2
-
3
  Inspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
4
  """
5
  import math
@@ -25,16 +24,24 @@ except:
25
  pass
26
  Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
27
 
 
28
  class MPTPreTrainedModel(PreTrainedModel):
29
  config_class = MPTConfig
30
  base_model_prefix = 'model'
31
- _no_split_modules = ['MPTBlock']
 
 
 
 
 
 
32
 
33
  class MPTModel(MPTPreTrainedModel):
34
 
35
  def __init__(self, config: MPTConfig):
36
  config._validate_config()
37
  super().__init__(config)
 
38
  self.attn_impl = config.attn_config['attn_impl']
39
  self.prefix_lm = config.attn_config['prefix_lm']
40
  self.attn_uses_sequence_id = config.attn_config['attn_uses_sequence_id']
@@ -140,7 +147,43 @@ class MPTModel(MPTPreTrainedModel):
140
  attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
141
  return attn_bias
142
 
143
- def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None, inputs_embeds: Optional[torch.Tensor]=None):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
144
  return_dict = return_dict if return_dict is not None else self.config.return_dict
145
  use_cache = use_cache if use_cache is not None else self.config.use_cache
146
  if attention_mask is not None:
@@ -156,13 +199,15 @@ class MPTModel(MPTPreTrainedModel):
156
  raise NotImplementedError('MPT does not support training with left padding.')
157
  if self.prefix_lm and prefix_mask is None:
158
  raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
159
- if inputs_embeds is not None:
160
- raise NotImplementedError('inputs_embeds is not implemented for MPT.')
161
  if self.training:
162
  if self.attn_uses_sequence_id and sequence_id is None:
163
  raise ValueError('sequence_id is a required argument when MPT is configured with attn_uses_sequence_id=True ' + 'and the model is in train mode.')
164
  elif self.attn_uses_sequence_id is False and sequence_id is not None:
165
  warnings.warn('MPT received non-None input for `sequence_id` but is configured with attn_uses_sequence_id=False. ' + 'This input will be ignored. If you want the model to use `sequence_id`, set attn_uses_sequence_id to True.')
 
 
 
 
166
  S = input_ids.size(1)
167
  assert S <= self.config.max_seq_len, f'Cannot forward input with seq_len={S}, this model only supports seq_len<={self.config.max_seq_len}'
168
  tok_emb = self.wte(input_ids)
@@ -199,7 +244,27 @@ class MPTModel(MPTPreTrainedModel):
199
  assert all_hidden_states is not None
200
  all_hidden_states = all_hidden_states + (x,)
201
  past_key_value = past_key_values[b_idx] if past_key_values is not None else None
202
- (x, attn_weights, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=self.is_causal)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
203
  if past_key_values is not None:
204
  past_key_values[b_idx] = past_key_value
205
  if output_attentions:
@@ -227,7 +292,6 @@ class MPTForCausalLM(MPTPreTrainedModel):
227
  super().__init__(config)
228
  if not config.tie_word_embeddings:
229
  raise ValueError('MPTForCausalLM only supports tied word embeddings')
230
- print(f'Instantiating an MPTForCausalLM model from {__file__}')
231
  self.transformer = MPTModel(config)
232
  for child in self.transformer.children():
233
  if isinstance(child, torch.nn.ModuleList):
@@ -262,12 +326,12 @@ class MPTForCausalLM(MPTPreTrainedModel):
262
  def get_decoder(self):
263
  return self.transformer
264
 
265
- def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, labels: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None, inputs_embeds: Optional[torch.FloatTensor]=None):
 
266
  return_dict = return_dict if return_dict is not None else self.config.return_dict
267
  use_cache = use_cache if use_cache is not None else self.config.use_cache
268
- if inputs_embeds is not None:
269
- raise NotImplementedError('inputs_embeds has to be None (for hf/peft support).')
270
- outputs = self.transformer(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id, return_dict=return_dict, output_attentions=output_attentions, output_hidden_states=output_hidden_states, use_cache=use_cache)
271
  logits = self.transformer.wte(outputs.last_hidden_state.to(self.transformer.wte.weight.device), True)
272
  if self.logit_scale is not None:
273
  if self.logit_scale == 0:
@@ -313,11 +377,10 @@ class MPTForCausalLM(MPTPreTrainedModel):
313
  @staticmethod
314
  def _reorder_cache(past_key_values, beam_idx):
315
  """Used by HuggingFace generate when using beam search with kv-caching.
316
-
317
  See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133
318
  for an example in transformers.
319
  """
320
  reordered_past = []
321
  for layer_past in past_key_values:
322
  reordered_past += [tuple((past_state.index_select(0, beam_idx) for past_state in layer_past))]
323
- return reordered_past
 
1
  """A simple, flexible implementation of a GPT model.
 
2
  Inspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
3
  """
4
  import math
 
24
  pass
25
  Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
26
 
27
+
28
  class MPTPreTrainedModel(PreTrainedModel):
29
  config_class = MPTConfig
30
  base_model_prefix = 'model'
31
+ _no_split_modules = ["MPTBlock"]
32
+ supports_gradient_checkpointing = True
33
+
34
+ def _set_gradient_checkpointing(self, module, value=False):
35
+ if isinstance(module, MPTModel):
36
+ module.gradient_checkpointing = value
37
+
38
 
39
  class MPTModel(MPTPreTrainedModel):
40
 
41
  def __init__(self, config: MPTConfig):
42
  config._validate_config()
43
  super().__init__(config)
44
+ self.gradient_checkpointing = False
45
  self.attn_impl = config.attn_config['attn_impl']
46
  self.prefix_lm = config.attn_config['prefix_lm']
47
  self.attn_uses_sequence_id = config.attn_config['attn_uses_sequence_id']
 
147
  attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
148
  return attn_bias
149
 
150
+ def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]] = None,
151
+ attention_mask: Optional[torch.ByteTensor] = None, prefix_mask: Optional[torch.ByteTensor] = None,
152
+ sequence_id: Optional[torch.LongTensor] = None, return_dict: Optional[bool] = None,
153
+ output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None,
154
+ use_cache: Optional[bool] = None, inputs_embeds: Optional[torch.FloatTensor] = None):
155
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
156
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
157
+ if self.gradient_checkpointing and self.training:
158
+ if use_cache:
159
+ use_cache = False
160
+ if input_ids is not None and inputs_embeds is not None:
161
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
162
+ elif input_ids is not None:
163
+ batch_size, seq_length = input_ids.shape
164
+ elif inputs_embeds is not None:
165
+ batch_size, seq_length, _ = inputs_embeds.shape
166
+ else:
167
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
168
+
169
+ seq_length_with_past = seq_length
170
+ past_key_values_length = 0
171
+
172
+ if past_key_values is not None:
173
+ past_key_values_length = past_key_values[0][0].shape[2]
174
+ seq_length_with_past = seq_length_with_past + past_key_values_length
175
+
176
+ if attention_mask is not None:
177
+ attention_mask = attention_mask.bool()
178
+ else:
179
+ attention_mask = torch.ones(
180
+ (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
181
+ )
182
+
183
+ if inputs_embeds is None:
184
+ tok_emb = self.wte(input_ids)
185
+ else:
186
+ tok_emb = inputs_embeds
187
  return_dict = return_dict if return_dict is not None else self.config.return_dict
188
  use_cache = use_cache if use_cache is not None else self.config.use_cache
189
  if attention_mask is not None:
 
199
  raise NotImplementedError('MPT does not support training with left padding.')
200
  if self.prefix_lm and prefix_mask is None:
201
  raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
 
 
202
  if self.training:
203
  if self.attn_uses_sequence_id and sequence_id is None:
204
  raise ValueError('sequence_id is a required argument when MPT is configured with attn_uses_sequence_id=True ' + 'and the model is in train mode.')
205
  elif self.attn_uses_sequence_id is False and sequence_id is not None:
206
  warnings.warn('MPT received non-None input for `sequence_id` but is configured with attn_uses_sequence_id=False. ' + 'This input will be ignored. If you want the model to use `sequence_id`, set attn_uses_sequence_id to True.')
207
+ if self.gradient_checkpointing and self.training:
208
+ if use_cache:
209
+ use_cache = False
210
+
211
  S = input_ids.size(1)
212
  assert S <= self.config.max_seq_len, f'Cannot forward input with seq_len={S}, this model only supports seq_len<={self.config.max_seq_len}'
213
  tok_emb = self.wte(input_ids)
 
244
  assert all_hidden_states is not None
245
  all_hidden_states = all_hidden_states + (x,)
246
  past_key_value = past_key_values[b_idx] if past_key_values is not None else None
247
+ if self.gradient_checkpointing and self.training:
248
+
249
+ def create_custom_forward(module):
250
+ def custom_forward(*inputs):
251
+ # None for past_key_value
252
+ return module(*inputs)
253
+
254
+ return custom_forward
255
+
256
+ (x, attn_weights, past_key_value) = torch.utils.checkpoint.checkpoint(
257
+ create_custom_forward(block),
258
+ x,
259
+ past_key_value,
260
+ attn_bias,
261
+ attention_mask,
262
+ self.is_causal,
263
+ )
264
+ else:
265
+ (x, attn_weights, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias,
266
+ attention_mask=attention_mask, is_causal=self.is_causal)
267
+
268
  if past_key_values is not None:
269
  past_key_values[b_idx] = past_key_value
270
  if output_attentions:
 
292
  super().__init__(config)
293
  if not config.tie_word_embeddings:
294
  raise ValueError('MPTForCausalLM only supports tied word embeddings')
 
295
  self.transformer = MPTModel(config)
296
  for child in self.transformer.children():
297
  if isinstance(child, torch.nn.ModuleList):
 
326
  def get_decoder(self):
327
  return self.transformer
328
 
329
+ def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, labels: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None,
330
+ inputs_embeds: Optional[torch.FloatTensor]=None):
331
  return_dict = return_dict if return_dict is not None else self.config.return_dict
332
  use_cache = use_cache if use_cache is not None else self.config.use_cache
333
+ outputs = self.transformer(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id, return_dict=return_dict, output_attentions=output_attentions, output_hidden_states=output_hidden_states, use_cache=use_cache,
334
+ inputs_embeds=inputs_embeds)
 
335
  logits = self.transformer.wte(outputs.last_hidden_state.to(self.transformer.wte.weight.device), True)
336
  if self.logit_scale is not None:
337
  if self.logit_scale == 0:
 
377
  @staticmethod
378
  def _reorder_cache(past_key_values, beam_idx):
379
  """Used by HuggingFace generate when using beam search with kv-caching.
 
380
  See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133
381
  for an example in transformers.
382
  """
383
  reordered_past = []
384
  for layer_past in past_key_values:
385
  reordered_past += [tuple((past_state.index_select(0, beam_idx) for past_state in layer_past))]
386
+ return reordered_past