LLM-foundry update February 07, 2024 19:44:25 (#9)
Browse files- LLM-foundry update February 07, 2024 19:44:25 (f62258af3f01cd1c6981bd8530ef7cacfe324583)
Co-authored-by: Irene Dea <[email protected]>
- attention.py +109 -59
- blocks.py +19 -5
- configuration_mpt.py +61 -18
- ffn.py +72 -14
- modeling_mpt.py +251 -59
- 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,
|
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(
|
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:
|
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[
|
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:
|
|
|
|
|
|
|
|
|
|
|
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.
|
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 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
|
|
|
|
|
|
|
|
137 |
query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
|
138 |
-
|
139 |
key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=kv_n_heads)
|
140 |
-
|
141 |
value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=kv_n_heads)
|
142 |
-
if kv_n_heads
|
143 |
-
|
144 |
-
|
145 |
-
|
146 |
-
|
147 |
-
|
|
|
|
|
|
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
154 |
else:
|
155 |
-
raise RuntimeError('flash-attn==1.0.9 or flash-attn==2.
|
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:
|
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 |
-
|
251 |
-
self.
|
|
|
|
|
|
|
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[
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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[
|
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 =
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
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',
|
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 |
-
|
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['
|
120 |
-
|
121 |
-
if self.attn_config['
|
122 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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
|
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'] == '
|
|
|
|
|
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 |
-
"""
|
2 |
-
|
|
|
|
|
|
|
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 |
-
|
|
|
16 |
if fc_type != 'te':
|
17 |
-
fc_kwargs['device'] = device
|
18 |
-
self.up_proj = FC_CLASS_REGISTRY[fc_type](d_model,
|
19 |
-
self.act =
|
20 |
-
self.down_proj = FC_CLASS_REGISTRY[fc_type](
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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
|
33 |
if len(kwargs) > 0:
|
34 |
-
raise ValueError(f'MPTMLP got an unexpected keyword argument: {kwargs}')
|
35 |
-
return
|
36 |
elif ffn_type == 'te_ln_mlp':
|
37 |
assert te is not None
|
38 |
-
|
|
|
|
|
|
|
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 .
|
|
|
|
|
|
|
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
|
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 =
|
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,
|
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
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
172 |
-
|
173 |
-
|
174 |
-
if past_key_values
|
175 |
-
|
176 |
-
|
177 |
-
|
178 |
-
|
179 |
-
|
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 |
-
|
183 |
-
|
184 |
-
|
185 |
-
|
186 |
-
|
187 |
-
|
188 |
-
|
|
|
|
|
|
|
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
|
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.
|
253 |
|
254 |
def set_input_embeddings(self, value: Union[SharedEmbedding, nn.Embedding]) -> None:
|
255 |
-
self.transformer.
|
|
|
|
|
|
|
|
|
|
|
256 |
|
257 |
-
def
|
258 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
259 |
|
260 |
-
def
|
261 |
-
self.
|
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 |
-
|
273 |
-
|
274 |
-
|
275 |
-
|
|
|
|
|
|
|
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
|
293 |
|
294 |
def activation_checkpointing_fn(self, module: nn.Module) -> bool:
|
295 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
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}.')
|