irenedea commited on
Commit
5f68131
1 Parent(s): 56bcbea

LLM-foundry update February 07, 2024 19:44:25

Browse files
Files changed (6) hide show
  1. attention.py +109 -59
  2. blocks.py +19 -5
  3. configuration_mpt.py +61 -18
  4. ffn.py +72 -14
  5. modeling_mpt.py +251 -59
  6. warnings.py +22 -0
attention.py CHANGED
@@ -1,21 +1,23 @@
1
  """Attention layers."""
2
  import math
3
  import warnings
4
- from typing import Any, List, Optional, Tuple
5
  import torch
6
  import torch.nn as nn
 
7
  from einops import rearrange
8
  from packaging import version
9
  from torch import nn
10
  from .fc import FC_CLASS_REGISTRY
11
  from .norm import NORM_CLASS_REGISTRY
12
 
13
- def is_flash_v2_installed():
 
14
  try:
15
  import flash_attn as flash_attn
16
  except:
17
  return False
18
- return version.parse(flash_attn.__version__) >= version.parse('2.0.0')
19
 
20
  def is_flash_v1_installed():
21
  try:
@@ -24,6 +26,16 @@ def is_flash_v1_installed():
24
  return False
25
  return version.parse(flash_attn.__version__) < version.parse('2.0.0')
26
 
 
 
 
 
 
 
 
 
 
 
27
  def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool) -> bool:
28
  if original_is_causal and num_query_tokens != num_key_tokens:
29
  if num_query_tokens != 1:
@@ -45,13 +57,7 @@ def repeat_kv_for_gqa(hidden: torch.Tensor, n_rep: int) -> torch.Tensor:
45
  hidden = hidden[:, :, :, None, :].expand(b, s, kv_n_heads, n_rep, d)
46
  return hidden.reshape(b, s, kv_n_heads * n_rep, d)
47
 
48
- def scaled_multihead_dot_product_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: Optional[int]=None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False, multiquery: bool=False) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
49
- if multiquery:
50
- warnings.warn(DeprecationWarning('The direct use of the multiquery arg is deprecated. Setting kv_n_heads=1 automatically. Please set kv_n_heads=1 explicitly to remove this warning.'))
51
- kv_n_heads = 1
52
- elif kv_n_heads is None:
53
- warnings.warn(DeprecationWarning('Not specifying a value for the kv_n_heads arg is deprecated. Setting kv_n_heads=n_heads automatically. Please set kv_n_heads=n_heads explicitly to remove this warning.'))
54
- kv_n_heads = n_heads
55
  q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)
56
  k = rearrange(key, 'b s (h d) -> b h d s', h=kv_n_heads)
57
  v = rearrange(value, 'b s (h d) -> b h s d', h=kv_n_heads)
@@ -97,7 +103,7 @@ def scaled_multihead_dot_product_attention(query: torch.Tensor, key: torch.Tenso
97
  return (out, attn_weight, past_key_value)
98
  return (out, None, past_key_value)
99
 
100
- def check_valid_inputs(*tensors: torch.Tensor, valid_dtypes: Optional[List[torch.dtype]]=None):
101
  if valid_dtypes is None:
102
  valid_dtypes = [torch.float16, torch.bfloat16]
103
  for tensor in tensors:
@@ -106,57 +112,64 @@ def check_valid_inputs(*tensors: torch.Tensor, valid_dtypes: Optional[List[torch
106
  if not tensor.is_cuda:
107
  raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
108
 
109
- def flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: Optional[int]=None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False, multiquery: bool=False) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
 
 
 
 
 
110
  try:
111
  from flash_attn import bert_padding, flash_attn_interface
112
  except:
113
- raise RuntimeError('Please install flash-attn==1.0.9 or flash-attn==2.3.2')
114
  check_valid_inputs(query, key, value)
115
- if multiquery:
116
- warnings.warn(DeprecationWarning('The direct use of the multiquery arg is deprecated. Setting kv_n_heads=1 automatically. Please set kv_n_heads=1 explicitly to remove this warning.'))
117
- kv_n_heads = 1
118
- elif kv_n_heads is None:
119
- warnings.warn(DeprecationWarning('Not specifying a value for the kv_n_heads arg is deprecated. Setting kv_n_heads=n_heads automatically. Please set kv_n_heads=n_heads explicitly to remove this warning.'))
120
- kv_n_heads = n_heads
121
  if past_key_value is not None:
122
  if len(past_key_value) != 0:
123
  key = torch.cat([past_key_value[0], key], dim=1)
124
  value = torch.cat([past_key_value[1], value], dim=1)
125
  past_key_value = (key, value)
126
- if attn_bias is not None:
127
- _s_q = max(0, attn_bias.size(2) - query.size(1))
128
- _s_k = max(0, attn_bias.size(3) - key.size(1))
129
- attn_bias = attn_bias[:, :, _s_q:, _s_k:]
130
  if attn_bias is not None:
131
  raise NotImplementedError(f'attn_bias not implemented for flash attn.')
132
  (batch_size, seqlen) = query.shape[:2]
133
- if key_padding_mask is None:
134
- key_padding_mask = torch.ones_like(key[:, :, 0], dtype=torch.bool)
135
- query_padding_mask = key_padding_mask[:, -query.size(1):]
136
- (query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(query, query_padding_mask)
 
 
 
 
137
  query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
138
- (key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(key, key_padding_mask)
139
  key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=kv_n_heads)
140
- (value_unpad, _, _, _) = bert_padding.unpad_input(value, key_padding_mask)
141
  value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=kv_n_heads)
142
- if kv_n_heads == 1:
143
- key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size(-1))
144
- value_unpad = value_unpad.expand(value_unpad.size(0), n_heads, value_unpad.size(-1))
145
- elif kv_n_heads < n_heads:
146
- key_unpad = repeat_kv_for_gqa(key_unpad.view(batch_size, seqlen, kv_n_heads, -1), n_heads // kv_n_heads).view(batch_size * seqlen, n_heads, -1)
147
- value_unpad = repeat_kv_for_gqa(value_unpad.view(batch_size, seqlen, kv_n_heads, -1), n_heads // kv_n_heads).view(batch_size * seqlen, n_heads, -1)
 
 
 
148
  dropout_p = dropout_p if training else 0.0
149
  reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
150
  if is_flash_v1_installed():
151
  output_unpad = flash_attn_interface.flash_attn_unpadded_func(q=query_unpad, k=key_unpad, v=value_unpad, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
152
  elif is_flash_v2_installed():
153
- output_unpad = flash_attn_interface.flash_attn_varlen_func(q=query_unpad, k=key_unpad, v=value_unpad, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
 
 
 
 
 
154
  else:
155
- raise RuntimeError('flash-attn==1.0.9 or flash-attn==2.3.2 is required.')
156
  output = bert_padding.pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size, seqlen)
157
  return (output, None, past_key_value)
158
 
159
- def triton_flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: Optional[int]=None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False, multiquery: bool=False) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
160
  try:
161
  from .flash_attn_triton import flash_attn_func
162
  except:
@@ -170,12 +183,6 @@ def triton_flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Te
170
  if not _installed:
171
  raise RuntimeError('Requirements for `attn_impl: triton` not installed. Either (1) have a CUDA-compatible GPU ' + 'and `pip install .[gpu]` if installing from llm-foundry source or ' + '`pip install triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir#subdirectory=python` ' + 'if installing from pypi, or (2) use torch attn model.attn_config.attn_impl=torch (torch attn_impl will be slow). ' + 'Note: (1) requires you have CMake and PyTorch already installed.')
172
  check_valid_inputs(query, key, value)
173
- if multiquery:
174
- warnings.warn(DeprecationWarning('The direct use of the multiquery arg is deprecated. Setting kv_n_heads=1 automatically. Please set kv_n_heads=1 explicitly to remove this warning.'))
175
- kv_n_heads = 1
176
- elif kv_n_heads is None:
177
- warnings.warn(DeprecationWarning('Not specifying a value for the kv_n_heads arg is deprecated. Setting kv_n_heads=n_heads automatically. Please set kv_n_heads=n_heads explicitly to remove this warning.'))
178
- kv_n_heads = n_heads
179
  if past_key_value is not None:
180
  if len(past_key_value) != 0:
181
  key = torch.cat([past_key_value[0], key], dim=1)
@@ -220,14 +227,16 @@ class GroupedQueryAttention(nn.Module):
220
  implementation enables user to also use additive bias.
221
  """
222
 
223
- def __init__(self, d_model: int, n_heads: int, kv_n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True):
224
  super().__init__()
225
  self.attn_impl = attn_impl
226
  self.clip_qkv = clip_qkv
227
  self.qk_ln = qk_ln
 
228
  self.d_model = d_model
229
  self.n_heads = n_heads
230
  self.kv_n_heads = kv_n_heads
 
231
  self.head_dim = d_model // n_heads
232
  if self.kv_n_heads <= 0:
233
  raise ValueError('kv_n_heads should be greater than zero.')
@@ -235,6 +244,8 @@ class GroupedQueryAttention(nn.Module):
235
  raise ValueError('The number of KV heads should be less than or equal to Q heads.')
236
  if self.n_heads % self.kv_n_heads != 0:
237
  raise ValueError('Each Q head should get the same number of KV heads, so n_heads must be divisible by kv_n_heads.')
 
 
238
  self.softmax_scale = softmax_scale
239
  if self.softmax_scale is None:
240
  self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
@@ -245,10 +256,13 @@ class GroupedQueryAttention(nn.Module):
245
  self.Wqkv = FC_CLASS_REGISTRY[fc_type](self.d_model, self.d_model + 2 * self.kv_n_heads * self.head_dim, **fc_kwargs)
246
  fuse_splits = [i * self.head_dim for i in range(1, self.n_heads + 2 * self.kv_n_heads)]
247
  self.Wqkv._fused = (0, fuse_splits)
248
- if self.qk_ln:
249
  norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
250
- self.q_ln = norm_class(self.d_model, device=device)
251
- self.k_ln = norm_class(self.kv_n_heads * self.head_dim, device=device)
 
 
 
252
  if self.attn_impl == 'flash':
253
  self.attn_fn = flash_attn_fn
254
  elif self.attn_impl == 'triton':
@@ -260,17 +274,51 @@ class GroupedQueryAttention(nn.Module):
260
  self.out_proj = FC_CLASS_REGISTRY[fc_type](self.d_model, self.d_model, **fc_kwargs)
261
  self.out_proj._is_residual = True
262
 
263
- def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, is_causal: bool=True, needs_weights: bool=False) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
264
  qkv = self.Wqkv(x)
265
  if self.clip_qkv:
266
  qkv = qkv.clamp(min=-self.clip_qkv, max=self.clip_qkv)
267
  (query, key, value) = qkv.split([self.d_model, self.kv_n_heads * self.head_dim, self.kv_n_heads * self.head_dim], dim=2)
268
  key_padding_mask = attention_mask
269
- if self.qk_ln:
 
 
 
 
 
270
  dtype = query.dtype
271
- query = self.q_ln(query).to(dtype)
272
- key = self.k_ln(key).to(dtype)
273
- (context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, self.kv_n_heads, past_key_value=past_key_value, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
274
  return (self.out_proj(context), attn_weights, past_key_value)
275
 
276
  class MultiheadAttention(GroupedQueryAttention):
@@ -280,8 +328,8 @@ class MultiheadAttention(GroupedQueryAttention):
280
  additive bias.
281
  """
282
 
283
- def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True):
284
- super().__init__(d_model=d_model, n_heads=n_heads, kv_n_heads=n_heads, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device, bias=bias)
285
 
286
  class MultiQueryAttention(GroupedQueryAttention):
287
  """Multi-Query self attention.
@@ -290,10 +338,10 @@ class MultiQueryAttention(GroupedQueryAttention):
290
  additive bias.
291
  """
292
 
293
- def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True):
294
- super().__init__(d_model=d_model, n_heads=n_heads, kv_n_heads=1, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device, bias=bias)
295
 
296
- def attn_bias_shape(attn_impl: str, n_heads: int, seq_len: int, alibi: bool, prefix_lm: bool, causal: bool, use_sequence_id: bool) -> Optional[Tuple[int, int, int, int]]:
297
  if attn_impl == 'flash':
298
  return None
299
  elif attn_impl in ['torch', 'triton']:
@@ -318,13 +366,15 @@ def build_attn_bias(attn_impl: str, attn_bias: torch.Tensor, n_heads: int, seq_l
318
  else:
319
  raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
320
 
321
- def gen_slopes(n_heads: int, alibi_bias_max: int=8, device: Optional[torch.device]=None) -> torch.Tensor:
322
  _n_heads = 2 ** math.ceil(math.log2(n_heads))
323
  m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)
324
  m = m.mul(alibi_bias_max / _n_heads)
325
  slopes = 1.0 / torch.pow(2, m)
326
  if _n_heads != n_heads:
327
  slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
 
 
328
  return slopes.view(1, n_heads, 1, 1)
329
 
330
  def build_alibi_bias(n_heads: int, seq_len: int, full: bool=False, alibi_bias_max: int=8, device: Optional[torch.device]=None, dtype: Optional[torch.dtype]=None) -> torch.Tensor:
 
1
  """Attention layers."""
2
  import math
3
  import warnings
4
+ from typing import Any, Optional
5
  import torch
6
  import torch.nn as nn
7
+ import transformers
8
  from einops import rearrange
9
  from packaging import version
10
  from torch import nn
11
  from .fc import FC_CLASS_REGISTRY
12
  from .norm import NORM_CLASS_REGISTRY
13
 
14
+ def is_flash_v2_installed(v2_version: str='2.0.0'):
15
+ assert version.parse(v2_version) >= version.parse('2.0.0')
16
  try:
17
  import flash_attn as flash_attn
18
  except:
19
  return False
20
+ return version.parse(flash_attn.__version__) >= version.parse(v2_version)
21
 
22
  def is_flash_v1_installed():
23
  try:
 
26
  return False
27
  return version.parse(flash_attn.__version__) < version.parse('2.0.0')
28
 
29
+ def is_transformers_version_gte(hf_version: str) -> bool:
30
+ return version.parse(transformers.__version__) >= version.parse(hf_version)
31
+
32
+ def check_alibi_support(attention_impl: str) -> bool:
33
+ return attention_impl != 'flash' or is_flash_v2_installed(v2_version='v2.4.2')
34
+ if is_flash_v1_installed():
35
+ import transformers
36
+ transformers.utils.is_flash_attn_available = lambda : False
37
+ from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
38
+
39
  def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool) -> bool:
40
  if original_is_causal and num_query_tokens != num_key_tokens:
41
  if num_query_tokens != 1:
 
57
  hidden = hidden[:, :, :, None, :].expand(b, s, kv_n_heads, n_rep, d)
58
  return hidden.reshape(b, s, kv_n_heads * n_rep, d)
59
 
60
+ def scaled_multihead_dot_product_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: int, past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor, torch.Tensor]]]:
 
 
 
 
 
 
61
  q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)
62
  k = rearrange(key, 'b s (h d) -> b h d s', h=kv_n_heads)
63
  v = rearrange(value, 'b s (h d) -> b h s d', h=kv_n_heads)
 
103
  return (out, attn_weight, past_key_value)
104
  return (out, None, past_key_value)
105
 
106
+ def check_valid_inputs(*tensors: torch.Tensor, valid_dtypes: Optional[list[torch.dtype]]=None):
107
  if valid_dtypes is None:
108
  valid_dtypes = [torch.float16, torch.bfloat16]
109
  for tensor in tensors:
 
112
  if not tensor.is_cuda:
113
  raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
114
 
115
+ def flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: int, past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False, multiquery: bool=False, should_repeat_kv_for_gqa: Optional[bool]=True, sliding_window_size: int=-1, alibi_slopes: Optional[torch.Tensor]=None, flash_attn_padding_info: Optional[dict[str, torch.Tensor]]=None) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor, torch.Tensor]]]:
116
+ if key_padding_mask is not None:
117
+ raise ValueError('key_padding_mask should be None for flash attn.')
118
+ del key_padding_mask
119
+ if flash_attn_padding_info is None:
120
+ raise ValueError('flash_attn_padding_info is required for flash attn.')
121
  try:
122
  from flash_attn import bert_padding, flash_attn_interface
123
  except:
124
+ raise RuntimeError('Please install flash-attn==1.0.9 or flash-attn==2.3.6')
125
  check_valid_inputs(query, key, value)
 
 
 
 
 
 
126
  if past_key_value is not None:
127
  if len(past_key_value) != 0:
128
  key = torch.cat([past_key_value[0], key], dim=1)
129
  value = torch.cat([past_key_value[1], value], dim=1)
130
  past_key_value = (key, value)
 
 
 
 
131
  if attn_bias is not None:
132
  raise NotImplementedError(f'attn_bias not implemented for flash attn.')
133
  (batch_size, seqlen) = query.shape[:2]
134
+ indices_q = flash_attn_padding_info['indices_q']
135
+ indices_k = flash_attn_padding_info['indices_k']
136
+ indices_v = flash_attn_padding_info['indices_v']
137
+ cu_seqlens_q = flash_attn_padding_info['cu_seqlens_q']
138
+ cu_seqlens_k = flash_attn_padding_info['cu_seqlens_k']
139
+ max_seqlen_q = flash_attn_padding_info['max_seqlen_q']
140
+ max_seqlen_k = flash_attn_padding_info['max_seqlen_k']
141
+ query_unpad = bert_padding.index_first_axis(rearrange(query, 'b s ... -> (b s) ...'), indices_q)
142
  query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
143
+ key_unpad = bert_padding.index_first_axis(rearrange(key, 'b s ... -> (b s) ...'), indices_k)
144
  key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=kv_n_heads)
145
+ value_unpad = bert_padding.index_first_axis(rearrange(value, 'b s ... -> (b s) ...'), indices_v)
146
  value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=kv_n_heads)
147
+ if kv_n_heads < n_heads and (not is_flash_v2_installed()) and (not should_repeat_kv_for_gqa):
148
+ raise ValueError('For Grouped Query Attention or Multi Query Attention, should_repeat_kv_for_gqa should be set to True if not using Flash Attention v2.')
149
+ if should_repeat_kv_for_gqa:
150
+ if kv_n_heads == 1:
151
+ key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size(-1))
152
+ value_unpad = value_unpad.expand(value_unpad.size(0), n_heads, value_unpad.size(-1))
153
+ elif kv_n_heads < n_heads:
154
+ key_unpad = repeat_kv_for_gqa(key_unpad.view(1, key_unpad.size(0), kv_n_heads, -1), n_heads // kv_n_heads).view(key_unpad.size(0), n_heads, -1)
155
+ value_unpad = repeat_kv_for_gqa(value_unpad.view(1, value_unpad.size(0), kv_n_heads, -1), n_heads // kv_n_heads).view(value_unpad.size(0), n_heads, -1)
156
  dropout_p = dropout_p if training else 0.0
157
  reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
158
  if is_flash_v1_installed():
159
  output_unpad = flash_attn_interface.flash_attn_unpadded_func(q=query_unpad, k=key_unpad, v=value_unpad, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
160
  elif is_flash_v2_installed():
161
+ alibi_kwargs = {}
162
+ if check_alibi_support('flash'):
163
+ alibi_kwargs = {'alibi_slopes': alibi_slopes}
164
+ elif alibi_slopes is not None:
165
+ raise ValueError('alibi_slopes is only supported for flash-attn>=2.4.2')
166
+ output_unpad = flash_attn_interface.flash_attn_varlen_func(q=query_unpad, k=key_unpad, v=value_unpad, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights, window_size=(sliding_window_size, sliding_window_size), **alibi_kwargs)
167
  else:
168
+ raise RuntimeError('flash-attn==1.0.9 or flash-attn==2.4.2 is required.')
169
  output = bert_padding.pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size, seqlen)
170
  return (output, None, past_key_value)
171
 
172
+ def triton_flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: int, past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor, torch.Tensor]]]:
173
  try:
174
  from .flash_attn_triton import flash_attn_func
175
  except:
 
183
  if not _installed:
184
  raise RuntimeError('Requirements for `attn_impl: triton` not installed. Either (1) have a CUDA-compatible GPU ' + 'and `pip install .[gpu]` if installing from llm-foundry source or ' + '`pip install triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir#subdirectory=python` ' + 'if installing from pypi, or (2) use torch attn model.attn_config.attn_impl=torch (torch attn_impl will be slow). ' + 'Note: (1) requires you have CMake and PyTorch already installed.')
185
  check_valid_inputs(query, key, value)
 
 
 
 
 
 
186
  if past_key_value is not None:
187
  if len(past_key_value) != 0:
188
  key = torch.cat([past_key_value[0], key], dim=1)
 
227
  implementation enables user to also use additive bias.
228
  """
229
 
230
+ def __init__(self, d_model: int, n_heads: int, kv_n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, qk_gn: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True, sliding_window_size: int=-1):
231
  super().__init__()
232
  self.attn_impl = attn_impl
233
  self.clip_qkv = clip_qkv
234
  self.qk_ln = qk_ln
235
+ self.qk_gn = qk_gn
236
  self.d_model = d_model
237
  self.n_heads = n_heads
238
  self.kv_n_heads = kv_n_heads
239
+ self.sliding_window_size = sliding_window_size
240
  self.head_dim = d_model // n_heads
241
  if self.kv_n_heads <= 0:
242
  raise ValueError('kv_n_heads should be greater than zero.')
 
244
  raise ValueError('The number of KV heads should be less than or equal to Q heads.')
245
  if self.n_heads % self.kv_n_heads != 0:
246
  raise ValueError('Each Q head should get the same number of KV heads, so n_heads must be divisible by kv_n_heads.')
247
+ if qk_ln and qk_gn:
248
+ raise ValueError('Only one of qk_ln and qk_gn can be set to True.')
249
  self.softmax_scale = softmax_scale
250
  if self.softmax_scale is None:
251
  self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
 
256
  self.Wqkv = FC_CLASS_REGISTRY[fc_type](self.d_model, self.d_model + 2 * self.kv_n_heads * self.head_dim, **fc_kwargs)
257
  fuse_splits = [i * self.head_dim for i in range(1, self.n_heads + 2 * self.kv_n_heads)]
258
  self.Wqkv._fused = (0, fuse_splits)
259
+ if self.qk_ln or self.qk_gn:
260
  norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
261
+ norm_size = self.head_dim if qk_gn else d_model
262
+ self.q_ln = norm_class(norm_size, device=device)
263
+ if qk_ln:
264
+ norm_size = self.head_dim * kv_n_heads
265
+ self.k_ln = norm_class(norm_size, device=device)
266
  if self.attn_impl == 'flash':
267
  self.attn_fn = flash_attn_fn
268
  elif self.attn_impl == 'triton':
 
274
  self.out_proj = FC_CLASS_REGISTRY[fc_type](self.d_model, self.d_model, **fc_kwargs)
275
  self.out_proj._is_residual = True
276
 
277
+ def forward(self, x: torch.Tensor, past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, rotary_emb_w_meta_info: Optional[dict]=None, is_causal: bool=True, needs_weights: bool=False, alibi_slopes: Optional[torch.Tensor]=None, flash_attn_padding_info: Optional[dict[str, torch.Tensor]]=None) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor, torch.Tensor]]]:
278
  qkv = self.Wqkv(x)
279
  if self.clip_qkv:
280
  qkv = qkv.clamp(min=-self.clip_qkv, max=self.clip_qkv)
281
  (query, key, value) = qkv.split([self.d_model, self.kv_n_heads * self.head_dim, self.kv_n_heads * self.head_dim], dim=2)
282
  key_padding_mask = attention_mask
283
+ if self.qk_ln or self.qk_gn:
284
+ (q_shape, k_shape) = (query.shape, key.shape)
285
+ if self.qk_gn:
286
+ (b, s) = query.shape[:2]
287
+ query = query.view(b, s, self.n_heads, -1)
288
+ key = key.view(b, s, self.kv_n_heads, -1)
289
  dtype = query.dtype
290
+ query = self.q_ln(query).to(dtype).view(q_shape)
291
+ key = self.k_ln(key).to(dtype).view(k_shape)
292
+ if rotary_emb_w_meta_info is not None:
293
+ rotary_emb = rotary_emb_w_meta_info['rotary_emb']
294
+ seq_len = rotary_emb_w_meta_info['seq_len']
295
+ offset_info = rotary_emb_w_meta_info['offset_info']
296
+ (bsz, seqlen) = query.shape[:2]
297
+ query = query.view(bsz, seqlen, -1, self.head_dim)
298
+ key = key.view(bsz, seqlen, -1, self.head_dim)
299
+ if rotary_emb_w_meta_info['impl'] == 'dail':
300
+ value = value.view(bsz, seqlen, -1, self.head_dim)
301
+ kv = torch.stack([key, value], dim=2)
302
+ (query, kv) = rotary_emb(query, kv, seqlen_offset=offset_info, max_seqlen=seq_len)
303
+ [key, value] = torch.unbind(kv, dim=2)
304
+ value = value.view(bsz, seqlen, self.kv_n_heads * self.head_dim)
305
+ elif rotary_emb_w_meta_info['impl'] == 'hf':
306
+ (cos, sin) = rotary_emb(value, seq_len)
307
+ if is_transformers_version_gte('4.36'):
308
+ (query, key) = apply_rotary_pos_emb(query, key, cos, sin, offset_info, unsqueeze_dim=2)
309
+ else:
310
+ query = query.transpose(1, 2)
311
+ key = key.transpose(1, 2)
312
+ (query, key) = apply_rotary_pos_emb(query, key, cos, sin, offset_info)
313
+ query = query.transpose(1, 2)
314
+ key = key.transpose(1, 2)
315
+ query = query.view(bsz, seqlen, self.d_model)
316
+ key = key.view(bsz, seqlen, self.kv_n_heads * self.head_dim)
317
+ extra_attn_kwargs = {}
318
+ if self.attn_impl == 'flash':
319
+ key_padding_mask = None
320
+ extra_attn_kwargs = {'should_repeat_kv_for_gqa': not is_flash_v2_installed(), 'sliding_window_size': self.sliding_window_size, 'alibi_slopes': alibi_slopes, 'flash_attn_padding_info': flash_attn_padding_info}
321
+ (context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, self.kv_n_heads, past_key_value=past_key_value, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights, **extra_attn_kwargs)
322
  return (self.out_proj(context), attn_weights, past_key_value)
323
 
324
  class MultiheadAttention(GroupedQueryAttention):
 
328
  additive bias.
329
  """
330
 
331
+ def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, qk_gn: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True, sliding_window_size: int=-1):
332
+ super().__init__(d_model=d_model, n_heads=n_heads, kv_n_heads=n_heads, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, qk_gn=qk_gn, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device, bias=bias, sliding_window_size=sliding_window_size)
333
 
334
  class MultiQueryAttention(GroupedQueryAttention):
335
  """Multi-Query self attention.
 
338
  additive bias.
339
  """
340
 
341
+ def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, qk_gn: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True, sliding_window_size: int=-1):
342
+ super().__init__(d_model=d_model, n_heads=n_heads, kv_n_heads=1, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, qk_gn=qk_gn, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device, bias=bias, sliding_window_size=sliding_window_size)
343
 
344
+ def attn_bias_shape(attn_impl: str, n_heads: int, seq_len: int, alibi: bool, prefix_lm: bool, causal: bool, use_sequence_id: bool) -> Optional[tuple[int, int, int, int]]:
345
  if attn_impl == 'flash':
346
  return None
347
  elif attn_impl in ['torch', 'triton']:
 
366
  else:
367
  raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
368
 
369
+ def gen_slopes(n_heads: int, alibi_bias_max: int=8, device: Optional[torch.device]=None, return_1d: bool=False) -> torch.Tensor:
370
  _n_heads = 2 ** math.ceil(math.log2(n_heads))
371
  m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)
372
  m = m.mul(alibi_bias_max / _n_heads)
373
  slopes = 1.0 / torch.pow(2, m)
374
  if _n_heads != n_heads:
375
  slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
376
+ if return_1d:
377
+ return slopes
378
  return slopes.view(1, n_heads, 1, 1)
379
 
380
  def build_alibi_bias(n_heads: int, seq_len: int, full: bool=False, alibi_bias_max: int=8, device: Optional[torch.device]=None, dtype: Optional[torch.dtype]=None) -> torch.Tensor:
blocks.py CHANGED
@@ -5,12 +5,17 @@ import torch.nn as nn
5
  from .attention import ATTN_CLASS_REGISTRY
6
  from .ffn import FFN_CLASS_REGISTRY, build_ffn
7
  from .norm import NORM_CLASS_REGISTRY
 
 
 
 
 
8
 
9
  class MPTBlock(nn.Module):
10
 
11
- def __init__(self, d_model: int, n_heads: int, expansion_ratio: int, attn_config: Optional[Dict]=None, ffn_config: Optional[Dict]=None, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, no_bias: bool=False, **kwargs: Any):
12
  if attn_config is None:
13
- attn_config = {'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}
14
  if ffn_config is None:
15
  ffn_config = {'ffn_type': 'mptmlp'}
16
  del kwargs
@@ -18,7 +23,7 @@ class MPTBlock(nn.Module):
18
  norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
19
  assert isinstance(attn_config['attn_type'], str)
20
  attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]
21
- args_to_exclude_in_attn_class = {'attn_type', 'prefix_lm', 'alibi', 'attn_uses_sequence_id', 'alibi_bias_max'}
22
  attn_config_subset_for_attn_class = {k: v for (k, v) in attn_config.items() if k not in args_to_exclude_in_attn_class}
23
  self.norm_1 = norm_class(d_model, device=device)
24
  self.attn = attn_class(d_model=d_model, n_heads=n_heads, fc_type=fc_type, device=device, **attn_config_subset_for_attn_class, bias=not no_bias)
@@ -28,14 +33,23 @@ class MPTBlock(nn.Module):
28
  self.ffn = build_ffn(d_model=d_model, expansion_ratio=expansion_ratio, device=device, bias=not no_bias, **ffn_config)
29
  self.resid_attn_dropout = nn.Dropout(resid_pdrop)
30
  self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
 
31
 
32
- def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.ByteTensor]=None, is_causal: bool=True, output_attentions: bool=False) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
33
  a = self.norm_1(x)
34
- (b, attn_weights, past_key_value) = self.attn(a, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=is_causal, needs_weights=output_attentions)
35
  x = x + self.resid_attn_dropout(b)
36
  m = x
37
  if self.norm_2 is not None:
38
  m = self.norm_2(x)
 
 
 
 
 
39
  n = self.ffn(m)
 
 
 
40
  x = x + self.resid_ffn_dropout(n)
41
  return (x, attn_weights, past_key_value)
 
5
  from .attention import ATTN_CLASS_REGISTRY
6
  from .ffn import FFN_CLASS_REGISTRY, build_ffn
7
  from .norm import NORM_CLASS_REGISTRY
8
+ try:
9
+ from flash_attn.bert_padding import unpad_input, pad_input
10
+ except:
11
+ (unpad_input, pad_input) = (None, None)
12
+ attn_config_defaults: Dict = {'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'qk_gn': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'sliding_window_size': -1, 'alibi': False, 'alibi_bias_max': 8, 'rope': False, 'rope_theta': 10000, 'rope_impl': 'dail', 'rope_dail_config': {'type': 'original', 'pos_idx_in_fp32': True, 'xpos_scale_base': 512}, 'rope_hf_config': {'type': 'no_scaling', 'factor': 1.0}}
13
 
14
  class MPTBlock(nn.Module):
15
 
16
+ def __init__(self, d_model: int, n_heads: int, expansion_ratio: int, attn_config: Optional[Dict]=None, ffn_config: Optional[Dict]=None, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, no_bias: bool=False, use_pad_tok_in_ffn: bool=True, **kwargs: Any):
17
  if attn_config is None:
18
+ attn_config = attn_config_defaults
19
  if ffn_config is None:
20
  ffn_config = {'ffn_type': 'mptmlp'}
21
  del kwargs
 
23
  norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
24
  assert isinstance(attn_config['attn_type'], str)
25
  attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]
26
+ args_to_exclude_in_attn_class = {'attn_type', 'prefix_lm', 'alibi', 'attn_uses_sequence_id', 'alibi_bias_max', 'rope', 'rope_theta', 'rope_impl', 'rope_dail_config', 'rope_hf_config'}
27
  attn_config_subset_for_attn_class = {k: v for (k, v) in attn_config.items() if k not in args_to_exclude_in_attn_class}
28
  self.norm_1 = norm_class(d_model, device=device)
29
  self.attn = attn_class(d_model=d_model, n_heads=n_heads, fc_type=fc_type, device=device, **attn_config_subset_for_attn_class, bias=not no_bias)
 
33
  self.ffn = build_ffn(d_model=d_model, expansion_ratio=expansion_ratio, device=device, bias=not no_bias, **ffn_config)
34
  self.resid_attn_dropout = nn.Dropout(resid_pdrop)
35
  self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
36
+ self.use_pad_tok_in_ffn = use_pad_tok_in_ffn
37
 
38
+ def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, rotary_emb_w_meta_info: Optional[Dict]=None, attention_mask: Optional[torch.ByteTensor]=None, is_causal: bool=True, output_attentions: bool=False, alibi_slopes: Optional[torch.Tensor]=None, flash_attn_padding_info: Optional[dict[str, torch.Tensor]]=None) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
39
  a = self.norm_1(x)
40
+ (b, attn_weights, past_key_value) = self.attn(a, 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=is_causal, needs_weights=output_attentions, alibi_slopes=alibi_slopes, flash_attn_padding_info=flash_attn_padding_info)
41
  x = x + self.resid_attn_dropout(b)
42
  m = x
43
  if self.norm_2 is not None:
44
  m = self.norm_2(x)
45
+ (batch_size, seq_len) = m.size()[:2]
46
+ indices = None
47
+ if not self.use_pad_tok_in_ffn:
48
+ assert unpad_input is not None
49
+ (m, indices, _, _) = unpad_input(m, attention_mask)
50
  n = self.ffn(m)
51
+ if not self.use_pad_tok_in_ffn:
52
+ assert pad_input is not None
53
+ n = pad_input(n, indices, batch_size, seq_len)
54
  x = x + self.resid_ffn_dropout(n)
55
  return (x, attn_weights, past_key_value)
configuration_mpt.py CHANGED
@@ -2,21 +2,26 @@
2
  import warnings
3
  from typing import Any, Dict, Optional, Union
4
  from transformers import PretrainedConfig
5
- attn_config_defaults: Dict = {'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}
 
 
 
 
 
6
  ffn_config_defaults: Dict = {'ffn_type': 'mptmlp'}
7
  init_config_defaults: Dict = {'name': 'kaiming_normal_', 'fan_mode': 'fan_in', 'init_nonlinearity': 'relu', 'init_div_is_residual': True, 'emb_init_std': None, 'emb_init_uniform_lim': None, 'init_std': None, 'init_gain': 0.0}
8
 
9
  class MPTConfig(PretrainedConfig):
10
  model_type = 'mpt'
11
 
12
- def __init__(self, d_model: int=2048, n_heads: int=16, n_layers: int=24, expansion_ratio: int=4, max_seq_len: int=2048, vocab_size: int=50368, resid_pdrop: float=0.0, emb_pdrop: float=0.0, learned_pos_emb: bool=True, attn_config: Dict=attn_config_defaults, ffn_config: Dict=ffn_config_defaults, init_device: str='cpu', logit_scale: Optional[Union[float, str]]=None, no_bias: bool=False, embedding_fraction: float=1.0, norm_type: str='low_precision_layernorm', use_cache: bool=False, init_config: Dict=init_config_defaults, fc_type: str='torch', verbose: Optional[int]=None, **kwargs: Any):
13
  """The MPT configuration class.
14
 
15
  Args:
16
  d_model (int): The size of the embedding dimension of the model.
17
  n_heads (int): The number of attention heads.
18
  n_layers (int): The number of layers in the model.
19
- expansion_ratio (int): The ratio of the up/down scale in the ffn.
20
  max_seq_len (int): The maximum sequence length of the model.
21
  vocab_size (int): The size of the vocabulary.
22
  resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.
@@ -27,6 +32,7 @@ class MPTConfig(PretrainedConfig):
27
  attn_pdrop (float): The dropout probability for the attention layers.
28
  attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.
29
  qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.
 
30
  clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to
31
  this value.
32
  softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,
@@ -38,15 +44,25 @@ class MPTConfig(PretrainedConfig):
38
  When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates
39
  which sub-sequence each token belongs to.
40
  Defaults to ``False`` meaning any provided `sequence_id` will be ignored.
 
41
  alibi (bool): Whether to use the alibi bias instead of position embeddings.
42
  alibi_bias_max (int): The maximum value of the alibi bias.
 
 
 
 
 
 
 
 
 
 
43
  kv_n_heads (Optional[int]): For grouped_query_attention only, allow user to specify number of kv heads.
44
  ffn_config (Dict): A dictionary used to configure the model's ffn module:
45
- ffn_type (str): type of ffn to use. Options: mptmlp, te_ln_mlp
46
  init_device (str): The device to use for parameter initialization.
47
  logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.
48
  no_bias (bool): Whether to use bias in all layers.
49
- verbose (int): The verbosity level. 0 is silent.
50
  embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.
51
  norm_type (str): choose type of norm to use
52
  use_cache (bool): Whether or not the model should return the last key/values attentions
@@ -66,6 +82,8 @@ class MPTConfig(PretrainedConfig):
66
  ---
67
  See llmfoundry.models.utils.param_init_fns.py for info on other param init config options
68
  fc_type (str): choose fc layer implementation. Options: torch and te. te layers support fp8 when using H100 GPUs.
 
 
69
  """
70
  self.d_model = d_model
71
  self.n_heads = n_heads
@@ -86,22 +104,23 @@ class MPTConfig(PretrainedConfig):
86
  self.use_cache = use_cache
87
  self.init_config = init_config
88
  self.fc_type = fc_type
89
- if verbose is not None:
90
- warnings.warn(DeprecationWarning('verbose argument for MPTConfig is now ignored and will be removed. Use python_log_level instead.'))
91
  if 'name' in kwargs:
92
  del kwargs['name']
93
  if 'loss_fn' in kwargs:
94
  del kwargs['loss_fn']
95
- if self.attn_config.get('alibi', False):
96
  self.learned_pos_emb = False
97
- warnings.warn(f'alibi is turned on, setting `learned_pos_emb` to `False.`')
98
- super().__init__(**kwargs)
99
  self._validate_config()
100
 
101
  def _set_config_defaults(self, config: Dict[str, Any], config_defaults: Dict[str, Any]) -> Dict[str, Any]:
102
  for (k, v) in config_defaults.items():
103
  if k not in config:
104
  config[k] = v
 
 
105
  return config
106
 
107
  def _validate_config(self) -> None:
@@ -116,25 +135,49 @@ class MPTConfig(PretrainedConfig):
116
  raise ValueError(f"Unknown attn_impl={self.attn_config['attn_impl']}")
117
  if self.attn_config['prefix_lm'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
118
  raise NotImplementedError('prefix_lm only implemented with torch and triton attention.')
119
- if self.attn_config['alibi'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
120
- raise NotImplementedError('alibi only implemented with torch and triton attention.')
121
- if self.attn_config['attn_uses_sequence_id'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
122
- raise NotImplementedError('attn_uses_sequence_id only implemented with torch and triton attention.')
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
123
  if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
124
  raise ValueError('model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!')
125
  if isinstance(self.logit_scale, str) and self.logit_scale != 'inv_sqrt_d_model':
126
  raise ValueError(f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
127
  if self.init_config.get('name', None) is None:
128
  raise ValueError(f"self.init_config={self.init_config!r} 'name' needs to be set.")
129
- if not self.learned_pos_emb and (not self.attn_config['alibi']):
130
- warnings.warn(f'Positional information not being provided to the model using either learned_pos_emb or alibi.')
131
  if self.fc_type == 'te' or self.ffn_config['ffn_type'] == 'te_ln_mlp':
132
  try:
133
  import transformer_engine.pytorch as te
134
  del te
135
  except:
136
  raise ImportError('TransformerEngine import fail. `fc_type: te` requires TransformerEngine be installed. ' + 'The required version of transformer_engine also requires FlashAttention v1.0.6 is installed:\n' + 'pip install flash-attn==1.0.6 --no-build-isolation \n' + 'pip install git+https://github.com/NVIDIA/TransformerEngine.git@144e4888b2cdd60bd52e706d5b7a79cb9c1a7156')
137
- if self.ffn_config['ffn_type'] == 'mptmlp':
 
 
138
  self.ffn_config['fc_type'] = self.fc_type
139
  elif self.ffn_config['ffn_type'] == 'te_ln_mlp':
140
- self.ffn_config['bias'] = not self.no_bias
 
 
 
 
 
 
 
 
2
  import warnings
3
  from typing import Any, Dict, Optional, Union
4
  from transformers import PretrainedConfig
5
+ from .attention import check_alibi_support, is_flash_v1_installed, is_flash_v2_installed
6
+ from .blocks import attn_config_defaults
7
+ from .fc import FC_CLASS_REGISTRY
8
+ from .norm import LPLayerNorm
9
+ from .ffn import FFN_CLASS_REGISTRY
10
+ from .warnings import VersionedDeprecationWarning
11
  ffn_config_defaults: Dict = {'ffn_type': 'mptmlp'}
12
  init_config_defaults: Dict = {'name': 'kaiming_normal_', 'fan_mode': 'fan_in', 'init_nonlinearity': 'relu', 'init_div_is_residual': True, 'emb_init_std': None, 'emb_init_uniform_lim': None, 'init_std': None, 'init_gain': 0.0}
13
 
14
  class MPTConfig(PretrainedConfig):
15
  model_type = 'mpt'
16
 
17
+ def __init__(self, d_model: int=2048, n_heads: int=16, n_layers: int=24, expansion_ratio: Union[int, float]=4, max_seq_len: int=2048, vocab_size: int=50368, resid_pdrop: float=0.0, emb_pdrop: float=0.0, learned_pos_emb: bool=True, attn_config: Dict=attn_config_defaults, ffn_config: Dict=ffn_config_defaults, init_device: str='cpu', logit_scale: Optional[Union[float, str]]=None, no_bias: bool=False, embedding_fraction: float=1.0, norm_type: str='low_precision_layernorm', use_cache: bool=False, init_config: Dict=init_config_defaults, fc_type: str='torch', tie_word_embeddings: bool=True, use_pad_tok_in_ffn: bool=True, **kwargs: Any):
18
  """The MPT configuration class.
19
 
20
  Args:
21
  d_model (int): The size of the embedding dimension of the model.
22
  n_heads (int): The number of attention heads.
23
  n_layers (int): The number of layers in the model.
24
+ expansion_ratio (Union[int, float]): The ratio of the up/down scale in the ffn.
25
  max_seq_len (int): The maximum sequence length of the model.
26
  vocab_size (int): The size of the vocabulary.
27
  resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.
 
32
  attn_pdrop (float): The dropout probability for the attention layers.
33
  attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.
34
  qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.
35
+ qk_gn (bool): Whether to apply group normalization to the queries and keys in the attention layer.
36
  clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to
37
  this value.
38
  softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,
 
44
  When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates
45
  which sub-sequence each token belongs to.
46
  Defaults to ``False`` meaning any provided `sequence_id` will be ignored.
47
+ sliding_window_size (int): Window size for sliding window local attention. Defaults to -1, which means no sliding window. Query at position i will only attend to keys between [i + seqlen_k - seqlen_q - window_size, i + seqlen_k - seqlen_q + window_size] inclusive. Only works for flash attention v2.3.0 or higher.
48
  alibi (bool): Whether to use the alibi bias instead of position embeddings.
49
  alibi_bias_max (int): The maximum value of the alibi bias.
50
+ rope (bool): Whether to use rotary positional embeddings.
51
+ rope_theta (int): The base frequency for rope.
52
+ rope_impl (str): The implementation of rope to use. One of 'hf' (to use the implementation from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py) or 'dail' (to use the implementation from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/layers/rotary.py).
53
+ rope_dail_config (Dict): The configuration for the dail implementation of rope.
54
+ type (str): The type of rotary position embedding to use. Options: 'original' (for https://arxiv.org/pdf/2104.09864.pdf), 'xpos' (for https://arxiv.org/pdf/2212.10554.pdf).
55
+ pos_idx_in_fp32 (bool): If True, the position indices [0, ..., seqlen - 1] are in fp32, otherwise they might be in lower precision. A consequence could be, for example, that bf16 rounds position 1995 to 2000, which leads to them having the same positional embedding.
56
+ xpos_scale_base (float): The scale base for XPos (if using XPos).
57
+ rope_hf_config (Dict): A dictionary used to configure rope's scaling behavior (when scaling beyond the training length).
58
+ type (str): Can be one of 'no_scaling', 'linear', or 'dynamic'. 'no_scaling' uses the default implementation for rotary embeddings, 'linear' uses linear scaling as proposed by the Reddit user /u/kaiokendev, and 'dynamic' uses Dynamic NTK scaling as proposed by the Reddit users /u/bloc97 and /u/emozilla.
59
+ factor (float): Scaling factor to use if using 'linear' or 'dynamic' as rope_scaling.type.
60
  kv_n_heads (Optional[int]): For grouped_query_attention only, allow user to specify number of kv heads.
61
  ffn_config (Dict): A dictionary used to configure the model's ffn module:
62
+ ffn_type (str): type of ffn to use. Options: mptmlp, mptglu, te_ln_mlp
63
  init_device (str): The device to use for parameter initialization.
64
  logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.
65
  no_bias (bool): Whether to use bias in all layers.
 
66
  embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.
67
  norm_type (str): choose type of norm to use
68
  use_cache (bool): Whether or not the model should return the last key/values attentions
 
82
  ---
83
  See llmfoundry.models.utils.param_init_fns.py for info on other param init config options
84
  fc_type (str): choose fc layer implementation. Options: torch and te. te layers support fp8 when using H100 GPUs.
85
+ tie_word_embeddings (bool): Whether to tie the input embedding and output layers.
86
+ use_pad_tok_in_ffn (bool): Whether to forward the pad token in the feedforward networks.
87
  """
88
  self.d_model = d_model
89
  self.n_heads = n_heads
 
104
  self.use_cache = use_cache
105
  self.init_config = init_config
106
  self.fc_type = fc_type
107
+ self.use_pad_tok_in_ffn = use_pad_tok_in_ffn
 
108
  if 'name' in kwargs:
109
  del kwargs['name']
110
  if 'loss_fn' in kwargs:
111
  del kwargs['loss_fn']
112
+ if self.attn_config.get('alibi', False) or self.attn_config.get('rope', False):
113
  self.learned_pos_emb = False
114
+ warnings.warn(f'alibi or rope is turned on, setting `learned_pos_emb` to `False.`')
115
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
116
  self._validate_config()
117
 
118
  def _set_config_defaults(self, config: Dict[str, Any], config_defaults: Dict[str, Any]) -> Dict[str, Any]:
119
  for (k, v) in config_defaults.items():
120
  if k not in config:
121
  config[k] = v
122
+ elif isinstance(v, dict):
123
+ config[k] = self._set_config_defaults(config[k] if config[k] is not None else {}, v)
124
  return config
125
 
126
  def _validate_config(self) -> None:
 
135
  raise ValueError(f"Unknown attn_impl={self.attn_config['attn_impl']}")
136
  if self.attn_config['prefix_lm'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
137
  raise NotImplementedError('prefix_lm only implemented with torch and triton attention.')
138
+ if self.attn_config['attn_impl'] == 'flash' and is_flash_v1_installed():
139
+ warnings.warn(VersionedDeprecationWarning('Support for Flash Attention v1 is deprecated. Please upgrade to Flash Attention v2.4.2. To install Flash Attention v2.4.2, please run `pip install -e ".[gpu-flash2]"` from the root directory of the llm-foundry repository.', remove_version='0.6.0'))
140
+ if self.attn_config['attn_impl'] == 'triton' and (not self.attn_config['prefix_lm']):
141
+ warnings.warn(UserWarning('If not using a Prefix Language Model, we recommend setting "attn_impl" to "flash" instead of "triton".'))
142
+ if self.attn_config['alibi'] and (not check_alibi_support(self.attn_config['attn_impl'])):
143
+ raise NotImplementedError('alibi only implemented with torch, triton, and flash (v2.4.2 or higher) attention.')
144
+ if self.attn_config['attn_uses_sequence_id'] and (not (self.attn_config['attn_impl'] in ['torch', 'triton'] or (self.attn_config['attn_impl'] == 'flash' and is_flash_v2_installed(v2_version='v2.1.2')))):
145
+ raise NotImplementedError('attn_uses_sequence_id only implemented with torch, triton, and flash (v2.1.2 or higher) attention.')
146
+ if self.attn_config['rope'] and self.attn_config['rope_impl'] not in ['dail', 'hf']:
147
+ raise ValueError('If rope is being used then rope_impl should be either "dail", or "hf".')
148
+ if self.attn_config['rope'] and self.attn_config['rope_impl'] == 'hf' and (self.attn_config['rope_hf_config']['type'] not in ['no_scaling', 'linear', 'dynamic']):
149
+ raise ValueError('If using hf implementation of rope, the type should be one of "no_scaling", "linear" or "dynamic".')
150
+ if self.attn_config['rope'] and self.attn_config['rope_impl'] == 'dail':
151
+ if self.attn_config['rope_dail_config']['type'] not in ['original', 'xpos']:
152
+ raise ValueError('If using the dail implementation of rope, the type should be one of "original" or "xpos".')
153
+ if not is_flash_v2_installed(v2_version='2.0.1'):
154
+ raise ImportError('If using the dail implementation of rope, the flash_attn library v2.0.1 or higher must be installed. Please check the instructions at https://github.com/mosaicml/llm-foundry/blob/main/TUTORIAL.md#what-kinds-of-positional-embeddings-does-llm-foundry-support')
155
+ if self.attn_config['sliding_window_size'] != -1 and (not (self.attn_config['attn_impl'] == 'flash' and is_flash_v2_installed(v2_version='v2.3.0'))):
156
+ raise NotImplementedError('sliding window only implemented with flash attention v2.3.0 or higher.')
157
  if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
158
  raise ValueError('model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!')
159
  if isinstance(self.logit_scale, str) and self.logit_scale != 'inv_sqrt_d_model':
160
  raise ValueError(f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
161
  if self.init_config.get('name', None) is None:
162
  raise ValueError(f"self.init_config={self.init_config!r} 'name' needs to be set.")
163
+ if not (self.learned_pos_emb or self.attn_config['alibi'] or self.attn_config['rope']):
164
+ warnings.warn(f'Positional information not being provided to the model using either learned_pos_emb or alibi or rope.')
165
  if self.fc_type == 'te' or self.ffn_config['ffn_type'] == 'te_ln_mlp':
166
  try:
167
  import transformer_engine.pytorch as te
168
  del te
169
  except:
170
  raise ImportError('TransformerEngine import fail. `fc_type: te` requires TransformerEngine be installed. ' + 'The required version of transformer_engine also requires FlashAttention v1.0.6 is installed:\n' + 'pip install flash-attn==1.0.6 --no-build-isolation \n' + 'pip install git+https://github.com/NVIDIA/TransformerEngine.git@144e4888b2cdd60bd52e706d5b7a79cb9c1a7156')
171
+ if self.ffn_config['ffn_type'] == 'mptgeglu':
172
+ raise ValueError('API CHANGE: `ffn_type=="mptgeglu"` changed to `ffn_type=="mptglu"`. ' + 'See [#829](https://github.com/mosaicml/llm-foundry/pull/829) for details.')
173
+ elif self.ffn_config['ffn_type'] in ['mptmlp', 'mptglu']:
174
  self.ffn_config['fc_type'] = self.fc_type
175
  elif self.ffn_config['ffn_type'] == 'te_ln_mlp':
176
+ self.ffn_config['bias'] = not self.no_bias
177
+ if 'ffn_act_fn' in self.ffn_config.keys():
178
+ raise ValueError(f'Transformer Engine block does not support custom activation functions.')
179
+ if not self.use_pad_tok_in_ffn:
180
+ try:
181
+ from flash_attn.bert_padding import unpad_input, pad_input
182
+ except:
183
+ raise ImportError('In order to set `use_pad_tok_in_ffn=False`, please install flash-attn==1.0.9 or flash-attn==2.3.6')
ffn.py CHANGED
@@ -1,5 +1,8 @@
1
- """GPT Blocks used for the GPT Model."""
2
- from typing import Any, Optional
 
 
 
3
  import torch
4
  import torch.nn as nn
5
  from .fc import FC_CLASS_REGISTRY
@@ -7,33 +10,88 @@ try:
7
  import transformer_engine.pytorch as te
8
  except:
9
  te = None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10
 
11
  class MPTMLP(nn.Module):
12
 
13
- def __init__(self, d_model: int, expansion_ratio: int, fc_type: str='torch', device: Optional[str]=None, bias: bool=True):
14
  super().__init__()
15
- fc_kwargs: dict[str, Any] = {'bias': bias}
 
16
  if fc_type != 'te':
17
- fc_kwargs['device'] = device
18
- self.up_proj = FC_CLASS_REGISTRY[fc_type](d_model, expansion_ratio * d_model, **fc_kwargs)
19
- self.act = nn.GELU(approximate='none')
20
- self.down_proj = FC_CLASS_REGISTRY[fc_type](expansion_ratio * d_model, d_model, **fc_kwargs)
21
  self.down_proj._is_residual = True
22
 
23
  def forward(self, x: torch.Tensor) -> torch.Tensor:
24
  return self.down_proj(self.act(self.up_proj(x)))
25
- FFN_CLASS_REGISTRY = {'mptmlp': MPTMLP}
 
 
 
 
 
 
 
 
 
26
  if te is not None:
27
  te.LayerNormMLP._has_norm = True
28
  FFN_CLASS_REGISTRY['te_ln_mlp'] = te.LayerNormMLP
29
 
30
- def build_ffn(d_model: int, expansion_ratio: int, fc_type: str='torch', device: Optional[str]=None, bias: bool=True, **kwargs: Any) -> nn.Module:
31
  ffn_type = kwargs.pop('ffn_type')
32
- if ffn_type == 'mptmlp':
33
  if len(kwargs) > 0:
34
- raise ValueError(f'MPTMLP got an unexpected keyword argument: {kwargs}')
35
- return MPTMLP(d_model=d_model, expansion_ratio=expansion_ratio, fc_type=fc_type, device=device, bias=bias)
36
  elif ffn_type == 'te_ln_mlp':
37
  assert te is not None
38
- return te.LayerNormMLP(hidden_size=d_model, ffn_hidden_size=d_model * expansion_ratio, bias=bias, **kwargs)
 
 
 
39
  raise ValueError(f'ffn_type={ffn_type!r} not recognized.')
 
1
+ """MPT Blocks used for the MPT Model."""
2
+ import logging
3
+ from copy import deepcopy
4
+ from functools import partial
5
+ from typing import Any, Callable, Optional, Union
6
  import torch
7
  import torch.nn as nn
8
  from .fc import FC_CLASS_REGISTRY
 
10
  import transformer_engine.pytorch as te
11
  except:
12
  te = None
13
+ log = logging.getLogger(__name__)
14
+ _FFN_ACT_FN_DEFAULT = {'name': 'gelu', 'approximate': 'none'}
15
+
16
+ def resolve_ffn_act_fn(config: Optional[dict]=None) -> Callable[[torch.Tensor], torch.Tensor]:
17
+ """Resolve the activation function for the feed-forward network.
18
+
19
+ Args:
20
+ config (Optional[dict]): The configuration dictionary for the activation function.
21
+ The dict config must specify the 'name' of a torch.nn.functional activation
22
+ function. All of other key values pairs are bound to the function as a partial.
23
+
24
+ Returns:
25
+ Callable[[torch.Tensor], torch.Tensor]: The activation function.
26
+ """
27
+ if config is None:
28
+ config = _FFN_ACT_FN_DEFAULT
29
+ config = deepcopy(config)
30
+ name = config.pop('name')
31
+ if not hasattr(torch.nn.functional, name):
32
+ raise ValueError(f'Unrecognised activation function name ({name}).')
33
+ act = getattr(torch.nn.functional, name)
34
+ return partial(act, **config)
35
+ _DEFAULT_ACT_FN = resolve_ffn_act_fn(_FFN_ACT_FN_DEFAULT)
36
+
37
+ def resolve_ffn_hidden_size(d_model: int, expansion_ratio: Union[int, float], ffn_hidden_size: Optional[int]=None) -> int:
38
+ """Resolve the hidden size of the feed-forward network.
39
+
40
+ Args:
41
+ d_model (int): The dimension of the input and output of the feed-forward network.
42
+ expansion_ratio (Union[int, float]): The expansion ratio of the feed-forward network.
43
+ ffn_hidden_size (Optional[int]): The hidden size of the feed-forward network.
44
+
45
+ Returns:
46
+ int: The hidden size of the feed-forward network.
47
+ """
48
+ if ffn_hidden_size is not None:
49
+ log.info(f'`expansion_ratio` (={expansion_ratio}) ignored when `ffn_hidden_size` (={ffn_hidden_size}) is specified.')
50
+ else:
51
+ ffn_hidden_size = int(d_model * expansion_ratio)
52
+ if ffn_hidden_size != d_model * expansion_ratio:
53
+ raise ValueError(f'`d_model * expansion_ratio` must be an integer (d_model={d_model!r}; expansion_ratio={expansion_ratio!r}; d_model * expansion_ratio={d_model * expansion_ratio!r}).')
54
+ return ffn_hidden_size
55
 
56
  class MPTMLP(nn.Module):
57
 
58
+ def __init__(self, d_model: int, expansion_ratio: Union[int, float], fc_type: str='torch', ffn_hidden_size: Optional[int]=None, act_fn: Callable[[torch.Tensor], torch.Tensor]=_DEFAULT_ACT_FN, device: Optional[str]=None, bias: bool=True):
59
  super().__init__()
60
+ ffn_hidden_size = resolve_ffn_hidden_size(d_model, expansion_ratio, ffn_hidden_size)
61
+ self.fc_kwargs: dict[str, Any] = {'bias': bias}
62
  if fc_type != 'te':
63
+ self.fc_kwargs['device'] = device
64
+ self.up_proj = FC_CLASS_REGISTRY[fc_type](d_model, ffn_hidden_size, **self.fc_kwargs)
65
+ self.act = act_fn
66
+ self.down_proj = FC_CLASS_REGISTRY[fc_type](ffn_hidden_size, d_model, **self.fc_kwargs)
67
  self.down_proj._is_residual = True
68
 
69
  def forward(self, x: torch.Tensor) -> torch.Tensor:
70
  return self.down_proj(self.act(self.up_proj(x)))
71
+
72
+ class MPTGLU(MPTMLP):
73
+
74
+ def __init__(self, d_model: int, expansion_ratio: Union[int, float], fc_type: str='torch', ffn_hidden_size: Optional[int]=None, act_fn: Callable[[torch.Tensor], torch.Tensor]=_DEFAULT_ACT_FN, device: Optional[str]=None, bias: bool=True):
75
+ super().__init__(d_model=d_model, expansion_ratio=expansion_ratio, fc_type=fc_type, ffn_hidden_size=ffn_hidden_size, act_fn=act_fn, device=device, bias=bias)
76
+ self.gate_proj = FC_CLASS_REGISTRY[fc_type](d_model, self.up_proj.out_features, **self.fc_kwargs)
77
+
78
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
79
+ return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))
80
+ FFN_CLASS_REGISTRY = {'mptmlp': MPTMLP, 'mptglu': MPTGLU}
81
  if te is not None:
82
  te.LayerNormMLP._has_norm = True
83
  FFN_CLASS_REGISTRY['te_ln_mlp'] = te.LayerNormMLP
84
 
85
+ def build_ffn(d_model: int, expansion_ratio: Union[int, float], fc_type: str='torch', ffn_hidden_size: Optional[int]=None, ffn_act_fn: Optional[dict]=None, device: Optional[str]=None, bias: bool=True, **kwargs: Any) -> nn.Module:
86
  ffn_type = kwargs.pop('ffn_type')
87
+ if ffn_type in ['mptmlp', 'mptglu']:
88
  if len(kwargs) > 0:
89
+ raise ValueError(f'MPTMLP (or MPTGLU) got an unexpected keyword argument: {kwargs}')
90
+ return FFN_CLASS_REGISTRY[ffn_type](d_model=d_model, expansion_ratio=expansion_ratio, fc_type=fc_type, act_fn=resolve_ffn_act_fn(ffn_act_fn), ffn_hidden_size=ffn_hidden_size, device=device, bias=bias)
91
  elif ffn_type == 'te_ln_mlp':
92
  assert te is not None
93
+ ffn_hidden_size = resolve_ffn_hidden_size(d_model, expansion_ratio, ffn_hidden_size)
94
+ if ffn_act_fn is not None:
95
+ raise ValueError(f'Transformer Engine block does not support custom activation functions.')
96
+ return te.LayerNormMLP(hidden_size=d_model, ffn_hidden_size=ffn_hidden_size, bias=bias, **kwargs)
97
  raise ValueError(f'ffn_type={ffn_type!r} not recognized.')
modeling_mpt.py CHANGED
@@ -2,15 +2,31 @@
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 +46,130 @@ 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 +197,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,18 +212,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
  if hasattr(module, 'use_bias'):
78
- log.info(f'Setting use_bias=False for {module}.')
79
  module.use_bias = False
80
  log.debug(self)
81
  log.debug(f"Using {self.config.init_config['name']} initialization.")
82
 
83
- def get_input_embeddings(self) -> nn.Embedding:
84
  return self.wte
85
 
86
- def set_input_embeddings(self, value: nn.Embedding) -> None:
87
  self.wte = value
88
 
89
  @torch.no_grad()
@@ -104,7 +244,7 @@ class MPTModel(MPTPreTrainedModel):
104
  attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)
105
  if self.attn_uses_sequence_id and sequence_id is not None:
106
  assert isinstance(attn_bias, torch.Tensor)
107
- attn_bias = self._apply_sequence_id(attn_bias, sequence_id)
108
  if attention_mask is not None:
109
  s_k = attention_mask.shape[-1]
110
  if attn_bias is None:
@@ -116,7 +256,7 @@ class MPTModel(MPTPreTrainedModel):
116
  raise ValueError(f'attention_mask shape={attention_mask.shape} ' + f'and prefix_mask shape={prefix_mask.shape} are not equal.')
117
  min_val = torch.finfo(attn_bias.dtype).min
118
  attn_bias = attn_bias.masked_fill(~attention_mask.view(-1, 1, 1, s_k), min_val)
119
- return (attn_bias, None)
120
 
121
  def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor) -> torch.Tensor:
122
  (s_k, s_q) = attn_bias.shape[-2:]
@@ -133,17 +273,7 @@ class MPTModel(MPTPreTrainedModel):
133
  attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
134
  return attn_bias
135
 
136
- def _apply_sequence_id(self, attn_bias: torch.Tensor, sequence_id: torch.LongTensor) -> torch.Tensor:
137
- seq_len = sequence_id.shape[-1]
138
- if seq_len > self.config.max_seq_len:
139
- raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
140
- attn_bias = attn_bias[..., :seq_len, :seq_len]
141
- cannot_attend = torch.logical_not(torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))).unsqueeze(1)
142
- min_val = torch.finfo(attn_bias.dtype).min
143
- attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
144
- return attn_bias
145
-
146
- 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:
147
  return_dict = return_dict if return_dict is not None else self.config.return_dict
148
  use_cache = use_cache if use_cache is not None else self.config.use_cache
149
  if attention_mask is not None:
@@ -159,33 +289,47 @@ class MPTModel(MPTPreTrainedModel):
159
  raise NotImplementedError('MPT does not support training with left padding.')
160
  if self.prefix_lm and prefix_mask is None:
161
  raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
162
- if inputs_embeds is not None:
163
- raise NotImplementedError('inputs_embeds is not implemented for MPT.')
164
  if self.training:
165
  if self.attn_uses_sequence_id and sequence_id is None:
166
  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.')
167
  elif self.attn_uses_sequence_id is False and sequence_id is not None:
168
  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.')
169
- S = input_ids.size(1)
 
 
 
 
 
 
 
 
 
 
 
 
 
170
  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}'
171
- tok_emb = self.wte(input_ids)
172
- if self.learned_pos_emb:
173
- past_position = 0
174
- if past_key_values is not None:
175
- if len(past_key_values) != self.config.n_layers:
176
- 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}).')
177
- past_position = past_key_values[0][0].size(1)
178
- if self.attn_impl == 'torch':
179
- past_position = past_key_values[0][0].size(3)
180
- if S + past_position > self.config.max_seq_len:
181
  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}.')
182
- pos = torch.arange(past_position, S + past_position, dtype=torch.long, device=input_ids.device).unsqueeze(0)
183
- if attention_mask is not None:
184
- pos = torch.clamp(pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[:, past_position:], min=0)
185
- pos_emb = self.wpe(pos)
186
- x = tok_emb + pos_emb
187
- else:
188
- x = tok_emb
 
 
 
189
  if self.embedding_fraction == 1:
190
  x = self.emb_drop(x)
191
  else:
@@ -193,17 +337,24 @@ class MPTModel(MPTPreTrainedModel):
193
  assert isinstance(self.emb_drop, nn.Module)
194
  x = self.emb_drop(x_shrunk)
195
  (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)
 
 
 
 
196
  presents = () if use_cache else None
197
  if use_cache and past_key_values is None:
198
  past_key_values = [() for _ in range(self.config.n_layers)]
199
  all_hidden_states = () if output_hidden_states else None
200
  all_self_attns = () if output_attentions else None
 
 
 
201
  for (b_idx, block) in enumerate(self.blocks):
202
  if output_hidden_states:
203
  assert all_hidden_states is not None
204
  all_hidden_states = all_hidden_states + (x,)
205
  past_key_value = past_key_values[b_idx] if past_key_values is not None else None
206
- (x, attn_weights, present) = 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))
207
  if presents is not None:
208
  presents += (present,)
209
  if output_attentions:
@@ -220,7 +371,7 @@ class MPTModel(MPTPreTrainedModel):
220
  MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
221
 
222
  def fsdp_wrap_fn(self, module: nn.Module) -> bool:
223
- return isinstance(module, MPTBlock)
224
 
225
  def activation_checkpointing_fn(self, module: nn.Module) -> bool:
226
  return isinstance(module, MPTBlock)
@@ -229,10 +380,12 @@ class MPTForCausalLM(MPTPreTrainedModel):
229
 
230
  def __init__(self, config: MPTConfig):
231
  super().__init__(config)
232
- if not config.tie_word_embeddings:
233
- raise ValueError('MPTForCausalLM only supports tied word embeddings')
234
  log.info(f'Instantiating an MPTForCausalLM model from {__file__}')
235
  self.transformer: MPTModel = MPTModel(config)
 
 
 
 
236
  for child in self.transformer.children():
237
  if isinstance(child, torch.nn.ModuleList):
238
  continue
@@ -248,17 +401,28 @@ class MPTForCausalLM(MPTPreTrainedModel):
248
  raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
249
  self.logit_scale = logit_scale
250
 
251
- def get_input_embeddings(self) -> nn.Embedding:
252
- return self.transformer.wte
253
 
254
  def set_input_embeddings(self, value: Union[SharedEmbedding, nn.Embedding]) -> None:
255
- self.transformer.wte = value
 
 
 
 
 
256
 
257
- def get_output_embeddings(self) -> nn.Embedding:
258
- return self.transformer.wte
 
 
 
 
 
 
259
 
260
- def set_output_embeddings(self, new_embeddings: Union[SharedEmbedding, nn.Embedding]) -> None:
261
- self.transformer.wte = new_embeddings
262
 
263
  def set_decoder(self, decoder: MPTModel) -> None:
264
  self.transformer = decoder
@@ -266,13 +430,16 @@ class MPTForCausalLM(MPTPreTrainedModel):
266
  def get_decoder(self) -> MPTModel:
267
  return self.transformer
268
 
269
- 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:
270
  return_dict = return_dict if return_dict is not None else self.config.return_dict
271
  use_cache = use_cache if use_cache is not None else self.config.use_cache
272
- if inputs_embeds is not None:
273
- raise NotImplementedError('inputs_embeds has to be None (for hf/peft support).')
274
- 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)
275
- logits = self.transformer.wte(outputs.last_hidden_state.to(self.transformer.wte.weight.device), True)
 
 
 
276
  if self.logit_scale is not None:
277
  if self.logit_scale == 0:
278
  warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
@@ -289,14 +456,34 @@ class MPTForCausalLM(MPTPreTrainedModel):
289
  MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
290
 
291
  def fsdp_wrap_fn(self, module: nn.Module) -> bool:
292
- return isinstance(module, MPTBlock)
293
 
294
  def activation_checkpointing_fn(self, module: nn.Module) -> bool:
295
- return isinstance(module, MPTBlock)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
296
 
297
  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]:
298
- if inputs_embeds is not None:
299
- raise NotImplementedError('inputs_embeds is not implemented for MPT yet')
300
  attention_mask = kwargs['attention_mask'].bool()
301
  if attention_mask[:, -1].sum() != attention_mask.shape[0]:
302
  raise NotImplementedError('MPT does not support generation with right padding.')
@@ -312,7 +499,12 @@ class MPTForCausalLM(MPTPreTrainedModel):
312
  raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
313
  else:
314
  prefix_mask = None
315
- 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)}
 
 
 
 
 
316
 
317
  @staticmethod
318
  def _reorder_cache(past_key_values: List[Tuple[torch.Tensor, torch.Tensor]], beam_idx: torch.LongTensor) -> List[Tuple[torch.Tensor, ...]]:
 
2
 
3
  Inspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
4
  """
5
+ from __future__ import annotations
6
  import math
7
  import warnings
8
  from typing import Any, Dict, List, Mapping, MutableMapping, Optional, Tuple, Union
9
  import torch
10
  import torch.nn as nn
11
  import torch.nn.functional as F
12
+ from .attention import is_flash_v1_installed, is_flash_v2_installed
13
+ if is_flash_v2_installed():
14
+ try:
15
+ from flash_attn import bert_padding
16
+ from flash_attn.layers.rotary import RotaryEmbedding as DAILRotaryEmbedding
17
+ except Exception as e:
18
+ raise e
19
+ if is_flash_v1_installed():
20
+ try:
21
+ from flash_attn import bert_padding
22
+ except Exception as e:
23
+ raise e
24
  from transformers import PreTrainedModel, PreTrainedTokenizerBase
25
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
26
+ from transformers.models.llama.modeling_llama import LlamaDynamicNTKScalingRotaryEmbedding as HFDynamicNTKScalingRotaryEmbedding
27
+ from transformers.models.llama.modeling_llama import LlamaLinearScalingRotaryEmbedding as HFLinearScalingRotaryEmbedding
28
+ from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding as HFRotaryEmbedding
29
+ from .attention import ATTN_CLASS_REGISTRY, attn_bias_shape, build_attn_bias, gen_slopes
30
  from .blocks import MPTBlock
31
  from .custom_embedding import SharedEmbedding
32
  from .fc import FC_CLASS_REGISTRY as FC_CLASS_REGISTRY
 
46
  import logging
47
  log = logging.getLogger(__name__)
48
 
49
+ 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):
50
+ if rope_impl == 'dail':
51
+ 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')
52
+ elif rope_impl == 'hf':
53
+ if rope_hf_config['type'] == 'no_scaling':
54
+ return HFRotaryEmbedding(rope_head_dim, max_position_embeddings=max_seq_len, base=rope_theta, device='cpu')
55
+ elif rope_hf_config['type'] == 'linear':
56
+ return HFLinearScalingRotaryEmbedding(rope_head_dim, max_position_embeddings=max_seq_len, base=rope_theta, scaling_factor=rope_hf_config['factor'], device='cpu')
57
+ elif rope_hf_config['type'] == 'dynamic':
58
+ return HFDynamicNTKScalingRotaryEmbedding(rope_head_dim, max_position_embeddings=max_seq_len, base=rope_theta, scaling_factor=rope_hf_config['factor'], device='cpu')
59
+ raise ValueError('rope_impl needs to be either dail or hf')
60
+
61
+ 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]):
62
+ """Generates the attention mask used for sequence masking in FA v2.
63
+
64
+ Only supports sequence id based sparse attention for no attention masking or attention masking with right padding.
65
+ In case of left padding:
66
+ 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).
67
+ 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.
68
+
69
+ Args:
70
+ sequence_id (Union[None, torch.Tensor]): Tensor containing the sequence id for each token. Shape (batch_size, seq_len).
71
+ S (int): Sequence length
72
+ attn_uses_sequence_id (bool): Whether the attention uses sequence id based masking.
73
+ attn_impl (str): Attention implementation. This function is only creates attention_mask_in_length for flash attention.
74
+ attention_mask (Union[torch.Tensor, None]): Attention mask tensor of shape (batch_size, seq_len)
75
+
76
+ Returns:
77
+ 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:
78
+ ```
79
+ [
80
+ [2, 3, 0, 0, 0, 0],
81
+ [3, 2, 0, 0, 0, 0],
82
+ [6, 0, 0, 0, 0, 0]
83
+ ]
84
+ ```
85
+ , which refers to the 3D-attention mask:
86
+ ```
87
+ [
88
+ [
89
+ [1, 0, 0, 0, 0, 0],
90
+ [1, 1, 0, 0, 0, 0],
91
+ [0, 0, 1, 0, 0, 0],
92
+ [0, 0, 1, 1, 0, 0],
93
+ [0, 0, 1, 1, 1, 0],
94
+ [0, 0, 0, 0, 0, 1]
95
+ ],
96
+ [
97
+ [1, 0, 0, 0, 0, 0],
98
+ [1, 1, 0, 0, 0, 0],
99
+ [1, 1, 1, 0, 0, 0],
100
+ [0, 0, 0, 1, 0, 0],
101
+ [0, 0, 0, 1, 1, 0],
102
+ [0, 0, 0, 0, 0, 1]
103
+ ],
104
+ [
105
+ [1, 0, 0, 0, 0, 0],
106
+ [1, 1, 0, 0, 0, 0],
107
+ [1, 1, 1, 0, 0, 0],
108
+ [1, 1, 1, 1, 0, 0],
109
+ [1, 1, 1, 1, 1, 0],
110
+ [1, 1, 1, 1, 1, 1]
111
+ ]
112
+ ]
113
+ ```.
114
+ (The description above is taken verbatim from https://github.com/Dao-AILab/flash-attention/blob/9356a1c0389660d7e231ff3163c1ac17d9e3824a/flash_attn/bert_padding.py#L125 .)
115
+ """
116
+ attention_mask_in_length = None
117
+ if sequence_id is not None and attn_uses_sequence_id and (attn_impl == 'flash'):
118
+ if attention_mask is not None and attention_mask[:, 0].sum() != attention_mask.shape[0]:
119
+ raise NotImplementedError('Left padding is not supported with flash attention when attn_uses_sequence_id is set to True.')
120
+ if S != sequence_id.shape[-1]:
121
+ raise ValueError(f'Sequence length ({S}) does not match length of sequences in sequence_id ({sequence_id.shape[-1]}).')
122
+ if attention_mask is not None:
123
+ sequence_id = sequence_id.masked_fill(~attention_mask, 0)
124
+ attention_mask_in_length = torch.nn.functional.one_hot(sequence_id)
125
+ if attention_mask is not None:
126
+ attention_mask_in_length = attention_mask_in_length.masked_fill(~attention_mask.unsqueeze(-1), 0)
127
+ attention_mask_in_length = attention_mask_in_length.sum(dim=1)
128
+ attention_mask_in_length = torch.nn.functional.pad(attention_mask_in_length, (0, S - attention_mask_in_length.shape[-1]), mode='constant', value=0)
129
+ return attention_mask_in_length
130
+
131
+ 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):
132
+ flash_attn_padding_info = {}
133
+ if attention_mask_in_length is None:
134
+ key_padding_mask = attention_mask
135
+ if key_padding_mask is None:
136
+ key_padding_mask = torch.ones((bsz, past_key_len + S), dtype=torch.bool, device=device)
137
+ query_padding_mask = key_padding_mask[:, -S:]
138
+ unpadding_function = bert_padding.unpad_input
139
+ else:
140
+ key_padding_mask = attention_mask_in_length
141
+ query_padding_mask = attention_mask_in_length
142
+ unpadding_function = bert_padding.unpad_input_for_concatenated_sequences
143
+ (_, indices_q, cu_seqlens_q, max_seqlen_q) = unpadding_function(torch.empty(bsz, S, 1, device=device), query_padding_mask)
144
+ (_, indices_k, cu_seqlens_k, max_seqlen_k) = unpadding_function(torch.empty(bsz, past_key_len + S, 1, device=device), key_padding_mask)
145
+ (_, indices_v, _, _) = unpadding_function(torch.empty(bsz, past_key_len + S, 1, device=device), key_padding_mask)
146
+ flash_attn_padding_info['indices_q'] = indices_q
147
+ flash_attn_padding_info['indices_k'] = indices_k
148
+ flash_attn_padding_info['indices_v'] = indices_v
149
+ flash_attn_padding_info['cu_seqlens_q'] = cu_seqlens_q
150
+ flash_attn_padding_info['cu_seqlens_k'] = cu_seqlens_k
151
+ flash_attn_padding_info['max_seqlen_q'] = max_seqlen_q
152
+ flash_attn_padding_info['max_seqlen_k'] = max_seqlen_k
153
+ return flash_attn_padding_info
154
+
155
+ def apply_sequence_id(attn_bias: torch.Tensor, sequence_id: torch.LongTensor, max_seq_len: int) -> torch.Tensor:
156
+ seq_len = sequence_id.shape[-1]
157
+ if seq_len > max_seq_len:
158
+ raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={max_seq_len}')
159
+ attn_bias = attn_bias[..., :seq_len, :seq_len]
160
+ cannot_attend = torch.logical_not(torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))).unsqueeze(1)
161
+ min_val = torch.finfo(attn_bias.dtype).min
162
+ attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
163
+ return attn_bias
164
+
165
  class MPTPreTrainedModel(PreTrainedModel):
166
  config_class = MPTConfig
167
  base_model_prefix = 'model'
168
  _no_split_modules = ['MPTBlock']
169
 
170
+ def _fsdp_wrap_fn(self: Union[MPTModel, MPTForCausalLM], module: nn.Module) -> bool:
171
+ return isinstance(module, MPTBlock)
172
+
173
  class MPTModel(MPTPreTrainedModel):
174
 
175
  def __init__(self, config: MPTConfig):
 
197
  self.emb_drop = nn.Dropout(config.emb_pdrop)
198
  self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
199
  self.norm_f = norm_class(config.d_model, device=config.init_device)
200
+ self.rope = config.attn_config['rope']
201
+ self.rope_impl = None
202
+ if self.rope:
203
+ self.rope_impl = config.attn_config['rope_impl']
204
+ 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)
205
  if config.init_device != 'meta':
206
  log.info(f'We recommend using config.init_device="meta" with Composer + FSDP for faster initialization.')
207
  self.apply(self.param_init_fn)
 
212
  if config.no_bias:
213
  for module in self.modules():
214
  if hasattr(module, 'bias') and isinstance(module.bias, nn.Parameter):
215
+ log.info(f'Removing bias from module={module!r}.')
216
  module.register_parameter('bias', None)
217
  if hasattr(module, 'use_bias'):
218
+ log.info(f'Setting use_bias=False for module={module!r}.')
219
  module.use_bias = False
220
  log.debug(self)
221
  log.debug(f"Using {self.config.init_config['name']} initialization.")
222
 
223
+ def get_input_embeddings(self) -> Union[SharedEmbedding, nn.Embedding]:
224
  return self.wte
225
 
226
+ def set_input_embeddings(self, value: Union[SharedEmbedding, nn.Embedding]) -> None:
227
  self.wte = value
228
 
229
  @torch.no_grad()
 
244
  attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)
245
  if self.attn_uses_sequence_id and sequence_id is not None:
246
  assert isinstance(attn_bias, torch.Tensor)
247
+ attn_bias = apply_sequence_id(attn_bias, sequence_id, self.config.max_seq_len)
248
  if attention_mask is not None:
249
  s_k = attention_mask.shape[-1]
250
  if attn_bias is None:
 
256
  raise ValueError(f'attention_mask shape={attention_mask.shape} ' + f'and prefix_mask shape={prefix_mask.shape} are not equal.')
257
  min_val = torch.finfo(attn_bias.dtype).min
258
  attn_bias = attn_bias.masked_fill(~attention_mask.view(-1, 1, 1, s_k), min_val)
259
+ return (attn_bias, attention_mask)
260
 
261
  def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor) -> torch.Tensor:
262
  (s_k, s_q) = attn_bias.shape[-2:]
 
273
  attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
274
  return attn_bias
275
 
276
+ 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:
 
 
 
 
 
 
 
 
 
 
277
  return_dict = return_dict if return_dict is not None else self.config.return_dict
278
  use_cache = use_cache if use_cache is not None else self.config.use_cache
279
  if attention_mask is not None:
 
289
  raise NotImplementedError('MPT does not support training with left padding.')
290
  if self.prefix_lm and prefix_mask is None:
291
  raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
 
 
292
  if self.training:
293
  if self.attn_uses_sequence_id and sequence_id is None:
294
  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.')
295
  elif self.attn_uses_sequence_id is False and sequence_id is not None:
296
  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.')
297
+ if input_ids is not None and inputs_embeds is not None:
298
+ raise ValueError('You cannot specify both input_ids and inputs_embeds.')
299
+ elif input_ids is not None:
300
+ bsz = input_ids.size(0)
301
+ S = input_ids.size(1)
302
+ x = self.wte(input_ids)
303
+ input_device = input_ids.device
304
+ elif inputs_embeds is not None:
305
+ bsz = inputs_embeds.size(0)
306
+ S = inputs_embeds.size(1)
307
+ x = inputs_embeds
308
+ input_device = inputs_embeds.device
309
+ else:
310
+ raise ValueError('You must specify input_ids or inputs_embeds')
311
  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}'
312
+ rotary_emb_w_meta_info = None
313
+ past_position = 0
314
+ if past_key_values is not None:
315
+ if len(past_key_values) != self.config.n_layers:
316
+ 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}).')
317
+ past_position = past_key_values[0][0].size(1)
318
+ if self.attn_impl == 'torch':
319
+ past_position = past_key_values[0][0].size(3)
320
+ if self.learned_pos_emb or self.rope:
321
+ if self.learned_pos_emb and S + past_position > self.config.max_seq_len:
322
  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}.')
323
+ if self.learned_pos_emb or (self.rope and self.rope_impl == 'hf'):
324
+ pos = torch.arange(past_position, S + past_position, dtype=torch.long, device=input_device).unsqueeze(0)
325
+ if attention_mask is not None:
326
+ pos = torch.clamp(pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[:, past_position:], min=0)
327
+ if self.learned_pos_emb:
328
+ x = x + self.wpe(pos)
329
+ elif self.rope and self.rope_impl == 'hf':
330
+ rotary_emb_w_meta_info = {'impl': self.rope_impl, 'rotary_emb': self.rotary_embedding, 'offset_info': pos, 'seq_len': S + past_position}
331
+ elif self.rope and self.rope_impl == 'dail':
332
+ rotary_emb_w_meta_info = {'impl': self.rope_impl, 'rotary_emb': self.rotary_embedding, 'offset_info': past_position, 'seq_len': S + past_position}
333
  if self.embedding_fraction == 1:
334
  x = self.emb_drop(x)
335
  else:
 
337
  assert isinstance(self.emb_drop, nn.Module)
338
  x = self.emb_drop(x_shrunk)
339
  (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)
340
+ 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)
341
+ alibi_slopes = None
342
+ if self.alibi and self.attn_impl == 'flash':
343
+ alibi_slopes = gen_slopes(n_heads=self.config.n_heads, alibi_bias_max=self.alibi_bias_max, device=x.device, return_1d=True)
344
  presents = () if use_cache else None
345
  if use_cache and past_key_values is None:
346
  past_key_values = [() for _ in range(self.config.n_layers)]
347
  all_hidden_states = () if output_hidden_states else None
348
  all_self_attns = () if output_attentions else None
349
+ flash_attn_padding_info = {}
350
+ if self.attn_impl == 'flash':
351
+ flash_attn_padding_info = gen_flash_attn_padding_info(bsz, S, past_position, x.device, attention_mask_in_length, attention_mask)
352
  for (b_idx, block) in enumerate(self.blocks):
353
  if output_hidden_states:
354
  assert all_hidden_states is not None
355
  all_hidden_states = all_hidden_states + (x,)
356
  past_key_value = past_key_values[b_idx] if past_key_values is not None else None
357
+ (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)
358
  if presents is not None:
359
  presents += (present,)
360
  if output_attentions:
 
371
  MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
372
 
373
  def fsdp_wrap_fn(self, module: nn.Module) -> bool:
374
+ return _fsdp_wrap_fn(self, module)
375
 
376
  def activation_checkpointing_fn(self, module: nn.Module) -> bool:
377
  return isinstance(module, MPTBlock)
 
380
 
381
  def __init__(self, config: MPTConfig):
382
  super().__init__(config)
 
 
383
  log.info(f'Instantiating an MPTForCausalLM model from {__file__}')
384
  self.transformer: MPTModel = MPTModel(config)
385
+ self.lm_head = None
386
+ if not config.tie_word_embeddings:
387
+ self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False, device=config.init_device)
388
+ self.lm_head._fsdp_wrap = True
389
  for child in self.transformer.children():
390
  if isinstance(child, torch.nn.ModuleList):
391
  continue
 
401
  raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
402
  self.logit_scale = logit_scale
403
 
404
+ def get_input_embeddings(self) -> Union[SharedEmbedding, nn.Embedding]:
405
+ return self.transformer.get_input_embeddings()
406
 
407
  def set_input_embeddings(self, value: Union[SharedEmbedding, nn.Embedding]) -> None:
408
+ self.transformer.set_input_embeddings(value)
409
+
410
+ def get_output_embeddings(self) -> Union[SharedEmbedding, nn.Embedding, nn.Linear]:
411
+ if self.lm_head is not None:
412
+ return self.lm_head
413
+ return self.transformer.get_input_embeddings()
414
 
415
+ def set_output_embeddings(self, new_embeddings: Union[SharedEmbedding, nn.Embedding, nn.Linear]) -> None:
416
+ if self.lm_head is not None:
417
+ self.lm_head = new_embeddings
418
+ else:
419
+ if not isinstance(new_embeddings, (SharedEmbedding, nn.Embedding)):
420
+ raise ValueError('new_embeddings must be an instance of SharedEmbedding ' + f'or nn.Embedding, but got {type(new_embeddings)}.')
421
+ 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`.')
422
+ self.transformer.set_input_embeddings(new_embeddings)
423
 
424
+ def tie_weights(self) -> None:
425
+ self.lm_head = None
426
 
427
  def set_decoder(self, decoder: MPTModel) -> None:
428
  self.transformer = decoder
 
430
  def get_decoder(self) -> MPTModel:
431
  return self.transformer
432
 
433
+ 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:
434
  return_dict = return_dict if return_dict is not None else self.config.return_dict
435
  use_cache = use_cache if use_cache is not None else self.config.use_cache
436
+ 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)
437
+ if self.lm_head is not None:
438
+ logits = self.lm_head(outputs.last_hidden_state)
439
+ else:
440
+ out = outputs.last_hidden_state
441
+ out = out.to(self.transformer.wte.weight.device)
442
+ logits = self.transformer.wte(out, True)
443
  if self.logit_scale is not None:
444
  if self.logit_scale == 0:
445
  warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
 
456
  MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
457
 
458
  def fsdp_wrap_fn(self, module: nn.Module) -> bool:
459
+ return _fsdp_wrap_fn(self, module)
460
 
461
  def activation_checkpointing_fn(self, module: nn.Module) -> bool:
462
+ act_ckpt_list = getattr(self.config, 'activation_checkpointing_target', None) or ['MPTBlock']
463
+ if isinstance(act_ckpt_list, str):
464
+ act_ckpt_list = [act_ckpt_list]
465
+ elif not isinstance(act_ckpt_list, list):
466
+ raise ValueError(f'activation_checkpointing_target must be either a single string or a list, but got {type(act_ckpt_list)}')
467
+ if 'MPTBlock' in act_ckpt_list or 'mptblock' in act_ckpt_list:
468
+ if len(act_ckpt_list) > 1:
469
+ log.info('Activation checkpointing MPTBlock only (ignoring other sub-block modules specified in activation_checkpointing_target).')
470
+ return isinstance(module, MPTBlock)
471
+ mod_types = ()
472
+ for mod_name in act_ckpt_list:
473
+ if mod_name.lower() == 'mptblock':
474
+ mod_types += (MPTBlock,)
475
+ elif mod_name in ATTN_CLASS_REGISTRY:
476
+ mod_types += (ATTN_CLASS_REGISTRY[mod_name],)
477
+ elif mod_name in FFN_CLASS_REGISTRY:
478
+ mod_types += (FFN_CLASS_REGISTRY[mod_name],)
479
+ elif mod_name in NORM_CLASS_REGISTRY:
480
+ mod_types += (NORM_CLASS_REGISTRY[mod_name],)
481
+ else:
482
+ msg = ', '.join(list(ATTN_CLASS_REGISTRY.keys()) + list(FFN_CLASS_REGISTRY.keys()) + list(NORM_CLASS_REGISTRY.keys()) + ['MPTBlock'])
483
+ raise ValueError(f'{mod_name} (specified in activation_checkpointing_target) is not a recognized option out of available options {msg}.')
484
+ return isinstance(module, mod_types)
485
 
486
  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]:
 
 
487
  attention_mask = kwargs['attention_mask'].bool()
488
  if attention_mask[:, -1].sum() != attention_mask.shape[0]:
489
  raise NotImplementedError('MPT does not support generation with right padding.')
 
499
  raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
500
  else:
501
  prefix_mask = None
502
+ if inputs_embeds is not None and past_key_values is None:
503
+ model_inputs = {'inputs_embeds': inputs_embeds}
504
+ else:
505
+ model_inputs = {'input_ids': input_ids}
506
+ 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)})
507
+ return model_inputs
508
 
509
  @staticmethod
510
  def _reorder_cache(past_key_values: List[Tuple[torch.Tensor, torch.Tensor]], beam_idx: torch.LongTensor) -> List[Tuple[torch.Tensor, ...]]:
warnings.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ class VersionedDeprecationWarning(DeprecationWarning):
2
+ """A custom deprecation warning class that includes version information.
3
+
4
+ Attributes:
5
+ message (str): The deprecation message describing why the feature is deprecated.
6
+ remove_version (str): The version in which the feature will be removed.
7
+
8
+ Example:
9
+ >>> def deprecated_function():
10
+ ... warnings.warn(
11
+ ... VersionedDeprecationWarning(
12
+ ... "Function XYZ is deprecated.",
13
+ ... after_version="2.0.0"
14
+ ... )
15
+ ... )
16
+ ...
17
+ >>> deprecated_function()
18
+ DeprecationWarning: Function XYZ is deprecated. It will be removed in version 2.0.0.
19
+ """
20
+
21
+ def __init__(self, message: str, remove_version: str) -> None:
22
+ super().__init__(message + f' It will be removed in version {remove_version}.')