Update modeling_mpt.py
Browse files- 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 .
|
|
|
|
|
|
|
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
|
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 =
|
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,
|
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
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
169 |
-
|
170 |
-
|
171 |
-
if past_key_values
|
172 |
-
|
173 |
-
|
174 |
-
|
175 |
-
|
176 |
-
|
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 |
-
|
180 |
-
|
181 |
-
|
182 |
-
|
183 |
-
|
184 |
-
|
185 |
-
|
|
|
|
|
|
|
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,
|
203 |
-
if
|
204 |
-
|
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=
|
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
|
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.
|
249 |
|
250 |
def set_input_embeddings(self, value: Union[SharedEmbedding, nn.Embedding]) -> None:
|
251 |
-
self.transformer.
|
252 |
|
253 |
-
def get_output_embeddings(self) -> nn.Embedding:
|
254 |
-
|
|
|
|
|
255 |
|
256 |
-
def set_output_embeddings(self, new_embeddings: Union[SharedEmbedding, nn.Embedding]) -> None:
|
257 |
-
self.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
269 |
-
|
270 |
-
|
271 |
-
|
|
|
|
|
|
|
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
|
289 |
|
290 |
def activation_checkpointing_fn(self, module: nn.Module) -> bool:
|
291 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
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 |
"""
|