kaizen9 commited on
Commit
7aaf32c
·
verified ·
1 Parent(s): 44f8bb2

Update modeling_mpt.py

Browse files
Files changed (1) hide show
  1. modeling_mpt.py +254 -63
modeling_mpt.py CHANGED
@@ -1,16 +1,31 @@
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
6
  import warnings
7
  from typing import Any, Dict, List, Mapping, MutableMapping, Optional, Tuple, Union
8
  import torch
9
  import torch.nn as nn
10
  import torch.nn.functional as F
 
 
 
 
 
 
 
 
 
 
 
 
11
  from transformers import PreTrainedModel, PreTrainedTokenizerBase
12
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
13
- from .attention import attn_bias_shape, build_attn_bias
 
 
 
14
  from .blocks import MPTBlock
15
  from .custom_embedding import SharedEmbedding
16
  from .fc import FC_CLASS_REGISTRY as FC_CLASS_REGISTRY
@@ -30,11 +45,127 @@ except:
30
  import logging
31
  log = logging.getLogger(__name__)
32
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
33
  class MPTPreTrainedModel(PreTrainedModel):
34
  config_class = MPTConfig
35
  base_model_prefix = 'model'
36
  _no_split_modules = ['MPTBlock']
37
 
 
 
 
38
  class MPTModel(MPTPreTrainedModel):
39
 
40
  def __init__(self, config: MPTConfig):
@@ -62,6 +193,11 @@ class MPTModel(MPTPreTrainedModel):
62
  self.emb_drop = nn.Dropout(config.emb_pdrop)
63
  self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
64
  self.norm_f = norm_class(config.d_model, device=config.init_device)
 
 
 
 
 
65
  if config.init_device != 'meta':
66
  log.info(f'We recommend using config.init_device="meta" with Composer + FSDP for faster initialization.')
67
  self.apply(self.param_init_fn)
@@ -72,15 +208,18 @@ class MPTModel(MPTPreTrainedModel):
72
  if config.no_bias:
73
  for module in self.modules():
74
  if hasattr(module, 'bias') and isinstance(module.bias, nn.Parameter):
75
- log.info(f'Removing bias ({module.bias}) from {module}.')
76
  module.register_parameter('bias', None)
 
 
 
77
  log.debug(self)
78
  log.debug(f"Using {self.config.init_config['name']} initialization.")
79
 
80
- def get_input_embeddings(self) -> nn.Embedding:
81
  return self.wte
82
 
83
- def set_input_embeddings(self, value: nn.Embedding) -> None:
84
  self.wte = value
85
 
86
  @torch.no_grad()
@@ -101,7 +240,7 @@ class MPTModel(MPTPreTrainedModel):
101
  attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)
102
  if self.attn_uses_sequence_id and sequence_id is not None:
103
  assert isinstance(attn_bias, torch.Tensor)
104
- attn_bias = self._apply_sequence_id(attn_bias, sequence_id)
105
  if attention_mask is not None:
106
  s_k = attention_mask.shape[-1]
107
  if attn_bias is None:
@@ -113,7 +252,7 @@ class MPTModel(MPTPreTrainedModel):
113
  raise ValueError(f'attention_mask shape={attention_mask.shape} ' + f'and prefix_mask shape={prefix_mask.shape} are not equal.')
114
  min_val = torch.finfo(attn_bias.dtype).min
115
  attn_bias = attn_bias.masked_fill(~attention_mask.view(-1, 1, 1, s_k), min_val)
116
- return (attn_bias, None)
117
 
118
  def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor) -> torch.Tensor:
119
  (s_k, s_q) = attn_bias.shape[-2:]
@@ -130,17 +269,7 @@ class MPTModel(MPTPreTrainedModel):
130
  attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
131
  return attn_bias
132
 
133
- def _apply_sequence_id(self, attn_bias: torch.Tensor, sequence_id: torch.LongTensor) -> torch.Tensor:
134
- seq_len = sequence_id.shape[-1]
135
- if seq_len > self.config.max_seq_len:
136
- raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
137
- attn_bias = attn_bias[..., :seq_len, :seq_len]
138
- cannot_attend = torch.logical_not(torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))).unsqueeze(1)
139
- min_val = torch.finfo(attn_bias.dtype).min
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) -> BaseModelOutputWithPast:
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,33 +285,47 @@ 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)
169
- if self.learned_pos_emb:
170
- past_position = 0
171
- if past_key_values is not None:
172
- if len(past_key_values) != self.config.n_layers:
173
- raise ValueError(f'past_key_values must provide a past_key_value for each attention ' + f'layer in the network (len(past_key_values)={len(past_key_values)!r}; self.config.n_layers={self.config.n_layers!r}).')
174
- past_position = past_key_values[0][0].size(1)
175
- if self.attn_impl == 'torch':
176
- past_position = past_key_values[0][0].size(3)
177
- if S + past_position > self.config.max_seq_len:
178
  raise ValueError(f'Cannot forward input with past sequence length {past_position} and current sequence length ' + f'{S + 1}, this model only supports total sequence length <= {self.config.max_seq_len}.')
179
- pos = torch.arange(past_position, S + past_position, dtype=torch.long, device=input_ids.device).unsqueeze(0)
180
- if attention_mask is not None:
181
- pos = torch.clamp(pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[:, past_position:], min=0)
182
- pos_emb = self.wpe(pos)
183
- x = tok_emb + pos_emb
184
- else:
185
- x = tok_emb
 
 
 
186
  if self.embedding_fraction == 1:
187
  x = self.emb_drop(x)
188
  else:
@@ -190,18 +333,26 @@ class MPTModel(MPTPreTrainedModel):
190
  assert isinstance(self.emb_drop, nn.Module)
191
  x = self.emb_drop(x_shrunk)
192
  (attn_bias, attention_mask) = self._attn_bias(device=x.device, dtype=torch.float32, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id)
 
 
 
 
 
193
  if use_cache and past_key_values is None:
194
  past_key_values = [() for _ in range(self.config.n_layers)]
195
  all_hidden_states = () if output_hidden_states else None
196
  all_self_attns = () if output_attentions else None
 
 
 
197
  for (b_idx, block) in enumerate(self.blocks):
198
  if output_hidden_states:
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, output_attentions=bool(output_attentions))
203
- if past_key_values is not None:
204
- past_key_values[b_idx] = past_key_value
205
  if output_attentions:
206
  assert all_self_attns is not None
207
  all_self_attns = all_self_attns + (attn_weights,)
@@ -209,14 +360,14 @@ class MPTModel(MPTPreTrainedModel):
209
  if output_hidden_states:
210
  assert all_hidden_states is not None
211
  all_hidden_states = all_hidden_states + (x,)
212
- return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values, hidden_states=all_hidden_states, attentions=all_self_attns)
213
 
214
  def param_init_fn(self, module: nn.Module) -> None:
215
  init_fn_name = self.config.init_config['name']
216
  MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
217
 
218
  def fsdp_wrap_fn(self, module: nn.Module) -> bool:
219
- return isinstance(module, MPTBlock)
220
 
221
  def activation_checkpointing_fn(self, module: nn.Module) -> bool:
222
  return isinstance(module, MPTBlock)
@@ -225,10 +376,12 @@ class MPTForCausalLM(MPTPreTrainedModel):
225
 
226
  def __init__(self, config: MPTConfig):
227
  super().__init__(config)
228
- if not config.tie_word_embeddings:
229
- raise ValueError('MPTForCausalLM only supports tied word embeddings')
230
  log.info(f'Instantiating an MPTForCausalLM model from {__file__}')
231
  self.transformer: MPTModel = MPTModel(config)
 
 
 
 
232
  for child in self.transformer.children():
233
  if isinstance(child, torch.nn.ModuleList):
234
  continue
@@ -244,17 +397,28 @@ class MPTForCausalLM(MPTPreTrainedModel):
244
  raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
245
  self.logit_scale = logit_scale
246
 
247
- def get_input_embeddings(self) -> nn.Embedding:
248
- return self.transformer.wte
249
 
250
  def set_input_embeddings(self, value: Union[SharedEmbedding, nn.Embedding]) -> None:
251
- self.transformer.wte = value
252
 
253
- def get_output_embeddings(self) -> nn.Embedding:
254
- return self.transformer.wte
 
 
255
 
256
- def set_output_embeddings(self, new_embeddings: Union[SharedEmbedding, nn.Embedding]) -> None:
257
- self.transformer.wte = new_embeddings
 
 
 
 
 
 
 
 
 
258
 
259
  def set_decoder(self, decoder: MPTModel) -> None:
260
  self.transformer = decoder
@@ -262,13 +426,16 @@ class MPTForCausalLM(MPTPreTrainedModel):
262
  def get_decoder(self) -> MPTModel:
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) -> CausalLMOutputWithPast:
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:
274
  warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
@@ -285,14 +452,34 @@ class MPTForCausalLM(MPTPreTrainedModel):
285
  MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
286
 
287
  def fsdp_wrap_fn(self, module: nn.Module) -> bool:
288
- return isinstance(module, MPTBlock)
289
 
290
  def activation_checkpointing_fn(self, module: nn.Module) -> bool:
291
- return isinstance(module, MPTBlock)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
292
 
293
  def prepare_inputs_for_generation(self, input_ids: torch.Tensor, past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]]=None, inputs_embeds: Optional[torch.Tensor]=None, **kwargs: Any) -> Dict[str, Any]:
294
- if inputs_embeds is not None:
295
- raise NotImplementedError('inputs_embeds is not implemented for MPT yet')
296
  attention_mask = kwargs['attention_mask'].bool()
297
  if attention_mask[:, -1].sum() != attention_mask.shape[0]:
298
  raise NotImplementedError('MPT does not support generation with right padding.')
@@ -308,12 +495,16 @@ class MPTForCausalLM(MPTPreTrainedModel):
308
  raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
309
  else:
310
  prefix_mask = None
311
- return {'input_ids': input_ids, 'attention_mask': attention_mask, 'prefix_mask': prefix_mask, 'sequence_id': sequence_id, 'past_key_values': past_key_values, 'use_cache': kwargs.get('use_cache', True)}
 
 
 
 
 
312
 
313
  @staticmethod
314
  def _reorder_cache(past_key_values: List[Tuple[torch.Tensor, torch.Tensor]], beam_idx: torch.LongTensor) -> List[Tuple[torch.Tensor, ...]]:
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
  """
 
1
  """A simple, flexible implementation of a GPT model.
 
2
  Inspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
3
  """
4
+ from __future__ import annotations
5
  import math
6
  import warnings
7
  from typing import Any, Dict, List, Mapping, MutableMapping, Optional, Tuple, Union
8
  import torch
9
  import torch.nn as nn
10
  import torch.nn.functional as F
11
+ from .attention import is_flash_v1_installed, is_flash_v2_installed
12
+ if is_flash_v2_installed():
13
+ try:
14
+ from flash_attn import bert_padding
15
+ from flash_attn.layers.rotary import RotaryEmbedding as DAILRotaryEmbedding
16
+ except Exception as e:
17
+ raise e
18
+ if is_flash_v1_installed():
19
+ try:
20
+ from flash_attn import bert_padding
21
+ except Exception as e:
22
+ raise e
23
  from transformers import PreTrainedModel, PreTrainedTokenizerBase
24
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
25
+ from transformers.models.llama.modeling_llama import LlamaDynamicNTKScalingRotaryEmbedding as HFDynamicNTKScalingRotaryEmbedding
26
+ from transformers.models.llama.modeling_llama import LlamaLinearScalingRotaryEmbedding as HFLinearScalingRotaryEmbedding
27
+ from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding as HFRotaryEmbedding
28
+ from .attention import ATTN_CLASS_REGISTRY, attn_bias_shape, build_attn_bias, gen_slopes
29
  from .blocks import MPTBlock
30
  from .custom_embedding import SharedEmbedding
31
  from .fc import FC_CLASS_REGISTRY as FC_CLASS_REGISTRY
 
45
  import logging
46
  log = logging.getLogger(__name__)
47
 
48
+ def gen_rotary_embedding(rope_head_dim: int, rope_impl: str, rope_theta: int, rope_dail_config: dict, rope_hf_config: dict, max_seq_len: int):
49
+ if rope_impl == 'dail':
50
+ return DAILRotaryEmbedding(dim=rope_head_dim, base=rope_theta, interleaved=False, scale_base=rope_dail_config['xpos_scale_base'] if rope_dail_config['type'] == 'xpos' else None, pos_idx_in_fp32=rope_dail_config['pos_idx_in_fp32'], device='cpu')
51
+ elif rope_impl == 'hf':
52
+ if rope_hf_config['type'] == 'no_scaling':
53
+ return HFRotaryEmbedding(rope_head_dim, max_position_embeddings=max_seq_len, base=rope_theta, device='cpu')
54
+ elif rope_hf_config['type'] == 'linear':
55
+ return HFLinearScalingRotaryEmbedding(rope_head_dim, max_position_embeddings=max_seq_len, base=rope_theta, scaling_factor=rope_hf_config['factor'], device='cpu')
56
+ elif rope_hf_config['type'] == 'dynamic':
57
+ return HFDynamicNTKScalingRotaryEmbedding(rope_head_dim, max_position_embeddings=max_seq_len, base=rope_theta, scaling_factor=rope_hf_config['factor'], device='cpu')
58
+ raise ValueError('rope_impl needs to be either dail or hf')
59
+
60
+ def gen_attention_mask_in_length(sequence_id: Union[None, torch.Tensor], S: int, attn_uses_sequence_id: bool, attn_impl: str, attention_mask: Union[torch.Tensor, None]):
61
+ """Generates the attention mask used for sequence masking in FA v2.
62
+ Only supports sequence id based sparse attention for no attention masking or attention masking with right padding.
63
+ In case of left padding:
64
+ 1. Training with left padding is not supported in MPT (see https://github.com/mosaicml/llm-foundry/blob/1eecd4cb8e734499f77f6a35f657b8b20c0adfcb/llmfoundry/models/mpt/modeling_mpt.py#L407).
65
+ 2. For generation with left padding, we only have a single sequence id per sample, so we don't need sequence id based sparse attention.
66
+ Args:
67
+ sequence_id (Union[None, torch.Tensor]): Tensor containing the sequence id for each token. Shape (batch_size, seq_len).
68
+ S (int): Sequence length
69
+ attn_uses_sequence_id (bool): Whether the attention uses sequence id based masking.
70
+ attn_impl (str): Attention implementation. This function is only creates attention_mask_in_length for flash attention.
71
+ attention_mask (Union[torch.Tensor, None]): Attention mask tensor of shape (batch_size, seq_len)
72
+ Returns:
73
+ attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is:
74
+ ```
75
+ [
76
+ [2, 3, 0, 0, 0, 0],
77
+ [3, 2, 0, 0, 0, 0],
78
+ [6, 0, 0, 0, 0, 0]
79
+ ]
80
+ ```
81
+ , which refers to the 3D-attention mask:
82
+ ```
83
+ [
84
+ [
85
+ [1, 0, 0, 0, 0, 0],
86
+ [1, 1, 0, 0, 0, 0],
87
+ [0, 0, 1, 0, 0, 0],
88
+ [0, 0, 1, 1, 0, 0],
89
+ [0, 0, 1, 1, 1, 0],
90
+ [0, 0, 0, 0, 0, 1]
91
+ ],
92
+ [
93
+ [1, 0, 0, 0, 0, 0],
94
+ [1, 1, 0, 0, 0, 0],
95
+ [1, 1, 1, 0, 0, 0],
96
+ [0, 0, 0, 1, 0, 0],
97
+ [0, 0, 0, 1, 1, 0],
98
+ [0, 0, 0, 0, 0, 1]
99
+ ],
100
+ [
101
+ [1, 0, 0, 0, 0, 0],
102
+ [1, 1, 0, 0, 0, 0],
103
+ [1, 1, 1, 0, 0, 0],
104
+ [1, 1, 1, 1, 0, 0],
105
+ [1, 1, 1, 1, 1, 0],
106
+ [1, 1, 1, 1, 1, 1]
107
+ ]
108
+ ]
109
+ ```.
110
+ (The description above is taken verbatim from https://github.com/Dao-AILab/flash-attention/blob/9356a1c0389660d7e231ff3163c1ac17d9e3824a/flash_attn/bert_padding.py#L125 .)
111
+ """
112
+ attention_mask_in_length = None
113
+ if sequence_id is not None and attn_uses_sequence_id and (attn_impl == 'flash'):
114
+ if attention_mask is not None and attention_mask[:, 0].sum() != attention_mask.shape[0]:
115
+ raise NotImplementedError('Left padding is not supported with flash attention when attn_uses_sequence_id is set to True.')
116
+ if S != sequence_id.shape[-1]:
117
+ raise ValueError(f'Sequence length ({S}) does not match length of sequences in sequence_id ({sequence_id.shape[-1]}).')
118
+ if attention_mask is not None:
119
+ sequence_id = sequence_id.masked_fill(~attention_mask, 0)
120
+ attention_mask_in_length = torch.nn.functional.one_hot(sequence_id)
121
+ if attention_mask is not None:
122
+ attention_mask_in_length = attention_mask_in_length.masked_fill(~attention_mask.unsqueeze(-1), 0)
123
+ attention_mask_in_length = attention_mask_in_length.sum(dim=1)
124
+ attention_mask_in_length = torch.nn.functional.pad(attention_mask_in_length, (0, S - attention_mask_in_length.shape[-1]), mode='constant', value=0)
125
+ return attention_mask_in_length
126
+
127
+ def gen_flash_attn_padding_info(bsz: int, S: int, past_key_len: int, device: torch.device, attention_mask_in_length: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None):
128
+ flash_attn_padding_info = {}
129
+ if attention_mask_in_length is None:
130
+ key_padding_mask = attention_mask
131
+ if key_padding_mask is None:
132
+ key_padding_mask = torch.ones((bsz, past_key_len + S), dtype=torch.bool, device=device)
133
+ query_padding_mask = key_padding_mask[:, -S:]
134
+ unpadding_function = bert_padding.unpad_input
135
+ else:
136
+ key_padding_mask = attention_mask_in_length
137
+ query_padding_mask = attention_mask_in_length
138
+ unpadding_function = bert_padding.unpad_input_for_concatenated_sequences
139
+ (_, indices_q, cu_seqlens_q, max_seqlen_q) = unpadding_function(torch.empty(bsz, S, 1, device=device), query_padding_mask)
140
+ (_, indices_k, cu_seqlens_k, max_seqlen_k) = unpadding_function(torch.empty(bsz, past_key_len + S, 1, device=device), key_padding_mask)
141
+ (_, indices_v, _, _) = unpadding_function(torch.empty(bsz, past_key_len + S, 1, device=device), key_padding_mask)
142
+ flash_attn_padding_info['indices_q'] = indices_q
143
+ flash_attn_padding_info['indices_k'] = indices_k
144
+ flash_attn_padding_info['indices_v'] = indices_v
145
+ flash_attn_padding_info['cu_seqlens_q'] = cu_seqlens_q
146
+ flash_attn_padding_info['cu_seqlens_k'] = cu_seqlens_k
147
+ flash_attn_padding_info['max_seqlen_q'] = max_seqlen_q
148
+ flash_attn_padding_info['max_seqlen_k'] = max_seqlen_k
149
+ return flash_attn_padding_info
150
+
151
+ def apply_sequence_id(attn_bias: torch.Tensor, sequence_id: torch.LongTensor, max_seq_len: int) -> torch.Tensor:
152
+ seq_len = sequence_id.shape[-1]
153
+ if seq_len > max_seq_len:
154
+ raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={max_seq_len}')
155
+ attn_bias = attn_bias[..., :seq_len, :seq_len]
156
+ cannot_attend = torch.logical_not(torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))).unsqueeze(1)
157
+ min_val = torch.finfo(attn_bias.dtype).min
158
+ attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
159
+ return attn_bias
160
+
161
  class MPTPreTrainedModel(PreTrainedModel):
162
  config_class = MPTConfig
163
  base_model_prefix = 'model'
164
  _no_split_modules = ['MPTBlock']
165
 
166
+ def _fsdp_wrap_fn(self: Union[MPTModel, MPTForCausalLM], module: nn.Module) -> bool:
167
+ return isinstance(module, MPTBlock)
168
+
169
  class MPTModel(MPTPreTrainedModel):
170
 
171
  def __init__(self, config: MPTConfig):
 
193
  self.emb_drop = nn.Dropout(config.emb_pdrop)
194
  self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
195
  self.norm_f = norm_class(config.d_model, device=config.init_device)
196
+ self.rope = config.attn_config['rope']
197
+ self.rope_impl = None
198
+ if self.rope:
199
+ self.rope_impl = config.attn_config['rope_impl']
200
+ self.rotary_embedding = gen_rotary_embedding(rope_head_dim=config.d_model // config.n_heads, rope_impl=self.rope_impl, rope_theta=config.attn_config['rope_theta'], rope_dail_config=config.attn_config['rope_dail_config'], rope_hf_config=config.attn_config['rope_hf_config'], max_seq_len=self.config.max_seq_len)
201
  if config.init_device != 'meta':
202
  log.info(f'We recommend using config.init_device="meta" with Composer + FSDP for faster initialization.')
203
  self.apply(self.param_init_fn)
 
208
  if config.no_bias:
209
  for module in self.modules():
210
  if hasattr(module, 'bias') and isinstance(module.bias, nn.Parameter):
211
+ log.info(f'Removing bias from module={module!r}.')
212
  module.register_parameter('bias', None)
213
+ if hasattr(module, 'use_bias'):
214
+ log.info(f'Setting use_bias=False for module={module!r}.')
215
+ module.use_bias = False
216
  log.debug(self)
217
  log.debug(f"Using {self.config.init_config['name']} initialization.")
218
 
219
+ def get_input_embeddings(self) -> Union[SharedEmbedding, nn.Embedding]:
220
  return self.wte
221
 
222
+ def set_input_embeddings(self, value: Union[SharedEmbedding, nn.Embedding]) -> None:
223
  self.wte = value
224
 
225
  @torch.no_grad()
 
240
  attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)
241
  if self.attn_uses_sequence_id and sequence_id is not None:
242
  assert isinstance(attn_bias, torch.Tensor)
243
+ attn_bias = apply_sequence_id(attn_bias, sequence_id, self.config.max_seq_len)
244
  if attention_mask is not None:
245
  s_k = attention_mask.shape[-1]
246
  if attn_bias is None:
 
252
  raise ValueError(f'attention_mask shape={attention_mask.shape} ' + f'and prefix_mask shape={prefix_mask.shape} are not equal.')
253
  min_val = torch.finfo(attn_bias.dtype).min
254
  attn_bias = attn_bias.masked_fill(~attention_mask.view(-1, 1, 1, s_k), min_val)
255
+ return (attn_bias, attention_mask)
256
 
257
  def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor) -> torch.Tensor:
258
  (s_k, s_q) = attn_bias.shape[-2:]
 
269
  attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
270
  return attn_bias
271
 
272
+ def forward(self, input_ids: Optional[torch.LongTensor]=None, 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) -> BaseModelOutputWithPast:
 
 
 
 
 
 
 
 
 
 
273
  return_dict = return_dict if return_dict is not None else self.config.return_dict
274
  use_cache = use_cache if use_cache is not None else self.config.use_cache
275
  if attention_mask is not None:
 
285
  raise NotImplementedError('MPT does not support training with left padding.')
286
  if self.prefix_lm and prefix_mask is None:
287
  raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
 
 
288
  if self.training:
289
  if self.attn_uses_sequence_id and sequence_id is None:
290
  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.')
291
  elif self.attn_uses_sequence_id is False and sequence_id is not None:
292
  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.')
293
+ if input_ids is not None and inputs_embeds is not None:
294
+ raise ValueError('You cannot specify both input_ids and inputs_embeds.')
295
+ elif input_ids is not None:
296
+ bsz = input_ids.size(0)
297
+ S = input_ids.size(1)
298
+ x = self.wte(input_ids)
299
+ input_device = input_ids.device
300
+ elif inputs_embeds is not None:
301
+ bsz = inputs_embeds.size(0)
302
+ S = inputs_embeds.size(1)
303
+ x = inputs_embeds
304
+ input_device = inputs_embeds.device
305
+ else:
306
+ raise ValueError('You must specify input_ids or inputs_embeds')
307
  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}'
308
+ rotary_emb_w_meta_info = None
309
+ past_position = 0
310
+ if past_key_values is not None:
311
+ if len(past_key_values) != self.config.n_layers:
312
+ raise ValueError(f'past_key_values must provide a past_key_value for each attention ' + f'layer in the network (len(past_key_values)={len(past_key_values)!r}; self.config.n_layers={self.config.n_layers!r}).')
313
+ past_position = past_key_values[0][0].size(1)
314
+ if self.attn_impl == 'torch':
315
+ past_position = past_key_values[0][0].size(3)
316
+ if self.learned_pos_emb or self.rope:
317
+ if self.learned_pos_emb and S + past_position > self.config.max_seq_len:
318
  raise ValueError(f'Cannot forward input with past sequence length {past_position} and current sequence length ' + f'{S + 1}, this model only supports total sequence length <= {self.config.max_seq_len}.')
319
+ if self.learned_pos_emb or (self.rope and self.rope_impl == 'hf'):
320
+ pos = torch.arange(past_position, S + past_position, dtype=torch.long, device=input_device).unsqueeze(0)
321
+ if attention_mask is not None:
322
+ pos = torch.clamp(pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[:, past_position:], min=0)
323
+ if self.learned_pos_emb:
324
+ x = x + self.wpe(pos)
325
+ elif self.rope and self.rope_impl == 'hf':
326
+ rotary_emb_w_meta_info = {'impl': self.rope_impl, 'rotary_emb': self.rotary_embedding, 'offset_info': pos, 'seq_len': S + past_position}
327
+ elif self.rope and self.rope_impl == 'dail':
328
+ rotary_emb_w_meta_info = {'impl': self.rope_impl, 'rotary_emb': self.rotary_embedding, 'offset_info': past_position, 'seq_len': S + past_position}
329
  if self.embedding_fraction == 1:
330
  x = self.emb_drop(x)
331
  else:
 
333
  assert isinstance(self.emb_drop, nn.Module)
334
  x = self.emb_drop(x_shrunk)
335
  (attn_bias, attention_mask) = self._attn_bias(device=x.device, dtype=torch.float32, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id)
336
+ attention_mask_in_length = gen_attention_mask_in_length(sequence_id=sequence_id, S=S, attn_uses_sequence_id=self.attn_uses_sequence_id, attn_impl=self.attn_impl, attention_mask=attention_mask)
337
+ alibi_slopes = None
338
+ if self.alibi and self.attn_impl == 'flash':
339
+ alibi_slopes = gen_slopes(n_heads=self.config.n_heads, alibi_bias_max=self.alibi_bias_max, device=x.device, return_1d=True)
340
+ presents = () if use_cache else None
341
  if use_cache and past_key_values is None:
342
  past_key_values = [() for _ in range(self.config.n_layers)]
343
  all_hidden_states = () if output_hidden_states else None
344
  all_self_attns = () if output_attentions else None
345
+ flash_attn_padding_info = {}
346
+ if self.attn_impl == 'flash':
347
+ flash_attn_padding_info = gen_flash_attn_padding_info(bsz, S, past_position, x.device, attention_mask_in_length, attention_mask)
348
  for (b_idx, block) in enumerate(self.blocks):
349
  if output_hidden_states:
350
  assert all_hidden_states is not None
351
  all_hidden_states = all_hidden_states + (x,)
352
  past_key_value = past_key_values[b_idx] if past_key_values is not None else None
353
+ (x, attn_weights, present) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, rotary_emb_w_meta_info=rotary_emb_w_meta_info, attention_mask=attention_mask, is_causal=self.is_causal, output_attentions=bool(output_attentions), alibi_slopes=alibi_slopes, flash_attn_padding_info=flash_attn_padding_info)
354
+ if presents is not None:
355
+ presents += (present,)
356
  if output_attentions:
357
  assert all_self_attns is not None
358
  all_self_attns = all_self_attns + (attn_weights,)
 
360
  if output_hidden_states:
361
  assert all_hidden_states is not None
362
  all_hidden_states = all_hidden_states + (x,)
363
+ return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attns)
364
 
365
  def param_init_fn(self, module: nn.Module) -> None:
366
  init_fn_name = self.config.init_config['name']
367
  MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
368
 
369
  def fsdp_wrap_fn(self, module: nn.Module) -> bool:
370
+ return _fsdp_wrap_fn(self, module)
371
 
372
  def activation_checkpointing_fn(self, module: nn.Module) -> bool:
373
  return isinstance(module, MPTBlock)
 
376
 
377
  def __init__(self, config: MPTConfig):
378
  super().__init__(config)
 
 
379
  log.info(f'Instantiating an MPTForCausalLM model from {__file__}')
380
  self.transformer: MPTModel = MPTModel(config)
381
+ self.lm_head = None
382
+ if not config.tie_word_embeddings:
383
+ self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False, device=config.init_device)
384
+ self.lm_head._fsdp_wrap = True
385
  for child in self.transformer.children():
386
  if isinstance(child, torch.nn.ModuleList):
387
  continue
 
397
  raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
398
  self.logit_scale = logit_scale
399
 
400
+ def get_input_embeddings(self) -> Union[SharedEmbedding, nn.Embedding]:
401
+ return self.transformer.get_input_embeddings()
402
 
403
  def set_input_embeddings(self, value: Union[SharedEmbedding, nn.Embedding]) -> None:
404
+ self.transformer.set_input_embeddings(value)
405
 
406
+ def get_output_embeddings(self) -> Union[SharedEmbedding, nn.Embedding, nn.Linear]:
407
+ if self.lm_head is not None:
408
+ return self.lm_head
409
+ return self.transformer.get_input_embeddings()
410
 
411
+ def set_output_embeddings(self, new_embeddings: Union[SharedEmbedding, nn.Embedding, nn.Linear]) -> None:
412
+ if self.lm_head is not None:
413
+ self.lm_head = new_embeddings
414
+ else:
415
+ if not isinstance(new_embeddings, (SharedEmbedding, nn.Embedding)):
416
+ raise ValueError('new_embeddings must be an instance of SharedEmbedding ' + f'or nn.Embedding, but got {type(new_embeddings)}.')
417
+ warnings.warn('Using `set_output_embeddings` to set the embedding layer of ' + 'MPTForCausalLM with tied weights. Given weights are tied, ' + 'using `set_input_embeddings` is recommended over using ' + '`set_output_embeddings`.')
418
+ self.transformer.set_input_embeddings(new_embeddings)
419
+
420
+ def tie_weights(self) -> None:
421
+ self.lm_head = None
422
 
423
  def set_decoder(self, decoder: MPTModel) -> None:
424
  self.transformer = decoder
 
426
  def get_decoder(self) -> MPTModel:
427
  return self.transformer
428
 
429
+ def forward(self, input_ids: Optional[torch.LongTensor]=None, 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) -> CausalLMOutputWithPast:
430
  return_dict = return_dict if return_dict is not None else self.config.return_dict
431
  use_cache = use_cache if use_cache is not None else self.config.use_cache
432
+ 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, inputs_embeds=inputs_embeds)
433
+ if self.lm_head is not None:
434
+ logits = self.lm_head(outputs.last_hidden_state)
435
+ else:
436
+ out = outputs.last_hidden_state
437
+ out = out.to(self.transformer.wte.weight.device)
438
+ logits = self.transformer.wte(out, True)
439
  if self.logit_scale is not None:
440
  if self.logit_scale == 0:
441
  warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
 
452
  MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
453
 
454
  def fsdp_wrap_fn(self, module: nn.Module) -> bool:
455
+ return _fsdp_wrap_fn(self, module)
456
 
457
  def activation_checkpointing_fn(self, module: nn.Module) -> bool:
458
+ act_ckpt_list = getattr(self.config, 'activation_checkpointing_target', None) or ['MPTBlock']
459
+ if isinstance(act_ckpt_list, str):
460
+ act_ckpt_list = [act_ckpt_list]
461
+ elif not isinstance(act_ckpt_list, list):
462
+ raise ValueError(f'activation_checkpointing_target must be either a single string or a list, but got {type(act_ckpt_list)}')
463
+ if 'MPTBlock' in act_ckpt_list or 'mptblock' in act_ckpt_list:
464
+ if len(act_ckpt_list) > 1:
465
+ log.info('Activation checkpointing MPTBlock only (ignoring other sub-block modules specified in activation_checkpointing_target).')
466
+ return isinstance(module, MPTBlock)
467
+ mod_types = ()
468
+ for mod_name in act_ckpt_list:
469
+ if mod_name.lower() == 'mptblock':
470
+ mod_types += (MPTBlock,)
471
+ elif mod_name in ATTN_CLASS_REGISTRY:
472
+ mod_types += (ATTN_CLASS_REGISTRY[mod_name],)
473
+ elif mod_name in FFN_CLASS_REGISTRY:
474
+ mod_types += (FFN_CLASS_REGISTRY[mod_name],)
475
+ elif mod_name in NORM_CLASS_REGISTRY:
476
+ mod_types += (NORM_CLASS_REGISTRY[mod_name],)
477
+ else:
478
+ msg = ', '.join(list(ATTN_CLASS_REGISTRY.keys()) + list(FFN_CLASS_REGISTRY.keys()) + list(NORM_CLASS_REGISTRY.keys()) + ['MPTBlock'])
479
+ raise ValueError(f'{mod_name} (specified in activation_checkpointing_target) is not a recognized option out of available options {msg}.')
480
+ return isinstance(module, mod_types)
481
 
482
  def prepare_inputs_for_generation(self, input_ids: torch.Tensor, past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]]=None, inputs_embeds: Optional[torch.Tensor]=None, **kwargs: Any) -> Dict[str, Any]:
 
 
483
  attention_mask = kwargs['attention_mask'].bool()
484
  if attention_mask[:, -1].sum() != attention_mask.shape[0]:
485
  raise NotImplementedError('MPT does not support generation with right padding.')
 
495
  raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
496
  else:
497
  prefix_mask = None
498
+ if inputs_embeds is not None and past_key_values is None:
499
+ model_inputs = {'inputs_embeds': inputs_embeds}
500
+ else:
501
+ model_inputs = {'input_ids': input_ids}
502
+ model_inputs.update({'attention_mask': attention_mask, 'prefix_mask': prefix_mask, 'sequence_id': sequence_id, 'past_key_values': past_key_values, 'use_cache': kwargs.get('use_cache', True)})
503
+ return model_inputs
504
 
505
  @staticmethod
506
  def _reorder_cache(past_key_values: List[Tuple[torch.Tensor, torch.Tensor]], beam_idx: torch.LongTensor) -> List[Tuple[torch.Tensor, ...]]:
507
  """Used by HuggingFace generate when using beam search with kv-caching.
 
508
  See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133
509
  for an example in transformers.
510
  """