lvkaokao commited on
Commit
dac7b46
1 Parent(s): ae15e66

add model.

Browse files
adapt_tokenizer.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Union
2
+ from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
3
+ Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
4
+ NUM_SENTINEL_TOKENS: int = 100
5
+
6
+ def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):
7
+ """Adds sentinel tokens and padding token (if missing).
8
+
9
+ Expands the tokenizer vocabulary to include sentinel tokens
10
+ used in mixture-of-denoiser tasks as well as a padding token.
11
+
12
+ All added tokens are added as special tokens. No tokens are
13
+ added if sentinel tokens and padding token already exist.
14
+ """
15
+ sentinels_to_add = [f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)]
16
+ tokenizer.add_tokens(sentinels_to_add, special_tokens=True)
17
+ if tokenizer.pad_token is None:
18
+ tokenizer.add_tokens('<pad>', special_tokens=True)
19
+ tokenizer.pad_token = '<pad>'
20
+ assert tokenizer.pad_token_id is not None
21
+ sentinels = ''.join([f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)])
22
+ _sentinel_token_ids = tokenizer(sentinels, add_special_tokens=False).input_ids
23
+ tokenizer.sentinel_token_ids = _sentinel_token_ids
24
+
25
+ class AutoTokenizerForMOD(AutoTokenizer):
26
+ """AutoTokenizer + Adaptation for MOD.
27
+
28
+ A simple wrapper around AutoTokenizer to make instantiating
29
+ an MOD-adapted tokenizer a bit easier.
30
+
31
+ MOD-adapted tokenizers have sentinel tokens (e.g., <extra_id_0>),
32
+ a padding token, and a property to get the token ids of the
33
+ sentinel tokens.
34
+ """
35
+
36
+ @classmethod
37
+ def from_pretrained(cls, *args, **kwargs):
38
+ """See `AutoTokenizer.from_pretrained` docstring."""
39
+ tokenizer = super().from_pretrained(*args, **kwargs)
40
+ adapt_tokenizer_for_denoising(tokenizer)
41
+ return tokenizer
attention.py ADDED
@@ -0,0 +1,276 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Attention layers."""
2
+ import math
3
+ import warnings
4
+ from typing import Optional
5
+ import torch
6
+ import torch.nn as nn
7
+ from einops import rearrange
8
+ from torch import nn
9
+ from .norm import LPLayerNorm
10
+
11
+ def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool):
12
+ if original_is_causal and num_query_tokens != num_key_tokens:
13
+ if num_query_tokens != 1:
14
+ raise NotImplementedError('MPT does not support query and key with different number of tokens, unless number of query tokens is 1.')
15
+ else:
16
+ return False
17
+ return original_is_causal
18
+
19
+ def scaled_multihead_dot_product_attention(query, key, value, n_heads, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
20
+ q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)
21
+ k = rearrange(key, 'b s (h d) -> b h d s', h=1 if multiquery else n_heads)
22
+ v = rearrange(value, 'b s (h d) -> b h s d', h=1 if multiquery else n_heads)
23
+ min_val = torch.finfo(q.dtype).min
24
+ (b, _, s_q, d) = q.shape
25
+ s_k = k.size(-1)
26
+ if softmax_scale is None:
27
+ softmax_scale = 1 / math.sqrt(d)
28
+ attn_weight = q.matmul(k) * softmax_scale
29
+ if attn_bias is not None:
30
+ if attn_bias.size(-1) != 1 and attn_bias.size(-1) != s_k or (attn_bias.size(-2) != 1 and attn_bias.size(-2) != s_q):
31
+ raise RuntimeError(f'attn_bias (shape: {attn_bias.shape}) is expected to broadcast to shape: {attn_weight.shape}.')
32
+ attn_weight = attn_weight + attn_bias
33
+ if key_padding_mask is not None:
34
+ if attn_bias is not None:
35
+ warnings.warn('Propogating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ' + 'unneccessary computation/memory usage. Consider integrating ' + 'into attn_bias once and passing that to each attention ' + 'module instead.')
36
+ attn_weight = attn_weight.masked_fill(~key_padding_mask.view((b, 1, 1, s_k)), min_val)
37
+ if is_causal:
38
+ s = max(s_q, s_k)
39
+ causal_mask = attn_weight.new_ones(s, s, dtype=torch.float16)
40
+ causal_mask = causal_mask.tril()
41
+ causal_mask = causal_mask.to(torch.bool)
42
+ causal_mask = ~causal_mask
43
+ causal_mask = causal_mask[-s_q:, -s_k:]
44
+ attn_weight = attn_weight.masked_fill(causal_mask.view(1, 1, s_q, s_k), min_val)
45
+ attn_weight = torch.softmax(attn_weight, dim=-1)
46
+ if dropout_p:
47
+ attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p, training=training, inplace=True)
48
+ out = attn_weight.matmul(v)
49
+ out = rearrange(out, 'b h s d -> b s (h d)')
50
+ if needs_weights:
51
+ return (out, attn_weight)
52
+ return (out, None)
53
+
54
+ def check_valid_inputs(*tensors, valid_dtypes=[torch.float16, torch.bfloat16]):
55
+ for tensor in tensors:
56
+ if tensor.dtype not in valid_dtypes:
57
+ raise TypeError(f'tensor.dtype={tensor.dtype!r} must be in valid_dtypes={valid_dtypes!r}.')
58
+ if not tensor.is_cuda:
59
+ raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
60
+
61
+ def flash_attn_fn(query, key, value, n_heads, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
62
+ try:
63
+ from flash_attn import bert_padding, flash_attn_interface
64
+ except:
65
+ raise RuntimeError('Please install flash-attn==1.0.3.post0')
66
+ check_valid_inputs(query, key, value)
67
+ if attn_bias is not None:
68
+ raise NotImplementedError(f'attn_bias not implemented for flash attn.')
69
+ (batch_size, seqlen) = query.shape[:2]
70
+ if key_padding_mask is None:
71
+ key_padding_mask = torch.ones_like(key[:, :, 0], dtype=torch.bool)
72
+ query_padding_mask = key_padding_mask[:, -query.size(1):]
73
+ (query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(query, query_padding_mask)
74
+ query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
75
+ (key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(key, key_padding_mask)
76
+ key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=1 if multiquery else n_heads)
77
+ (value_unpad, _, _, _) = bert_padding.unpad_input(value, key_padding_mask)
78
+ value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=1 if multiquery else n_heads)
79
+ if multiquery:
80
+ key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size(-1))
81
+ value_unpad = value_unpad.expand(value_unpad.size(0), n_heads, value_unpad.size(-1))
82
+ dropout_p = dropout_p if training else 0.0
83
+ reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
84
+ output_unpad = flash_attn_interface.flash_attn_unpadded_func(query_unpad, key_unpad, value_unpad, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
85
+ output = bert_padding.pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size, seqlen)
86
+ return (output, None)
87
+
88
+ def triton_flash_attn_fn(query, key, value, n_heads, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
89
+ try:
90
+ from flash_attn import flash_attn_triton
91
+ except:
92
+ raise RuntimeError('Please install flash-attn==1.0.3.post0 and triton==2.0.0.dev20221202')
93
+ check_valid_inputs(query, key, value)
94
+ if dropout_p:
95
+ raise NotImplementedError(f'Dropout not implemented for attn_impl: triton.')
96
+ if needs_weights:
97
+ raise NotImplementedError(f'attn_impl: triton cannot return attn weights.')
98
+ if key_padding_mask is not None:
99
+ warnings.warn('Propagating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ' + 'unnecessary computation/memory usage. Consider integrating ' + 'into attn_bias once and passing that to each attention ' + 'module instead.')
100
+ (b_size, s_k) = key_padding_mask.shape[:2]
101
+ if attn_bias is None:
102
+ attn_bias = query.new_zeros(b_size, 1, 1, s_k)
103
+ attn_bias = attn_bias.masked_fill(~key_padding_mask.view((b_size, 1, 1, s_k)), torch.finfo(query.dtype).min)
104
+ query = rearrange(query, 'b s (h d) -> b s h d', h=n_heads)
105
+ key = rearrange(key, 'b s (h d) -> b s h d', h=1 if multiquery else n_heads)
106
+ value = rearrange(value, 'b s (h d) -> b s h d', h=1 if multiquery else n_heads)
107
+ if multiquery:
108
+ key = key.expand(*key.shape[:2], n_heads, key.size(-1))
109
+ value = value.expand(*value.shape[:2], n_heads, value.size(-1))
110
+ reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
111
+ attn_output = flash_attn_triton.flash_attn_func(query, key, value, attn_bias, reset_is_causal, softmax_scale)
112
+ output = attn_output.view(*attn_output.shape[:2], -1)
113
+ return (output, None)
114
+
115
+ class MultiheadAttention(nn.Module):
116
+ """Multi-head self attention.
117
+
118
+ Using torch or triton attention implemetation enables user to also use
119
+ additive bias.
120
+ """
121
+
122
+ 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, low_precision_layernorm: bool=False, device: Optional[str]=None):
123
+ super().__init__()
124
+ self.attn_impl = attn_impl
125
+ self.clip_qkv = clip_qkv
126
+ self.qk_ln = qk_ln
127
+ self.d_model = d_model
128
+ self.n_heads = n_heads
129
+ self.softmax_scale = softmax_scale
130
+ if self.softmax_scale is None:
131
+ self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
132
+ self.attn_dropout_p = attn_pdrop
133
+ self.Wqkv = nn.Linear(self.d_model, 3 * self.d_model, device=device)
134
+ fuse_splits = (d_model, 2 * d_model)
135
+ self.Wqkv._fused = (0, fuse_splits)
136
+ if self.qk_ln:
137
+ layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
138
+ self.q_ln = layernorm_class(self.d_model, device=device)
139
+ self.k_ln = layernorm_class(self.d_model, device=device)
140
+ if self.attn_impl == 'flash':
141
+ self.attn_fn = flash_attn_fn
142
+ elif self.attn_impl == 'triton':
143
+ self.attn_fn = triton_flash_attn_fn
144
+ warnings.warn('While `attn_impl: triton` can be faster than `attn_impl: flash` ' + 'it uses more memory. When training larger models this can trigger ' + 'alloc retries which hurts performance. If encountered, we recommend ' + 'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.')
145
+ elif self.attn_impl == 'torch':
146
+ self.attn_fn = scaled_multihead_dot_product_attention
147
+ if torch.cuda.is_available():
148
+ warnings.warn('Using `attn_impl: torch`. If your model does not use `alibi` or ' + '`prefix_lm` we recommend using `attn_impl: flash` otherwise ' + 'we recommend using `attn_impl: triton`.')
149
+ else:
150
+ raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
151
+ self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
152
+ self.out_proj._is_residual = True
153
+
154
+ def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
155
+ qkv = self.Wqkv(x)
156
+ if self.clip_qkv:
157
+ qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
158
+ (query, key, value) = qkv.chunk(3, dim=2)
159
+ key_padding_mask = attention_mask
160
+ if self.qk_ln:
161
+ dtype = query.dtype
162
+ query = self.q_ln(query).to(dtype)
163
+ key = self.k_ln(key).to(dtype)
164
+ if past_key_value is not None:
165
+ if len(past_key_value) != 0:
166
+ key = torch.cat([past_key_value[0], key], dim=1)
167
+ value = torch.cat([past_key_value[1], value], dim=1)
168
+ past_key_value = (key, value)
169
+ if attn_bias is not None:
170
+ attn_bias = attn_bias[:, :, -query.size(1):, -key.size(1):]
171
+ (context, attn_weights) = self.attn_fn(query, key, value, self.n_heads, 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)
172
+ return (self.out_proj(context), attn_weights, past_key_value)
173
+
174
+ class MultiQueryAttention(nn.Module):
175
+ """Multi-Query self attention.
176
+
177
+ Using torch or triton attention implemetation enables user to also use
178
+ additive bias.
179
+ """
180
+
181
+ 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, low_precision_layernorm: bool=False, device: Optional[str]=None):
182
+ super().__init__()
183
+ self.attn_impl = attn_impl
184
+ self.clip_qkv = clip_qkv
185
+ self.qk_ln = qk_ln
186
+ self.d_model = d_model
187
+ self.n_heads = n_heads
188
+ self.head_dim = d_model // n_heads
189
+ self.softmax_scale = softmax_scale
190
+ if self.softmax_scale is None:
191
+ self.softmax_scale = 1 / math.sqrt(self.head_dim)
192
+ self.attn_dropout_p = attn_pdrop
193
+ self.Wqkv = nn.Linear(d_model, d_model + 2 * self.head_dim, device=device)
194
+ fuse_splits = (d_model, d_model + self.head_dim)
195
+ self.Wqkv._fused = (0, fuse_splits)
196
+ if self.qk_ln:
197
+ layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
198
+ self.q_ln = layernorm_class(d_model, device=device)
199
+ self.k_ln = layernorm_class(self.head_dim, device=device)
200
+ if self.attn_impl == 'flash':
201
+ self.attn_fn = flash_attn_fn
202
+ elif self.attn_impl == 'triton':
203
+ self.attn_fn = triton_flash_attn_fn
204
+ warnings.warn('While `attn_impl: triton` can be faster than `attn_impl: flash` ' + 'it uses more memory. When training larger models this can trigger ' + 'alloc retries which hurts performance. If encountered, we recommend ' + 'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.')
205
+ elif self.attn_impl == 'torch':
206
+ self.attn_fn = scaled_multihead_dot_product_attention
207
+ if torch.cuda.is_available():
208
+ warnings.warn('Using `attn_impl: torch`. If your model does not use `alibi` or ' + '`prefix_lm` we recommend using `attn_impl: flash` otherwise ' + 'we recommend using `attn_impl: triton`.')
209
+ else:
210
+ raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
211
+ self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
212
+ self.out_proj._is_residual = True
213
+
214
+ def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
215
+ qkv = self.Wqkv(x)
216
+ if self.clip_qkv:
217
+ qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
218
+ (query, key, value) = qkv.split([self.d_model, self.head_dim, self.head_dim], dim=2)
219
+ key_padding_mask = attention_mask
220
+ if self.qk_ln:
221
+ dtype = query.dtype
222
+ query = self.q_ln(query).to(dtype)
223
+ key = self.k_ln(key).to(dtype)
224
+ if past_key_value is not None:
225
+ if len(past_key_value) != 0:
226
+ key = torch.cat([past_key_value[0], key], dim=1)
227
+ value = torch.cat([past_key_value[1], value], dim=1)
228
+ past_key_value = (key, value)
229
+ if attn_bias is not None:
230
+ attn_bias = attn_bias[:, :, -query.size(1):, -key.size(1):]
231
+ (context, attn_weights) = self.attn_fn(query, key, value, self.n_heads, 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, multiquery=True)
232
+ return (self.out_proj(context), attn_weights, past_key_value)
233
+
234
+ def attn_bias_shape(attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_sequence_id):
235
+ if attn_impl == 'flash':
236
+ return None
237
+ elif attn_impl in ['torch', 'triton']:
238
+ if alibi:
239
+ if (prefix_lm or not causal) or use_sequence_id:
240
+ return (1, n_heads, seq_len, seq_len)
241
+ return (1, n_heads, 1, seq_len)
242
+ elif prefix_lm or use_sequence_id:
243
+ return (1, 1, seq_len, seq_len)
244
+ return None
245
+ else:
246
+ raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
247
+
248
+ def build_attn_bias(attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=False, alibi_bias_max=8):
249
+ if attn_impl == 'flash':
250
+ return None
251
+ elif attn_impl in ['torch', 'triton']:
252
+ if alibi:
253
+ (device, dtype) = (attn_bias.device, attn_bias.dtype)
254
+ attn_bias = attn_bias.add(build_alibi_bias(n_heads, seq_len, full=not causal, alibi_bias_max=alibi_bias_max, device=device, dtype=dtype))
255
+ return attn_bias
256
+ else:
257
+ raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
258
+
259
+ def gen_slopes(n_heads, alibi_bias_max=8, device=None):
260
+ _n_heads = 2 ** math.ceil(math.log2(n_heads))
261
+ m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)
262
+ m = m.mul(alibi_bias_max / _n_heads)
263
+ slopes = 1.0 / torch.pow(2, m)
264
+ if _n_heads != n_heads:
265
+ slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
266
+ return slopes.view(1, n_heads, 1, 1)
267
+
268
+ def build_alibi_bias(n_heads, seq_len, full=False, alibi_bias_max=8, device=None, dtype=None):
269
+ alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, 1, seq_len)
270
+ if full:
271
+ alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, seq_len, 1)
272
+ alibi_bias = alibi_bias.abs().mul(-1)
273
+ slopes = gen_slopes(n_heads, alibi_bias_max, device=device)
274
+ alibi_bias = alibi_bias * slopes
275
+ return alibi_bias.to(dtype=dtype)
276
+ ATTN_CLASS_REGISTRY = {'multihead_attention': MultiheadAttention, 'multiquery_attention': MultiQueryAttention}
blocks.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """GPT Blocks used for the GPT Model."""
2
+ from typing import Dict, Optional, Tuple
3
+ import torch
4
+ import torch.nn as nn
5
+ from .attention import ATTN_CLASS_REGISTRY
6
+ from .norm import NORM_CLASS_REGISTRY
7
+
8
+ class MPTMLP(nn.Module):
9
+
10
+ def __init__(self, d_model: int, expansion_ratio: int, device: Optional[str]=None):
11
+ super().__init__()
12
+ self.up_proj = nn.Linear(d_model, expansion_ratio * d_model, device=device)
13
+ self.act = nn.GELU(approximate='none')
14
+ self.down_proj = nn.Linear(expansion_ratio * d_model, d_model, device=device)
15
+ self.down_proj._is_residual = True
16
+
17
+ def forward(self, x):
18
+ return self.down_proj(self.act(self.up_proj(x)))
19
+
20
+ class MPTBlock(nn.Module):
21
+
22
+ def __init__(self, d_model: int, n_heads: int, expansion_ratio: int, attn_config: Dict={'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', device: Optional[str]=None, **kwargs):
23
+ del kwargs
24
+ super().__init__()
25
+ norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
26
+ attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]
27
+ self.norm_1 = norm_class(d_model, device=device)
28
+ self.attn = attn_class(attn_impl=attn_config['attn_impl'], clip_qkv=attn_config['clip_qkv'], qk_ln=attn_config['qk_ln'], softmax_scale=attn_config['softmax_scale'], attn_pdrop=attn_config['attn_pdrop'], d_model=d_model, n_heads=n_heads, device=device)
29
+ self.norm_2 = norm_class(d_model, device=device)
30
+ self.ffn = MPTMLP(d_model=d_model, expansion_ratio=expansion_ratio, device=device)
31
+ self.resid_attn_dropout = nn.Dropout(resid_pdrop)
32
+ self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
33
+
34
+ def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.ByteTensor]=None, is_causal: bool=True) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor]]]:
35
+ a = self.norm_1(x)
36
+ (b, _, past_key_value) = self.attn(a, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=is_causal)
37
+ x = x + self.resid_attn_dropout(b)
38
+ m = self.norm_2(x)
39
+ n = self.ffn(m)
40
+ x = x + self.resid_ffn_dropout(n)
41
+ return (x, past_key_value)
config.json ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "mosaicml/mpt-7b",
3
+ "architectures": [
4
+ "MPTForCausalLM"
5
+ ],
6
+ "attn_config": {
7
+ "alibi": true,
8
+ "alibi_bias_max": 8,
9
+ "attn_impl": "torch",
10
+ "attn_pdrop": 0,
11
+ "attn_type": "multihead_attention",
12
+ "attn_uses_sequence_id": false,
13
+ "clip_qkv": null,
14
+ "prefix_lm": false,
15
+ "qk_ln": false,
16
+ "softmax_scale": null
17
+ },
18
+ "auto_map": {
19
+ "AutoConfig": "mosaicml/mpt-7b--configuration_mpt.MPTConfig",
20
+ "AutoModelForCausalLM": "mosaicml/mpt-7b--modeling_mpt.MPTForCausalLM"
21
+ },
22
+ "d_model": 4096,
23
+ "emb_pdrop": 0,
24
+ "embedding_fraction": 1.0,
25
+ "expansion_ratio": 4,
26
+ "init_config": {
27
+ "emb_init_std": null,
28
+ "emb_init_uniform_lim": null,
29
+ "fan_mode": "fan_in",
30
+ "init_div_is_residual": true,
31
+ "init_gain": 0,
32
+ "init_nonlinearity": "relu",
33
+ "init_std": 0.02,
34
+ "name": "kaiming_normal_",
35
+ "verbose": 0
36
+ },
37
+ "init_device": "cpu",
38
+ "learned_pos_emb": true,
39
+ "logit_scale": null,
40
+ "max_seq_len": 2048,
41
+ "model_type": "mpt",
42
+ "n_heads": 32,
43
+ "n_layers": 32,
44
+ "no_bias": true,
45
+ "norm_type": "low_precision_layernorm",
46
+ "resid_pdrop": 0,
47
+ "tokenizer_name": "EleutherAI/gpt-neox-20b",
48
+ "torch_dtype": "bfloat16",
49
+ "transformers_version": "4.30.2",
50
+ "use_cache": false,
51
+ "verbose": 0,
52
+ "vocab_size": 50279
53
+ }
configuration_mpt.py ADDED
@@ -0,0 +1,118 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """A HuggingFace-style model configuration."""
2
+ from typing import Dict, Optional, Union
3
+ from transformers import PretrainedConfig
4
+ attn_config_defaults: Dict = {'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}
5
+ init_config_defaults: Dict = {'name': 'kaiming_normal_', 'fan_mode': 'fan_in', 'init_nonlinearity': 'relu'}
6
+
7
+ class MPTConfig(PretrainedConfig):
8
+ model_type = 'mpt'
9
+
10
+ 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, init_device: str='cpu', logit_scale: Optional[Union[float, str]]=None, no_bias: bool=False, verbose: int=0, embedding_fraction: float=1.0, norm_type: str='low_precision_layernorm', use_cache: bool=False, init_config: Dict=init_config_defaults, **kwargs):
11
+ """The MPT configuration class.
12
+
13
+ Args:
14
+ d_model (int): The size of the embedding dimension of the model.
15
+ n_heads (int): The number of attention heads.
16
+ n_layers (int): The number of layers in the model.
17
+ expansion_ratio (int): The ratio of the up/down scale in the MLP.
18
+ max_seq_len (int): The maximum sequence length of the model.
19
+ vocab_size (int): The size of the vocabulary.
20
+ resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.
21
+ emb_pdrop (float): The dropout probability for the embedding layer.
22
+ learned_pos_emb (bool): Whether to use learned positional embeddings
23
+ attn_config (Dict): A dictionary used to configure the model's attention module:
24
+ attn_type (str): type of attention to use. Options: multihead_attention, multiquery_attention
25
+ attn_pdrop (float): The dropout probability for the attention layers.
26
+ attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.
27
+ qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.
28
+ clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to
29
+ this value.
30
+ softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,
31
+ use the default scale of ``1/sqrt(d_keys)``.
32
+ prefix_lm (Optional[bool]): Whether the model should operate as a Prefix LM. This requires passing an
33
+ extra `prefix_mask` argument which indicates which tokens belong to the prefix. Tokens in the prefix
34
+ can attend to one another bi-directionally. Tokens outside the prefix use causal attention.
35
+ attn_uses_sequence_id (Optional[bool]): Whether to restrict attention to tokens that have the same sequence_id.
36
+ When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates
37
+ which sub-sequence each token belongs to.
38
+ Defaults to ``False`` meaning any provided `sequence_id` will be ignored.
39
+ alibi (bool): Whether to use the alibi bias instead of position embeddings.
40
+ alibi_bias_max (int): The maximum value of the alibi bias.
41
+ init_device (str): The device to use for parameter initialization.
42
+ logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.
43
+ no_bias (bool): Whether to use bias in all layers.
44
+ verbose (int): The verbosity level. 0 is silent.
45
+ embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.
46
+ norm_type (str): choose type of norm to use
47
+ multiquery_attention (bool): Whether to use multiquery attention implementation.
48
+ use_cache (bool): Whether or not the model should return the last key/values attentions
49
+ init_config (Dict): A dictionary used to configure the model initialization:
50
+ init_config.name: The parameter initialization scheme to use. Options: 'default_', 'baseline_',
51
+ 'kaiming_uniform_', 'kaiming_normal_', 'neox_init_', 'small_init_', 'xavier_uniform_', or
52
+ 'xavier_normal_'. These mimic the parameter initialization methods in PyTorch.
53
+ init_div_is_residual (Union[int, float, str, bool]): Value to divide initial weights by if ``module._is_residual`` is True.
54
+ emb_init_std (Optional[float]): The standard deviation of the normal distribution used to initialize the embedding layer.
55
+ emb_init_uniform_lim (Optional[Union[Tuple[float, float], float]]): The lower and upper limits of the uniform distribution
56
+ used to initialize the embedding layer. Mutually exclusive with ``emb_init_std``.
57
+ init_std (float): The standard deviation of the normal distribution used to initialize the model,
58
+ if using the baseline_ parameter initialization scheme.
59
+ init_gain (float): The gain to use for parameter initialization with kaiming or xavier initialization schemes.
60
+ fan_mode (str): The fan mode to use for parameter initialization with kaiming initialization schemes.
61
+ init_nonlinearity (str): The nonlinearity to use for parameter initialization with kaiming initialization schemes.
62
+ ---
63
+ See llmfoundry.models.utils.param_init_fns.py for info on other param init config options
64
+ """
65
+ self.d_model = d_model
66
+ self.n_heads = n_heads
67
+ self.n_layers = n_layers
68
+ self.expansion_ratio = expansion_ratio
69
+ self.max_seq_len = max_seq_len
70
+ self.vocab_size = vocab_size
71
+ self.resid_pdrop = resid_pdrop
72
+ self.emb_pdrop = emb_pdrop
73
+ self.learned_pos_emb = learned_pos_emb
74
+ self.attn_config = attn_config
75
+ self.init_device = init_device
76
+ self.logit_scale = logit_scale
77
+ self.no_bias = no_bias
78
+ self.verbose = verbose
79
+ self.embedding_fraction = embedding_fraction
80
+ self.norm_type = norm_type
81
+ self.use_cache = use_cache
82
+ self.init_config = init_config
83
+ if 'name' in kwargs:
84
+ del kwargs['name']
85
+ if 'loss_fn' in kwargs:
86
+ del kwargs['loss_fn']
87
+ super().__init__(**kwargs)
88
+ self._validate_config()
89
+
90
+ def _set_config_defaults(self, config, config_defaults):
91
+ for (k, v) in config_defaults.items():
92
+ if k not in config:
93
+ config[k] = v
94
+ return config
95
+
96
+ def _validate_config(self):
97
+ self.attn_config = self._set_config_defaults(self.attn_config, attn_config_defaults)
98
+ self.init_config = self._set_config_defaults(self.init_config, init_config_defaults)
99
+ if self.d_model % self.n_heads != 0:
100
+ raise ValueError('d_model must be divisible by n_heads')
101
+ if any((prob < 0 or prob > 1 for prob in [self.attn_config['attn_pdrop'], self.resid_pdrop, self.emb_pdrop])):
102
+ raise ValueError("self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1")
103
+ if self.attn_config['attn_impl'] not in ['torch', 'flash', 'triton']:
104
+ raise ValueError(f"Unknown attn_impl={self.attn_config['attn_impl']}")
105
+ if self.attn_config['prefix_lm'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
106
+ raise NotImplementedError('prefix_lm only implemented with torch and triton attention.')
107
+ if self.attn_config['alibi'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
108
+ raise NotImplementedError('alibi only implemented with torch and triton attention.')
109
+ if self.attn_config['attn_uses_sequence_id'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
110
+ raise NotImplementedError('attn_uses_sequence_id only implemented with torch and triton attention.')
111
+ if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
112
+ raise ValueError('model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!')
113
+ if isinstance(self.logit_scale, str) and self.logit_scale != 'inv_sqrt_d_model':
114
+ raise ValueError(f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
115
+ if self.init_config.get('name', None) is None:
116
+ raise ValueError(f"self.init_config={self.init_config!r} 'name' needs to be set.")
117
+ if not self.learned_pos_emb and (not self.attn_config['alibi']):
118
+ raise ValueError(f'Positional information must be provided to the model using either learned_pos_emb or alibi.')
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": 0,
4
+ "transformers_version": "4.30.2",
5
+ "use_cache": false
6
+ }
hf_prefixlm_converter.py ADDED
@@ -0,0 +1,415 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Converts Huggingface Causal LM to Prefix LM.
2
+
3
+ Conversion does lightweight surgery on a HuggingFace
4
+ Causal LM to convert it to a Prefix LM.
5
+
6
+ Prefix LMs accepts a `bidirectional_mask` input in `forward`
7
+ and treat the input prompt as the prefix in `generate`.
8
+ """
9
+ import math
10
+ import warnings
11
+ from types import MethodType
12
+ from typing import Any, Dict, List, Optional, Tuple, Union
13
+ import torch
14
+ from transformers.models.bloom.modeling_bloom import BaseModelOutputWithPastAndCrossAttentions, BloomForCausalLM, BloomModel, CausalLMOutputWithCrossAttentions, CrossEntropyLoss
15
+ from transformers.models.bloom.modeling_bloom import _expand_mask as _expand_mask_bloom
16
+ from transformers.models.bloom.modeling_bloom import _make_causal_mask as _make_causal_mask_bloom
17
+ from transformers.models.bloom.modeling_bloom import logging
18
+ from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
19
+ from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoForCausalLM
20
+ from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
21
+ from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
22
+ from transformers.models.opt.modeling_opt import OPTForCausalLM
23
+ from transformers.models.opt.modeling_opt import _expand_mask as _expand_mask_opt
24
+ from transformers.models.opt.modeling_opt import _make_causal_mask as _make_causal_mask_opt
25
+ logger = logging.get_logger(__name__)
26
+ _SUPPORTED_GPT_MODELS = (GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM)
27
+ CAUSAL_GPT_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM]
28
+
29
+ def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_TYPES:
30
+ """Converts a GPT-style Causal LM to a Prefix LM.
31
+
32
+ Supported HuggingFace model classes:
33
+ - `GPT2LMHeadModel`
34
+ - `GPTNeoForCausalLM`
35
+ - `GPTNeoXForCausalLM`
36
+ - `GPTJForCausalLM`
37
+
38
+ See `convert_hf_causal_lm_to_prefix_lm` for more details.
39
+ """
40
+ if hasattr(model, '_prefix_lm_converted'):
41
+ return model
42
+ assert isinstance(model, _SUPPORTED_GPT_MODELS)
43
+ assert model.config.add_cross_attention == False, 'Only supports GPT-style decoder-only models'
44
+
45
+ def _get_attn_modules(model: CAUSAL_GPT_TYPES) -> List[torch.nn.Module]:
46
+ """Helper that gets a list of the model's attention modules.
47
+
48
+ Each module has a `bias` buffer used for causal masking. The Prefix LM
49
+ conversion adds logic to dynamically manipulate these biases to support
50
+ Prefix LM attention masking.
51
+ """
52
+ attn_modules = []
53
+ if isinstance(model, GPTNeoXForCausalLM):
54
+ blocks = model.gpt_neox.layers
55
+ else:
56
+ blocks = model.transformer.h
57
+ for block in blocks:
58
+ if isinstance(model, GPTNeoForCausalLM):
59
+ if block.attn.attention_type != 'global':
60
+ continue
61
+ attn_module = block.attn.attention
62
+ elif isinstance(model, GPTNeoXForCausalLM):
63
+ attn_module = block.attention
64
+ else:
65
+ attn_module = block.attn
66
+ attn_modules.append(attn_module)
67
+ return attn_modules
68
+ setattr(model, '_original_forward', getattr(model, 'forward'))
69
+ setattr(model, '_original_generate', getattr(model, 'generate'))
70
+
71
+ def forward(self: CAUSAL_GPT_TYPES, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]]=None, attention_mask: Optional[torch.FloatTensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
72
+ """Wraps original forward to enable PrefixLM attention."""
73
+
74
+ def call_og_forward():
75
+ if isinstance(self, GPTNeoXForCausalLM):
76
+ return self._original_forward(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
77
+ else:
78
+ return self._original_forward(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
79
+ if bidirectional_mask is None:
80
+ return call_og_forward()
81
+ assert isinstance(bidirectional_mask, torch.Tensor)
82
+ attn_modules = _get_attn_modules(model)
83
+ (b, s) = bidirectional_mask.shape
84
+ max_length = attn_modules[0].bias.shape[-1]
85
+ if s > max_length:
86
+ raise ValueError(f'bidirectional_mask sequence length (={s}) exceeds the ' + f'max length allowed by the model ({max_length}).')
87
+ assert s <= max_length
88
+ if s < max_length:
89
+ pad = torch.zeros((int(b), int(max_length - s)), dtype=bidirectional_mask.dtype, device=bidirectional_mask.device)
90
+ bidirectional_mask = torch.cat([bidirectional_mask, pad], dim=1)
91
+ bidirectional = bidirectional_mask.unsqueeze(1).unsqueeze(1)
92
+ for attn_module in attn_modules:
93
+ attn_module.bias.data = torch.logical_or(attn_module.bias.data, bidirectional)
94
+ output = call_og_forward()
95
+ for attn_module in attn_modules:
96
+ attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]
97
+ return output
98
+
99
+ def generate(self: CAUSAL_GPT_TYPES, *args: tuple, **kwargs: Dict[str, Any]):
100
+ """Wraps original generate to enable PrefixLM attention."""
101
+ attn_modules = _get_attn_modules(model)
102
+ for attn_module in attn_modules:
103
+ attn_module.bias.data[:] = 1
104
+ output = self._original_generate(*args, **kwargs)
105
+ for attn_module in attn_modules:
106
+ attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]
107
+ return output
108
+ setattr(model, 'forward', MethodType(forward, model))
109
+ setattr(model, 'generate', MethodType(generate, model))
110
+ setattr(model, '_prefix_lm_converted', True)
111
+ return model
112
+
113
+ def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCausalLM:
114
+ """Converts a BLOOM Causal LM to a Prefix LM.
115
+
116
+ Supported HuggingFace model classes:
117
+ - `BloomForCausalLM`
118
+
119
+ See `convert_hf_causal_lm_to_prefix_lm` for more details.
120
+ """
121
+ if hasattr(model, '_prefix_lm_converted'):
122
+ return model
123
+ assert isinstance(model, BloomForCausalLM)
124
+ assert model.config.add_cross_attention == False, 'Only supports BLOOM decoder-only models'
125
+
126
+ def _prepare_attn_mask(self: BloomModel, attention_mask: torch.Tensor, bidirectional_mask: Optional[torch.Tensor], input_shape: Tuple[int, int], past_key_values_length: int) -> torch.BoolTensor:
127
+ combined_attention_mask = None
128
+ device = attention_mask.device
129
+ (_, src_length) = input_shape
130
+ if src_length > 1:
131
+ combined_attention_mask = _make_causal_mask_bloom(input_shape, device=device, past_key_values_length=past_key_values_length)
132
+ if bidirectional_mask is not None:
133
+ assert attention_mask.shape == bidirectional_mask.shape
134
+ expanded_bidirectional_mask = _expand_mask_bloom(bidirectional_mask, tgt_length=src_length)
135
+ combined_attention_mask = torch.logical_and(combined_attention_mask, expanded_bidirectional_mask)
136
+ expanded_attn_mask = _expand_mask_bloom(attention_mask, tgt_length=src_length)
137
+ combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask
138
+ return combined_attention_mask
139
+
140
+ def _build_alibi_tensor(self: BloomModel, batch_size: int, query_length: int, key_length: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
141
+ num_heads = self.config.n_head
142
+ closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
143
+ base = torch.tensor(2 ** (-2 ** (-(math.log2(closest_power_of_2) - 3))), device=device, dtype=torch.float32)
144
+ powers = torch.arange(1, 1 + closest_power_of_2, device=device, dtype=torch.int32)
145
+ slopes = torch.pow(base, powers)
146
+ if closest_power_of_2 != num_heads:
147
+ extra_base = torch.tensor(2 ** (-2 ** (-(math.log2(2 * closest_power_of_2) - 3))), device=device, dtype=torch.float32)
148
+ num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
149
+ extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=device, dtype=torch.int32)
150
+ slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
151
+ qa = torch.arange(query_length, device=device, dtype=torch.int32).view(-1, 1)
152
+ ka = torch.arange(key_length, device=device, dtype=torch.int32).view(1, -1)
153
+ diffs = qa - ka + key_length - query_length
154
+ diffs = -diffs.abs()
155
+ alibi = slopes.view(1, num_heads, 1, 1) * diffs.view(1, 1, query_length, key_length)
156
+ alibi = alibi.expand(batch_size, -1, -1, -1).reshape(-1, query_length, key_length)
157
+ return alibi.to(dtype)
158
+ KeyValueT = Tuple[torch.Tensor, torch.Tensor]
159
+
160
+ def forward(self: BloomModel, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[KeyValueT, ...]]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.LongTensor]=None, inputs_embeds: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **deprecated_arguments) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
161
+ if deprecated_arguments.pop('position_ids', False) is not False:
162
+ warnings.warn('`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. ' + 'You can safely ignore passing `position_ids`.', FutureWarning)
163
+ if len(deprecated_arguments) > 0:
164
+ raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
165
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
166
+ output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
167
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
168
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
169
+ if input_ids is not None and inputs_embeds is not None:
170
+ raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
171
+ elif input_ids is not None:
172
+ (batch_size, seq_length) = input_ids.shape
173
+ elif inputs_embeds is not None:
174
+ (batch_size, seq_length, _) = inputs_embeds.shape
175
+ else:
176
+ raise ValueError('You have to specify either input_ids or inputs_embeds')
177
+ if past_key_values is None:
178
+ past_key_values = tuple([None] * len(self.h))
179
+ head_mask = self.get_head_mask(head_mask, self.config.n_layer)
180
+ if inputs_embeds is None:
181
+ inputs_embeds = self.word_embeddings(input_ids)
182
+ hidden_states = self.word_embeddings_layernorm(inputs_embeds)
183
+ presents = () if use_cache else None
184
+ all_self_attentions = () if output_attentions else None
185
+ all_hidden_states = () if output_hidden_states else None
186
+ seq_length_with_past = seq_length
187
+ past_key_values_length = 0
188
+ if past_key_values[0] is not None:
189
+ tmp = past_key_values[0][0]
190
+ past_key_values_length = tmp.shape[2]
191
+ seq_length_with_past = seq_length_with_past + past_key_values_length
192
+ if attention_mask is None:
193
+ attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
194
+ else:
195
+ attention_mask = attention_mask.to(hidden_states.device)
196
+ alibi = self._build_alibi_tensor(batch_size=batch_size, query_length=seq_length, key_length=seq_length_with_past, dtype=hidden_states.dtype, device=hidden_states.device)
197
+ causal_mask = self._prepare_attn_mask(attention_mask, bidirectional_mask, input_shape=(batch_size, seq_length), past_key_values_length=past_key_values_length)
198
+ for (i, (block, layer_past)) in enumerate(zip(self.h, past_key_values)):
199
+ if output_hidden_states:
200
+ hst = (hidden_states,)
201
+ all_hidden_states = all_hidden_states + hst
202
+ if self.gradient_checkpointing and self.training:
203
+ if use_cache:
204
+ logger.warning('`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...')
205
+ use_cache = False
206
+
207
+ def create_custom_forward(module):
208
+
209
+ def custom_forward(*inputs):
210
+ return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)
211
+ return custom_forward
212
+ outputs = torch.utils.checkpoint.checkpoint(create_custom_forward(block), hidden_states, alibi, causal_mask, head_mask[i])
213
+ else:
214
+ outputs = block(hidden_states, layer_past=layer_past, attention_mask=causal_mask, head_mask=head_mask[i], use_cache=use_cache, output_attentions=output_attentions, alibi=alibi)
215
+ hidden_states = outputs[0]
216
+ if use_cache is True:
217
+ presents = presents + (outputs[1],)
218
+ if output_attentions:
219
+ oa = (outputs[2 if use_cache else 1],)
220
+ all_self_attentions = all_self_attentions + oa
221
+ hidden_states = self.ln_f(hidden_states)
222
+ if output_hidden_states:
223
+ hst = (hidden_states,)
224
+ all_hidden_states = all_hidden_states + hst
225
+ if not return_dict:
226
+ return tuple((v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None))
227
+ return BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=hidden_states, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attentions)
228
+ setattr(model.transformer, '_prepare_attn_mask', MethodType(_prepare_attn_mask, model.transformer))
229
+ setattr(model.transformer, '_build_alibi_tensor', MethodType(_build_alibi_tensor, model.transformer))
230
+ setattr(model.transformer, 'forward', MethodType(forward, model.transformer))
231
+ KeyValueT = Tuple[torch.Tensor, torch.Tensor]
232
+
233
+ def forward(self: BloomForCausalLM, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[KeyValueT, ...]]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, labels: Optional[torch.Tensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **deprecated_arguments) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
234
+ """Replacement forward method for BloomCausalLM."""
235
+ if deprecated_arguments.pop('position_ids', False) is not False:
236
+ warnings.warn('`position_ids` have no functionality in BLOOM and will be removed ' + 'in v5.0.0. You can safely ignore passing `position_ids`.', FutureWarning)
237
+ if len(deprecated_arguments) > 0:
238
+ raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
239
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
240
+ transformer_outputs = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask, bidirectional_mask=bidirectional_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
241
+ hidden_states = transformer_outputs[0]
242
+ lm_logits = self.lm_head(hidden_states)
243
+ loss = None
244
+ if labels is not None:
245
+ shift_logits = lm_logits[..., :-1, :].contiguous()
246
+ shift_labels = labels[..., 1:].contiguous()
247
+ (batch_size, seq_length, vocab_size) = shift_logits.shape
248
+ loss_fct = CrossEntropyLoss()
249
+ loss = loss_fct(shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length))
250
+ if not return_dict:
251
+ output = (lm_logits,) + transformer_outputs[1:]
252
+ return (loss,) + output if loss is not None else output
253
+ return CausalLMOutputWithCrossAttentions(loss=loss, logits=lm_logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions)
254
+
255
+ def prepare_inputs_for_generation(self: BloomForCausalLM, input_ids: torch.LongTensor, past: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, **kwargs) -> dict:
256
+ if past:
257
+ input_ids = input_ids[:, -1].unsqueeze(-1)
258
+ bidirectional_mask = None
259
+ if past[0][0].shape[0] == input_ids.shape[0]:
260
+ past = self._convert_to_bloom_cache(past)
261
+ else:
262
+ bidirectional_mask = torch.ones_like(input_ids)
263
+ return {'input_ids': input_ids, 'past_key_values': past, 'use_cache': True, 'attention_mask': attention_mask, 'bidirectional_mask': bidirectional_mask}
264
+ setattr(model, 'forward', MethodType(forward, model))
265
+ setattr(model, 'prepare_inputs_for_generation', MethodType(prepare_inputs_for_generation, model))
266
+ setattr(model, '_prefix_lm_converted', True)
267
+ return model
268
+
269
+ def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM:
270
+ """Converts an OPT Causal LM to a Prefix LM.
271
+
272
+ Supported HuggingFace model classes:
273
+ - `OPTForCausalLM`
274
+
275
+ See `convert_hf_causal_lm_to_prefix_lm` for more details.
276
+ """
277
+ if hasattr(model, '_prefix_lm_converted'):
278
+ return model
279
+ assert isinstance(model, OPTForCausalLM)
280
+ assert model.config.add_cross_attention == False, 'Only supports OPT decoder-only models'
281
+ setattr(model, '_original_forward', getattr(model, 'forward'))
282
+ setattr(model, '_original_generate', getattr(model, 'generate'))
283
+ model.model.decoder.bidirectional_mask = None
284
+
285
+ def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
286
+ combined_attention_mask = None
287
+ if input_shape[-1] > 1:
288
+ if self.bidirectional_mask == 'g':
289
+ (bsz, src_length) = input_shape
290
+ combined_attention_mask = torch.zeros((bsz, 1, src_length, src_length + past_key_values_length), dtype=inputs_embeds.dtype, device=inputs_embeds.device)
291
+ else:
292
+ combined_attention_mask = _make_causal_mask_opt(input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length).to(inputs_embeds.device)
293
+ if self.bidirectional_mask is not None:
294
+ assert attention_mask.shape == self.bidirectional_mask.shape
295
+ expanded_bidirectional_mask = _expand_mask_opt(self.bidirectional_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
296
+ combined_attention_mask = torch.maximum(expanded_bidirectional_mask, combined_attention_mask)
297
+ if attention_mask is not None:
298
+ expanded_attn_mask = _expand_mask_opt(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
299
+ combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
300
+ return combined_attention_mask
301
+ setattr(model.model.decoder, '_prepare_decoder_attention_mask', MethodType(_prepare_decoder_attention_mask, model.model.decoder))
302
+
303
+ def forward(self: OPTForCausalLM, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.ByteTensor]=None, head_mask: Optional[torch.Tensor]=None, past_key_values: Optional[List[torch.FloatTensor]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
304
+
305
+ def call_og_forward():
306
+ return self._original_forward(input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
307
+ if bidirectional_mask is None:
308
+ return call_og_forward()
309
+ self.model.decoder.bidirectional_mask = bidirectional_mask
310
+ try:
311
+ outputs = call_og_forward()
312
+ except:
313
+ self.model.decoder.bidirectional_mask = None
314
+ raise
315
+ self.model.decoder.bidirectional_mask = None
316
+ return outputs
317
+
318
+ def generate(self: OPTForCausalLM, *args: tuple, **kwargs: Dict[str, Any]):
319
+ """Wraps original generate to enable PrefixLM-style attention."""
320
+ self.model.decoder.bidirectional_mask = 'g'
321
+ try:
322
+ output = self._original_generate(*args, **kwargs)
323
+ except:
324
+ self.model.decoder.bidirectional_mask = None
325
+ raise
326
+ self.model.decoder.bidirectional_mask = None
327
+ return output
328
+ setattr(model, 'forward', MethodType(forward, model))
329
+ setattr(model, 'generate', MethodType(generate, model))
330
+ setattr(model, '_prefix_lm_converted', True)
331
+ return model
332
+ _SUPPORTED_HF_MODELS = _SUPPORTED_GPT_MODELS + (BloomForCausalLM, OPTForCausalLM)
333
+ CAUSAL_LM_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM, BloomForCausalLM, OPTForCausalLM]
334
+
335
+ def convert_hf_causal_lm_to_prefix_lm(model: CAUSAL_LM_TYPES) -> CAUSAL_LM_TYPES:
336
+ """Converts a HuggingFace Causal LM to a Prefix LM.
337
+
338
+ Supported HuggingFace model classes:
339
+ - `GPT2LMHeadModel`
340
+ - `GPTNeoForCausalLM`
341
+ - `GPTNeoXForCausalLM`
342
+ - `GPTJForCausalLM`
343
+ - `BloomForCausalLM`
344
+ - `OPTForCausalLM`
345
+
346
+ Conversion to a Prefix LM is done by modifying the `forward` method, and possibly also the
347
+ `generate` method and/or select underlying methods depending on the model class.
348
+
349
+ These changes preserve the model API, but add a new input to `forward`: "bidirectional_mask".
350
+
351
+ Notes on training:
352
+ To actually train the converted model as a Prefix LM, training batches will need to indicate
353
+ the prefix/target structure by including `bidirectional_mask` as part of the batch inputs.
354
+
355
+ **This is not a standard input and requires custom layers either within or after your dataloader.**
356
+
357
+ In addition to adding `bidirectional_mask` to the batch, this custom code should modify `labels`
358
+ such that `batch['labels'][batch['bidirectional_mask'] == 1] == -100`.
359
+ That is, the prefix portion of the sequence should not generate any loss. Loss should only be
360
+ generated by the target portion of the sequence.
361
+
362
+ Notes on `GPTNeoForCausalLM`:
363
+ To simplify the implementation, "global" and "local" attention layers are handled differently.
364
+ For "global" layers, we handle conversion as described above. For "local" layers, which use a
365
+ causal attention mask within a restricted local window, we do not alter the masking.
366
+
367
+ Notes on `forward` method conversion:
368
+ After conversion, the `forward` method will handle a new input, `bidirectional_mask`,
369
+ which should be a [batch_size, seq_length] byte tensor, where 1 indicates token positions
370
+ belonging to the prefix (prefix tokens can attend to one another bidirectionally), and
371
+ 0 indicates token positions belonging to the target.
372
+
373
+ The new `forward` method will incorporate `bidirectional_mask` (if supplied) into the existing
374
+ causal mask, call the original `forward` method, and (if the causal mask is a buffer) reset
375
+ the causal masks before returning the result.
376
+
377
+ Notes on `generate` method conversion:
378
+ After conversion, the `generate` method will have the same signature but will internally
379
+ convert all causal masks to be purely bidirectional, call the original `generate` method, and
380
+ (where appropriate) reset the causal masks before returning the result.
381
+
382
+ This works thanks to the logic of the HuggingFace `generate` API, which first encodes the token
383
+ "prompt" passed to `generate` (which is treated as the prefix) and then sequentially generates
384
+ each new token. Encodings are cached as generation happens, so all prefix tokens can attend to one
385
+ another (as expected in a Prefix LM) and generated tokens can only attend to prefix tokens and
386
+ previously-generated tokens (also as expected in a Prefix LM).
387
+
388
+ To preserve the API, the original methods are renamed to `_original_forward` and
389
+ `_original_generate`, and replaced with new `forward` and `generate` methods that wrap
390
+ them, respectively. Although implementation details vary by model class.
391
+ """
392
+ if isinstance(model, _SUPPORTED_GPT_MODELS):
393
+ return _convert_gpt_causal_lm_to_prefix_lm(model)
394
+ elif isinstance(model, BloomForCausalLM):
395
+ return _convert_bloom_causal_lm_to_prefix_lm(model)
396
+ elif isinstance(model, OPTForCausalLM):
397
+ return _convert_opt_causal_lm_to_prefix_lm(model)
398
+ else:
399
+ raise TypeError(f'Cannot convert model to Prefix LM. ' + f'Model does not belong to set of supported HF models:' + f'\n{_SUPPORTED_HF_MODELS}')
400
+
401
+ def add_bidirectional_mask_if_missing(batch: Dict[str, Any]):
402
+ """Attempts to add bidirectional_mask to batch if missing.
403
+
404
+ Raises:
405
+ KeyError if bidirectional_mask is missing and can't be inferred
406
+ """
407
+ if 'bidirectional_mask' not in batch:
408
+ if batch.get('mode', None) == 'icl_task':
409
+ batch['bidirectional_mask'] = batch['attention_mask'].clone()
410
+ for (i, continuation_indices) in enumerate(batch['continuation_indices']):
411
+ batch['bidirectional_mask'][i, continuation_indices] = 0
412
+ elif 'labels' in batch and 'attention_mask' in batch:
413
+ batch['bidirectional_mask'] = torch.logical_and(torch.eq(batch['attention_mask'], 1), torch.eq(batch['labels'], -100)).type_as(batch['attention_mask'])
414
+ else:
415
+ raise KeyError('No bidirectional_mask in batch and not sure how to construct one.')
meta_init_context.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from contextlib import contextmanager
2
+ import torch
3
+ import torch.nn as nn
4
+
5
+ @contextmanager
6
+ def init_empty_weights(include_buffers: bool=False):
7
+ """Meta initialization context manager.
8
+
9
+ A context manager under which models are initialized with all parameters
10
+ on the meta device, therefore creating an empty model. Useful when just
11
+ initializing the model would blow the available RAM.
12
+
13
+ Args:
14
+ include_buffers (`bool`, *optional*, defaults to `False`): Whether or
15
+ not to also put all buffers on the meta device while initializing.
16
+
17
+ Example:
18
+ ```python
19
+ import torch.nn as nn
20
+
21
+ # Initialize a model with 100 billions parameters in no time and without using any RAM.
22
+ with init_empty_weights():
23
+ tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
24
+ ```
25
+
26
+ <Tip warning={true}>
27
+
28
+ Any model created under this context manager has no weights. As such you can't do something like
29
+ `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].
30
+
31
+ </Tip>
32
+ """
33
+ with init_on_device(torch.device('meta'), include_buffers=include_buffers) as f:
34
+ yield f
35
+
36
+ @contextmanager
37
+ def init_on_device(device: torch.device, include_buffers: bool=False):
38
+ """Device initialization context manager.
39
+
40
+ A context manager under which models are initialized with all parameters
41
+ on the specified device.
42
+
43
+ Args:
44
+ device (`torch.device`): Device to initialize all parameters on.
45
+ include_buffers (`bool`, *optional*, defaults to `False`): Whether or
46
+ not to also put all buffers on the meta device while initializing.
47
+
48
+ Example:
49
+ ```python
50
+ import torch.nn as nn
51
+
52
+ with init_on_device(device=torch.device("cuda")):
53
+ tst = nn.Liner(100, 100) # on `cuda` device
54
+ ```
55
+ """
56
+ old_register_parameter = nn.Module.register_parameter
57
+ if include_buffers:
58
+ old_register_buffer = nn.Module.register_buffer
59
+
60
+ def register_empty_parameter(module, name, param):
61
+ old_register_parameter(module, name, param)
62
+ if param is not None:
63
+ param_cls = type(module._parameters[name])
64
+ kwargs = module._parameters[name].__dict__
65
+ module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
66
+
67
+ def register_empty_buffer(module, name, buffer):
68
+ old_register_buffer(module, name, buffer)
69
+ if buffer is not None:
70
+ module._buffers[name] = module._buffers[name].to(device)
71
+ if include_buffers:
72
+ tensor_constructors_to_patch = {torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ['empty', 'zeros', 'ones', 'full']}
73
+ else:
74
+ tensor_constructors_to_patch = {}
75
+
76
+ def patch_tensor_constructor(fn):
77
+
78
+ def wrapper(*args, **kwargs):
79
+ kwargs['device'] = device
80
+ return fn(*args, **kwargs)
81
+ return wrapper
82
+ try:
83
+ nn.Module.register_parameter = register_empty_parameter
84
+ if include_buffers:
85
+ nn.Module.register_buffer = register_empty_buffer
86
+ for torch_function_name in tensor_constructors_to_patch.keys():
87
+ setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
88
+ yield
89
+ finally:
90
+ nn.Module.register_parameter = old_register_parameter
91
+ if include_buffers:
92
+ nn.Module.register_buffer = old_register_buffer
93
+ for (torch_function_name, old_torch_function) in tensor_constructors_to_patch.items():
94
+ setattr(torch, torch_function_name, old_torch_function)
modeling_mpt.py ADDED
@@ -0,0 +1,290 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """A simple, flexible implementation of a GPT model.
2
+
3
+ Inspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
4
+ """
5
+ import math
6
+ import warnings
7
+ from typing import List, Optional, Tuple, Union
8
+ import torch
9
+ import torch.nn as nn
10
+ import torch.nn.functional as F
11
+ from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
12
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
13
+ from .attention import attn_bias_shape, build_attn_bias
14
+ from .blocks import MPTBlock
15
+ from .norm import NORM_CLASS_REGISTRY
16
+ from .configuration_mpt import MPTConfig
17
+ from .adapt_tokenizer import AutoTokenizerForMOD, adapt_tokenizer_for_denoising
18
+ from .hf_prefixlm_converter import add_bidirectional_mask_if_missing, convert_hf_causal_lm_to_prefix_lm
19
+ from .meta_init_context import init_empty_weights
20
+ from .param_init_fns import MODEL_INIT_REGISTRY, generic_param_init_fn_
21
+ Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
22
+
23
+ class MPTPreTrainedModel(PreTrainedModel):
24
+ config_class = MPTConfig
25
+ base_model_prefix = 'model'
26
+
27
+ class MPTModel(MPTPreTrainedModel):
28
+
29
+ def __init__(self, config: MPTConfig):
30
+ config._validate_config()
31
+ super().__init__(config)
32
+ self.attn_impl = config.attn_config['attn_impl']
33
+ self.prefix_lm = config.attn_config['prefix_lm']
34
+ self.attn_uses_sequence_id = config.attn_config['attn_uses_sequence_id']
35
+ self.alibi = config.attn_config['alibi']
36
+ self.alibi_bias_max = config.attn_config['alibi_bias_max']
37
+ if config.norm_type.lower() not in NORM_CLASS_REGISTRY.keys():
38
+ norm_options = ' | '.join(NORM_CLASS_REGISTRY.keys())
39
+ raise NotImplementedError(f'Requested norm type ({config.norm_type}) is not implemented within this repo (Options: {norm_options}).')
40
+ norm_class = NORM_CLASS_REGISTRY[config.norm_type.lower()]
41
+ self.embedding_fraction = config.embedding_fraction
42
+ self.wte = nn.Embedding(config.vocab_size, config.d_model, device=config.init_device)
43
+ if not self.alibi:
44
+ self.wpe = nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
45
+ self.emb_drop = nn.Dropout(config.emb_pdrop)
46
+ self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
47
+ self.norm_f = norm_class(config.d_model, device=config.init_device)
48
+ if config.init_device != 'meta':
49
+ self.apply(self.param_init_fn)
50
+ self.is_causal = not self.prefix_lm
51
+ self._attn_bias_initialized = False
52
+ self.attn_bias = None
53
+ self.attn_bias_shape = attn_bias_shape(self.attn_impl, config.n_heads, config.max_seq_len, self.alibi, prefix_lm=self.prefix_lm, causal=self.is_causal, use_sequence_id=self.attn_uses_sequence_id)
54
+ if config.no_bias:
55
+ for module in self.modules():
56
+ if hasattr(module, 'bias') and isinstance(module.bias, nn.Parameter):
57
+ if config.verbose:
58
+ warnings.warn(f'Removing bias ({module.bias}) from {module}.')
59
+ module.register_parameter('bias', None)
60
+ if config.verbose and config.verbose > 2:
61
+ print(self)
62
+ if 'verbose' not in self.config.init_config:
63
+ self.config.init_config['verbose'] = self.config.verbose
64
+ if self.config.init_config['verbose'] > 1:
65
+ init_fn_name = self.config.init_config['name']
66
+ warnings.warn(f'Using {init_fn_name} initialization.')
67
+
68
+ def get_input_embeddings(self):
69
+ return self.wte
70
+
71
+ def set_input_embeddings(self, value):
72
+ self.wte = value
73
+
74
+ @torch.no_grad()
75
+ def _attn_bias(self, device, dtype, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None):
76
+ if not self._attn_bias_initialized:
77
+ if self.attn_bias_shape:
78
+ self.attn_bias = torch.zeros(self.attn_bias_shape, device=device, dtype=dtype)
79
+ self.attn_bias = build_attn_bias(self.attn_impl, self.attn_bias, self.config.n_heads, self.config.max_seq_len, causal=self.is_causal, alibi=self.alibi, alibi_bias_max=self.alibi_bias_max)
80
+ self._attn_bias_initialized = True
81
+ if self.attn_impl == 'flash':
82
+ return (self.attn_bias, attention_mask)
83
+ if self.attn_bias is not None:
84
+ self.attn_bias = self.attn_bias.to(dtype=dtype, device=device)
85
+ attn_bias = self.attn_bias
86
+ if self.prefix_lm:
87
+ assert isinstance(attn_bias, torch.Tensor)
88
+ assert isinstance(prefix_mask, torch.Tensor)
89
+ attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)
90
+ if self.attn_uses_sequence_id and sequence_id is not None:
91
+ assert isinstance(attn_bias, torch.Tensor)
92
+ attn_bias = self._apply_sequence_id(attn_bias, sequence_id)
93
+ if attention_mask is not None:
94
+ s_k = attention_mask.shape[-1]
95
+ if attn_bias is None:
96
+ attn_bias = torch.zeros((1, 1, 1, s_k), device=device, dtype=dtype)
97
+ else:
98
+ attn_bias = attn_bias[:, :, :, -s_k:]
99
+ if prefix_mask is not None and attention_mask.shape != prefix_mask.shape:
100
+ raise ValueError(f'attention_mask shape={attention_mask.shape} ' + f'and prefix_mask shape={prefix_mask.shape} are not equal.')
101
+ min_val = torch.finfo(attn_bias.dtype).min
102
+ attn_bias = attn_bias.masked_fill(~attention_mask.view(-1, 1, 1, s_k), min_val)
103
+ return (attn_bias, None)
104
+
105
+ def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor):
106
+ (s_k, s_q) = attn_bias.shape[-2:]
107
+ if s_k != self.config.max_seq_len or s_q != self.config.max_seq_len:
108
+ raise ValueError('attn_bias does not match the expected shape. ' + f'The last two dimensions should both be {self.config.max_length} ' + f'but are {s_k} and {s_q}.')
109
+ seq_len = prefix_mask.shape[-1]
110
+ if seq_len > self.config.max_seq_len:
111
+ raise ValueError(f'prefix_mask sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
112
+ attn_bias = attn_bias[..., :seq_len, :seq_len]
113
+ causal = torch.tril(torch.ones((seq_len, seq_len), dtype=torch.bool, device=prefix_mask.device)).view(1, 1, seq_len, seq_len)
114
+ prefix = prefix_mask.view(-1, 1, 1, seq_len)
115
+ cannot_attend = ~torch.logical_or(causal, prefix.bool())
116
+ min_val = torch.finfo(attn_bias.dtype).min
117
+ attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
118
+ return attn_bias
119
+
120
+ def _apply_sequence_id(self, attn_bias: torch.Tensor, sequence_id: torch.LongTensor):
121
+ seq_len = sequence_id.shape[-1]
122
+ if seq_len > self.config.max_seq_len:
123
+ raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
124
+ attn_bias = attn_bias[..., :seq_len, :seq_len]
125
+ cannot_attend = torch.logical_not(torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))).unsqueeze(1)
126
+ min_val = torch.finfo(attn_bias.dtype).min
127
+ attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
128
+ return attn_bias
129
+
130
+ 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):
131
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
132
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
133
+ if attention_mask is not None:
134
+ attention_mask = attention_mask.bool()
135
+ if prefix_mask is not None:
136
+ prefix_mask = prefix_mask.bool()
137
+ if not return_dict:
138
+ raise NotImplementedError('return_dict False is not implemented yet for MPT')
139
+ if output_attentions:
140
+ raise NotImplementedError('output_attentions is not implemented yet for MPT')
141
+ if attention_mask is not None and attention_mask[:, 0].sum() != attention_mask.shape[0] and self.training:
142
+ raise NotImplementedError('MPT does not support training with left padding.')
143
+ if self.prefix_lm and prefix_mask is None:
144
+ raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
145
+ if self.training:
146
+ if self.attn_uses_sequence_id and sequence_id is None:
147
+ 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.')
148
+ elif self.attn_uses_sequence_id is False and sequence_id is not None:
149
+ 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.')
150
+ S = input_ids.size(1)
151
+ 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}'
152
+ tok_emb = self.wte(input_ids)
153
+ if self.alibi:
154
+ x = tok_emb
155
+ else:
156
+ past_position = 0
157
+ if past_key_values is not None:
158
+ if len(past_key_values) != self.config.n_layers:
159
+ 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}).')
160
+ past_position = past_key_values[0][0].size(1)
161
+ if S + past_position > self.config.max_seq_len:
162
+ raise ValueError(f'Cannot forward input with past sequence length {past_position} and current sequence length {S + 1}, this model only supports total sequence length <= {self.config.max_seq_len}.')
163
+ pos = torch.arange(past_position, S + past_position, dtype=torch.long, device=input_ids.device).unsqueeze(0)
164
+ if attention_mask is not None:
165
+ pos = torch.clamp(pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[:, past_position:], min=0)
166
+ pos_emb = self.wpe(pos)
167
+ x = tok_emb + pos_emb
168
+ if self.embedding_fraction == 1:
169
+ x = self.emb_drop(x)
170
+ else:
171
+ x_shrunk = x * self.embedding_fraction + x.detach() * (1 - self.embedding_fraction)
172
+ assert isinstance(self.emb_drop, nn.Module)
173
+ x = self.emb_drop(x_shrunk)
174
+ (attn_bias, attention_mask) = self._attn_bias(device=x.device, dtype=x.dtype, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id)
175
+ if use_cache and past_key_values is None:
176
+ past_key_values = [() for _ in range(self.config.n_layers)]
177
+ all_hidden_states = () if output_hidden_states else None
178
+ for (b_idx, block) in enumerate(self.blocks):
179
+ if output_hidden_states:
180
+ assert all_hidden_states is not None
181
+ all_hidden_states = all_hidden_states + (x,)
182
+ past_key_value = past_key_values[b_idx] if past_key_values is not None else None
183
+ (x, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=self.is_causal)
184
+ if past_key_values is not None:
185
+ past_key_values[b_idx] = past_key_value
186
+ x = self.norm_f(x)
187
+ return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values, hidden_states=all_hidden_states)
188
+
189
+ def param_init_fn(self, module):
190
+ init_fn_name = self.config.init_config['name']
191
+ MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
192
+
193
+ def fsdp_wrap_fn(self, module):
194
+ return isinstance(module, MPTBlock)
195
+
196
+ def activation_checkpointing_fn(self, module):
197
+ return isinstance(module, MPTBlock)
198
+
199
+ class MPTForCausalLM(MPTPreTrainedModel):
200
+
201
+ def __init__(self, config: MPTConfig):
202
+ super().__init__(config)
203
+ if not config.tie_word_embeddings:
204
+ raise ValueError('MPTForCausalLM only supports tied word embeddings')
205
+ self.transformer = MPTModel(config)
206
+ self.logit_scale = None
207
+ if config.logit_scale is not None:
208
+ logit_scale = config.logit_scale
209
+ if isinstance(logit_scale, str):
210
+ if logit_scale == 'inv_sqrt_d_model':
211
+ logit_scale = 1 / math.sqrt(config.d_model)
212
+ else:
213
+ raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
214
+ self.logit_scale = logit_scale
215
+
216
+ def get_input_embeddings(self):
217
+ return self.transformer.wte
218
+
219
+ def set_input_embeddings(self, value):
220
+ self.transformer.wte = value
221
+
222
+ def get_output_embeddings(self):
223
+ return self.transformer.wte
224
+
225
+ def set_output_embeddings(self, new_embeddings):
226
+ self.transformer.wte = new_embeddings
227
+
228
+ def set_decoder(self, decoder):
229
+ self.transformer = decoder
230
+
231
+ def get_decoder(self):
232
+ return self.transformer
233
+
234
+ 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):
235
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
236
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
237
+ 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)
238
+ logits = F.linear(outputs.last_hidden_state, self.transformer.wte.weight)
239
+ if self.logit_scale is not None:
240
+ if self.logit_scale == 0:
241
+ warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
242
+ logits *= self.logit_scale
243
+ loss = None
244
+ if labels is not None:
245
+ labels = torch.roll(labels, shifts=-1)
246
+ labels[:, -1] = -100
247
+ loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.to(logits.device).view(-1))
248
+ return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states)
249
+
250
+ def param_init_fn(self, module):
251
+ init_fn_name = self.config.init_config['name']
252
+ MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
253
+
254
+ def fsdp_wrap_fn(self, module):
255
+ return isinstance(module, MPTBlock)
256
+
257
+ def activation_checkpointing_fn(self, module):
258
+ return isinstance(module, MPTBlock)
259
+
260
+ def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
261
+ if inputs_embeds is not None:
262
+ raise NotImplementedError('inputs_embeds is not implemented for MPT yet')
263
+ attention_mask = kwargs['attention_mask'].bool()
264
+ if attention_mask[:, -1].sum() != attention_mask.shape[0]:
265
+ raise NotImplementedError('MPT does not support generation with right padding.')
266
+ if self.transformer.attn_uses_sequence_id and self.training:
267
+ sequence_id = torch.zeros_like(input_ids[:1])
268
+ else:
269
+ sequence_id = None
270
+ if past_key_values is not None:
271
+ input_ids = input_ids[:, -1].unsqueeze(-1)
272
+ if self.transformer.prefix_lm:
273
+ prefix_mask = torch.ones_like(attention_mask)
274
+ if kwargs.get('use_cache') == False:
275
+ raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
276
+ else:
277
+ prefix_mask = None
278
+ return {'input_ids': input_ids, 'attention_mask': attention_mask, 'prefix_mask': prefix_mask, 'sequence_id': sequence_id, 'past_key_values': past_key_values, 'use_cache': kwargs.get('use_cache', True)}
279
+
280
+ @staticmethod
281
+ def _reorder_cache(past_key_values, beam_idx):
282
+ """Used by HuggingFace generate when using beam search with kv-caching.
283
+
284
+ See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133
285
+ for an example in transformers.
286
+ """
287
+ reordered_past = []
288
+ for layer_past in past_key_values:
289
+ reordered_past += [tuple((past_state.index_select(0, beam_idx) for past_state in layer_past))]
290
+ return reordered_past
norm.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ def _cast_if_autocast_enabled(tensor):
4
+ if torch.is_autocast_enabled():
5
+ if tensor.device.type == 'cuda':
6
+ dtype = torch.get_autocast_gpu_dtype()
7
+ elif tensor.device.type == 'cpu':
8
+ dtype = torch.get_autocast_cpu_dtype()
9
+ else:
10
+ raise NotImplementedError()
11
+ return tensor.to(dtype=dtype)
12
+ return tensor
13
+
14
+ class LPLayerNorm(torch.nn.LayerNorm):
15
+
16
+ def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None):
17
+ super().__init__(normalized_shape=normalized_shape, eps=eps, elementwise_affine=elementwise_affine, device=device, dtype=dtype)
18
+
19
+ def forward(self, x):
20
+ module_device = x.device
21
+ downcast_x = _cast_if_autocast_enabled(x)
22
+ downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
23
+ downcast_bias = _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias
24
+ with torch.autocast(enabled=False, device_type=module_device.type):
25
+ return torch.nn.functional.layer_norm(downcast_x, self.normalized_shape, downcast_weight, downcast_bias, self.eps)
26
+
27
+ def rms_norm(x, weight=None, eps=1e-05):
28
+ output = x / torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
29
+ if weight is not None:
30
+ return output * weight
31
+ return output
32
+
33
+ class RMSNorm(torch.nn.Module):
34
+
35
+ def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
36
+ super().__init__()
37
+ self.eps = eps
38
+ if weight:
39
+ self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))
40
+ else:
41
+ self.register_parameter('weight', None)
42
+
43
+ def forward(self, x):
44
+ return rms_norm(x.float(), self.weight, self.eps).to(dtype=x.dtype)
45
+
46
+ class LPRMSNorm(RMSNorm):
47
+
48
+ def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
49
+ super().__init__(normalized_shape=normalized_shape, eps=eps, weight=weight, dtype=dtype, device=device)
50
+
51
+ def forward(self, x):
52
+ downcast_x = _cast_if_autocast_enabled(x)
53
+ downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
54
+ with torch.autocast(enabled=False, device_type=x.device.type):
55
+ return rms_norm(downcast_x, downcast_weight, self.eps).to(dtype=x.dtype)
56
+ NORM_CLASS_REGISTRY = {'layernorm': torch.nn.LayerNorm, 'low_precision_layernorm': LPLayerNorm, 'rmsnorm': RMSNorm, 'low_precision_rmsnorm': LPRMSNorm}
param_init_fns.py ADDED
@@ -0,0 +1,181 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import warnings
3
+ from collections.abc import Sequence
4
+ from functools import partial
5
+ from typing import Optional, Tuple, Union
6
+ import torch
7
+ from torch import nn
8
+ from .norm import NORM_CLASS_REGISTRY
9
+
10
+ def torch_default_param_init_fn_(module: nn.Module, verbose: int=0, **kwargs):
11
+ del kwargs
12
+ if verbose > 1:
13
+ warnings.warn(f"Initializing network using module's reset_parameters attribute")
14
+ if hasattr(module, 'reset_parameters'):
15
+ module.reset_parameters()
16
+
17
+ def fused_init_helper_(module: nn.Module, init_fn_):
18
+ _fused = getattr(module, '_fused', None)
19
+ if _fused is None:
20
+ raise RuntimeError(f'Internal logic error')
21
+ (dim, splits) = _fused
22
+ splits = (0, *splits, module.weight.size(dim))
23
+ for (s, e) in zip(splits[:-1], splits[1:]):
24
+ slice_indices = [slice(None)] * module.weight.ndim
25
+ slice_indices[dim] = slice(s, e)
26
+ init_fn_(module.weight[slice_indices])
27
+
28
+ def generic_param_init_fn_(module: nn.Module, init_fn_, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
29
+ del kwargs
30
+ if verbose > 1:
31
+ warnings.warn(f'If model has bias parameters they are initialized to 0.')
32
+ init_div_is_residual = init_div_is_residual
33
+ if init_div_is_residual is False:
34
+ div_is_residual = 1.0
35
+ elif init_div_is_residual is True:
36
+ div_is_residual = math.sqrt(2 * n_layers)
37
+ elif isinstance(init_div_is_residual, float) or isinstance(init_div_is_residual, int):
38
+ div_is_residual = init_div_is_residual
39
+ elif isinstance(init_div_is_residual, str) and init_div_is_residual.isnumeric():
40
+ div_is_residual = float(init_div_is_residual)
41
+ else:
42
+ div_is_residual = 1.0
43
+ raise ValueError(f'Expected init_div_is_residual to be boolean or numeric, got {init_div_is_residual}')
44
+ if init_div_is_residual is not False:
45
+ if verbose > 1:
46
+ warnings.warn(f'Initializing _is_residual layers then dividing them by {div_is_residual:.3f}. ' + f'Set `init_div_is_residual: false` in init config to disable this.')
47
+ if isinstance(module, nn.Linear):
48
+ if hasattr(module, '_fused'):
49
+ fused_init_helper_(module, init_fn_)
50
+ else:
51
+ init_fn_(module.weight)
52
+ if module.bias is not None:
53
+ torch.nn.init.zeros_(module.bias)
54
+ if init_div_is_residual is not False and getattr(module, '_is_residual', False):
55
+ with torch.no_grad():
56
+ module.weight.div_(div_is_residual)
57
+ elif isinstance(module, nn.Embedding):
58
+ if emb_init_std is not None:
59
+ std = emb_init_std
60
+ if std == 0:
61
+ warnings.warn(f'Embedding layer initialized to 0.')
62
+ emb_init_fn_ = partial(torch.nn.init.normal_, mean=0.0, std=std)
63
+ if verbose > 1:
64
+ warnings.warn(f'Embedding layer initialized using normal distribution with mean=0 and std={std!r}.')
65
+ elif emb_init_uniform_lim is not None:
66
+ lim = emb_init_uniform_lim
67
+ if isinstance(lim, Sequence):
68
+ if len(lim) > 2:
69
+ raise ValueError(f'Uniform init requires a min and a max limit. User input: {lim}.')
70
+ if lim[0] == lim[1]:
71
+ warnings.warn(f'Embedding layer initialized to {lim[0]}.')
72
+ else:
73
+ if lim == 0:
74
+ warnings.warn(f'Embedding layer initialized to 0.')
75
+ lim = [-lim, lim]
76
+ (a, b) = lim
77
+ emb_init_fn_ = partial(torch.nn.init.uniform_, a=a, b=b)
78
+ if verbose > 1:
79
+ warnings.warn(f'Embedding layer initialized using uniform distribution in range {lim}.')
80
+ else:
81
+ emb_init_fn_ = init_fn_
82
+ emb_init_fn_(module.weight)
83
+ elif isinstance(module, tuple(set(NORM_CLASS_REGISTRY.values()))):
84
+ if verbose > 1:
85
+ warnings.warn(f'Norm weights are set to 1. If norm layer has a bias it is initialized to 0.')
86
+ if hasattr(module, 'weight') and module.weight is not None:
87
+ torch.nn.init.ones_(module.weight)
88
+ if hasattr(module, 'bias') and module.bias is not None:
89
+ torch.nn.init.zeros_(module.bias)
90
+ elif isinstance(module, nn.MultiheadAttention):
91
+ if module._qkv_same_embed_dim:
92
+ assert module.in_proj_weight is not None
93
+ assert module.q_proj_weight is None and module.k_proj_weight is None and (module.v_proj_weight is None)
94
+ assert d_model is not None
95
+ _d = d_model
96
+ splits = (0, _d, 2 * _d, 3 * _d)
97
+ for (s, e) in zip(splits[:-1], splits[1:]):
98
+ init_fn_(module.in_proj_weight[s:e])
99
+ else:
100
+ assert module.q_proj_weight is not None and module.k_proj_weight is not None and (module.v_proj_weight is not None)
101
+ assert module.in_proj_weight is None
102
+ init_fn_(module.q_proj_weight)
103
+ init_fn_(module.k_proj_weight)
104
+ init_fn_(module.v_proj_weight)
105
+ if module.in_proj_bias is not None:
106
+ torch.nn.init.zeros_(module.in_proj_bias)
107
+ if module.bias_k is not None:
108
+ torch.nn.init.zeros_(module.bias_k)
109
+ if module.bias_v is not None:
110
+ torch.nn.init.zeros_(module.bias_v)
111
+ init_fn_(module.out_proj.weight)
112
+ if init_div_is_residual is not False and getattr(module.out_proj, '_is_residual', False):
113
+ with torch.no_grad():
114
+ module.out_proj.weight.div_(div_is_residual)
115
+ if module.out_proj.bias is not None:
116
+ torch.nn.init.zeros_(module.out_proj.bias)
117
+ else:
118
+ for _ in module.parameters(recurse=False):
119
+ raise NotImplementedError(f'{module.__class__.__name__} parameters are not initialized by param_init_fn.')
120
+
121
+ def _normal_init_(std, mean=0.0):
122
+ return partial(torch.nn.init.normal_, mean=mean, std=std)
123
+
124
+ def _normal_param_init_fn_(module: nn.Module, std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
125
+ del kwargs
126
+ init_fn_ = _normal_init_(std=std)
127
+ if verbose > 1:
128
+ warnings.warn(f'Using torch.nn.init.normal_ init fn mean=0.0, std={std}')
129
+ generic_param_init_fn_(module=module, init_fn_=init_fn_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
130
+
131
+ def baseline_param_init_fn_(module: nn.Module, init_std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
132
+ del kwargs
133
+ if init_std is None:
134
+ raise ValueError("You must set model.init_config['init_std'] to a float value to use the default initialization scheme.")
135
+ _normal_param_init_fn_(module=module, std=init_std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
136
+
137
+ def small_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
138
+ del kwargs
139
+ std = math.sqrt(2 / (5 * d_model))
140
+ _normal_param_init_fn_(module=module, std=std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
141
+
142
+ def neox_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
143
+ """From section 2.3.1 of GPT-NeoX-20B:
144
+
145
+ An Open-Source AutoregressiveLanguage Model — Black et. al. (2022)
146
+ see https://github.com/EleutherAI/gpt-neox/blob/9610391ab319403cef079b438edd016a2443af54/megatron/model/init_functions.py#L151
147
+ and https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/transformer.py
148
+ """
149
+ del kwargs
150
+ residual_div = n_layers / math.sqrt(10)
151
+ if verbose > 1:
152
+ warnings.warn(f'setting init_div_is_residual to {residual_div}')
153
+ small_param_init_fn_(module=module, d_model=d_model, n_layers=n_layers, init_div_is_residual=residual_div, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
154
+
155
+ def kaiming_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', verbose: int=0, **kwargs):
156
+ del kwargs
157
+ if verbose > 1:
158
+ warnings.warn(f'Using nn.init.kaiming_uniform_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}')
159
+ kaiming_uniform_ = partial(nn.init.kaiming_uniform_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
160
+ generic_param_init_fn_(module=module, init_fn_=kaiming_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
161
+
162
+ def kaiming_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', verbose: int=0, **kwargs):
163
+ del kwargs
164
+ if verbose > 1:
165
+ warnings.warn(f'Using nn.init.kaiming_normal_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}')
166
+ kaiming_normal_ = partial(torch.nn.init.kaiming_normal_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
167
+ generic_param_init_fn_(module=module, init_fn_=kaiming_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
168
+
169
+ def xavier_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, verbose: int=0, **kwargs):
170
+ del kwargs
171
+ xavier_uniform_ = partial(torch.nn.init.xavier_uniform_, gain=init_gain)
172
+ if verbose > 1:
173
+ warnings.warn(f'Using torch.nn.init.xavier_uniform_ init fn with parameters: ' + f'gain={init_gain}')
174
+ generic_param_init_fn_(module=module, init_fn_=xavier_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
175
+
176
+ def xavier_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, verbose: int=0, **kwargs):
177
+ xavier_normal_ = partial(torch.nn.init.xavier_normal_, gain=init_gain)
178
+ if verbose > 1:
179
+ warnings.warn(f'Using torch.nn.init.xavier_normal_ init fn with parameters: ' + f'gain={init_gain}')
180
+ generic_param_init_fn_(module=module, init_fn_=xavier_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
181
+ MODEL_INIT_REGISTRY = {'default_': torch_default_param_init_fn_, 'baseline_': baseline_param_init_fn_, 'kaiming_uniform_': kaiming_uniform_param_init_fn_, 'kaiming_normal_': kaiming_normal_param_init_fn_, 'neox_init_': neox_param_init_fn_, 'small_init_': small_param_init_fn_, 'xavier_uniform_': xavier_uniform_param_init_fn_, 'xavier_normal_': xavier_normal_param_init_fn_}
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:119a68ab92d3200f185ed35e1ae1758148af5487b90354f876652a555cc9ec6a
3
+ size 26594704201
special_tokens_map.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "additional_special_tokens": [
3
+ "<|im_start|>",
4
+ "<|im_end|>"
5
+ ],
6
+ "bos_token": "<|endoftext|>",
7
+ "eos_token": "<|endoftext|>",
8
+ "unk_token": "<|endoftext|>"
9
+ }
test_itex_warm.py ADDED
@@ -0,0 +1,128 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from intel_extension_for_transformers.backends.neural_engine.compile import compile
2
+ # graph = compile("/home/ubuntu/mengfeil/IR/llama-7b-bf16/")
3
+ # graph = compile("/home/ubuntu/mengfeil/IR/llama-7b-int8/")
4
+ # graph = compile("./llama-7b-hf-conv-itrex-bf16")
5
+ graph = compile("/home/ubuntu/neuralchat_server/frameworks.ai.nlp-toolkit.intel-nlp-toolkit/examples/huggingface/pytorch/text-generation/deployment/bf16ir")
6
+ from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaTokenizer, AutoModel, AutoConfig
7
+ import torch
8
+ import numpy as np
9
+ import torch.nn.functional as F
10
+ import time
11
+
12
+ model_name = "../mengfeil/llama-7b"
13
+ tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
14
+ config = AutoConfig.from_pretrained(model_name)
15
+
16
+ #prompt = "Once upon a time, there existed a little girl who liked to have adventures." + \
17
+ # " She wanted to go to places and meet new people, and have fun"
18
+ # prompt = "Once upon a time, there existed a little girl, who liked to have adventures. She wanted to go to places and meet new people, and have fun."
19
+ prompt = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n Human: tell me something about China.\n Assistant:"
20
+
21
+ print(prompt)
22
+ init_input_ids = tokenizer(prompt, return_tensors="pt").input_ids[0]
23
+
24
+
25
+ input_ids = init_input_ids.clone()
26
+
27
+ attention_mask = torch.ones(len(input_ids)+1)
28
+ attention_mask[0] = 0
29
+ position_ids = torch.arange(len(input_ids))
30
+ past_key_value = tuple([(torch.zeros([1,32,1,128]), torch.zeros([1,32,1,128])) for i in range(32)])
31
+
32
+ input_ids = input_ids.unsqueeze(0)
33
+ attention_mask = attention_mask.unsqueeze(0)
34
+ position_ids = position_ids.unsqueeze(0)
35
+ all_input_ids = input_ids.clone()
36
+
37
+
38
+ # input_ids_1 = input_ids.cpu().numpy().astype(np.int32)
39
+ # attention_mask_1 = attention_mask.cpu().numpy().astype(np.int32)
40
+ # past_k_v = [past_key_value[i][j].cpu().numpy() for i in range(32) for j in range(2)]
41
+
42
+ max_new_tokens = 32
43
+ temperature = 0.9
44
+
45
+
46
+ def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
47
+ """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
48
+ Args:
49
+ logits: logits distribution shape (vocabulary size)
50
+ top_k >0: keep only top k tokens with highest probability (top-k filtering).
51
+ top_p >0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
52
+ Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
53
+ """
54
+ assert logits.dim() == 1 # batch size 1 for now - could be updated for more but the code would be less clear
55
+ top_k = min(top_k, logits.size(-1)) # Safety check
56
+ if top_k > 0:
57
+ # Remove all tokens with a probability less than the last token of the top-k
58
+ indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
59
+ logits[indices_to_remove] = filter_value
60
+
61
+ if top_p > 0.0:
62
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True)
63
+ cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
64
+
65
+ # Remove tokens with cumulative probability above the threshold
66
+ sorted_indices_to_remove = cumulative_probs > top_p
67
+ # Shift the indices to the right to keep also the first token above the threshold
68
+ sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
69
+ sorted_indices_to_remove[..., 0] = 0
70
+
71
+ indices_to_remove = sorted_indices[sorted_indices_to_remove]
72
+ logits[indices_to_remove] = filter_value
73
+ return logits
74
+
75
+
76
+ # start
77
+ total_time = 0.0
78
+ num_iter = 10
79
+ num_warmup = 3
80
+
81
+ for i in range(num_iter):
82
+
83
+ input_ids_1 = input_ids.cpu().numpy().astype(np.int32)
84
+ attention_mask_1 = attention_mask.cpu().numpy().astype(np.int32)
85
+ past_k_v = [past_key_value[i][j].cpu().numpy() for i in range(32) for j in range(2)]
86
+ output_ids = list(init_input_ids)
87
+
88
+ tic = time.time()
89
+
90
+ for step in range(max_new_tokens):
91
+ a = time.time()
92
+ predictions = graph.inference([input_ids_1, attention_mask_1] + past_k_v)
93
+ # predictions = graph.inference([input_ids_1] + past_k_v + [attention_mask_1])
94
+ print(time.time() - a)
95
+
96
+ outs = []
97
+ for key in predictions:
98
+ outs.append(predictions[key])
99
+
100
+ logits = outs[0]
101
+ past_k_v = outs[1:]
102
+ logits = torch.from_numpy(logits)
103
+
104
+ next_token_logits = logits[:, -1, :]
105
+ probs = torch.nn.functional.softmax(next_token_logits, dim=-1)
106
+ token = int(torch.argmax(probs, dim=-1))
107
+
108
+ """
109
+ last_token_logits = logits[0][-1]
110
+ logits = logits[0, -1, :] / temperature
111
+ filtered_logits = top_k_top_p_filtering(logits, top_k=10, top_p=0.8)
112
+ probabilities = F.softmax(filtered_logits, dim=-1)
113
+ token = int(torch.multinomial(probabilities, 1))
114
+ """
115
+
116
+ output_ids.append(token)
117
+ input_ids_1 = torch.tensor([[token]])
118
+ attention_mask_1 = torch.cat([torch.from_numpy(attention_mask_1), torch.ones([1, 1])], dim=-1)
119
+ input_ids_1 = input_ids_1.cpu().numpy().astype(np.int32)
120
+ attention_mask_1 = attention_mask_1.cpu().numpy().astype(np.int32)
121
+
122
+ toc = time.time()
123
+ if i >= num_warmup:
124
+ total_time += (toc - tic)
125
+ # print(output_ids)
126
+ print(tokenizer.decode(output_ids, skip_special_tokens=True))
127
+
128
+ print("Inference latency: %.3f ms." % (total_time / (num_iter - num_warmup) * 1000))
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "bos_token": "<|endoftext|>",
4
+ "clean_up_tokenization_spaces": true,
5
+ "eos_token": "<|endoftext|>",
6
+ "model_max_length": 2048,
7
+ "tokenizer_class": "GPTNeoXTokenizer",
8
+ "unk_token": "<|endoftext|>"
9
+ }
trainer_state.json ADDED
@@ -0,0 +1,3616 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "best_metric": null,
3
+ "best_model_checkpoint": null,
4
+ "epoch": 0.3713417039942442,
5
+ "global_step": 6000,
6
+ "is_hyper_param_search": false,
7
+ "is_local_process_zero": true,
8
+ "is_world_process_zero": true,
9
+ "log_history": [
10
+ {
11
+ "epoch": 0.0,
12
+ "learning_rate": 9.278350515463919e-08,
13
+ "loss": 2.5166,
14
+ "step": 10
15
+ },
16
+ {
17
+ "epoch": 0.0,
18
+ "learning_rate": 1.9587628865979384e-07,
19
+ "loss": 1.9059,
20
+ "step": 20
21
+ },
22
+ {
23
+ "epoch": 0.0,
24
+ "learning_rate": 2.989690721649485e-07,
25
+ "loss": 1.7869,
26
+ "step": 30
27
+ },
28
+ {
29
+ "epoch": 0.0,
30
+ "learning_rate": 4.0206185567010316e-07,
31
+ "loss": 1.7863,
32
+ "step": 40
33
+ },
34
+ {
35
+ "epoch": 0.0,
36
+ "learning_rate": 5.051546391752578e-07,
37
+ "loss": 1.6946,
38
+ "step": 50
39
+ },
40
+ {
41
+ "epoch": 0.0,
42
+ "learning_rate": 6.082474226804124e-07,
43
+ "loss": 1.7419,
44
+ "step": 60
45
+ },
46
+ {
47
+ "epoch": 0.0,
48
+ "learning_rate": 7.11340206185567e-07,
49
+ "loss": 1.645,
50
+ "step": 70
51
+ },
52
+ {
53
+ "epoch": 0.0,
54
+ "learning_rate": 8.144329896907217e-07,
55
+ "loss": 1.6611,
56
+ "step": 80
57
+ },
58
+ {
59
+ "epoch": 0.01,
60
+ "learning_rate": 9.175257731958763e-07,
61
+ "loss": 1.5985,
62
+ "step": 90
63
+ },
64
+ {
65
+ "epoch": 0.01,
66
+ "learning_rate": 1.020618556701031e-06,
67
+ "loss": 1.6198,
68
+ "step": 100
69
+ },
70
+ {
71
+ "epoch": 0.01,
72
+ "learning_rate": 1.1237113402061856e-06,
73
+ "loss": 1.6259,
74
+ "step": 110
75
+ },
76
+ {
77
+ "epoch": 0.01,
78
+ "learning_rate": 1.2268041237113403e-06,
79
+ "loss": 1.5973,
80
+ "step": 120
81
+ },
82
+ {
83
+ "epoch": 0.01,
84
+ "learning_rate": 1.329896907216495e-06,
85
+ "loss": 1.5941,
86
+ "step": 130
87
+ },
88
+ {
89
+ "epoch": 0.01,
90
+ "learning_rate": 1.4329896907216496e-06,
91
+ "loss": 1.5597,
92
+ "step": 140
93
+ },
94
+ {
95
+ "epoch": 0.01,
96
+ "learning_rate": 1.5360824742268042e-06,
97
+ "loss": 1.5672,
98
+ "step": 150
99
+ },
100
+ {
101
+ "epoch": 0.01,
102
+ "learning_rate": 1.639175257731959e-06,
103
+ "loss": 1.5372,
104
+ "step": 160
105
+ },
106
+ {
107
+ "epoch": 0.01,
108
+ "learning_rate": 1.7422680412371134e-06,
109
+ "loss": 1.5715,
110
+ "step": 170
111
+ },
112
+ {
113
+ "epoch": 0.01,
114
+ "learning_rate": 1.8453608247422682e-06,
115
+ "loss": 1.5389,
116
+ "step": 180
117
+ },
118
+ {
119
+ "epoch": 0.01,
120
+ "learning_rate": 1.948453608247423e-06,
121
+ "loss": 1.5525,
122
+ "step": 190
123
+ },
124
+ {
125
+ "epoch": 0.01,
126
+ "learning_rate": 2.0515463917525773e-06,
127
+ "loss": 1.5871,
128
+ "step": 200
129
+ },
130
+ {
131
+ "epoch": 0.01,
132
+ "learning_rate": 2.1546391752577322e-06,
133
+ "loss": 1.5442,
134
+ "step": 210
135
+ },
136
+ {
137
+ "epoch": 0.01,
138
+ "learning_rate": 2.2577319587628867e-06,
139
+ "loss": 1.5335,
140
+ "step": 220
141
+ },
142
+ {
143
+ "epoch": 0.01,
144
+ "learning_rate": 2.3608247422680415e-06,
145
+ "loss": 1.5103,
146
+ "step": 230
147
+ },
148
+ {
149
+ "epoch": 0.01,
150
+ "learning_rate": 2.463917525773196e-06,
151
+ "loss": 1.5016,
152
+ "step": 240
153
+ },
154
+ {
155
+ "epoch": 0.02,
156
+ "learning_rate": 2.5670103092783504e-06,
157
+ "loss": 1.5101,
158
+ "step": 250
159
+ },
160
+ {
161
+ "epoch": 0.02,
162
+ "learning_rate": 2.6701030927835053e-06,
163
+ "loss": 1.5323,
164
+ "step": 260
165
+ },
166
+ {
167
+ "epoch": 0.02,
168
+ "learning_rate": 2.77319587628866e-06,
169
+ "loss": 1.4785,
170
+ "step": 270
171
+ },
172
+ {
173
+ "epoch": 0.02,
174
+ "learning_rate": 2.8762886597938146e-06,
175
+ "loss": 1.4132,
176
+ "step": 280
177
+ },
178
+ {
179
+ "epoch": 0.02,
180
+ "learning_rate": 2.979381443298969e-06,
181
+ "loss": 1.4733,
182
+ "step": 290
183
+ },
184
+ {
185
+ "epoch": 0.02,
186
+ "learning_rate": 3.082474226804124e-06,
187
+ "loss": 1.4258,
188
+ "step": 300
189
+ },
190
+ {
191
+ "epoch": 0.02,
192
+ "learning_rate": 3.1855670103092784e-06,
193
+ "loss": 1.4557,
194
+ "step": 310
195
+ },
196
+ {
197
+ "epoch": 0.02,
198
+ "learning_rate": 3.2886597938144333e-06,
199
+ "loss": 1.44,
200
+ "step": 320
201
+ },
202
+ {
203
+ "epoch": 0.02,
204
+ "learning_rate": 3.391752577319588e-06,
205
+ "loss": 1.4348,
206
+ "step": 330
207
+ },
208
+ {
209
+ "epoch": 0.02,
210
+ "learning_rate": 3.494845360824742e-06,
211
+ "loss": 1.406,
212
+ "step": 340
213
+ },
214
+ {
215
+ "epoch": 0.02,
216
+ "learning_rate": 3.597938144329897e-06,
217
+ "loss": 1.4239,
218
+ "step": 350
219
+ },
220
+ {
221
+ "epoch": 0.02,
222
+ "learning_rate": 3.701030927835052e-06,
223
+ "loss": 1.4185,
224
+ "step": 360
225
+ },
226
+ {
227
+ "epoch": 0.02,
228
+ "learning_rate": 3.8041237113402064e-06,
229
+ "loss": 1.3954,
230
+ "step": 370
231
+ },
232
+ {
233
+ "epoch": 0.02,
234
+ "learning_rate": 3.907216494845361e-06,
235
+ "loss": 1.3759,
236
+ "step": 380
237
+ },
238
+ {
239
+ "epoch": 0.02,
240
+ "learning_rate": 4.010309278350516e-06,
241
+ "loss": 1.3231,
242
+ "step": 390
243
+ },
244
+ {
245
+ "epoch": 0.02,
246
+ "learning_rate": 4.11340206185567e-06,
247
+ "loss": 1.3059,
248
+ "step": 400
249
+ },
250
+ {
251
+ "epoch": 0.03,
252
+ "learning_rate": 4.216494845360825e-06,
253
+ "loss": 1.3258,
254
+ "step": 410
255
+ },
256
+ {
257
+ "epoch": 0.03,
258
+ "learning_rate": 4.31958762886598e-06,
259
+ "loss": 1.3722,
260
+ "step": 420
261
+ },
262
+ {
263
+ "epoch": 0.03,
264
+ "learning_rate": 4.422680412371134e-06,
265
+ "loss": 1.2917,
266
+ "step": 430
267
+ },
268
+ {
269
+ "epoch": 0.03,
270
+ "learning_rate": 4.525773195876289e-06,
271
+ "loss": 1.2893,
272
+ "step": 440
273
+ },
274
+ {
275
+ "epoch": 0.03,
276
+ "learning_rate": 4.628865979381444e-06,
277
+ "loss": 1.2865,
278
+ "step": 450
279
+ },
280
+ {
281
+ "epoch": 0.03,
282
+ "learning_rate": 4.731958762886599e-06,
283
+ "loss": 1.28,
284
+ "step": 460
285
+ },
286
+ {
287
+ "epoch": 0.03,
288
+ "learning_rate": 4.835051546391753e-06,
289
+ "loss": 1.299,
290
+ "step": 470
291
+ },
292
+ {
293
+ "epoch": 0.03,
294
+ "learning_rate": 4.9381443298969075e-06,
295
+ "loss": 1.2949,
296
+ "step": 480
297
+ },
298
+ {
299
+ "epoch": 0.03,
300
+ "learning_rate": 5.041237113402062e-06,
301
+ "loss": 1.3078,
302
+ "step": 490
303
+ },
304
+ {
305
+ "epoch": 0.03,
306
+ "learning_rate": 5.144329896907216e-06,
307
+ "loss": 1.2242,
308
+ "step": 500
309
+ },
310
+ {
311
+ "epoch": 0.03,
312
+ "learning_rate": 5.247422680412372e-06,
313
+ "loss": 1.2714,
314
+ "step": 510
315
+ },
316
+ {
317
+ "epoch": 0.03,
318
+ "learning_rate": 5.350515463917526e-06,
319
+ "loss": 1.2482,
320
+ "step": 520
321
+ },
322
+ {
323
+ "epoch": 0.03,
324
+ "learning_rate": 5.45360824742268e-06,
325
+ "loss": 1.2703,
326
+ "step": 530
327
+ },
328
+ {
329
+ "epoch": 0.03,
330
+ "learning_rate": 5.556701030927836e-06,
331
+ "loss": 1.2578,
332
+ "step": 540
333
+ },
334
+ {
335
+ "epoch": 0.03,
336
+ "learning_rate": 5.65979381443299e-06,
337
+ "loss": 1.2989,
338
+ "step": 550
339
+ },
340
+ {
341
+ "epoch": 0.03,
342
+ "learning_rate": 5.762886597938144e-06,
343
+ "loss": 1.2853,
344
+ "step": 560
345
+ },
346
+ {
347
+ "epoch": 0.04,
348
+ "learning_rate": 5.8659793814433e-06,
349
+ "loss": 1.3022,
350
+ "step": 570
351
+ },
352
+ {
353
+ "epoch": 0.04,
354
+ "learning_rate": 5.969072164948454e-06,
355
+ "loss": 1.2871,
356
+ "step": 580
357
+ },
358
+ {
359
+ "epoch": 0.04,
360
+ "learning_rate": 6.0721649484536086e-06,
361
+ "loss": 1.2679,
362
+ "step": 590
363
+ },
364
+ {
365
+ "epoch": 0.04,
366
+ "learning_rate": 6.1752577319587634e-06,
367
+ "loss": 1.2732,
368
+ "step": 600
369
+ },
370
+ {
371
+ "epoch": 0.04,
372
+ "learning_rate": 6.278350515463918e-06,
373
+ "loss": 1.2641,
374
+ "step": 610
375
+ },
376
+ {
377
+ "epoch": 0.04,
378
+ "learning_rate": 6.381443298969072e-06,
379
+ "loss": 1.2919,
380
+ "step": 620
381
+ },
382
+ {
383
+ "epoch": 0.04,
384
+ "learning_rate": 6.484536082474227e-06,
385
+ "loss": 1.2594,
386
+ "step": 630
387
+ },
388
+ {
389
+ "epoch": 0.04,
390
+ "learning_rate": 6.587628865979382e-06,
391
+ "loss": 1.2592,
392
+ "step": 640
393
+ },
394
+ {
395
+ "epoch": 0.04,
396
+ "learning_rate": 6.690721649484536e-06,
397
+ "loss": 1.2652,
398
+ "step": 650
399
+ },
400
+ {
401
+ "epoch": 0.04,
402
+ "learning_rate": 6.793814432989692e-06,
403
+ "loss": 1.3274,
404
+ "step": 660
405
+ },
406
+ {
407
+ "epoch": 0.04,
408
+ "learning_rate": 6.896907216494846e-06,
409
+ "loss": 1.2155,
410
+ "step": 670
411
+ },
412
+ {
413
+ "epoch": 0.04,
414
+ "learning_rate": 7e-06,
415
+ "loss": 1.2837,
416
+ "step": 680
417
+ },
418
+ {
419
+ "epoch": 0.04,
420
+ "learning_rate": 7.103092783505156e-06,
421
+ "loss": 1.265,
422
+ "step": 690
423
+ },
424
+ {
425
+ "epoch": 0.04,
426
+ "learning_rate": 7.20618556701031e-06,
427
+ "loss": 1.2427,
428
+ "step": 700
429
+ },
430
+ {
431
+ "epoch": 0.04,
432
+ "learning_rate": 7.309278350515464e-06,
433
+ "loss": 1.2688,
434
+ "step": 710
435
+ },
436
+ {
437
+ "epoch": 0.04,
438
+ "learning_rate": 7.412371134020619e-06,
439
+ "loss": 1.3071,
440
+ "step": 720
441
+ },
442
+ {
443
+ "epoch": 0.05,
444
+ "learning_rate": 7.515463917525773e-06,
445
+ "loss": 1.2346,
446
+ "step": 730
447
+ },
448
+ {
449
+ "epoch": 0.05,
450
+ "learning_rate": 7.618556701030928e-06,
451
+ "loss": 1.2246,
452
+ "step": 740
453
+ },
454
+ {
455
+ "epoch": 0.05,
456
+ "learning_rate": 7.721649484536083e-06,
457
+ "loss": 1.2604,
458
+ "step": 750
459
+ },
460
+ {
461
+ "epoch": 0.05,
462
+ "learning_rate": 7.824742268041238e-06,
463
+ "loss": 1.2589,
464
+ "step": 760
465
+ },
466
+ {
467
+ "epoch": 0.05,
468
+ "learning_rate": 7.927835051546391e-06,
469
+ "loss": 1.2512,
470
+ "step": 770
471
+ },
472
+ {
473
+ "epoch": 0.05,
474
+ "learning_rate": 8.030927835051548e-06,
475
+ "loss": 1.2229,
476
+ "step": 780
477
+ },
478
+ {
479
+ "epoch": 0.05,
480
+ "learning_rate": 8.134020618556701e-06,
481
+ "loss": 1.2326,
482
+ "step": 790
483
+ },
484
+ {
485
+ "epoch": 0.05,
486
+ "learning_rate": 8.237113402061856e-06,
487
+ "loss": 1.3097,
488
+ "step": 800
489
+ },
490
+ {
491
+ "epoch": 0.05,
492
+ "learning_rate": 8.34020618556701e-06,
493
+ "loss": 1.2358,
494
+ "step": 810
495
+ },
496
+ {
497
+ "epoch": 0.05,
498
+ "learning_rate": 8.443298969072166e-06,
499
+ "loss": 1.2746,
500
+ "step": 820
501
+ },
502
+ {
503
+ "epoch": 0.05,
504
+ "learning_rate": 8.54639175257732e-06,
505
+ "loss": 1.3063,
506
+ "step": 830
507
+ },
508
+ {
509
+ "epoch": 0.05,
510
+ "learning_rate": 8.649484536082475e-06,
511
+ "loss": 1.2702,
512
+ "step": 840
513
+ },
514
+ {
515
+ "epoch": 0.05,
516
+ "learning_rate": 8.75257731958763e-06,
517
+ "loss": 1.2421,
518
+ "step": 850
519
+ },
520
+ {
521
+ "epoch": 0.05,
522
+ "learning_rate": 8.855670103092783e-06,
523
+ "loss": 1.3047,
524
+ "step": 860
525
+ },
526
+ {
527
+ "epoch": 0.05,
528
+ "learning_rate": 8.95876288659794e-06,
529
+ "loss": 1.2452,
530
+ "step": 870
531
+ },
532
+ {
533
+ "epoch": 0.05,
534
+ "learning_rate": 9.061855670103093e-06,
535
+ "loss": 1.2753,
536
+ "step": 880
537
+ },
538
+ {
539
+ "epoch": 0.06,
540
+ "learning_rate": 9.164948453608248e-06,
541
+ "loss": 1.2532,
542
+ "step": 890
543
+ },
544
+ {
545
+ "epoch": 0.06,
546
+ "learning_rate": 9.268041237113403e-06,
547
+ "loss": 1.2379,
548
+ "step": 900
549
+ },
550
+ {
551
+ "epoch": 0.06,
552
+ "learning_rate": 9.371134020618558e-06,
553
+ "loss": 1.2891,
554
+ "step": 910
555
+ },
556
+ {
557
+ "epoch": 0.06,
558
+ "learning_rate": 9.474226804123711e-06,
559
+ "loss": 1.2773,
560
+ "step": 920
561
+ },
562
+ {
563
+ "epoch": 0.06,
564
+ "learning_rate": 9.577319587628868e-06,
565
+ "loss": 1.2753,
566
+ "step": 930
567
+ },
568
+ {
569
+ "epoch": 0.06,
570
+ "learning_rate": 9.68041237113402e-06,
571
+ "loss": 1.1976,
572
+ "step": 940
573
+ },
574
+ {
575
+ "epoch": 0.06,
576
+ "learning_rate": 9.783505154639176e-06,
577
+ "loss": 1.2721,
578
+ "step": 950
579
+ },
580
+ {
581
+ "epoch": 0.06,
582
+ "learning_rate": 9.88659793814433e-06,
583
+ "loss": 1.3058,
584
+ "step": 960
585
+ },
586
+ {
587
+ "epoch": 0.06,
588
+ "learning_rate": 9.989690721649485e-06,
589
+ "loss": 1.2383,
590
+ "step": 970
591
+ },
592
+ {
593
+ "epoch": 0.06,
594
+ "learning_rate": 9.998105303046253e-06,
595
+ "loss": 1.2356,
596
+ "step": 980
597
+ },
598
+ {
599
+ "epoch": 0.06,
600
+ "learning_rate": 9.996000084208755e-06,
601
+ "loss": 1.2639,
602
+ "step": 990
603
+ },
604
+ {
605
+ "epoch": 0.06,
606
+ "learning_rate": 9.993894865371256e-06,
607
+ "loss": 1.2689,
608
+ "step": 1000
609
+ },
610
+ {
611
+ "epoch": 0.06,
612
+ "learning_rate": 9.991789646533758e-06,
613
+ "loss": 1.2764,
614
+ "step": 1010
615
+ },
616
+ {
617
+ "epoch": 0.06,
618
+ "learning_rate": 9.989684427696261e-06,
619
+ "loss": 1.2602,
620
+ "step": 1020
621
+ },
622
+ {
623
+ "epoch": 0.06,
624
+ "learning_rate": 9.987579208858762e-06,
625
+ "loss": 1.3034,
626
+ "step": 1030
627
+ },
628
+ {
629
+ "epoch": 0.06,
630
+ "learning_rate": 9.985473990021264e-06,
631
+ "loss": 1.2318,
632
+ "step": 1040
633
+ },
634
+ {
635
+ "epoch": 0.06,
636
+ "learning_rate": 9.983368771183765e-06,
637
+ "loss": 1.252,
638
+ "step": 1050
639
+ },
640
+ {
641
+ "epoch": 0.07,
642
+ "learning_rate": 9.981263552346267e-06,
643
+ "loss": 1.3111,
644
+ "step": 1060
645
+ },
646
+ {
647
+ "epoch": 0.07,
648
+ "learning_rate": 9.97915833350877e-06,
649
+ "loss": 1.255,
650
+ "step": 1070
651
+ },
652
+ {
653
+ "epoch": 0.07,
654
+ "learning_rate": 9.97705311467127e-06,
655
+ "loss": 1.2736,
656
+ "step": 1080
657
+ },
658
+ {
659
+ "epoch": 0.07,
660
+ "learning_rate": 9.974947895833773e-06,
661
+ "loss": 1.3128,
662
+ "step": 1090
663
+ },
664
+ {
665
+ "epoch": 0.07,
666
+ "learning_rate": 9.972842676996276e-06,
667
+ "loss": 1.2349,
668
+ "step": 1100
669
+ },
670
+ {
671
+ "epoch": 0.07,
672
+ "learning_rate": 9.970737458158777e-06,
673
+ "loss": 1.2797,
674
+ "step": 1110
675
+ },
676
+ {
677
+ "epoch": 0.07,
678
+ "learning_rate": 9.968632239321279e-06,
679
+ "loss": 1.2698,
680
+ "step": 1120
681
+ },
682
+ {
683
+ "epoch": 0.07,
684
+ "learning_rate": 9.96652702048378e-06,
685
+ "loss": 1.2714,
686
+ "step": 1130
687
+ },
688
+ {
689
+ "epoch": 0.07,
690
+ "learning_rate": 9.964421801646282e-06,
691
+ "loss": 1.2965,
692
+ "step": 1140
693
+ },
694
+ {
695
+ "epoch": 0.07,
696
+ "learning_rate": 9.962316582808785e-06,
697
+ "loss": 1.2764,
698
+ "step": 1150
699
+ },
700
+ {
701
+ "epoch": 0.07,
702
+ "learning_rate": 9.960211363971286e-06,
703
+ "loss": 1.2294,
704
+ "step": 1160
705
+ },
706
+ {
707
+ "epoch": 0.07,
708
+ "learning_rate": 9.958106145133788e-06,
709
+ "loss": 1.2623,
710
+ "step": 1170
711
+ },
712
+ {
713
+ "epoch": 0.07,
714
+ "learning_rate": 9.956000926296289e-06,
715
+ "loss": 1.2814,
716
+ "step": 1180
717
+ },
718
+ {
719
+ "epoch": 0.07,
720
+ "learning_rate": 9.953895707458791e-06,
721
+ "loss": 1.2793,
722
+ "step": 1190
723
+ },
724
+ {
725
+ "epoch": 0.07,
726
+ "learning_rate": 9.951790488621294e-06,
727
+ "loss": 1.2749,
728
+ "step": 1200
729
+ },
730
+ {
731
+ "epoch": 0.07,
732
+ "learning_rate": 9.949685269783795e-06,
733
+ "loss": 1.2689,
734
+ "step": 1210
735
+ },
736
+ {
737
+ "epoch": 0.08,
738
+ "learning_rate": 9.947580050946297e-06,
739
+ "loss": 1.2913,
740
+ "step": 1220
741
+ },
742
+ {
743
+ "epoch": 0.08,
744
+ "learning_rate": 9.9454748321088e-06,
745
+ "loss": 1.2386,
746
+ "step": 1230
747
+ },
748
+ {
749
+ "epoch": 0.08,
750
+ "learning_rate": 9.9433696132713e-06,
751
+ "loss": 1.2519,
752
+ "step": 1240
753
+ },
754
+ {
755
+ "epoch": 0.08,
756
+ "learning_rate": 9.941264394433803e-06,
757
+ "loss": 1.2893,
758
+ "step": 1250
759
+ },
760
+ {
761
+ "epoch": 0.08,
762
+ "learning_rate": 9.939159175596304e-06,
763
+ "loss": 1.2398,
764
+ "step": 1260
765
+ },
766
+ {
767
+ "epoch": 0.08,
768
+ "learning_rate": 9.937053956758805e-06,
769
+ "loss": 1.2836,
770
+ "step": 1270
771
+ },
772
+ {
773
+ "epoch": 0.08,
774
+ "learning_rate": 9.934948737921307e-06,
775
+ "loss": 1.2326,
776
+ "step": 1280
777
+ },
778
+ {
779
+ "epoch": 0.08,
780
+ "learning_rate": 9.93284351908381e-06,
781
+ "loss": 1.2768,
782
+ "step": 1290
783
+ },
784
+ {
785
+ "epoch": 0.08,
786
+ "learning_rate": 9.93073830024631e-06,
787
+ "loss": 1.2584,
788
+ "step": 1300
789
+ },
790
+ {
791
+ "epoch": 0.08,
792
+ "learning_rate": 9.928633081408813e-06,
793
+ "loss": 1.2738,
794
+ "step": 1310
795
+ },
796
+ {
797
+ "epoch": 0.08,
798
+ "learning_rate": 9.926527862571314e-06,
799
+ "loss": 1.2329,
800
+ "step": 1320
801
+ },
802
+ {
803
+ "epoch": 0.08,
804
+ "learning_rate": 9.924422643733816e-06,
805
+ "loss": 1.2623,
806
+ "step": 1330
807
+ },
808
+ {
809
+ "epoch": 0.08,
810
+ "learning_rate": 9.922317424896319e-06,
811
+ "loss": 1.2839,
812
+ "step": 1340
813
+ },
814
+ {
815
+ "epoch": 0.08,
816
+ "learning_rate": 9.92021220605882e-06,
817
+ "loss": 1.2806,
818
+ "step": 1350
819
+ },
820
+ {
821
+ "epoch": 0.08,
822
+ "learning_rate": 9.918106987221322e-06,
823
+ "loss": 1.232,
824
+ "step": 1360
825
+ },
826
+ {
827
+ "epoch": 0.08,
828
+ "learning_rate": 9.916001768383825e-06,
829
+ "loss": 1.256,
830
+ "step": 1370
831
+ },
832
+ {
833
+ "epoch": 0.09,
834
+ "learning_rate": 9.913896549546325e-06,
835
+ "loss": 1.225,
836
+ "step": 1380
837
+ },
838
+ {
839
+ "epoch": 0.09,
840
+ "learning_rate": 9.911791330708828e-06,
841
+ "loss": 1.2247,
842
+ "step": 1390
843
+ },
844
+ {
845
+ "epoch": 0.09,
846
+ "learning_rate": 9.909686111871329e-06,
847
+ "loss": 1.2422,
848
+ "step": 1400
849
+ },
850
+ {
851
+ "epoch": 0.09,
852
+ "learning_rate": 9.907580893033831e-06,
853
+ "loss": 1.2499,
854
+ "step": 1410
855
+ },
856
+ {
857
+ "epoch": 0.09,
858
+ "learning_rate": 9.905475674196334e-06,
859
+ "loss": 1.2569,
860
+ "step": 1420
861
+ },
862
+ {
863
+ "epoch": 0.09,
864
+ "learning_rate": 9.903370455358834e-06,
865
+ "loss": 1.2771,
866
+ "step": 1430
867
+ },
868
+ {
869
+ "epoch": 0.09,
870
+ "learning_rate": 9.901265236521337e-06,
871
+ "loss": 1.2789,
872
+ "step": 1440
873
+ },
874
+ {
875
+ "epoch": 0.09,
876
+ "learning_rate": 9.89916001768384e-06,
877
+ "loss": 1.2213,
878
+ "step": 1450
879
+ },
880
+ {
881
+ "epoch": 0.09,
882
+ "learning_rate": 9.89705479884634e-06,
883
+ "loss": 1.2393,
884
+ "step": 1460
885
+ },
886
+ {
887
+ "epoch": 0.09,
888
+ "learning_rate": 9.894949580008843e-06,
889
+ "loss": 1.257,
890
+ "step": 1470
891
+ },
892
+ {
893
+ "epoch": 0.09,
894
+ "learning_rate": 9.892844361171344e-06,
895
+ "loss": 1.2572,
896
+ "step": 1480
897
+ },
898
+ {
899
+ "epoch": 0.09,
900
+ "learning_rate": 9.890739142333846e-06,
901
+ "loss": 1.2503,
902
+ "step": 1490
903
+ },
904
+ {
905
+ "epoch": 0.09,
906
+ "learning_rate": 9.888633923496349e-06,
907
+ "loss": 1.2404,
908
+ "step": 1500
909
+ },
910
+ {
911
+ "epoch": 0.09,
912
+ "learning_rate": 9.88652870465885e-06,
913
+ "loss": 1.2847,
914
+ "step": 1510
915
+ },
916
+ {
917
+ "epoch": 0.09,
918
+ "learning_rate": 9.884423485821352e-06,
919
+ "loss": 1.2551,
920
+ "step": 1520
921
+ },
922
+ {
923
+ "epoch": 0.09,
924
+ "learning_rate": 9.882318266983854e-06,
925
+ "loss": 1.2741,
926
+ "step": 1530
927
+ },
928
+ {
929
+ "epoch": 0.1,
930
+ "learning_rate": 9.880213048146355e-06,
931
+ "loss": 1.2735,
932
+ "step": 1540
933
+ },
934
+ {
935
+ "epoch": 0.1,
936
+ "learning_rate": 9.878107829308858e-06,
937
+ "loss": 1.273,
938
+ "step": 1550
939
+ },
940
+ {
941
+ "epoch": 0.1,
942
+ "learning_rate": 9.876002610471358e-06,
943
+ "loss": 1.2357,
944
+ "step": 1560
945
+ },
946
+ {
947
+ "epoch": 0.1,
948
+ "learning_rate": 9.873897391633861e-06,
949
+ "loss": 1.1864,
950
+ "step": 1570
951
+ },
952
+ {
953
+ "epoch": 0.1,
954
+ "learning_rate": 9.871792172796363e-06,
955
+ "loss": 1.277,
956
+ "step": 1580
957
+ },
958
+ {
959
+ "epoch": 0.1,
960
+ "learning_rate": 9.869686953958864e-06,
961
+ "loss": 1.2375,
962
+ "step": 1590
963
+ },
964
+ {
965
+ "epoch": 0.1,
966
+ "learning_rate": 9.867581735121367e-06,
967
+ "loss": 1.2754,
968
+ "step": 1600
969
+ },
970
+ {
971
+ "epoch": 0.1,
972
+ "learning_rate": 9.865476516283868e-06,
973
+ "loss": 1.2051,
974
+ "step": 1610
975
+ },
976
+ {
977
+ "epoch": 0.1,
978
+ "learning_rate": 9.86337129744637e-06,
979
+ "loss": 1.2579,
980
+ "step": 1620
981
+ },
982
+ {
983
+ "epoch": 0.1,
984
+ "learning_rate": 9.861266078608873e-06,
985
+ "loss": 1.2216,
986
+ "step": 1630
987
+ },
988
+ {
989
+ "epoch": 0.1,
990
+ "learning_rate": 9.859160859771373e-06,
991
+ "loss": 1.2529,
992
+ "step": 1640
993
+ },
994
+ {
995
+ "epoch": 0.1,
996
+ "learning_rate": 9.857055640933876e-06,
997
+ "loss": 1.207,
998
+ "step": 1650
999
+ },
1000
+ {
1001
+ "epoch": 0.1,
1002
+ "learning_rate": 9.854950422096378e-06,
1003
+ "loss": 1.2275,
1004
+ "step": 1660
1005
+ },
1006
+ {
1007
+ "epoch": 0.1,
1008
+ "learning_rate": 9.85284520325888e-06,
1009
+ "loss": 1.2769,
1010
+ "step": 1670
1011
+ },
1012
+ {
1013
+ "epoch": 0.1,
1014
+ "learning_rate": 9.850739984421382e-06,
1015
+ "loss": 1.2165,
1016
+ "step": 1680
1017
+ },
1018
+ {
1019
+ "epoch": 0.1,
1020
+ "learning_rate": 9.848634765583883e-06,
1021
+ "loss": 1.2903,
1022
+ "step": 1690
1023
+ },
1024
+ {
1025
+ "epoch": 0.11,
1026
+ "learning_rate": 9.846529546746385e-06,
1027
+ "loss": 1.2548,
1028
+ "step": 1700
1029
+ },
1030
+ {
1031
+ "epoch": 0.11,
1032
+ "learning_rate": 9.844424327908888e-06,
1033
+ "loss": 1.2652,
1034
+ "step": 1710
1035
+ },
1036
+ {
1037
+ "epoch": 0.11,
1038
+ "learning_rate": 9.842319109071388e-06,
1039
+ "loss": 1.2718,
1040
+ "step": 1720
1041
+ },
1042
+ {
1043
+ "epoch": 0.11,
1044
+ "learning_rate": 9.84021389023389e-06,
1045
+ "loss": 1.269,
1046
+ "step": 1730
1047
+ },
1048
+ {
1049
+ "epoch": 0.11,
1050
+ "learning_rate": 9.838108671396393e-06,
1051
+ "loss": 1.2362,
1052
+ "step": 1740
1053
+ },
1054
+ {
1055
+ "epoch": 0.11,
1056
+ "learning_rate": 9.836003452558894e-06,
1057
+ "loss": 1.205,
1058
+ "step": 1750
1059
+ },
1060
+ {
1061
+ "epoch": 0.11,
1062
+ "learning_rate": 9.833898233721397e-06,
1063
+ "loss": 1.2649,
1064
+ "step": 1760
1065
+ },
1066
+ {
1067
+ "epoch": 0.11,
1068
+ "learning_rate": 9.831793014883897e-06,
1069
+ "loss": 1.2517,
1070
+ "step": 1770
1071
+ },
1072
+ {
1073
+ "epoch": 0.11,
1074
+ "learning_rate": 9.8296877960464e-06,
1075
+ "loss": 1.2015,
1076
+ "step": 1780
1077
+ },
1078
+ {
1079
+ "epoch": 0.11,
1080
+ "learning_rate": 9.827582577208902e-06,
1081
+ "loss": 1.2382,
1082
+ "step": 1790
1083
+ },
1084
+ {
1085
+ "epoch": 0.11,
1086
+ "learning_rate": 9.825477358371403e-06,
1087
+ "loss": 1.2236,
1088
+ "step": 1800
1089
+ },
1090
+ {
1091
+ "epoch": 0.11,
1092
+ "learning_rate": 9.823372139533906e-06,
1093
+ "loss": 1.2503,
1094
+ "step": 1810
1095
+ },
1096
+ {
1097
+ "epoch": 0.11,
1098
+ "learning_rate": 9.821266920696407e-06,
1099
+ "loss": 1.2409,
1100
+ "step": 1820
1101
+ },
1102
+ {
1103
+ "epoch": 0.11,
1104
+ "learning_rate": 9.819161701858909e-06,
1105
+ "loss": 1.2236,
1106
+ "step": 1830
1107
+ },
1108
+ {
1109
+ "epoch": 0.11,
1110
+ "learning_rate": 9.817056483021412e-06,
1111
+ "loss": 1.2246,
1112
+ "step": 1840
1113
+ },
1114
+ {
1115
+ "epoch": 0.11,
1116
+ "learning_rate": 9.814951264183912e-06,
1117
+ "loss": 1.234,
1118
+ "step": 1850
1119
+ },
1120
+ {
1121
+ "epoch": 0.12,
1122
+ "learning_rate": 9.812846045346415e-06,
1123
+ "loss": 1.2384,
1124
+ "step": 1860
1125
+ },
1126
+ {
1127
+ "epoch": 0.12,
1128
+ "learning_rate": 9.810740826508917e-06,
1129
+ "loss": 1.2103,
1130
+ "step": 1870
1131
+ },
1132
+ {
1133
+ "epoch": 0.12,
1134
+ "learning_rate": 9.808635607671418e-06,
1135
+ "loss": 1.2398,
1136
+ "step": 1880
1137
+ },
1138
+ {
1139
+ "epoch": 0.12,
1140
+ "learning_rate": 9.80653038883392e-06,
1141
+ "loss": 1.2246,
1142
+ "step": 1890
1143
+ },
1144
+ {
1145
+ "epoch": 0.12,
1146
+ "learning_rate": 9.804425169996421e-06,
1147
+ "loss": 1.1894,
1148
+ "step": 1900
1149
+ },
1150
+ {
1151
+ "epoch": 0.12,
1152
+ "learning_rate": 9.802319951158924e-06,
1153
+ "loss": 1.262,
1154
+ "step": 1910
1155
+ },
1156
+ {
1157
+ "epoch": 0.12,
1158
+ "learning_rate": 9.800214732321426e-06,
1159
+ "loss": 1.2595,
1160
+ "step": 1920
1161
+ },
1162
+ {
1163
+ "epoch": 0.12,
1164
+ "learning_rate": 9.798109513483927e-06,
1165
+ "loss": 1.1954,
1166
+ "step": 1930
1167
+ },
1168
+ {
1169
+ "epoch": 0.12,
1170
+ "learning_rate": 9.79600429464643e-06,
1171
+ "loss": 1.2578,
1172
+ "step": 1940
1173
+ },
1174
+ {
1175
+ "epoch": 0.12,
1176
+ "learning_rate": 9.793899075808932e-06,
1177
+ "loss": 1.2286,
1178
+ "step": 1950
1179
+ },
1180
+ {
1181
+ "epoch": 0.12,
1182
+ "learning_rate": 9.791793856971433e-06,
1183
+ "loss": 1.2455,
1184
+ "step": 1960
1185
+ },
1186
+ {
1187
+ "epoch": 0.12,
1188
+ "learning_rate": 9.789688638133936e-06,
1189
+ "loss": 1.214,
1190
+ "step": 1970
1191
+ },
1192
+ {
1193
+ "epoch": 0.12,
1194
+ "learning_rate": 9.787583419296436e-06,
1195
+ "loss": 1.2488,
1196
+ "step": 1980
1197
+ },
1198
+ {
1199
+ "epoch": 0.12,
1200
+ "learning_rate": 9.785478200458939e-06,
1201
+ "loss": 1.266,
1202
+ "step": 1990
1203
+ },
1204
+ {
1205
+ "epoch": 0.12,
1206
+ "learning_rate": 9.783372981621441e-06,
1207
+ "loss": 1.2139,
1208
+ "step": 2000
1209
+ },
1210
+ {
1211
+ "epoch": 0.12,
1212
+ "learning_rate": 9.781267762783942e-06,
1213
+ "loss": 1.2516,
1214
+ "step": 2010
1215
+ },
1216
+ {
1217
+ "epoch": 0.13,
1218
+ "learning_rate": 9.779162543946445e-06,
1219
+ "loss": 1.2878,
1220
+ "step": 2020
1221
+ },
1222
+ {
1223
+ "epoch": 0.13,
1224
+ "learning_rate": 9.777057325108945e-06,
1225
+ "loss": 1.2494,
1226
+ "step": 2030
1227
+ },
1228
+ {
1229
+ "epoch": 0.13,
1230
+ "learning_rate": 9.774952106271448e-06,
1231
+ "loss": 1.207,
1232
+ "step": 2040
1233
+ },
1234
+ {
1235
+ "epoch": 0.13,
1236
+ "learning_rate": 9.77284688743395e-06,
1237
+ "loss": 1.2796,
1238
+ "step": 2050
1239
+ },
1240
+ {
1241
+ "epoch": 0.13,
1242
+ "learning_rate": 9.770741668596451e-06,
1243
+ "loss": 1.2285,
1244
+ "step": 2060
1245
+ },
1246
+ {
1247
+ "epoch": 0.13,
1248
+ "learning_rate": 9.768636449758954e-06,
1249
+ "loss": 1.2467,
1250
+ "step": 2070
1251
+ },
1252
+ {
1253
+ "epoch": 0.13,
1254
+ "learning_rate": 9.766531230921456e-06,
1255
+ "loss": 1.1801,
1256
+ "step": 2080
1257
+ },
1258
+ {
1259
+ "epoch": 0.13,
1260
+ "learning_rate": 9.764426012083957e-06,
1261
+ "loss": 1.2399,
1262
+ "step": 2090
1263
+ },
1264
+ {
1265
+ "epoch": 0.13,
1266
+ "learning_rate": 9.76232079324646e-06,
1267
+ "loss": 1.2359,
1268
+ "step": 2100
1269
+ },
1270
+ {
1271
+ "epoch": 0.13,
1272
+ "learning_rate": 9.76021557440896e-06,
1273
+ "loss": 1.2074,
1274
+ "step": 2110
1275
+ },
1276
+ {
1277
+ "epoch": 0.13,
1278
+ "learning_rate": 9.758110355571463e-06,
1279
+ "loss": 1.2601,
1280
+ "step": 2120
1281
+ },
1282
+ {
1283
+ "epoch": 0.13,
1284
+ "learning_rate": 9.756005136733965e-06,
1285
+ "loss": 1.2456,
1286
+ "step": 2130
1287
+ },
1288
+ {
1289
+ "epoch": 0.13,
1290
+ "learning_rate": 9.753899917896466e-06,
1291
+ "loss": 1.2479,
1292
+ "step": 2140
1293
+ },
1294
+ {
1295
+ "epoch": 0.13,
1296
+ "learning_rate": 9.751794699058969e-06,
1297
+ "loss": 1.2593,
1298
+ "step": 2150
1299
+ },
1300
+ {
1301
+ "epoch": 0.13,
1302
+ "learning_rate": 9.749689480221471e-06,
1303
+ "loss": 1.1856,
1304
+ "step": 2160
1305
+ },
1306
+ {
1307
+ "epoch": 0.13,
1308
+ "learning_rate": 9.747584261383972e-06,
1309
+ "loss": 1.2634,
1310
+ "step": 2170
1311
+ },
1312
+ {
1313
+ "epoch": 0.13,
1314
+ "learning_rate": 9.745479042546474e-06,
1315
+ "loss": 1.2046,
1316
+ "step": 2180
1317
+ },
1318
+ {
1319
+ "epoch": 0.14,
1320
+ "learning_rate": 9.743373823708975e-06,
1321
+ "loss": 1.2753,
1322
+ "step": 2190
1323
+ },
1324
+ {
1325
+ "epoch": 0.14,
1326
+ "learning_rate": 9.741268604871478e-06,
1327
+ "loss": 1.2393,
1328
+ "step": 2200
1329
+ },
1330
+ {
1331
+ "epoch": 0.14,
1332
+ "learning_rate": 9.739163386033979e-06,
1333
+ "loss": 1.224,
1334
+ "step": 2210
1335
+ },
1336
+ {
1337
+ "epoch": 0.14,
1338
+ "learning_rate": 9.737058167196481e-06,
1339
+ "loss": 1.2767,
1340
+ "step": 2220
1341
+ },
1342
+ {
1343
+ "epoch": 0.14,
1344
+ "learning_rate": 9.734952948358982e-06,
1345
+ "loss": 1.2584,
1346
+ "step": 2230
1347
+ },
1348
+ {
1349
+ "epoch": 0.14,
1350
+ "learning_rate": 9.732847729521484e-06,
1351
+ "loss": 1.2717,
1352
+ "step": 2240
1353
+ },
1354
+ {
1355
+ "epoch": 0.14,
1356
+ "learning_rate": 9.730742510683985e-06,
1357
+ "loss": 1.2364,
1358
+ "step": 2250
1359
+ },
1360
+ {
1361
+ "epoch": 0.14,
1362
+ "learning_rate": 9.728637291846488e-06,
1363
+ "loss": 1.2354,
1364
+ "step": 2260
1365
+ },
1366
+ {
1367
+ "epoch": 0.14,
1368
+ "learning_rate": 9.72653207300899e-06,
1369
+ "loss": 1.2347,
1370
+ "step": 2270
1371
+ },
1372
+ {
1373
+ "epoch": 0.14,
1374
+ "learning_rate": 9.724426854171491e-06,
1375
+ "loss": 1.2537,
1376
+ "step": 2280
1377
+ },
1378
+ {
1379
+ "epoch": 0.14,
1380
+ "learning_rate": 9.722321635333993e-06,
1381
+ "loss": 1.2186,
1382
+ "step": 2290
1383
+ },
1384
+ {
1385
+ "epoch": 0.14,
1386
+ "learning_rate": 9.720216416496496e-06,
1387
+ "loss": 1.2351,
1388
+ "step": 2300
1389
+ },
1390
+ {
1391
+ "epoch": 0.14,
1392
+ "learning_rate": 9.718111197658997e-06,
1393
+ "loss": 1.2325,
1394
+ "step": 2310
1395
+ },
1396
+ {
1397
+ "epoch": 0.14,
1398
+ "learning_rate": 9.7160059788215e-06,
1399
+ "loss": 1.1996,
1400
+ "step": 2320
1401
+ },
1402
+ {
1403
+ "epoch": 0.14,
1404
+ "learning_rate": 9.713900759984e-06,
1405
+ "loss": 1.2023,
1406
+ "step": 2330
1407
+ },
1408
+ {
1409
+ "epoch": 0.14,
1410
+ "learning_rate": 9.711795541146503e-06,
1411
+ "loss": 1.2527,
1412
+ "step": 2340
1413
+ },
1414
+ {
1415
+ "epoch": 0.15,
1416
+ "learning_rate": 9.709690322309005e-06,
1417
+ "loss": 1.2281,
1418
+ "step": 2350
1419
+ },
1420
+ {
1421
+ "epoch": 0.15,
1422
+ "learning_rate": 9.707585103471506e-06,
1423
+ "loss": 1.2382,
1424
+ "step": 2360
1425
+ },
1426
+ {
1427
+ "epoch": 0.15,
1428
+ "learning_rate": 9.705479884634008e-06,
1429
+ "loss": 1.2405,
1430
+ "step": 2370
1431
+ },
1432
+ {
1433
+ "epoch": 0.15,
1434
+ "learning_rate": 9.70337466579651e-06,
1435
+ "loss": 1.248,
1436
+ "step": 2380
1437
+ },
1438
+ {
1439
+ "epoch": 0.15,
1440
+ "learning_rate": 9.701269446959012e-06,
1441
+ "loss": 1.224,
1442
+ "step": 2390
1443
+ },
1444
+ {
1445
+ "epoch": 0.15,
1446
+ "learning_rate": 9.699164228121514e-06,
1447
+ "loss": 1.22,
1448
+ "step": 2400
1449
+ },
1450
+ {
1451
+ "epoch": 0.15,
1452
+ "learning_rate": 9.697059009284015e-06,
1453
+ "loss": 1.219,
1454
+ "step": 2410
1455
+ },
1456
+ {
1457
+ "epoch": 0.15,
1458
+ "learning_rate": 9.694953790446517e-06,
1459
+ "loss": 1.2518,
1460
+ "step": 2420
1461
+ },
1462
+ {
1463
+ "epoch": 0.15,
1464
+ "learning_rate": 9.69284857160902e-06,
1465
+ "loss": 1.219,
1466
+ "step": 2430
1467
+ },
1468
+ {
1469
+ "epoch": 0.15,
1470
+ "learning_rate": 9.69074335277152e-06,
1471
+ "loss": 1.2168,
1472
+ "step": 2440
1473
+ },
1474
+ {
1475
+ "epoch": 0.15,
1476
+ "learning_rate": 9.688638133934023e-06,
1477
+ "loss": 1.2469,
1478
+ "step": 2450
1479
+ },
1480
+ {
1481
+ "epoch": 0.15,
1482
+ "learning_rate": 9.686532915096524e-06,
1483
+ "loss": 1.2381,
1484
+ "step": 2460
1485
+ },
1486
+ {
1487
+ "epoch": 0.15,
1488
+ "learning_rate": 9.684427696259027e-06,
1489
+ "loss": 1.2001,
1490
+ "step": 2470
1491
+ },
1492
+ {
1493
+ "epoch": 0.15,
1494
+ "learning_rate": 9.682322477421529e-06,
1495
+ "loss": 1.2004,
1496
+ "step": 2480
1497
+ },
1498
+ {
1499
+ "epoch": 0.15,
1500
+ "learning_rate": 9.68021725858403e-06,
1501
+ "loss": 1.2409,
1502
+ "step": 2490
1503
+ },
1504
+ {
1505
+ "epoch": 0.15,
1506
+ "learning_rate": 9.678112039746532e-06,
1507
+ "loss": 1.2389,
1508
+ "step": 2500
1509
+ },
1510
+ {
1511
+ "epoch": 0.16,
1512
+ "learning_rate": 9.676006820909035e-06,
1513
+ "loss": 1.242,
1514
+ "step": 2510
1515
+ },
1516
+ {
1517
+ "epoch": 0.16,
1518
+ "learning_rate": 9.673901602071536e-06,
1519
+ "loss": 1.2372,
1520
+ "step": 2520
1521
+ },
1522
+ {
1523
+ "epoch": 0.16,
1524
+ "learning_rate": 9.671796383234038e-06,
1525
+ "loss": 1.2223,
1526
+ "step": 2530
1527
+ },
1528
+ {
1529
+ "epoch": 0.16,
1530
+ "learning_rate": 9.669691164396539e-06,
1531
+ "loss": 1.2506,
1532
+ "step": 2540
1533
+ },
1534
+ {
1535
+ "epoch": 0.16,
1536
+ "learning_rate": 9.667585945559041e-06,
1537
+ "loss": 1.2093,
1538
+ "step": 2550
1539
+ },
1540
+ {
1541
+ "epoch": 0.16,
1542
+ "learning_rate": 9.665480726721544e-06,
1543
+ "loss": 1.2171,
1544
+ "step": 2560
1545
+ },
1546
+ {
1547
+ "epoch": 0.16,
1548
+ "learning_rate": 9.663375507884045e-06,
1549
+ "loss": 1.2363,
1550
+ "step": 2570
1551
+ },
1552
+ {
1553
+ "epoch": 0.16,
1554
+ "learning_rate": 9.661270289046547e-06,
1555
+ "loss": 1.2978,
1556
+ "step": 2580
1557
+ },
1558
+ {
1559
+ "epoch": 0.16,
1560
+ "learning_rate": 9.65916507020905e-06,
1561
+ "loss": 1.2216,
1562
+ "step": 2590
1563
+ },
1564
+ {
1565
+ "epoch": 0.16,
1566
+ "learning_rate": 9.65705985137155e-06,
1567
+ "loss": 1.1937,
1568
+ "step": 2600
1569
+ },
1570
+ {
1571
+ "epoch": 0.16,
1572
+ "learning_rate": 9.654954632534053e-06,
1573
+ "loss": 1.2366,
1574
+ "step": 2610
1575
+ },
1576
+ {
1577
+ "epoch": 0.16,
1578
+ "learning_rate": 9.652849413696554e-06,
1579
+ "loss": 1.2465,
1580
+ "step": 2620
1581
+ },
1582
+ {
1583
+ "epoch": 0.16,
1584
+ "learning_rate": 9.650744194859056e-06,
1585
+ "loss": 1.2704,
1586
+ "step": 2630
1587
+ },
1588
+ {
1589
+ "epoch": 0.16,
1590
+ "learning_rate": 9.648638976021559e-06,
1591
+ "loss": 1.2113,
1592
+ "step": 2640
1593
+ },
1594
+ {
1595
+ "epoch": 0.16,
1596
+ "learning_rate": 9.64653375718406e-06,
1597
+ "loss": 1.2679,
1598
+ "step": 2650
1599
+ },
1600
+ {
1601
+ "epoch": 0.16,
1602
+ "learning_rate": 9.644428538346562e-06,
1603
+ "loss": 1.2005,
1604
+ "step": 2660
1605
+ },
1606
+ {
1607
+ "epoch": 0.17,
1608
+ "learning_rate": 9.642323319509063e-06,
1609
+ "loss": 1.2474,
1610
+ "step": 2670
1611
+ },
1612
+ {
1613
+ "epoch": 0.17,
1614
+ "learning_rate": 9.640218100671565e-06,
1615
+ "loss": 1.2308,
1616
+ "step": 2680
1617
+ },
1618
+ {
1619
+ "epoch": 0.17,
1620
+ "learning_rate": 9.638112881834068e-06,
1621
+ "loss": 1.2391,
1622
+ "step": 2690
1623
+ },
1624
+ {
1625
+ "epoch": 0.17,
1626
+ "learning_rate": 9.636007662996569e-06,
1627
+ "loss": 1.1968,
1628
+ "step": 2700
1629
+ },
1630
+ {
1631
+ "epoch": 0.17,
1632
+ "learning_rate": 9.633902444159071e-06,
1633
+ "loss": 1.2001,
1634
+ "step": 2710
1635
+ },
1636
+ {
1637
+ "epoch": 0.17,
1638
+ "learning_rate": 9.631797225321574e-06,
1639
+ "loss": 1.2688,
1640
+ "step": 2720
1641
+ },
1642
+ {
1643
+ "epoch": 0.17,
1644
+ "learning_rate": 9.629692006484075e-06,
1645
+ "loss": 1.2646,
1646
+ "step": 2730
1647
+ },
1648
+ {
1649
+ "epoch": 0.17,
1650
+ "learning_rate": 9.627586787646577e-06,
1651
+ "loss": 1.2606,
1652
+ "step": 2740
1653
+ },
1654
+ {
1655
+ "epoch": 0.17,
1656
+ "learning_rate": 9.625481568809078e-06,
1657
+ "loss": 1.1915,
1658
+ "step": 2750
1659
+ },
1660
+ {
1661
+ "epoch": 0.17,
1662
+ "learning_rate": 9.62337634997158e-06,
1663
+ "loss": 1.204,
1664
+ "step": 2760
1665
+ },
1666
+ {
1667
+ "epoch": 0.17,
1668
+ "learning_rate": 9.621271131134083e-06,
1669
+ "loss": 1.2128,
1670
+ "step": 2770
1671
+ },
1672
+ {
1673
+ "epoch": 0.17,
1674
+ "learning_rate": 9.619165912296584e-06,
1675
+ "loss": 1.2116,
1676
+ "step": 2780
1677
+ },
1678
+ {
1679
+ "epoch": 0.17,
1680
+ "learning_rate": 9.617060693459086e-06,
1681
+ "loss": 1.2287,
1682
+ "step": 2790
1683
+ },
1684
+ {
1685
+ "epoch": 0.17,
1686
+ "learning_rate": 9.614955474621589e-06,
1687
+ "loss": 1.2443,
1688
+ "step": 2800
1689
+ },
1690
+ {
1691
+ "epoch": 0.17,
1692
+ "learning_rate": 9.61285025578409e-06,
1693
+ "loss": 1.2926,
1694
+ "step": 2810
1695
+ },
1696
+ {
1697
+ "epoch": 0.17,
1698
+ "learning_rate": 9.610745036946592e-06,
1699
+ "loss": 1.2195,
1700
+ "step": 2820
1701
+ },
1702
+ {
1703
+ "epoch": 0.18,
1704
+ "learning_rate": 9.608639818109093e-06,
1705
+ "loss": 1.2345,
1706
+ "step": 2830
1707
+ },
1708
+ {
1709
+ "epoch": 0.18,
1710
+ "learning_rate": 9.606534599271595e-06,
1711
+ "loss": 1.2588,
1712
+ "step": 2840
1713
+ },
1714
+ {
1715
+ "epoch": 0.18,
1716
+ "learning_rate": 9.604429380434098e-06,
1717
+ "loss": 1.2392,
1718
+ "step": 2850
1719
+ },
1720
+ {
1721
+ "epoch": 0.18,
1722
+ "learning_rate": 9.602324161596599e-06,
1723
+ "loss": 1.2529,
1724
+ "step": 2860
1725
+ },
1726
+ {
1727
+ "epoch": 0.18,
1728
+ "learning_rate": 9.600218942759101e-06,
1729
+ "loss": 1.2119,
1730
+ "step": 2870
1731
+ },
1732
+ {
1733
+ "epoch": 0.18,
1734
+ "learning_rate": 9.598113723921602e-06,
1735
+ "loss": 1.2416,
1736
+ "step": 2880
1737
+ },
1738
+ {
1739
+ "epoch": 0.18,
1740
+ "learning_rate": 9.596008505084104e-06,
1741
+ "loss": 1.2111,
1742
+ "step": 2890
1743
+ },
1744
+ {
1745
+ "epoch": 0.18,
1746
+ "learning_rate": 9.593903286246607e-06,
1747
+ "loss": 1.2493,
1748
+ "step": 2900
1749
+ },
1750
+ {
1751
+ "epoch": 0.18,
1752
+ "learning_rate": 9.591798067409108e-06,
1753
+ "loss": 1.2481,
1754
+ "step": 2910
1755
+ },
1756
+ {
1757
+ "epoch": 0.18,
1758
+ "learning_rate": 9.58969284857161e-06,
1759
+ "loss": 1.2265,
1760
+ "step": 2920
1761
+ },
1762
+ {
1763
+ "epoch": 0.18,
1764
+ "learning_rate": 9.587587629734113e-06,
1765
+ "loss": 1.2549,
1766
+ "step": 2930
1767
+ },
1768
+ {
1769
+ "epoch": 0.18,
1770
+ "learning_rate": 9.585482410896613e-06,
1771
+ "loss": 1.2474,
1772
+ "step": 2940
1773
+ },
1774
+ {
1775
+ "epoch": 0.18,
1776
+ "learning_rate": 9.583377192059116e-06,
1777
+ "loss": 1.1773,
1778
+ "step": 2950
1779
+ },
1780
+ {
1781
+ "epoch": 0.18,
1782
+ "learning_rate": 9.581271973221617e-06,
1783
+ "loss": 1.2612,
1784
+ "step": 2960
1785
+ },
1786
+ {
1787
+ "epoch": 0.18,
1788
+ "learning_rate": 9.57916675438412e-06,
1789
+ "loss": 1.2247,
1790
+ "step": 2970
1791
+ },
1792
+ {
1793
+ "epoch": 0.18,
1794
+ "learning_rate": 9.577061535546622e-06,
1795
+ "loss": 1.2075,
1796
+ "step": 2980
1797
+ },
1798
+ {
1799
+ "epoch": 0.19,
1800
+ "learning_rate": 9.574956316709123e-06,
1801
+ "loss": 1.1812,
1802
+ "step": 2990
1803
+ },
1804
+ {
1805
+ "epoch": 0.19,
1806
+ "learning_rate": 9.572851097871625e-06,
1807
+ "loss": 1.2058,
1808
+ "step": 3000
1809
+ },
1810
+ {
1811
+ "epoch": 0.19,
1812
+ "learning_rate": 9.570745879034128e-06,
1813
+ "loss": 1.2781,
1814
+ "step": 3010
1815
+ },
1816
+ {
1817
+ "epoch": 0.19,
1818
+ "learning_rate": 9.568640660196628e-06,
1819
+ "loss": 1.2572,
1820
+ "step": 3020
1821
+ },
1822
+ {
1823
+ "epoch": 0.19,
1824
+ "learning_rate": 9.566535441359131e-06,
1825
+ "loss": 1.2794,
1826
+ "step": 3030
1827
+ },
1828
+ {
1829
+ "epoch": 0.19,
1830
+ "learning_rate": 9.564430222521632e-06,
1831
+ "loss": 1.2136,
1832
+ "step": 3040
1833
+ },
1834
+ {
1835
+ "epoch": 0.19,
1836
+ "learning_rate": 9.562325003684134e-06,
1837
+ "loss": 1.2632,
1838
+ "step": 3050
1839
+ },
1840
+ {
1841
+ "epoch": 0.19,
1842
+ "learning_rate": 9.560219784846637e-06,
1843
+ "loss": 1.2584,
1844
+ "step": 3060
1845
+ },
1846
+ {
1847
+ "epoch": 0.19,
1848
+ "learning_rate": 9.558114566009137e-06,
1849
+ "loss": 1.286,
1850
+ "step": 3070
1851
+ },
1852
+ {
1853
+ "epoch": 0.19,
1854
+ "learning_rate": 9.55600934717164e-06,
1855
+ "loss": 1.247,
1856
+ "step": 3080
1857
+ },
1858
+ {
1859
+ "epoch": 0.19,
1860
+ "learning_rate": 9.55390412833414e-06,
1861
+ "loss": 1.2715,
1862
+ "step": 3090
1863
+ },
1864
+ {
1865
+ "epoch": 0.19,
1866
+ "learning_rate": 9.551798909496643e-06,
1867
+ "loss": 1.2184,
1868
+ "step": 3100
1869
+ },
1870
+ {
1871
+ "epoch": 0.19,
1872
+ "learning_rate": 9.549693690659146e-06,
1873
+ "loss": 1.261,
1874
+ "step": 3110
1875
+ },
1876
+ {
1877
+ "epoch": 0.19,
1878
+ "learning_rate": 9.547588471821647e-06,
1879
+ "loss": 1.2183,
1880
+ "step": 3120
1881
+ },
1882
+ {
1883
+ "epoch": 0.19,
1884
+ "learning_rate": 9.545483252984149e-06,
1885
+ "loss": 1.1887,
1886
+ "step": 3130
1887
+ },
1888
+ {
1889
+ "epoch": 0.19,
1890
+ "learning_rate": 9.543378034146652e-06,
1891
+ "loss": 1.2405,
1892
+ "step": 3140
1893
+ },
1894
+ {
1895
+ "epoch": 0.19,
1896
+ "learning_rate": 9.541272815309152e-06,
1897
+ "loss": 1.2499,
1898
+ "step": 3150
1899
+ },
1900
+ {
1901
+ "epoch": 0.2,
1902
+ "learning_rate": 9.539167596471653e-06,
1903
+ "loss": 1.2164,
1904
+ "step": 3160
1905
+ },
1906
+ {
1907
+ "epoch": 0.2,
1908
+ "learning_rate": 9.537062377634156e-06,
1909
+ "loss": 1.2614,
1910
+ "step": 3170
1911
+ },
1912
+ {
1913
+ "epoch": 0.2,
1914
+ "learning_rate": 9.534957158796657e-06,
1915
+ "loss": 1.2475,
1916
+ "step": 3180
1917
+ },
1918
+ {
1919
+ "epoch": 0.2,
1920
+ "learning_rate": 9.532851939959159e-06,
1921
+ "loss": 1.2559,
1922
+ "step": 3190
1923
+ },
1924
+ {
1925
+ "epoch": 0.2,
1926
+ "learning_rate": 9.530746721121662e-06,
1927
+ "loss": 1.2457,
1928
+ "step": 3200
1929
+ },
1930
+ {
1931
+ "epoch": 0.2,
1932
+ "learning_rate": 9.528641502284162e-06,
1933
+ "loss": 1.2228,
1934
+ "step": 3210
1935
+ },
1936
+ {
1937
+ "epoch": 0.2,
1938
+ "learning_rate": 9.526536283446665e-06,
1939
+ "loss": 1.219,
1940
+ "step": 3220
1941
+ },
1942
+ {
1943
+ "epoch": 0.2,
1944
+ "learning_rate": 9.524431064609166e-06,
1945
+ "loss": 1.2255,
1946
+ "step": 3230
1947
+ },
1948
+ {
1949
+ "epoch": 0.2,
1950
+ "learning_rate": 9.522325845771668e-06,
1951
+ "loss": 1.1923,
1952
+ "step": 3240
1953
+ },
1954
+ {
1955
+ "epoch": 0.2,
1956
+ "learning_rate": 9.52022062693417e-06,
1957
+ "loss": 1.1996,
1958
+ "step": 3250
1959
+ },
1960
+ {
1961
+ "epoch": 0.2,
1962
+ "learning_rate": 9.518115408096671e-06,
1963
+ "loss": 1.2186,
1964
+ "step": 3260
1965
+ },
1966
+ {
1967
+ "epoch": 0.2,
1968
+ "learning_rate": 9.516010189259174e-06,
1969
+ "loss": 1.2384,
1970
+ "step": 3270
1971
+ },
1972
+ {
1973
+ "epoch": 0.2,
1974
+ "learning_rate": 9.513904970421676e-06,
1975
+ "loss": 1.2119,
1976
+ "step": 3280
1977
+ },
1978
+ {
1979
+ "epoch": 0.2,
1980
+ "learning_rate": 9.511799751584177e-06,
1981
+ "loss": 1.2455,
1982
+ "step": 3290
1983
+ },
1984
+ {
1985
+ "epoch": 0.2,
1986
+ "learning_rate": 9.50969453274668e-06,
1987
+ "loss": 1.2314,
1988
+ "step": 3300
1989
+ },
1990
+ {
1991
+ "epoch": 0.2,
1992
+ "learning_rate": 9.50758931390918e-06,
1993
+ "loss": 1.1995,
1994
+ "step": 3310
1995
+ },
1996
+ {
1997
+ "epoch": 0.21,
1998
+ "learning_rate": 9.505484095071683e-06,
1999
+ "loss": 1.2308,
2000
+ "step": 3320
2001
+ },
2002
+ {
2003
+ "epoch": 0.21,
2004
+ "learning_rate": 9.503378876234186e-06,
2005
+ "loss": 1.1957,
2006
+ "step": 3330
2007
+ },
2008
+ {
2009
+ "epoch": 0.21,
2010
+ "learning_rate": 9.501273657396686e-06,
2011
+ "loss": 1.2557,
2012
+ "step": 3340
2013
+ },
2014
+ {
2015
+ "epoch": 0.21,
2016
+ "learning_rate": 9.499168438559189e-06,
2017
+ "loss": 1.2351,
2018
+ "step": 3350
2019
+ },
2020
+ {
2021
+ "epoch": 0.21,
2022
+ "learning_rate": 9.497063219721691e-06,
2023
+ "loss": 1.2085,
2024
+ "step": 3360
2025
+ },
2026
+ {
2027
+ "epoch": 0.21,
2028
+ "learning_rate": 9.494958000884192e-06,
2029
+ "loss": 1.2241,
2030
+ "step": 3370
2031
+ },
2032
+ {
2033
+ "epoch": 0.21,
2034
+ "learning_rate": 9.492852782046695e-06,
2035
+ "loss": 1.1909,
2036
+ "step": 3380
2037
+ },
2038
+ {
2039
+ "epoch": 0.21,
2040
+ "learning_rate": 9.490747563209195e-06,
2041
+ "loss": 1.1886,
2042
+ "step": 3390
2043
+ },
2044
+ {
2045
+ "epoch": 0.21,
2046
+ "learning_rate": 9.488642344371698e-06,
2047
+ "loss": 1.2161,
2048
+ "step": 3400
2049
+ },
2050
+ {
2051
+ "epoch": 0.21,
2052
+ "learning_rate": 9.4865371255342e-06,
2053
+ "loss": 1.2718,
2054
+ "step": 3410
2055
+ },
2056
+ {
2057
+ "epoch": 0.21,
2058
+ "learning_rate": 9.484431906696701e-06,
2059
+ "loss": 1.2007,
2060
+ "step": 3420
2061
+ },
2062
+ {
2063
+ "epoch": 0.21,
2064
+ "learning_rate": 9.482326687859204e-06,
2065
+ "loss": 1.2038,
2066
+ "step": 3430
2067
+ },
2068
+ {
2069
+ "epoch": 0.21,
2070
+ "learning_rate": 9.480221469021705e-06,
2071
+ "loss": 1.2153,
2072
+ "step": 3440
2073
+ },
2074
+ {
2075
+ "epoch": 0.21,
2076
+ "learning_rate": 9.478116250184207e-06,
2077
+ "loss": 1.1932,
2078
+ "step": 3450
2079
+ },
2080
+ {
2081
+ "epoch": 0.21,
2082
+ "learning_rate": 9.47601103134671e-06,
2083
+ "loss": 1.255,
2084
+ "step": 3460
2085
+ },
2086
+ {
2087
+ "epoch": 0.21,
2088
+ "learning_rate": 9.47390581250921e-06,
2089
+ "loss": 1.2738,
2090
+ "step": 3470
2091
+ },
2092
+ {
2093
+ "epoch": 0.22,
2094
+ "learning_rate": 9.471800593671713e-06,
2095
+ "loss": 1.2533,
2096
+ "step": 3480
2097
+ },
2098
+ {
2099
+ "epoch": 0.22,
2100
+ "learning_rate": 9.469695374834215e-06,
2101
+ "loss": 1.2408,
2102
+ "step": 3490
2103
+ },
2104
+ {
2105
+ "epoch": 0.22,
2106
+ "learning_rate": 9.467590155996716e-06,
2107
+ "loss": 1.2331,
2108
+ "step": 3500
2109
+ },
2110
+ {
2111
+ "epoch": 0.22,
2112
+ "learning_rate": 9.465484937159219e-06,
2113
+ "loss": 1.2205,
2114
+ "step": 3510
2115
+ },
2116
+ {
2117
+ "epoch": 0.22,
2118
+ "learning_rate": 9.46337971832172e-06,
2119
+ "loss": 1.2569,
2120
+ "step": 3520
2121
+ },
2122
+ {
2123
+ "epoch": 0.22,
2124
+ "learning_rate": 9.461274499484222e-06,
2125
+ "loss": 1.2245,
2126
+ "step": 3530
2127
+ },
2128
+ {
2129
+ "epoch": 0.22,
2130
+ "learning_rate": 9.459169280646724e-06,
2131
+ "loss": 1.2192,
2132
+ "step": 3540
2133
+ },
2134
+ {
2135
+ "epoch": 0.22,
2136
+ "learning_rate": 9.457064061809225e-06,
2137
+ "loss": 1.2518,
2138
+ "step": 3550
2139
+ },
2140
+ {
2141
+ "epoch": 0.22,
2142
+ "learning_rate": 9.454958842971728e-06,
2143
+ "loss": 1.2619,
2144
+ "step": 3560
2145
+ },
2146
+ {
2147
+ "epoch": 0.22,
2148
+ "learning_rate": 9.45285362413423e-06,
2149
+ "loss": 1.1859,
2150
+ "step": 3570
2151
+ },
2152
+ {
2153
+ "epoch": 0.22,
2154
+ "learning_rate": 9.450748405296731e-06,
2155
+ "loss": 1.2216,
2156
+ "step": 3580
2157
+ },
2158
+ {
2159
+ "epoch": 0.22,
2160
+ "learning_rate": 9.448643186459234e-06,
2161
+ "loss": 1.221,
2162
+ "step": 3590
2163
+ },
2164
+ {
2165
+ "epoch": 0.22,
2166
+ "learning_rate": 9.446537967621734e-06,
2167
+ "loss": 1.2224,
2168
+ "step": 3600
2169
+ },
2170
+ {
2171
+ "epoch": 0.22,
2172
+ "learning_rate": 9.444432748784237e-06,
2173
+ "loss": 1.2324,
2174
+ "step": 3610
2175
+ },
2176
+ {
2177
+ "epoch": 0.22,
2178
+ "learning_rate": 9.44232752994674e-06,
2179
+ "loss": 1.1747,
2180
+ "step": 3620
2181
+ },
2182
+ {
2183
+ "epoch": 0.22,
2184
+ "learning_rate": 9.44022231110924e-06,
2185
+ "loss": 1.1962,
2186
+ "step": 3630
2187
+ },
2188
+ {
2189
+ "epoch": 0.23,
2190
+ "learning_rate": 9.438117092271743e-06,
2191
+ "loss": 1.2235,
2192
+ "step": 3640
2193
+ },
2194
+ {
2195
+ "epoch": 0.23,
2196
+ "learning_rate": 9.436011873434245e-06,
2197
+ "loss": 1.2081,
2198
+ "step": 3650
2199
+ },
2200
+ {
2201
+ "epoch": 0.23,
2202
+ "learning_rate": 9.433906654596746e-06,
2203
+ "loss": 1.2398,
2204
+ "step": 3660
2205
+ },
2206
+ {
2207
+ "epoch": 0.23,
2208
+ "learning_rate": 9.431801435759248e-06,
2209
+ "loss": 1.248,
2210
+ "step": 3670
2211
+ },
2212
+ {
2213
+ "epoch": 0.23,
2214
+ "learning_rate": 9.42969621692175e-06,
2215
+ "loss": 1.1936,
2216
+ "step": 3680
2217
+ },
2218
+ {
2219
+ "epoch": 0.23,
2220
+ "learning_rate": 9.427590998084252e-06,
2221
+ "loss": 1.2317,
2222
+ "step": 3690
2223
+ },
2224
+ {
2225
+ "epoch": 0.23,
2226
+ "learning_rate": 9.425485779246754e-06,
2227
+ "loss": 1.2139,
2228
+ "step": 3700
2229
+ },
2230
+ {
2231
+ "epoch": 0.23,
2232
+ "learning_rate": 9.423380560409255e-06,
2233
+ "loss": 1.1601,
2234
+ "step": 3710
2235
+ },
2236
+ {
2237
+ "epoch": 0.23,
2238
+ "learning_rate": 9.421275341571758e-06,
2239
+ "loss": 1.2127,
2240
+ "step": 3720
2241
+ },
2242
+ {
2243
+ "epoch": 0.23,
2244
+ "learning_rate": 9.419170122734258e-06,
2245
+ "loss": 1.2082,
2246
+ "step": 3730
2247
+ },
2248
+ {
2249
+ "epoch": 0.23,
2250
+ "learning_rate": 9.41706490389676e-06,
2251
+ "loss": 1.1971,
2252
+ "step": 3740
2253
+ },
2254
+ {
2255
+ "epoch": 0.23,
2256
+ "learning_rate": 9.414959685059263e-06,
2257
+ "loss": 1.2289,
2258
+ "step": 3750
2259
+ },
2260
+ {
2261
+ "epoch": 0.23,
2262
+ "learning_rate": 9.412854466221764e-06,
2263
+ "loss": 1.2133,
2264
+ "step": 3760
2265
+ },
2266
+ {
2267
+ "epoch": 0.23,
2268
+ "learning_rate": 9.410749247384267e-06,
2269
+ "loss": 1.2111,
2270
+ "step": 3770
2271
+ },
2272
+ {
2273
+ "epoch": 0.23,
2274
+ "learning_rate": 9.408644028546769e-06,
2275
+ "loss": 1.2342,
2276
+ "step": 3780
2277
+ },
2278
+ {
2279
+ "epoch": 0.23,
2280
+ "learning_rate": 9.40653880970927e-06,
2281
+ "loss": 1.217,
2282
+ "step": 3790
2283
+ },
2284
+ {
2285
+ "epoch": 0.24,
2286
+ "learning_rate": 9.404433590871772e-06,
2287
+ "loss": 1.2651,
2288
+ "step": 3800
2289
+ },
2290
+ {
2291
+ "epoch": 0.24,
2292
+ "learning_rate": 9.402328372034273e-06,
2293
+ "loss": 1.2259,
2294
+ "step": 3810
2295
+ },
2296
+ {
2297
+ "epoch": 0.24,
2298
+ "learning_rate": 9.400223153196776e-06,
2299
+ "loss": 1.2434,
2300
+ "step": 3820
2301
+ },
2302
+ {
2303
+ "epoch": 0.24,
2304
+ "learning_rate": 9.398117934359278e-06,
2305
+ "loss": 1.2199,
2306
+ "step": 3830
2307
+ },
2308
+ {
2309
+ "epoch": 0.24,
2310
+ "learning_rate": 9.396012715521779e-06,
2311
+ "loss": 1.2299,
2312
+ "step": 3840
2313
+ },
2314
+ {
2315
+ "epoch": 0.24,
2316
+ "learning_rate": 9.393907496684282e-06,
2317
+ "loss": 1.2156,
2318
+ "step": 3850
2319
+ },
2320
+ {
2321
+ "epoch": 0.24,
2322
+ "learning_rate": 9.391802277846784e-06,
2323
+ "loss": 1.2402,
2324
+ "step": 3860
2325
+ },
2326
+ {
2327
+ "epoch": 0.24,
2328
+ "learning_rate": 9.389697059009285e-06,
2329
+ "loss": 1.237,
2330
+ "step": 3870
2331
+ },
2332
+ {
2333
+ "epoch": 0.24,
2334
+ "learning_rate": 9.387591840171787e-06,
2335
+ "loss": 1.2141,
2336
+ "step": 3880
2337
+ },
2338
+ {
2339
+ "epoch": 0.24,
2340
+ "learning_rate": 9.385486621334288e-06,
2341
+ "loss": 1.2253,
2342
+ "step": 3890
2343
+ },
2344
+ {
2345
+ "epoch": 0.24,
2346
+ "learning_rate": 9.38338140249679e-06,
2347
+ "loss": 1.217,
2348
+ "step": 3900
2349
+ },
2350
+ {
2351
+ "epoch": 0.24,
2352
+ "learning_rate": 9.381276183659293e-06,
2353
+ "loss": 1.1919,
2354
+ "step": 3910
2355
+ },
2356
+ {
2357
+ "epoch": 0.24,
2358
+ "learning_rate": 9.379170964821794e-06,
2359
+ "loss": 1.1663,
2360
+ "step": 3920
2361
+ },
2362
+ {
2363
+ "epoch": 0.24,
2364
+ "learning_rate": 9.377065745984296e-06,
2365
+ "loss": 1.2338,
2366
+ "step": 3930
2367
+ },
2368
+ {
2369
+ "epoch": 0.24,
2370
+ "learning_rate": 9.374960527146797e-06,
2371
+ "loss": 1.2399,
2372
+ "step": 3940
2373
+ },
2374
+ {
2375
+ "epoch": 0.24,
2376
+ "learning_rate": 9.3728553083093e-06,
2377
+ "loss": 1.1608,
2378
+ "step": 3950
2379
+ },
2380
+ {
2381
+ "epoch": 0.25,
2382
+ "learning_rate": 9.370750089471802e-06,
2383
+ "loss": 1.1752,
2384
+ "step": 3960
2385
+ },
2386
+ {
2387
+ "epoch": 0.25,
2388
+ "learning_rate": 9.368644870634303e-06,
2389
+ "loss": 1.2364,
2390
+ "step": 3970
2391
+ },
2392
+ {
2393
+ "epoch": 0.25,
2394
+ "learning_rate": 9.366539651796806e-06,
2395
+ "loss": 1.2053,
2396
+ "step": 3980
2397
+ },
2398
+ {
2399
+ "epoch": 0.25,
2400
+ "learning_rate": 9.364434432959308e-06,
2401
+ "loss": 1.2431,
2402
+ "step": 3990
2403
+ },
2404
+ {
2405
+ "epoch": 0.25,
2406
+ "learning_rate": 9.362329214121809e-06,
2407
+ "loss": 1.1948,
2408
+ "step": 4000
2409
+ },
2410
+ {
2411
+ "epoch": 0.25,
2412
+ "learning_rate": 9.360223995284311e-06,
2413
+ "loss": 1.2248,
2414
+ "step": 4010
2415
+ },
2416
+ {
2417
+ "epoch": 0.25,
2418
+ "learning_rate": 9.358118776446812e-06,
2419
+ "loss": 1.2057,
2420
+ "step": 4020
2421
+ },
2422
+ {
2423
+ "epoch": 0.25,
2424
+ "learning_rate": 9.356013557609315e-06,
2425
+ "loss": 1.2373,
2426
+ "step": 4030
2427
+ },
2428
+ {
2429
+ "epoch": 0.25,
2430
+ "learning_rate": 9.353908338771817e-06,
2431
+ "loss": 1.1993,
2432
+ "step": 4040
2433
+ },
2434
+ {
2435
+ "epoch": 0.25,
2436
+ "learning_rate": 9.351803119934318e-06,
2437
+ "loss": 1.1474,
2438
+ "step": 4050
2439
+ },
2440
+ {
2441
+ "epoch": 0.25,
2442
+ "learning_rate": 9.34969790109682e-06,
2443
+ "loss": 1.2084,
2444
+ "step": 4060
2445
+ },
2446
+ {
2447
+ "epoch": 0.25,
2448
+ "learning_rate": 9.347592682259323e-06,
2449
+ "loss": 1.224,
2450
+ "step": 4070
2451
+ },
2452
+ {
2453
+ "epoch": 0.25,
2454
+ "learning_rate": 9.345487463421824e-06,
2455
+ "loss": 1.206,
2456
+ "step": 4080
2457
+ },
2458
+ {
2459
+ "epoch": 0.25,
2460
+ "learning_rate": 9.343382244584326e-06,
2461
+ "loss": 1.2225,
2462
+ "step": 4090
2463
+ },
2464
+ {
2465
+ "epoch": 0.25,
2466
+ "learning_rate": 9.341277025746827e-06,
2467
+ "loss": 1.2189,
2468
+ "step": 4100
2469
+ },
2470
+ {
2471
+ "epoch": 0.25,
2472
+ "learning_rate": 9.339171806909328e-06,
2473
+ "loss": 1.25,
2474
+ "step": 4110
2475
+ },
2476
+ {
2477
+ "epoch": 0.25,
2478
+ "learning_rate": 9.33706658807183e-06,
2479
+ "loss": 1.251,
2480
+ "step": 4120
2481
+ },
2482
+ {
2483
+ "epoch": 0.26,
2484
+ "learning_rate": 9.334961369234333e-06,
2485
+ "loss": 1.2048,
2486
+ "step": 4130
2487
+ },
2488
+ {
2489
+ "epoch": 0.26,
2490
+ "learning_rate": 9.332856150396834e-06,
2491
+ "loss": 1.2369,
2492
+ "step": 4140
2493
+ },
2494
+ {
2495
+ "epoch": 0.26,
2496
+ "learning_rate": 9.330750931559336e-06,
2497
+ "loss": 1.2427,
2498
+ "step": 4150
2499
+ },
2500
+ {
2501
+ "epoch": 0.26,
2502
+ "learning_rate": 9.328645712721837e-06,
2503
+ "loss": 1.2873,
2504
+ "step": 4160
2505
+ },
2506
+ {
2507
+ "epoch": 0.26,
2508
+ "learning_rate": 9.32654049388434e-06,
2509
+ "loss": 1.1579,
2510
+ "step": 4170
2511
+ },
2512
+ {
2513
+ "epoch": 0.26,
2514
+ "learning_rate": 9.324435275046842e-06,
2515
+ "loss": 1.2025,
2516
+ "step": 4180
2517
+ },
2518
+ {
2519
+ "epoch": 0.26,
2520
+ "learning_rate": 9.322330056209343e-06,
2521
+ "loss": 1.209,
2522
+ "step": 4190
2523
+ },
2524
+ {
2525
+ "epoch": 0.26,
2526
+ "learning_rate": 9.320224837371845e-06,
2527
+ "loss": 1.2015,
2528
+ "step": 4200
2529
+ },
2530
+ {
2531
+ "epoch": 0.26,
2532
+ "learning_rate": 9.318119618534348e-06,
2533
+ "loss": 1.2509,
2534
+ "step": 4210
2535
+ },
2536
+ {
2537
+ "epoch": 0.26,
2538
+ "learning_rate": 9.316014399696849e-06,
2539
+ "loss": 1.2696,
2540
+ "step": 4220
2541
+ },
2542
+ {
2543
+ "epoch": 0.26,
2544
+ "learning_rate": 9.313909180859351e-06,
2545
+ "loss": 1.2281,
2546
+ "step": 4230
2547
+ },
2548
+ {
2549
+ "epoch": 0.26,
2550
+ "learning_rate": 9.311803962021852e-06,
2551
+ "loss": 1.2089,
2552
+ "step": 4240
2553
+ },
2554
+ {
2555
+ "epoch": 0.26,
2556
+ "learning_rate": 9.309698743184354e-06,
2557
+ "loss": 1.2831,
2558
+ "step": 4250
2559
+ },
2560
+ {
2561
+ "epoch": 0.26,
2562
+ "learning_rate": 9.307593524346857e-06,
2563
+ "loss": 1.2757,
2564
+ "step": 4260
2565
+ },
2566
+ {
2567
+ "epoch": 0.26,
2568
+ "learning_rate": 9.305488305509358e-06,
2569
+ "loss": 1.2267,
2570
+ "step": 4270
2571
+ },
2572
+ {
2573
+ "epoch": 0.26,
2574
+ "learning_rate": 9.30338308667186e-06,
2575
+ "loss": 1.1837,
2576
+ "step": 4280
2577
+ },
2578
+ {
2579
+ "epoch": 0.27,
2580
+ "learning_rate": 9.301277867834361e-06,
2581
+ "loss": 1.2245,
2582
+ "step": 4290
2583
+ },
2584
+ {
2585
+ "epoch": 0.27,
2586
+ "learning_rate": 9.299172648996863e-06,
2587
+ "loss": 1.157,
2588
+ "step": 4300
2589
+ },
2590
+ {
2591
+ "epoch": 0.27,
2592
+ "learning_rate": 9.297067430159366e-06,
2593
+ "loss": 1.215,
2594
+ "step": 4310
2595
+ },
2596
+ {
2597
+ "epoch": 0.27,
2598
+ "learning_rate": 9.294962211321867e-06,
2599
+ "loss": 1.2421,
2600
+ "step": 4320
2601
+ },
2602
+ {
2603
+ "epoch": 0.27,
2604
+ "learning_rate": 9.29285699248437e-06,
2605
+ "loss": 1.2581,
2606
+ "step": 4330
2607
+ },
2608
+ {
2609
+ "epoch": 0.27,
2610
+ "learning_rate": 9.290751773646872e-06,
2611
+ "loss": 1.1966,
2612
+ "step": 4340
2613
+ },
2614
+ {
2615
+ "epoch": 0.27,
2616
+ "learning_rate": 9.288646554809373e-06,
2617
+ "loss": 1.2494,
2618
+ "step": 4350
2619
+ },
2620
+ {
2621
+ "epoch": 0.27,
2622
+ "learning_rate": 9.286541335971875e-06,
2623
+ "loss": 1.1633,
2624
+ "step": 4360
2625
+ },
2626
+ {
2627
+ "epoch": 0.27,
2628
+ "learning_rate": 9.284436117134376e-06,
2629
+ "loss": 1.2258,
2630
+ "step": 4370
2631
+ },
2632
+ {
2633
+ "epoch": 0.27,
2634
+ "learning_rate": 9.282330898296878e-06,
2635
+ "loss": 1.2703,
2636
+ "step": 4380
2637
+ },
2638
+ {
2639
+ "epoch": 0.27,
2640
+ "learning_rate": 9.280225679459381e-06,
2641
+ "loss": 1.1973,
2642
+ "step": 4390
2643
+ },
2644
+ {
2645
+ "epoch": 0.27,
2646
+ "learning_rate": 9.278120460621882e-06,
2647
+ "loss": 1.2614,
2648
+ "step": 4400
2649
+ },
2650
+ {
2651
+ "epoch": 0.27,
2652
+ "learning_rate": 9.276015241784384e-06,
2653
+ "loss": 1.243,
2654
+ "step": 4410
2655
+ },
2656
+ {
2657
+ "epoch": 0.27,
2658
+ "learning_rate": 9.273910022946887e-06,
2659
+ "loss": 1.2473,
2660
+ "step": 4420
2661
+ },
2662
+ {
2663
+ "epoch": 0.27,
2664
+ "learning_rate": 9.271804804109387e-06,
2665
+ "loss": 1.2269,
2666
+ "step": 4430
2667
+ },
2668
+ {
2669
+ "epoch": 0.27,
2670
+ "learning_rate": 9.26969958527189e-06,
2671
+ "loss": 1.2466,
2672
+ "step": 4440
2673
+ },
2674
+ {
2675
+ "epoch": 0.28,
2676
+ "learning_rate": 9.26759436643439e-06,
2677
+ "loss": 1.2362,
2678
+ "step": 4450
2679
+ },
2680
+ {
2681
+ "epoch": 0.28,
2682
+ "learning_rate": 9.265489147596893e-06,
2683
+ "loss": 1.2277,
2684
+ "step": 4460
2685
+ },
2686
+ {
2687
+ "epoch": 0.28,
2688
+ "learning_rate": 9.263383928759396e-06,
2689
+ "loss": 1.1939,
2690
+ "step": 4470
2691
+ },
2692
+ {
2693
+ "epoch": 0.28,
2694
+ "learning_rate": 9.261278709921897e-06,
2695
+ "loss": 1.2013,
2696
+ "step": 4480
2697
+ },
2698
+ {
2699
+ "epoch": 0.28,
2700
+ "learning_rate": 9.259173491084399e-06,
2701
+ "loss": 1.2057,
2702
+ "step": 4490
2703
+ },
2704
+ {
2705
+ "epoch": 0.28,
2706
+ "learning_rate": 9.257068272246902e-06,
2707
+ "loss": 1.2276,
2708
+ "step": 4500
2709
+ },
2710
+ {
2711
+ "epoch": 0.28,
2712
+ "learning_rate": 9.254963053409402e-06,
2713
+ "loss": 1.2029,
2714
+ "step": 4510
2715
+ },
2716
+ {
2717
+ "epoch": 0.28,
2718
+ "learning_rate": 9.252857834571905e-06,
2719
+ "loss": 1.2285,
2720
+ "step": 4520
2721
+ },
2722
+ {
2723
+ "epoch": 0.28,
2724
+ "learning_rate": 9.250752615734406e-06,
2725
+ "loss": 1.2078,
2726
+ "step": 4530
2727
+ },
2728
+ {
2729
+ "epoch": 0.28,
2730
+ "learning_rate": 9.248647396896908e-06,
2731
+ "loss": 1.2317,
2732
+ "step": 4540
2733
+ },
2734
+ {
2735
+ "epoch": 0.28,
2736
+ "learning_rate": 9.24654217805941e-06,
2737
+ "loss": 1.2266,
2738
+ "step": 4550
2739
+ },
2740
+ {
2741
+ "epoch": 0.28,
2742
+ "learning_rate": 9.244436959221911e-06,
2743
+ "loss": 1.212,
2744
+ "step": 4560
2745
+ },
2746
+ {
2747
+ "epoch": 0.28,
2748
+ "learning_rate": 9.242331740384414e-06,
2749
+ "loss": 1.1849,
2750
+ "step": 4570
2751
+ },
2752
+ {
2753
+ "epoch": 0.28,
2754
+ "learning_rate": 9.240226521546915e-06,
2755
+ "loss": 1.2238,
2756
+ "step": 4580
2757
+ },
2758
+ {
2759
+ "epoch": 0.28,
2760
+ "learning_rate": 9.238121302709417e-06,
2761
+ "loss": 1.221,
2762
+ "step": 4590
2763
+ },
2764
+ {
2765
+ "epoch": 0.28,
2766
+ "learning_rate": 9.23601608387192e-06,
2767
+ "loss": 1.2356,
2768
+ "step": 4600
2769
+ },
2770
+ {
2771
+ "epoch": 0.29,
2772
+ "learning_rate": 9.23391086503442e-06,
2773
+ "loss": 1.2287,
2774
+ "step": 4610
2775
+ },
2776
+ {
2777
+ "epoch": 0.29,
2778
+ "learning_rate": 9.231805646196923e-06,
2779
+ "loss": 1.2226,
2780
+ "step": 4620
2781
+ },
2782
+ {
2783
+ "epoch": 0.29,
2784
+ "learning_rate": 9.229700427359426e-06,
2785
+ "loss": 1.159,
2786
+ "step": 4630
2787
+ },
2788
+ {
2789
+ "epoch": 0.29,
2790
+ "learning_rate": 9.227595208521926e-06,
2791
+ "loss": 1.2239,
2792
+ "step": 4640
2793
+ },
2794
+ {
2795
+ "epoch": 0.29,
2796
+ "learning_rate": 9.225489989684429e-06,
2797
+ "loss": 1.2547,
2798
+ "step": 4650
2799
+ },
2800
+ {
2801
+ "epoch": 0.29,
2802
+ "learning_rate": 9.22338477084693e-06,
2803
+ "loss": 1.1689,
2804
+ "step": 4660
2805
+ },
2806
+ {
2807
+ "epoch": 0.29,
2808
+ "learning_rate": 9.221279552009432e-06,
2809
+ "loss": 1.1546,
2810
+ "step": 4670
2811
+ },
2812
+ {
2813
+ "epoch": 0.29,
2814
+ "learning_rate": 9.219174333171935e-06,
2815
+ "loss": 1.2197,
2816
+ "step": 4680
2817
+ },
2818
+ {
2819
+ "epoch": 0.29,
2820
+ "learning_rate": 9.217069114334436e-06,
2821
+ "loss": 1.1826,
2822
+ "step": 4690
2823
+ },
2824
+ {
2825
+ "epoch": 0.29,
2826
+ "learning_rate": 9.214963895496938e-06,
2827
+ "loss": 1.2543,
2828
+ "step": 4700
2829
+ },
2830
+ {
2831
+ "epoch": 0.29,
2832
+ "learning_rate": 9.21285867665944e-06,
2833
+ "loss": 1.1747,
2834
+ "step": 4710
2835
+ },
2836
+ {
2837
+ "epoch": 0.29,
2838
+ "learning_rate": 9.210753457821941e-06,
2839
+ "loss": 1.2486,
2840
+ "step": 4720
2841
+ },
2842
+ {
2843
+ "epoch": 0.29,
2844
+ "learning_rate": 9.208648238984444e-06,
2845
+ "loss": 1.2506,
2846
+ "step": 4730
2847
+ },
2848
+ {
2849
+ "epoch": 0.29,
2850
+ "learning_rate": 9.206543020146945e-06,
2851
+ "loss": 1.2257,
2852
+ "step": 4740
2853
+ },
2854
+ {
2855
+ "epoch": 0.29,
2856
+ "learning_rate": 9.204437801309447e-06,
2857
+ "loss": 1.183,
2858
+ "step": 4750
2859
+ },
2860
+ {
2861
+ "epoch": 0.29,
2862
+ "learning_rate": 9.20233258247195e-06,
2863
+ "loss": 1.2092,
2864
+ "step": 4760
2865
+ },
2866
+ {
2867
+ "epoch": 0.3,
2868
+ "learning_rate": 9.20022736363445e-06,
2869
+ "loss": 1.1907,
2870
+ "step": 4770
2871
+ },
2872
+ {
2873
+ "epoch": 0.3,
2874
+ "learning_rate": 9.198122144796953e-06,
2875
+ "loss": 1.1898,
2876
+ "step": 4780
2877
+ },
2878
+ {
2879
+ "epoch": 0.3,
2880
+ "learning_rate": 9.196016925959454e-06,
2881
+ "loss": 1.1834,
2882
+ "step": 4790
2883
+ },
2884
+ {
2885
+ "epoch": 0.3,
2886
+ "learning_rate": 9.193911707121956e-06,
2887
+ "loss": 1.2316,
2888
+ "step": 4800
2889
+ },
2890
+ {
2891
+ "epoch": 0.3,
2892
+ "learning_rate": 9.191806488284459e-06,
2893
+ "loss": 1.2018,
2894
+ "step": 4810
2895
+ },
2896
+ {
2897
+ "epoch": 0.3,
2898
+ "learning_rate": 9.18970126944696e-06,
2899
+ "loss": 1.2066,
2900
+ "step": 4820
2901
+ },
2902
+ {
2903
+ "epoch": 0.3,
2904
+ "learning_rate": 9.187596050609462e-06,
2905
+ "loss": 1.182,
2906
+ "step": 4830
2907
+ },
2908
+ {
2909
+ "epoch": 0.3,
2910
+ "learning_rate": 9.185490831771965e-06,
2911
+ "loss": 1.2089,
2912
+ "step": 4840
2913
+ },
2914
+ {
2915
+ "epoch": 0.3,
2916
+ "learning_rate": 9.183385612934465e-06,
2917
+ "loss": 1.1919,
2918
+ "step": 4850
2919
+ },
2920
+ {
2921
+ "epoch": 0.3,
2922
+ "learning_rate": 9.181280394096968e-06,
2923
+ "loss": 1.2734,
2924
+ "step": 4860
2925
+ },
2926
+ {
2927
+ "epoch": 0.3,
2928
+ "learning_rate": 9.179175175259469e-06,
2929
+ "loss": 1.25,
2930
+ "step": 4870
2931
+ },
2932
+ {
2933
+ "epoch": 0.3,
2934
+ "learning_rate": 9.177069956421971e-06,
2935
+ "loss": 1.213,
2936
+ "step": 4880
2937
+ },
2938
+ {
2939
+ "epoch": 0.3,
2940
+ "learning_rate": 9.174964737584474e-06,
2941
+ "loss": 1.2126,
2942
+ "step": 4890
2943
+ },
2944
+ {
2945
+ "epoch": 0.3,
2946
+ "learning_rate": 9.172859518746974e-06,
2947
+ "loss": 1.203,
2948
+ "step": 4900
2949
+ },
2950
+ {
2951
+ "epoch": 0.3,
2952
+ "learning_rate": 9.170754299909477e-06,
2953
+ "loss": 1.2718,
2954
+ "step": 4910
2955
+ },
2956
+ {
2957
+ "epoch": 0.3,
2958
+ "learning_rate": 9.16864908107198e-06,
2959
+ "loss": 1.1845,
2960
+ "step": 4920
2961
+ },
2962
+ {
2963
+ "epoch": 0.31,
2964
+ "learning_rate": 9.16654386223448e-06,
2965
+ "loss": 1.1991,
2966
+ "step": 4930
2967
+ },
2968
+ {
2969
+ "epoch": 0.31,
2970
+ "learning_rate": 9.164438643396983e-06,
2971
+ "loss": 1.2049,
2972
+ "step": 4940
2973
+ },
2974
+ {
2975
+ "epoch": 0.31,
2976
+ "learning_rate": 9.162333424559484e-06,
2977
+ "loss": 1.2345,
2978
+ "step": 4950
2979
+ },
2980
+ {
2981
+ "epoch": 0.31,
2982
+ "learning_rate": 9.160228205721986e-06,
2983
+ "loss": 1.2284,
2984
+ "step": 4960
2985
+ },
2986
+ {
2987
+ "epoch": 0.31,
2988
+ "learning_rate": 9.158122986884489e-06,
2989
+ "loss": 1.2253,
2990
+ "step": 4970
2991
+ },
2992
+ {
2993
+ "epoch": 0.31,
2994
+ "learning_rate": 9.15601776804699e-06,
2995
+ "loss": 1.2038,
2996
+ "step": 4980
2997
+ },
2998
+ {
2999
+ "epoch": 0.31,
3000
+ "learning_rate": 9.153912549209492e-06,
3001
+ "loss": 1.2353,
3002
+ "step": 4990
3003
+ },
3004
+ {
3005
+ "epoch": 0.31,
3006
+ "learning_rate": 9.151807330371993e-06,
3007
+ "loss": 1.2029,
3008
+ "step": 5000
3009
+ },
3010
+ {
3011
+ "epoch": 0.31,
3012
+ "learning_rate": 9.149702111534495e-06,
3013
+ "loss": 1.1451,
3014
+ "step": 5010
3015
+ },
3016
+ {
3017
+ "epoch": 0.31,
3018
+ "learning_rate": 9.147596892696998e-06,
3019
+ "loss": 1.186,
3020
+ "step": 5020
3021
+ },
3022
+ {
3023
+ "epoch": 0.31,
3024
+ "learning_rate": 9.145491673859498e-06,
3025
+ "loss": 1.2406,
3026
+ "step": 5030
3027
+ },
3028
+ {
3029
+ "epoch": 0.31,
3030
+ "learning_rate": 9.143386455022e-06,
3031
+ "loss": 1.1957,
3032
+ "step": 5040
3033
+ },
3034
+ {
3035
+ "epoch": 0.31,
3036
+ "learning_rate": 9.141281236184502e-06,
3037
+ "loss": 1.19,
3038
+ "step": 5050
3039
+ },
3040
+ {
3041
+ "epoch": 0.31,
3042
+ "learning_rate": 9.139176017347004e-06,
3043
+ "loss": 1.2007,
3044
+ "step": 5060
3045
+ },
3046
+ {
3047
+ "epoch": 0.31,
3048
+ "learning_rate": 9.137070798509505e-06,
3049
+ "loss": 1.2259,
3050
+ "step": 5070
3051
+ },
3052
+ {
3053
+ "epoch": 0.31,
3054
+ "learning_rate": 9.134965579672008e-06,
3055
+ "loss": 1.2204,
3056
+ "step": 5080
3057
+ },
3058
+ {
3059
+ "epoch": 0.32,
3060
+ "learning_rate": 9.132860360834508e-06,
3061
+ "loss": 1.2021,
3062
+ "step": 5090
3063
+ },
3064
+ {
3065
+ "epoch": 0.32,
3066
+ "learning_rate": 9.13075514199701e-06,
3067
+ "loss": 1.2208,
3068
+ "step": 5100
3069
+ },
3070
+ {
3071
+ "epoch": 0.32,
3072
+ "learning_rate": 9.128649923159513e-06,
3073
+ "loss": 1.1702,
3074
+ "step": 5110
3075
+ },
3076
+ {
3077
+ "epoch": 0.32,
3078
+ "learning_rate": 9.126544704322014e-06,
3079
+ "loss": 1.2513,
3080
+ "step": 5120
3081
+ },
3082
+ {
3083
+ "epoch": 0.32,
3084
+ "learning_rate": 9.124439485484517e-06,
3085
+ "loss": 1.1855,
3086
+ "step": 5130
3087
+ },
3088
+ {
3089
+ "epoch": 0.32,
3090
+ "learning_rate": 9.122334266647017e-06,
3091
+ "loss": 1.1868,
3092
+ "step": 5140
3093
+ },
3094
+ {
3095
+ "epoch": 0.32,
3096
+ "learning_rate": 9.12022904780952e-06,
3097
+ "loss": 1.1861,
3098
+ "step": 5150
3099
+ },
3100
+ {
3101
+ "epoch": 0.32,
3102
+ "learning_rate": 9.118123828972022e-06,
3103
+ "loss": 1.184,
3104
+ "step": 5160
3105
+ },
3106
+ {
3107
+ "epoch": 0.32,
3108
+ "learning_rate": 9.116018610134523e-06,
3109
+ "loss": 1.2227,
3110
+ "step": 5170
3111
+ },
3112
+ {
3113
+ "epoch": 0.32,
3114
+ "learning_rate": 9.113913391297026e-06,
3115
+ "loss": 1.1739,
3116
+ "step": 5180
3117
+ },
3118
+ {
3119
+ "epoch": 0.32,
3120
+ "learning_rate": 9.111808172459528e-06,
3121
+ "loss": 1.1705,
3122
+ "step": 5190
3123
+ },
3124
+ {
3125
+ "epoch": 0.32,
3126
+ "learning_rate": 9.109702953622029e-06,
3127
+ "loss": 1.1681,
3128
+ "step": 5200
3129
+ },
3130
+ {
3131
+ "epoch": 0.32,
3132
+ "learning_rate": 9.107597734784532e-06,
3133
+ "loss": 1.2116,
3134
+ "step": 5210
3135
+ },
3136
+ {
3137
+ "epoch": 0.32,
3138
+ "learning_rate": 9.105492515947032e-06,
3139
+ "loss": 1.2158,
3140
+ "step": 5220
3141
+ },
3142
+ {
3143
+ "epoch": 0.32,
3144
+ "learning_rate": 9.103387297109535e-06,
3145
+ "loss": 1.1452,
3146
+ "step": 5230
3147
+ },
3148
+ {
3149
+ "epoch": 0.32,
3150
+ "learning_rate": 9.101282078272037e-06,
3151
+ "loss": 1.1695,
3152
+ "step": 5240
3153
+ },
3154
+ {
3155
+ "epoch": 0.32,
3156
+ "learning_rate": 9.099176859434538e-06,
3157
+ "loss": 1.2099,
3158
+ "step": 5250
3159
+ },
3160
+ {
3161
+ "epoch": 0.33,
3162
+ "learning_rate": 9.09707164059704e-06,
3163
+ "loss": 1.2359,
3164
+ "step": 5260
3165
+ },
3166
+ {
3167
+ "epoch": 0.33,
3168
+ "learning_rate": 9.094966421759543e-06,
3169
+ "loss": 1.207,
3170
+ "step": 5270
3171
+ },
3172
+ {
3173
+ "epoch": 0.33,
3174
+ "learning_rate": 9.092861202922044e-06,
3175
+ "loss": 1.2119,
3176
+ "step": 5280
3177
+ },
3178
+ {
3179
+ "epoch": 0.33,
3180
+ "learning_rate": 9.090755984084546e-06,
3181
+ "loss": 1.1408,
3182
+ "step": 5290
3183
+ },
3184
+ {
3185
+ "epoch": 0.33,
3186
+ "learning_rate": 9.088650765247047e-06,
3187
+ "loss": 1.1842,
3188
+ "step": 5300
3189
+ },
3190
+ {
3191
+ "epoch": 0.33,
3192
+ "learning_rate": 9.08654554640955e-06,
3193
+ "loss": 1.1786,
3194
+ "step": 5310
3195
+ },
3196
+ {
3197
+ "epoch": 0.33,
3198
+ "learning_rate": 9.084440327572052e-06,
3199
+ "loss": 1.2016,
3200
+ "step": 5320
3201
+ },
3202
+ {
3203
+ "epoch": 0.33,
3204
+ "learning_rate": 9.082335108734553e-06,
3205
+ "loss": 1.2011,
3206
+ "step": 5330
3207
+ },
3208
+ {
3209
+ "epoch": 0.33,
3210
+ "learning_rate": 9.080229889897056e-06,
3211
+ "loss": 1.235,
3212
+ "step": 5340
3213
+ },
3214
+ {
3215
+ "epoch": 0.33,
3216
+ "learning_rate": 9.078124671059556e-06,
3217
+ "loss": 1.2116,
3218
+ "step": 5350
3219
+ },
3220
+ {
3221
+ "epoch": 0.33,
3222
+ "learning_rate": 9.076019452222059e-06,
3223
+ "loss": 1.2083,
3224
+ "step": 5360
3225
+ },
3226
+ {
3227
+ "epoch": 0.33,
3228
+ "learning_rate": 9.073914233384561e-06,
3229
+ "loss": 1.2152,
3230
+ "step": 5370
3231
+ },
3232
+ {
3233
+ "epoch": 0.33,
3234
+ "learning_rate": 9.071809014547062e-06,
3235
+ "loss": 1.239,
3236
+ "step": 5380
3237
+ },
3238
+ {
3239
+ "epoch": 0.33,
3240
+ "learning_rate": 9.069703795709565e-06,
3241
+ "loss": 1.1989,
3242
+ "step": 5390
3243
+ },
3244
+ {
3245
+ "epoch": 0.33,
3246
+ "learning_rate": 9.067598576872067e-06,
3247
+ "loss": 1.2193,
3248
+ "step": 5400
3249
+ },
3250
+ {
3251
+ "epoch": 0.33,
3252
+ "learning_rate": 9.065493358034568e-06,
3253
+ "loss": 1.2051,
3254
+ "step": 5410
3255
+ },
3256
+ {
3257
+ "epoch": 0.34,
3258
+ "learning_rate": 9.06338813919707e-06,
3259
+ "loss": 1.1981,
3260
+ "step": 5420
3261
+ },
3262
+ {
3263
+ "epoch": 0.34,
3264
+ "learning_rate": 9.061282920359571e-06,
3265
+ "loss": 1.2218,
3266
+ "step": 5430
3267
+ },
3268
+ {
3269
+ "epoch": 0.34,
3270
+ "learning_rate": 9.059177701522074e-06,
3271
+ "loss": 1.2669,
3272
+ "step": 5440
3273
+ },
3274
+ {
3275
+ "epoch": 0.34,
3276
+ "learning_rate": 9.057072482684576e-06,
3277
+ "loss": 1.2127,
3278
+ "step": 5450
3279
+ },
3280
+ {
3281
+ "epoch": 0.34,
3282
+ "learning_rate": 9.054967263847077e-06,
3283
+ "loss": 1.1538,
3284
+ "step": 5460
3285
+ },
3286
+ {
3287
+ "epoch": 0.34,
3288
+ "learning_rate": 9.05286204500958e-06,
3289
+ "loss": 1.1664,
3290
+ "step": 5470
3291
+ },
3292
+ {
3293
+ "epoch": 0.34,
3294
+ "learning_rate": 9.050756826172082e-06,
3295
+ "loss": 1.1976,
3296
+ "step": 5480
3297
+ },
3298
+ {
3299
+ "epoch": 0.34,
3300
+ "learning_rate": 9.048651607334583e-06,
3301
+ "loss": 1.2223,
3302
+ "step": 5490
3303
+ },
3304
+ {
3305
+ "epoch": 0.34,
3306
+ "learning_rate": 9.046546388497085e-06,
3307
+ "loss": 1.201,
3308
+ "step": 5500
3309
+ },
3310
+ {
3311
+ "epoch": 0.34,
3312
+ "learning_rate": 9.044441169659586e-06,
3313
+ "loss": 1.2372,
3314
+ "step": 5510
3315
+ },
3316
+ {
3317
+ "epoch": 0.34,
3318
+ "learning_rate": 9.042335950822089e-06,
3319
+ "loss": 1.1758,
3320
+ "step": 5520
3321
+ },
3322
+ {
3323
+ "epoch": 0.34,
3324
+ "learning_rate": 9.040230731984591e-06,
3325
+ "loss": 1.1793,
3326
+ "step": 5530
3327
+ },
3328
+ {
3329
+ "epoch": 0.34,
3330
+ "learning_rate": 9.038125513147092e-06,
3331
+ "loss": 1.1706,
3332
+ "step": 5540
3333
+ },
3334
+ {
3335
+ "epoch": 0.34,
3336
+ "learning_rate": 9.036020294309594e-06,
3337
+ "loss": 1.2267,
3338
+ "step": 5550
3339
+ },
3340
+ {
3341
+ "epoch": 0.34,
3342
+ "learning_rate": 9.033915075472097e-06,
3343
+ "loss": 1.218,
3344
+ "step": 5560
3345
+ },
3346
+ {
3347
+ "epoch": 0.34,
3348
+ "learning_rate": 9.031809856634598e-06,
3349
+ "loss": 1.1789,
3350
+ "step": 5570
3351
+ },
3352
+ {
3353
+ "epoch": 0.35,
3354
+ "learning_rate": 9.0297046377971e-06,
3355
+ "loss": 1.2317,
3356
+ "step": 5580
3357
+ },
3358
+ {
3359
+ "epoch": 0.35,
3360
+ "learning_rate": 9.027599418959601e-06,
3361
+ "loss": 1.23,
3362
+ "step": 5590
3363
+ },
3364
+ {
3365
+ "epoch": 0.35,
3366
+ "learning_rate": 9.025494200122104e-06,
3367
+ "loss": 1.2058,
3368
+ "step": 5600
3369
+ },
3370
+ {
3371
+ "epoch": 0.35,
3372
+ "learning_rate": 9.023388981284606e-06,
3373
+ "loss": 1.276,
3374
+ "step": 5610
3375
+ },
3376
+ {
3377
+ "epoch": 0.35,
3378
+ "learning_rate": 9.021283762447107e-06,
3379
+ "loss": 1.1758,
3380
+ "step": 5620
3381
+ },
3382
+ {
3383
+ "epoch": 0.35,
3384
+ "learning_rate": 9.01917854360961e-06,
3385
+ "loss": 1.182,
3386
+ "step": 5630
3387
+ },
3388
+ {
3389
+ "epoch": 0.35,
3390
+ "learning_rate": 9.01707332477211e-06,
3391
+ "loss": 1.2027,
3392
+ "step": 5640
3393
+ },
3394
+ {
3395
+ "epoch": 0.35,
3396
+ "learning_rate": 9.014968105934613e-06,
3397
+ "loss": 1.2442,
3398
+ "step": 5650
3399
+ },
3400
+ {
3401
+ "epoch": 0.35,
3402
+ "learning_rate": 9.012862887097115e-06,
3403
+ "loss": 1.1509,
3404
+ "step": 5660
3405
+ },
3406
+ {
3407
+ "epoch": 0.35,
3408
+ "learning_rate": 9.010757668259616e-06,
3409
+ "loss": 1.2369,
3410
+ "step": 5670
3411
+ },
3412
+ {
3413
+ "epoch": 0.35,
3414
+ "learning_rate": 9.008652449422118e-06,
3415
+ "loss": 1.225,
3416
+ "step": 5680
3417
+ },
3418
+ {
3419
+ "epoch": 0.35,
3420
+ "learning_rate": 9.006547230584621e-06,
3421
+ "loss": 1.2575,
3422
+ "step": 5690
3423
+ },
3424
+ {
3425
+ "epoch": 0.35,
3426
+ "learning_rate": 9.004442011747122e-06,
3427
+ "loss": 1.1801,
3428
+ "step": 5700
3429
+ },
3430
+ {
3431
+ "epoch": 0.35,
3432
+ "learning_rate": 9.002336792909624e-06,
3433
+ "loss": 1.1817,
3434
+ "step": 5710
3435
+ },
3436
+ {
3437
+ "epoch": 0.35,
3438
+ "learning_rate": 9.000231574072125e-06,
3439
+ "loss": 1.2392,
3440
+ "step": 5720
3441
+ },
3442
+ {
3443
+ "epoch": 0.35,
3444
+ "learning_rate": 8.998126355234628e-06,
3445
+ "loss": 1.1718,
3446
+ "step": 5730
3447
+ },
3448
+ {
3449
+ "epoch": 0.36,
3450
+ "learning_rate": 8.99602113639713e-06,
3451
+ "loss": 1.2155,
3452
+ "step": 5740
3453
+ },
3454
+ {
3455
+ "epoch": 0.36,
3456
+ "learning_rate": 8.993915917559631e-06,
3457
+ "loss": 1.1976,
3458
+ "step": 5750
3459
+ },
3460
+ {
3461
+ "epoch": 0.36,
3462
+ "learning_rate": 8.991810698722133e-06,
3463
+ "loss": 1.1715,
3464
+ "step": 5760
3465
+ },
3466
+ {
3467
+ "epoch": 0.36,
3468
+ "learning_rate": 8.989705479884636e-06,
3469
+ "loss": 1.1555,
3470
+ "step": 5770
3471
+ },
3472
+ {
3473
+ "epoch": 0.36,
3474
+ "learning_rate": 8.987600261047137e-06,
3475
+ "loss": 1.2071,
3476
+ "step": 5780
3477
+ },
3478
+ {
3479
+ "epoch": 0.36,
3480
+ "learning_rate": 8.98549504220964e-06,
3481
+ "loss": 1.2062,
3482
+ "step": 5790
3483
+ },
3484
+ {
3485
+ "epoch": 0.36,
3486
+ "learning_rate": 8.98338982337214e-06,
3487
+ "loss": 1.1978,
3488
+ "step": 5800
3489
+ },
3490
+ {
3491
+ "epoch": 0.36,
3492
+ "learning_rate": 8.981284604534642e-06,
3493
+ "loss": 1.2125,
3494
+ "step": 5810
3495
+ },
3496
+ {
3497
+ "epoch": 0.36,
3498
+ "learning_rate": 8.979179385697145e-06,
3499
+ "loss": 1.1887,
3500
+ "step": 5820
3501
+ },
3502
+ {
3503
+ "epoch": 0.36,
3504
+ "learning_rate": 8.977074166859646e-06,
3505
+ "loss": 1.2384,
3506
+ "step": 5830
3507
+ },
3508
+ {
3509
+ "epoch": 0.36,
3510
+ "learning_rate": 8.974968948022148e-06,
3511
+ "loss": 1.2708,
3512
+ "step": 5840
3513
+ },
3514
+ {
3515
+ "epoch": 0.36,
3516
+ "learning_rate": 8.972863729184649e-06,
3517
+ "loss": 1.1985,
3518
+ "step": 5850
3519
+ },
3520
+ {
3521
+ "epoch": 0.36,
3522
+ "learning_rate": 8.970758510347152e-06,
3523
+ "loss": 1.2202,
3524
+ "step": 5860
3525
+ },
3526
+ {
3527
+ "epoch": 0.36,
3528
+ "learning_rate": 8.968653291509654e-06,
3529
+ "loss": 1.2281,
3530
+ "step": 5870
3531
+ },
3532
+ {
3533
+ "epoch": 0.36,
3534
+ "learning_rate": 8.966548072672155e-06,
3535
+ "loss": 1.2158,
3536
+ "step": 5880
3537
+ },
3538
+ {
3539
+ "epoch": 0.36,
3540
+ "learning_rate": 8.964442853834657e-06,
3541
+ "loss": 1.1912,
3542
+ "step": 5890
3543
+ },
3544
+ {
3545
+ "epoch": 0.37,
3546
+ "learning_rate": 8.96233763499716e-06,
3547
+ "loss": 1.1822,
3548
+ "step": 5900
3549
+ },
3550
+ {
3551
+ "epoch": 0.37,
3552
+ "learning_rate": 8.96023241615966e-06,
3553
+ "loss": 1.2181,
3554
+ "step": 5910
3555
+ },
3556
+ {
3557
+ "epoch": 0.37,
3558
+ "learning_rate": 8.958127197322163e-06,
3559
+ "loss": 1.1901,
3560
+ "step": 5920
3561
+ },
3562
+ {
3563
+ "epoch": 0.37,
3564
+ "learning_rate": 8.956021978484664e-06,
3565
+ "loss": 1.1742,
3566
+ "step": 5930
3567
+ },
3568
+ {
3569
+ "epoch": 0.37,
3570
+ "learning_rate": 8.953916759647166e-06,
3571
+ "loss": 1.1675,
3572
+ "step": 5940
3573
+ },
3574
+ {
3575
+ "epoch": 0.37,
3576
+ "learning_rate": 8.951811540809669e-06,
3577
+ "loss": 1.249,
3578
+ "step": 5950
3579
+ },
3580
+ {
3581
+ "epoch": 0.37,
3582
+ "learning_rate": 8.94970632197217e-06,
3583
+ "loss": 1.1683,
3584
+ "step": 5960
3585
+ },
3586
+ {
3587
+ "epoch": 0.37,
3588
+ "learning_rate": 8.947601103134672e-06,
3589
+ "loss": 1.1558,
3590
+ "step": 5970
3591
+ },
3592
+ {
3593
+ "epoch": 0.37,
3594
+ "learning_rate": 8.945495884297175e-06,
3595
+ "loss": 1.1685,
3596
+ "step": 5980
3597
+ },
3598
+ {
3599
+ "epoch": 0.37,
3600
+ "learning_rate": 8.943390665459674e-06,
3601
+ "loss": 1.1734,
3602
+ "step": 5990
3603
+ },
3604
+ {
3605
+ "epoch": 0.37,
3606
+ "learning_rate": 8.941285446622176e-06,
3607
+ "loss": 1.1478,
3608
+ "step": 6000
3609
+ }
3610
+ ],
3611
+ "max_steps": 48471,
3612
+ "num_train_epochs": 3,
3613
+ "total_flos": 3.0400555051175117e+19,
3614
+ "trial_name": null,
3615
+ "trial_params": null
3616
+ }
training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2e46570d1be944689c0056f202607801965879286f93708fe363bbf196a88dc0
3
+ size 5371
zero_to_fp32.py ADDED
@@ -0,0 +1,578 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+
3
+ # Copyright (c) Microsoft Corporation.
4
+ # SPDX-License-Identifier: Apache-2.0
5
+
6
+ # DeepSpeed Team
7
+
8
+ # This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
9
+ # copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
10
+ # the future. Once extracted, the weights don't require DeepSpeed and can be used in any
11
+ # application.
12
+ #
13
+ # example: python zero_to_fp32.py . pytorch_model.bin
14
+
15
+ import argparse
16
+ import torch
17
+ import glob
18
+ import math
19
+ import os
20
+ import re
21
+ from collections import OrderedDict
22
+ from dataclasses import dataclass
23
+
24
+ # while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
25
+ # DeepSpeed data structures it has to be available in the current python environment.
26
+ from deepspeed.utils import logger
27
+ from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
28
+ FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
29
+ FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
30
+
31
+
32
+ @dataclass
33
+ class zero_model_state:
34
+ buffers: dict()
35
+ param_shapes: dict()
36
+ shared_params: list
37
+ ds_version: int
38
+ frozen_param_shapes: dict()
39
+ frozen_param_fragments: dict()
40
+
41
+
42
+ debug = 0
43
+
44
+ # load to cpu
45
+ device = torch.device('cpu')
46
+
47
+
48
+ def atoi(text):
49
+ return int(text) if text.isdigit() else text
50
+
51
+
52
+ def natural_keys(text):
53
+ '''
54
+ alist.sort(key=natural_keys) sorts in human order
55
+ http://nedbatchelder.com/blog/200712/human_sorting.html
56
+ (See Toothy's implementation in the comments)
57
+ '''
58
+ return [atoi(c) for c in re.split(r'(\d+)', text)]
59
+
60
+
61
+ def get_model_state_file(checkpoint_dir, zero_stage):
62
+ if not os.path.isdir(checkpoint_dir):
63
+ raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
64
+
65
+ # there should be only one file
66
+ if zero_stage == 2:
67
+ file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
68
+ elif zero_stage == 3:
69
+ file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
70
+
71
+ if not os.path.exists(file):
72
+ raise FileNotFoundError(f"can't find model states file at '{file}'")
73
+
74
+ return file
75
+
76
+
77
+ def get_checkpoint_files(checkpoint_dir, glob_pattern):
78
+ # XXX: need to test that this simple glob rule works for multi-node setup too
79
+ ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
80
+
81
+ if len(ckpt_files) == 0:
82
+ raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
83
+
84
+ return ckpt_files
85
+
86
+
87
+ def get_optim_files(checkpoint_dir):
88
+ return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
89
+
90
+
91
+ def get_model_state_files(checkpoint_dir):
92
+ return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
93
+
94
+
95
+ def parse_model_states(files):
96
+ zero_model_states = []
97
+ for file in files:
98
+ state_dict = torch.load(file, map_location=device)
99
+
100
+ if BUFFER_NAMES not in state_dict:
101
+ raise ValueError(f"{file} is not a model state checkpoint")
102
+ buffer_names = state_dict[BUFFER_NAMES]
103
+ if debug:
104
+ print("Found buffers:", buffer_names)
105
+
106
+ # recover just the buffers while restoring them to fp32 if they were saved in fp16
107
+ buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
108
+ param_shapes = state_dict[PARAM_SHAPES]
109
+
110
+ # collect parameters that are included in param_shapes
111
+ param_names = []
112
+ for s in param_shapes:
113
+ for name in s.keys():
114
+ param_names.append(name)
115
+
116
+ # update with frozen parameters
117
+ frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
118
+ if frozen_param_shapes is not None:
119
+ if debug:
120
+ print(f"Found frozen_param_shapes: {frozen_param_shapes}")
121
+ param_names += list(frozen_param_shapes.keys())
122
+
123
+ # handle shared params
124
+ shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
125
+
126
+ ds_version = state_dict.get(DS_VERSION, None)
127
+
128
+ frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
129
+
130
+ z_model_state = zero_model_state(buffers=buffers,
131
+ param_shapes=param_shapes,
132
+ shared_params=shared_params,
133
+ ds_version=ds_version,
134
+ frozen_param_shapes=frozen_param_shapes,
135
+ frozen_param_fragments=frozen_param_fragments)
136
+ zero_model_states.append(z_model_state)
137
+
138
+ return zero_model_states
139
+
140
+
141
+ def parse_optim_states(files, ds_checkpoint_dir):
142
+
143
+ total_files = len(files)
144
+ state_dicts = []
145
+ for f in files:
146
+ state_dicts.append(torch.load(f, map_location=device))
147
+
148
+ if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
149
+ raise ValueError(f"{files[0]} is not a zero checkpoint")
150
+ zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
151
+ world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
152
+
153
+ # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
154
+ # parameters can be different from data parallelism for non-expert parameters. So we can just
155
+ # use the max of the partition_count to get the dp world_size.
156
+
157
+ if type(world_size) is list:
158
+ world_size = max(world_size)
159
+
160
+ if world_size != total_files:
161
+ raise ValueError(
162
+ f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
163
+ "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
164
+ )
165
+
166
+ # the groups are named differently in each stage
167
+ if zero_stage == 2:
168
+ fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
169
+ elif zero_stage == 3:
170
+ fp32_groups_key = FP32_FLAT_GROUPS
171
+ else:
172
+ raise ValueError(f"unknown zero stage {zero_stage}")
173
+
174
+ if zero_stage == 2:
175
+ fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
176
+ elif zero_stage == 3:
177
+ # if there is more than one param group, there will be multiple flattened tensors - one
178
+ # flattened tensor per group - for simplicity merge them into a single tensor
179
+ #
180
+ # XXX: could make the script more memory efficient for when there are multiple groups - it
181
+ # will require matching the sub-lists of param_shapes for each param group flattened tensor
182
+
183
+ fp32_flat_groups = [
184
+ torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) for i in range(len(state_dicts))
185
+ ]
186
+
187
+ return zero_stage, world_size, fp32_flat_groups
188
+
189
+
190
+ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
191
+ """
192
+ Returns fp32 state_dict reconstructed from ds checkpoint
193
+
194
+ Args:
195
+ - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
196
+
197
+ """
198
+ print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
199
+
200
+ optim_files = get_optim_files(ds_checkpoint_dir)
201
+ zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
202
+ print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
203
+
204
+ model_files = get_model_state_files(ds_checkpoint_dir)
205
+
206
+ zero_model_states = parse_model_states(model_files)
207
+ print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
208
+
209
+ if zero_stage == 2:
210
+ return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states)
211
+ elif zero_stage == 3:
212
+ return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states)
213
+
214
+
215
+ def _zero2_merge_frozen_params(state_dict, zero_model_states):
216
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
217
+ return
218
+
219
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
220
+ frozen_param_fragments = zero_model_states[0].frozen_param_fragments
221
+
222
+ if debug:
223
+ num_elem = sum(s.numel() for s in frozen_param_shapes.values())
224
+ print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
225
+
226
+ wanted_params = len(frozen_param_shapes)
227
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
228
+ avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
229
+ print(f'Frozen params: Have {avail_numel} numels to process.')
230
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
231
+
232
+ total_params = 0
233
+ total_numel = 0
234
+ for name, shape in frozen_param_shapes.items():
235
+ total_params += 1
236
+ unpartitioned_numel = shape.numel()
237
+ total_numel += unpartitioned_numel
238
+
239
+ state_dict[name] = frozen_param_fragments[name]
240
+
241
+ if debug:
242
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
243
+
244
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
245
+
246
+
247
+ def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
248
+ param_shapes = zero_model_states[0].param_shapes
249
+
250
+ # Reconstruction protocol:
251
+ #
252
+ # XXX: document this
253
+
254
+ if debug:
255
+ for i in range(world_size):
256
+ for j in range(len(fp32_flat_groups[0])):
257
+ print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
258
+
259
+ # XXX: memory usage doubles here (zero2)
260
+ num_param_groups = len(fp32_flat_groups[0])
261
+ merged_single_partition_of_fp32_groups = []
262
+ for i in range(num_param_groups):
263
+ merged_partitions = [sd[i] for sd in fp32_flat_groups]
264
+ full_single_fp32_vector = torch.cat(merged_partitions, 0)
265
+ merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
266
+ avail_numel = sum(
267
+ [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
268
+
269
+ if debug:
270
+ wanted_params = sum([len(shapes) for shapes in param_shapes])
271
+ wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
272
+ # not asserting if there is a mismatch due to possible padding
273
+ print(f"Have {avail_numel} numels to process.")
274
+ print(f"Need {wanted_numel} numels in {wanted_params} params.")
275
+
276
+ # params
277
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
278
+ # out-of-core computing solution
279
+ total_numel = 0
280
+ total_params = 0
281
+ for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
282
+ offset = 0
283
+ avail_numel = full_single_fp32_vector.numel()
284
+ for name, shape in shapes.items():
285
+
286
+ unpartitioned_numel = shape.numel()
287
+ total_numel += unpartitioned_numel
288
+ total_params += 1
289
+
290
+ if debug:
291
+ print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
292
+ state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
293
+ offset += unpartitioned_numel
294
+
295
+ # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
296
+ # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
297
+ # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
298
+ # live optimizer object, so we are checking that the numbers are within the right range
299
+ align_to = 2 * world_size
300
+
301
+ def zero2_align(x):
302
+ return align_to * math.ceil(x / align_to)
303
+
304
+ if debug:
305
+ print(f"original offset={offset}, avail_numel={avail_numel}")
306
+
307
+ offset = zero2_align(offset)
308
+ avail_numel = zero2_align(avail_numel)
309
+
310
+ if debug:
311
+ print(f"aligned offset={offset}, avail_numel={avail_numel}")
312
+
313
+ # Sanity check
314
+ if offset != avail_numel:
315
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
316
+
317
+ print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
318
+
319
+
320
+ def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states):
321
+ state_dict = OrderedDict()
322
+
323
+ # buffers
324
+ buffers = zero_model_states[0].buffers
325
+ state_dict.update(buffers)
326
+ if debug:
327
+ print(f"added {len(buffers)} buffers")
328
+
329
+ _zero2_merge_frozen_params(state_dict, zero_model_states)
330
+
331
+ _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
332
+
333
+ # recover shared parameters
334
+ for pair in zero_model_states[0].shared_params:
335
+ if pair[1] in state_dict:
336
+ state_dict[pair[0]] = state_dict[pair[1]]
337
+
338
+ return state_dict
339
+
340
+
341
+ def zero3_partitioned_param_info(unpartitioned_numel, world_size):
342
+ remainder = unpartitioned_numel % world_size
343
+ padding_numel = (world_size - remainder) if remainder else 0
344
+ partitioned_numel = math.ceil(unpartitioned_numel / world_size)
345
+ return partitioned_numel, padding_numel
346
+
347
+
348
+ def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
349
+ if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
350
+ return
351
+
352
+ if debug:
353
+ for i in range(world_size):
354
+ num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
355
+ print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
356
+
357
+ frozen_param_shapes = zero_model_states[0].frozen_param_shapes
358
+ wanted_params = len(frozen_param_shapes)
359
+ wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
360
+ avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
361
+ print(f'Frozen params: Have {avail_numel} numels to process.')
362
+ print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
363
+
364
+ total_params = 0
365
+ total_numel = 0
366
+ for name, shape in zero_model_states[0].frozen_param_shapes.items():
367
+ total_params += 1
368
+ unpartitioned_numel = shape.numel()
369
+ total_numel += unpartitioned_numel
370
+
371
+ param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
372
+ state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
373
+
374
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
375
+
376
+ if debug:
377
+ print(
378
+ f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
379
+ )
380
+
381
+ print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
382
+
383
+
384
+ def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
385
+ param_shapes = zero_model_states[0].param_shapes
386
+ avail_numel = fp32_flat_groups[0].numel() * world_size
387
+ # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
388
+ # param, re-consolidating each param, while dealing with padding if any
389
+
390
+ # merge list of dicts, preserving order
391
+ param_shapes = {k: v for d in param_shapes for k, v in d.items()}
392
+
393
+ if debug:
394
+ for i in range(world_size):
395
+ print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
396
+
397
+ wanted_params = len(param_shapes)
398
+ wanted_numel = sum(shape.numel() for shape in param_shapes.values())
399
+ # not asserting if there is a mismatch due to possible padding
400
+ avail_numel = fp32_flat_groups[0].numel() * world_size
401
+ print(f"Trainable params: Have {avail_numel} numels to process.")
402
+ print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
403
+
404
+ # params
405
+ # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
406
+ # out-of-core computing solution
407
+ offset = 0
408
+ total_numel = 0
409
+ total_params = 0
410
+ for name, shape in param_shapes.items():
411
+
412
+ unpartitioned_numel = shape.numel()
413
+ total_numel += unpartitioned_numel
414
+ total_params += 1
415
+
416
+ partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
417
+
418
+ if debug:
419
+ print(
420
+ f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
421
+ )
422
+
423
+ # XXX: memory usage doubles here
424
+ state_dict[name] = torch.cat(
425
+ tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)),
426
+ 0).narrow(0, 0, unpartitioned_numel).view(shape)
427
+ offset += partitioned_numel
428
+
429
+ offset *= world_size
430
+
431
+ # Sanity check
432
+ if offset != avail_numel:
433
+ raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
434
+
435
+ print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
436
+
437
+
438
+ def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states):
439
+ state_dict = OrderedDict()
440
+
441
+ # buffers
442
+ buffers = zero_model_states[0].buffers
443
+ state_dict.update(buffers)
444
+ if debug:
445
+ print(f"added {len(buffers)} buffers")
446
+
447
+ _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
448
+
449
+ _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
450
+
451
+ # recover shared parameters
452
+ for pair in zero_model_states[0].shared_params:
453
+ if pair[1] in state_dict:
454
+ state_dict[pair[0]] = state_dict[pair[1]]
455
+
456
+ return state_dict
457
+
458
+
459
+ def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
460
+ """
461
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
462
+ ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
463
+ via a model hub.
464
+
465
+ Args:
466
+ - ``checkpoint_dir``: path to the desired checkpoint folder
467
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
468
+
469
+ Returns:
470
+ - pytorch ``state_dict``
471
+
472
+ Note: this approach may not work if your application doesn't have sufficient free CPU memory and
473
+ you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
474
+ the checkpoint.
475
+
476
+ A typical usage might be ::
477
+
478
+ from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
479
+ # do the training and checkpoint saving
480
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
481
+ model = model.cpu() # move to cpu
482
+ model.load_state_dict(state_dict)
483
+ # submit to model hub or save the model to share with others
484
+
485
+ In this example the ``model`` will no longer be usable in the deepspeed context of the same
486
+ application. i.e. you will need to re-initialize the deepspeed engine, since
487
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
488
+
489
+ If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
490
+
491
+ """
492
+ if tag is None:
493
+ latest_path = os.path.join(checkpoint_dir, 'latest')
494
+ if os.path.isfile(latest_path):
495
+ with open(latest_path, 'r') as fd:
496
+ tag = fd.read().strip()
497
+ else:
498
+ raise ValueError(f"Unable to find 'latest' file at {latest_path}")
499
+
500
+ ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
501
+
502
+ if not os.path.isdir(ds_checkpoint_dir):
503
+ raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
504
+
505
+ return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
506
+
507
+
508
+ def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
509
+ """
510
+ Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
511
+ loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
512
+
513
+ Args:
514
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
515
+ - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
516
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
517
+ """
518
+
519
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
520
+ print(f"Saving fp32 state dict to {output_file}")
521
+ torch.save(state_dict, output_file)
522
+
523
+
524
+ def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
525
+ """
526
+ 1. Put the provided model to cpu
527
+ 2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
528
+ 3. Load it into the provided model
529
+
530
+ Args:
531
+ - ``model``: the model object to update
532
+ - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
533
+ - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
534
+
535
+ Returns:
536
+ - ``model`: modified model
537
+
538
+ Make sure you have plenty of CPU memory available before you call this function. If you don't
539
+ have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
540
+ conveniently placed for you in the checkpoint folder.
541
+
542
+ A typical usage might be ::
543
+
544
+ from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
545
+ model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
546
+ # submit to model hub or save the model to share with others
547
+
548
+ Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
549
+ of the same application. i.e. you will need to re-initialize the deepspeed engine, since
550
+ ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
551
+
552
+ """
553
+ logger.info(f"Extracting fp32 weights")
554
+ state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
555
+
556
+ logger.info(f"Overwriting model with fp32 weights")
557
+ model = model.cpu()
558
+ model.load_state_dict(state_dict, strict=False)
559
+
560
+ return model
561
+
562
+
563
+ if __name__ == "__main__":
564
+
565
+ parser = argparse.ArgumentParser()
566
+ parser.add_argument("checkpoint_dir",
567
+ type=str,
568
+ help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
569
+ parser.add_argument(
570
+ "output_file",
571
+ type=str,
572
+ help="path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)")
573
+ parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
574
+ args = parser.parse_args()
575
+
576
+ debug = args.debug
577
+
578
+ convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)