Upload folder using huggingface_hub

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/config.json

@@ -0,0 +1,30 @@
+{
+  "_name_or_path": "/mnt/data/tungtran/output_model/irish_llama2_data_v3/checkpoint-2200",
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 13824,
+  "max_position_embeddings": 4096,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 40,
+  "num_hidden_layers": 40,
+  "num_key_value_heads": 40,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.46.3",
+  "use_cache": true,
+  "vocab_size": 35483
+}

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/generation_config.json

@@ -0,0 +1,10 @@
+{
+  "bos_token_id": 1,
+  "do_sample": true,
+  "eos_token_id": 2,
+  "max_length": 4096,
+  "pad_token_id": 0,
+  "temperature": 0.6,
+  "top_p": 0.9,
+  "transformers_version": "4.46.3"
+}

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/latest

@@ -0,0 +1 @@
+global_step600

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/model-00001-of-00006.safetensors

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/model-00002-of-00006.safetensors

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/model-00003-of-00006.safetensors

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/model-00004-of-00006.safetensors

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/model-00005-of-00006.safetensors

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/model-00006-of-00006.safetensors

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/model.safetensors.index.json

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/rng_state_0.pth

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/rng_state_1.pth

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/rng_state_2.pth

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/rng_state_3.pth

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/rng_state_4.pth

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/rng_state_5.pth

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/rng_state_6.pth

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/rng_state_7.pth

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/scheduler.pt

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "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
+  }
+}

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/tokenizer.model

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/tokenizer_config.json

@@ -0,0 +1,43 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": true,
+  "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
+    }
+  },
+  "bos_token": "<s>",
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = 'You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.' %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' '  + content.strip() + ' ' + eos_token }}{% endif %}{% endfor %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "</s>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/trainer_state.json

@@ -0,0 +1,2140 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 4.0,
+  "eval_steps": 500,
+  "global_step": 600,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.006666666666666667,
+      "grad_norm": 2.1791141124998874,
+      "learning_rate": 3.3333333333333333e-06,
+      "loss": 1.6756,
+      "step": 1
+    },
+    {
+      "epoch": 0.013333333333333334,
+      "grad_norm": 2.21212674396317,
+      "learning_rate": 6.666666666666667e-06,
+      "loss": 1.7119,
+      "step": 2
+    },
+    {
+      "epoch": 0.02666666666666667,
+      "grad_norm": 1.052610355531555,
+      "learning_rate": 1.3333333333333333e-05,
+      "loss": 1.5794,
+      "step": 4
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 0.40843681896688516,
+      "learning_rate": 2e-05,
+      "loss": 1.406,
+      "step": 6
+    },
+    {
+      "epoch": 0.05333333333333334,
+      "grad_norm": 0.5064473426489812,
+      "learning_rate": 2.6666666666666667e-05,
+      "loss": 1.3868,
+      "step": 8
+    },
+    {
+      "epoch": 0.06666666666666667,
+      "grad_norm": 3.3113519775496347,
+      "learning_rate": 3.3333333333333335e-05,
+      "loss": 1.5195,
+      "step": 10
+    },
+    {
+      "epoch": 0.08,
+      "grad_norm": 0.3604597146473128,
+      "learning_rate": 4e-05,
+      "loss": 1.3777,
+      "step": 12
+    },
+    {
+      "epoch": 0.09333333333333334,
+      "grad_norm": 0.23235777614167963,
+      "learning_rate": 4.666666666666667e-05,
+      "loss": 1.2747,
+      "step": 14
+    },
+    {
+      "epoch": 0.10666666666666667,
+      "grad_norm": 0.24107331430994436,
+      "learning_rate": 5.333333333333333e-05,
+      "loss": 1.2307,
+      "step": 16
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 0.19828951527897512,
+      "learning_rate": 6e-05,
+      "loss": 1.2127,
+      "step": 18
+    },
+    {
+      "epoch": 0.13333333333333333,
+      "grad_norm": 0.12945101422872216,
+      "learning_rate": 6.666666666666667e-05,
+      "loss": 1.2027,
+      "step": 20
+    },
+    {
+      "epoch": 0.14666666666666667,
+      "grad_norm": 0.12079427155353203,
+      "learning_rate": 7.333333333333333e-05,
+      "loss": 1.165,
+      "step": 22
+    },
+    {
+      "epoch": 0.16,
+      "grad_norm": 0.11223243327491182,
+      "learning_rate": 8e-05,
+      "loss": 1.1332,
+      "step": 24
+    },
+    {
+      "epoch": 0.17333333333333334,
+      "grad_norm": 0.11700825567943814,
+      "learning_rate": 8.666666666666667e-05,
+      "loss": 1.1353,
+      "step": 26
+    },
+    {
+      "epoch": 0.18666666666666668,
+      "grad_norm": 0.10577860008538974,
+      "learning_rate": 9.333333333333334e-05,
+      "loss": 1.1126,
+      "step": 28
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 0.08770959641951533,
+      "learning_rate": 0.0001,
+      "loss": 1.116,
+      "step": 30
+    },
+    {
+      "epoch": 0.21333333333333335,
+      "grad_norm": 0.08173970603071407,
+      "learning_rate": 9.999696229471716e-05,
+      "loss": 1.09,
+      "step": 32
+    },
+    {
+      "epoch": 0.22666666666666666,
+      "grad_norm": 0.07804539860851849,
+      "learning_rate": 9.998784954797474e-05,
+      "loss": 1.0622,
+      "step": 34
+    },
+    {
+      "epoch": 0.24,
+      "grad_norm": 0.07305109208219707,
+      "learning_rate": 9.997266286704631e-05,
+      "loss": 1.0522,
+      "step": 36
+    },
+    {
+      "epoch": 0.25333333333333335,
+      "grad_norm": 0.08178925315694836,
+      "learning_rate": 9.99514040972383e-05,
+      "loss": 1.0539,
+      "step": 38
+    },
+    {
+      "epoch": 0.26666666666666666,
+      "grad_norm": 0.07827810776223704,
+      "learning_rate": 9.992407582166581e-05,
+      "loss": 1.0715,
+      "step": 40
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 0.0844202707800422,
+      "learning_rate": 9.989068136093873e-05,
+      "loss": 1.0274,
+      "step": 42
+    },
+    {
+      "epoch": 0.29333333333333333,
+      "grad_norm": 0.06818226761844356,
+      "learning_rate": 9.985122477275824e-05,
+      "loss": 1.0338,
+      "step": 44
+    },
+    {
+      "epoch": 0.30666666666666664,
+      "grad_norm": 0.07583484586118773,
+      "learning_rate": 9.980571085142381e-05,
+      "loss": 1.0323,
+      "step": 46
+    },
+    {
+      "epoch": 0.32,
+      "grad_norm": 0.0755020533085722,
+      "learning_rate": 9.975414512725057e-05,
+      "loss": 1.015,
+      "step": 48
+    },
+    {
+      "epoch": 0.3333333333333333,
+      "grad_norm": 0.08406998549299528,
+      "learning_rate": 9.969653386589748e-05,
+      "loss": 0.9656,
+      "step": 50
+    },
+    {
+      "epoch": 0.3466666666666667,
+      "grad_norm": 0.07101342121045431,
+      "learning_rate": 9.963288406760582e-05,
+      "loss": 0.9657,
+      "step": 52
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 0.07417386295183337,
+      "learning_rate": 9.956320346634876e-05,
+      "loss": 0.9937,
+      "step": 54
+    },
+    {
+      "epoch": 0.37333333333333335,
+      "grad_norm": 0.07154338853868397,
+      "learning_rate": 9.94875005288915e-05,
+      "loss": 0.9589,
+      "step": 56
+    },
+    {
+      "epoch": 0.38666666666666666,
+      "grad_norm": 0.07218740352233513,
+      "learning_rate": 9.940578445376258e-05,
+      "loss": 0.9685,
+      "step": 58
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 0.07100653973675006,
+      "learning_rate": 9.931806517013612e-05,
+      "loss": 0.9573,
+      "step": 60
+    },
+    {
+      "epoch": 0.41333333333333333,
+      "grad_norm": 0.06711082249359601,
+      "learning_rate": 9.922435333662536e-05,
+      "loss": 0.9364,
+      "step": 62
+    },
+    {
+      "epoch": 0.4266666666666667,
+      "grad_norm": 0.07153154856840055,
+      "learning_rate": 9.912466033998757e-05,
+      "loss": 0.9374,
+      "step": 64
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 0.06963514347132123,
+      "learning_rate": 9.901899829374047e-05,
+      "loss": 0.9663,
+      "step": 66
+    },
+    {
+      "epoch": 0.4533333333333333,
+      "grad_norm": 0.0659376914599576,
+      "learning_rate": 9.890738003669029e-05,
+      "loss": 0.9108,
+      "step": 68
+    },
+    {
+      "epoch": 0.4666666666666667,
+      "grad_norm": 0.06958378585339599,
+      "learning_rate": 9.878981913137179e-05,
+      "loss": 0.903,
+      "step": 70
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 0.0764469187644936,
+      "learning_rate": 9.86663298624003e-05,
+      "loss": 0.9008,
+      "step": 72
+    },
+    {
+      "epoch": 0.49333333333333335,
+      "grad_norm": 0.07709623493461817,
+      "learning_rate": 9.8536927234736e-05,
+      "loss": 0.9305,
+      "step": 74
+    },
+    {
+      "epoch": 0.5066666666666667,
+      "grad_norm": 0.06882033015566792,
+      "learning_rate": 9.840162697186075e-05,
+      "loss": 0.9153,
+      "step": 76
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 0.06984284125927373,
+      "learning_rate": 9.826044551386744e-05,
+      "loss": 0.8861,
+      "step": 78
+    },
+    {
+      "epoch": 0.5333333333333333,
+      "grad_norm": 11.87344671430505,
+      "learning_rate": 9.811340001546251e-05,
+      "loss": 0.8805,
+      "step": 80
+    },
+    {
+      "epoch": 0.5466666666666666,
+      "grad_norm": 0.07471362876432724,
+      "learning_rate": 9.796050834388149e-05,
+      "loss": 0.8571,
+      "step": 82
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 0.20377795789303127,
+      "learning_rate": 9.780178907671789e-05,
+      "loss": 0.8725,
+      "step": 84
+    },
+    {
+      "epoch": 0.5733333333333334,
+      "grad_norm": 0.09817542250054903,
+      "learning_rate": 9.763726149966596e-05,
+      "loss": 0.8883,
+      "step": 86
+    },
+    {
+      "epoch": 0.5866666666666667,
+      "grad_norm": 0.10979752851172697,
+      "learning_rate": 9.746694560417731e-05,
+      "loss": 0.8708,
+      "step": 88
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 0.09854482601183034,
+      "learning_rate": 9.729086208503174e-05,
+      "loss": 0.8883,
+      "step": 90
+    },
+    {
+      "epoch": 0.6133333333333333,
+      "grad_norm": 0.08153129504069773,
+      "learning_rate": 9.710903233782272e-05,
+      "loss": 0.8385,
+      "step": 92
+    },
+    {
+      "epoch": 0.6266666666666667,
+      "grad_norm": 0.07933734102384941,
+      "learning_rate": 9.692147845635761e-05,
+      "loss": 0.8658,
+      "step": 94
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 0.08703768235227666,
+      "learning_rate": 9.672822322997305e-05,
+      "loss": 0.8314,
+      "step": 96
+    },
+    {
+      "epoch": 0.6533333333333333,
+      "grad_norm": 0.07185119301522967,
+      "learning_rate": 9.652929014076593e-05,
+      "loss": 0.7829,
+      "step": 98
+    },
+    {
+      "epoch": 0.6666666666666666,
+      "grad_norm": 0.08362387977855704,
+      "learning_rate": 9.632470336074009e-05,
+      "loss": 0.7932,
+      "step": 100
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 0.07857330769987199,
+      "learning_rate": 9.611448774886924e-05,
+      "loss": 0.8094,
+      "step": 102
+    },
+    {
+      "epoch": 0.6933333333333334,
+      "grad_norm": 0.07470080527568211,
+      "learning_rate": 9.589866884807635e-05,
+      "loss": 0.8153,
+      "step": 104
+    },
+    {
+      "epoch": 0.7066666666666667,
+      "grad_norm": 0.07991307787992707,
+      "learning_rate": 9.567727288213005e-05,
+      "loss": 0.7736,
+      "step": 106
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.0733968002848949,
+      "learning_rate": 9.545032675245813e-05,
+      "loss": 0.78,
+      "step": 108
+    },
+    {
+      "epoch": 0.7333333333333333,
+      "grad_norm": 0.07751184405863311,
+      "learning_rate": 9.521785803487889e-05,
+      "loss": 0.7959,
+      "step": 110
+    },
+    {
+      "epoch": 0.7466666666666667,
+      "grad_norm": 0.08355612624474511,
+      "learning_rate": 9.497989497625035e-05,
+      "loss": 0.7751,
+      "step": 112
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 0.0929055766139149,
+      "learning_rate": 9.473646649103818e-05,
+      "loss": 0.7848,
+      "step": 114
+    },
+    {
+      "epoch": 0.7733333333333333,
+      "grad_norm": 0.09585556290161484,
+      "learning_rate": 9.448760215780217e-05,
+      "loss": 0.7786,
+      "step": 116
+    },
+    {
+      "epoch": 0.7866666666666666,
+      "grad_norm": 0.0836807061462442,
+      "learning_rate": 9.42333322156023e-05,
+      "loss": 0.7334,
+      "step": 118
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 0.07708063741525999,
+      "learning_rate": 9.397368756032445e-05,
+      "loss": 0.7602,
+      "step": 120
+    },
+    {
+      "epoch": 0.8133333333333334,
+      "grad_norm": 1.0840349168421557,
+      "learning_rate": 9.370869974092629e-05,
+      "loss": 0.7948,
+      "step": 122
+    },
+    {
+      "epoch": 0.8266666666666667,
+      "grad_norm": 0.6248060432126027,
+      "learning_rate": 9.343840095560372e-05,
+      "loss": 0.7419,
+      "step": 124
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 0.14470369129957764,
+      "learning_rate": 9.316282404787871e-05,
+      "loss": 0.7623,
+      "step": 126
+    },
+    {
+      "epoch": 0.8533333333333334,
+      "grad_norm": 0.131124623959581,
+      "learning_rate": 9.288200250260836e-05,
+      "loss": 0.7591,
+      "step": 128
+    },
+    {
+      "epoch": 0.8666666666666667,
+      "grad_norm": 0.11725616877030448,
+      "learning_rate": 9.259597044191636e-05,
+      "loss": 0.7409,
+      "step": 130
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 0.12199467676115175,
+      "learning_rate": 9.230476262104677e-05,
+      "loss": 0.6848,
+      "step": 132
+    },
+    {
+      "epoch": 0.8933333333333333,
+      "grad_norm": 0.08692329616701802,
+      "learning_rate": 9.200841442414106e-05,
+      "loss": 0.7042,
+      "step": 134
+    },
+    {
+      "epoch": 0.9066666666666666,
+      "grad_norm": 0.07518959531670097,
+      "learning_rate": 9.17069618599385e-05,
+      "loss": 0.692,
+      "step": 136
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 0.09269304556873778,
+      "learning_rate": 9.140044155740101e-05,
+      "loss": 0.719,
+      "step": 138
+    },
+    {
+      "epoch": 0.9333333333333333,
+      "grad_norm": 0.0794274462938484,
+      "learning_rate": 9.108889076126226e-05,
+      "loss": 0.7442,
+      "step": 140
+    },
+    {
+      "epoch": 0.9466666666666667,
+      "grad_norm": 0.07208605668488438,
+      "learning_rate": 9.077234732750224e-05,
+      "loss": 0.6918,
+      "step": 142
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 0.08494128110838685,
+      "learning_rate": 9.045084971874738e-05,
+      "loss": 0.6942,
+      "step": 144
+    },
+    {
+      "epoch": 0.9733333333333334,
+      "grad_norm": 0.07966340996836804,
+      "learning_rate": 9.012443699959705e-05,
+      "loss": 0.6471,
+      "step": 146
+    },
+    {
+      "epoch": 0.9866666666666667,
+      "grad_norm": 0.07509382026953668,
+      "learning_rate": 8.979314883187693e-05,
+      "loss": 0.6641,
+      "step": 148
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 0.09671488019621963,
+      "learning_rate": 8.945702546981969e-05,
+      "loss": 0.6563,
+      "step": 150
+    },
+    {
+      "epoch": 1.0133333333333334,
+      "grad_norm": 0.07306699627281733,
+      "learning_rate": 8.911610775517382e-05,
+      "loss": 0.4728,
+      "step": 152
+    },
+    {
+      "epoch": 1.0266666666666666,
+      "grad_norm": 0.07113580005854522,
+      "learning_rate": 8.877043711224108e-05,
+      "loss": 0.4537,
+      "step": 154
+    },
+    {
+      "epoch": 1.04,
+      "grad_norm": 0.08242409293384922,
+      "learning_rate": 8.842005554284296e-05,
+      "loss": 0.4754,
+      "step": 156
+    },
+    {
+      "epoch": 1.0533333333333332,
+      "grad_norm": 0.08558759822156714,
+      "learning_rate": 8.806500562121723e-05,
+      "loss": 0.4409,
+      "step": 158
+    },
+    {
+      "epoch": 1.0666666666666667,
+      "grad_norm": 0.08093797796678245,
+      "learning_rate": 8.770533048884482e-05,
+      "loss": 0.4629,
+      "step": 160
+    },
+    {
+      "epoch": 1.08,
+      "grad_norm": 0.07276419262446279,
+      "learning_rate": 8.73410738492077e-05,
+      "loss": 0.4509,
+      "step": 162
+    },
+    {
+      "epoch": 1.0933333333333333,
+      "grad_norm": 0.13007281279313135,
+      "learning_rate": 8.697227996247861e-05,
+      "loss": 0.4929,
+      "step": 164
+    },
+    {
+      "epoch": 1.1066666666666667,
+      "grad_norm": 0.08797907567452011,
+      "learning_rate": 8.659899364014309e-05,
+      "loss": 0.4567,
+      "step": 166
+    },
+    {
+      "epoch": 1.12,
+      "grad_norm": 0.08399912199961071,
+      "learning_rate": 8.622126023955446e-05,
+      "loss": 0.4561,
+      "step": 168
+    },
+    {
+      "epoch": 1.1333333333333333,
+      "grad_norm": 0.14999451040489625,
+      "learning_rate": 8.583912565842257e-05,
+      "loss": 0.461,
+      "step": 170
+    },
+    {
+      "epoch": 1.1466666666666667,
+      "grad_norm": 0.07239352221098809,
+      "learning_rate": 8.545263632923687e-05,
+      "loss": 0.4386,
+      "step": 172
+    },
+    {
+      "epoch": 1.16,
+      "grad_norm": 0.07274293681830099,
+      "learning_rate": 8.506183921362443e-05,
+      "loss": 0.425,
+      "step": 174
+    },
+    {
+      "epoch": 1.1733333333333333,
+      "grad_norm": 0.06567961209669387,
+      "learning_rate": 8.466678179664379e-05,
+      "loss": 0.4316,
+      "step": 176
+    },
+    {
+      "epoch": 1.1866666666666668,
+      "grad_norm": 0.06800506030151886,
+      "learning_rate": 8.4267512081015e-05,
+      "loss": 0.4262,
+      "step": 178
+    },
+    {
+      "epoch": 1.2,
+      "grad_norm": 0.06310710452459224,
+      "learning_rate": 8.386407858128706e-05,
+      "loss": 0.4393,
+      "step": 180
+    },
+    {
+      "epoch": 1.2133333333333334,
+      "grad_norm": 0.07077155581372561,
+      "learning_rate": 8.345653031794292e-05,
+      "loss": 0.4394,
+      "step": 182
+    },
+    {
+      "epoch": 1.2266666666666666,
+      "grad_norm": 0.06505271915956287,
+      "learning_rate": 8.304491681144306e-05,
+      "loss": 0.4486,
+      "step": 184
+    },
+    {
+      "epoch": 1.24,
+      "grad_norm": 0.059685754017888774,
+      "learning_rate": 8.262928807620843e-05,
+      "loss": 0.4449,
+      "step": 186
+    },
+    {
+      "epoch": 1.2533333333333334,
+      "grad_norm": 0.061596408555461855,
+      "learning_rate": 8.220969461454322e-05,
+      "loss": 0.4241,
+      "step": 188
+    },
+    {
+      "epoch": 1.2666666666666666,
+      "grad_norm": 0.06365147365999292,
+      "learning_rate": 8.178618741049842e-05,
+      "loss": 0.4516,
+      "step": 190
+    },
+    {
+      "epoch": 1.28,
+      "grad_norm": 0.06440909119431051,
+      "learning_rate": 8.135881792367686e-05,
+      "loss": 0.4047,
+      "step": 192
+    },
+    {
+      "epoch": 1.2933333333333334,
+      "grad_norm": 0.06734344900628648,
+      "learning_rate": 8.092763808298048e-05,
+      "loss": 0.4169,
+      "step": 194
+    },
+    {
+      "epoch": 1.3066666666666666,
+      "grad_norm": 0.05832916804134293,
+      "learning_rate": 8.049270028030046e-05,
+      "loss": 0.4313,
+      "step": 196
+    },
+    {
+      "epoch": 1.32,
+      "grad_norm": 0.060432835217476055,
+      "learning_rate": 8.005405736415126e-05,
+      "loss": 0.4181,
+      "step": 198
+    },
+    {
+      "epoch": 1.3333333333333333,
+      "grad_norm": 0.3189524387852971,
+      "learning_rate": 7.961176263324901e-05,
+      "loss": 0.4671,
+      "step": 200
+    },
+    {
+      "epoch": 1.3466666666666667,
+      "grad_norm": 0.0592761178725548,
+      "learning_rate": 7.916586983003533e-05,
+      "loss": 0.4399,
+      "step": 202
+    },
+    {
+      "epoch": 1.3599999999999999,
+      "grad_norm": 0.0709909293526123,
+      "learning_rate": 7.871643313414718e-05,
+      "loss": 0.4297,
+      "step": 204
+    },
+    {
+      "epoch": 1.3733333333333333,
+      "grad_norm": 0.06198999527057115,
+      "learning_rate": 7.82635071558336e-05,
+      "loss": 0.4251,
+      "step": 206
+    },
+    {
+      "epoch": 1.3866666666666667,
+      "grad_norm": 0.058077551229268935,
+      "learning_rate": 7.780714692932002e-05,
+      "loss": 0.4105,
+      "step": 208
+    },
+    {
+      "epoch": 1.4,
+      "grad_norm": 0.06776940923998968,
+      "learning_rate": 7.734740790612136e-05,
+      "loss": 0.4445,
+      "step": 210
+    },
+    {
+      "epoch": 1.4133333333333333,
+      "grad_norm": 0.05980241043639003,
+      "learning_rate": 7.688434594830392e-05,
+      "loss": 0.4049,
+      "step": 212
+    },
+    {
+      "epoch": 1.4266666666666667,
+      "grad_norm": 0.06075570212356431,
+      "learning_rate": 7.641801732169795e-05,
+      "loss": 0.4049,
+      "step": 214
+    },
+    {
+      "epoch": 1.44,
+      "grad_norm": 0.0585436609373402,
+      "learning_rate": 7.594847868906076e-05,
+      "loss": 0.4265,
+      "step": 216
+    },
+    {
+      "epoch": 1.4533333333333334,
+      "grad_norm": 0.05886231084821202,
+      "learning_rate": 7.547578710319174e-05,
+      "loss": 0.4361,
+      "step": 218
+    },
+    {
+      "epoch": 1.4666666666666668,
+      "grad_norm": 0.05758877919907456,
+      "learning_rate": 7.500000000000001e-05,
+      "loss": 0.4273,
+      "step": 220
+    },
+    {
+      "epoch": 1.48,
+      "grad_norm": 0.057186857880065896,
+      "learning_rate": 7.452117519152542e-05,
+      "loss": 0.3918,
+      "step": 222
+    },
+    {
+      "epoch": 1.4933333333333334,
+      "grad_norm": 0.060986952045000244,
+      "learning_rate": 7.403937085891397e-05,
+      "loss": 0.4116,
+      "step": 224
+    },
+    {
+      "epoch": 1.5066666666666668,
+      "grad_norm": 0.06094291896870185,
+      "learning_rate": 7.355464554534837e-05,
+      "loss": 0.3968,
+      "step": 226
+    },
+    {
+      "epoch": 1.52,
+      "grad_norm": 0.058322590431556104,
+      "learning_rate": 7.30670581489344e-05,
+      "loss": 0.3876,
+      "step": 228
+    },
+    {
+      "epoch": 1.5333333333333332,
+      "grad_norm": 0.05822076237291155,
+      "learning_rate": 7.257666791554448e-05,
+      "loss": 0.4272,
+      "step": 230
+    },
+    {
+      "epoch": 1.5466666666666666,
+      "grad_norm": 0.055094265304100165,
+      "learning_rate": 7.20835344316187e-05,
+      "loss": 0.4133,
+      "step": 232
+    },
+    {
+      "epoch": 1.56,
+      "grad_norm": 0.06044900185707112,
+      "learning_rate": 7.158771761692464e-05,
+      "loss": 0.4061,
+      "step": 234
+    },
+    {
+      "epoch": 1.5733333333333333,
+      "grad_norm": 0.05589749063648132,
+      "learning_rate": 7.108927771727661e-05,
+      "loss": 0.407,
+      "step": 236
+    },
+    {
+      "epoch": 1.5866666666666667,
+      "grad_norm": 0.05920526918754276,
+      "learning_rate": 7.058827529721525e-05,
+      "loss": 0.3993,
+      "step": 238
+    },
+    {
+      "epoch": 1.6,
+      "grad_norm": 0.08772211610388515,
+      "learning_rate": 7.008477123264848e-05,
+      "loss": 0.4208,
+      "step": 240
+    },
+    {
+      "epoch": 1.6133333333333333,
+      "grad_norm": 0.05570158803300833,
+      "learning_rate": 6.957882670345458e-05,
+      "loss": 0.405,
+      "step": 242
+    },
+    {
+      "epoch": 1.6266666666666667,
+      "grad_norm": 0.06563516734346801,
+      "learning_rate": 6.90705031860483e-05,
+      "loss": 0.4098,
+      "step": 244
+    },
+    {
+      "epoch": 1.6400000000000001,
+      "grad_norm": 0.05631976426536162,
+      "learning_rate": 6.855986244591104e-05,
+      "loss": 0.3877,
+      "step": 246
+    },
+    {
+      "epoch": 1.6533333333333333,
+      "grad_norm": 0.05644442739953386,
+      "learning_rate": 6.804696653008575e-05,
+      "loss": 0.4134,
+      "step": 248
+    },
+    {
+      "epoch": 1.6666666666666665,
+      "grad_norm": 0.05275767947492662,
+      "learning_rate": 6.753187775963773e-05,
+      "loss": 0.4075,
+      "step": 250
+    },
+    {
+      "epoch": 1.6800000000000002,
+      "grad_norm": 0.05873640121915926,
+      "learning_rate": 6.701465872208216e-05,
+      "loss": 0.4103,
+      "step": 252
+    },
+    {
+      "epoch": 1.6933333333333334,
+      "grad_norm": 0.0596209632813038,
+      "learning_rate": 6.649537226377915e-05,
+      "loss": 0.4368,
+      "step": 254
+    },
+    {
+      "epoch": 1.7066666666666666,
+      "grad_norm": 0.0550209071778787,
+      "learning_rate": 6.59740814822974e-05,
+      "loss": 0.4079,
+      "step": 256
+    },
+    {
+      "epoch": 1.72,
+      "grad_norm": 0.056297483765733,
+      "learning_rate": 6.545084971874738e-05,
+      "loss": 0.441,
+      "step": 258
+    },
+    {
+      "epoch": 1.7333333333333334,
+      "grad_norm": 0.15265154617055526,
+      "learning_rate": 6.492574055008473e-05,
+      "loss": 0.4041,
+      "step": 260
+    },
+    {
+      "epoch": 1.7466666666666666,
+      "grad_norm": 0.07752909356786238,
+      "learning_rate": 6.439881778138531e-05,
+      "loss": 0.4041,
+      "step": 262
+    },
+    {
+      "epoch": 1.76,
+      "grad_norm": 0.08024865113518131,
+      "learning_rate": 6.387014543809223e-05,
+      "loss": 0.3986,
+      "step": 264
+    },
+    {
+      "epoch": 1.7733333333333334,
+      "grad_norm": 0.10615717597280923,
+      "learning_rate": 6.333978775823631e-05,
+      "loss": 0.4134,
+      "step": 266
+    },
+    {
+      "epoch": 1.7866666666666666,
+      "grad_norm": 0.06710384327076761,
+      "learning_rate": 6.280780918463057e-05,
+      "loss": 0.3705,
+      "step": 268
+    },
+    {
+      "epoch": 1.8,
+      "grad_norm": 0.0633674315054036,
+      "learning_rate": 6.227427435703997e-05,
+      "loss": 0.4157,
+      "step": 270
+    },
+    {
+      "epoch": 1.8133333333333335,
+      "grad_norm": 0.060851228058216876,
+      "learning_rate": 6.173924810432705e-05,
+      "loss": 0.4065,
+      "step": 272
+    },
+    {
+      "epoch": 1.8266666666666667,
+      "grad_norm": 0.058897518477155845,
+      "learning_rate": 6.12027954365748e-05,
+      "loss": 0.4019,
+      "step": 274
+    },
+    {
+      "epoch": 1.8399999999999999,
+      "grad_norm": 0.05632347613072352,
+      "learning_rate": 6.066498153718735e-05,
+      "loss": 0.4142,
+      "step": 276
+    },
+    {
+      "epoch": 1.8533333333333335,
+      "grad_norm": 0.05684584086680513,
+      "learning_rate": 6.012587175496961e-05,
+      "loss": 0.4019,
+      "step": 278
+    },
+    {
+      "epoch": 1.8666666666666667,
+      "grad_norm": 0.054019444337334764,
+      "learning_rate": 5.958553159618693e-05,
+      "loss": 0.4006,
+      "step": 280
+    },
+    {
+      "epoch": 1.88,
+      "grad_norm": 0.05665016004137828,
+      "learning_rate": 5.90440267166055e-05,
+      "loss": 0.411,
+      "step": 282
+    },
+    {
+      "epoch": 1.8933333333333333,
+      "grad_norm": 0.057316286082019274,
+      "learning_rate": 5.850142291351466e-05,
+      "loss": 0.4085,
+      "step": 284
+    },
+    {
+      "epoch": 1.9066666666666667,
+      "grad_norm": 0.08470811257422295,
+      "learning_rate": 5.795778611773197e-05,
+      "loss": 0.3819,
+      "step": 286
+    },
+    {
+      "epoch": 1.92,
+      "grad_norm": 0.05784580025752463,
+      "learning_rate": 5.74131823855921e-05,
+      "loss": 0.3821,
+      "step": 288
+    },
+    {
+      "epoch": 1.9333333333333333,
+      "grad_norm": 0.05663311823528816,
+      "learning_rate": 5.686767789092041e-05,
+      "loss": 0.4154,
+      "step": 290
+    },
+    {
+      "epoch": 1.9466666666666668,
+      "grad_norm": 0.057073777066581295,
+      "learning_rate": 5.6321338916992315e-05,
+      "loss": 0.4135,
+      "step": 292
+    },
+    {
+      "epoch": 1.96,
+      "grad_norm": 0.0574979098693646,
+      "learning_rate": 5.577423184847932e-05,
+      "loss": 0.408,
+      "step": 294
+    },
+    {
+      "epoch": 1.9733333333333334,
+      "grad_norm": 0.05504450544637458,
+      "learning_rate": 5.522642316338268e-05,
+      "loss": 0.4034,
+      "step": 296
+    },
+    {
+      "epoch": 1.9866666666666668,
+      "grad_norm": 0.057266688427656184,
+      "learning_rate": 5.467797942495589e-05,
+      "loss": 0.3741,
+      "step": 298
+    },
+    {
+      "epoch": 2.0,
+      "grad_norm": 0.05613792113707641,
+      "learning_rate": 5.4128967273616625e-05,
+      "loss": 0.3944,
+      "step": 300
+    },
+    {
+      "epoch": 2.013333333333333,
+      "grad_norm": 0.07519534578570336,
+      "learning_rate": 5.357945341884936e-05,
+      "loss": 0.2236,
+      "step": 302
+    },
+    {
+      "epoch": 2.026666666666667,
+      "grad_norm": 0.18176383216046363,
+      "learning_rate": 5.3029504631099694e-05,
+      "loss": 0.2312,
+      "step": 304
+    },
+    {
+      "epoch": 2.04,
+      "grad_norm": 0.0608749623913546,
+      "learning_rate": 5.247918773366112e-05,
+      "loss": 0.2194,
+      "step": 306
+    },
+    {
+      "epoch": 2.0533333333333332,
+      "grad_norm": 0.06549466194373386,
+      "learning_rate": 5.1928569594555524e-05,
+      "loss": 0.2038,
+      "step": 308
+    },
+    {
+      "epoch": 2.066666666666667,
+      "grad_norm": 0.06841206457163714,
+      "learning_rate": 5.1377717118408105e-05,
+      "loss": 0.2103,
+      "step": 310
+    },
+    {
+      "epoch": 2.08,
+      "grad_norm": 0.06431960601426534,
+      "learning_rate": 5.0826697238317935e-05,
+      "loss": 0.1899,
+      "step": 312
+    },
+    {
+      "epoch": 2.0933333333333333,
+      "grad_norm": 0.061164578906699046,
+      "learning_rate": 5.027557690772503e-05,
+      "loss": 0.1955,
+      "step": 314
+    },
+    {
+      "epoch": 2.1066666666666665,
+      "grad_norm": 0.05806737719217297,
+      "learning_rate": 4.972442309227498e-05,
+      "loss": 0.1871,
+      "step": 316
+    },
+    {
+      "epoch": 2.12,
+      "grad_norm": 0.10638328645893626,
+      "learning_rate": 4.917330276168208e-05,
+      "loss": 0.1926,
+      "step": 318
+    },
+    {
+      "epoch": 2.1333333333333333,
+      "grad_norm": 0.055513863609778974,
+      "learning_rate": 4.8622282881591906e-05,
+      "loss": 0.2008,
+      "step": 320
+    },
+    {
+      "epoch": 2.1466666666666665,
+      "grad_norm": 0.05781185580562024,
+      "learning_rate": 4.8071430405444474e-05,
+      "loss": 0.2001,
+      "step": 322
+    },
+    {
+      "epoch": 2.16,
+      "grad_norm": 0.058407688794793834,
+      "learning_rate": 4.7520812266338885e-05,
+      "loss": 0.2091,
+      "step": 324
+    },
+    {
+      "epoch": 2.1733333333333333,
+      "grad_norm": 0.05455993945014677,
+      "learning_rate": 4.697049536890033e-05,
+      "loss": 0.1943,
+      "step": 326
+    },
+    {
+      "epoch": 2.1866666666666665,
+      "grad_norm": 0.10418377423186777,
+      "learning_rate": 4.642054658115067e-05,
+      "loss": 0.2101,
+      "step": 328
+    },
+    {
+      "epoch": 2.2,
+      "grad_norm": 0.058762085061184646,
+      "learning_rate": 4.5871032726383386e-05,
+      "loss": 0.199,
+      "step": 330
+    },
+    {
+      "epoch": 2.2133333333333334,
+      "grad_norm": 0.05961716033762566,
+      "learning_rate": 4.5322020575044114e-05,
+      "loss": 0.1945,
+      "step": 332
+    },
+    {
+      "epoch": 2.2266666666666666,
+      "grad_norm": 0.06064491694409506,
+      "learning_rate": 4.477357683661734e-05,
+      "loss": 0.2023,
+      "step": 334
+    },
+    {
+      "epoch": 2.24,
+      "grad_norm": 0.12224153955185721,
+      "learning_rate": 4.4225768151520694e-05,
+      "loss": 0.2372,
+      "step": 336
+    },
+    {
+      "epoch": 2.2533333333333334,
+      "grad_norm": 0.06273415054348162,
+      "learning_rate": 4.367866108300769e-05,
+      "loss": 0.202,
+      "step": 338
+    },
+    {
+      "epoch": 2.2666666666666666,
+      "grad_norm": 0.05439313011616356,
+      "learning_rate": 4.3132322109079596e-05,
+      "loss": 0.1838,
+      "step": 340
+    },
+    {
+      "epoch": 2.2800000000000002,
+      "grad_norm": 0.05408796882713708,
+      "learning_rate": 4.2586817614407895e-05,
+      "loss": 0.2041,
+      "step": 342
+    },
+    {
+      "epoch": 2.2933333333333334,
+      "grad_norm": 0.14961779038669223,
+      "learning_rate": 4.2042213882268025e-05,
+      "loss": 0.217,
+      "step": 344
+    },
+    {
+      "epoch": 2.3066666666666666,
+      "grad_norm": 0.06357828772862752,
+      "learning_rate": 4.149857708648536e-05,
+      "loss": 0.2003,
+      "step": 346
+    },
+    {
+      "epoch": 2.32,
+      "grad_norm": 0.054464894143134636,
+      "learning_rate": 4.095597328339452e-05,
+      "loss": 0.1909,
+      "step": 348
+    },
+    {
+      "epoch": 2.3333333333333335,
+      "grad_norm": 0.066782998686637,
+      "learning_rate": 4.0414468403813095e-05,
+      "loss": 0.205,
+      "step": 350
+    },
+    {
+      "epoch": 2.3466666666666667,
+      "grad_norm": 0.10946146387161931,
+      "learning_rate": 3.9874128245030404e-05,
+      "loss": 0.204,
+      "step": 352
+    },
+    {
+      "epoch": 2.36,
+      "grad_norm": 0.09072556796395603,
+      "learning_rate": 3.933501846281267e-05,
+      "loss": 0.2021,
+      "step": 354
+    },
+    {
+      "epoch": 2.3733333333333335,
+      "grad_norm": 0.08641355840276257,
+      "learning_rate": 3.879720456342521e-05,
+      "loss": 0.2204,
+      "step": 356
+    },
+    {
+      "epoch": 2.3866666666666667,
+      "grad_norm": 0.0733857764459229,
+      "learning_rate": 3.826075189567296e-05,
+      "loss": 0.2012,
+      "step": 358
+    },
+    {
+      "epoch": 2.4,
+      "grad_norm": 0.0626333235413589,
+      "learning_rate": 3.772572564296005e-05,
+      "loss": 0.2053,
+      "step": 360
+    },
+    {
+      "epoch": 2.413333333333333,
+      "grad_norm": 0.06594691744680978,
+      "learning_rate": 3.719219081536942e-05,
+      "loss": 0.1971,
+      "step": 362
+    },
+    {
+      "epoch": 2.4266666666666667,
+      "grad_norm": 0.06571408902511346,
+      "learning_rate": 3.666021224176369e-05,
+      "loss": 0.1977,
+      "step": 364
+    },
+    {
+      "epoch": 2.44,
+      "grad_norm": 0.0890355148563959,
+      "learning_rate": 3.612985456190778e-05,
+      "loss": 0.1974,
+      "step": 366
+    },
+    {
+      "epoch": 2.453333333333333,
+      "grad_norm": 0.05984188182013961,
+      "learning_rate": 3.56011822186147e-05,
+      "loss": 0.2045,
+      "step": 368
+    },
+    {
+      "epoch": 2.466666666666667,
+      "grad_norm": 0.05840811324090727,
+      "learning_rate": 3.5074259449915284e-05,
+      "loss": 0.1971,
+      "step": 370
+    },
+    {
+      "epoch": 2.48,
+      "grad_norm": 0.05496682361577455,
+      "learning_rate": 3.4549150281252636e-05,
+      "loss": 0.2088,
+      "step": 372
+    },
+    {
+      "epoch": 2.493333333333333,
+      "grad_norm": 0.056504022577883756,
+      "learning_rate": 3.40259185177026e-05,
+      "loss": 0.2017,
+      "step": 374
+    },
+    {
+      "epoch": 2.506666666666667,
+      "grad_norm": 0.06329898582468307,
+      "learning_rate": 3.350462773622086e-05,
+      "loss": 0.2064,
+      "step": 376
+    },
+    {
+      "epoch": 2.52,
+      "grad_norm": 0.05916890402008683,
+      "learning_rate": 3.298534127791785e-05,
+      "loss": 0.211,
+      "step": 378
+    },
+    {
+      "epoch": 2.533333333333333,
+      "grad_norm": 0.05979626765953065,
+      "learning_rate": 3.2468122240362284e-05,
+      "loss": 0.2044,
+      "step": 380
+    },
+    {
+      "epoch": 2.546666666666667,
+      "grad_norm": 0.06045023498620747,
+      "learning_rate": 3.1953033469914276e-05,
+      "loss": 0.2074,
+      "step": 382
+    },
+    {
+      "epoch": 2.56,
+      "grad_norm": 0.06281675933045895,
+      "learning_rate": 3.144013755408895e-05,
+      "loss": 0.2146,
+      "step": 384
+    },
+    {
+      "epoch": 2.5733333333333333,
+      "grad_norm": 0.07056631468249065,
+      "learning_rate": 3.0929496813951694e-05,
+      "loss": 0.1853,
+      "step": 386
+    },
+    {
+      "epoch": 2.586666666666667,
+      "grad_norm": 0.06942161032122679,
+      "learning_rate": 3.042117329654544e-05,
+      "loss": 0.1916,
+      "step": 388
+    },
+    {
+      "epoch": 2.6,
+      "grad_norm": 0.06027569169276504,
+      "learning_rate": 2.991522876735154e-05,
+      "loss": 0.1946,
+      "step": 390
+    },
+    {
+      "epoch": 2.6133333333333333,
+      "grad_norm": 0.05374331020259679,
+      "learning_rate": 2.9411724702784758e-05,
+      "loss": 0.1948,
+      "step": 392
+    },
+    {
+      "epoch": 2.626666666666667,
+      "grad_norm": 0.0542758123431381,
+      "learning_rate": 2.89107222827234e-05,
+      "loss": 0.1912,
+      "step": 394
+    },
+    {
+      "epoch": 2.64,
+      "grad_norm": 0.0571540947080613,
+      "learning_rate": 2.8412282383075363e-05,
+      "loss": 0.1984,
+      "step": 396
+    },
+    {
+      "epoch": 2.6533333333333333,
+      "grad_norm": 0.05627882046395902,
+      "learning_rate": 2.79164655683813e-05,
+      "loss": 0.2158,
+      "step": 398
+    },
+    {
+      "epoch": 2.6666666666666665,
+      "grad_norm": 0.05259956524533675,
+      "learning_rate": 2.7423332084455544e-05,
+      "loss": 0.1747,
+      "step": 400
+    },
+    {
+      "epoch": 2.68,
+      "grad_norm": 0.05093786217012115,
+      "learning_rate": 2.693294185106562e-05,
+      "loss": 0.2026,
+      "step": 402
+    },
+    {
+      "epoch": 2.6933333333333334,
+      "grad_norm": 0.05277547570213214,
+      "learning_rate": 2.644535445465164e-05,
+      "loss": 0.1813,
+      "step": 404
+    },
+    {
+      "epoch": 2.7066666666666666,
+      "grad_norm": 0.05303823378689789,
+      "learning_rate": 2.5960629141086012e-05,
+      "loss": 0.1819,
+      "step": 406
+    },
+    {
+      "epoch": 2.7199999999999998,
+      "grad_norm": 0.08344410667249319,
+      "learning_rate": 2.547882480847461e-05,
+      "loss": 0.2231,
+      "step": 408
+    },
+    {
+      "epoch": 2.7333333333333334,
+      "grad_norm": 0.05610742314142284,
+      "learning_rate": 2.500000000000001e-05,
+      "loss": 0.1906,
+      "step": 410
+    },
+    {
+      "epoch": 2.7466666666666666,
+      "grad_norm": 0.05772628938455535,
+      "learning_rate": 2.4524212896808263e-05,
+      "loss": 0.1994,
+      "step": 412
+    },
+    {
+      "epoch": 2.76,
+      "grad_norm": 0.05103502451047731,
+      "learning_rate": 2.405152131093926e-05,
+      "loss": 0.1798,
+      "step": 414
+    },
+    {
+      "epoch": 2.7733333333333334,
+      "grad_norm": 0.05325999726886616,
+      "learning_rate": 2.3581982678302063e-05,
+      "loss": 0.1914,
+      "step": 416
+    },
+    {
+      "epoch": 2.7866666666666666,
+      "grad_norm": 0.059490514144748714,
+      "learning_rate": 2.3115654051696095e-05,
+      "loss": 0.1951,
+      "step": 418
+    },
+    {
+      "epoch": 2.8,
+      "grad_norm": 0.04984418568151086,
+      "learning_rate": 2.2652592093878666e-05,
+      "loss": 0.1808,
+      "step": 420
+    },
+    {
+      "epoch": 2.8133333333333335,
+      "grad_norm": 0.05810409569821001,
+      "learning_rate": 2.219285307067997e-05,
+      "loss": 0.203,
+      "step": 422
+    },
+    {
+      "epoch": 2.8266666666666667,
+      "grad_norm": 0.06033075206383965,
+      "learning_rate": 2.1736492844166407e-05,
+      "loss": 0.2046,
+      "step": 424
+    },
+    {
+      "epoch": 2.84,
+      "grad_norm": 0.05658423500567751,
+      "learning_rate": 2.128356686585282e-05,
+      "loss": 0.1889,
+      "step": 426
+    },
+    {
+      "epoch": 2.8533333333333335,
+      "grad_norm": 0.05504606313698527,
+      "learning_rate": 2.0834130169964692e-05,
+      "loss": 0.2025,
+      "step": 428
+    },
+    {
+      "epoch": 2.8666666666666667,
+      "grad_norm": 0.05661096276616458,
+      "learning_rate": 2.0388237366751006e-05,
+      "loss": 0.2063,
+      "step": 430
+    },
+    {
+      "epoch": 2.88,
+      "grad_norm": 0.0491417808278452,
+      "learning_rate": 1.9945942635848748e-05,
+      "loss": 0.1878,
+      "step": 432
+    },
+    {
+      "epoch": 2.8933333333333335,
+      "grad_norm": 0.05320288379451015,
+      "learning_rate": 1.950729971969955e-05,
+      "loss": 0.1934,
+      "step": 434
+    },
+    {
+      "epoch": 2.9066666666666667,
+      "grad_norm": 0.06582732049676977,
+      "learning_rate": 1.9072361917019536e-05,
+      "loss": 0.1912,
+      "step": 436
+    },
+    {
+      "epoch": 2.92,
+      "grad_norm": 0.05333512123441002,
+      "learning_rate": 1.8641182076323148e-05,
+      "loss": 0.2072,
+      "step": 438
+    },
+    {
+      "epoch": 2.9333333333333336,
+      "grad_norm": 0.05868667850689424,
+      "learning_rate": 1.821381258950161e-05,
+      "loss": 0.206,
+      "step": 440
+    },
+    {
+      "epoch": 2.9466666666666668,
+      "grad_norm": 0.05637285498837677,
+      "learning_rate": 1.7790305385456795e-05,
+      "loss": 0.1879,
+      "step": 442
+    },
+    {
+      "epoch": 2.96,
+      "grad_norm": 0.05386393235541858,
+      "learning_rate": 1.7370711923791567e-05,
+      "loss": 0.1958,
+      "step": 444
+    },
+    {
+      "epoch": 2.9733333333333336,
+      "grad_norm": 0.05299769413689916,
+      "learning_rate": 1.6955083188556947e-05,
+      "loss": 0.1925,
+      "step": 446
+    },
+    {
+      "epoch": 2.986666666666667,
+      "grad_norm": 0.05133518316661849,
+      "learning_rate": 1.6543469682057106e-05,
+      "loss": 0.1915,
+      "step": 448
+    },
+    {
+      "epoch": 3.0,
+      "grad_norm": 0.04885194532692743,
+      "learning_rate": 1.6135921418712956e-05,
+      "loss": 0.1739,
+      "step": 450
+    },
+    {
+      "epoch": 3.013333333333333,
+      "grad_norm": 0.06986956461906896,
+      "learning_rate": 1.5732487918985018e-05,
+      "loss": 0.1038,
+      "step": 452
+    },
+    {
+      "epoch": 3.026666666666667,
+      "grad_norm": 0.04645876637991321,
+      "learning_rate": 1.5333218203356243e-05,
+      "loss": 0.1015,
+      "step": 454
+    },
+    {
+      "epoch": 3.04,
+      "grad_norm": 0.07283995125204705,
+      "learning_rate": 1.4938160786375572e-05,
+      "loss": 0.0925,
+      "step": 456
+    },
+    {
+      "epoch": 3.0533333333333332,
+      "grad_norm": 0.07910413235666625,
+      "learning_rate": 1.4547363670763137e-05,
+      "loss": 0.0924,
+      "step": 458
+    },
+    {
+      "epoch": 3.066666666666667,
+      "grad_norm": 0.045001557247815696,
+      "learning_rate": 1.4160874341577446e-05,
+      "loss": 0.0803,
+      "step": 460
+    },
+    {
+      "epoch": 3.08,
+      "grad_norm": 0.06093464723725437,
+      "learning_rate": 1.3778739760445552e-05,
+      "loss": 0.1234,
+      "step": 462
+    },
+    {
+      "epoch": 3.0933333333333333,
+      "grad_norm": 0.05094873498049632,
+      "learning_rate": 1.3401006359856915e-05,
+      "loss": 0.0995,
+      "step": 464
+    },
+    {
+      "epoch": 3.1066666666666665,
+      "grad_norm": 0.047212663662106066,
+      "learning_rate": 1.3027720037521397e-05,
+      "loss": 0.0946,
+      "step": 466
+    },
+    {
+      "epoch": 3.12,
+      "grad_norm": 0.04483003730663884,
+      "learning_rate": 1.2658926150792322e-05,
+      "loss": 0.0878,
+      "step": 468
+    },
+    {
+      "epoch": 3.1333333333333333,
+      "grad_norm": 0.052921476974263,
+      "learning_rate": 1.2294669511155193e-05,
+      "loss": 0.0912,
+      "step": 470
+    },
+    {
+      "epoch": 3.1466666666666665,
+      "grad_norm": 0.06418992888305905,
+      "learning_rate": 1.1934994378782772e-05,
+      "loss": 0.0883,
+      "step": 472
+    },
+    {
+      "epoch": 3.16,
+      "grad_norm": 0.046014320213057874,
+      "learning_rate": 1.157994445715706e-05,
+      "loss": 0.0841,
+      "step": 474
+    },
+    {
+      "epoch": 3.1733333333333333,
+      "grad_norm": 0.04373558000132941,
+      "learning_rate": 1.1229562887758926e-05,
+      "loss": 0.086,
+      "step": 476
+    },
+    {
+      "epoch": 3.1866666666666665,
+      "grad_norm": 0.04549644276317006,
+      "learning_rate": 1.0883892244826172e-05,
+      "loss": 0.0841,
+      "step": 478
+    },
+    {
+      "epoch": 3.2,
+      "grad_norm": 0.04499841039501619,
+      "learning_rate": 1.0542974530180327e-05,
+      "loss": 0.0854,
+      "step": 480
+    },
+    {
+      "epoch": 3.2133333333333334,
+      "grad_norm": 0.04458813671986281,
+      "learning_rate": 1.0206851168123077e-05,
+      "loss": 0.0806,
+      "step": 482
+    },
+    {
+      "epoch": 3.2266666666666666,
+      "grad_norm": 0.04538958754859135,
+      "learning_rate": 9.875563000402948e-06,
+      "loss": 0.0909,
+      "step": 484
+    },
+    {
+      "epoch": 3.24,
+      "grad_norm": 0.04676100575004549,
+      "learning_rate": 9.549150281252633e-06,
+      "loss": 0.0867,
+      "step": 486
+    },
+    {
+      "epoch": 3.2533333333333334,
+      "grad_norm": 0.04775667156033876,
+      "learning_rate": 9.227652672497761e-06,
+      "loss": 0.0837,
+      "step": 488
+    },
+    {
+      "epoch": 3.2666666666666666,
+      "grad_norm": 0.04871933401332665,
+      "learning_rate": 8.911109238737747e-06,
+      "loss": 0.0905,
+      "step": 490
+    },
+    {
+      "epoch": 3.2800000000000002,
+      "grad_norm": 0.048251061213709735,
+      "learning_rate": 8.599558442598998e-06,
+      "loss": 0.0835,
+      "step": 492
+    },
+    {
+      "epoch": 3.2933333333333334,
+      "grad_norm": 0.044194676928720145,
+      "learning_rate": 8.293038140061515e-06,
+      "loss": 0.0821,
+      "step": 494
+    },
+    {
+      "epoch": 3.3066666666666666,
+      "grad_norm": 0.041166919617420104,
+      "learning_rate": 7.991585575858961e-06,
+      "loss": 0.0771,
+      "step": 496
+    },
+    {
+      "epoch": 3.32,
+      "grad_norm": 0.04340450144455652,
+      "learning_rate": 7.695237378953223e-06,
+      "loss": 0.0856,
+      "step": 498
+    },
+    {
+      "epoch": 3.3333333333333335,
+      "grad_norm": 0.04482578603744737,
+      "learning_rate": 7.404029558083653e-06,
+      "loss": 0.091,
+      "step": 500
+    },
+    {
+      "epoch": 3.3466666666666667,
+      "grad_norm": 0.0423271388459505,
+      "learning_rate": 7.1179974973916486e-06,
+      "loss": 0.0982,
+      "step": 502
+    },
+    {
+      "epoch": 3.36,
+      "grad_norm": 0.04729150303046179,
+      "learning_rate": 6.837175952121306e-06,
+      "loss": 0.0897,
+      "step": 504
+    },
+    {
+      "epoch": 3.3733333333333335,
+      "grad_norm": 0.043663930554440095,
+      "learning_rate": 6.561599044396288e-06,
+      "loss": 0.0805,
+      "step": 506
+    },
+    {
+      "epoch": 3.3866666666666667,
+      "grad_norm": 0.047396034103113265,
+      "learning_rate": 6.291300259073724e-06,
+      "loss": 0.0873,
+      "step": 508
+    },
+    {
+      "epoch": 3.4,
+      "grad_norm": 0.041607037767040306,
+      "learning_rate": 6.026312439675552e-06,
+      "loss": 0.0877,
+      "step": 510
+    },
+    {
+      "epoch": 3.413333333333333,
+      "grad_norm": 0.04570055978863104,
+      "learning_rate": 5.766667784397706e-06,
+      "loss": 0.0842,
+      "step": 512
+    },
+    {
+      "epoch": 3.4266666666666667,
+      "grad_norm": 0.047084915674548235,
+      "learning_rate": 5.512397842197847e-06,
+      "loss": 0.0827,
+      "step": 514
+    },
+    {
+      "epoch": 3.44,
+      "grad_norm": 0.04248308044077154,
+      "learning_rate": 5.263533508961827e-06,
+      "loss": 0.0872,
+      "step": 516
+    },
+    {
+      "epoch": 3.453333333333333,
+      "grad_norm": 0.0431950346712709,
+      "learning_rate": 5.020105023749644e-06,
+      "loss": 0.0799,
+      "step": 518
+    },
+    {
+      "epoch": 3.466666666666667,
+      "grad_norm": 0.042080373582692975,
+      "learning_rate": 4.782141965121128e-06,
+      "loss": 0.0848,
+      "step": 520
+    },
+    {
+      "epoch": 3.48,
+      "grad_norm": 0.04306063391695332,
+      "learning_rate": 4.549673247541875e-06,
+      "loss": 0.0856,
+      "step": 522
+    },
+    {
+      "epoch": 3.493333333333333,
+      "grad_norm": 0.04068878744398977,
+      "learning_rate": 4.322727117869951e-06,
+      "loss": 0.0839,
+      "step": 524
+    },
+    {
+      "epoch": 3.506666666666667,
+      "grad_norm": 0.04265900688561506,
+      "learning_rate": 4.101331151923649e-06,
+      "loss": 0.0835,
+      "step": 526
+    },
+    {
+      "epoch": 3.52,
+      "grad_norm": 0.04420106611259069,
+      "learning_rate": 3.885512251130763e-06,
+      "loss": 0.081,
+      "step": 528
+    },
+    {
+      "epoch": 3.533333333333333,
+      "grad_norm": 0.04448511152443269,
+      "learning_rate": 3.675296639259912e-06,
+      "loss": 0.0781,
+      "step": 530
+    },
+    {
+      "epoch": 3.546666666666667,
+      "grad_norm": 0.0497215492121585,
+      "learning_rate": 3.470709859234084e-06,
+      "loss": 0.0861,
+      "step": 532
+    },
+    {
+      "epoch": 3.56,
+      "grad_norm": 0.04327001143276467,
+      "learning_rate": 3.271776770026963e-06,
+      "loss": 0.0821,
+      "step": 534
+    },
+    {
+      "epoch": 3.5733333333333333,
+      "grad_norm": 0.04640151743831662,
+      "learning_rate": 3.078521543642399e-06,
+      "loss": 0.0832,
+      "step": 536
+    },
+    {
+      "epoch": 3.586666666666667,
+      "grad_norm": 0.04561248756380693,
+      "learning_rate": 2.890967662177285e-06,
+      "loss": 0.0825,
+      "step": 538
+    },
+    {
+      "epoch": 3.6,
+      "grad_norm": 0.044823431840851795,
+      "learning_rate": 2.7091379149682685e-06,
+      "loss": 0.0855,
+      "step": 540
+    },
+    {
+      "epoch": 3.6133333333333333,
+      "grad_norm": 0.04467387500603864,
+      "learning_rate": 2.533054395822704e-06,
+      "loss": 0.0808,
+      "step": 542
+    },
+    {
+      "epoch": 3.626666666666667,
+      "grad_norm": 0.04309386698570381,
+      "learning_rate": 2.362738500334055e-06,
+      "loss": 0.0854,
+      "step": 544
+    },
+    {
+      "epoch": 3.64,
+      "grad_norm": 0.044664454082356274,
+      "learning_rate": 2.1982109232821178e-06,
+      "loss": 0.0789,
+      "step": 546
+    },
+    {
+      "epoch": 3.6533333333333333,
+      "grad_norm": 0.04716609021143504,
+      "learning_rate": 2.0394916561185083e-06,
+      "loss": 0.0816,
+      "step": 548
+    },
+    {
+      "epoch": 3.6666666666666665,
+      "grad_norm": 0.04680056577205041,
+      "learning_rate": 1.8865999845374793e-06,
+      "loss": 0.0852,
+      "step": 550
+    },
+    {
+      "epoch": 3.68,
+      "grad_norm": 0.04094899680731965,
+      "learning_rate": 1.7395544861325718e-06,
+      "loss": 0.0825,
+      "step": 552
+    },
+    {
+      "epoch": 3.6933333333333334,
+      "grad_norm": 0.04513598017091709,
+      "learning_rate": 1.5983730281392662e-06,
+      "loss": 0.0844,
+      "step": 554
+    },
+    {
+      "epoch": 3.7066666666666666,
+      "grad_norm": 0.04447335757534387,
+      "learning_rate": 1.463072765264001e-06,
+      "loss": 0.0795,
+      "step": 556
+    },
+    {
+      "epoch": 3.7199999999999998,
+      "grad_norm": 0.05017281858396558,
+      "learning_rate": 1.333670137599713e-06,
+      "loss": 0.0863,
+      "step": 558
+    },
+    {
+      "epoch": 3.7333333333333334,
+      "grad_norm": 0.04423261507584507,
+      "learning_rate": 1.210180868628219e-06,
+      "loss": 0.0754,
+      "step": 560
+    },
+    {
+      "epoch": 3.7466666666666666,
+      "grad_norm": 0.04570867789002808,
+      "learning_rate": 1.0926199633097157e-06,
+      "loss": 0.0849,
+      "step": 562
+    },
+    {
+      "epoch": 3.76,
+      "grad_norm": 0.04507390589448999,
+      "learning_rate": 9.810017062595322e-07,
+      "loss": 0.0806,
+      "step": 564
+    },
+    {
+      "epoch": 3.7733333333333334,
+      "grad_norm": 0.04677585146978537,
+      "learning_rate": 8.753396600124253e-07,
+      "loss": 0.0888,
+      "step": 566
+    },
+    {
+      "epoch": 3.7866666666666666,
+      "grad_norm": 0.04312088768971057,
+      "learning_rate": 7.756466633746407e-07,
+      "loss": 0.0796,
+      "step": 568
+    },
+    {
+      "epoch": 3.8,
+      "grad_norm": 0.0394326107141489,
+      "learning_rate": 6.819348298638839e-07,
+      "loss": 0.0788,
+      "step": 570
+    },
+    {
+      "epoch": 3.8133333333333335,
+      "grad_norm": 0.042695812138021566,
+      "learning_rate": 5.9421554623742e-07,
+      "loss": 0.0754,
+      "step": 572
+    },
+    {
+      "epoch": 3.8266666666666667,
+      "grad_norm": 0.0439412303235543,
+      "learning_rate": 5.124994711084963e-07,
+      "loss": 0.0846,
+      "step": 574
+    },
+    {
+      "epoch": 3.84,
+      "grad_norm": 0.039293878670403605,
+      "learning_rate": 4.367965336512403e-07,
+      "loss": 0.0776,
+      "step": 576
+    },
+    {
+      "epoch": 3.8533333333333335,
+      "grad_norm": 0.043627292403458834,
+      "learning_rate": 3.6711593239417973e-07,
+      "loss": 0.0782,
+      "step": 578
+    },
+    {
+      "epoch": 3.8666666666666667,
+      "grad_norm": 0.04213940119937548,
+      "learning_rate": 3.034661341025258e-07,
+      "loss": 0.0861,
+      "step": 580
+    },
+    {
+      "epoch": 3.88,
+      "grad_norm": 0.04373997508644139,
+      "learning_rate": 2.458548727494292e-07,
+      "loss": 0.0891,
+      "step": 582
+    },
+    {
+      "epoch": 3.8933333333333335,
+      "grad_norm": 0.047501079719067495,
+      "learning_rate": 1.942891485762044e-07,
+      "loss": 0.0862,
+      "step": 584
+    },
+    {
+      "epoch": 3.9066666666666667,
+      "grad_norm": 0.04330962300407197,
+      "learning_rate": 1.4877522724175973e-07,
+      "loss": 0.0804,
+      "step": 586
+    },
+    {
+      "epoch": 3.92,
+      "grad_norm": 0.051714606824807656,
+      "learning_rate": 1.0931863906127327e-07,
+      "loss": 0.0934,
+      "step": 588
+    },
+    {
+      "epoch": 3.9333333333333336,
+      "grad_norm": 0.04613614120222887,
+      "learning_rate": 7.59241783341913e-08,
+      "loss": 0.0948,
+      "step": 590
+    },
+    {
+      "epoch": 3.9466666666666668,
+      "grad_norm": 0.16351246299670308,
+      "learning_rate": 4.859590276170556e-08,
+      "loss": 0.0887,
+      "step": 592
+    },
+    {
+      "epoch": 3.96,
+      "grad_norm": 0.04256032208605194,
+      "learning_rate": 2.7337132953697554e-08,
+      "loss": 0.0866,
+      "step": 594
+    },
+    {
+      "epoch": 3.9733333333333336,
+      "grad_norm": 0.04791037317445386,
+      "learning_rate": 1.215045202527243e-08,
+      "loss": 0.0906,
+      "step": 596
+    },
+    {
+      "epoch": 3.986666666666667,
+      "grad_norm": 0.04423655674990359,
+      "learning_rate": 3.0377052828489683e-09,
+      "loss": 0.0832,
+      "step": 598
+    },
+    {
+      "epoch": 4.0,
+      "grad_norm": 0.0445911005271618,
+      "learning_rate": 0.0,
+      "loss": 0.0851,
+      "step": 600
+    }
+  ],
+  "logging_steps": 2,
+  "max_steps": 600,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 4,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": true
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 2.4291042241019904e+19,
+  "train_batch_size": 2,
+  "trial_name": null,
+  "trial_params": null
+}

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/training_args.bin

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

uccix_v2_instruct_191224_no_parallel_mixture_lr1e-4/checkpoint-600/zero_to_fp32.py

@@ -0,0 +1,592 @@
+#!/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 _has_callable(obj, fn):
+    attr = getattr(obj, fn, None)
+    return callable(attr)
+
+
+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() if _has_callable(shape, 'numel') else math.prod(shape)
+            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)