Upload folder using huggingface_hub

config.json

@@ -0,0 +1,26 @@
+{
+  "_name_or_path": "/workspace/models/llama2-13b-cb",
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 13824,
+  "max_position_embeddings": 4096,
+  "model_type": "llama",
+  "num_attention_heads": 40,
+  "num_hidden_layers": 40,
+  "num_key_value_heads": 40,
+  "pad_token_id": 0,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.31.0",
+  "use_cache": false,
+  "vocab_size": 32000
+}

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "temperature": 0.9,
+  "top_p": 0.6,
+  "transformers_version": "4.31.0"
+}

latest

@@ -0,0 +1 @@
+global_step597

pytorch_model-00001-of-00003.bin

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

pytorch_model-00002-of-00003.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8e5175bcb45696a42b625ef454d60e86f68d1ee1fa8b80343da166c966f7903e
+size 9904162976

pytorch_model-00003-of-00003.bin

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

pytorch_model.bin.index.json

@@ -0,0 +1,410 @@
+{
+  "metadata": {
+    "total_size": 26031733760
+  },
+  "weight_map": {
+    "lm_head.weight": "pytorch_model-00003-of-00003.bin",
+    "model.embed_tokens.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.0.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.0.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.0.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.0.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.0.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.0.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.0.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.0.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.0.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.0.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.1.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.1.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.1.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.1.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.1.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.1.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.1.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.1.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.1.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.1.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.10.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.10.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.10.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.10.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.10.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.10.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.10.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.10.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.10.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.10.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.11.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.11.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.11.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.11.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.11.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.11.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.11.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.11.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.11.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.11.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.12.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.12.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.12.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.12.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.12.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.12.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.12.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.12.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.12.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.12.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.13.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.13.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.13.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.13.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.13.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.13.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.13.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.13.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.13.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.13.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.14.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.14.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.14.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.14.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.14.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.14.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.14.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.14.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.14.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.14.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.15.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.15.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.15.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.15.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.15.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.15.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.15.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.15.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.15.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.15.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.16.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.16.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.16.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.16.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.16.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.16.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.16.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.16.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.16.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.16.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.17.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.17.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.17.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.17.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.17.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.17.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.17.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.17.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.17.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.17.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.18.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.18.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.18.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.18.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.18.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.18.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.18.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.18.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.18.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.18.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.19.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.19.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.19.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.19.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.19.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.19.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.19.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.19.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.19.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.19.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.2.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.2.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.2.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.2.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.2.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.2.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.2.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.2.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.2.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.2.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.20.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.20.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.20.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.20.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.20.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.20.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.20.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.20.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.20.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.20.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.21.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.21.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.21.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.21.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.21.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.21.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.21.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.21.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.21.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.21.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.22.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.22.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.22.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.22.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.22.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.22.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.22.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.22.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.22.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.22.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.23.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.23.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.23.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.23.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.23.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.23.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.23.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.23.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.23.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.23.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.24.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.24.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.24.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.24.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.24.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.24.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.24.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.24.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.24.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.24.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.25.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.25.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.25.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.25.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.25.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.25.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.25.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.25.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.25.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.25.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.26.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.26.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.26.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.26.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.26.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.26.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.26.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.26.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.26.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.26.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.27.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.27.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.27.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.27.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.27.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.27.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.27.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.27.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.27.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.27.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.28.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.28.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.28.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.28.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.28.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.28.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.28.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.28.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.28.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.28.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.29.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.29.mlp.down_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.29.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.29.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.29.post_attention_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.29.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.29.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.29.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.29.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.29.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.3.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.3.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.3.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.3.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.3.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.3.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.3.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.3.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.3.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.3.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.30.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.30.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.30.mlp.gate_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.30.mlp.up_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.30.post_attention_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.30.self_attn.k_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.30.self_attn.o_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.30.self_attn.q_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.30.self_attn.rotary_emb.inv_freq": "pytorch_model-00002-of-00003.bin",
+    "model.layers.30.self_attn.v_proj.weight": "pytorch_model-00002-of-00003.bin",
+    "model.layers.31.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.31.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.31.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.31.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.31.post_attention_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.31.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.31.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.31.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.31.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.31.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.32.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.32.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.32.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.32.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.32.post_attention_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.32.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.32.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.32.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.32.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.32.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.33.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.33.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.33.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.33.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.33.post_attention_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.33.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.33.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.33.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.33.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.33.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.34.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.34.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.34.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.34.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.34.post_attention_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.34.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.34.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.34.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.34.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.34.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.35.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.35.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.35.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.35.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.35.post_attention_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.35.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.35.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.35.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.35.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.35.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.36.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.36.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.36.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.36.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.36.post_attention_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.36.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.36.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.36.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.36.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.36.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.37.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.37.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.37.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.37.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.37.post_attention_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.37.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.37.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.37.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.37.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.37.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.38.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.38.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.38.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.38.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.38.post_attention_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.38.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.38.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.38.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.38.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.38.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.39.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.39.mlp.down_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.39.mlp.gate_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.39.mlp.up_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.39.post_attention_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.39.self_attn.k_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.39.self_attn.o_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.39.self_attn.q_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.39.self_attn.rotary_emb.inv_freq": "pytorch_model-00003-of-00003.bin",
+    "model.layers.39.self_attn.v_proj.weight": "pytorch_model-00003-of-00003.bin",
+    "model.layers.4.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.4.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.4.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.4.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.4.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.4.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.4.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.4.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.4.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.4.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.5.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.5.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.5.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.5.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.5.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.5.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.5.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.5.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.5.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.5.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.6.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.6.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.6.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.6.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.6.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.6.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.6.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.6.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.6.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.6.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.7.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.7.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.7.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.7.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.7.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.7.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.7.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.7.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.7.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.7.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.8.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.8.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.8.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.8.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.8.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.8.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.8.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.8.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.8.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.8.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.9.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.9.mlp.down_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.9.mlp.gate_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.9.mlp.up_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.9.post_attention_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.9.self_attn.k_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.9.self_attn.o_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.9.self_attn.q_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.layers.9.self_attn.rotary_emb.inv_freq": "pytorch_model-00001-of-00003.bin",
+    "model.layers.9.self_attn.v_proj.weight": "pytorch_model-00001-of-00003.bin",
+    "model.norm.weight": "pytorch_model-00003-of-00003.bin"
+  }
+}

rng_state_0.pth

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

rng_state_1.pth

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

rng_state_2.pth

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

rng_state_3.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f537c78a9ca12294baa835f33188b772f2cbf238559c0f36ab3505e19304755
+size 17655

special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<unk>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723

tokenizer_config.json

@@ -0,0 +1,35 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "legacy": false,
+  "model_max_length": 4096,
+  "pad_token": null,
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

trainer_state.json

@@ -0,0 +1,3598 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 2.992481203007519,
+  "global_step": 597,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.0,
+      "loss": 1.6269,
+      "step": 1
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 4.362085839710631e-06,
+      "loss": 1.3291,
+      "step": 2
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 6.9137424808681075e-06,
+      "loss": 1.288,
+      "step": 3
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 8.724171679421262e-06,
+      "loss": 1.2228,
+      "step": 4
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.0128449663534445e-05,
+      "loss": 1.1311,
+      "step": 5
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1.1275828320578736e-05,
+      "loss": 1.1302,
+      "step": 6
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.2245923152549418e-05,
+      "loss": 1.1106,
+      "step": 7
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 1.3086257519131893e-05,
+      "loss": 1.0906,
+      "step": 8
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.3827484961736215e-05,
+      "loss": 1.0762,
+      "step": 9
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 1.449053550324508e-05,
+      "loss": 1.1165,
+      "step": 10
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.5090337677104985e-05,
+      "loss": 1.0933,
+      "step": 11
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 1.563791416028937e-05,
+      "loss": 1.0669,
+      "step": 12
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.6141635695206056e-05,
+      "loss": 1.0529,
+      "step": 13
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 1.660800899226005e-05,
+      "loss": 1.0585,
+      "step": 14
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.7042192144402555e-05,
+      "loss": 1.0577,
+      "step": 15
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 1.7448343358842524e-05,
+      "loss": 1.0163,
+      "step": 16
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.782986378016149e-05,
+      "loss": 1.0157,
+      "step": 17
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 1.8189570801446844e-05,
+      "loss": 1.0188,
+      "step": 18
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.8529824454509456e-05,
+      "loss": 1.011,
+      "step": 19
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.885262134295571e-05,
+      "loss": 0.9944,
+      "step": 20
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9159665633417526e-05,
+      "loss": 1.0111,
+      "step": 21
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 1.9452423516815614e-05,
+      "loss": 1.0075,
+      "step": 22
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 1.9732165553570022e-05,
+      "loss": 0.977,
+      "step": 23
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 2e-05,
+      "loss": 0.9518,
+      "step": 24
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 2e-05,
+      "loss": 1.0217,
+      "step": 25
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 1.9965095986038395e-05,
+      "loss": 0.9755,
+      "step": 26
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9930191972076792e-05,
+      "loss": 0.9597,
+      "step": 27
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 1.9895287958115186e-05,
+      "loss": 0.9954,
+      "step": 28
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.986038394415358e-05,
+      "loss": 0.9797,
+      "step": 29
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 1.9825479930191973e-05,
+      "loss": 0.9659,
+      "step": 30
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9790575916230367e-05,
+      "loss": 0.9604,
+      "step": 31
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 1.9755671902268764e-05,
+      "loss": 0.9531,
+      "step": 32
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.9720767888307158e-05,
+      "loss": 0.9474,
+      "step": 33
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 1.968586387434555e-05,
+      "loss": 0.9707,
+      "step": 34
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.9650959860383945e-05,
+      "loss": 0.9531,
+      "step": 35
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 1.961605584642234e-05,
+      "loss": 0.9464,
+      "step": 36
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.9581151832460736e-05,
+      "loss": 0.9683,
+      "step": 37
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 1.954624781849913e-05,
+      "loss": 0.9791,
+      "step": 38
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.9511343804537523e-05,
+      "loss": 0.9639,
+      "step": 39
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 1.947643979057592e-05,
+      "loss": 0.9698,
+      "step": 40
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.9441535776614313e-05,
+      "loss": 0.9494,
+      "step": 41
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 1.9406631762652707e-05,
+      "loss": 0.9133,
+      "step": 42
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.93717277486911e-05,
+      "loss": 0.957,
+      "step": 43
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 1.9336823734729494e-05,
+      "loss": 0.9716,
+      "step": 44
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.930191972076789e-05,
+      "loss": 0.9417,
+      "step": 45
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 1.9267015706806285e-05,
+      "loss": 0.9314,
+      "step": 46
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.923211169284468e-05,
+      "loss": 0.9466,
+      "step": 47
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 1.9197207678883076e-05,
+      "loss": 0.9356,
+      "step": 48
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.9162303664921466e-05,
+      "loss": 0.9515,
+      "step": 49
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 1.9127399650959863e-05,
+      "loss": 0.9493,
+      "step": 50
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.9092495636998257e-05,
+      "loss": 0.9238,
+      "step": 51
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 1.905759162303665e-05,
+      "loss": 0.8964,
+      "step": 52
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.9022687609075044e-05,
+      "loss": 0.9171,
+      "step": 53
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 1.898778359511344e-05,
+      "loss": 0.9376,
+      "step": 54
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.895287958115183e-05,
+      "loss": 0.9099,
+      "step": 55
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 1.891797556719023e-05,
+      "loss": 0.9079,
+      "step": 56
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.8883071553228622e-05,
+      "loss": 0.9179,
+      "step": 57
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 1.8848167539267016e-05,
+      "loss": 0.927,
+      "step": 58
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.8813263525305413e-05,
+      "loss": 0.9034,
+      "step": 59
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 1.8778359511343806e-05,
+      "loss": 0.9394,
+      "step": 60
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.87434554973822e-05,
+      "loss": 0.9212,
+      "step": 61
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 1.8708551483420594e-05,
+      "loss": 0.9245,
+      "step": 62
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.8673647469458987e-05,
+      "loss": 0.9168,
+      "step": 63
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 1.8638743455497384e-05,
+      "loss": 0.9483,
+      "step": 64
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.8603839441535778e-05,
+      "loss": 0.9529,
+      "step": 65
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 1.856893542757417e-05,
+      "loss": 0.9324,
+      "step": 66
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.853403141361257e-05,
+      "loss": 0.9584,
+      "step": 67
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 1.8499127399650962e-05,
+      "loss": 0.9238,
+      "step": 68
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.8464223385689356e-05,
+      "loss": 0.9264,
+      "step": 69
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 1.842931937172775e-05,
+      "loss": 0.9102,
+      "step": 70
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.8394415357766143e-05,
+      "loss": 0.9337,
+      "step": 71
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 1.835951134380454e-05,
+      "loss": 0.9241,
+      "step": 72
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.8324607329842934e-05,
+      "loss": 0.9169,
+      "step": 73
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 1.8289703315881327e-05,
+      "loss": 0.9277,
+      "step": 74
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.825479930191972e-05,
+      "loss": 0.9354,
+      "step": 75
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 1.8219895287958115e-05,
+      "loss": 0.908,
+      "step": 76
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.8184991273996512e-05,
+      "loss": 0.9055,
+      "step": 77
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 1.8150087260034905e-05,
+      "loss": 0.9371,
+      "step": 78
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.81151832460733e-05,
+      "loss": 0.9266,
+      "step": 79
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 1.8080279232111696e-05,
+      "loss": 0.917,
+      "step": 80
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.804537521815009e-05,
+      "loss": 0.9042,
+      "step": 81
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 1.8010471204188483e-05,
+      "loss": 0.9092,
+      "step": 82
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.7975567190226877e-05,
+      "loss": 0.9096,
+      "step": 83
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 1.794066317626527e-05,
+      "loss": 0.9432,
+      "step": 84
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.7905759162303668e-05,
+      "loss": 0.8991,
+      "step": 85
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 1.787085514834206e-05,
+      "loss": 0.9307,
+      "step": 86
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.7835951134380455e-05,
+      "loss": 0.906,
+      "step": 87
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 1.7801047120418852e-05,
+      "loss": 0.8982,
+      "step": 88
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.7766143106457242e-05,
+      "loss": 0.9108,
+      "step": 89
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 1.773123909249564e-05,
+      "loss": 0.9366,
+      "step": 90
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.7696335078534033e-05,
+      "loss": 0.9144,
+      "step": 91
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 1.7661431064572427e-05,
+      "loss": 0.9078,
+      "step": 92
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.7626527050610824e-05,
+      "loss": 0.9023,
+      "step": 93
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 1.7591623036649217e-05,
+      "loss": 0.8962,
+      "step": 94
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.755671902268761e-05,
+      "loss": 0.8982,
+      "step": 95
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 1.7521815008726005e-05,
+      "loss": 0.8806,
+      "step": 96
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.7486910994764398e-05,
+      "loss": 0.9171,
+      "step": 97
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 1.7452006980802795e-05,
+      "loss": 0.909,
+      "step": 98
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.741710296684119e-05,
+      "loss": 0.9284,
+      "step": 99
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 1.7382198952879583e-05,
+      "loss": 0.8931,
+      "step": 100
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.734729493891798e-05,
+      "loss": 0.9087,
+      "step": 101
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 1.731239092495637e-05,
+      "loss": 0.9029,
+      "step": 102
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 1.7277486910994767e-05,
+      "loss": 0.8965,
+      "step": 103
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 1.724258289703316e-05,
+      "loss": 0.9018,
+      "step": 104
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 1.7207678883071554e-05,
+      "loss": 0.9086,
+      "step": 105
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 1.7172774869109948e-05,
+      "loss": 0.9094,
+      "step": 106
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 1.7137870855148345e-05,
+      "loss": 0.9108,
+      "step": 107
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 1.710296684118674e-05,
+      "loss": 0.9219,
+      "step": 108
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 1.7068062827225132e-05,
+      "loss": 0.8886,
+      "step": 109
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 1.7033158813263526e-05,
+      "loss": 0.8884,
+      "step": 110
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 1.699825479930192e-05,
+      "loss": 0.8941,
+      "step": 111
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 1.6963350785340316e-05,
+      "loss": 0.9083,
+      "step": 112
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 1.692844677137871e-05,
+      "loss": 0.8894,
+      "step": 113
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 1.6893542757417104e-05,
+      "loss": 0.9183,
+      "step": 114
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 1.6858638743455497e-05,
+      "loss": 0.9147,
+      "step": 115
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 1.682373472949389e-05,
+      "loss": 0.9384,
+      "step": 116
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 1.6788830715532288e-05,
+      "loss": 0.8927,
+      "step": 117
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 1.675392670157068e-05,
+      "loss": 0.8963,
+      "step": 118
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 1.6719022687609075e-05,
+      "loss": 0.8825,
+      "step": 119
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 1.6684118673647472e-05,
+      "loss": 0.9027,
+      "step": 120
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 1.6649214659685866e-05,
+      "loss": 0.8963,
+      "step": 121
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 1.661431064572426e-05,
+      "loss": 0.9076,
+      "step": 122
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 1.6579406631762653e-05,
+      "loss": 0.8914,
+      "step": 123
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 1.6544502617801047e-05,
+      "loss": 0.8643,
+      "step": 124
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 1.6509598603839444e-05,
+      "loss": 0.8928,
+      "step": 125
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 1.6474694589877838e-05,
+      "loss": 0.8884,
+      "step": 126
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 1.643979057591623e-05,
+      "loss": 0.8869,
+      "step": 127
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 1.6404886561954628e-05,
+      "loss": 0.8969,
+      "step": 128
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 1.636998254799302e-05,
+      "loss": 0.8697,
+      "step": 129
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 1.6335078534031416e-05,
+      "loss": 0.9247,
+      "step": 130
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 1.630017452006981e-05,
+      "loss": 0.8845,
+      "step": 131
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 1.6265270506108203e-05,
+      "loss": 0.9019,
+      "step": 132
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 1.62303664921466e-05,
+      "loss": 0.8775,
+      "step": 133
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 1.6195462478184994e-05,
+      "loss": 0.8651,
+      "step": 134
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 1.6160558464223387e-05,
+      "loss": 0.9006,
+      "step": 135
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 1.612565445026178e-05,
+      "loss": 0.8892,
+      "step": 136
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 1.6090750436300174e-05,
+      "loss": 0.8842,
+      "step": 137
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 1.605584642233857e-05,
+      "loss": 0.8817,
+      "step": 138
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 1.6020942408376965e-05,
+      "loss": 0.9012,
+      "step": 139
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 1.598603839441536e-05,
+      "loss": 0.8772,
+      "step": 140
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 1.5951134380453756e-05,
+      "loss": 0.8845,
+      "step": 141
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 1.5916230366492146e-05,
+      "loss": 0.9037,
+      "step": 142
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 1.5881326352530543e-05,
+      "loss": 0.8907,
+      "step": 143
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 1.5846422338568937e-05,
+      "loss": 0.8891,
+      "step": 144
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 1.581151832460733e-05,
+      "loss": 0.8905,
+      "step": 145
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 1.5776614310645727e-05,
+      "loss": 0.8923,
+      "step": 146
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 1.574171029668412e-05,
+      "loss": 0.9181,
+      "step": 147
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 1.5706806282722515e-05,
+      "loss": 0.8782,
+      "step": 148
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 1.567190226876091e-05,
+      "loss": 0.8783,
+      "step": 149
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 1.5636998254799302e-05,
+      "loss": 0.8941,
+      "step": 150
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 1.56020942408377e-05,
+      "loss": 0.9005,
+      "step": 151
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 1.5567190226876093e-05,
+      "loss": 0.8856,
+      "step": 152
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 1.5532286212914486e-05,
+      "loss": 0.9212,
+      "step": 153
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 1.5497382198952883e-05,
+      "loss": 0.8792,
+      "step": 154
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 1.5462478184991274e-05,
+      "loss": 0.8824,
+      "step": 155
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 1.5427574171029667e-05,
+      "loss": 0.8836,
+      "step": 156
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 1.5392670157068064e-05,
+      "loss": 0.8796,
+      "step": 157
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 1.5357766143106458e-05,
+      "loss": 0.8618,
+      "step": 158
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.532286212914485e-05,
+      "loss": 0.9057,
+      "step": 159
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.528795811518325e-05,
+      "loss": 0.9108,
+      "step": 160
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.5253054101221642e-05,
+      "loss": 0.8861,
+      "step": 161
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 1.5218150087260036e-05,
+      "loss": 0.8619,
+      "step": 162
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.518324607329843e-05,
+      "loss": 0.8902,
+      "step": 163
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 1.5148342059336825e-05,
+      "loss": 0.8787,
+      "step": 164
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.511343804537522e-05,
+      "loss": 0.881,
+      "step": 165
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 1.5078534031413614e-05,
+      "loss": 0.8711,
+      "step": 166
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.504363001745201e-05,
+      "loss": 0.8762,
+      "step": 167
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 1.5008726003490403e-05,
+      "loss": 0.8819,
+      "step": 168
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.4973821989528796e-05,
+      "loss": 0.9147,
+      "step": 169
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 1.493891797556719e-05,
+      "loss": 0.9049,
+      "step": 170
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.4904013961605585e-05,
+      "loss": 0.8568,
+      "step": 171
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 1.486910994764398e-05,
+      "loss": 0.8879,
+      "step": 172
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 1.4834205933682374e-05,
+      "loss": 0.9097,
+      "step": 173
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 1.479930191972077e-05,
+      "loss": 0.9034,
+      "step": 174
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 1.4764397905759162e-05,
+      "loss": 0.8992,
+      "step": 175
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 1.4729493891797557e-05,
+      "loss": 0.8815,
+      "step": 176
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 1.4694589877835952e-05,
+      "loss": 0.8874,
+      "step": 177
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 1.4659685863874346e-05,
+      "loss": 0.8754,
+      "step": 178
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 1.4624781849912741e-05,
+      "loss": 0.8926,
+      "step": 179
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 1.4589877835951137e-05,
+      "loss": 0.8894,
+      "step": 180
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 1.455497382198953e-05,
+      "loss": 0.8738,
+      "step": 181
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 1.4520069808027924e-05,
+      "loss": 0.9274,
+      "step": 182
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 1.4485165794066318e-05,
+      "loss": 0.8875,
+      "step": 183
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 1.4450261780104713e-05,
+      "loss": 0.8803,
+      "step": 184
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 1.4415357766143108e-05,
+      "loss": 0.881,
+      "step": 185
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 1.4380453752181502e-05,
+      "loss": 0.8816,
+      "step": 186
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.4345549738219897e-05,
+      "loss": 0.8668,
+      "step": 187
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 1.4310645724258293e-05,
+      "loss": 0.8928,
+      "step": 188
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.4275741710296685e-05,
+      "loss": 0.8654,
+      "step": 189
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 1.424083769633508e-05,
+      "loss": 0.8997,
+      "step": 190
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 1.4205933682373474e-05,
+      "loss": 0.8917,
+      "step": 191
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 1.4171029668411869e-05,
+      "loss": 0.8658,
+      "step": 192
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 1.4136125654450264e-05,
+      "loss": 0.8815,
+      "step": 193
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 1.4101221640488658e-05,
+      "loss": 0.8945,
+      "step": 194
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.4066317626527052e-05,
+      "loss": 0.9159,
+      "step": 195
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.4031413612565445e-05,
+      "loss": 0.8828,
+      "step": 196
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.399650959860384e-05,
+      "loss": 0.8701,
+      "step": 197
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.3961605584642234e-05,
+      "loss": 0.8688,
+      "step": 198
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 1.392670157068063e-05,
+      "loss": 0.8803,
+      "step": 199
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 1.3891797556719025e-05,
+      "loss": 0.7831,
+      "step": 200
+    },
+    {
+      "epoch": 1.01,
+      "learning_rate": 1.3856893542757418e-05,
+      "loss": 0.773,
+      "step": 201
+    },
+    {
+      "epoch": 1.01,
+      "learning_rate": 1.3821989528795812e-05,
+      "loss": 0.7513,
+      "step": 202
+    },
+    {
+      "epoch": 1.02,
+      "learning_rate": 1.3787085514834206e-05,
+      "loss": 0.7496,
+      "step": 203
+    },
+    {
+      "epoch": 1.02,
+      "learning_rate": 1.3752181500872601e-05,
+      "loss": 0.76,
+      "step": 204
+    },
+    {
+      "epoch": 1.03,
+      "learning_rate": 1.3717277486910996e-05,
+      "loss": 0.7374,
+      "step": 205
+    },
+    {
+      "epoch": 1.03,
+      "learning_rate": 1.368237347294939e-05,
+      "loss": 0.7383,
+      "step": 206
+    },
+    {
+      "epoch": 1.04,
+      "learning_rate": 1.3647469458987785e-05,
+      "loss": 0.7312,
+      "step": 207
+    },
+    {
+      "epoch": 1.04,
+      "learning_rate": 1.361256544502618e-05,
+      "loss": 0.7314,
+      "step": 208
+    },
+    {
+      "epoch": 1.05,
+      "learning_rate": 1.3577661431064573e-05,
+      "loss": 0.7364,
+      "step": 209
+    },
+    {
+      "epoch": 1.05,
+      "learning_rate": 1.3542757417102968e-05,
+      "loss": 0.7344,
+      "step": 210
+    },
+    {
+      "epoch": 1.06,
+      "learning_rate": 1.3507853403141362e-05,
+      "loss": 0.7202,
+      "step": 211
+    },
+    {
+      "epoch": 1.06,
+      "learning_rate": 1.3472949389179757e-05,
+      "loss": 0.7173,
+      "step": 212
+    },
+    {
+      "epoch": 1.07,
+      "learning_rate": 1.3438045375218152e-05,
+      "loss": 0.7294,
+      "step": 213
+    },
+    {
+      "epoch": 1.07,
+      "learning_rate": 1.3403141361256546e-05,
+      "loss": 0.7284,
+      "step": 214
+    },
+    {
+      "epoch": 1.08,
+      "learning_rate": 1.336823734729494e-05,
+      "loss": 0.7147,
+      "step": 215
+    },
+    {
+      "epoch": 1.08,
+      "learning_rate": 1.3333333333333333e-05,
+      "loss": 0.7171,
+      "step": 216
+    },
+    {
+      "epoch": 1.09,
+      "learning_rate": 1.3298429319371729e-05,
+      "loss": 0.7396,
+      "step": 217
+    },
+    {
+      "epoch": 1.09,
+      "learning_rate": 1.3263525305410124e-05,
+      "loss": 0.7201,
+      "step": 218
+    },
+    {
+      "epoch": 1.1,
+      "learning_rate": 1.3228621291448518e-05,
+      "loss": 0.7477,
+      "step": 219
+    },
+    {
+      "epoch": 1.1,
+      "learning_rate": 1.3193717277486913e-05,
+      "loss": 0.7191,
+      "step": 220
+    },
+    {
+      "epoch": 1.11,
+      "learning_rate": 1.3158813263525307e-05,
+      "loss": 0.7346,
+      "step": 221
+    },
+    {
+      "epoch": 1.11,
+      "learning_rate": 1.31239092495637e-05,
+      "loss": 0.7366,
+      "step": 222
+    },
+    {
+      "epoch": 1.12,
+      "learning_rate": 1.3089005235602094e-05,
+      "loss": 0.6966,
+      "step": 223
+    },
+    {
+      "epoch": 1.12,
+      "learning_rate": 1.305410122164049e-05,
+      "loss": 0.7437,
+      "step": 224
+    },
+    {
+      "epoch": 1.13,
+      "learning_rate": 1.3019197207678885e-05,
+      "loss": 0.7295,
+      "step": 225
+    },
+    {
+      "epoch": 1.13,
+      "learning_rate": 1.2984293193717278e-05,
+      "loss": 0.7344,
+      "step": 226
+    },
+    {
+      "epoch": 1.14,
+      "learning_rate": 1.2949389179755674e-05,
+      "loss": 0.7391,
+      "step": 227
+    },
+    {
+      "epoch": 1.14,
+      "learning_rate": 1.2914485165794069e-05,
+      "loss": 0.7167,
+      "step": 228
+    },
+    {
+      "epoch": 1.15,
+      "learning_rate": 1.287958115183246e-05,
+      "loss": 0.7286,
+      "step": 229
+    },
+    {
+      "epoch": 1.15,
+      "learning_rate": 1.2844677137870856e-05,
+      "loss": 0.7146,
+      "step": 230
+    },
+    {
+      "epoch": 1.16,
+      "learning_rate": 1.280977312390925e-05,
+      "loss": 0.733,
+      "step": 231
+    },
+    {
+      "epoch": 1.16,
+      "learning_rate": 1.2774869109947645e-05,
+      "loss": 0.7264,
+      "step": 232
+    },
+    {
+      "epoch": 1.17,
+      "learning_rate": 1.273996509598604e-05,
+      "loss": 0.7355,
+      "step": 233
+    },
+    {
+      "epoch": 1.17,
+      "learning_rate": 1.2705061082024434e-05,
+      "loss": 0.7531,
+      "step": 234
+    },
+    {
+      "epoch": 1.18,
+      "learning_rate": 1.2670157068062828e-05,
+      "loss": 0.7303,
+      "step": 235
+    },
+    {
+      "epoch": 1.18,
+      "learning_rate": 1.2635253054101221e-05,
+      "loss": 0.7381,
+      "step": 236
+    },
+    {
+      "epoch": 1.19,
+      "learning_rate": 1.2600349040139617e-05,
+      "loss": 0.7378,
+      "step": 237
+    },
+    {
+      "epoch": 1.19,
+      "learning_rate": 1.2565445026178012e-05,
+      "loss": 0.7516,
+      "step": 238
+    },
+    {
+      "epoch": 1.2,
+      "learning_rate": 1.2530541012216406e-05,
+      "loss": 0.7523,
+      "step": 239
+    },
+    {
+      "epoch": 1.2,
+      "learning_rate": 1.2495636998254801e-05,
+      "loss": 0.7167,
+      "step": 240
+    },
+    {
+      "epoch": 1.21,
+      "learning_rate": 1.2460732984293196e-05,
+      "loss": 0.7255,
+      "step": 241
+    },
+    {
+      "epoch": 1.21,
+      "learning_rate": 1.2425828970331588e-05,
+      "loss": 0.7396,
+      "step": 242
+    },
+    {
+      "epoch": 1.22,
+      "learning_rate": 1.2390924956369984e-05,
+      "loss": 0.719,
+      "step": 243
+    },
+    {
+      "epoch": 1.22,
+      "learning_rate": 1.2356020942408377e-05,
+      "loss": 0.7457,
+      "step": 244
+    },
+    {
+      "epoch": 1.23,
+      "learning_rate": 1.2321116928446773e-05,
+      "loss": 0.7341,
+      "step": 245
+    },
+    {
+      "epoch": 1.23,
+      "learning_rate": 1.2286212914485168e-05,
+      "loss": 0.7342,
+      "step": 246
+    },
+    {
+      "epoch": 1.24,
+      "learning_rate": 1.2251308900523562e-05,
+      "loss": 0.7193,
+      "step": 247
+    },
+    {
+      "epoch": 1.24,
+      "learning_rate": 1.2216404886561957e-05,
+      "loss": 0.7323,
+      "step": 248
+    },
+    {
+      "epoch": 1.25,
+      "learning_rate": 1.2181500872600349e-05,
+      "loss": 0.739,
+      "step": 249
+    },
+    {
+      "epoch": 1.25,
+      "learning_rate": 1.2146596858638744e-05,
+      "loss": 0.7503,
+      "step": 250
+    },
+    {
+      "epoch": 1.26,
+      "learning_rate": 1.2111692844677138e-05,
+      "loss": 0.7408,
+      "step": 251
+    },
+    {
+      "epoch": 1.26,
+      "learning_rate": 1.2076788830715533e-05,
+      "loss": 0.7302,
+      "step": 252
+    },
+    {
+      "epoch": 1.27,
+      "learning_rate": 1.2041884816753929e-05,
+      "loss": 0.7347,
+      "step": 253
+    },
+    {
+      "epoch": 1.27,
+      "learning_rate": 1.2006980802792322e-05,
+      "loss": 0.724,
+      "step": 254
+    },
+    {
+      "epoch": 1.28,
+      "learning_rate": 1.1972076788830716e-05,
+      "loss": 0.7312,
+      "step": 255
+    },
+    {
+      "epoch": 1.28,
+      "learning_rate": 1.193717277486911e-05,
+      "loss": 0.7115,
+      "step": 256
+    },
+    {
+      "epoch": 1.29,
+      "learning_rate": 1.1902268760907505e-05,
+      "loss": 0.7483,
+      "step": 257
+    },
+    {
+      "epoch": 1.29,
+      "learning_rate": 1.18673647469459e-05,
+      "loss": 0.7235,
+      "step": 258
+    },
+    {
+      "epoch": 1.3,
+      "learning_rate": 1.1832460732984294e-05,
+      "loss": 0.7373,
+      "step": 259
+    },
+    {
+      "epoch": 1.3,
+      "learning_rate": 1.179755671902269e-05,
+      "loss": 0.7455,
+      "step": 260
+    },
+    {
+      "epoch": 1.31,
+      "learning_rate": 1.1762652705061085e-05,
+      "loss": 0.7436,
+      "step": 261
+    },
+    {
+      "epoch": 1.31,
+      "learning_rate": 1.1727748691099476e-05,
+      "loss": 0.7396,
+      "step": 262
+    },
+    {
+      "epoch": 1.32,
+      "learning_rate": 1.1692844677137872e-05,
+      "loss": 0.7305,
+      "step": 263
+    },
+    {
+      "epoch": 1.32,
+      "learning_rate": 1.1657940663176265e-05,
+      "loss": 0.7156,
+      "step": 264
+    },
+    {
+      "epoch": 1.33,
+      "learning_rate": 1.162303664921466e-05,
+      "loss": 0.7453,
+      "step": 265
+    },
+    {
+      "epoch": 1.33,
+      "learning_rate": 1.1588132635253056e-05,
+      "loss": 0.7052,
+      "step": 266
+    },
+    {
+      "epoch": 1.34,
+      "learning_rate": 1.155322862129145e-05,
+      "loss": 0.7193,
+      "step": 267
+    },
+    {
+      "epoch": 1.34,
+      "learning_rate": 1.1518324607329845e-05,
+      "loss": 0.7203,
+      "step": 268
+    },
+    {
+      "epoch": 1.35,
+      "learning_rate": 1.1483420593368237e-05,
+      "loss": 0.7124,
+      "step": 269
+    },
+    {
+      "epoch": 1.35,
+      "learning_rate": 1.1448516579406632e-05,
+      "loss": 0.7442,
+      "step": 270
+    },
+    {
+      "epoch": 1.36,
+      "learning_rate": 1.1413612565445028e-05,
+      "loss": 0.7334,
+      "step": 271
+    },
+    {
+      "epoch": 1.36,
+      "learning_rate": 1.1378708551483421e-05,
+      "loss": 0.7376,
+      "step": 272
+    },
+    {
+      "epoch": 1.37,
+      "learning_rate": 1.1343804537521817e-05,
+      "loss": 0.7278,
+      "step": 273
+    },
+    {
+      "epoch": 1.37,
+      "learning_rate": 1.130890052356021e-05,
+      "loss": 0.7569,
+      "step": 274
+    },
+    {
+      "epoch": 1.38,
+      "learning_rate": 1.1273996509598604e-05,
+      "loss": 0.7298,
+      "step": 275
+    },
+    {
+      "epoch": 1.38,
+      "learning_rate": 1.1239092495636998e-05,
+      "loss": 0.731,
+      "step": 276
+    },
+    {
+      "epoch": 1.39,
+      "learning_rate": 1.1204188481675393e-05,
+      "loss": 0.7218,
+      "step": 277
+    },
+    {
+      "epoch": 1.39,
+      "learning_rate": 1.1169284467713788e-05,
+      "loss": 0.7405,
+      "step": 278
+    },
+    {
+      "epoch": 1.4,
+      "learning_rate": 1.1134380453752182e-05,
+      "loss": 0.7411,
+      "step": 279
+    },
+    {
+      "epoch": 1.4,
+      "learning_rate": 1.1099476439790577e-05,
+      "loss": 0.7355,
+      "step": 280
+    },
+    {
+      "epoch": 1.41,
+      "learning_rate": 1.1064572425828973e-05,
+      "loss": 0.7161,
+      "step": 281
+    },
+    {
+      "epoch": 1.41,
+      "learning_rate": 1.1029668411867365e-05,
+      "loss": 0.7179,
+      "step": 282
+    },
+    {
+      "epoch": 1.42,
+      "learning_rate": 1.099476439790576e-05,
+      "loss": 0.7442,
+      "step": 283
+    },
+    {
+      "epoch": 1.42,
+      "learning_rate": 1.0959860383944154e-05,
+      "loss": 0.756,
+      "step": 284
+    },
+    {
+      "epoch": 1.43,
+      "learning_rate": 1.0924956369982549e-05,
+      "loss": 0.7473,
+      "step": 285
+    },
+    {
+      "epoch": 1.43,
+      "learning_rate": 1.0890052356020944e-05,
+      "loss": 0.7354,
+      "step": 286
+    },
+    {
+      "epoch": 1.44,
+      "learning_rate": 1.0855148342059338e-05,
+      "loss": 0.733,
+      "step": 287
+    },
+    {
+      "epoch": 1.44,
+      "learning_rate": 1.0820244328097733e-05,
+      "loss": 0.7411,
+      "step": 288
+    },
+    {
+      "epoch": 1.45,
+      "learning_rate": 1.0785340314136125e-05,
+      "loss": 0.7489,
+      "step": 289
+    },
+    {
+      "epoch": 1.45,
+      "learning_rate": 1.075043630017452e-05,
+      "loss": 0.7401,
+      "step": 290
+    },
+    {
+      "epoch": 1.46,
+      "learning_rate": 1.0715532286212916e-05,
+      "loss": 0.7402,
+      "step": 291
+    },
+    {
+      "epoch": 1.46,
+      "learning_rate": 1.068062827225131e-05,
+      "loss": 0.7315,
+      "step": 292
+    },
+    {
+      "epoch": 1.47,
+      "learning_rate": 1.0645724258289705e-05,
+      "loss": 0.7617,
+      "step": 293
+    },
+    {
+      "epoch": 1.47,
+      "learning_rate": 1.06108202443281e-05,
+      "loss": 0.7505,
+      "step": 294
+    },
+    {
+      "epoch": 1.48,
+      "learning_rate": 1.0575916230366492e-05,
+      "loss": 0.7291,
+      "step": 295
+    },
+    {
+      "epoch": 1.48,
+      "learning_rate": 1.0541012216404887e-05,
+      "loss": 0.7406,
+      "step": 296
+    },
+    {
+      "epoch": 1.49,
+      "learning_rate": 1.0506108202443281e-05,
+      "loss": 0.7487,
+      "step": 297
+    },
+    {
+      "epoch": 1.49,
+      "learning_rate": 1.0471204188481676e-05,
+      "loss": 0.7284,
+      "step": 298
+    },
+    {
+      "epoch": 1.5,
+      "learning_rate": 1.043630017452007e-05,
+      "loss": 0.7369,
+      "step": 299
+    },
+    {
+      "epoch": 1.5,
+      "learning_rate": 1.0401396160558465e-05,
+      "loss": 0.7416,
+      "step": 300
+    },
+    {
+      "epoch": 1.51,
+      "learning_rate": 1.036649214659686e-05,
+      "loss": 0.733,
+      "step": 301
+    },
+    {
+      "epoch": 1.51,
+      "learning_rate": 1.0331588132635253e-05,
+      "loss": 0.7415,
+      "step": 302
+    },
+    {
+      "epoch": 1.52,
+      "learning_rate": 1.0296684118673648e-05,
+      "loss": 0.7393,
+      "step": 303
+    },
+    {
+      "epoch": 1.52,
+      "learning_rate": 1.0261780104712042e-05,
+      "loss": 0.7148,
+      "step": 304
+    },
+    {
+      "epoch": 1.53,
+      "learning_rate": 1.0226876090750437e-05,
+      "loss": 0.7399,
+      "step": 305
+    },
+    {
+      "epoch": 1.53,
+      "learning_rate": 1.0191972076788832e-05,
+      "loss": 0.7424,
+      "step": 306
+    },
+    {
+      "epoch": 1.54,
+      "learning_rate": 1.0157068062827226e-05,
+      "loss": 0.7453,
+      "step": 307
+    },
+    {
+      "epoch": 1.54,
+      "learning_rate": 1.0122164048865621e-05,
+      "loss": 0.7239,
+      "step": 308
+    },
+    {
+      "epoch": 1.55,
+      "learning_rate": 1.0087260034904013e-05,
+      "loss": 0.724,
+      "step": 309
+    },
+    {
+      "epoch": 1.55,
+      "learning_rate": 1.0052356020942409e-05,
+      "loss": 0.723,
+      "step": 310
+    },
+    {
+      "epoch": 1.56,
+      "learning_rate": 1.0017452006980804e-05,
+      "loss": 0.7421,
+      "step": 311
+    },
+    {
+      "epoch": 1.56,
+      "learning_rate": 9.982547993019198e-06,
+      "loss": 0.729,
+      "step": 312
+    },
+    {
+      "epoch": 1.57,
+      "learning_rate": 9.947643979057593e-06,
+      "loss": 0.7457,
+      "step": 313
+    },
+    {
+      "epoch": 1.57,
+      "learning_rate": 9.912739965095987e-06,
+      "loss": 0.7413,
+      "step": 314
+    },
+    {
+      "epoch": 1.58,
+      "learning_rate": 9.877835951134382e-06,
+      "loss": 0.7206,
+      "step": 315
+    },
+    {
+      "epoch": 1.58,
+      "learning_rate": 9.842931937172776e-06,
+      "loss": 0.731,
+      "step": 316
+    },
+    {
+      "epoch": 1.59,
+      "learning_rate": 9.80802792321117e-06,
+      "loss": 0.7395,
+      "step": 317
+    },
+    {
+      "epoch": 1.59,
+      "learning_rate": 9.773123909249565e-06,
+      "loss": 0.7377,
+      "step": 318
+    },
+    {
+      "epoch": 1.6,
+      "learning_rate": 9.73821989528796e-06,
+      "loss": 0.7589,
+      "step": 319
+    },
+    {
+      "epoch": 1.6,
+      "learning_rate": 9.703315881326354e-06,
+      "loss": 0.7416,
+      "step": 320
+    },
+    {
+      "epoch": 1.61,
+      "learning_rate": 9.668411867364747e-06,
+      "loss": 0.729,
+      "step": 321
+    },
+    {
+      "epoch": 1.61,
+      "learning_rate": 9.633507853403143e-06,
+      "loss": 0.7399,
+      "step": 322
+    },
+    {
+      "epoch": 1.62,
+      "learning_rate": 9.598603839441538e-06,
+      "loss": 0.725,
+      "step": 323
+    },
+    {
+      "epoch": 1.62,
+      "learning_rate": 9.563699825479932e-06,
+      "loss": 0.7292,
+      "step": 324
+    },
+    {
+      "epoch": 1.63,
+      "learning_rate": 9.528795811518325e-06,
+      "loss": 0.7538,
+      "step": 325
+    },
+    {
+      "epoch": 1.63,
+      "learning_rate": 9.49389179755672e-06,
+      "loss": 0.7275,
+      "step": 326
+    },
+    {
+      "epoch": 1.64,
+      "learning_rate": 9.458987783595114e-06,
+      "loss": 0.7176,
+      "step": 327
+    },
+    {
+      "epoch": 1.64,
+      "learning_rate": 9.424083769633508e-06,
+      "loss": 0.7359,
+      "step": 328
+    },
+    {
+      "epoch": 1.65,
+      "learning_rate": 9.389179755671903e-06,
+      "loss": 0.7267,
+      "step": 329
+    },
+    {
+      "epoch": 1.65,
+      "learning_rate": 9.354275741710297e-06,
+      "loss": 0.7684,
+      "step": 330
+    },
+    {
+      "epoch": 1.66,
+      "learning_rate": 9.319371727748692e-06,
+      "loss": 0.7479,
+      "step": 331
+    },
+    {
+      "epoch": 1.66,
+      "learning_rate": 9.284467713787086e-06,
+      "loss": 0.729,
+      "step": 332
+    },
+    {
+      "epoch": 1.67,
+      "learning_rate": 9.249563699825481e-06,
+      "loss": 0.7442,
+      "step": 333
+    },
+    {
+      "epoch": 1.67,
+      "learning_rate": 9.214659685863875e-06,
+      "loss": 0.726,
+      "step": 334
+    },
+    {
+      "epoch": 1.68,
+      "learning_rate": 9.17975567190227e-06,
+      "loss": 0.7368,
+      "step": 335
+    },
+    {
+      "epoch": 1.68,
+      "learning_rate": 9.144851657940664e-06,
+      "loss": 0.7198,
+      "step": 336
+    },
+    {
+      "epoch": 1.69,
+      "learning_rate": 9.109947643979057e-06,
+      "loss": 0.7173,
+      "step": 337
+    },
+    {
+      "epoch": 1.69,
+      "learning_rate": 9.075043630017453e-06,
+      "loss": 0.7129,
+      "step": 338
+    },
+    {
+      "epoch": 1.7,
+      "learning_rate": 9.040139616055848e-06,
+      "loss": 0.7412,
+      "step": 339
+    },
+    {
+      "epoch": 1.7,
+      "learning_rate": 9.005235602094242e-06,
+      "loss": 0.742,
+      "step": 340
+    },
+    {
+      "epoch": 1.71,
+      "learning_rate": 8.970331588132635e-06,
+      "loss": 0.7335,
+      "step": 341
+    },
+    {
+      "epoch": 1.71,
+      "learning_rate": 8.93542757417103e-06,
+      "loss": 0.7235,
+      "step": 342
+    },
+    {
+      "epoch": 1.72,
+      "learning_rate": 8.900523560209426e-06,
+      "loss": 0.7321,
+      "step": 343
+    },
+    {
+      "epoch": 1.72,
+      "learning_rate": 8.86561954624782e-06,
+      "loss": 0.6992,
+      "step": 344
+    },
+    {
+      "epoch": 1.73,
+      "learning_rate": 8.830715532286213e-06,
+      "loss": 0.7379,
+      "step": 345
+    },
+    {
+      "epoch": 1.73,
+      "learning_rate": 8.795811518324609e-06,
+      "loss": 0.742,
+      "step": 346
+    },
+    {
+      "epoch": 1.74,
+      "learning_rate": 8.760907504363002e-06,
+      "loss": 0.7271,
+      "step": 347
+    },
+    {
+      "epoch": 1.74,
+      "learning_rate": 8.726003490401398e-06,
+      "loss": 0.7312,
+      "step": 348
+    },
+    {
+      "epoch": 1.75,
+      "learning_rate": 8.691099476439791e-06,
+      "loss": 0.7292,
+      "step": 349
+    },
+    {
+      "epoch": 1.75,
+      "learning_rate": 8.656195462478185e-06,
+      "loss": 0.722,
+      "step": 350
+    },
+    {
+      "epoch": 1.76,
+      "learning_rate": 8.62129144851658e-06,
+      "loss": 0.706,
+      "step": 351
+    },
+    {
+      "epoch": 1.76,
+      "learning_rate": 8.586387434554974e-06,
+      "loss": 0.7335,
+      "step": 352
+    },
+    {
+      "epoch": 1.77,
+      "learning_rate": 8.55148342059337e-06,
+      "loss": 0.7317,
+      "step": 353
+    },
+    {
+      "epoch": 1.77,
+      "learning_rate": 8.516579406631763e-06,
+      "loss": 0.7313,
+      "step": 354
+    },
+    {
+      "epoch": 1.78,
+      "learning_rate": 8.481675392670158e-06,
+      "loss": 0.7308,
+      "step": 355
+    },
+    {
+      "epoch": 1.78,
+      "learning_rate": 8.446771378708552e-06,
+      "loss": 0.7521,
+      "step": 356
+    },
+    {
+      "epoch": 1.79,
+      "learning_rate": 8.411867364746945e-06,
+      "loss": 0.7206,
+      "step": 357
+    },
+    {
+      "epoch": 1.79,
+      "learning_rate": 8.37696335078534e-06,
+      "loss": 0.7184,
+      "step": 358
+    },
+    {
+      "epoch": 1.8,
+      "learning_rate": 8.342059336823736e-06,
+      "loss": 0.7268,
+      "step": 359
+    },
+    {
+      "epoch": 1.8,
+      "learning_rate": 8.30715532286213e-06,
+      "loss": 0.7283,
+      "step": 360
+    },
+    {
+      "epoch": 1.81,
+      "learning_rate": 8.272251308900523e-06,
+      "loss": 0.7294,
+      "step": 361
+    },
+    {
+      "epoch": 1.81,
+      "learning_rate": 8.237347294938919e-06,
+      "loss": 0.7103,
+      "step": 362
+    },
+    {
+      "epoch": 1.82,
+      "learning_rate": 8.202443280977314e-06,
+      "loss": 0.7449,
+      "step": 363
+    },
+    {
+      "epoch": 1.82,
+      "learning_rate": 8.167539267015708e-06,
+      "loss": 0.73,
+      "step": 364
+    },
+    {
+      "epoch": 1.83,
+      "learning_rate": 8.132635253054101e-06,
+      "loss": 0.7418,
+      "step": 365
+    },
+    {
+      "epoch": 1.83,
+      "learning_rate": 8.097731239092497e-06,
+      "loss": 0.7366,
+      "step": 366
+    },
+    {
+      "epoch": 1.84,
+      "learning_rate": 8.06282722513089e-06,
+      "loss": 0.7339,
+      "step": 367
+    },
+    {
+      "epoch": 1.84,
+      "learning_rate": 8.027923211169286e-06,
+      "loss": 0.707,
+      "step": 368
+    },
+    {
+      "epoch": 1.85,
+      "learning_rate": 7.99301919720768e-06,
+      "loss": 0.7308,
+      "step": 369
+    },
+    {
+      "epoch": 1.85,
+      "learning_rate": 7.958115183246073e-06,
+      "loss": 0.7418,
+      "step": 370
+    },
+    {
+      "epoch": 1.86,
+      "learning_rate": 7.923211169284468e-06,
+      "loss": 0.723,
+      "step": 371
+    },
+    {
+      "epoch": 1.86,
+      "learning_rate": 7.888307155322864e-06,
+      "loss": 0.7444,
+      "step": 372
+    },
+    {
+      "epoch": 1.87,
+      "learning_rate": 7.853403141361257e-06,
+      "loss": 0.7612,
+      "step": 373
+    },
+    {
+      "epoch": 1.87,
+      "learning_rate": 7.818499127399651e-06,
+      "loss": 0.7148,
+      "step": 374
+    },
+    {
+      "epoch": 1.88,
+      "learning_rate": 7.783595113438046e-06,
+      "loss": 0.7247,
+      "step": 375
+    },
+    {
+      "epoch": 1.88,
+      "learning_rate": 7.748691099476442e-06,
+      "loss": 0.7161,
+      "step": 376
+    },
+    {
+      "epoch": 1.89,
+      "learning_rate": 7.713787085514834e-06,
+      "loss": 0.7299,
+      "step": 377
+    },
+    {
+      "epoch": 1.89,
+      "learning_rate": 7.678883071553229e-06,
+      "loss": 0.7415,
+      "step": 378
+    },
+    {
+      "epoch": 1.9,
+      "learning_rate": 7.643979057591624e-06,
+      "loss": 0.7202,
+      "step": 379
+    },
+    {
+      "epoch": 1.9,
+      "learning_rate": 7.609075043630018e-06,
+      "loss": 0.7249,
+      "step": 380
+    },
+    {
+      "epoch": 1.91,
+      "learning_rate": 7.5741710296684124e-06,
+      "loss": 0.7068,
+      "step": 381
+    },
+    {
+      "epoch": 1.91,
+      "learning_rate": 7.539267015706807e-06,
+      "loss": 0.7654,
+      "step": 382
+    },
+    {
+      "epoch": 1.92,
+      "learning_rate": 7.504363001745201e-06,
+      "loss": 0.7133,
+      "step": 383
+    },
+    {
+      "epoch": 1.92,
+      "learning_rate": 7.469458987783595e-06,
+      "loss": 0.7313,
+      "step": 384
+    },
+    {
+      "epoch": 1.93,
+      "learning_rate": 7.43455497382199e-06,
+      "loss": 0.7407,
+      "step": 385
+    },
+    {
+      "epoch": 1.93,
+      "learning_rate": 7.399650959860385e-06,
+      "loss": 0.724,
+      "step": 386
+    },
+    {
+      "epoch": 1.94,
+      "learning_rate": 7.3647469458987785e-06,
+      "loss": 0.748,
+      "step": 387
+    },
+    {
+      "epoch": 1.94,
+      "learning_rate": 7.329842931937173e-06,
+      "loss": 0.7445,
+      "step": 388
+    },
+    {
+      "epoch": 1.95,
+      "learning_rate": 7.294938917975568e-06,
+      "loss": 0.7191,
+      "step": 389
+    },
+    {
+      "epoch": 1.95,
+      "learning_rate": 7.260034904013962e-06,
+      "loss": 0.7436,
+      "step": 390
+    },
+    {
+      "epoch": 1.96,
+      "learning_rate": 7.2251308900523565e-06,
+      "loss": 0.7303,
+      "step": 391
+    },
+    {
+      "epoch": 1.96,
+      "learning_rate": 7.190226876090751e-06,
+      "loss": 0.7373,
+      "step": 392
+    },
+    {
+      "epoch": 1.97,
+      "learning_rate": 7.155322862129146e-06,
+      "loss": 0.7105,
+      "step": 393
+    },
+    {
+      "epoch": 1.97,
+      "learning_rate": 7.12041884816754e-06,
+      "loss": 0.7306,
+      "step": 394
+    },
+    {
+      "epoch": 1.98,
+      "learning_rate": 7.0855148342059345e-06,
+      "loss": 0.7046,
+      "step": 395
+    },
+    {
+      "epoch": 1.98,
+      "learning_rate": 7.050610820244329e-06,
+      "loss": 0.7337,
+      "step": 396
+    },
+    {
+      "epoch": 1.99,
+      "learning_rate": 7.015706806282723e-06,
+      "loss": 0.7273,
+      "step": 397
+    },
+    {
+      "epoch": 1.99,
+      "learning_rate": 6.980802792321117e-06,
+      "loss": 0.7332,
+      "step": 398
+    },
+    {
+      "epoch": 2.0,
+      "learning_rate": 6.945898778359512e-06,
+      "loss": 0.6763,
+      "step": 399
+    },
+    {
+      "epoch": 2.01,
+      "learning_rate": 6.910994764397906e-06,
+      "loss": 0.5935,
+      "step": 400
+    },
+    {
+      "epoch": 2.01,
+      "learning_rate": 6.8760907504363005e-06,
+      "loss": 0.5917,
+      "step": 401
+    },
+    {
+      "epoch": 2.02,
+      "learning_rate": 6.841186736474695e-06,
+      "loss": 0.5795,
+      "step": 402
+    },
+    {
+      "epoch": 2.02,
+      "learning_rate": 6.80628272251309e-06,
+      "loss": 0.5724,
+      "step": 403
+    },
+    {
+      "epoch": 2.03,
+      "learning_rate": 6.771378708551484e-06,
+      "loss": 0.5707,
+      "step": 404
+    },
+    {
+      "epoch": 2.03,
+      "learning_rate": 6.7364746945898785e-06,
+      "loss": 0.5947,
+      "step": 405
+    },
+    {
+      "epoch": 2.04,
+      "learning_rate": 6.701570680628273e-06,
+      "loss": 0.5856,
+      "step": 406
+    },
+    {
+      "epoch": 2.04,
+      "learning_rate": 6.666666666666667e-06,
+      "loss": 0.5873,
+      "step": 407
+    },
+    {
+      "epoch": 2.05,
+      "learning_rate": 6.631762652705062e-06,
+      "loss": 0.5581,
+      "step": 408
+    },
+    {
+      "epoch": 2.05,
+      "learning_rate": 6.5968586387434565e-06,
+      "loss": 0.5864,
+      "step": 409
+    },
+    {
+      "epoch": 2.06,
+      "learning_rate": 6.56195462478185e-06,
+      "loss": 0.5866,
+      "step": 410
+    },
+    {
+      "epoch": 2.06,
+      "learning_rate": 6.527050610820245e-06,
+      "loss": 0.6144,
+      "step": 411
+    },
+    {
+      "epoch": 2.07,
+      "learning_rate": 6.492146596858639e-06,
+      "loss": 0.5727,
+      "step": 412
+    },
+    {
+      "epoch": 2.07,
+      "learning_rate": 6.4572425828970344e-06,
+      "loss": 0.5816,
+      "step": 413
+    },
+    {
+      "epoch": 2.08,
+      "learning_rate": 6.422338568935428e-06,
+      "loss": 0.5662,
+      "step": 414
+    },
+    {
+      "epoch": 2.08,
+      "learning_rate": 6.3874345549738226e-06,
+      "loss": 0.57,
+      "step": 415
+    },
+    {
+      "epoch": 2.09,
+      "learning_rate": 6.352530541012217e-06,
+      "loss": 0.5965,
+      "step": 416
+    },
+    {
+      "epoch": 2.09,
+      "learning_rate": 6.317626527050611e-06,
+      "loss": 0.5402,
+      "step": 417
+    },
+    {
+      "epoch": 2.1,
+      "learning_rate": 6.282722513089006e-06,
+      "loss": 0.5626,
+      "step": 418
+    },
+    {
+      "epoch": 2.1,
+      "learning_rate": 6.2478184991274005e-06,
+      "loss": 0.5822,
+      "step": 419
+    },
+    {
+      "epoch": 2.11,
+      "learning_rate": 6.212914485165794e-06,
+      "loss": 0.5609,
+      "step": 420
+    },
+    {
+      "epoch": 2.11,
+      "learning_rate": 6.178010471204189e-06,
+      "loss": 0.5886,
+      "step": 421
+    },
+    {
+      "epoch": 2.12,
+      "learning_rate": 6.143106457242584e-06,
+      "loss": 0.5705,
+      "step": 422
+    },
+    {
+      "epoch": 2.12,
+      "learning_rate": 6.1082024432809785e-06,
+      "loss": 0.5777,
+      "step": 423
+    },
+    {
+      "epoch": 2.13,
+      "learning_rate": 6.073298429319372e-06,
+      "loss": 0.5606,
+      "step": 424
+    },
+    {
+      "epoch": 2.13,
+      "learning_rate": 6.038394415357767e-06,
+      "loss": 0.5772,
+      "step": 425
+    },
+    {
+      "epoch": 2.14,
+      "learning_rate": 6.003490401396161e-06,
+      "loss": 0.5626,
+      "step": 426
+    },
+    {
+      "epoch": 2.14,
+      "learning_rate": 5.968586387434555e-06,
+      "loss": 0.5919,
+      "step": 427
+    },
+    {
+      "epoch": 2.15,
+      "learning_rate": 5.93368237347295e-06,
+      "loss": 0.5642,
+      "step": 428
+    },
+    {
+      "epoch": 2.15,
+      "learning_rate": 5.898778359511345e-06,
+      "loss": 0.5702,
+      "step": 429
+    },
+    {
+      "epoch": 2.16,
+      "learning_rate": 5.863874345549738e-06,
+      "loss": 0.5746,
+      "step": 430
+    },
+    {
+      "epoch": 2.16,
+      "learning_rate": 5.828970331588133e-06,
+      "loss": 0.579,
+      "step": 431
+    },
+    {
+      "epoch": 2.17,
+      "learning_rate": 5.794066317626528e-06,
+      "loss": 0.5748,
+      "step": 432
+    },
+    {
+      "epoch": 2.17,
+      "learning_rate": 5.7591623036649226e-06,
+      "loss": 0.5669,
+      "step": 433
+    },
+    {
+      "epoch": 2.18,
+      "learning_rate": 5.724258289703316e-06,
+      "loss": 0.5597,
+      "step": 434
+    },
+    {
+      "epoch": 2.18,
+      "learning_rate": 5.689354275741711e-06,
+      "loss": 0.5823,
+      "step": 435
+    },
+    {
+      "epoch": 2.19,
+      "learning_rate": 5.654450261780105e-06,
+      "loss": 0.5701,
+      "step": 436
+    },
+    {
+      "epoch": 2.19,
+      "learning_rate": 5.619546247818499e-06,
+      "loss": 0.5648,
+      "step": 437
+    },
+    {
+      "epoch": 2.2,
+      "learning_rate": 5.584642233856894e-06,
+      "loss": 0.5789,
+      "step": 438
+    },
+    {
+      "epoch": 2.2,
+      "learning_rate": 5.549738219895289e-06,
+      "loss": 0.5732,
+      "step": 439
+    },
+    {
+      "epoch": 2.21,
+      "learning_rate": 5.514834205933682e-06,
+      "loss": 0.5816,
+      "step": 440
+    },
+    {
+      "epoch": 2.21,
+      "learning_rate": 5.479930191972077e-06,
+      "loss": 0.5864,
+      "step": 441
+    },
+    {
+      "epoch": 2.22,
+      "learning_rate": 5.445026178010472e-06,
+      "loss": 0.5709,
+      "step": 442
+    },
+    {
+      "epoch": 2.22,
+      "learning_rate": 5.410122164048867e-06,
+      "loss": 0.5881,
+      "step": 443
+    },
+    {
+      "epoch": 2.23,
+      "learning_rate": 5.37521815008726e-06,
+      "loss": 0.573,
+      "step": 444
+    },
+    {
+      "epoch": 2.23,
+      "learning_rate": 5.340314136125655e-06,
+      "loss": 0.5668,
+      "step": 445
+    },
+    {
+      "epoch": 2.24,
+      "learning_rate": 5.30541012216405e-06,
+      "loss": 0.5864,
+      "step": 446
+    },
+    {
+      "epoch": 2.24,
+      "learning_rate": 5.270506108202444e-06,
+      "loss": 0.5532,
+      "step": 447
+    },
+    {
+      "epoch": 2.25,
+      "learning_rate": 5.235602094240838e-06,
+      "loss": 0.5632,
+      "step": 448
+    },
+    {
+      "epoch": 2.25,
+      "learning_rate": 5.200698080279233e-06,
+      "loss": 0.569,
+      "step": 449
+    },
+    {
+      "epoch": 2.26,
+      "learning_rate": 5.165794066317626e-06,
+      "loss": 0.5748,
+      "step": 450
+    },
+    {
+      "epoch": 2.26,
+      "learning_rate": 5.130890052356021e-06,
+      "loss": 0.561,
+      "step": 451
+    },
+    {
+      "epoch": 2.27,
+      "learning_rate": 5.095986038394416e-06,
+      "loss": 0.5661,
+      "step": 452
+    },
+    {
+      "epoch": 2.27,
+      "learning_rate": 5.061082024432811e-06,
+      "loss": 0.5458,
+      "step": 453
+    },
+    {
+      "epoch": 2.28,
+      "learning_rate": 5.026178010471204e-06,
+      "loss": 0.5632,
+      "step": 454
+    },
+    {
+      "epoch": 2.28,
+      "learning_rate": 4.991273996509599e-06,
+      "loss": 0.5593,
+      "step": 455
+    },
+    {
+      "epoch": 2.29,
+      "learning_rate": 4.956369982547993e-06,
+      "loss": 0.5788,
+      "step": 456
+    },
+    {
+      "epoch": 2.29,
+      "learning_rate": 4.921465968586388e-06,
+      "loss": 0.5673,
+      "step": 457
+    },
+    {
+      "epoch": 2.3,
+      "learning_rate": 4.886561954624782e-06,
+      "loss": 0.5817,
+      "step": 458
+    },
+    {
+      "epoch": 2.3,
+      "learning_rate": 4.851657940663177e-06,
+      "loss": 0.6004,
+      "step": 459
+    },
+    {
+      "epoch": 2.31,
+      "learning_rate": 4.816753926701571e-06,
+      "loss": 0.5476,
+      "step": 460
+    },
+    {
+      "epoch": 2.31,
+      "learning_rate": 4.781849912739966e-06,
+      "loss": 0.5845,
+      "step": 461
+    },
+    {
+      "epoch": 2.32,
+      "learning_rate": 4.74694589877836e-06,
+      "loss": 0.568,
+      "step": 462
+    },
+    {
+      "epoch": 2.32,
+      "learning_rate": 4.712041884816754e-06,
+      "loss": 0.5712,
+      "step": 463
+    },
+    {
+      "epoch": 2.33,
+      "learning_rate": 4.677137870855148e-06,
+      "loss": 0.583,
+      "step": 464
+    },
+    {
+      "epoch": 2.33,
+      "learning_rate": 4.642233856893543e-06,
+      "loss": 0.5612,
+      "step": 465
+    },
+    {
+      "epoch": 2.34,
+      "learning_rate": 4.607329842931937e-06,
+      "loss": 0.5544,
+      "step": 466
+    },
+    {
+      "epoch": 2.34,
+      "learning_rate": 4.572425828970332e-06,
+      "loss": 0.5843,
+      "step": 467
+    },
+    {
+      "epoch": 2.35,
+      "learning_rate": 4.537521815008726e-06,
+      "loss": 0.5771,
+      "step": 468
+    },
+    {
+      "epoch": 2.35,
+      "learning_rate": 4.502617801047121e-06,
+      "loss": 0.5871,
+      "step": 469
+    },
+    {
+      "epoch": 2.36,
+      "learning_rate": 4.467713787085515e-06,
+      "loss": 0.5876,
+      "step": 470
+    },
+    {
+      "epoch": 2.36,
+      "learning_rate": 4.43280977312391e-06,
+      "loss": 0.5789,
+      "step": 471
+    },
+    {
+      "epoch": 2.37,
+      "learning_rate": 4.397905759162304e-06,
+      "loss": 0.5898,
+      "step": 472
+    },
+    {
+      "epoch": 2.37,
+      "learning_rate": 4.363001745200699e-06,
+      "loss": 0.5548,
+      "step": 473
+    },
+    {
+      "epoch": 2.38,
+      "learning_rate": 4.3280977312390925e-06,
+      "loss": 0.5752,
+      "step": 474
+    },
+    {
+      "epoch": 2.38,
+      "learning_rate": 4.293193717277487e-06,
+      "loss": 0.5549,
+      "step": 475
+    },
+    {
+      "epoch": 2.39,
+      "learning_rate": 4.2582897033158814e-06,
+      "loss": 0.5744,
+      "step": 476
+    },
+    {
+      "epoch": 2.39,
+      "learning_rate": 4.223385689354276e-06,
+      "loss": 0.5949,
+      "step": 477
+    },
+    {
+      "epoch": 2.4,
+      "learning_rate": 4.18848167539267e-06,
+      "loss": 0.5817,
+      "step": 478
+    },
+    {
+      "epoch": 2.4,
+      "learning_rate": 4.153577661431065e-06,
+      "loss": 0.5775,
+      "step": 479
+    },
+    {
+      "epoch": 2.41,
+      "learning_rate": 4.118673647469459e-06,
+      "loss": 0.5728,
+      "step": 480
+    },
+    {
+      "epoch": 2.41,
+      "learning_rate": 4.083769633507854e-06,
+      "loss": 0.5971,
+      "step": 481
+    },
+    {
+      "epoch": 2.42,
+      "learning_rate": 4.048865619546248e-06,
+      "loss": 0.5743,
+      "step": 482
+    },
+    {
+      "epoch": 2.42,
+      "learning_rate": 4.013961605584643e-06,
+      "loss": 0.5767,
+      "step": 483
+    },
+    {
+      "epoch": 2.43,
+      "learning_rate": 3.9790575916230365e-06,
+      "loss": 0.5834,
+      "step": 484
+    },
+    {
+      "epoch": 2.43,
+      "learning_rate": 3.944153577661432e-06,
+      "loss": 0.5666,
+      "step": 485
+    },
+    {
+      "epoch": 2.44,
+      "learning_rate": 3.9092495636998255e-06,
+      "loss": 0.5744,
+      "step": 486
+    },
+    {
+      "epoch": 2.44,
+      "learning_rate": 3.874345549738221e-06,
+      "loss": 0.5704,
+      "step": 487
+    },
+    {
+      "epoch": 2.45,
+      "learning_rate": 3.8394415357766145e-06,
+      "loss": 0.5805,
+      "step": 488
+    },
+    {
+      "epoch": 2.45,
+      "learning_rate": 3.804537521815009e-06,
+      "loss": 0.5494,
+      "step": 489
+    },
+    {
+      "epoch": 2.46,
+      "learning_rate": 3.7696335078534035e-06,
+      "loss": 0.5564,
+      "step": 490
+    },
+    {
+      "epoch": 2.46,
+      "learning_rate": 3.7347294938917975e-06,
+      "loss": 0.5449,
+      "step": 491
+    },
+    {
+      "epoch": 2.47,
+      "learning_rate": 3.6998254799301924e-06,
+      "loss": 0.5635,
+      "step": 492
+    },
+    {
+      "epoch": 2.47,
+      "learning_rate": 3.6649214659685865e-06,
+      "loss": 0.567,
+      "step": 493
+    },
+    {
+      "epoch": 2.48,
+      "learning_rate": 3.630017452006981e-06,
+      "loss": 0.586,
+      "step": 494
+    },
+    {
+      "epoch": 2.48,
+      "learning_rate": 3.5951134380453755e-06,
+      "loss": 0.569,
+      "step": 495
+    },
+    {
+      "epoch": 2.49,
+      "learning_rate": 3.56020942408377e-06,
+      "loss": 0.5607,
+      "step": 496
+    },
+    {
+      "epoch": 2.49,
+      "learning_rate": 3.5253054101221645e-06,
+      "loss": 0.5737,
+      "step": 497
+    },
+    {
+      "epoch": 2.5,
+      "learning_rate": 3.4904013961605585e-06,
+      "loss": 0.5813,
+      "step": 498
+    },
+    {
+      "epoch": 2.5,
+      "learning_rate": 3.455497382198953e-06,
+      "loss": 0.5906,
+      "step": 499
+    },
+    {
+      "epoch": 2.51,
+      "learning_rate": 3.4205933682373475e-06,
+      "loss": 0.5948,
+      "step": 500
+    },
+    {
+      "epoch": 2.51,
+      "learning_rate": 3.385689354275742e-06,
+      "loss": 0.572,
+      "step": 501
+    },
+    {
+      "epoch": 2.52,
+      "learning_rate": 3.3507853403141365e-06,
+      "loss": 0.5577,
+      "step": 502
+    },
+    {
+      "epoch": 2.52,
+      "learning_rate": 3.315881326352531e-06,
+      "loss": 0.5611,
+      "step": 503
+    },
+    {
+      "epoch": 2.53,
+      "learning_rate": 3.280977312390925e-06,
+      "loss": 0.5753,
+      "step": 504
+    },
+    {
+      "epoch": 2.53,
+      "learning_rate": 3.2460732984293196e-06,
+      "loss": 0.5748,
+      "step": 505
+    },
+    {
+      "epoch": 2.54,
+      "learning_rate": 3.211169284467714e-06,
+      "loss": 0.556,
+      "step": 506
+    },
+    {
+      "epoch": 2.54,
+      "learning_rate": 3.1762652705061085e-06,
+      "loss": 0.5784,
+      "step": 507
+    },
+    {
+      "epoch": 2.55,
+      "learning_rate": 3.141361256544503e-06,
+      "loss": 0.5715,
+      "step": 508
+    },
+    {
+      "epoch": 2.55,
+      "learning_rate": 3.106457242582897e-06,
+      "loss": 0.5888,
+      "step": 509
+    },
+    {
+      "epoch": 2.56,
+      "learning_rate": 3.071553228621292e-06,
+      "loss": 0.5729,
+      "step": 510
+    },
+    {
+      "epoch": 2.56,
+      "learning_rate": 3.036649214659686e-06,
+      "loss": 0.5759,
+      "step": 511
+    },
+    {
+      "epoch": 2.57,
+      "learning_rate": 3.0017452006980806e-06,
+      "loss": 0.5763,
+      "step": 512
+    },
+    {
+      "epoch": 2.57,
+      "learning_rate": 2.966841186736475e-06,
+      "loss": 0.5574,
+      "step": 513
+    },
+    {
+      "epoch": 2.58,
+      "learning_rate": 2.931937172774869e-06,
+      "loss": 0.5927,
+      "step": 514
+    },
+    {
+      "epoch": 2.58,
+      "learning_rate": 2.897033158813264e-06,
+      "loss": 0.5717,
+      "step": 515
+    },
+    {
+      "epoch": 2.59,
+      "learning_rate": 2.862129144851658e-06,
+      "loss": 0.5553,
+      "step": 516
+    },
+    {
+      "epoch": 2.59,
+      "learning_rate": 2.8272251308900526e-06,
+      "loss": 0.5712,
+      "step": 517
+    },
+    {
+      "epoch": 2.6,
+      "learning_rate": 2.792321116928447e-06,
+      "loss": 0.5931,
+      "step": 518
+    },
+    {
+      "epoch": 2.6,
+      "learning_rate": 2.757417102966841e-06,
+      "loss": 0.5523,
+      "step": 519
+    },
+    {
+      "epoch": 2.61,
+      "learning_rate": 2.722513089005236e-06,
+      "loss": 0.5782,
+      "step": 520
+    },
+    {
+      "epoch": 2.61,
+      "learning_rate": 2.68760907504363e-06,
+      "loss": 0.5701,
+      "step": 521
+    },
+    {
+      "epoch": 2.62,
+      "learning_rate": 2.652705061082025e-06,
+      "loss": 0.5819,
+      "step": 522
+    },
+    {
+      "epoch": 2.62,
+      "learning_rate": 2.617801047120419e-06,
+      "loss": 0.583,
+      "step": 523
+    },
+    {
+      "epoch": 2.63,
+      "learning_rate": 2.582897033158813e-06,
+      "loss": 0.5425,
+      "step": 524
+    },
+    {
+      "epoch": 2.63,
+      "learning_rate": 2.547993019197208e-06,
+      "loss": 0.5527,
+      "step": 525
+    },
+    {
+      "epoch": 2.64,
+      "learning_rate": 2.513089005235602e-06,
+      "loss": 0.5525,
+      "step": 526
+    },
+    {
+      "epoch": 2.64,
+      "learning_rate": 2.4781849912739967e-06,
+      "loss": 0.5654,
+      "step": 527
+    },
+    {
+      "epoch": 2.65,
+      "learning_rate": 2.443280977312391e-06,
+      "loss": 0.5693,
+      "step": 528
+    },
+    {
+      "epoch": 2.65,
+      "learning_rate": 2.4083769633507856e-06,
+      "loss": 0.5796,
+      "step": 529
+    },
+    {
+      "epoch": 2.66,
+      "learning_rate": 2.37347294938918e-06,
+      "loss": 0.5918,
+      "step": 530
+    },
+    {
+      "epoch": 2.66,
+      "learning_rate": 2.338568935427574e-06,
+      "loss": 0.5669,
+      "step": 531
+    },
+    {
+      "epoch": 2.67,
+      "learning_rate": 2.3036649214659687e-06,
+      "loss": 0.5651,
+      "step": 532
+    },
+    {
+      "epoch": 2.67,
+      "learning_rate": 2.268760907504363e-06,
+      "loss": 0.5861,
+      "step": 533
+    },
+    {
+      "epoch": 2.68,
+      "learning_rate": 2.2338568935427577e-06,
+      "loss": 0.5591,
+      "step": 534
+    },
+    {
+      "epoch": 2.68,
+      "learning_rate": 2.198952879581152e-06,
+      "loss": 0.5594,
+      "step": 535
+    },
+    {
+      "epoch": 2.69,
+      "learning_rate": 2.1640488656195462e-06,
+      "loss": 0.5816,
+      "step": 536
+    },
+    {
+      "epoch": 2.69,
+      "learning_rate": 2.1291448516579407e-06,
+      "loss": 0.5988,
+      "step": 537
+    },
+    {
+      "epoch": 2.7,
+      "learning_rate": 2.094240837696335e-06,
+      "loss": 0.5642,
+      "step": 538
+    },
+    {
+      "epoch": 2.7,
+      "learning_rate": 2.0593368237347297e-06,
+      "loss": 0.579,
+      "step": 539
+    },
+    {
+      "epoch": 2.71,
+      "learning_rate": 2.024432809773124e-06,
+      "loss": 0.5742,
+      "step": 540
+    },
+    {
+      "epoch": 2.71,
+      "learning_rate": 1.9895287958115183e-06,
+      "loss": 0.5604,
+      "step": 541
+    },
+    {
+      "epoch": 2.72,
+      "learning_rate": 1.9546247818499127e-06,
+      "loss": 0.5751,
+      "step": 542
+    },
+    {
+      "epoch": 2.72,
+      "learning_rate": 1.9197207678883072e-06,
+      "loss": 0.5675,
+      "step": 543
+    },
+    {
+      "epoch": 2.73,
+      "learning_rate": 1.8848167539267017e-06,
+      "loss": 0.5614,
+      "step": 544
+    },
+    {
+      "epoch": 2.73,
+      "learning_rate": 1.8499127399650962e-06,
+      "loss": 0.5637,
+      "step": 545
+    },
+    {
+      "epoch": 2.74,
+      "learning_rate": 1.8150087260034905e-06,
+      "loss": 0.5702,
+      "step": 546
+    },
+    {
+      "epoch": 2.74,
+      "learning_rate": 1.780104712041885e-06,
+      "loss": 0.5569,
+      "step": 547
+    },
+    {
+      "epoch": 2.75,
+      "learning_rate": 1.7452006980802793e-06,
+      "loss": 0.5658,
+      "step": 548
+    },
+    {
+      "epoch": 2.75,
+      "learning_rate": 1.7102966841186738e-06,
+      "loss": 0.5593,
+      "step": 549
+    },
+    {
+      "epoch": 2.76,
+      "learning_rate": 1.6753926701570683e-06,
+      "loss": 0.5683,
+      "step": 550
+    },
+    {
+      "epoch": 2.76,
+      "learning_rate": 1.6404886561954625e-06,
+      "loss": 0.5894,
+      "step": 551
+    },
+    {
+      "epoch": 2.77,
+      "learning_rate": 1.605584642233857e-06,
+      "loss": 0.5648,
+      "step": 552
+    },
+    {
+      "epoch": 2.77,
+      "learning_rate": 1.5706806282722515e-06,
+      "loss": 0.5686,
+      "step": 553
+    },
+    {
+      "epoch": 2.78,
+      "learning_rate": 1.535776614310646e-06,
+      "loss": 0.5577,
+      "step": 554
+    },
+    {
+      "epoch": 2.78,
+      "learning_rate": 1.5008726003490403e-06,
+      "loss": 0.6011,
+      "step": 555
+    },
+    {
+      "epoch": 2.79,
+      "learning_rate": 1.4659685863874346e-06,
+      "loss": 0.5758,
+      "step": 556
+    },
+    {
+      "epoch": 2.79,
+      "learning_rate": 1.431064572425829e-06,
+      "loss": 0.5687,
+      "step": 557
+    },
+    {
+      "epoch": 2.8,
+      "learning_rate": 1.3961605584642235e-06,
+      "loss": 0.5788,
+      "step": 558
+    },
+    {
+      "epoch": 2.8,
+      "learning_rate": 1.361256544502618e-06,
+      "loss": 0.5724,
+      "step": 559
+    },
+    {
+      "epoch": 2.81,
+      "learning_rate": 1.3263525305410125e-06,
+      "loss": 0.5622,
+      "step": 560
+    },
+    {
+      "epoch": 2.81,
+      "learning_rate": 1.2914485165794066e-06,
+      "loss": 0.5632,
+      "step": 561
+    },
+    {
+      "epoch": 2.82,
+      "learning_rate": 1.256544502617801e-06,
+      "loss": 0.5673,
+      "step": 562
+    },
+    {
+      "epoch": 2.82,
+      "learning_rate": 1.2216404886561956e-06,
+      "loss": 0.5551,
+      "step": 563
+    },
+    {
+      "epoch": 2.83,
+      "learning_rate": 1.18673647469459e-06,
+      "loss": 0.5643,
+      "step": 564
+    },
+    {
+      "epoch": 2.83,
+      "learning_rate": 1.1518324607329843e-06,
+      "loss": 0.6011,
+      "step": 565
+    },
+    {
+      "epoch": 2.84,
+      "learning_rate": 1.1169284467713788e-06,
+      "loss": 0.572,
+      "step": 566
+    },
+    {
+      "epoch": 2.84,
+      "learning_rate": 1.0820244328097731e-06,
+      "loss": 0.5787,
+      "step": 567
+    },
+    {
+      "epoch": 2.85,
+      "learning_rate": 1.0471204188481676e-06,
+      "loss": 0.5684,
+      "step": 568
+    },
+    {
+      "epoch": 2.85,
+      "learning_rate": 1.012216404886562e-06,
+      "loss": 0.5653,
+      "step": 569
+    },
+    {
+      "epoch": 2.86,
+      "learning_rate": 9.773123909249564e-07,
+      "loss": 0.5756,
+      "step": 570
+    },
+    {
+      "epoch": 2.86,
+      "learning_rate": 9.424083769633509e-07,
+      "loss": 0.5688,
+      "step": 571
+    },
+    {
+      "epoch": 2.87,
+      "learning_rate": 9.075043630017452e-07,
+      "loss": 0.5673,
+      "step": 572
+    },
+    {
+      "epoch": 2.87,
+      "learning_rate": 8.726003490401396e-07,
+      "loss": 0.5678,
+      "step": 573
+    },
+    {
+      "epoch": 2.88,
+      "learning_rate": 8.376963350785341e-07,
+      "loss": 0.556,
+      "step": 574
+    },
+    {
+      "epoch": 2.88,
+      "learning_rate": 8.027923211169285e-07,
+      "loss": 0.5686,
+      "step": 575
+    },
+    {
+      "epoch": 2.89,
+      "learning_rate": 7.67888307155323e-07,
+      "loss": 0.5588,
+      "step": 576
+    },
+    {
+      "epoch": 2.89,
+      "learning_rate": 7.329842931937173e-07,
+      "loss": 0.5581,
+      "step": 577
+    },
+    {
+      "epoch": 2.9,
+      "learning_rate": 6.980802792321118e-07,
+      "loss": 0.5607,
+      "step": 578
+    },
+    {
+      "epoch": 2.9,
+      "learning_rate": 6.631762652705063e-07,
+      "loss": 0.5529,
+      "step": 579
+    },
+    {
+      "epoch": 2.91,
+      "learning_rate": 6.282722513089005e-07,
+      "loss": 0.5775,
+      "step": 580
+    },
+    {
+      "epoch": 2.91,
+      "learning_rate": 5.93368237347295e-07,
+      "loss": 0.5794,
+      "step": 581
+    },
+    {
+      "epoch": 2.92,
+      "learning_rate": 5.584642233856894e-07,
+      "loss": 0.5803,
+      "step": 582
+    },
+    {
+      "epoch": 2.92,
+      "learning_rate": 5.235602094240838e-07,
+      "loss": 0.5547,
+      "step": 583
+    },
+    {
+      "epoch": 2.93,
+      "learning_rate": 4.886561954624782e-07,
+      "loss": 0.5578,
+      "step": 584
+    },
+    {
+      "epoch": 2.93,
+      "learning_rate": 4.537521815008726e-07,
+      "loss": 0.5787,
+      "step": 585
+    },
+    {
+      "epoch": 2.94,
+      "learning_rate": 4.1884816753926706e-07,
+      "loss": 0.5916,
+      "step": 586
+    },
+    {
+      "epoch": 2.94,
+      "learning_rate": 3.839441535776615e-07,
+      "loss": 0.571,
+      "step": 587
+    },
+    {
+      "epoch": 2.95,
+      "learning_rate": 3.490401396160559e-07,
+      "loss": 0.5364,
+      "step": 588
+    },
+    {
+      "epoch": 2.95,
+      "learning_rate": 3.1413612565445027e-07,
+      "loss": 0.5824,
+      "step": 589
+    },
+    {
+      "epoch": 2.96,
+      "learning_rate": 2.792321116928447e-07,
+      "loss": 0.5747,
+      "step": 590
+    },
+    {
+      "epoch": 2.96,
+      "learning_rate": 2.443280977312391e-07,
+      "loss": 0.5588,
+      "step": 591
+    },
+    {
+      "epoch": 2.97,
+      "learning_rate": 2.0942408376963353e-07,
+      "loss": 0.5661,
+      "step": 592
+    },
+    {
+      "epoch": 2.97,
+      "learning_rate": 1.7452006980802794e-07,
+      "loss": 0.5612,
+      "step": 593
+    },
+    {
+      "epoch": 2.98,
+      "learning_rate": 1.3961605584642235e-07,
+      "loss": 0.5759,
+      "step": 594
+    },
+    {
+      "epoch": 2.98,
+      "learning_rate": 1.0471204188481677e-07,
+      "loss": 0.5689,
+      "step": 595
+    },
+    {
+      "epoch": 2.99,
+      "learning_rate": 6.980802792321118e-08,
+      "loss": 0.5575,
+      "step": 596
+    },
+    {
+      "epoch": 2.99,
+      "learning_rate": 3.490401396160559e-08,
+      "loss": 0.5573,
+      "step": 597
+    }
+  ],
+  "max_steps": 597,
+  "num_train_epochs": 3,
+  "total_flos": 826828684001280.0,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:15c7f6bae13d726117b591744d6f47b6e671352bdcb4dc9c6fcfa9aa75727c90
+size 6331

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 1, 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)