Upload folder using huggingface_hub

added_tokens.json

@@ -0,0 +1,6 @@
+{
+  "<|content|>": 32000,
+  "<|from|>": 32002,
+  "<|recipient|>": 32001,
+  "<|stop|>": 32003
+}

config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "/workspace/Mistral-7B-v0.1",
+  "architectures": [
+    "MistralForCausalLM"
+  ],
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 14336,
+  "max_position_embeddings": 32768,
+  "model_type": "mistral",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "sliding_window": 8192,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.35.2",
+  "use_cache": false,
+  "vocab_size": 32004
+}

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "transformers_version": "4.35.2"
+}

latest

@@ -0,0 +1 @@
+global_step391

model-00001-of-00003.safetensors

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

model-00002-of-00003.safetensors

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

model-00003-of-00003.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,298 @@
+{
+  "metadata": {
+    "total_size": 14483529728
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00003-of-00003.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.29.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.30.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.31.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.norm.weight": "model-00003-of-00003.safetensors"
+  }
+}

rng_state_0.pth

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

rng_state_1.pth

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

rng_state_2.pth

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

rng_state_3.pth

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

scheduler.pt

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

special_tokens_map.json

@@ -0,0 +1,54 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<|content|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|recipient|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|from|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|stop|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "</s>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.model

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

tokenizer_config.json

@@ -0,0 +1,79 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32000": {
+      "content": "<|content|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32001": {
+      "content": "<|recipient|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32002": {
+      "content": "<|from|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32003": {
+      "content": "<|stop|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|content|>",
+    "<|recipient|>",
+    "<|from|>",
+    "<|stop|>"
+  ],
+  "bos_token": "<s>",
+  "chat_template": "{% for message in messages %}\n{% if message['role'] == 'user' or message['role'] == 'system' %}\n{{ '<|from|>' + message['role'] + '\n<|recipient|>all\n<|content|>' + message['content'] + '\n' }}{% elif message['role'] == 'tool' %}\n{{ '<|from|>' + message['name'] + '\n<|recipient|>all\n<|content|>' + message['content'] + '\n' }}{% else %}\n{% set contain_content='no'%}\n{% if message['content'] is not none %}\n{{ '<|from|>assistant\n<|recipient|>all\n<|content|>' + message['content'] }}{% set contain_content='yes'%}\n{% endif %}\n{% if 'tool_calls' in message and message['tool_calls'] is not none %}\n{% for tool_call in message['tool_calls'] %}\n{% set prompt='<|from|>assistant\n<|recipient|>' + tool_call['function']['name'] + '\n<|content|>' + tool_call['function']['arguments'] %}\n{% if loop.index == 1 and contain_content == \"no\" %}\n{{ prompt }}{% else %}\n{{ '\n' + prompt}}{% endif %}\n{% endfor %}\n{% endif %}\n{{ '<|stop|>\n' }}{% endif %}\n{% endfor %}\n{% if add_generation_prompt %}{{ '<|from|>assistant\n<|recipient|>' }}{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": true,
+  "model_max_length": 8192,
+  "pad_token": "</s>",
+  "padding_side": "left",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

trainer_state.json

@@ -0,0 +1,2455 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 1.0,
+  "eval_steps": 75,
+  "global_step": 391,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0,
+      "learning_rate": 7.5e-07,
+      "loss": 1.0384,
+      "step": 1
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.5e-06,
+      "loss": 1.0344,
+      "step": 2
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 2.25e-06,
+      "loss": 0.8877,
+      "step": 3
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 3e-06,
+      "loss": 0.9107,
+      "step": 4
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 3.75e-06,
+      "loss": 0.8865,
+      "step": 5
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 4.5e-06,
+      "loss": 1.0755,
+      "step": 6
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 5.2500000000000006e-06,
+      "loss": 0.9334,
+      "step": 7
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 6e-06,
+      "loss": 0.8215,
+      "step": 8
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 6.75e-06,
+      "loss": 0.9601,
+      "step": 9
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 7.5e-06,
+      "loss": 0.7402,
+      "step": 10
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 8.25e-06,
+      "loss": 0.8968,
+      "step": 11
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 9e-06,
+      "loss": 0.7669,
+      "step": 12
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 8.999845402895058e-06,
+      "loss": 0.7714,
+      "step": 13
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 8.999381622202572e-06,
+      "loss": 0.6609,
+      "step": 14
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 8.998608689788832e-06,
+      "loss": 0.8098,
+      "step": 15
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 8.997526658761886e-06,
+      "loss": 0.8355,
+      "step": 16
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 8.996135603467899e-06,
+      "loss": 0.8318,
+      "step": 17
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.994435619486036e-06,
+      "loss": 0.9067,
+      "step": 18
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.992426823621897e-06,
+      "loss": 0.7113,
+      "step": 19
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.99010935389949e-06,
+      "loss": 0.804,
+      "step": 20
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 8.987483369551757e-06,
+      "loss": 0.776,
+      "step": 21
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 8.984549051009623e-06,
+      "loss": 0.8339,
+      "step": 22
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 8.981306599889595e-06,
+      "loss": 0.833,
+      "step": 23
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 8.977756238979921e-06,
+      "loss": 0.7614,
+      "step": 24
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 8.973898212225277e-06,
+      "loss": 0.7553,
+      "step": 25
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 8.969732784710005e-06,
+      "loss": 0.6437,
+      "step": 26
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 8.965260242639897e-06,
+      "loss": 0.7507,
+      "step": 27
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 8.960480893322532e-06,
+      "loss": 0.722,
+      "step": 28
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 8.955395065146165e-06,
+      "loss": 0.775,
+      "step": 29
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 8.950003107557156e-06,
+      "loss": 0.6426,
+      "step": 30
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 8.944305391035962e-06,
+      "loss": 0.5657,
+      "step": 31
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 8.938302307071689e-06,
+      "loss": 0.7168,
+      "step": 32
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 8.93199426813518e-06,
+      "loss": 0.6611,
+      "step": 33
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 8.92538170765068e-06,
+      "loss": 0.5408,
+      "step": 34
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 8.918465079966063e-06,
+      "loss": 0.7239,
+      "step": 35
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 8.9112448603216e-06,
+      "loss": 0.7005,
+      "step": 36
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 8.90372154481732e-06,
+      "loss": 0.5097,
+      "step": 37
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 8.895895650378904e-06,
+      "loss": 0.584,
+      "step": 38
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 8.887767714722187e-06,
+      "loss": 0.7294,
+      "step": 39
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 8.879338296316203e-06,
+      "loss": 0.7048,
+      "step": 40
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 8.870607974344808e-06,
+      "loss": 0.7611,
+      "step": 41
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 8.861577348666893e-06,
+      "loss": 0.5931,
+      "step": 42
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 8.852247039775163e-06,
+      "loss": 0.7365,
+      "step": 43
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 8.842617688753502e-06,
+      "loss": 0.6636,
+      "step": 44
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 8.832689957232929e-06,
+      "loss": 0.6412,
+      "step": 45
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 8.822464527346134e-06,
+      "loss": 0.6989,
+      "step": 46
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 8.81194210168061e-06,
+      "loss": 0.7198,
+      "step": 47
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 8.801123403230374e-06,
+      "loss": 0.6605,
+      "step": 48
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 8.7900091753463e-06,
+      "loss": 0.6546,
+      "step": 49
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 8.778600181685031e-06,
+      "loss": 0.7965,
+      "step": 50
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 8.766897206156524e-06,
+      "loss": 0.6213,
+      "step": 51
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 8.754901052870166e-06,
+      "loss": 0.6297,
+      "step": 52
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 8.742612546079549e-06,
+      "loss": 0.78,
+      "step": 53
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 8.730032530125813e-06,
+      "loss": 0.6817,
+      "step": 54
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 8.717161869379647e-06,
+      "loss": 0.6632,
+      "step": 55
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 8.70400144818189e-06,
+      "loss": 0.5498,
+      "step": 56
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 8.690552170782772e-06,
+      "loss": 0.6081,
+      "step": 57
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 8.676814961279782e-06,
+      "loss": 0.4716,
+      "step": 58
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 8.662790763554175e-06,
+      "loss": 0.5989,
+      "step": 59
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 8.64848054120611e-06,
+      "loss": 0.535,
+      "step": 60
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 8.633885277488455e-06,
+      "loss": 0.5284,
+      "step": 61
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 8.619005975239218e-06,
+      "loss": 0.7458,
+      "step": 62
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 8.603843656812642e-06,
+      "loss": 0.6758,
+      "step": 63
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 8.588399364008963e-06,
+      "loss": 0.6736,
+      "step": 64
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 8.57267415800283e-06,
+      "loss": 0.4975,
+      "step": 65
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 8.556669119270387e-06,
+      "loss": 0.5586,
+      "step": 66
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 8.540385347515032e-06,
+      "loss": 0.6974,
+      "step": 67
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 8.523823961591867e-06,
+      "loss": 0.6778,
+      "step": 68
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 8.506986099430807e-06,
+      "loss": 0.6791,
+      "step": 69
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 8.48987291795841e-06,
+      "loss": 0.5207,
+      "step": 70
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 8.472485593018366e-06,
+      "loss": 0.5651,
+      "step": 71
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 8.45482531929072e-06,
+      "loss": 0.6487,
+      "step": 72
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 8.436893310209779e-06,
+      "loss": 0.6202,
+      "step": 73
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 8.418690797880737e-06,
+      "loss": 0.5263,
+      "step": 74
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 8.400219032995022e-06,
+      "loss": 0.7257,
+      "step": 75
+    },
+    {
+      "epoch": 0.19,
+      "eval_accuracy": 0.814565802103578,
+      "eval_accuracy_<|content|>": 0.9681089145841104,
+      "eval_accuracy_<|from|>": 0.9962073324905183,
+      "eval_accuracy_<|recipient|>": 0.9987357774968394,
+      "eval_accuracy_<|stop|>": 0.8882521489971347,
+      "eval_accuracy_total_num_<|content|>": 5362,
+      "eval_accuracy_total_num_<|from|>": 791,
+      "eval_accuracy_total_num_<|recipient|>": 791,
+      "eval_accuracy_total_num_<|stop|>": 4537,
+      "eval_loss": 0.618374228477478,
+      "eval_perplexity": 1.0682635492781798,
+      "eval_runtime": 290.3169,
+      "eval_samples_per_second": 4.736,
+      "eval_steps_per_second": 0.148,
+      "step": 75
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 8.381479284744354e-06,
+      "loss": 0.6557,
+      "step": 76
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 8.362472840733548e-06,
+      "loss": 0.6288,
+      "step": 77
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 8.343201006892032e-06,
+      "loss": 0.5379,
+      "step": 78
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 8.323665107384127e-06,
+      "loss": 0.6197,
+      "step": 79
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 8.303866484518059e-06,
+      "loss": 0.5107,
+      "step": 80
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 8.283806498653725e-06,
+      "loss": 0.67,
+      "step": 81
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 8.263486528109237e-06,
+      "loss": 0.6076,
+      "step": 82
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 8.242907969066198e-06,
+      "loss": 0.6506,
+      "step": 83
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 8.22207223547379e-06,
+      "loss": 0.6114,
+      "step": 84
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 8.200980758951609e-06,
+      "loss": 0.4775,
+      "step": 85
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 8.179634988691303e-06,
+      "loss": 0.6625,
+      "step": 86
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 8.158036391357e-06,
+      "loss": 0.5161,
+      "step": 87
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 8.136186450984527e-06,
+      "loss": 0.4956,
+      "step": 88
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 8.114086668879454e-06,
+      "loss": 0.6379,
+      "step": 89
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 8.09173856351393e-06,
+      "loss": 0.6708,
+      "step": 90
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 8.069143670422347e-06,
+      "loss": 0.5111,
+      "step": 91
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 8.046303542095846e-06,
+      "loss": 0.5079,
+      "step": 92
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 8.02321974787563e-06,
+      "loss": 0.5029,
+      "step": 93
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 7.999893873845152e-06,
+      "loss": 0.4098,
+      "step": 94
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 7.976327522721114e-06,
+      "loss": 0.4683,
+      "step": 95
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 7.952522313743371e-06,
+      "loss": 0.5922,
+      "step": 96
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 7.928479882563648e-06,
+      "loss": 0.4946,
+      "step": 97
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 7.904201881133171e-06,
+      "loss": 0.4977,
+      "step": 98
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 7.879689977589154e-06,
+      "loss": 0.6153,
+      "step": 99
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 7.85494585614019e-06,
+      "loss": 0.5741,
+      "step": 100
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 7.829971216950513e-06,
+      "loss": 0.4949,
+      "step": 101
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 7.804767776023202e-06,
+      "loss": 0.5977,
+      "step": 102
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 7.779337265082256e-06,
+      "loss": 0.6465,
+      "step": 103
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 7.753681431453614e-06,
+      "loss": 0.6386,
+      "step": 104
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 7.727802037945105e-06,
+      "loss": 0.6208,
+      "step": 105
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 7.70170086272531e-06,
+      "loss": 0.6551,
+      "step": 106
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 7.675379699201395e-06,
+      "loss": 0.662,
+      "step": 107
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 7.648840355895885e-06,
+      "loss": 0.5964,
+      "step": 108
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 7.6220846563224e-06,
+      "loss": 0.6025,
+      "step": 109
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 7.595114438860358e-06,
+      "loss": 0.5055,
+      "step": 110
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 7.567931556628665e-06,
+      "loss": 0.5483,
+      "step": 111
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 7.540537877358389e-06,
+      "loss": 0.6338,
+      "step": 112
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 7.51293528326442e-06,
+      "loss": 0.6748,
+      "step": 113
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 7.485125670916155e-06,
+      "loss": 0.5733,
+      "step": 114
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 7.4571109511071714e-06,
+      "loss": 0.6164,
+      "step": 115
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 7.428893048723952e-06,
+      "loss": 0.6496,
+      "step": 116
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 7.400473902613611e-06,
+      "loss": 0.6056,
+      "step": 117
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 7.371855465450694e-06,
+      "loss": 0.4862,
+      "step": 118
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 7.343039703602988e-06,
+      "loss": 0.5984,
+      "step": 119
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 7.314028596996431e-06,
+      "loss": 0.5891,
+      "step": 120
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 7.284824138979066e-06,
+      "loss": 0.5345,
+      "step": 121
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 7.255428336184075e-06,
+      "loss": 0.7941,
+      "step": 122
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 7.2258432083919064e-06,
+      "loss": 0.552,
+      "step": 123
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 7.196070788391497e-06,
+      "loss": 0.599,
+      "step": 124
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 7.166113121840595e-06,
+      "loss": 0.4727,
+      "step": 125
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 7.135972267125212e-06,
+      "loss": 0.545,
+      "step": 126
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 7.1056502952181815e-06,
+      "loss": 0.6106,
+      "step": 127
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 7.075149289536871e-06,
+      "loss": 0.4726,
+      "step": 128
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 7.044471345800024e-06,
+      "loss": 0.5416,
+      "step": 129
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 7.0136185718837685e-06,
+      "loss": 0.6889,
+      "step": 130
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 6.9825930876767834e-06,
+      "loss": 0.594,
+      "step": 131
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 6.951397024934641e-06,
+      "loss": 0.5456,
+      "step": 132
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 6.920032527133334e-06,
+      "loss": 0.4282,
+      "step": 133
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 6.888501749322002e-06,
+      "loss": 0.6089,
+      "step": 134
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 6.856806857974848e-06,
+      "loss": 0.6507,
+      "step": 135
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 6.824950030842293e-06,
+      "loss": 0.6411,
+      "step": 136
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 6.792933456801339e-06,
+      "loss": 0.5493,
+      "step": 137
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 6.760759335705163e-06,
+      "loss": 0.4705,
+      "step": 138
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 6.728429878231978e-06,
+      "loss": 0.4708,
+      "step": 139
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 6.695947305733131e-06,
+      "loss": 0.5714,
+      "step": 140
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 6.6633138500804735e-06,
+      "loss": 0.5466,
+      "step": 141
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 6.630531753513014e-06,
+      "loss": 0.4242,
+      "step": 142
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 6.597603268482853e-06,
+      "loss": 0.6997,
+      "step": 143
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 6.5645306575004145e-06,
+      "loss": 0.6938,
+      "step": 144
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 6.531316192978991e-06,
+      "loss": 0.7704,
+      "step": 145
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 6.497962157078611e-06,
+      "loss": 0.5127,
+      "step": 146
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 6.4644708415492205e-06,
+      "loss": 0.6113,
+      "step": 147
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 6.4308445475732315e-06,
+      "loss": 0.563,
+      "step": 148
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 6.397085585607401e-06,
+      "loss": 0.6154,
+      "step": 149
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 6.3631962752240746e-06,
+      "loss": 0.5561,
+      "step": 150
+    },
+    {
+      "epoch": 0.38,
+      "eval_accuracy": 0.819768782794211,
+      "eval_accuracy_<|content|>": 1.0,
+      "eval_accuracy_<|from|>": 0.9393173198482933,
+      "eval_accuracy_<|recipient|>": 1.0,
+      "eval_accuracy_<|stop|>": 0.8756887811329072,
+      "eval_accuracy_total_num_<|content|>": 5362,
+      "eval_accuracy_total_num_<|from|>": 791,
+      "eval_accuracy_total_num_<|recipient|>": 791,
+      "eval_accuracy_total_num_<|stop|>": 4537,
+      "eval_loss": 0.586061418056488,
+      "eval_perplexity": 1.0650794931769392,
+      "eval_runtime": 223.6772,
+      "eval_samples_per_second": 6.147,
+      "eval_steps_per_second": 0.192,
+      "step": 150
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 6.32917894495182e-06,
+      "loss": 0.4405,
+      "step": 151
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 6.295035932115428e-06,
+      "loss": 0.3631,
+      "step": 152
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 6.260769582675315e-06,
+      "loss": 0.6282,
+      "step": 153
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 6.226382251066333e-06,
+      "loss": 0.7472,
+      "step": 154
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 6.191876300036003e-06,
+      "loss": 0.5803,
+      "step": 155
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 6.1572541004821585e-06,
+      "loss": 0.5861,
+      "step": 156
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 6.122518031290052e-06,
+      "loss": 0.6237,
+      "step": 157
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 6.0876704791689e-06,
+      "loss": 0.5705,
+      "step": 158
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 6.0527138384878885e-06,
+      "loss": 0.6131,
+      "step": 159
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 6.017650511111662e-06,
+      "loss": 0.6446,
+      "step": 160
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 5.98248290623529e-06,
+      "loss": 0.5882,
+      "step": 161
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 5.947213440218726e-06,
+      "loss": 0.5134,
+      "step": 162
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 5.911844536420788e-06,
+      "loss": 0.6261,
+      "step": 163
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 5.8763786250326455e-06,
+      "loss": 0.5609,
+      "step": 164
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 5.8408181429108425e-06,
+      "loss": 0.5729,
+      "step": 165
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 5.805165533409863e-06,
+      "loss": 0.5378,
+      "step": 166
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 5.769423246214248e-06,
+      "loss": 0.5914,
+      "step": 167
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 5.733593737170271e-06,
+      "loss": 0.5126,
+      "step": 168
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 5.6976794681172075e-06,
+      "loss": 0.5309,
+      "step": 169
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 5.661682906718181e-06,
+      "loss": 0.5544,
+      "step": 170
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 5.625606526290604e-06,
+      "loss": 0.552,
+      "step": 171
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 5.5894528056362406e-06,
+      "loss": 0.5787,
+      "step": 172
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 5.55322422887089e-06,
+      "loss": 0.5153,
+      "step": 173
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 5.516923285253701e-06,
+      "loss": 0.4472,
+      "step": 174
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 5.4805524690161325e-06,
+      "loss": 0.6333,
+      "step": 175
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 5.444114279190586e-06,
+      "loss": 0.4386,
+      "step": 176
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 5.407611219438685e-06,
+      "loss": 0.613,
+      "step": 177
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 5.371045797879255e-06,
+      "loss": 0.5707,
+      "step": 178
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 5.334420526915993e-06,
+      "loss": 0.5866,
+      "step": 179
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 5.297737923064836e-06,
+      "loss": 0.4922,
+      "step": 180
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 5.261000506781051e-06,
+      "loss": 0.5826,
+      "step": 181
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 5.224210802286064e-06,
+      "loss": 0.6741,
+      "step": 182
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 5.187371337394009e-06,
+      "loss": 0.3813,
+      "step": 183
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 5.150484643338051e-06,
+      "loss": 0.6125,
+      "step": 184
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 5.113553254596463e-06,
+      "loss": 0.66,
+      "step": 185
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 5.076579708718481e-06,
+      "loss": 0.5914,
+      "step": 186
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 5.039566546149946e-06,
+      "loss": 0.4527,
+      "step": 187
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 5.002516310058766e-06,
+      "loss": 0.5896,
+      "step": 188
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 4.965431546160153e-06,
+      "loss": 0.4518,
+      "step": 189
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 4.928314802541726e-06,
+      "loss": 0.5798,
+      "step": 190
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 4.891168629488419e-06,
+      "loss": 0.5669,
+      "step": 191
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 4.853995579307262e-06,
+      "loss": 0.6075,
+      "step": 192
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 4.816798206152006e-06,
+      "loss": 0.6057,
+      "step": 193
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 4.77957906584763e-06,
+      "loss": 0.5076,
+      "step": 194
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 4.742340715714727e-06,
+      "loss": 0.5997,
+      "step": 195
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 4.705085714393797e-06,
+      "loss": 0.5709,
+      "step": 196
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 4.667816621669442e-06,
+      "loss": 0.5029,
+      "step": 197
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 4.630535998294477e-06,
+      "loss": 0.5885,
+      "step": 198
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 4.5932464058139885e-06,
+      "loss": 0.5776,
+      "step": 199
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 4.55595040638933e-06,
+      "loss": 0.4949,
+      "step": 200
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 4.5186505626220725e-06,
+      "loss": 0.4202,
+      "step": 201
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 4.48134943737793e-06,
+      "loss": 0.5012,
+      "step": 202
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 4.444049593610671e-06,
+      "loss": 0.613,
+      "step": 203
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 4.406753594186011e-06,
+      "loss": 0.4806,
+      "step": 204
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 4.369464001705524e-06,
+      "loss": 0.6567,
+      "step": 205
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 4.332183378330558e-06,
+      "loss": 0.3533,
+      "step": 206
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 4.294914285606203e-06,
+      "loss": 0.7252,
+      "step": 207
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 4.257659284285274e-06,
+      "loss": 0.4848,
+      "step": 208
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 4.220420934152371e-06,
+      "loss": 0.3598,
+      "step": 209
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 4.1832017938479936e-06,
+      "loss": 0.5295,
+      "step": 210
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 4.146004420692739e-06,
+      "loss": 0.5091,
+      "step": 211
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 4.108831370511581e-06,
+      "loss": 0.5788,
+      "step": 212
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 4.071685197458274e-06,
+      "loss": 0.462,
+      "step": 213
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 4.034568453839847e-06,
+      "loss": 0.5832,
+      "step": 214
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 3.997483689941234e-06,
+      "loss": 0.4365,
+      "step": 215
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 3.960433453850053e-06,
+      "loss": 0.6247,
+      "step": 216
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 3.923420291281522e-06,
+      "loss": 0.538,
+      "step": 217
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 3.886446745403538e-06,
+      "loss": 0.5798,
+      "step": 218
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 3.849515356661949e-06,
+      "loss": 0.568,
+      "step": 219
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 3.8126286626059916e-06,
+      "loss": 0.4764,
+      "step": 220
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 3.7757891977139374e-06,
+      "loss": 0.5134,
+      "step": 221
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 3.738999493218949e-06,
+      "loss": 0.5493,
+      "step": 222
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 3.7022620769351665e-06,
+      "loss": 0.5529,
+      "step": 223
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 3.665579473084008e-06,
+      "loss": 0.4234,
+      "step": 224
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 3.6289542021207454e-06,
+      "loss": 0.6216,
+      "step": 225
+    },
+    {
+      "epoch": 0.58,
+      "eval_accuracy": 0.8257179536751563,
+      "eval_accuracy_<|content|>": 1.0,
+      "eval_accuracy_<|from|>": 0.9924146649810367,
+      "eval_accuracy_<|recipient|>": 1.0,
+      "eval_accuracy_<|stop|>": 0.9440158695173022,
+      "eval_accuracy_total_num_<|content|>": 5362,
+      "eval_accuracy_total_num_<|from|>": 791,
+      "eval_accuracy_total_num_<|recipient|>": 791,
+      "eval_accuracy_total_num_<|stop|>": 4537,
+      "eval_loss": 0.5591428279876709,
+      "eval_perplexity": 1.0623002333585039,
+      "eval_runtime": 223.7635,
+      "eval_samples_per_second": 6.145,
+      "eval_steps_per_second": 0.192,
+      "step": 225
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 3.5923887805613165e-06,
+      "loss": 0.6102,
+      "step": 226
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 3.5558857208094145e-06,
+      "loss": 0.5378,
+      "step": 227
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 3.5194475309838677e-06,
+      "loss": 0.6289,
+      "step": 228
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 3.4830767147463014e-06,
+      "loss": 0.3814,
+      "step": 229
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 3.4467757711291108e-06,
+      "loss": 0.5506,
+      "step": 230
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 3.4105471943637592e-06,
+      "loss": 0.5942,
+      "step": 231
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 3.374393473709396e-06,
+      "loss": 0.5092,
+      "step": 232
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 3.3383170932818184e-06,
+      "loss": 0.5533,
+      "step": 233
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 3.302320531882792e-06,
+      "loss": 0.454,
+      "step": 234
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 3.266406262829731e-06,
+      "loss": 0.5039,
+      "step": 235
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 3.2305767537857534e-06,
+      "loss": 0.5447,
+      "step": 236
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 3.194834466590136e-06,
+      "loss": 0.4985,
+      "step": 237
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 3.1591818570891573e-06,
+      "loss": 0.6181,
+      "step": 238
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 3.1236213749673555e-06,
+      "loss": 0.5648,
+      "step": 239
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 3.0881554635792124e-06,
+      "loss": 0.6169,
+      "step": 240
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 3.0527865597812755e-06,
+      "loss": 0.6264,
+      "step": 241
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 3.017517093764711e-06,
+      "loss": 0.5292,
+      "step": 242
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 2.982349488888337e-06,
+      "loss": 0.5578,
+      "step": 243
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 2.9472861615121117e-06,
+      "loss": 0.4992,
+      "step": 244
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 2.9123295208311007e-06,
+      "loss": 0.5403,
+      "step": 245
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 2.877481968709948e-06,
+      "loss": 0.5385,
+      "step": 246
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 2.842745899517843e-06,
+      "loss": 0.4517,
+      "step": 247
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 2.8081236999639976e-06,
+      "loss": 0.5326,
+      "step": 248
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 2.7736177489336662e-06,
+      "loss": 0.4631,
+      "step": 249
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 2.739230417324686e-06,
+      "loss": 0.5639,
+      "step": 250
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 2.7049640678845723e-06,
+      "loss": 0.4956,
+      "step": 251
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 2.6708210550481794e-06,
+      "loss": 0.5281,
+      "step": 252
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 2.636803724775927e-06,
+      "loss": 0.2976,
+      "step": 253
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 2.6029144143926003e-06,
+      "loss": 0.5906,
+      "step": 254
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 2.569155452426767e-06,
+      "loss": 0.5625,
+      "step": 255
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 2.5355291584507814e-06,
+      "loss": 0.5878,
+      "step": 256
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 2.5020378429213904e-06,
+      "loss": 0.595,
+      "step": 257
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 2.4686838070210094e-06,
+      "loss": 0.3659,
+      "step": 258
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 2.435469342499587e-06,
+      "loss": 0.5602,
+      "step": 259
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 2.402396731517147e-06,
+      "loss": 0.4733,
+      "step": 260
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 2.3694682464869856e-06,
+      "loss": 0.5969,
+      "step": 261
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 2.336686149919527e-06,
+      "loss": 0.4799,
+      "step": 262
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 2.3040526942668707e-06,
+      "loss": 0.3591,
+      "step": 263
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 2.271570121768021e-06,
+      "loss": 0.4472,
+      "step": 264
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 2.239240664294837e-06,
+      "loss": 0.4683,
+      "step": 265
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 2.2070665431986617e-06,
+      "loss": 0.4283,
+      "step": 266
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 2.1750499691577055e-06,
+      "loss": 0.6666,
+      "step": 267
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 2.1431931420251543e-06,
+      "loss": 0.5802,
+      "step": 268
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 2.1114982506779997e-06,
+      "loss": 0.5939,
+      "step": 269
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 2.0799674728666665e-06,
+      "loss": 0.5326,
+      "step": 270
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 2.0486029750653605e-06,
+      "loss": 0.6165,
+      "step": 271
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 2.017406912323217e-06,
+      "loss": 0.557,
+      "step": 272
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 1.986381428116232e-06,
+      "loss": 0.5167,
+      "step": 273
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 1.9555286541999766e-06,
+      "loss": 0.6048,
+      "step": 274
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 1.92485071046313e-06,
+      "loss": 0.4407,
+      "step": 275
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 1.8943497047818179e-06,
+      "loss": 0.6113,
+      "step": 276
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 1.864027732874788e-06,
+      "loss": 0.7606,
+      "step": 277
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 1.8338868781594052e-06,
+      "loss": 0.6256,
+      "step": 278
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 1.8039292116085027e-06,
+      "loss": 0.6035,
+      "step": 279
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 1.7741567916080946e-06,
+      "loss": 0.5717,
+      "step": 280
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 1.7445716638159248e-06,
+      "loss": 0.475,
+      "step": 281
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 1.715175861020934e-06,
+      "loss": 0.4888,
+      "step": 282
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 1.6859714030035694e-06,
+      "loss": 0.5913,
+      "step": 283
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 1.6569602963970123e-06,
+      "loss": 0.3934,
+      "step": 284
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 1.628144534549307e-06,
+      "loss": 0.5487,
+      "step": 285
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 1.5995260973863885e-06,
+      "loss": 0.494,
+      "step": 286
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 1.5711069512760497e-06,
+      "loss": 0.4315,
+      "step": 287
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 1.5428890488928284e-06,
+      "loss": 0.4848,
+      "step": 288
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 1.5148743290838455e-06,
+      "loss": 0.4264,
+      "step": 289
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 1.4870647167355795e-06,
+      "loss": 0.6395,
+      "step": 290
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 1.4594621226416103e-06,
+      "loss": 0.5726,
+      "step": 291
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 1.4320684433713354e-06,
+      "loss": 0.5543,
+      "step": 292
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 1.4048855611396424e-06,
+      "loss": 0.6457,
+      "step": 293
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 1.3779153436776006e-06,
+      "loss": 0.3915,
+      "step": 294
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 1.351159644104115e-06,
+      "loss": 0.5187,
+      "step": 295
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 1.3246203007986048e-06,
+      "loss": 0.6015,
+      "step": 296
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 1.298299137274692e-06,
+      "loss": 0.6067,
+      "step": 297
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 1.2721979620548955e-06,
+      "loss": 0.4431,
+      "step": 298
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 1.2463185685463859e-06,
+      "loss": 0.3615,
+      "step": 299
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 1.220662734917746e-06,
+      "loss": 0.5003,
+      "step": 300
+    },
+    {
+      "epoch": 0.77,
+      "eval_accuracy": 0.8336343532396525,
+      "eval_accuracy_<|content|>": 1.0,
+      "eval_accuracy_<|from|>": 0.9873577749683944,
+      "eval_accuracy_<|recipient|>": 1.0,
+      "eval_accuracy_<|stop|>": 0.91359929468812,
+      "eval_accuracy_total_num_<|content|>": 5362,
+      "eval_accuracy_total_num_<|from|>": 791,
+      "eval_accuracy_total_num_<|recipient|>": 791,
+      "eval_accuracy_total_num_<|stop|>": 4537,
+      "eval_loss": 0.5285482406616211,
+      "eval_perplexity": 1.0587839466417142,
+      "eval_runtime": 223.59,
+      "eval_samples_per_second": 6.15,
+      "eval_steps_per_second": 0.192,
+      "step": 300
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 1.1952322239767983e-06,
+      "loss": 0.3402,
+      "step": 301
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 1.170028783049487e-06,
+      "loss": 0.6256,
+      "step": 302
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 1.1450541438598118e-06,
+      "loss": 0.5378,
+      "step": 303
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 1.1203100224108464e-06,
+      "loss": 0.5334,
+      "step": 304
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 1.09579811886683e-06,
+      "loss": 0.4257,
+      "step": 305
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 1.0715201174363525e-06,
+      "loss": 0.4915,
+      "step": 306
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 1.0474776862566299e-06,
+      "loss": 0.4102,
+      "step": 307
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 1.0236724772788846e-06,
+      "loss": 0.3463,
+      "step": 308
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 1.00010612615485e-06,
+      "loss": 0.5686,
+      "step": 309
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 9.76780252124369e-07,
+      "loss": 0.5914,
+      "step": 310
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 9.536964579041548e-07,
+      "loss": 0.6401,
+      "step": 311
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 9.308563295776531e-07,
+      "loss": 0.5474,
+      "step": 312
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 9.082614364860701e-07,
+      "loss": 0.4519,
+      "step": 313
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 8.859133311205453e-07,
+      "loss": 0.5629,
+      "step": 314
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 8.638135490154735e-07,
+      "loss": 0.4641,
+      "step": 315
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 8.419636086430022e-07,
+      "loss": 0.4696,
+      "step": 316
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 8.203650113086972e-07,
+      "loss": 0.538,
+      "step": 317
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 7.990192410483916e-07,
+      "loss": 0.6092,
+      "step": 318
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 7.779277645262103e-07,
+      "loss": 0.407,
+      "step": 319
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 7.570920309338014e-07,
+      "loss": 0.5779,
+      "step": 320
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 7.365134718907647e-07,
+      "loss": 0.5572,
+      "step": 321
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 7.161935013462746e-07,
+      "loss": 0.6165,
+      "step": 322
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 6.961335154819422e-07,
+      "loss": 0.5908,
+      "step": 323
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 6.763348926158732e-07,
+      "loss": 0.5739,
+      "step": 324
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 6.567989931079675e-07,
+      "loss": 0.5375,
+      "step": 325
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 6.375271592664525e-07,
+      "loss": 0.5726,
+      "step": 326
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 6.18520715255646e-07,
+      "loss": 0.5559,
+      "step": 327
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 5.997809670049795e-07,
+      "loss": 0.5309,
+      "step": 328
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 5.813092021192639e-07,
+      "loss": 0.5068,
+      "step": 329
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 5.631066897902227e-07,
+      "loss": 0.4953,
+      "step": 330
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 5.451746807092811e-07,
+      "loss": 0.5307,
+      "step": 331
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 5.275144069816338e-07,
+      "loss": 0.5182,
+      "step": 332
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 5.101270820415908e-07,
+      "loss": 0.5599,
+      "step": 333
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 4.930139005691915e-07,
+      "loss": 0.4452,
+      "step": 334
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 4.761760384081335e-07,
+      "loss": 0.5605,
+      "step": 335
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 4.596146524849682e-07,
+      "loss": 0.3468,
+      "step": 336
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 4.433308807296139e-07,
+      "loss": 0.5684,
+      "step": 337
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 4.2732584199717023e-07,
+      "loss": 0.3672,
+      "step": 338
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 4.116006359910375e-07,
+      "loss": 0.5228,
+      "step": 339
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 3.961563431873597e-07,
+      "loss": 0.4777,
+      "step": 340
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 3.809940247607826e-07,
+      "loss": 0.4455,
+      "step": 341
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 3.661147225115447e-07,
+      "loss": 0.5601,
+      "step": 342
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 3.515194587938898e-07,
+      "loss": 0.5627,
+      "step": 343
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 3.3720923644582554e-07,
+      "loss": 0.4795,
+      "step": 344
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 3.23185038720218e-07,
+      "loss": 0.481,
+      "step": 345
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 3.0944782921722787e-07,
+      "loss": 0.4925,
+      "step": 346
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 2.959985518181108e-07,
+      "loss": 0.481,
+      "step": 347
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 2.8283813062035433e-07,
+      "loss": 0.462,
+      "step": 348
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 2.699674698741872e-07,
+      "loss": 0.5985,
+      "step": 349
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 2.5738745392045136e-07,
+      "loss": 0.4534,
+      "step": 350
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 2.4509894712983346e-07,
+      "loss": 0.4966,
+      "step": 351
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 2.3310279384347738e-07,
+      "loss": 0.4829,
+      "step": 352
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 2.2139981831496808e-07,
+      "loss": 0.5734,
+      "step": 353
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 2.099908246537011e-07,
+      "loss": 0.3882,
+      "step": 354
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 1.9887659676962695e-07,
+      "loss": 0.508,
+      "step": 355
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 1.8805789831939146e-07,
+      "loss": 0.3879,
+      "step": 356
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 1.7753547265386706e-07,
+      "loss": 0.4153,
+      "step": 357
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 1.6731004276707185e-07,
+      "loss": 0.5335,
+      "step": 358
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 1.5738231124649955e-07,
+      "loss": 0.5547,
+      "step": 359
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 1.4775296022483897e-07,
+      "loss": 0.4444,
+      "step": 360
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 1.3842265133310762e-07,
+      "loss": 0.536,
+      "step": 361
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 1.2939202565519243e-07,
+      "loss": 0.4738,
+      "step": 362
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 1.2066170368379763e-07,
+      "loss": 0.4359,
+      "step": 363
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 1.1223228527781281e-07,
+      "loss": 0.4132,
+      "step": 364
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 1.0410434962109633e-07,
+      "loss": 0.514,
+      "step": 365
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 9.627845518268069e-08,
+      "loss": 0.4843,
+      "step": 366
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 8.875513967839826e-08,
+      "loss": 0.4589,
+      "step": 367
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 8.15349200339363e-08,
+      "loss": 0.5233,
+      "step": 368
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 7.461829234931988e-08,
+      "loss": 0.3706,
+      "step": 369
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 6.800573186482111e-08,
+      "loss": 0.4884,
+      "step": 370
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 6.169769292831134e-08,
+      "loss": 0.4554,
+      "step": 371
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 5.569460896403755e-08,
+      "loss": 0.4819,
+      "step": 372
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 4.999689244284472e-08,
+      "loss": 0.49,
+      "step": 373
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 4.4604934853834986e-08,
+      "loss": 0.6322,
+      "step": 374
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 3.9519106677467664e-08,
+      "loss": 0.4662,
+      "step": 375
+    },
+    {
+      "epoch": 0.96,
+      "eval_accuracy": 0.835875467640784,
+      "eval_accuracy_<|content|>": 1.0,
+      "eval_accuracy_<|from|>": 0.9949431099873578,
+      "eval_accuracy_<|recipient|>": 1.0,
+      "eval_accuracy_<|stop|>": 0.9279259422525898,
+      "eval_accuracy_total_num_<|content|>": 5362,
+      "eval_accuracy_total_num_<|from|>": 791,
+      "eval_accuracy_total_num_<|recipient|>": 791,
+      "eval_accuracy_total_num_<|stop|>": 4537,
+      "eval_loss": 0.5191700458526611,
+      "eval_perplexity": 1.057695465829475,
+      "eval_runtime": 223.7003,
+      "eval_samples_per_second": 6.147,
+      "eval_steps_per_second": 0.192,
+      "step": 375
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 3.4739757360103594e-08,
+      "loss": 0.4331,
+      "step": 376
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 3.026721528999487e-08,
+      "loss": 0.4685,
+      "step": 377
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 2.6101787774722885e-08,
+      "loss": 0.479,
+      "step": 378
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 2.224376102007919e-08,
+      "loss": 0.5054,
+      "step": 379
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 1.8693400110406267e-08,
+      "loss": 0.5825,
+      "step": 380
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 1.5450948990378077e-08,
+      "loss": 0.4259,
+      "step": 381
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.2516630448241973e-08,
+      "loss": 0.4667,
+      "step": 382
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 9.890646100509937e-09,
+      "loss": 0.5553,
+      "step": 383
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 7.573176378104674e-09,
+      "loss": 0.3993,
+      "step": 384
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 5.564380513964518e-09,
+      "loss": 0.4619,
+      "step": 385
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 3.864396532100689e-09,
+      "loss": 0.4039,
+      "step": 386
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 2.4733412381143796e-09,
+      "loss": 0.4518,
+      "step": 387
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 1.3913102111691722e-09,
+      "loss": 0.4452,
+      "step": 388
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 6.183777974287929e-10,
+      "loss": 0.3681,
+      "step": 389
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 1.5459710494220015e-10,
+      "loss": 0.3691,
+      "step": 390
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 0.0,
+      "loss": 0.3283,
+      "step": 391
+    }
+  ],
+  "logging_steps": 1.0,
+  "max_steps": 391,
+  "num_train_epochs": 1,
+  "save_steps": 200.0,
+  "total_flos": 163734890741760.0,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

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

zero_to_fp32.py

@@ -0,0 +1,587 @@
+#!/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_dict = torch.load(f, map_location=device)
+        # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+        # and also handle the case where it was already removed by another helper script
+        state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+        state_dicts.append(state_dict)
+
+    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("-t",
+                        "--tag",
+                        type=str,
+                        default=None,
+                        help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+    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, tag=args.tag)