add model.

adapt_tokenizer.py

@@ -0,0 +1,41 @@
+from typing import Union
+from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
+Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
+NUM_SENTINEL_TOKENS: int = 100
+
+def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):
+    """Adds sentinel tokens and padding token (if missing).
+
+    Expands the tokenizer vocabulary to include sentinel tokens
+    used in mixture-of-denoiser tasks as well as a padding token.
+
+    All added tokens are added as special tokens. No tokens are
+    added if sentinel tokens and padding token already exist.
+    """
+    sentinels_to_add = [f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)]
+    tokenizer.add_tokens(sentinels_to_add, special_tokens=True)
+    if tokenizer.pad_token is None:
+        tokenizer.add_tokens('<pad>', special_tokens=True)
+        tokenizer.pad_token = '<pad>'
+        assert tokenizer.pad_token_id is not None
+    sentinels = ''.join([f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)])
+    _sentinel_token_ids = tokenizer(sentinels, add_special_tokens=False).input_ids
+    tokenizer.sentinel_token_ids = _sentinel_token_ids
+
+class AutoTokenizerForMOD(AutoTokenizer):
+    """AutoTokenizer + Adaptation for MOD.
+
+    A simple wrapper around AutoTokenizer to make instantiating
+    an MOD-adapted tokenizer a bit easier.
+
+    MOD-adapted tokenizers have sentinel tokens (e.g., <extra_id_0>),
+    a padding token, and a property to get the token ids of the
+    sentinel tokens.
+    """
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        """See `AutoTokenizer.from_pretrained` docstring."""
+        tokenizer = super().from_pretrained(*args, **kwargs)
+        adapt_tokenizer_for_denoising(tokenizer)
+        return tokenizer

attention.py

@@ -0,0 +1,276 @@
+"""Attention layers."""
+import math
+import warnings
+from typing import Optional
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch import nn
+from .norm import LPLayerNorm
+
+def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool):
+    if original_is_causal and num_query_tokens != num_key_tokens:
+        if num_query_tokens != 1:
+            raise NotImplementedError('MPT does not support query and key with different number of tokens, unless number of query tokens is 1.')
+        else:
+            return False
+    return original_is_causal
+
+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):
+    q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)
+    k = rearrange(key, 'b s (h d) -> b h d s', h=1 if multiquery else n_heads)
+    v = rearrange(value, 'b s (h d) -> b h s d', h=1 if multiquery else n_heads)
+    min_val = torch.finfo(q.dtype).min
+    (b, _, s_q, d) = q.shape
+    s_k = k.size(-1)
+    if softmax_scale is None:
+        softmax_scale = 1 / math.sqrt(d)
+    attn_weight = q.matmul(k) * softmax_scale
+    if attn_bias is not None:
+        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):
+            raise RuntimeError(f'attn_bias (shape: {attn_bias.shape}) is expected to broadcast to shape: {attn_weight.shape}.')
+        attn_weight = attn_weight + attn_bias
+    if key_padding_mask is not None:
+        if attn_bias is not None:
+            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.')
+        attn_weight = attn_weight.masked_fill(~key_padding_mask.view((b, 1, 1, s_k)), min_val)
+    if is_causal:
+        s = max(s_q, s_k)
+        causal_mask = attn_weight.new_ones(s, s, dtype=torch.float16)
+        causal_mask = causal_mask.tril()
+        causal_mask = causal_mask.to(torch.bool)
+        causal_mask = ~causal_mask
+        causal_mask = causal_mask[-s_q:, -s_k:]
+        attn_weight = attn_weight.masked_fill(causal_mask.view(1, 1, s_q, s_k), min_val)
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    if dropout_p:
+        attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p, training=training, inplace=True)
+    out = attn_weight.matmul(v)
+    out = rearrange(out, 'b h s d -> b s (h d)')
+    if needs_weights:
+        return (out, attn_weight)
+    return (out, None)
+
+def check_valid_inputs(*tensors, valid_dtypes=[torch.float16, torch.bfloat16]):
+    for tensor in tensors:
+        if tensor.dtype not in valid_dtypes:
+            raise TypeError(f'tensor.dtype={tensor.dtype!r} must be in valid_dtypes={valid_dtypes!r}.')
+        if not tensor.is_cuda:
+            raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
+
+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):
+    try:
+        from flash_attn import bert_padding, flash_attn_interface
+    except:
+        raise RuntimeError('Please install flash-attn==1.0.3.post0')
+    check_valid_inputs(query, key, value)
+    if attn_bias is not None:
+        raise NotImplementedError(f'attn_bias not implemented for flash attn.')
+    (batch_size, seqlen) = query.shape[:2]
+    if key_padding_mask is None:
+        key_padding_mask = torch.ones_like(key[:, :, 0], dtype=torch.bool)
+    query_padding_mask = key_padding_mask[:, -query.size(1):]
+    (query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(query, query_padding_mask)
+    query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
+    (key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(key, key_padding_mask)
+    key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=1 if multiquery else n_heads)
+    (value_unpad, _, _, _) = bert_padding.unpad_input(value, key_padding_mask)
+    value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=1 if multiquery else n_heads)
+    if multiquery:
+        key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size(-1))
+        value_unpad = value_unpad.expand(value_unpad.size(0), n_heads, value_unpad.size(-1))
+    dropout_p = dropout_p if training else 0.0
+    reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
+    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)
+    output = bert_padding.pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size, seqlen)
+    return (output, None)
+
+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):
+    try:
+        from flash_attn import flash_attn_triton
+    except:
+        raise RuntimeError('Please install flash-attn==1.0.3.post0 and triton==2.0.0.dev20221202')
+    check_valid_inputs(query, key, value)
+    if dropout_p:
+        raise NotImplementedError(f'Dropout not implemented for attn_impl: triton.')
+    if needs_weights:
+        raise NotImplementedError(f'attn_impl: triton cannot return attn weights.')
+    if key_padding_mask is not None:
+        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.')
+        (b_size, s_k) = key_padding_mask.shape[:2]
+        if attn_bias is None:
+            attn_bias = query.new_zeros(b_size, 1, 1, s_k)
+        attn_bias = attn_bias.masked_fill(~key_padding_mask.view((b_size, 1, 1, s_k)), torch.finfo(query.dtype).min)
+    query = rearrange(query, 'b s (h d) -> b s h d', h=n_heads)
+    key = rearrange(key, 'b s (h d) -> b s h d', h=1 if multiquery else n_heads)
+    value = rearrange(value, 'b s (h d) -> b s h d', h=1 if multiquery else n_heads)
+    if multiquery:
+        key = key.expand(*key.shape[:2], n_heads, key.size(-1))
+        value = value.expand(*value.shape[:2], n_heads, value.size(-1))
+    reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
+    attn_output = flash_attn_triton.flash_attn_func(query, key, value, attn_bias, reset_is_causal, softmax_scale)
+    output = attn_output.view(*attn_output.shape[:2], -1)
+    return (output, None)
+
+class MultiheadAttention(nn.Module):
+    """Multi-head self attention.
+
+    Using torch or triton attention implemetation enables user to also use
+    additive bias.
+    """
+
+    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):
+        super().__init__()
+        self.attn_impl = attn_impl
+        self.clip_qkv = clip_qkv
+        self.qk_ln = qk_ln
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.softmax_scale = softmax_scale
+        if self.softmax_scale is None:
+            self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
+        self.attn_dropout_p = attn_pdrop
+        self.Wqkv = nn.Linear(self.d_model, 3 * self.d_model, device=device)
+        fuse_splits = (d_model, 2 * d_model)
+        self.Wqkv._fused = (0, fuse_splits)
+        if self.qk_ln:
+            layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
+            self.q_ln = layernorm_class(self.d_model, device=device)
+            self.k_ln = layernorm_class(self.d_model, device=device)
+        if self.attn_impl == 'flash':
+            self.attn_fn = flash_attn_fn
+        elif self.attn_impl == 'triton':
+            self.attn_fn = triton_flash_attn_fn
+            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`.')
+        elif self.attn_impl == 'torch':
+            self.attn_fn = scaled_multihead_dot_product_attention
+            if torch.cuda.is_available():
+                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`.')
+        else:
+            raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+        self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
+        self.out_proj._is_residual = True
+
+    def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
+        qkv = self.Wqkv(x)
+        if self.clip_qkv:
+            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
+        (query, key, value) = qkv.chunk(3, dim=2)
+        key_padding_mask = attention_mask
+        if self.qk_ln:
+            dtype = query.dtype
+            query = self.q_ln(query).to(dtype)
+            key = self.k_ln(key).to(dtype)
+        if past_key_value is not None:
+            if len(past_key_value) != 0:
+                key = torch.cat([past_key_value[0], key], dim=1)
+                value = torch.cat([past_key_value[1], value], dim=1)
+            past_key_value = (key, value)
+        if attn_bias is not None:
+            attn_bias = attn_bias[:, :, -query.size(1):, -key.size(1):]
+        (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)
+        return (self.out_proj(context), attn_weights, past_key_value)
+
+class MultiQueryAttention(nn.Module):
+    """Multi-Query self attention.
+
+    Using torch or triton attention implemetation enables user to also use
+    additive bias.
+    """
+
+    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):
+        super().__init__()
+        self.attn_impl = attn_impl
+        self.clip_qkv = clip_qkv
+        self.qk_ln = qk_ln
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.softmax_scale = softmax_scale
+        if self.softmax_scale is None:
+            self.softmax_scale = 1 / math.sqrt(self.head_dim)
+        self.attn_dropout_p = attn_pdrop
+        self.Wqkv = nn.Linear(d_model, d_model + 2 * self.head_dim, device=device)
+        fuse_splits = (d_model, d_model + self.head_dim)
+        self.Wqkv._fused = (0, fuse_splits)
+        if self.qk_ln:
+            layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
+            self.q_ln = layernorm_class(d_model, device=device)
+            self.k_ln = layernorm_class(self.head_dim, device=device)
+        if self.attn_impl == 'flash':
+            self.attn_fn = flash_attn_fn
+        elif self.attn_impl == 'triton':
+            self.attn_fn = triton_flash_attn_fn
+            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`.')
+        elif self.attn_impl == 'torch':
+            self.attn_fn = scaled_multihead_dot_product_attention
+            if torch.cuda.is_available():
+                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`.')
+        else:
+            raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+        self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
+        self.out_proj._is_residual = True
+
+    def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
+        qkv = self.Wqkv(x)
+        if self.clip_qkv:
+            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
+        (query, key, value) = qkv.split([self.d_model, self.head_dim, self.head_dim], dim=2)
+        key_padding_mask = attention_mask
+        if self.qk_ln:
+            dtype = query.dtype
+            query = self.q_ln(query).to(dtype)
+            key = self.k_ln(key).to(dtype)
+        if past_key_value is not None:
+            if len(past_key_value) != 0:
+                key = torch.cat([past_key_value[0], key], dim=1)
+                value = torch.cat([past_key_value[1], value], dim=1)
+            past_key_value = (key, value)
+        if attn_bias is not None:
+            attn_bias = attn_bias[:, :, -query.size(1):, -key.size(1):]
+        (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)
+        return (self.out_proj(context), attn_weights, past_key_value)
+
+def attn_bias_shape(attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_sequence_id):
+    if attn_impl == 'flash':
+        return None
+    elif attn_impl in ['torch', 'triton']:
+        if alibi:
+            if (prefix_lm or not causal) or use_sequence_id:
+                return (1, n_heads, seq_len, seq_len)
+            return (1, n_heads, 1, seq_len)
+        elif prefix_lm or use_sequence_id:
+            return (1, 1, seq_len, seq_len)
+        return None
+    else:
+        raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+
+def build_attn_bias(attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=False, alibi_bias_max=8):
+    if attn_impl == 'flash':
+        return None
+    elif attn_impl in ['torch', 'triton']:
+        if alibi:
+            (device, dtype) = (attn_bias.device, attn_bias.dtype)
+            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))
+        return attn_bias
+    else:
+        raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+
+def gen_slopes(n_heads, alibi_bias_max=8, device=None):
+    _n_heads = 2 ** math.ceil(math.log2(n_heads))
+    m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)
+    m = m.mul(alibi_bias_max / _n_heads)
+    slopes = 1.0 / torch.pow(2, m)
+    if _n_heads != n_heads:
+        slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
+    return slopes.view(1, n_heads, 1, 1)
+
+def build_alibi_bias(n_heads, seq_len, full=False, alibi_bias_max=8, device=None, dtype=None):
+    alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, 1, seq_len)
+    if full:
+        alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, seq_len, 1)
+        alibi_bias = alibi_bias.abs().mul(-1)
+    slopes = gen_slopes(n_heads, alibi_bias_max, device=device)
+    alibi_bias = alibi_bias * slopes
+    return alibi_bias.to(dtype=dtype)
+ATTN_CLASS_REGISTRY = {'multihead_attention': MultiheadAttention, 'multiquery_attention': MultiQueryAttention}

blocks.py

@@ -0,0 +1,41 @@
+"""GPT Blocks used for the GPT Model."""
+from typing import Dict, Optional, Tuple
+import torch
+import torch.nn as nn
+from .attention import ATTN_CLASS_REGISTRY
+from .norm import NORM_CLASS_REGISTRY
+
+class MPTMLP(nn.Module):
+
+    def __init__(self, d_model: int, expansion_ratio: int, device: Optional[str]=None):
+        super().__init__()
+        self.up_proj = nn.Linear(d_model, expansion_ratio * d_model, device=device)
+        self.act = nn.GELU(approximate='none')
+        self.down_proj = nn.Linear(expansion_ratio * d_model, d_model, device=device)
+        self.down_proj._is_residual = True
+
+    def forward(self, x):
+        return self.down_proj(self.act(self.up_proj(x)))
+
+class MPTBlock(nn.Module):
+
+    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):
+        del kwargs
+        super().__init__()
+        norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
+        attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]
+        self.norm_1 = norm_class(d_model, device=device)
+        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)
+        self.norm_2 = norm_class(d_model, device=device)
+        self.ffn = MPTMLP(d_model=d_model, expansion_ratio=expansion_ratio, device=device)
+        self.resid_attn_dropout = nn.Dropout(resid_pdrop)
+        self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
+
+    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]]]:
+        a = self.norm_1(x)
+        (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)
+        x = x + self.resid_attn_dropout(b)
+        m = self.norm_2(x)
+        n = self.ffn(m)
+        x = x + self.resid_ffn_dropout(n)
+        return (x, past_key_value)

config.json

@@ -0,0 +1,53 @@
+{
+  "_name_or_path": "mosaicml/mpt-7b",
+  "architectures": [
+    "MPTForCausalLM"
+  ],
+  "attn_config": {
+    "alibi": true,
+    "alibi_bias_max": 8,
+    "attn_impl": "torch",
+    "attn_pdrop": 0,
+    "attn_type": "multihead_attention",
+    "attn_uses_sequence_id": false,
+    "clip_qkv": null,
+    "prefix_lm": false,
+    "qk_ln": false,
+    "softmax_scale": null
+  },
+  "auto_map": {
+    "AutoConfig": "mosaicml/mpt-7b--configuration_mpt.MPTConfig",
+    "AutoModelForCausalLM": "mosaicml/mpt-7b--modeling_mpt.MPTForCausalLM"
+  },
+  "d_model": 4096,
+  "emb_pdrop": 0,
+  "embedding_fraction": 1.0,
+  "expansion_ratio": 4,
+  "init_config": {
+    "emb_init_std": null,
+    "emb_init_uniform_lim": null,
+    "fan_mode": "fan_in",
+    "init_div_is_residual": true,
+    "init_gain": 0,
+    "init_nonlinearity": "relu",
+    "init_std": 0.02,
+    "name": "kaiming_normal_",
+    "verbose": 0
+  },
+  "init_device": "cpu",
+  "learned_pos_emb": true,
+  "logit_scale": null,
+  "max_seq_len": 2048,
+  "model_type": "mpt",
+  "n_heads": 32,
+  "n_layers": 32,
+  "no_bias": true,
+  "norm_type": "low_precision_layernorm",
+  "resid_pdrop": 0,
+  "tokenizer_name": "EleutherAI/gpt-neox-20b",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.30.2",
+  "use_cache": false,
+  "verbose": 0,
+  "vocab_size": 50279
+}

configuration_mpt.py

@@ -0,0 +1,118 @@
+"""A HuggingFace-style model configuration."""
+from typing import Dict, Optional, Union
+from transformers import PretrainedConfig
+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}
+init_config_defaults: Dict = {'name': 'kaiming_normal_', 'fan_mode': 'fan_in', 'init_nonlinearity': 'relu'}
+
+class MPTConfig(PretrainedConfig):
+    model_type = 'mpt'
+
+    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):
+        """The MPT configuration class.
+
+        Args:
+            d_model (int): The size of the embedding dimension of the model.
+            n_heads (int): The number of attention heads.
+            n_layers (int): The number of layers in the model.
+            expansion_ratio (int): The ratio of the up/down scale in the MLP.
+            max_seq_len (int): The maximum sequence length of the model.
+            vocab_size (int): The size of the vocabulary.
+            resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.
+            emb_pdrop (float): The dropout probability for the embedding layer.
+            learned_pos_emb (bool): Whether to use learned positional embeddings
+            attn_config (Dict):  A dictionary used to configure the model's attention module:
+                attn_type (str): type of attention to use. Options: multihead_attention, multiquery_attention
+                attn_pdrop (float): The dropout probability for the attention layers.
+                attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.
+                qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.
+                clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to
+                    this value.
+                softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,
+                    use the default scale of ``1/sqrt(d_keys)``.
+                prefix_lm (Optional[bool]): Whether the model should operate as a Prefix LM. This requires passing an
+                    extra `prefix_mask` argument which indicates which tokens belong to the prefix. Tokens in the prefix
+                    can attend to one another bi-directionally. Tokens outside the prefix use causal attention.
+                attn_uses_sequence_id (Optional[bool]): Whether to restrict attention to tokens that have the same sequence_id.
+                    When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates
+                    which sub-sequence each token belongs to.
+                    Defaults to ``False`` meaning any provided `sequence_id` will be ignored.
+                alibi (bool): Whether to use the alibi bias instead of position embeddings.
+                alibi_bias_max (int): The maximum value of the alibi bias.
+            init_device (str): The device to use for parameter initialization.
+            logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.
+            no_bias (bool): Whether to use bias in all layers.
+            verbose (int): The verbosity level. 0 is silent.
+            embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.
+            norm_type (str): choose type of norm to use
+            multiquery_attention (bool): Whether to use multiquery attention implementation.
+            use_cache (bool): Whether or not the model should return the last key/values attentions
+            init_config (Dict): A dictionary used to configure the model initialization:
+                init_config.name: The parameter initialization scheme to use. Options: 'default_', 'baseline_',
+                    'kaiming_uniform_', 'kaiming_normal_', 'neox_init_', 'small_init_', 'xavier_uniform_', or
+                    'xavier_normal_'. These mimic the parameter initialization methods in PyTorch.
+                init_div_is_residual (Union[int, float, str, bool]): Value to divide initial weights by if ``module._is_residual`` is True.
+                emb_init_std (Optional[float]): The standard deviation of the normal distribution used to initialize the embedding layer.
+                emb_init_uniform_lim (Optional[Union[Tuple[float, float], float]]): The lower and upper limits of the uniform distribution
+                    used to initialize the embedding layer. Mutually exclusive with ``emb_init_std``.
+                init_std (float): The standard deviation of the normal distribution used to initialize the model,
+                    if using the baseline_ parameter initialization scheme.
+                init_gain (float): The gain to use for parameter initialization with kaiming or xavier initialization schemes.
+                fan_mode (str): The fan mode to use for parameter initialization with kaiming initialization schemes.
+                init_nonlinearity (str): The nonlinearity to use for parameter initialization with kaiming initialization schemes.
+                ---
+                See llmfoundry.models.utils.param_init_fns.py for info on other param init config options
+        """
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.expansion_ratio = expansion_ratio
+        self.max_seq_len = max_seq_len
+        self.vocab_size = vocab_size
+        self.resid_pdrop = resid_pdrop
+        self.emb_pdrop = emb_pdrop
+        self.learned_pos_emb = learned_pos_emb
+        self.attn_config = attn_config
+        self.init_device = init_device
+        self.logit_scale = logit_scale
+        self.no_bias = no_bias
+        self.verbose = verbose
+        self.embedding_fraction = embedding_fraction
+        self.norm_type = norm_type
+        self.use_cache = use_cache
+        self.init_config = init_config
+        if 'name' in kwargs:
+            del kwargs['name']
+        if 'loss_fn' in kwargs:
+            del kwargs['loss_fn']
+        super().__init__(**kwargs)
+        self._validate_config()
+
+    def _set_config_defaults(self, config, config_defaults):
+        for (k, v) in config_defaults.items():
+            if k not in config:
+                config[k] = v
+        return config
+
+    def _validate_config(self):
+        self.attn_config = self._set_config_defaults(self.attn_config, attn_config_defaults)
+        self.init_config = self._set_config_defaults(self.init_config, init_config_defaults)
+        if self.d_model % self.n_heads != 0:
+            raise ValueError('d_model must be divisible by n_heads')
+        if any((prob < 0 or prob > 1 for prob in [self.attn_config['attn_pdrop'], self.resid_pdrop, self.emb_pdrop])):
+            raise ValueError("self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1")
+        if self.attn_config['attn_impl'] not in ['torch', 'flash', 'triton']:
+            raise ValueError(f"Unknown attn_impl={self.attn_config['attn_impl']}")
+        if self.attn_config['prefix_lm'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
+            raise NotImplementedError('prefix_lm only implemented with torch and triton attention.')
+        if self.attn_config['alibi'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
+            raise NotImplementedError('alibi only implemented with torch and triton attention.')
+        if self.attn_config['attn_uses_sequence_id'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
+            raise NotImplementedError('attn_uses_sequence_id only implemented with torch and triton attention.')
+        if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
+            raise ValueError('model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!')
+        if isinstance(self.logit_scale, str) and self.logit_scale != 'inv_sqrt_d_model':
+            raise ValueError(f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
+        if self.init_config.get('name', None) is None:
+            raise ValueError(f"self.init_config={self.init_config!r} 'name' needs to be set.")
+        if not self.learned_pos_emb and (not self.attn_config['alibi']):
+            raise ValueError(f'Positional information must be provided to the model using either learned_pos_emb or alibi.')

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 0,
+  "transformers_version": "4.30.2",
+  "use_cache": false
+}

hf_prefixlm_converter.py

@@ -0,0 +1,415 @@
+"""Converts Huggingface Causal LM to Prefix LM.
+
+Conversion does lightweight surgery on a HuggingFace
+Causal LM to convert it to a Prefix LM.
+
+Prefix LMs accepts a `bidirectional_mask` input in `forward`
+and treat the input prompt as the prefix in `generate`.
+"""
+import math
+import warnings
+from types import MethodType
+from typing import Any, Dict, List, Optional, Tuple, Union
+import torch
+from transformers.models.bloom.modeling_bloom import BaseModelOutputWithPastAndCrossAttentions, BloomForCausalLM, BloomModel, CausalLMOutputWithCrossAttentions, CrossEntropyLoss
+from transformers.models.bloom.modeling_bloom import _expand_mask as _expand_mask_bloom
+from transformers.models.bloom.modeling_bloom import _make_causal_mask as _make_causal_mask_bloom
+from transformers.models.bloom.modeling_bloom import logging
+from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
+from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoForCausalLM
+from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
+from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
+from transformers.models.opt.modeling_opt import OPTForCausalLM
+from transformers.models.opt.modeling_opt import _expand_mask as _expand_mask_opt
+from transformers.models.opt.modeling_opt import _make_causal_mask as _make_causal_mask_opt
+logger = logging.get_logger(__name__)
+_SUPPORTED_GPT_MODELS = (GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM)
+CAUSAL_GPT_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM]
+
+def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_TYPES:
+    """Converts a GPT-style Causal LM to a Prefix LM.
+
+    Supported HuggingFace model classes:
+        - `GPT2LMHeadModel`
+        - `GPTNeoForCausalLM`
+        - `GPTNeoXForCausalLM`
+        - `GPTJForCausalLM`
+
+    See `convert_hf_causal_lm_to_prefix_lm` for more details.
+    """
+    if hasattr(model, '_prefix_lm_converted'):
+        return model
+    assert isinstance(model, _SUPPORTED_GPT_MODELS)
+    assert model.config.add_cross_attention == False, 'Only supports GPT-style decoder-only models'
+
+    def _get_attn_modules(model: CAUSAL_GPT_TYPES) -> List[torch.nn.Module]:
+        """Helper that gets a list of the model's attention modules.
+
+        Each module has a `bias` buffer used for causal masking. The Prefix LM
+        conversion adds logic to dynamically manipulate these biases to support
+        Prefix LM attention masking.
+        """
+        attn_modules = []
+        if isinstance(model, GPTNeoXForCausalLM):
+            blocks = model.gpt_neox.layers
+        else:
+            blocks = model.transformer.h
+        for block in blocks:
+            if isinstance(model, GPTNeoForCausalLM):
+                if block.attn.attention_type != 'global':
+                    continue
+                attn_module = block.attn.attention
+            elif isinstance(model, GPTNeoXForCausalLM):
+                attn_module = block.attention
+            else:
+                attn_module = block.attn
+            attn_modules.append(attn_module)
+        return attn_modules
+    setattr(model, '_original_forward', getattr(model, 'forward'))
+    setattr(model, '_original_generate', getattr(model, 'generate'))
+
+    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):
+        """Wraps original forward to enable PrefixLM attention."""
+
+        def call_og_forward():
+            if isinstance(self, GPTNeoXForCausalLM):
+                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)
+            else:
+                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)
+        if bidirectional_mask is None:
+            return call_og_forward()
+        assert isinstance(bidirectional_mask, torch.Tensor)
+        attn_modules = _get_attn_modules(model)
+        (b, s) = bidirectional_mask.shape
+        max_length = attn_modules[0].bias.shape[-1]
+        if s > max_length:
+            raise ValueError(f'bidirectional_mask sequence length (={s}) exceeds the ' + f'max length allowed by the model ({max_length}).')
+        assert s <= max_length
+        if s < max_length:
+            pad = torch.zeros((int(b), int(max_length - s)), dtype=bidirectional_mask.dtype, device=bidirectional_mask.device)
+            bidirectional_mask = torch.cat([bidirectional_mask, pad], dim=1)
+        bidirectional = bidirectional_mask.unsqueeze(1).unsqueeze(1)
+        for attn_module in attn_modules:
+            attn_module.bias.data = torch.logical_or(attn_module.bias.data, bidirectional)
+        output = call_og_forward()
+        for attn_module in attn_modules:
+            attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]
+        return output
+
+    def generate(self: CAUSAL_GPT_TYPES, *args: tuple, **kwargs: Dict[str, Any]):
+        """Wraps original generate to enable PrefixLM attention."""
+        attn_modules = _get_attn_modules(model)
+        for attn_module in attn_modules:
+            attn_module.bias.data[:] = 1
+        output = self._original_generate(*args, **kwargs)
+        for attn_module in attn_modules:
+            attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]
+        return output
+    setattr(model, 'forward', MethodType(forward, model))
+    setattr(model, 'generate', MethodType(generate, model))
+    setattr(model, '_prefix_lm_converted', True)
+    return model
+
+def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCausalLM:
+    """Converts a BLOOM Causal LM to a Prefix LM.
+
+    Supported HuggingFace model classes:
+        - `BloomForCausalLM`
+
+    See `convert_hf_causal_lm_to_prefix_lm` for more details.
+    """
+    if hasattr(model, '_prefix_lm_converted'):
+        return model
+    assert isinstance(model, BloomForCausalLM)
+    assert model.config.add_cross_attention == False, 'Only supports BLOOM decoder-only models'
+
+    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:
+        combined_attention_mask = None
+        device = attention_mask.device
+        (_, src_length) = input_shape
+        if src_length > 1:
+            combined_attention_mask = _make_causal_mask_bloom(input_shape, device=device, past_key_values_length=past_key_values_length)
+            if bidirectional_mask is not None:
+                assert attention_mask.shape == bidirectional_mask.shape
+                expanded_bidirectional_mask = _expand_mask_bloom(bidirectional_mask, tgt_length=src_length)
+                combined_attention_mask = torch.logical_and(combined_attention_mask, expanded_bidirectional_mask)
+        expanded_attn_mask = _expand_mask_bloom(attention_mask, tgt_length=src_length)
+        combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask
+        return combined_attention_mask
+
+    def _build_alibi_tensor(self: BloomModel, batch_size: int, query_length: int, key_length: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+        num_heads = self.config.n_head
+        closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
+        base = torch.tensor(2 ** (-2 ** (-(math.log2(closest_power_of_2) - 3))), device=device, dtype=torch.float32)
+        powers = torch.arange(1, 1 + closest_power_of_2, device=device, dtype=torch.int32)
+        slopes = torch.pow(base, powers)
+        if closest_power_of_2 != num_heads:
+            extra_base = torch.tensor(2 ** (-2 ** (-(math.log2(2 * closest_power_of_2) - 3))), device=device, dtype=torch.float32)
+            num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
+            extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=device, dtype=torch.int32)
+            slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
+        qa = torch.arange(query_length, device=device, dtype=torch.int32).view(-1, 1)
+        ka = torch.arange(key_length, device=device, dtype=torch.int32).view(1, -1)
+        diffs = qa - ka + key_length - query_length
+        diffs = -diffs.abs()
+        alibi = slopes.view(1, num_heads, 1, 1) * diffs.view(1, 1, query_length, key_length)
+        alibi = alibi.expand(batch_size, -1, -1, -1).reshape(-1, query_length, key_length)
+        return alibi.to(dtype)
+    KeyValueT = Tuple[torch.Tensor, torch.Tensor]
+
+    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]:
+        if deprecated_arguments.pop('position_ids', False) is not False:
+            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)
+        if len(deprecated_arguments) > 0:
+            raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
+        elif input_ids is not None:
+            (batch_size, seq_length) = input_ids.shape
+        elif inputs_embeds is not None:
+            (batch_size, seq_length, _) = inputs_embeds.shape
+        else:
+            raise ValueError('You have to specify either input_ids or inputs_embeds')
+        if past_key_values is None:
+            past_key_values = tuple([None] * len(self.h))
+        head_mask = self.get_head_mask(head_mask, self.config.n_layer)
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+        hidden_states = self.word_embeddings_layernorm(inputs_embeds)
+        presents = () if use_cache else None
+        all_self_attentions = () if output_attentions else None
+        all_hidden_states = () if output_hidden_states else None
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+        if past_key_values[0] is not None:
+            tmp = past_key_values[0][0]
+            past_key_values_length = tmp.shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+        if attention_mask is None:
+            attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
+        else:
+            attention_mask = attention_mask.to(hidden_states.device)
+        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)
+        causal_mask = self._prepare_attn_mask(attention_mask, bidirectional_mask, input_shape=(batch_size, seq_length), past_key_values_length=past_key_values_length)
+        for (i, (block, layer_past)) in enumerate(zip(self.h, past_key_values)):
+            if output_hidden_states:
+                hst = (hidden_states,)
+                all_hidden_states = all_hidden_states + hst
+            if self.gradient_checkpointing and self.training:
+                if use_cache:
+                    logger.warning('`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...')
+                    use_cache = False
+
+                def create_custom_forward(module):
+
+                    def custom_forward(*inputs):
+                        return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)
+                    return custom_forward
+                outputs = torch.utils.checkpoint.checkpoint(create_custom_forward(block), hidden_states, alibi, causal_mask, head_mask[i])
+            else:
+                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)
+            hidden_states = outputs[0]
+            if use_cache is True:
+                presents = presents + (outputs[1],)
+            if output_attentions:
+                oa = (outputs[2 if use_cache else 1],)
+                all_self_attentions = all_self_attentions + oa
+        hidden_states = self.ln_f(hidden_states)
+        if output_hidden_states:
+            hst = (hidden_states,)
+            all_hidden_states = all_hidden_states + hst
+        if not return_dict:
+            return tuple((v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None))
+        return BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=hidden_states, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attentions)
+    setattr(model.transformer, '_prepare_attn_mask', MethodType(_prepare_attn_mask, model.transformer))
+    setattr(model.transformer, '_build_alibi_tensor', MethodType(_build_alibi_tensor, model.transformer))
+    setattr(model.transformer, 'forward', MethodType(forward, model.transformer))
+    KeyValueT = Tuple[torch.Tensor, torch.Tensor]
+
+    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]:
+        """Replacement forward method for BloomCausalLM."""
+        if deprecated_arguments.pop('position_ids', False) is not False:
+            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)
+        if len(deprecated_arguments) > 0:
+            raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        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)
+        hidden_states = transformer_outputs[0]
+        lm_logits = self.lm_head(hidden_states)
+        loss = None
+        if labels is not None:
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            (batch_size, seq_length, vocab_size) = shift_logits.shape
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length))
+        if not return_dict:
+            output = (lm_logits,) + transformer_outputs[1:]
+            return (loss,) + output if loss is not None else output
+        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)
+
+    def prepare_inputs_for_generation(self: BloomForCausalLM, input_ids: torch.LongTensor, past: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, **kwargs) -> dict:
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            bidirectional_mask = None
+            if past[0][0].shape[0] == input_ids.shape[0]:
+                past = self._convert_to_bloom_cache(past)
+        else:
+            bidirectional_mask = torch.ones_like(input_ids)
+        return {'input_ids': input_ids, 'past_key_values': past, 'use_cache': True, 'attention_mask': attention_mask, 'bidirectional_mask': bidirectional_mask}
+    setattr(model, 'forward', MethodType(forward, model))
+    setattr(model, 'prepare_inputs_for_generation', MethodType(prepare_inputs_for_generation, model))
+    setattr(model, '_prefix_lm_converted', True)
+    return model
+
+def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM:
+    """Converts an OPT Causal LM to a Prefix LM.
+
+    Supported HuggingFace model classes:
+        - `OPTForCausalLM`
+
+    See `convert_hf_causal_lm_to_prefix_lm` for more details.
+    """
+    if hasattr(model, '_prefix_lm_converted'):
+        return model
+    assert isinstance(model, OPTForCausalLM)
+    assert model.config.add_cross_attention == False, 'Only supports OPT decoder-only models'
+    setattr(model, '_original_forward', getattr(model, 'forward'))
+    setattr(model, '_original_generate', getattr(model, 'generate'))
+    model.model.decoder.bidirectional_mask = None
+
+    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            if self.bidirectional_mask == 'g':
+                (bsz, src_length) = input_shape
+                combined_attention_mask = torch.zeros((bsz, 1, src_length, src_length + past_key_values_length), dtype=inputs_embeds.dtype, device=inputs_embeds.device)
+            else:
+                combined_attention_mask = _make_causal_mask_opt(input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length).to(inputs_embeds.device)
+                if self.bidirectional_mask is not None:
+                    assert attention_mask.shape == self.bidirectional_mask.shape
+                    expanded_bidirectional_mask = _expand_mask_opt(self.bidirectional_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
+                    combined_attention_mask = torch.maximum(expanded_bidirectional_mask, combined_attention_mask)
+        if attention_mask is not None:
+            expanded_attn_mask = _expand_mask_opt(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
+            combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+        return combined_attention_mask
+    setattr(model.model.decoder, '_prepare_decoder_attention_mask', MethodType(_prepare_decoder_attention_mask, model.model.decoder))
+
+    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):
+
+        def call_og_forward():
+            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)
+        if bidirectional_mask is None:
+            return call_og_forward()
+        self.model.decoder.bidirectional_mask = bidirectional_mask
+        try:
+            outputs = call_og_forward()
+        except:
+            self.model.decoder.bidirectional_mask = None
+            raise
+        self.model.decoder.bidirectional_mask = None
+        return outputs
+
+    def generate(self: OPTForCausalLM, *args: tuple, **kwargs: Dict[str, Any]):
+        """Wraps original generate to enable PrefixLM-style attention."""
+        self.model.decoder.bidirectional_mask = 'g'
+        try:
+            output = self._original_generate(*args, **kwargs)
+        except:
+            self.model.decoder.bidirectional_mask = None
+            raise
+        self.model.decoder.bidirectional_mask = None
+        return output
+    setattr(model, 'forward', MethodType(forward, model))
+    setattr(model, 'generate', MethodType(generate, model))
+    setattr(model, '_prefix_lm_converted', True)
+    return model
+_SUPPORTED_HF_MODELS = _SUPPORTED_GPT_MODELS + (BloomForCausalLM, OPTForCausalLM)
+CAUSAL_LM_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM, BloomForCausalLM, OPTForCausalLM]
+
+def convert_hf_causal_lm_to_prefix_lm(model: CAUSAL_LM_TYPES) -> CAUSAL_LM_TYPES:
+    """Converts a HuggingFace Causal LM to a Prefix LM.
+
+    Supported HuggingFace model classes:
+        - `GPT2LMHeadModel`
+        - `GPTNeoForCausalLM`
+        - `GPTNeoXForCausalLM`
+        - `GPTJForCausalLM`
+        - `BloomForCausalLM`
+        - `OPTForCausalLM`
+
+    Conversion to a Prefix LM is done by modifying the `forward` method, and possibly also the
+    `generate` method and/or select underlying methods depending on the model class.
+
+    These changes preserve the model API, but add a new input to `forward`: "bidirectional_mask".
+
+    Notes on training:
+        To actually train the converted model as a Prefix LM, training batches will need to indicate
+        the prefix/target structure by including `bidirectional_mask` as part of the batch inputs.
+
+        **This is not a standard input and requires custom layers either within or after your dataloader.**
+
+        In addition to adding `bidirectional_mask` to the batch, this custom code should modify `labels`
+        such that `batch['labels'][batch['bidirectional_mask'] == 1] == -100`.
+        That is, the prefix portion of the sequence should not generate any loss. Loss should only be
+        generated by the target portion of the sequence.
+
+    Notes on `GPTNeoForCausalLM`:
+        To simplify the implementation, "global" and "local" attention layers are handled differently.
+        For "global" layers, we handle conversion as described above. For "local" layers, which use a
+        causal attention mask within a restricted local window, we do not alter the masking.
+
+    Notes on `forward` method conversion:
+        After conversion, the `forward` method will handle a new input, `bidirectional_mask`,
+        which should be a [batch_size, seq_length] byte tensor, where 1 indicates token positions
+        belonging to the prefix (prefix tokens can attend to one another bidirectionally), and
+        0 indicates token positions belonging to the target.
+
+        The new `forward` method will incorporate `bidirectional_mask` (if supplied) into the existing
+        causal mask, call the original `forward` method, and (if the causal mask is a buffer) reset
+        the causal masks before returning the result.
+
+    Notes on `generate` method conversion:
+        After conversion, the `generate` method will have the same signature but will internally
+        convert all causal masks to be purely bidirectional, call the original `generate` method, and
+        (where appropriate) reset the causal masks before returning the result.
+
+        This works thanks to the logic of the HuggingFace `generate` API, which first encodes the token
+        "prompt" passed to `generate` (which is treated as the prefix) and then sequentially generates
+        each new token. Encodings are cached as generation happens, so all prefix tokens can attend to one
+        another (as expected in a Prefix LM) and generated tokens can only attend to prefix tokens and
+        previously-generated tokens (also as expected in a Prefix LM).
+
+    To preserve the API, the original methods are renamed to `_original_forward` and
+    `_original_generate`, and replaced with new `forward` and `generate` methods that wrap
+    them, respectively. Although implementation details vary by model class.
+    """
+    if isinstance(model, _SUPPORTED_GPT_MODELS):
+        return _convert_gpt_causal_lm_to_prefix_lm(model)
+    elif isinstance(model, BloomForCausalLM):
+        return _convert_bloom_causal_lm_to_prefix_lm(model)
+    elif isinstance(model, OPTForCausalLM):
+        return _convert_opt_causal_lm_to_prefix_lm(model)
+    else:
+        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}')
+
+def add_bidirectional_mask_if_missing(batch: Dict[str, Any]):
+    """Attempts to add bidirectional_mask to batch if missing.
+
+    Raises:
+        KeyError if bidirectional_mask is missing and can't be inferred
+    """
+    if 'bidirectional_mask' not in batch:
+        if batch.get('mode', None) == 'icl_task':
+            batch['bidirectional_mask'] = batch['attention_mask'].clone()
+            for (i, continuation_indices) in enumerate(batch['continuation_indices']):
+                batch['bidirectional_mask'][i, continuation_indices] = 0
+        elif 'labels' in batch and 'attention_mask' in batch:
+            batch['bidirectional_mask'] = torch.logical_and(torch.eq(batch['attention_mask'], 1), torch.eq(batch['labels'], -100)).type_as(batch['attention_mask'])
+        else:
+            raise KeyError('No bidirectional_mask in batch and not sure how to construct one.')

meta_init_context.py

@@ -0,0 +1,94 @@
+from contextlib import contextmanager
+import torch
+import torch.nn as nn
+
+@contextmanager
+def init_empty_weights(include_buffers: bool=False):
+    """Meta initialization context manager.
+
+    A context manager under which models are initialized with all parameters
+    on the meta device, therefore creating an empty model. Useful when just
+    initializing the model would blow the available RAM.
+
+    Args:
+        include_buffers (`bool`, *optional*, defaults to `False`): Whether or
+            not to also put all buffers on the meta device while initializing.
+
+    Example:
+    ```python
+    import torch.nn as nn
+
+    # Initialize a model with 100 billions parameters in no time and without using any RAM.
+    with init_empty_weights():
+        tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])
+    ```
+
+    <Tip warning={true}>
+
+    Any model created under this context manager has no weights. As such you can't do something like
+    `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`].
+
+    </Tip>
+    """
+    with init_on_device(torch.device('meta'), include_buffers=include_buffers) as f:
+        yield f
+
+@contextmanager
+def init_on_device(device: torch.device, include_buffers: bool=False):
+    """Device initialization context manager.
+
+    A context manager under which models are initialized with all parameters
+    on the specified device.
+
+    Args:
+        device (`torch.device`): Device to initialize all parameters on.
+        include_buffers (`bool`, *optional*, defaults to `False`): Whether or
+            not to also put all buffers on the meta device while initializing.
+
+    Example:
+    ```python
+    import torch.nn as nn
+
+    with init_on_device(device=torch.device("cuda")):
+        tst = nn.Liner(100, 100)  # on `cuda` device
+    ```
+    """
+    old_register_parameter = nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = nn.Module.register_buffer
+
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
+
+    def register_empty_buffer(module, name, buffer):
+        old_register_buffer(module, name, buffer)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(device)
+    if include_buffers:
+        tensor_constructors_to_patch = {torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ['empty', 'zeros', 'ones', 'full']}
+    else:
+        tensor_constructors_to_patch = {}
+
+    def patch_tensor_constructor(fn):
+
+        def wrapper(*args, **kwargs):
+            kwargs['device'] = device
+            return fn(*args, **kwargs)
+        return wrapper
+    try:
+        nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            nn.Module.register_buffer = register_empty_buffer
+        for torch_function_name in tensor_constructors_to_patch.keys():
+            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
+        yield
+    finally:
+        nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            nn.Module.register_buffer = old_register_buffer
+        for (torch_function_name, old_torch_function) in tensor_constructors_to_patch.items():
+            setattr(torch, torch_function_name, old_torch_function)

modeling_mpt.py

@@ -0,0 +1,290 @@
+"""A simple, flexible implementation of a GPT model.
+
+Inspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
+"""
+import math
+import warnings
+from typing import List, Optional, Tuple, Union
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from .attention import attn_bias_shape, build_attn_bias
+from .blocks import MPTBlock
+from .norm import NORM_CLASS_REGISTRY
+from .configuration_mpt import MPTConfig
+from .adapt_tokenizer import AutoTokenizerForMOD, adapt_tokenizer_for_denoising
+from .hf_prefixlm_converter import add_bidirectional_mask_if_missing, convert_hf_causal_lm_to_prefix_lm
+from .meta_init_context import init_empty_weights
+from .param_init_fns import MODEL_INIT_REGISTRY, generic_param_init_fn_
+Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
+
+class MPTPreTrainedModel(PreTrainedModel):
+    config_class = MPTConfig
+    base_model_prefix = 'model'
+
+class MPTModel(MPTPreTrainedModel):
+
+    def __init__(self, config: MPTConfig):
+        config._validate_config()
+        super().__init__(config)
+        self.attn_impl = config.attn_config['attn_impl']
+        self.prefix_lm = config.attn_config['prefix_lm']
+        self.attn_uses_sequence_id = config.attn_config['attn_uses_sequence_id']
+        self.alibi = config.attn_config['alibi']
+        self.alibi_bias_max = config.attn_config['alibi_bias_max']
+        if config.norm_type.lower() not in NORM_CLASS_REGISTRY.keys():
+            norm_options = ' | '.join(NORM_CLASS_REGISTRY.keys())
+            raise NotImplementedError(f'Requested norm type ({config.norm_type}) is not implemented within this repo (Options: {norm_options}).')
+        norm_class = NORM_CLASS_REGISTRY[config.norm_type.lower()]
+        self.embedding_fraction = config.embedding_fraction
+        self.wte = nn.Embedding(config.vocab_size, config.d_model, device=config.init_device)
+        if not self.alibi:
+            self.wpe = nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
+        self.emb_drop = nn.Dropout(config.emb_pdrop)
+        self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
+        self.norm_f = norm_class(config.d_model, device=config.init_device)
+        if config.init_device != 'meta':
+            self.apply(self.param_init_fn)
+        self.is_causal = not self.prefix_lm
+        self._attn_bias_initialized = False
+        self.attn_bias = None
+        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)
+        if config.no_bias:
+            for module in self.modules():
+                if hasattr(module, 'bias') and isinstance(module.bias, nn.Parameter):
+                    if config.verbose:
+                        warnings.warn(f'Removing bias ({module.bias}) from {module}.')
+                    module.register_parameter('bias', None)
+        if config.verbose and config.verbose > 2:
+            print(self)
+        if 'verbose' not in self.config.init_config:
+            self.config.init_config['verbose'] = self.config.verbose
+        if self.config.init_config['verbose'] > 1:
+            init_fn_name = self.config.init_config['name']
+            warnings.warn(f'Using {init_fn_name} initialization.')
+
+    def get_input_embeddings(self):
+        return self.wte
+
+    def set_input_embeddings(self, value):
+        self.wte = value
+
+    @torch.no_grad()
+    def _attn_bias(self, device, dtype, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None):
+        if not self._attn_bias_initialized:
+            if self.attn_bias_shape:
+                self.attn_bias = torch.zeros(self.attn_bias_shape, device=device, dtype=dtype)
+                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)
+            self._attn_bias_initialized = True
+        if self.attn_impl == 'flash':
+            return (self.attn_bias, attention_mask)
+        if self.attn_bias is not None:
+            self.attn_bias = self.attn_bias.to(dtype=dtype, device=device)
+        attn_bias = self.attn_bias
+        if self.prefix_lm:
+            assert isinstance(attn_bias, torch.Tensor)
+            assert isinstance(prefix_mask, torch.Tensor)
+            attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)
+        if self.attn_uses_sequence_id and sequence_id is not None:
+            assert isinstance(attn_bias, torch.Tensor)
+            attn_bias = self._apply_sequence_id(attn_bias, sequence_id)
+        if attention_mask is not None:
+            s_k = attention_mask.shape[-1]
+            if attn_bias is None:
+                attn_bias = torch.zeros((1, 1, 1, s_k), device=device, dtype=dtype)
+            else:
+                attn_bias = attn_bias[:, :, :, -s_k:]
+            if prefix_mask is not None and attention_mask.shape != prefix_mask.shape:
+                raise ValueError(f'attention_mask shape={attention_mask.shape} ' + f'and prefix_mask shape={prefix_mask.shape} are not equal.')
+            min_val = torch.finfo(attn_bias.dtype).min
+            attn_bias = attn_bias.masked_fill(~attention_mask.view(-1, 1, 1, s_k), min_val)
+        return (attn_bias, None)
+
+    def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor):
+        (s_k, s_q) = attn_bias.shape[-2:]
+        if s_k != self.config.max_seq_len or s_q != self.config.max_seq_len:
+            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}.')
+        seq_len = prefix_mask.shape[-1]
+        if seq_len > self.config.max_seq_len:
+            raise ValueError(f'prefix_mask sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
+        attn_bias = attn_bias[..., :seq_len, :seq_len]
+        causal = torch.tril(torch.ones((seq_len, seq_len), dtype=torch.bool, device=prefix_mask.device)).view(1, 1, seq_len, seq_len)
+        prefix = prefix_mask.view(-1, 1, 1, seq_len)
+        cannot_attend = ~torch.logical_or(causal, prefix.bool())
+        min_val = torch.finfo(attn_bias.dtype).min
+        attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
+        return attn_bias
+
+    def _apply_sequence_id(self, attn_bias: torch.Tensor, sequence_id: torch.LongTensor):
+        seq_len = sequence_id.shape[-1]
+        if seq_len > self.config.max_seq_len:
+            raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
+        attn_bias = attn_bias[..., :seq_len, :seq_len]
+        cannot_attend = torch.logical_not(torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))).unsqueeze(1)
+        min_val = torch.finfo(attn_bias.dtype).min
+        attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
+        return attn_bias
+
+    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):
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        if attention_mask is not None:
+            attention_mask = attention_mask.bool()
+        if prefix_mask is not None:
+            prefix_mask = prefix_mask.bool()
+        if not return_dict:
+            raise NotImplementedError('return_dict False is not implemented yet for MPT')
+        if output_attentions:
+            raise NotImplementedError('output_attentions is not implemented yet for MPT')
+        if attention_mask is not None and attention_mask[:, 0].sum() != attention_mask.shape[0] and self.training:
+            raise NotImplementedError('MPT does not support training with left padding.')
+        if self.prefix_lm and prefix_mask is None:
+            raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
+        if self.training:
+            if self.attn_uses_sequence_id and sequence_id is None:
+                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.')
+            elif self.attn_uses_sequence_id is False and sequence_id is not None:
+                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.')
+        S = input_ids.size(1)
+        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}'
+        tok_emb = self.wte(input_ids)
+        if self.alibi:
+            x = tok_emb
+        else:
+            past_position = 0
+            if past_key_values is not None:
+                if len(past_key_values) != self.config.n_layers:
+                    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}).')
+                past_position = past_key_values[0][0].size(1)
+            if S + past_position > self.config.max_seq_len:
+                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}.')
+            pos = torch.arange(past_position, S + past_position, dtype=torch.long, device=input_ids.device).unsqueeze(0)
+            if attention_mask is not None:
+                pos = torch.clamp(pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[:, past_position:], min=0)
+            pos_emb = self.wpe(pos)
+            x = tok_emb + pos_emb
+        if self.embedding_fraction == 1:
+            x = self.emb_drop(x)
+        else:
+            x_shrunk = x * self.embedding_fraction + x.detach() * (1 - self.embedding_fraction)
+            assert isinstance(self.emb_drop, nn.Module)
+            x = self.emb_drop(x_shrunk)
+        (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)
+        if use_cache and past_key_values is None:
+            past_key_values = [() for _ in range(self.config.n_layers)]
+        all_hidden_states = () if output_hidden_states else None
+        for (b_idx, block) in enumerate(self.blocks):
+            if output_hidden_states:
+                assert all_hidden_states is not None
+                all_hidden_states = all_hidden_states + (x,)
+            past_key_value = past_key_values[b_idx] if past_key_values is not None else None
+            (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)
+            if past_key_values is not None:
+                past_key_values[b_idx] = past_key_value
+        x = self.norm_f(x)
+        return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values, hidden_states=all_hidden_states)
+
+    def param_init_fn(self, module):
+        init_fn_name = self.config.init_config['name']
+        MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
+
+    def fsdp_wrap_fn(self, module):
+        return isinstance(module, MPTBlock)
+
+    def activation_checkpointing_fn(self, module):
+        return isinstance(module, MPTBlock)
+
+class MPTForCausalLM(MPTPreTrainedModel):
+
+    def __init__(self, config: MPTConfig):
+        super().__init__(config)
+        if not config.tie_word_embeddings:
+            raise ValueError('MPTForCausalLM only supports tied word embeddings')
+        self.transformer = MPTModel(config)
+        self.logit_scale = None
+        if config.logit_scale is not None:
+            logit_scale = config.logit_scale
+            if isinstance(logit_scale, str):
+                if logit_scale == 'inv_sqrt_d_model':
+                    logit_scale = 1 / math.sqrt(config.d_model)
+                else:
+                    raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
+            self.logit_scale = logit_scale
+
+    def get_input_embeddings(self):
+        return self.transformer.wte
+
+    def set_input_embeddings(self, value):
+        self.transformer.wte = value
+
+    def get_output_embeddings(self):
+        return self.transformer.wte
+
+    def set_output_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.transformer = decoder
+
+    def get_decoder(self):
+        return self.transformer
+
+    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):
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        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)
+        logits = F.linear(outputs.last_hidden_state, self.transformer.wte.weight)
+        if self.logit_scale is not None:
+            if self.logit_scale == 0:
+                warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
+            logits *= self.logit_scale
+        loss = None
+        if labels is not None:
+            labels = torch.roll(labels, shifts=-1)
+            labels[:, -1] = -100
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.to(logits.device).view(-1))
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states)
+
+    def param_init_fn(self, module):
+        init_fn_name = self.config.init_config['name']
+        MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
+
+    def fsdp_wrap_fn(self, module):
+        return isinstance(module, MPTBlock)
+
+    def activation_checkpointing_fn(self, module):
+        return isinstance(module, MPTBlock)
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
+        if inputs_embeds is not None:
+            raise NotImplementedError('inputs_embeds is not implemented for MPT yet')
+        attention_mask = kwargs['attention_mask'].bool()
+        if attention_mask[:, -1].sum() != attention_mask.shape[0]:
+            raise NotImplementedError('MPT does not support generation with right padding.')
+        if self.transformer.attn_uses_sequence_id and self.training:
+            sequence_id = torch.zeros_like(input_ids[:1])
+        else:
+            sequence_id = None
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+        if self.transformer.prefix_lm:
+            prefix_mask = torch.ones_like(attention_mask)
+            if kwargs.get('use_cache') == False:
+                raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
+        else:
+            prefix_mask = None
+        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)}
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        """Used by HuggingFace generate when using beam search with kv-caching.
+
+        See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133
+        for an example in transformers.
+        """
+        reordered_past = []
+        for layer_past in past_key_values:
+            reordered_past += [tuple((past_state.index_select(0, beam_idx) for past_state in layer_past))]
+        return reordered_past

norm.py

@@ -0,0 +1,56 @@
+import torch
+
+def _cast_if_autocast_enabled(tensor):
+    if torch.is_autocast_enabled():
+        if tensor.device.type == 'cuda':
+            dtype = torch.get_autocast_gpu_dtype()
+        elif tensor.device.type == 'cpu':
+            dtype = torch.get_autocast_cpu_dtype()
+        else:
+            raise NotImplementedError()
+        return tensor.to(dtype=dtype)
+    return tensor
+
+class LPLayerNorm(torch.nn.LayerNorm):
+
+    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None):
+        super().__init__(normalized_shape=normalized_shape, eps=eps, elementwise_affine=elementwise_affine, device=device, dtype=dtype)
+
+    def forward(self, x):
+        module_device = x.device
+        downcast_x = _cast_if_autocast_enabled(x)
+        downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
+        downcast_bias = _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias
+        with torch.autocast(enabled=False, device_type=module_device.type):
+            return torch.nn.functional.layer_norm(downcast_x, self.normalized_shape, downcast_weight, downcast_bias, self.eps)
+
+def rms_norm(x, weight=None, eps=1e-05):
+    output = x / torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+    if weight is not None:
+        return output * weight
+    return output
+
+class RMSNorm(torch.nn.Module):
+
+    def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
+        super().__init__()
+        self.eps = eps
+        if weight:
+            self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))
+        else:
+            self.register_parameter('weight', None)
+
+    def forward(self, x):
+        return rms_norm(x.float(), self.weight, self.eps).to(dtype=x.dtype)
+
+class LPRMSNorm(RMSNorm):
+
+    def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
+        super().__init__(normalized_shape=normalized_shape, eps=eps, weight=weight, dtype=dtype, device=device)
+
+    def forward(self, x):
+        downcast_x = _cast_if_autocast_enabled(x)
+        downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
+        with torch.autocast(enabled=False, device_type=x.device.type):
+            return rms_norm(downcast_x, downcast_weight, self.eps).to(dtype=x.dtype)
+NORM_CLASS_REGISTRY = {'layernorm': torch.nn.LayerNorm, 'low_precision_layernorm': LPLayerNorm, 'rmsnorm': RMSNorm, 'low_precision_rmsnorm': LPRMSNorm}

param_init_fns.py

@@ -0,0 +1,181 @@
+import math
+import warnings
+from collections.abc import Sequence
+from functools import partial
+from typing import Optional, Tuple, Union
+import torch
+from torch import nn
+from .norm import NORM_CLASS_REGISTRY
+
+def torch_default_param_init_fn_(module: nn.Module, verbose: int=0, **kwargs):
+    del kwargs
+    if verbose > 1:
+        warnings.warn(f"Initializing network using module's reset_parameters attribute")
+    if hasattr(module, 'reset_parameters'):
+        module.reset_parameters()
+
+def fused_init_helper_(module: nn.Module, init_fn_):
+    _fused = getattr(module, '_fused', None)
+    if _fused is None:
+        raise RuntimeError(f'Internal logic error')
+    (dim, splits) = _fused
+    splits = (0, *splits, module.weight.size(dim))
+    for (s, e) in zip(splits[:-1], splits[1:]):
+        slice_indices = [slice(None)] * module.weight.ndim
+        slice_indices[dim] = slice(s, e)
+        init_fn_(module.weight[slice_indices])
+
+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):
+    del kwargs
+    if verbose > 1:
+        warnings.warn(f'If model has bias parameters they are initialized to 0.')
+    init_div_is_residual = init_div_is_residual
+    if init_div_is_residual is False:
+        div_is_residual = 1.0
+    elif init_div_is_residual is True:
+        div_is_residual = math.sqrt(2 * n_layers)
+    elif isinstance(init_div_is_residual, float) or isinstance(init_div_is_residual, int):
+        div_is_residual = init_div_is_residual
+    elif isinstance(init_div_is_residual, str) and init_div_is_residual.isnumeric():
+        div_is_residual = float(init_div_is_residual)
+    else:
+        div_is_residual = 1.0
+        raise ValueError(f'Expected init_div_is_residual to be boolean or numeric, got {init_div_is_residual}')
+    if init_div_is_residual is not False:
+        if verbose > 1:
+            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.')
+    if isinstance(module, nn.Linear):
+        if hasattr(module, '_fused'):
+            fused_init_helper_(module, init_fn_)
+        else:
+            init_fn_(module.weight)
+        if module.bias is not None:
+            torch.nn.init.zeros_(module.bias)
+        if init_div_is_residual is not False and getattr(module, '_is_residual', False):
+            with torch.no_grad():
+                module.weight.div_(div_is_residual)
+    elif isinstance(module, nn.Embedding):
+        if emb_init_std is not None:
+            std = emb_init_std
+            if std == 0:
+                warnings.warn(f'Embedding layer initialized to 0.')
+            emb_init_fn_ = partial(torch.nn.init.normal_, mean=0.0, std=std)
+            if verbose > 1:
+                warnings.warn(f'Embedding layer initialized using normal distribution with mean=0 and std={std!r}.')
+        elif emb_init_uniform_lim is not None:
+            lim = emb_init_uniform_lim
+            if isinstance(lim, Sequence):
+                if len(lim) > 2:
+                    raise ValueError(f'Uniform init requires a min and a max limit. User input: {lim}.')
+                if lim[0] == lim[1]:
+                    warnings.warn(f'Embedding layer initialized to {lim[0]}.')
+            else:
+                if lim == 0:
+                    warnings.warn(f'Embedding layer initialized to 0.')
+                lim = [-lim, lim]
+            (a, b) = lim
+            emb_init_fn_ = partial(torch.nn.init.uniform_, a=a, b=b)
+            if verbose > 1:
+                warnings.warn(f'Embedding layer initialized using uniform distribution in range {lim}.')
+        else:
+            emb_init_fn_ = init_fn_
+        emb_init_fn_(module.weight)
+    elif isinstance(module, tuple(set(NORM_CLASS_REGISTRY.values()))):
+        if verbose > 1:
+            warnings.warn(f'Norm weights are set to 1. If norm layer has a bias it is initialized to 0.')
+        if hasattr(module, 'weight') and module.weight is not None:
+            torch.nn.init.ones_(module.weight)
+        if hasattr(module, 'bias') and module.bias is not None:
+            torch.nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.MultiheadAttention):
+        if module._qkv_same_embed_dim:
+            assert module.in_proj_weight is not None
+            assert module.q_proj_weight is None and module.k_proj_weight is None and (module.v_proj_weight is None)
+            assert d_model is not None
+            _d = d_model
+            splits = (0, _d, 2 * _d, 3 * _d)
+            for (s, e) in zip(splits[:-1], splits[1:]):
+                init_fn_(module.in_proj_weight[s:e])
+        else:
+            assert module.q_proj_weight is not None and module.k_proj_weight is not None and (module.v_proj_weight is not None)
+            assert module.in_proj_weight is None
+            init_fn_(module.q_proj_weight)
+            init_fn_(module.k_proj_weight)
+            init_fn_(module.v_proj_weight)
+        if module.in_proj_bias is not None:
+            torch.nn.init.zeros_(module.in_proj_bias)
+        if module.bias_k is not None:
+            torch.nn.init.zeros_(module.bias_k)
+        if module.bias_v is not None:
+            torch.nn.init.zeros_(module.bias_v)
+        init_fn_(module.out_proj.weight)
+        if init_div_is_residual is not False and getattr(module.out_proj, '_is_residual', False):
+            with torch.no_grad():
+                module.out_proj.weight.div_(div_is_residual)
+        if module.out_proj.bias is not None:
+            torch.nn.init.zeros_(module.out_proj.bias)
+    else:
+        for _ in module.parameters(recurse=False):
+            raise NotImplementedError(f'{module.__class__.__name__} parameters are not initialized by param_init_fn.')
+
+def _normal_init_(std, mean=0.0):
+    return partial(torch.nn.init.normal_, mean=mean, std=std)
+
+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):
+    del kwargs
+    init_fn_ = _normal_init_(std=std)
+    if verbose > 1:
+        warnings.warn(f'Using torch.nn.init.normal_ init fn mean=0.0, std={std}')
+    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)
+
+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):
+    del kwargs
+    if init_std is None:
+        raise ValueError("You must set model.init_config['init_std'] to a float value to use the default initialization scheme.")
+    _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)
+
+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):
+    del kwargs
+    std = math.sqrt(2 / (5 * d_model))
+    _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)
+
+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):
+    """From section 2.3.1 of GPT-NeoX-20B:
+
+    An Open-Source AutoregressiveLanguage Model — Black et. al. (2022)
+    see https://github.com/EleutherAI/gpt-neox/blob/9610391ab319403cef079b438edd016a2443af54/megatron/model/init_functions.py#L151
+    and https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/transformer.py
+    """
+    del kwargs
+    residual_div = n_layers / math.sqrt(10)
+    if verbose > 1:
+        warnings.warn(f'setting init_div_is_residual to {residual_div}')
+    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)
+
+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):
+    del kwargs
+    if verbose > 1:
+        warnings.warn(f'Using nn.init.kaiming_uniform_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}')
+    kaiming_uniform_ = partial(nn.init.kaiming_uniform_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
+    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)
+
+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):
+    del kwargs
+    if verbose > 1:
+        warnings.warn(f'Using nn.init.kaiming_normal_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}')
+    kaiming_normal_ = partial(torch.nn.init.kaiming_normal_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
+    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)
+
+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):
+    del kwargs
+    xavier_uniform_ = partial(torch.nn.init.xavier_uniform_, gain=init_gain)
+    if verbose > 1:
+        warnings.warn(f'Using torch.nn.init.xavier_uniform_ init fn with parameters: ' + f'gain={init_gain}')
+    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)
+
+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):
+    xavier_normal_ = partial(torch.nn.init.xavier_normal_, gain=init_gain)
+    if verbose > 1:
+        warnings.warn(f'Using torch.nn.init.xavier_normal_ init fn with parameters: ' + f'gain={init_gain}')
+    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)
+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

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:119a68ab92d3200f185ed35e1ae1758148af5487b90354f876652a555cc9ec6a
+size 26594704201

special_tokens_map.json

@@ -0,0 +1,9 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

test_itex_warm.py

@@ -0,0 +1,128 @@
+from intel_extension_for_transformers.backends.neural_engine.compile import compile
+# graph = compile("/home/ubuntu/mengfeil/IR/llama-7b-bf16/")
+# graph = compile("/home/ubuntu/mengfeil/IR/llama-7b-int8/")
+# graph = compile("./llama-7b-hf-conv-itrex-bf16")
+graph = compile("/home/ubuntu/neuralchat_server/frameworks.ai.nlp-toolkit.intel-nlp-toolkit/examples/huggingface/pytorch/text-generation/deployment/bf16ir")
+from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaTokenizer, AutoModel, AutoConfig
+import torch
+import numpy as np
+import torch.nn.functional as F
+import time
+
+model_name = "../mengfeil/llama-7b"
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+config = AutoConfig.from_pretrained(model_name)
+
+#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"
+# 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."
+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:"
+
+print(prompt)
+init_input_ids = tokenizer(prompt, return_tensors="pt").input_ids[0]
+
+
+input_ids = init_input_ids.clone()
+
+attention_mask = torch.ones(len(input_ids)+1)
+attention_mask[0] = 0
+position_ids = torch.arange(len(input_ids))
+past_key_value = tuple([(torch.zeros([1,32,1,128]), torch.zeros([1,32,1,128])) for i in range(32)])
+
+input_ids = input_ids.unsqueeze(0)
+attention_mask = attention_mask.unsqueeze(0)
+position_ids = position_ids.unsqueeze(0)
+all_input_ids = input_ids.clone()
+
+
+# input_ids_1 = input_ids.cpu().numpy().astype(np.int32)
+# attention_mask_1 = attention_mask.cpu().numpy().astype(np.int32)
+# past_k_v = [past_key_value[i][j].cpu().numpy() for i in range(32) for j in range(2)]
+
+max_new_tokens = 32
+temperature = 0.9
+
+
+def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
+    """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+        Args:
+            logits: logits distribution shape (vocabulary size)
+            top_k >0: keep only top k tokens with highest probability (top-k filtering).
+            top_p >0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+                Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+    """
+    assert logits.dim() == 1  # batch size 1 for now - could be updated for more but the code would be less clear
+    top_k = min(top_k, logits.size(-1))  # Safety check
+    if top_k > 0:
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p > 0.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        indices_to_remove = sorted_indices[sorted_indices_to_remove]
+        logits[indices_to_remove] = filter_value
+    return logits
+
+
+# start
+total_time = 0.0
+num_iter = 10
+num_warmup = 3
+
+for i in range(num_iter):
+
+    input_ids_1 = input_ids.cpu().numpy().astype(np.int32)
+    attention_mask_1 = attention_mask.cpu().numpy().astype(np.int32)
+    past_k_v = [past_key_value[i][j].cpu().numpy() for i in range(32) for j in range(2)]
+    output_ids = list(init_input_ids)
+    
+    tic = time.time()
+
+    for step in range(max_new_tokens):
+        a = time.time()
+        predictions = graph.inference([input_ids_1, attention_mask_1] + past_k_v)
+        # predictions = graph.inference([input_ids_1] + past_k_v + [attention_mask_1])
+        print(time.time() - a)
+
+        outs = []
+        for key in predictions:
+            outs.append(predictions[key])
+
+        logits = outs[0]
+        past_k_v = outs[1:]
+        logits = torch.from_numpy(logits)
+
+        next_token_logits = logits[:, -1, :]
+        probs = torch.nn.functional.softmax(next_token_logits, dim=-1)
+        token = int(torch.argmax(probs, dim=-1))
+
+        """
+        last_token_logits = logits[0][-1]
+        logits = logits[0, -1, :] / temperature
+        filtered_logits = top_k_top_p_filtering(logits, top_k=10, top_p=0.8)
+        probabilities = F.softmax(filtered_logits, dim=-1)
+        token = int(torch.multinomial(probabilities, 1))
+        """
+
+        output_ids.append(token)
+        input_ids_1 = torch.tensor([[token]])
+        attention_mask_1 = torch.cat([torch.from_numpy(attention_mask_1), torch.ones([1, 1])], dim=-1)
+        input_ids_1 = input_ids_1.cpu().numpy().astype(np.int32)
+        attention_mask_1 = attention_mask_1.cpu().numpy().astype(np.int32)
+    
+    toc = time.time()
+    if i >= num_warmup:
+        total_time += (toc - tic)
+    # print(output_ids)
+    print(tokenizer.decode(output_ids, skip_special_tokens=True))
+
+print("Inference latency: %.3f ms." % (total_time / (num_iter - num_warmup) * 1000))

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "add_prefix_space": false,
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 2048,
+  "tokenizer_class": "GPTNeoXTokenizer",
+  "unk_token": "<|endoftext|>"
+}

trainer_state.json

@@ -0,0 +1,3616 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.3713417039942442,
+  "global_step": 6000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0,
+      "learning_rate": 9.278350515463919e-08,
+      "loss": 2.5166,
+      "step": 10
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 1.9587628865979384e-07,
+      "loss": 1.9059,
+      "step": 20
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 2.989690721649485e-07,
+      "loss": 1.7869,
+      "step": 30
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 4.0206185567010316e-07,
+      "loss": 1.7863,
+      "step": 40
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 5.051546391752578e-07,
+      "loss": 1.6946,
+      "step": 50
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 6.082474226804124e-07,
+      "loss": 1.7419,
+      "step": 60
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 7.11340206185567e-07,
+      "loss": 1.645,
+      "step": 70
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 8.144329896907217e-07,
+      "loss": 1.6611,
+      "step": 80
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 9.175257731958763e-07,
+      "loss": 1.5985,
+      "step": 90
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.020618556701031e-06,
+      "loss": 1.6198,
+      "step": 100
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.1237113402061856e-06,
+      "loss": 1.6259,
+      "step": 110
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.2268041237113403e-06,
+      "loss": 1.5973,
+      "step": 120
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.329896907216495e-06,
+      "loss": 1.5941,
+      "step": 130
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.4329896907216496e-06,
+      "loss": 1.5597,
+      "step": 140
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.5360824742268042e-06,
+      "loss": 1.5672,
+      "step": 150
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.639175257731959e-06,
+      "loss": 1.5372,
+      "step": 160
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.7422680412371134e-06,
+      "loss": 1.5715,
+      "step": 170
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.8453608247422682e-06,
+      "loss": 1.5389,
+      "step": 180
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.948453608247423e-06,
+      "loss": 1.5525,
+      "step": 190
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 2.0515463917525773e-06,
+      "loss": 1.5871,
+      "step": 200
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 2.1546391752577322e-06,
+      "loss": 1.5442,
+      "step": 210
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 2.2577319587628867e-06,
+      "loss": 1.5335,
+      "step": 220
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 2.3608247422680415e-06,
+      "loss": 1.5103,
+      "step": 230
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 2.463917525773196e-06,
+      "loss": 1.5016,
+      "step": 240
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 2.5670103092783504e-06,
+      "loss": 1.5101,
+      "step": 250
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 2.6701030927835053e-06,
+      "loss": 1.5323,
+      "step": 260
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 2.77319587628866e-06,
+      "loss": 1.4785,
+      "step": 270
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 2.8762886597938146e-06,
+      "loss": 1.4132,
+      "step": 280
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 2.979381443298969e-06,
+      "loss": 1.4733,
+      "step": 290
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 3.082474226804124e-06,
+      "loss": 1.4258,
+      "step": 300
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 3.1855670103092784e-06,
+      "loss": 1.4557,
+      "step": 310
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 3.2886597938144333e-06,
+      "loss": 1.44,
+      "step": 320
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 3.391752577319588e-06,
+      "loss": 1.4348,
+      "step": 330
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 3.494845360824742e-06,
+      "loss": 1.406,
+      "step": 340
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 3.597938144329897e-06,
+      "loss": 1.4239,
+      "step": 350
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 3.701030927835052e-06,
+      "loss": 1.4185,
+      "step": 360
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 3.8041237113402064e-06,
+      "loss": 1.3954,
+      "step": 370
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 3.907216494845361e-06,
+      "loss": 1.3759,
+      "step": 380
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 4.010309278350516e-06,
+      "loss": 1.3231,
+      "step": 390
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 4.11340206185567e-06,
+      "loss": 1.3059,
+      "step": 400
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 4.216494845360825e-06,
+      "loss": 1.3258,
+      "step": 410
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 4.31958762886598e-06,
+      "loss": 1.3722,
+      "step": 420
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 4.422680412371134e-06,
+      "loss": 1.2917,
+      "step": 430
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 4.525773195876289e-06,
+      "loss": 1.2893,
+      "step": 440
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 4.628865979381444e-06,
+      "loss": 1.2865,
+      "step": 450
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 4.731958762886599e-06,
+      "loss": 1.28,
+      "step": 460
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 4.835051546391753e-06,
+      "loss": 1.299,
+      "step": 470
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 4.9381443298969075e-06,
+      "loss": 1.2949,
+      "step": 480
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 5.041237113402062e-06,
+      "loss": 1.3078,
+      "step": 490
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 5.144329896907216e-06,
+      "loss": 1.2242,
+      "step": 500
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 5.247422680412372e-06,
+      "loss": 1.2714,
+      "step": 510
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 5.350515463917526e-06,
+      "loss": 1.2482,
+      "step": 520
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 5.45360824742268e-06,
+      "loss": 1.2703,
+      "step": 530
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 5.556701030927836e-06,
+      "loss": 1.2578,
+      "step": 540
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 5.65979381443299e-06,
+      "loss": 1.2989,
+      "step": 550
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 5.762886597938144e-06,
+      "loss": 1.2853,
+      "step": 560
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 5.8659793814433e-06,
+      "loss": 1.3022,
+      "step": 570
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 5.969072164948454e-06,
+      "loss": 1.2871,
+      "step": 580
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 6.0721649484536086e-06,
+      "loss": 1.2679,
+      "step": 590
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 6.1752577319587634e-06,
+      "loss": 1.2732,
+      "step": 600
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 6.278350515463918e-06,
+      "loss": 1.2641,
+      "step": 610
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 6.381443298969072e-06,
+      "loss": 1.2919,
+      "step": 620
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 6.484536082474227e-06,
+      "loss": 1.2594,
+      "step": 630
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 6.587628865979382e-06,
+      "loss": 1.2592,
+      "step": 640
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 6.690721649484536e-06,
+      "loss": 1.2652,
+      "step": 650
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 6.793814432989692e-06,
+      "loss": 1.3274,
+      "step": 660
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 6.896907216494846e-06,
+      "loss": 1.2155,
+      "step": 670
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 7e-06,
+      "loss": 1.2837,
+      "step": 680
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 7.103092783505156e-06,
+      "loss": 1.265,
+      "step": 690
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 7.20618556701031e-06,
+      "loss": 1.2427,
+      "step": 700
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 7.309278350515464e-06,
+      "loss": 1.2688,
+      "step": 710
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 7.412371134020619e-06,
+      "loss": 1.3071,
+      "step": 720
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 7.515463917525773e-06,
+      "loss": 1.2346,
+      "step": 730
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 7.618556701030928e-06,
+      "loss": 1.2246,
+      "step": 740
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 7.721649484536083e-06,
+      "loss": 1.2604,
+      "step": 750
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 7.824742268041238e-06,
+      "loss": 1.2589,
+      "step": 760
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 7.927835051546391e-06,
+      "loss": 1.2512,
+      "step": 770
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.030927835051548e-06,
+      "loss": 1.2229,
+      "step": 780
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.134020618556701e-06,
+      "loss": 1.2326,
+      "step": 790
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.237113402061856e-06,
+      "loss": 1.3097,
+      "step": 800
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.34020618556701e-06,
+      "loss": 1.2358,
+      "step": 810
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.443298969072166e-06,
+      "loss": 1.2746,
+      "step": 820
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.54639175257732e-06,
+      "loss": 1.3063,
+      "step": 830
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.649484536082475e-06,
+      "loss": 1.2702,
+      "step": 840
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.75257731958763e-06,
+      "loss": 1.2421,
+      "step": 850
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.855670103092783e-06,
+      "loss": 1.3047,
+      "step": 860
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.95876288659794e-06,
+      "loss": 1.2452,
+      "step": 870
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 9.061855670103093e-06,
+      "loss": 1.2753,
+      "step": 880
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.164948453608248e-06,
+      "loss": 1.2532,
+      "step": 890
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.268041237113403e-06,
+      "loss": 1.2379,
+      "step": 900
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.371134020618558e-06,
+      "loss": 1.2891,
+      "step": 910
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.474226804123711e-06,
+      "loss": 1.2773,
+      "step": 920
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.577319587628868e-06,
+      "loss": 1.2753,
+      "step": 930
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.68041237113402e-06,
+      "loss": 1.1976,
+      "step": 940
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.783505154639176e-06,
+      "loss": 1.2721,
+      "step": 950
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.88659793814433e-06,
+      "loss": 1.3058,
+      "step": 960
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.989690721649485e-06,
+      "loss": 1.2383,
+      "step": 970
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.998105303046253e-06,
+      "loss": 1.2356,
+      "step": 980
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.996000084208755e-06,
+      "loss": 1.2639,
+      "step": 990
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.993894865371256e-06,
+      "loss": 1.2689,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.991789646533758e-06,
+      "loss": 1.2764,
+      "step": 1010
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.989684427696261e-06,
+      "loss": 1.2602,
+      "step": 1020
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.987579208858762e-06,
+      "loss": 1.3034,
+      "step": 1030
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.985473990021264e-06,
+      "loss": 1.2318,
+      "step": 1040
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.983368771183765e-06,
+      "loss": 1.252,
+      "step": 1050
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.981263552346267e-06,
+      "loss": 1.3111,
+      "step": 1060
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.97915833350877e-06,
+      "loss": 1.255,
+      "step": 1070
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.97705311467127e-06,
+      "loss": 1.2736,
+      "step": 1080
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.974947895833773e-06,
+      "loss": 1.3128,
+      "step": 1090
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.972842676996276e-06,
+      "loss": 1.2349,
+      "step": 1100
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.970737458158777e-06,
+      "loss": 1.2797,
+      "step": 1110
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.968632239321279e-06,
+      "loss": 1.2698,
+      "step": 1120
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.96652702048378e-06,
+      "loss": 1.2714,
+      "step": 1130
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.964421801646282e-06,
+      "loss": 1.2965,
+      "step": 1140
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.962316582808785e-06,
+      "loss": 1.2764,
+      "step": 1150
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.960211363971286e-06,
+      "loss": 1.2294,
+      "step": 1160
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.958106145133788e-06,
+      "loss": 1.2623,
+      "step": 1170
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.956000926296289e-06,
+      "loss": 1.2814,
+      "step": 1180
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.953895707458791e-06,
+      "loss": 1.2793,
+      "step": 1190
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.951790488621294e-06,
+      "loss": 1.2749,
+      "step": 1200
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.949685269783795e-06,
+      "loss": 1.2689,
+      "step": 1210
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.947580050946297e-06,
+      "loss": 1.2913,
+      "step": 1220
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.9454748321088e-06,
+      "loss": 1.2386,
+      "step": 1230
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.9433696132713e-06,
+      "loss": 1.2519,
+      "step": 1240
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.941264394433803e-06,
+      "loss": 1.2893,
+      "step": 1250
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.939159175596304e-06,
+      "loss": 1.2398,
+      "step": 1260
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.937053956758805e-06,
+      "loss": 1.2836,
+      "step": 1270
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.934948737921307e-06,
+      "loss": 1.2326,
+      "step": 1280
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.93284351908381e-06,
+      "loss": 1.2768,
+      "step": 1290
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.93073830024631e-06,
+      "loss": 1.2584,
+      "step": 1300
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.928633081408813e-06,
+      "loss": 1.2738,
+      "step": 1310
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.926527862571314e-06,
+      "loss": 1.2329,
+      "step": 1320
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.924422643733816e-06,
+      "loss": 1.2623,
+      "step": 1330
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.922317424896319e-06,
+      "loss": 1.2839,
+      "step": 1340
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.92021220605882e-06,
+      "loss": 1.2806,
+      "step": 1350
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.918106987221322e-06,
+      "loss": 1.232,
+      "step": 1360
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.916001768383825e-06,
+      "loss": 1.256,
+      "step": 1370
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.913896549546325e-06,
+      "loss": 1.225,
+      "step": 1380
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.911791330708828e-06,
+      "loss": 1.2247,
+      "step": 1390
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.909686111871329e-06,
+      "loss": 1.2422,
+      "step": 1400
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.907580893033831e-06,
+      "loss": 1.2499,
+      "step": 1410
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.905475674196334e-06,
+      "loss": 1.2569,
+      "step": 1420
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.903370455358834e-06,
+      "loss": 1.2771,
+      "step": 1430
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.901265236521337e-06,
+      "loss": 1.2789,
+      "step": 1440
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.89916001768384e-06,
+      "loss": 1.2213,
+      "step": 1450
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.89705479884634e-06,
+      "loss": 1.2393,
+      "step": 1460
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.894949580008843e-06,
+      "loss": 1.257,
+      "step": 1470
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.892844361171344e-06,
+      "loss": 1.2572,
+      "step": 1480
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.890739142333846e-06,
+      "loss": 1.2503,
+      "step": 1490
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.888633923496349e-06,
+      "loss": 1.2404,
+      "step": 1500
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.88652870465885e-06,
+      "loss": 1.2847,
+      "step": 1510
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.884423485821352e-06,
+      "loss": 1.2551,
+      "step": 1520
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.882318266983854e-06,
+      "loss": 1.2741,
+      "step": 1530
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.880213048146355e-06,
+      "loss": 1.2735,
+      "step": 1540
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.878107829308858e-06,
+      "loss": 1.273,
+      "step": 1550
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.876002610471358e-06,
+      "loss": 1.2357,
+      "step": 1560
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.873897391633861e-06,
+      "loss": 1.1864,
+      "step": 1570
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.871792172796363e-06,
+      "loss": 1.277,
+      "step": 1580
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.869686953958864e-06,
+      "loss": 1.2375,
+      "step": 1590
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.867581735121367e-06,
+      "loss": 1.2754,
+      "step": 1600
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.865476516283868e-06,
+      "loss": 1.2051,
+      "step": 1610
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.86337129744637e-06,
+      "loss": 1.2579,
+      "step": 1620
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.861266078608873e-06,
+      "loss": 1.2216,
+      "step": 1630
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.859160859771373e-06,
+      "loss": 1.2529,
+      "step": 1640
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.857055640933876e-06,
+      "loss": 1.207,
+      "step": 1650
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.854950422096378e-06,
+      "loss": 1.2275,
+      "step": 1660
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.85284520325888e-06,
+      "loss": 1.2769,
+      "step": 1670
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.850739984421382e-06,
+      "loss": 1.2165,
+      "step": 1680
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.848634765583883e-06,
+      "loss": 1.2903,
+      "step": 1690
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.846529546746385e-06,
+      "loss": 1.2548,
+      "step": 1700
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.844424327908888e-06,
+      "loss": 1.2652,
+      "step": 1710
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.842319109071388e-06,
+      "loss": 1.2718,
+      "step": 1720
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.84021389023389e-06,
+      "loss": 1.269,
+      "step": 1730
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.838108671396393e-06,
+      "loss": 1.2362,
+      "step": 1740
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.836003452558894e-06,
+      "loss": 1.205,
+      "step": 1750
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.833898233721397e-06,
+      "loss": 1.2649,
+      "step": 1760
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.831793014883897e-06,
+      "loss": 1.2517,
+      "step": 1770
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.8296877960464e-06,
+      "loss": 1.2015,
+      "step": 1780
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.827582577208902e-06,
+      "loss": 1.2382,
+      "step": 1790
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.825477358371403e-06,
+      "loss": 1.2236,
+      "step": 1800
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.823372139533906e-06,
+      "loss": 1.2503,
+      "step": 1810
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.821266920696407e-06,
+      "loss": 1.2409,
+      "step": 1820
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.819161701858909e-06,
+      "loss": 1.2236,
+      "step": 1830
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.817056483021412e-06,
+      "loss": 1.2246,
+      "step": 1840
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.814951264183912e-06,
+      "loss": 1.234,
+      "step": 1850
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.812846045346415e-06,
+      "loss": 1.2384,
+      "step": 1860
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.810740826508917e-06,
+      "loss": 1.2103,
+      "step": 1870
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.808635607671418e-06,
+      "loss": 1.2398,
+      "step": 1880
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.80653038883392e-06,
+      "loss": 1.2246,
+      "step": 1890
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.804425169996421e-06,
+      "loss": 1.1894,
+      "step": 1900
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.802319951158924e-06,
+      "loss": 1.262,
+      "step": 1910
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.800214732321426e-06,
+      "loss": 1.2595,
+      "step": 1920
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.798109513483927e-06,
+      "loss": 1.1954,
+      "step": 1930
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.79600429464643e-06,
+      "loss": 1.2578,
+      "step": 1940
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.793899075808932e-06,
+      "loss": 1.2286,
+      "step": 1950
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.791793856971433e-06,
+      "loss": 1.2455,
+      "step": 1960
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.789688638133936e-06,
+      "loss": 1.214,
+      "step": 1970
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.787583419296436e-06,
+      "loss": 1.2488,
+      "step": 1980
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.785478200458939e-06,
+      "loss": 1.266,
+      "step": 1990
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.783372981621441e-06,
+      "loss": 1.2139,
+      "step": 2000
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.781267762783942e-06,
+      "loss": 1.2516,
+      "step": 2010
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.779162543946445e-06,
+      "loss": 1.2878,
+      "step": 2020
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.777057325108945e-06,
+      "loss": 1.2494,
+      "step": 2030
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.774952106271448e-06,
+      "loss": 1.207,
+      "step": 2040
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.77284688743395e-06,
+      "loss": 1.2796,
+      "step": 2050
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.770741668596451e-06,
+      "loss": 1.2285,
+      "step": 2060
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.768636449758954e-06,
+      "loss": 1.2467,
+      "step": 2070
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.766531230921456e-06,
+      "loss": 1.1801,
+      "step": 2080
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.764426012083957e-06,
+      "loss": 1.2399,
+      "step": 2090
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.76232079324646e-06,
+      "loss": 1.2359,
+      "step": 2100
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.76021557440896e-06,
+      "loss": 1.2074,
+      "step": 2110
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.758110355571463e-06,
+      "loss": 1.2601,
+      "step": 2120
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.756005136733965e-06,
+      "loss": 1.2456,
+      "step": 2130
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.753899917896466e-06,
+      "loss": 1.2479,
+      "step": 2140
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.751794699058969e-06,
+      "loss": 1.2593,
+      "step": 2150
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.749689480221471e-06,
+      "loss": 1.1856,
+      "step": 2160
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.747584261383972e-06,
+      "loss": 1.2634,
+      "step": 2170
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.745479042546474e-06,
+      "loss": 1.2046,
+      "step": 2180
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.743373823708975e-06,
+      "loss": 1.2753,
+      "step": 2190
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.741268604871478e-06,
+      "loss": 1.2393,
+      "step": 2200
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.739163386033979e-06,
+      "loss": 1.224,
+      "step": 2210
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.737058167196481e-06,
+      "loss": 1.2767,
+      "step": 2220
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.734952948358982e-06,
+      "loss": 1.2584,
+      "step": 2230
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.732847729521484e-06,
+      "loss": 1.2717,
+      "step": 2240
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.730742510683985e-06,
+      "loss": 1.2364,
+      "step": 2250
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.728637291846488e-06,
+      "loss": 1.2354,
+      "step": 2260
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.72653207300899e-06,
+      "loss": 1.2347,
+      "step": 2270
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.724426854171491e-06,
+      "loss": 1.2537,
+      "step": 2280
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.722321635333993e-06,
+      "loss": 1.2186,
+      "step": 2290
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.720216416496496e-06,
+      "loss": 1.2351,
+      "step": 2300
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.718111197658997e-06,
+      "loss": 1.2325,
+      "step": 2310
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.7160059788215e-06,
+      "loss": 1.1996,
+      "step": 2320
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.713900759984e-06,
+      "loss": 1.2023,
+      "step": 2330
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.711795541146503e-06,
+      "loss": 1.2527,
+      "step": 2340
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.709690322309005e-06,
+      "loss": 1.2281,
+      "step": 2350
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.707585103471506e-06,
+      "loss": 1.2382,
+      "step": 2360
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.705479884634008e-06,
+      "loss": 1.2405,
+      "step": 2370
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.70337466579651e-06,
+      "loss": 1.248,
+      "step": 2380
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.701269446959012e-06,
+      "loss": 1.224,
+      "step": 2390
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.699164228121514e-06,
+      "loss": 1.22,
+      "step": 2400
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.697059009284015e-06,
+      "loss": 1.219,
+      "step": 2410
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.694953790446517e-06,
+      "loss": 1.2518,
+      "step": 2420
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.69284857160902e-06,
+      "loss": 1.219,
+      "step": 2430
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.69074335277152e-06,
+      "loss": 1.2168,
+      "step": 2440
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.688638133934023e-06,
+      "loss": 1.2469,
+      "step": 2450
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.686532915096524e-06,
+      "loss": 1.2381,
+      "step": 2460
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.684427696259027e-06,
+      "loss": 1.2001,
+      "step": 2470
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.682322477421529e-06,
+      "loss": 1.2004,
+      "step": 2480
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.68021725858403e-06,
+      "loss": 1.2409,
+      "step": 2490
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.678112039746532e-06,
+      "loss": 1.2389,
+      "step": 2500
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.676006820909035e-06,
+      "loss": 1.242,
+      "step": 2510
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.673901602071536e-06,
+      "loss": 1.2372,
+      "step": 2520
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.671796383234038e-06,
+      "loss": 1.2223,
+      "step": 2530
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.669691164396539e-06,
+      "loss": 1.2506,
+      "step": 2540
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.667585945559041e-06,
+      "loss": 1.2093,
+      "step": 2550
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.665480726721544e-06,
+      "loss": 1.2171,
+      "step": 2560
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.663375507884045e-06,
+      "loss": 1.2363,
+      "step": 2570
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.661270289046547e-06,
+      "loss": 1.2978,
+      "step": 2580
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.65916507020905e-06,
+      "loss": 1.2216,
+      "step": 2590
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.65705985137155e-06,
+      "loss": 1.1937,
+      "step": 2600
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.654954632534053e-06,
+      "loss": 1.2366,
+      "step": 2610
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.652849413696554e-06,
+      "loss": 1.2465,
+      "step": 2620
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.650744194859056e-06,
+      "loss": 1.2704,
+      "step": 2630
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.648638976021559e-06,
+      "loss": 1.2113,
+      "step": 2640
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.64653375718406e-06,
+      "loss": 1.2679,
+      "step": 2650
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.644428538346562e-06,
+      "loss": 1.2005,
+      "step": 2660
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.642323319509063e-06,
+      "loss": 1.2474,
+      "step": 2670
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.640218100671565e-06,
+      "loss": 1.2308,
+      "step": 2680
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.638112881834068e-06,
+      "loss": 1.2391,
+      "step": 2690
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.636007662996569e-06,
+      "loss": 1.1968,
+      "step": 2700
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.633902444159071e-06,
+      "loss": 1.2001,
+      "step": 2710
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.631797225321574e-06,
+      "loss": 1.2688,
+      "step": 2720
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.629692006484075e-06,
+      "loss": 1.2646,
+      "step": 2730
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.627586787646577e-06,
+      "loss": 1.2606,
+      "step": 2740
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.625481568809078e-06,
+      "loss": 1.1915,
+      "step": 2750
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.62337634997158e-06,
+      "loss": 1.204,
+      "step": 2760
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.621271131134083e-06,
+      "loss": 1.2128,
+      "step": 2770
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.619165912296584e-06,
+      "loss": 1.2116,
+      "step": 2780
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.617060693459086e-06,
+      "loss": 1.2287,
+      "step": 2790
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.614955474621589e-06,
+      "loss": 1.2443,
+      "step": 2800
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.61285025578409e-06,
+      "loss": 1.2926,
+      "step": 2810
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.610745036946592e-06,
+      "loss": 1.2195,
+      "step": 2820
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.608639818109093e-06,
+      "loss": 1.2345,
+      "step": 2830
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.606534599271595e-06,
+      "loss": 1.2588,
+      "step": 2840
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.604429380434098e-06,
+      "loss": 1.2392,
+      "step": 2850
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.602324161596599e-06,
+      "loss": 1.2529,
+      "step": 2860
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.600218942759101e-06,
+      "loss": 1.2119,
+      "step": 2870
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.598113723921602e-06,
+      "loss": 1.2416,
+      "step": 2880
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.596008505084104e-06,
+      "loss": 1.2111,
+      "step": 2890
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.593903286246607e-06,
+      "loss": 1.2493,
+      "step": 2900
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.591798067409108e-06,
+      "loss": 1.2481,
+      "step": 2910
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.58969284857161e-06,
+      "loss": 1.2265,
+      "step": 2920
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.587587629734113e-06,
+      "loss": 1.2549,
+      "step": 2930
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.585482410896613e-06,
+      "loss": 1.2474,
+      "step": 2940
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.583377192059116e-06,
+      "loss": 1.1773,
+      "step": 2950
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.581271973221617e-06,
+      "loss": 1.2612,
+      "step": 2960
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.57916675438412e-06,
+      "loss": 1.2247,
+      "step": 2970
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.577061535546622e-06,
+      "loss": 1.2075,
+      "step": 2980
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.574956316709123e-06,
+      "loss": 1.1812,
+      "step": 2990
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.572851097871625e-06,
+      "loss": 1.2058,
+      "step": 3000
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.570745879034128e-06,
+      "loss": 1.2781,
+      "step": 3010
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.568640660196628e-06,
+      "loss": 1.2572,
+      "step": 3020
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.566535441359131e-06,
+      "loss": 1.2794,
+      "step": 3030
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.564430222521632e-06,
+      "loss": 1.2136,
+      "step": 3040
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.562325003684134e-06,
+      "loss": 1.2632,
+      "step": 3050
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.560219784846637e-06,
+      "loss": 1.2584,
+      "step": 3060
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.558114566009137e-06,
+      "loss": 1.286,
+      "step": 3070
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.55600934717164e-06,
+      "loss": 1.247,
+      "step": 3080
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.55390412833414e-06,
+      "loss": 1.2715,
+      "step": 3090
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.551798909496643e-06,
+      "loss": 1.2184,
+      "step": 3100
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.549693690659146e-06,
+      "loss": 1.261,
+      "step": 3110
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.547588471821647e-06,
+      "loss": 1.2183,
+      "step": 3120
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.545483252984149e-06,
+      "loss": 1.1887,
+      "step": 3130
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.543378034146652e-06,
+      "loss": 1.2405,
+      "step": 3140
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.541272815309152e-06,
+      "loss": 1.2499,
+      "step": 3150
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.539167596471653e-06,
+      "loss": 1.2164,
+      "step": 3160
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.537062377634156e-06,
+      "loss": 1.2614,
+      "step": 3170
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.534957158796657e-06,
+      "loss": 1.2475,
+      "step": 3180
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.532851939959159e-06,
+      "loss": 1.2559,
+      "step": 3190
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.530746721121662e-06,
+      "loss": 1.2457,
+      "step": 3200
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.528641502284162e-06,
+      "loss": 1.2228,
+      "step": 3210
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.526536283446665e-06,
+      "loss": 1.219,
+      "step": 3220
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.524431064609166e-06,
+      "loss": 1.2255,
+      "step": 3230
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.522325845771668e-06,
+      "loss": 1.1923,
+      "step": 3240
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.52022062693417e-06,
+      "loss": 1.1996,
+      "step": 3250
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.518115408096671e-06,
+      "loss": 1.2186,
+      "step": 3260
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.516010189259174e-06,
+      "loss": 1.2384,
+      "step": 3270
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.513904970421676e-06,
+      "loss": 1.2119,
+      "step": 3280
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.511799751584177e-06,
+      "loss": 1.2455,
+      "step": 3290
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.50969453274668e-06,
+      "loss": 1.2314,
+      "step": 3300
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.50758931390918e-06,
+      "loss": 1.1995,
+      "step": 3310
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.505484095071683e-06,
+      "loss": 1.2308,
+      "step": 3320
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.503378876234186e-06,
+      "loss": 1.1957,
+      "step": 3330
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.501273657396686e-06,
+      "loss": 1.2557,
+      "step": 3340
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.499168438559189e-06,
+      "loss": 1.2351,
+      "step": 3350
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.497063219721691e-06,
+      "loss": 1.2085,
+      "step": 3360
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.494958000884192e-06,
+      "loss": 1.2241,
+      "step": 3370
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.492852782046695e-06,
+      "loss": 1.1909,
+      "step": 3380
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.490747563209195e-06,
+      "loss": 1.1886,
+      "step": 3390
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.488642344371698e-06,
+      "loss": 1.2161,
+      "step": 3400
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.4865371255342e-06,
+      "loss": 1.2718,
+      "step": 3410
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.484431906696701e-06,
+      "loss": 1.2007,
+      "step": 3420
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.482326687859204e-06,
+      "loss": 1.2038,
+      "step": 3430
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.480221469021705e-06,
+      "loss": 1.2153,
+      "step": 3440
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.478116250184207e-06,
+      "loss": 1.1932,
+      "step": 3450
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.47601103134671e-06,
+      "loss": 1.255,
+      "step": 3460
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.47390581250921e-06,
+      "loss": 1.2738,
+      "step": 3470
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.471800593671713e-06,
+      "loss": 1.2533,
+      "step": 3480
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.469695374834215e-06,
+      "loss": 1.2408,
+      "step": 3490
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.467590155996716e-06,
+      "loss": 1.2331,
+      "step": 3500
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.465484937159219e-06,
+      "loss": 1.2205,
+      "step": 3510
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.46337971832172e-06,
+      "loss": 1.2569,
+      "step": 3520
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.461274499484222e-06,
+      "loss": 1.2245,
+      "step": 3530
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.459169280646724e-06,
+      "loss": 1.2192,
+      "step": 3540
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.457064061809225e-06,
+      "loss": 1.2518,
+      "step": 3550
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.454958842971728e-06,
+      "loss": 1.2619,
+      "step": 3560
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.45285362413423e-06,
+      "loss": 1.1859,
+      "step": 3570
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.450748405296731e-06,
+      "loss": 1.2216,
+      "step": 3580
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.448643186459234e-06,
+      "loss": 1.221,
+      "step": 3590
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.446537967621734e-06,
+      "loss": 1.2224,
+      "step": 3600
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.444432748784237e-06,
+      "loss": 1.2324,
+      "step": 3610
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.44232752994674e-06,
+      "loss": 1.1747,
+      "step": 3620
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.44022231110924e-06,
+      "loss": 1.1962,
+      "step": 3630
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.438117092271743e-06,
+      "loss": 1.2235,
+      "step": 3640
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.436011873434245e-06,
+      "loss": 1.2081,
+      "step": 3650
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.433906654596746e-06,
+      "loss": 1.2398,
+      "step": 3660
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.431801435759248e-06,
+      "loss": 1.248,
+      "step": 3670
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.42969621692175e-06,
+      "loss": 1.1936,
+      "step": 3680
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.427590998084252e-06,
+      "loss": 1.2317,
+      "step": 3690
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.425485779246754e-06,
+      "loss": 1.2139,
+      "step": 3700
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.423380560409255e-06,
+      "loss": 1.1601,
+      "step": 3710
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.421275341571758e-06,
+      "loss": 1.2127,
+      "step": 3720
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.419170122734258e-06,
+      "loss": 1.2082,
+      "step": 3730
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.41706490389676e-06,
+      "loss": 1.1971,
+      "step": 3740
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.414959685059263e-06,
+      "loss": 1.2289,
+      "step": 3750
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.412854466221764e-06,
+      "loss": 1.2133,
+      "step": 3760
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.410749247384267e-06,
+      "loss": 1.2111,
+      "step": 3770
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.408644028546769e-06,
+      "loss": 1.2342,
+      "step": 3780
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.40653880970927e-06,
+      "loss": 1.217,
+      "step": 3790
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.404433590871772e-06,
+      "loss": 1.2651,
+      "step": 3800
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.402328372034273e-06,
+      "loss": 1.2259,
+      "step": 3810
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.400223153196776e-06,
+      "loss": 1.2434,
+      "step": 3820
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.398117934359278e-06,
+      "loss": 1.2199,
+      "step": 3830
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.396012715521779e-06,
+      "loss": 1.2299,
+      "step": 3840
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.393907496684282e-06,
+      "loss": 1.2156,
+      "step": 3850
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.391802277846784e-06,
+      "loss": 1.2402,
+      "step": 3860
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.389697059009285e-06,
+      "loss": 1.237,
+      "step": 3870
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.387591840171787e-06,
+      "loss": 1.2141,
+      "step": 3880
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.385486621334288e-06,
+      "loss": 1.2253,
+      "step": 3890
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.38338140249679e-06,
+      "loss": 1.217,
+      "step": 3900
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.381276183659293e-06,
+      "loss": 1.1919,
+      "step": 3910
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.379170964821794e-06,
+      "loss": 1.1663,
+      "step": 3920
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.377065745984296e-06,
+      "loss": 1.2338,
+      "step": 3930
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.374960527146797e-06,
+      "loss": 1.2399,
+      "step": 3940
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 9.3728553083093e-06,
+      "loss": 1.1608,
+      "step": 3950
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.370750089471802e-06,
+      "loss": 1.1752,
+      "step": 3960
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.368644870634303e-06,
+      "loss": 1.2364,
+      "step": 3970
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.366539651796806e-06,
+      "loss": 1.2053,
+      "step": 3980
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.364434432959308e-06,
+      "loss": 1.2431,
+      "step": 3990
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.362329214121809e-06,
+      "loss": 1.1948,
+      "step": 4000
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.360223995284311e-06,
+      "loss": 1.2248,
+      "step": 4010
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.358118776446812e-06,
+      "loss": 1.2057,
+      "step": 4020
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.356013557609315e-06,
+      "loss": 1.2373,
+      "step": 4030
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.353908338771817e-06,
+      "loss": 1.1993,
+      "step": 4040
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.351803119934318e-06,
+      "loss": 1.1474,
+      "step": 4050
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.34969790109682e-06,
+      "loss": 1.2084,
+      "step": 4060
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.347592682259323e-06,
+      "loss": 1.224,
+      "step": 4070
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.345487463421824e-06,
+      "loss": 1.206,
+      "step": 4080
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.343382244584326e-06,
+      "loss": 1.2225,
+      "step": 4090
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.341277025746827e-06,
+      "loss": 1.2189,
+      "step": 4100
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.339171806909328e-06,
+      "loss": 1.25,
+      "step": 4110
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 9.33706658807183e-06,
+      "loss": 1.251,
+      "step": 4120
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.334961369234333e-06,
+      "loss": 1.2048,
+      "step": 4130
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.332856150396834e-06,
+      "loss": 1.2369,
+      "step": 4140
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.330750931559336e-06,
+      "loss": 1.2427,
+      "step": 4150
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.328645712721837e-06,
+      "loss": 1.2873,
+      "step": 4160
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.32654049388434e-06,
+      "loss": 1.1579,
+      "step": 4170
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.324435275046842e-06,
+      "loss": 1.2025,
+      "step": 4180
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.322330056209343e-06,
+      "loss": 1.209,
+      "step": 4190
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.320224837371845e-06,
+      "loss": 1.2015,
+      "step": 4200
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.318119618534348e-06,
+      "loss": 1.2509,
+      "step": 4210
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.316014399696849e-06,
+      "loss": 1.2696,
+      "step": 4220
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.313909180859351e-06,
+      "loss": 1.2281,
+      "step": 4230
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.311803962021852e-06,
+      "loss": 1.2089,
+      "step": 4240
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.309698743184354e-06,
+      "loss": 1.2831,
+      "step": 4250
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.307593524346857e-06,
+      "loss": 1.2757,
+      "step": 4260
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.305488305509358e-06,
+      "loss": 1.2267,
+      "step": 4270
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 9.30338308667186e-06,
+      "loss": 1.1837,
+      "step": 4280
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.301277867834361e-06,
+      "loss": 1.2245,
+      "step": 4290
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.299172648996863e-06,
+      "loss": 1.157,
+      "step": 4300
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.297067430159366e-06,
+      "loss": 1.215,
+      "step": 4310
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.294962211321867e-06,
+      "loss": 1.2421,
+      "step": 4320
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.29285699248437e-06,
+      "loss": 1.2581,
+      "step": 4330
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.290751773646872e-06,
+      "loss": 1.1966,
+      "step": 4340
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.288646554809373e-06,
+      "loss": 1.2494,
+      "step": 4350
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.286541335971875e-06,
+      "loss": 1.1633,
+      "step": 4360
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.284436117134376e-06,
+      "loss": 1.2258,
+      "step": 4370
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.282330898296878e-06,
+      "loss": 1.2703,
+      "step": 4380
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.280225679459381e-06,
+      "loss": 1.1973,
+      "step": 4390
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.278120460621882e-06,
+      "loss": 1.2614,
+      "step": 4400
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.276015241784384e-06,
+      "loss": 1.243,
+      "step": 4410
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.273910022946887e-06,
+      "loss": 1.2473,
+      "step": 4420
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.271804804109387e-06,
+      "loss": 1.2269,
+      "step": 4430
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 9.26969958527189e-06,
+      "loss": 1.2466,
+      "step": 4440
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.26759436643439e-06,
+      "loss": 1.2362,
+      "step": 4450
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.265489147596893e-06,
+      "loss": 1.2277,
+      "step": 4460
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.263383928759396e-06,
+      "loss": 1.1939,
+      "step": 4470
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.261278709921897e-06,
+      "loss": 1.2013,
+      "step": 4480
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.259173491084399e-06,
+      "loss": 1.2057,
+      "step": 4490
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.257068272246902e-06,
+      "loss": 1.2276,
+      "step": 4500
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.254963053409402e-06,
+      "loss": 1.2029,
+      "step": 4510
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.252857834571905e-06,
+      "loss": 1.2285,
+      "step": 4520
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.250752615734406e-06,
+      "loss": 1.2078,
+      "step": 4530
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.248647396896908e-06,
+      "loss": 1.2317,
+      "step": 4540
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.24654217805941e-06,
+      "loss": 1.2266,
+      "step": 4550
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.244436959221911e-06,
+      "loss": 1.212,
+      "step": 4560
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.242331740384414e-06,
+      "loss": 1.1849,
+      "step": 4570
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.240226521546915e-06,
+      "loss": 1.2238,
+      "step": 4580
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.238121302709417e-06,
+      "loss": 1.221,
+      "step": 4590
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 9.23601608387192e-06,
+      "loss": 1.2356,
+      "step": 4600
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.23391086503442e-06,
+      "loss": 1.2287,
+      "step": 4610
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.231805646196923e-06,
+      "loss": 1.2226,
+      "step": 4620
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.229700427359426e-06,
+      "loss": 1.159,
+      "step": 4630
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.227595208521926e-06,
+      "loss": 1.2239,
+      "step": 4640
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.225489989684429e-06,
+      "loss": 1.2547,
+      "step": 4650
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.22338477084693e-06,
+      "loss": 1.1689,
+      "step": 4660
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.221279552009432e-06,
+      "loss": 1.1546,
+      "step": 4670
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.219174333171935e-06,
+      "loss": 1.2197,
+      "step": 4680
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.217069114334436e-06,
+      "loss": 1.1826,
+      "step": 4690
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.214963895496938e-06,
+      "loss": 1.2543,
+      "step": 4700
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.21285867665944e-06,
+      "loss": 1.1747,
+      "step": 4710
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.210753457821941e-06,
+      "loss": 1.2486,
+      "step": 4720
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.208648238984444e-06,
+      "loss": 1.2506,
+      "step": 4730
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.206543020146945e-06,
+      "loss": 1.2257,
+      "step": 4740
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.204437801309447e-06,
+      "loss": 1.183,
+      "step": 4750
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 9.20233258247195e-06,
+      "loss": 1.2092,
+      "step": 4760
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.20022736363445e-06,
+      "loss": 1.1907,
+      "step": 4770
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.198122144796953e-06,
+      "loss": 1.1898,
+      "step": 4780
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.196016925959454e-06,
+      "loss": 1.1834,
+      "step": 4790
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.193911707121956e-06,
+      "loss": 1.2316,
+      "step": 4800
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.191806488284459e-06,
+      "loss": 1.2018,
+      "step": 4810
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.18970126944696e-06,
+      "loss": 1.2066,
+      "step": 4820
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.187596050609462e-06,
+      "loss": 1.182,
+      "step": 4830
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.185490831771965e-06,
+      "loss": 1.2089,
+      "step": 4840
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.183385612934465e-06,
+      "loss": 1.1919,
+      "step": 4850
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.181280394096968e-06,
+      "loss": 1.2734,
+      "step": 4860
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.179175175259469e-06,
+      "loss": 1.25,
+      "step": 4870
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.177069956421971e-06,
+      "loss": 1.213,
+      "step": 4880
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.174964737584474e-06,
+      "loss": 1.2126,
+      "step": 4890
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.172859518746974e-06,
+      "loss": 1.203,
+      "step": 4900
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.170754299909477e-06,
+      "loss": 1.2718,
+      "step": 4910
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 9.16864908107198e-06,
+      "loss": 1.1845,
+      "step": 4920
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.16654386223448e-06,
+      "loss": 1.1991,
+      "step": 4930
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.164438643396983e-06,
+      "loss": 1.2049,
+      "step": 4940
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.162333424559484e-06,
+      "loss": 1.2345,
+      "step": 4950
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.160228205721986e-06,
+      "loss": 1.2284,
+      "step": 4960
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.158122986884489e-06,
+      "loss": 1.2253,
+      "step": 4970
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.15601776804699e-06,
+      "loss": 1.2038,
+      "step": 4980
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.153912549209492e-06,
+      "loss": 1.2353,
+      "step": 4990
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.151807330371993e-06,
+      "loss": 1.2029,
+      "step": 5000
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.149702111534495e-06,
+      "loss": 1.1451,
+      "step": 5010
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.147596892696998e-06,
+      "loss": 1.186,
+      "step": 5020
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.145491673859498e-06,
+      "loss": 1.2406,
+      "step": 5030
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.143386455022e-06,
+      "loss": 1.1957,
+      "step": 5040
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.141281236184502e-06,
+      "loss": 1.19,
+      "step": 5050
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.139176017347004e-06,
+      "loss": 1.2007,
+      "step": 5060
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.137070798509505e-06,
+      "loss": 1.2259,
+      "step": 5070
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 9.134965579672008e-06,
+      "loss": 1.2204,
+      "step": 5080
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.132860360834508e-06,
+      "loss": 1.2021,
+      "step": 5090
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.13075514199701e-06,
+      "loss": 1.2208,
+      "step": 5100
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.128649923159513e-06,
+      "loss": 1.1702,
+      "step": 5110
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.126544704322014e-06,
+      "loss": 1.2513,
+      "step": 5120
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.124439485484517e-06,
+      "loss": 1.1855,
+      "step": 5130
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.122334266647017e-06,
+      "loss": 1.1868,
+      "step": 5140
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.12022904780952e-06,
+      "loss": 1.1861,
+      "step": 5150
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.118123828972022e-06,
+      "loss": 1.184,
+      "step": 5160
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.116018610134523e-06,
+      "loss": 1.2227,
+      "step": 5170
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.113913391297026e-06,
+      "loss": 1.1739,
+      "step": 5180
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.111808172459528e-06,
+      "loss": 1.1705,
+      "step": 5190
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.109702953622029e-06,
+      "loss": 1.1681,
+      "step": 5200
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.107597734784532e-06,
+      "loss": 1.2116,
+      "step": 5210
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.105492515947032e-06,
+      "loss": 1.2158,
+      "step": 5220
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.103387297109535e-06,
+      "loss": 1.1452,
+      "step": 5230
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.101282078272037e-06,
+      "loss": 1.1695,
+      "step": 5240
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 9.099176859434538e-06,
+      "loss": 1.2099,
+      "step": 5250
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.09707164059704e-06,
+      "loss": 1.2359,
+      "step": 5260
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.094966421759543e-06,
+      "loss": 1.207,
+      "step": 5270
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.092861202922044e-06,
+      "loss": 1.2119,
+      "step": 5280
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.090755984084546e-06,
+      "loss": 1.1408,
+      "step": 5290
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.088650765247047e-06,
+      "loss": 1.1842,
+      "step": 5300
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.08654554640955e-06,
+      "loss": 1.1786,
+      "step": 5310
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.084440327572052e-06,
+      "loss": 1.2016,
+      "step": 5320
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.082335108734553e-06,
+      "loss": 1.2011,
+      "step": 5330
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.080229889897056e-06,
+      "loss": 1.235,
+      "step": 5340
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.078124671059556e-06,
+      "loss": 1.2116,
+      "step": 5350
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.076019452222059e-06,
+      "loss": 1.2083,
+      "step": 5360
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.073914233384561e-06,
+      "loss": 1.2152,
+      "step": 5370
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.071809014547062e-06,
+      "loss": 1.239,
+      "step": 5380
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.069703795709565e-06,
+      "loss": 1.1989,
+      "step": 5390
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.067598576872067e-06,
+      "loss": 1.2193,
+      "step": 5400
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 9.065493358034568e-06,
+      "loss": 1.2051,
+      "step": 5410
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.06338813919707e-06,
+      "loss": 1.1981,
+      "step": 5420
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.061282920359571e-06,
+      "loss": 1.2218,
+      "step": 5430
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.059177701522074e-06,
+      "loss": 1.2669,
+      "step": 5440
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.057072482684576e-06,
+      "loss": 1.2127,
+      "step": 5450
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.054967263847077e-06,
+      "loss": 1.1538,
+      "step": 5460
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.05286204500958e-06,
+      "loss": 1.1664,
+      "step": 5470
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.050756826172082e-06,
+      "loss": 1.1976,
+      "step": 5480
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.048651607334583e-06,
+      "loss": 1.2223,
+      "step": 5490
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.046546388497085e-06,
+      "loss": 1.201,
+      "step": 5500
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.044441169659586e-06,
+      "loss": 1.2372,
+      "step": 5510
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.042335950822089e-06,
+      "loss": 1.1758,
+      "step": 5520
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.040230731984591e-06,
+      "loss": 1.1793,
+      "step": 5530
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.038125513147092e-06,
+      "loss": 1.1706,
+      "step": 5540
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.036020294309594e-06,
+      "loss": 1.2267,
+      "step": 5550
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.033915075472097e-06,
+      "loss": 1.218,
+      "step": 5560
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 9.031809856634598e-06,
+      "loss": 1.1789,
+      "step": 5570
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.0297046377971e-06,
+      "loss": 1.2317,
+      "step": 5580
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.027599418959601e-06,
+      "loss": 1.23,
+      "step": 5590
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.025494200122104e-06,
+      "loss": 1.2058,
+      "step": 5600
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.023388981284606e-06,
+      "loss": 1.276,
+      "step": 5610
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.021283762447107e-06,
+      "loss": 1.1758,
+      "step": 5620
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.01917854360961e-06,
+      "loss": 1.182,
+      "step": 5630
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.01707332477211e-06,
+      "loss": 1.2027,
+      "step": 5640
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.014968105934613e-06,
+      "loss": 1.2442,
+      "step": 5650
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.012862887097115e-06,
+      "loss": 1.1509,
+      "step": 5660
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.010757668259616e-06,
+      "loss": 1.2369,
+      "step": 5670
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.008652449422118e-06,
+      "loss": 1.225,
+      "step": 5680
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.006547230584621e-06,
+      "loss": 1.2575,
+      "step": 5690
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.004442011747122e-06,
+      "loss": 1.1801,
+      "step": 5700
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.002336792909624e-06,
+      "loss": 1.1817,
+      "step": 5710
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 9.000231574072125e-06,
+      "loss": 1.2392,
+      "step": 5720
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 8.998126355234628e-06,
+      "loss": 1.1718,
+      "step": 5730
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.99602113639713e-06,
+      "loss": 1.2155,
+      "step": 5740
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.993915917559631e-06,
+      "loss": 1.1976,
+      "step": 5750
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.991810698722133e-06,
+      "loss": 1.1715,
+      "step": 5760
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.989705479884636e-06,
+      "loss": 1.1555,
+      "step": 5770
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.987600261047137e-06,
+      "loss": 1.2071,
+      "step": 5780
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.98549504220964e-06,
+      "loss": 1.2062,
+      "step": 5790
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.98338982337214e-06,
+      "loss": 1.1978,
+      "step": 5800
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.981284604534642e-06,
+      "loss": 1.2125,
+      "step": 5810
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.979179385697145e-06,
+      "loss": 1.1887,
+      "step": 5820
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.977074166859646e-06,
+      "loss": 1.2384,
+      "step": 5830
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.974968948022148e-06,
+      "loss": 1.2708,
+      "step": 5840
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.972863729184649e-06,
+      "loss": 1.1985,
+      "step": 5850
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.970758510347152e-06,
+      "loss": 1.2202,
+      "step": 5860
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.968653291509654e-06,
+      "loss": 1.2281,
+      "step": 5870
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.966548072672155e-06,
+      "loss": 1.2158,
+      "step": 5880
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 8.964442853834657e-06,
+      "loss": 1.1912,
+      "step": 5890
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.96233763499716e-06,
+      "loss": 1.1822,
+      "step": 5900
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.96023241615966e-06,
+      "loss": 1.2181,
+      "step": 5910
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.958127197322163e-06,
+      "loss": 1.1901,
+      "step": 5920
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.956021978484664e-06,
+      "loss": 1.1742,
+      "step": 5930
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.953916759647166e-06,
+      "loss": 1.1675,
+      "step": 5940
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.951811540809669e-06,
+      "loss": 1.249,
+      "step": 5950
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.94970632197217e-06,
+      "loss": 1.1683,
+      "step": 5960
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.947601103134672e-06,
+      "loss": 1.1558,
+      "step": 5970
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.945495884297175e-06,
+      "loss": 1.1685,
+      "step": 5980
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.943390665459674e-06,
+      "loss": 1.1734,
+      "step": 5990
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 8.941285446622176e-06,
+      "loss": 1.1478,
+      "step": 6000
+    }
+  ],
+  "max_steps": 48471,
+  "num_train_epochs": 3,
+  "total_flos": 3.0400555051175117e+19,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e46570d1be944689c0056f202607801965879286f93708fe363bbf196a88dc0
+size 5371

zero_to_fp32.py

@@ -0,0 +1,578 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example: python zero_to_fp32.py . pytorch_model.bin
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage == 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device)
+
+        if BUFFER_NAMES not in state_dict:
+            raise ValueError(f"{file} is not a model state checkpoint")
+        buffer_names = state_dict[BUFFER_NAMES]
+        if debug:
+            print("Found buffers:", buffer_names)
+
+        # recover just the buffers while restoring them to fp32 if they were saved in fp16
+        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+        param_shapes = state_dict[PARAM_SHAPES]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # handle shared params
+        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+
+    total_files = len(files)
+    state_dicts = []
+    for f in files:
+        state_dicts.append(torch.load(f, map_location=device))
+
+    if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage == 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    if zero_stage == 2:
+        fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    elif zero_stage == 3:
+        # if there is more than one param group, there will be multiple flattened tensors - one
+        # flattened tensor per group - for simplicity merge them into a single tensor
+        #
+        # XXX: could make the script more memory efficient for when there are multiple groups - it
+        # will require matching the sub-lists of param_shapes for each param group flattened tensor
+
+        fp32_flat_groups = [
+            torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) for i in range(len(state_dicts))
+        ]
+
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage == 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel()
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = fp32_flat_groups[0].numel() * world_size
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    for name, shape in param_shapes.items():
+
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # XXX: memory usage doubles here
+        state_dict[name] = torch.cat(
+            tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)),
+            0).narrow(0, 0, unpartitioned_numel).view(shape)
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``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``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
+    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
+        - ``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``
+    """
+
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+    print(f"Saving fp32 state dict to {output_file}")
+    torch.save(state_dict, output_file)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``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``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info(f"Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info(f"Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint_dir",
+                        type=str,
+                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help="path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)")
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)