update model with approx 1.6 epochs training

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "pszemraj/opt-peter-1.3B-1E",
+  "_name_or_path": "pszemraj/opt-peter-1.3B",
   "activation_dropout": 0.0,
   "activation_function": "relu",
   "architectures": [

latest

@@ -0,0 +1 @@
+global_step1016

merges.txt

@@ -1,4 +1,4 @@
-#version: 0.2 - Trained by `huggingface/tokenizers`
+#version: 0.2
 Ġ t
 Ġ a
 h e

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0d2508e76de339f54093ad2f86fc540bfa5ac8c7a037fc5c0c51be088f877e80
-size 5263011731
+oid sha256:4cad2a7fef5856274323ae7e812eec18858a768fd7d3f83bcfc3c2e9d0d3d1de
+size 5263006227

tokenizer_config.json

@@ -1 +1 @@
-{"unk_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "bos_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "eos_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "add_prefix_space": false, "errors": "replace", "pad_token": {"content": "<pad>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "add_bos_token": true, "special_tokens_map_file": null, "name_or_path": "pszemraj/opt-peter-1.3B-1E", "model_max_length": 512}
+{"errors": "replace", "unk_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "bos_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "eos_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "pad_token": {"content": "<pad>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "add_prefix_space": false, "add_bos_token": true, "special_tokens_map_file": null, "name_or_path": "pszemraj/opt-peter-1.3B", "model_max_length": 512, "tokenizer_class": "GPT2Tokenizer"}

trainer_state.json

@@ -1,1776 +1,1242 @@
 {
   "best_metric": null,
   "best_model_checkpoint": null,
-  "epoch": 1.1828941378787112,
-  "global_step": 1461,
+  "epoch": 1.6457290244272493,
+  "global_step": 1016,
   "is_hyper_param_search": false,
   "is_local_process_zero": true,
   "is_world_process_zero": true,
   "log_history": [
     {
-      "epoch": 0.0,
-      "learning_rate": 6e-06,
-      "loss": 1.9846,
+      "epoch": 0.01,
+      "learning_rate": 2.4193548387096776e-06,
+      "loss": 3.4647,
       "step": 5
     },
     {
-      "epoch": 0.01,
-      "learning_rate": 1.2e-05,
-      "loss": 1.8345,
+      "epoch": 0.02,
+      "learning_rate": 4.838709677419355e-06,
+      "loss": 2.4108,
       "step": 10
     },
     {
-      "epoch": 0.01,
-      "learning_rate": 1.8e-05,
-      "loss": 1.7106,
+      "epoch": 0.02,
+      "learning_rate": 7.258064516129032e-06,
+      "loss": 1.5747,
       "step": 15
     },
     {
-      "epoch": 0.02,
-      "learning_rate": 2.4e-05,
-      "loss": 1.6318,
+      "epoch": 0.03,
+      "learning_rate": 9.67741935483871e-06,
+      "loss": 1.5005,
       "step": 20
     },
     {
-      "epoch": 0.02,
-      "learning_rate": 3e-05,
-      "loss": 1.5838,
+      "epoch": 0.04,
+      "learning_rate": 1.2096774193548387e-05,
+      "loss": 1.4694,
       "step": 25
     },
     {
-      "epoch": 0.02,
-      "learning_rate": 2.9999690442167746e-05,
-      "loss": 1.6664,
+      "epoch": 0.05,
+      "learning_rate": 1.4516129032258065e-05,
+      "loss": 1.4399,
       "step": 30
     },
     {
-      "epoch": 0.03,
-      "learning_rate": 2.999876178144779e-05,
-      "loss": 1.5785,
+      "epoch": 0.06,
+      "learning_rate": 1.6935483870967744e-05,
+      "loss": 1.4215,
       "step": 35
     },
     {
-      "epoch": 0.03,
-      "learning_rate": 2.9997214056170024e-05,
-      "loss": 1.6536,
+      "epoch": 0.06,
+      "learning_rate": 1.935483870967742e-05,
+      "loss": 1.4526,
       "step": 40
     },
     {
-      "epoch": 0.04,
-      "learning_rate": 2.9995047330215847e-05,
-      "loss": 1.5866,
+      "epoch": 0.07,
+      "learning_rate": 2.1774193548387097e-05,
+      "loss": 1.4458,
       "step": 45
     },
     {
-      "epoch": 0.04,
-      "learning_rate": 2.9992261693015524e-05,
-      "loss": 1.6324,
+      "epoch": 0.08,
+      "learning_rate": 2.4193548387096773e-05,
+      "loss": 1.4237,
       "step": 50
     },
     {
-      "epoch": 0.04,
-      "learning_rate": 2.9988857259544498e-05,
-      "loss": 1.6263,
+      "epoch": 0.09,
+      "learning_rate": 2.6612903225806453e-05,
+      "loss": 1.4339,
       "step": 55
     },
     {
-      "epoch": 0.05,
-      "learning_rate": 2.9984834170318635e-05,
-      "loss": 1.6481,
+      "epoch": 0.1,
+      "learning_rate": 2.903225806451613e-05,
+      "loss": 1.4294,
       "step": 60
     },
     {
-      "epoch": 0.05,
-      "learning_rate": 2.9980192591388452e-05,
-      "loss": 1.6703,
+      "epoch": 0.11,
+      "learning_rate": 2.999951499529191e-05,
+      "loss": 1.4968,
       "step": 65
     },
     {
-      "epoch": 0.06,
-      "learning_rate": 2.997493271433222e-05,
-      "loss": 1.6566,
+      "epoch": 0.11,
+      "learning_rate": 2.9996551191211948e-05,
+      "loss": 1.4261,
       "step": 70
     },
     {
-      "epoch": 0.06,
-      "learning_rate": 2.9969054756248093e-05,
-      "loss": 1.6168,
+      "epoch": 0.12,
+      "learning_rate": 2.9990893561853812e-05,
+      "loss": 1.4371,
       "step": 75
     },
     {
-      "epoch": 0.06,
-      "learning_rate": 2.9962558959745133e-05,
-      "loss": 1.6836,
+      "epoch": 0.13,
+      "learning_rate": 2.9982543123495507e-05,
+      "loss": 1.4412,
       "step": 80
     },
     {
-      "epoch": 0.07,
-      "learning_rate": 2.9955445592933296e-05,
-      "loss": 1.7089,
+      "epoch": 0.14,
+      "learning_rate": 2.9971501376123366e-05,
+      "loss": 1.4638,
       "step": 85
     },
     {
-      "epoch": 0.07,
-      "learning_rate": 2.994771494941238e-05,
-      "loss": 1.6573,
+      "epoch": 0.15,
+      "learning_rate": 2.9957770303162634e-05,
+      "loss": 1.4498,
       "step": 90
     },
     {
-      "epoch": 0.08,
-      "learning_rate": 2.9939367348259873e-05,
-      "loss": 1.5558,
+      "epoch": 0.15,
+      "learning_rate": 2.9941352371121173e-05,
+      "loss": 1.4393,
       "step": 95
     },
     {
-      "epoch": 0.08,
-      "learning_rate": 2.9930403134017823e-05,
-      "loss": 1.6069,
+      "epoch": 0.16,
+      "learning_rate": 2.992225052914641e-05,
+      "loss": 1.4291,
       "step": 100
     },
     {
-      "epoch": 0.08,
-      "learning_rate": 2.992082267667859e-05,
-      "loss": 1.6357,
+      "epoch": 0.17,
+      "learning_rate": 2.990046820849558e-05,
+      "loss": 1.4587,
       "step": 105
     },
     {
-      "epoch": 0.09,
-      "learning_rate": 2.9910626371669593e-05,
-      "loss": 1.5968,
+      "epoch": 0.18,
+      "learning_rate": 2.9876009321919372e-05,
+      "loss": 1.4272,
       "step": 110
     },
     {
-      "epoch": 0.09,
-      "learning_rate": 2.9899814639836972e-05,
-      "loss": 1.6244,
+      "epoch": 0.19,
+      "learning_rate": 2.9848878262959076e-05,
+      "loss": 1.4622,
       "step": 115
     },
     {
-      "epoch": 0.1,
-      "learning_rate": 2.9888387927428234e-05,
-      "loss": 1.5168,
+      "epoch": 0.19,
+      "learning_rate": 2.981907990515739e-05,
+      "loss": 1.4863,
       "step": 120
     },
     {
-      "epoch": 0.1,
-      "learning_rate": 2.987634670607381e-05,
-      "loss": 1.6162,
+      "epoch": 0.2,
+      "learning_rate": 2.9786619601182965e-05,
+      "loss": 1.4274,
       "step": 125
     },
     {
-      "epoch": 0.11,
-      "learning_rate": 2.9863691472767633e-05,
-      "loss": 1.6014,
+      "epoch": 0.21,
+      "learning_rate": 2.975150318186892e-05,
+      "loss": 1.4382,
       "step": 130
     },
     {
-      "epoch": 0.11,
-      "learning_rate": 2.9850422749846577e-05,
-      "loss": 1.6596,
+      "epoch": 0.22,
+      "learning_rate": 2.9713736955165456e-05,
+      "loss": 1.4413,
       "step": 135
     },
     {
-      "epoch": 0.11,
-      "learning_rate": 2.9836541084968914e-05,
-      "loss": 1.6175,
+      "epoch": 0.23,
+      "learning_rate": 2.9673327705006745e-05,
+      "loss": 1.431,
       "step": 140
     },
     {
-      "epoch": 0.12,
-      "learning_rate": 2.9822047051091735e-05,
-      "loss": 1.6219,
+      "epoch": 0.23,
+      "learning_rate": 2.963028269009235e-05,
+      "loss": 1.4659,
       "step": 145
     },
     {
-      "epoch": 0.12,
-      "learning_rate": 2.9806941246447258e-05,
-      "loss": 1.5605,
+      "epoch": 0.24,
+      "learning_rate": 2.9584609642583337e-05,
+      "loss": 1.4426,
       "step": 150
     },
     {
-      "epoch": 0.13,
-      "learning_rate": 2.9791224294518173e-05,
-      "loss": 1.6062,
+      "epoch": 0.25,
+      "learning_rate": 2.9536316766713357e-05,
+      "loss": 1.4733,
       "step": 155
     },
     {
-      "epoch": 0.13,
-      "learning_rate": 2.9774896844011887e-05,
-      "loss": 1.5561,
+      "epoch": 0.26,
+      "learning_rate": 2.9485412737314923e-05,
+      "loss": 1.4595,
       "step": 160
     },
     {
-      "epoch": 0.13,
-      "learning_rate": 2.9757959568833758e-05,
-      "loss": 1.5929,
+      "epoch": 0.27,
+      "learning_rate": 2.9431906698261136e-05,
+      "loss": 1.4845,
       "step": 165
     },
     {
-      "epoch": 0.14,
-      "learning_rate": 2.9740413168059278e-05,
-      "loss": 1.6087,
+      "epoch": 0.28,
+      "learning_rate": 2.9375808260823192e-05,
+      "loss": 1.5219,
       "step": 170
     },
     {
-      "epoch": 0.14,
-      "learning_rate": 2.9722258365905223e-05,
-      "loss": 1.603,
+      "epoch": 0.28,
+      "learning_rate": 2.931712750194392e-05,
+      "loss": 1.532,
       "step": 175
     },
     {
-      "epoch": 0.15,
-      "learning_rate": 2.9703495911699746e-05,
-      "loss": 1.5866,
+      "epoch": 0.29,
+      "learning_rate": 2.9255874962427638e-05,
+      "loss": 1.4632,
       "step": 180
     },
     {
-      "epoch": 0.15,
-      "learning_rate": 2.9684126579851468e-05,
-      "loss": 1.696,
+      "epoch": 0.3,
+      "learning_rate": 2.9192061645046724e-05,
+      "loss": 1.5057,
       "step": 185
     },
     {
-      "epoch": 0.15,
-      "learning_rate": 2.9664151169817515e-05,
-      "loss": 1.59,
+      "epoch": 0.31,
+      "learning_rate": 2.9125699012565204e-05,
+      "loss": 1.5023,
       "step": 190
     },
     {
-      "epoch": 0.16,
-      "learning_rate": 2.9643570506070493e-05,
-      "loss": 1.5724,
+      "epoch": 0.32,
+      "learning_rate": 2.9056798985679688e-05,
+      "loss": 1.4423,
       "step": 195
     },
     {
-      "epoch": 0.16,
-      "learning_rate": 2.9622385438064493e-05,
-      "loss": 1.5482,
+      "epoch": 0.32,
+      "learning_rate": 2.8985373940878053e-05,
+      "loss": 1.4968,
       "step": 200
     },
     {
-      "epoch": 0.17,
-      "learning_rate": 2.9600596840200022e-05,
-      "loss": 1.6198,
+      "epoch": 0.33,
+      "learning_rate": 2.8911436708216276e-05,
+      "loss": 1.4767,
       "step": 205
     },
     {
-      "epoch": 0.17,
-      "learning_rate": 2.9578205611787877e-05,
-      "loss": 1.6246,
+      "epoch": 0.34,
+      "learning_rate": 2.883500056901376e-05,
+      "loss": 1.4968,
       "step": 210
     },
     {
-      "epoch": 0.17,
-      "learning_rate": 2.955521267701209e-05,
-      "loss": 1.6953,
+      "epoch": 0.35,
+      "learning_rate": 2.875607925346762e-05,
+      "loss": 1.4842,
       "step": 215
     },
     {
-      "epoch": 0.18,
-      "learning_rate": 2.9531618984891737e-05,
-      "loss": 1.6397,
+      "epoch": 0.36,
+      "learning_rate": 2.867468693818634e-05,
+      "loss": 1.4921,
       "step": 220
     },
     {
-      "epoch": 0.18,
-      "learning_rate": 2.9507425509241757e-05,
-      "loss": 1.5797,
+      "epoch": 0.36,
+      "learning_rate": 2.859083824364323e-05,
+      "loss": 1.4969,
       "step": 225
     },
     {
-      "epoch": 0.19,
-      "learning_rate": 2.948263324863282e-05,
-      "loss": 1.63,
+      "epoch": 0.37,
+      "learning_rate": 2.8504548231550143e-05,
+      "loss": 1.4423,
       "step": 230
     },
     {
-      "epoch": 0.19,
-      "learning_rate": 2.945724322635004e-05,
-      "loss": 1.5258,
+      "epoch": 0.38,
+      "learning_rate": 2.8415832402151956e-05,
+      "loss": 1.5165,
       "step": 235
     },
     {
-      "epoch": 0.19,
-      "learning_rate": 2.9431256490350795e-05,
-      "loss": 1.6107,
+      "epoch": 0.39,
+      "learning_rate": 2.832470669144227e-05,
+      "loss": 1.4715,
       "step": 240
     },
     {
-      "epoch": 0.2,
-      "learning_rate": 2.9404674113221433e-05,
-      "loss": 1.5577,
+      "epoch": 0.4,
+      "learning_rate": 2.8231187468300836e-05,
+      "loss": 1.4747,
       "step": 245
     },
     {
-      "epoch": 0.2,
-      "learning_rate": 2.937749719213303e-05,
-      "loss": 1.6001,
+      "epoch": 0.4,
+      "learning_rate": 2.8135291531553192e-05,
+      "loss": 1.466,
       "step": 250
     },
     {
-      "epoch": 0.21,
-      "learning_rate": 2.9349726848796083e-05,
-      "loss": 1.6326,
+      "epoch": 0.41,
+      "learning_rate": 2.8037036106953134e-05,
+      "loss": 1.4609,
       "step": 255
     },
     {
-      "epoch": 0.21,
-      "learning_rate": 2.932136422941424e-05,
-      "loss": 1.5575,
+      "epoch": 0.42,
+      "learning_rate": 2.793643884408843e-05,
+      "loss": 1.4649,
       "step": 260
     },
     {
-      "epoch": 0.21,
-      "learning_rate": 2.929241050463696e-05,
-      "loss": 1.692,
+      "epoch": 0.43,
+      "learning_rate": 2.7833517813210437e-05,
+      "loss": 1.4923,
       "step": 265
     },
     {
-      "epoch": 0.22,
-      "learning_rate": 2.926286686951121e-05,
-      "loss": 1.5985,
+      "epoch": 0.44,
+      "learning_rate": 2.7728291501988173e-05,
+      "loss": 1.4968,
       "step": 270
     },
     {
-      "epoch": 0.22,
-      "learning_rate": 2.9232734543432146e-05,
-      "loss": 1.5635,
+      "epoch": 0.45,
+      "learning_rate": 2.7620778812187338e-05,
+      "loss": 1.4545,
       "step": 275
     },
     {
-      "epoch": 0.23,
-      "learning_rate": 2.920201477009277e-05,
-      "loss": 1.6745,
+      "epoch": 0.45,
+      "learning_rate": 2.7510999056275038e-05,
+      "loss": 1.4791,
       "step": 280
     },
     {
-      "epoch": 0.23,
-      "learning_rate": 2.9170708817432612e-05,
-      "loss": 1.5613,
+      "epoch": 0.46,
+      "learning_rate": 2.739897195395067e-05,
+      "loss": 1.4849,
       "step": 285
     },
     {
-      "epoch": 0.23,
-      "learning_rate": 2.9138817977585383e-05,
-      "loss": 1.6063,
+      "epoch": 0.47,
+      "learning_rate": 2.728471762860369e-05,
+      "loss": 1.4492,
       "step": 290
     },
     {
-      "epoch": 0.24,
-      "learning_rate": 2.910634356682565e-05,
-      "loss": 1.6118,
+      "epoch": 0.48,
+      "learning_rate": 2.716825660369885e-05,
+      "loss": 1.5052,
       "step": 295
     },
     {
-      "epoch": 0.24,
-      "learning_rate": 2.9073286925514504e-05,
-      "loss": 1.5912,
+      "epoch": 0.49,
+      "learning_rate": 2.704960979908957e-05,
+      "loss": 1.4701,
       "step": 300
     },
     {
-      "epoch": 0.25,
-      "learning_rate": 2.9039649418044247e-05,
-      "loss": 1.596,
+      "epoch": 0.49,
+      "learning_rate": 2.6928798527260127e-05,
+      "loss": 1.48,
       "step": 305
     },
     {
-      "epoch": 0.25,
-      "learning_rate": 2.900543243278206e-05,
-      "loss": 1.5501,
+      "epoch": 0.5,
+      "learning_rate": 2.680584448949729e-05,
+      "loss": 1.5158,
       "step": 310
     },
     {
-      "epoch": 0.25,
-      "learning_rate": 2.8970637382012714e-05,
-      "loss": 1.6077,
+      "epoch": 0.51,
+      "learning_rate": 2.6680769771992136e-05,
+      "loss": 1.5047,
       "step": 315
     },
     {
-      "epoch": 0.26,
-      "learning_rate": 2.8935265701880277e-05,
-      "loss": 1.5393,
+      "epoch": 0.52,
+      "learning_rate": 2.6553596841872682e-05,
+      "loss": 1.5023,
       "step": 320
     },
     {
-      "epoch": 0.26,
-      "learning_rate": 2.8899318852328833e-05,
-      "loss": 1.5622,
+      "epoch": 0.53,
+      "learning_rate": 2.6424348543168177e-05,
+      "loss": 1.4932,
       "step": 325
     },
     {
-      "epoch": 0.27,
-      "learning_rate": 2.8862798317042222e-05,
-      "loss": 1.5683,
+      "epoch": 0.53,
+      "learning_rate": 2.6293048092705586e-05,
+      "loss": 1.5241,
       "step": 330
     },
     {
-      "epoch": 0.27,
-      "learning_rate": 2.882570560338281e-05,
-      "loss": 1.5644,
+      "epoch": 0.54,
+      "learning_rate": 2.6159719075939196e-05,
+      "loss": 1.4969,
       "step": 335
     },
     {
-      "epoch": 0.28,
-      "learning_rate": 2.878804224232927e-05,
-      "loss": 1.6027,
+      "epoch": 0.55,
+      "learning_rate": 2.602438544271395e-05,
+      "loss": 1.4983,
       "step": 340
     },
     {
-      "epoch": 0.28,
-      "learning_rate": 2.8749809788413383e-05,
-      "loss": 1.5418,
+      "epoch": 0.56,
+      "learning_rate": 2.5887071502963338e-05,
+      "loss": 1.4801,
       "step": 345
     },
     {
-      "epoch": 0.28,
-      "learning_rate": 2.87110098196559e-05,
-      "loss": 1.5934,
+      "epoch": 0.57,
+      "learning_rate": 2.574780192234264e-05,
+      "loss": 1.4595,
       "step": 350
     },
     {
-      "epoch": 0.29,
-      "learning_rate": 2.8671643937501375e-05,
-      "loss": 1.5265,
+      "epoch": 0.57,
+      "learning_rate": 2.5606601717798212e-05,
+      "loss": 1.4901,
       "step": 355
     },
     {
-      "epoch": 0.29,
-      "learning_rate": 2.8631713766752097e-05,
-      "loss": 1.6044,
+      "epoch": 0.58,
+      "learning_rate": 2.5463496253073726e-05,
+      "loss": 1.4939,
       "step": 360
     },
     {
-      "epoch": 0.3,
-      "learning_rate": 2.859122095550101e-05,
-      "loss": 1.6244,
+      "epoch": 0.59,
+      "learning_rate": 2.531851123415406e-05,
+      "loss": 1.4791,
       "step": 365
     },
     {
-      "epoch": 0.3,
-      "learning_rate": 2.8550167175063705e-05,
-      "loss": 1.5721,
+      "epoch": 0.6,
+      "learning_rate": 2.5171672704647785e-05,
+      "loss": 1.448,
       "step": 370
     },
     {
-      "epoch": 0.3,
-      "learning_rate": 2.850855411990941e-05,
-      "loss": 1.5572,
+      "epoch": 0.61,
+      "learning_rate": 2.502300704110891e-05,
+      "loss": 1.4857,
       "step": 375
     },
     {
-      "epoch": 0.31,
-      "learning_rate": 2.8466383507591083e-05,
-      "loss": 1.5034,
+      "epoch": 0.61,
+      "learning_rate": 2.4872540948298913e-05,
+      "loss": 1.4829,
       "step": 380
     },
     {
-      "epoch": 0.31,
-      "learning_rate": 2.84236570786745e-05,
-      "loss": 1.6026,
+      "epoch": 0.62,
+      "learning_rate": 2.472030145438974e-05,
+      "loss": 1.4786,
       "step": 385
     },
     {
-      "epoch": 0.32,
-      "learning_rate": 2.8380376596666425e-05,
-      "loss": 1.6428,
+      "epoch": 0.63,
+      "learning_rate": 2.4566315906108772e-05,
+      "loss": 1.4118,
       "step": 390
     },
     {
-      "epoch": 0.32,
-      "learning_rate": 2.833654384794181e-05,
-      "loss": 1.5364,
+      "epoch": 0.64,
+      "learning_rate": 2.4410611963826522e-05,
+      "loss": 1.429,
       "step": 395
     },
     {
-      "epoch": 0.32,
-      "learning_rate": 2.8292160641670088e-05,
-      "loss": 1.5064,
+      "epoch": 0.65,
+      "learning_rate": 2.4253217596588036e-05,
+      "loss": 1.4719,
       "step": 400
     },
     {
-      "epoch": 0.33,
-      "learning_rate": 2.8247228809740468e-05,
-      "loss": 1.6338,
+      "epoch": 0.66,
+      "learning_rate": 2.4094161077088784e-05,
+      "loss": 1.5054,
       "step": 405
     },
     {
-      "epoch": 0.33,
-      "learning_rate": 2.820175020668635e-05,
-      "loss": 1.5189,
+      "epoch": 0.66,
+      "learning_rate": 2.3933470976596088e-05,
+      "loss": 1.4493,
       "step": 410
     },
     {
-      "epoch": 0.34,
-      "learning_rate": 2.8155726709608777e-05,
-      "loss": 1.5359,
+      "epoch": 0.67,
+      "learning_rate": 2.3771176159816846e-05,
+      "loss": 1.4957,
       "step": 415
     },
     {
-      "epoch": 0.34,
-      "learning_rate": 2.810916021809894e-05,
-      "loss": 1.6113,
+      "epoch": 0.68,
+      "learning_rate": 2.360730577971259e-05,
+      "loss": 1.4764,
       "step": 420
     },
     {
-      "epoch": 0.34,
-      "learning_rate": 2.8062052654159797e-05,
-      "loss": 1.5724,
+      "epoch": 0.69,
+      "learning_rate": 2.3441889272262742e-05,
+      "loss": 1.5056,
       "step": 425
     },
     {
-      "epoch": 0.35,
-      "learning_rate": 2.8014405962126735e-05,
-      "loss": 1.5174,
+      "epoch": 0.7,
+      "learning_rate": 2.3274956351177037e-05,
+      "loss": 1.4732,
       "step": 430
     },
     {
-      "epoch": 0.35,
-      "learning_rate": 2.7966222108587307e-05,
-      "loss": 1.5741,
+      "epoch": 0.7,
+      "learning_rate": 2.3106537002558074e-05,
+      "loss": 1.5047,
       "step": 435
     },
     {
-      "epoch": 0.36,
-      "learning_rate": 2.7917503082300086e-05,
-      "loss": 1.595,
+      "epoch": 0.71,
+      "learning_rate": 2.293666147951491e-05,
+      "loss": 1.5098,
       "step": 440
     },
     {
-      "epoch": 0.36,
-      "learning_rate": 2.7868250894112555e-05,
-      "loss": 1.6298,
+      "epoch": 0.72,
+      "learning_rate": 2.2765360296728697e-05,
+      "loss": 1.504,
       "step": 445
     },
     {
-      "epoch": 0.36,
-      "learning_rate": 2.7818467576878136e-05,
-      "loss": 1.5871,
+      "epoch": 0.73,
+      "learning_rate": 2.259266422497137e-05,
+      "loss": 1.4742,
       "step": 450
     },
     {
-      "epoch": 0.37,
-      "learning_rate": 2.776815518537226e-05,
-      "loss": 1.5841,
+      "epoch": 0.74,
+      "learning_rate": 2.2418604285578273e-05,
+      "loss": 1.482,
       "step": 455
     },
     {
-      "epoch": 0.37,
-      "learning_rate": 2.7717315796207576e-05,
-      "loss": 1.5128,
+      "epoch": 0.74,
+      "learning_rate": 2.2243211744875818e-05,
+      "loss": 1.4702,
       "step": 460
     },
     {
-      "epoch": 0.38,
-      "learning_rate": 2.7665951507748223e-05,
-      "loss": 1.5055,
+      "epoch": 0.75,
+      "learning_rate": 2.2066518108565137e-05,
+      "loss": 1.4839,
       "step": 465
     },
     {
-      "epoch": 0.38,
-      "learning_rate": 2.7614064440023254e-05,
-      "loss": 1.6022,
+      "epoch": 0.76,
+      "learning_rate": 2.18885551160627e-05,
+      "loss": 1.4942,
       "step": 470
     },
     {
-      "epoch": 0.38,
-      "learning_rate": 2.7561656734639085e-05,
-      "loss": 1.5673,
+      "epoch": 0.77,
+      "learning_rate": 2.1709354734798998e-05,
+      "loss": 1.4497,
       "step": 475
     },
     {
-      "epoch": 0.39,
-      "learning_rate": 2.7508730554691145e-05,
-      "loss": 1.5504,
+      "epoch": 0.78,
+      "learning_rate": 2.152894915447624e-05,
+      "loss": 1.4927,
       "step": 480
     },
     {
-      "epoch": 0.39,
-      "learning_rate": 2.7455288084674565e-05,
-      "loss": 1.6085,
+      "epoch": 0.78,
+      "learning_rate": 2.134737078128611e-05,
+      "loss": 1.4794,
       "step": 485
     },
     {
-      "epoch": 0.4,
-      "learning_rate": 2.7401331530394037e-05,
-      "loss": 1.5598,
+      "epoch": 0.79,
+      "learning_rate": 2.1164652232088674e-05,
+      "loss": 1.5034,
       "step": 490
     },
     {
-      "epoch": 0.4,
-      "learning_rate": 2.7346863118872766e-05,
-      "loss": 1.5559,
+      "epoch": 0.8,
+      "learning_rate": 2.0980826328553416e-05,
+      "loss": 1.5157,
       "step": 495
     },
     {
-      "epoch": 0.4,
-      "learning_rate": 2.729188509826053e-05,
-      "loss": 1.5366,
+      "epoch": 0.81,
+      "learning_rate": 2.0795926091263504e-05,
+      "loss": 1.4681,
       "step": 500
     },
     {
-      "epoch": 0.41,
-      "learning_rate": 2.7236399737740912e-05,
-      "loss": 1.5431,
+      "epoch": 0.82,
+      "learning_rate": 2.0609984733784287e-05,
+      "loss": 1.4732,
       "step": 505
     },
     {
-      "epoch": 0.41,
-      "learning_rate": 2.7180409327437648e-05,
-      "loss": 1.58,
+      "epoch": 0.83,
+      "learning_rate": 2.042303565669719e-05,
+      "loss": 1.5047,
       "step": 510
     },
     {
-      "epoch": 0.42,
-      "learning_rate": 2.712391617832006e-05,
-      "loss": 1.5656,
+      "epoch": 0.83,
+      "learning_rate": 2.0235112441599948e-05,
+      "loss": 1.5093,
       "step": 515
     },
     {
-      "epoch": 0.42,
-      "learning_rate": 2.7066922622107726e-05,
-      "loss": 1.5488,
+      "epoch": 0.84,
+      "learning_rate": 2.0046248845074373e-05,
+      "loss": 1.4515,
       "step": 520
     },
     {
-      "epoch": 0.42,
-      "learning_rate": 2.700943101117421e-05,
-      "loss": 1.563,
+      "epoch": 0.85,
+      "learning_rate": 1.9856478792622666e-05,
+      "loss": 1.5051,
       "step": 525
     },
     {
-      "epoch": 0.43,
-      "learning_rate": 2.6951443718449966e-05,
-      "loss": 1.5845,
+      "epoch": 0.86,
+      "learning_rate": 1.9665836372573397e-05,
+      "loss": 1.5073,
       "step": 530
     },
     {
-      "epoch": 0.43,
-      "learning_rate": 2.689296313732442e-05,
-      "loss": 1.5782,
+      "epoch": 0.87,
+      "learning_rate": 1.947435582995821e-05,
+      "loss": 1.4952,
       "step": 535
     },
     {
-      "epoch": 0.44,
-      "learning_rate": 2.6833991681547158e-05,
-      "loss": 1.6135,
+      "epoch": 0.87,
+      "learning_rate": 1.928207156036043e-05,
+      "loss": 1.4308,
       "step": 540
     },
     {
-      "epoch": 0.44,
-      "learning_rate": 2.677453178512831e-05,
-      "loss": 1.5745,
+      "epoch": 0.88,
+      "learning_rate": 1.9089018103736568e-05,
+      "loss": 1.4588,
       "step": 545
     },
     {
-      "epoch": 0.45,
-      "learning_rate": 2.6714585902238105e-05,
-      "loss": 1.6138,
+      "epoch": 0.89,
+      "learning_rate": 1.8895230138211942e-05,
+      "loss": 1.5477,
       "step": 550
     },
     {
-      "epoch": 0.45,
-      "learning_rate": 2.6654156507105543e-05,
-      "loss": 1.5663,
+      "epoch": 0.9,
+      "learning_rate": 1.870074247385144e-05,
+      "loss": 1.4979,
       "step": 555
     },
     {
-      "epoch": 0.45,
-      "learning_rate": 2.6593246093916307e-05,
-      "loss": 1.5842,
+      "epoch": 0.91,
+      "learning_rate": 1.8505590046406615e-05,
+      "loss": 1.4487,
       "step": 560
     },
     {
-      "epoch": 0.46,
-      "learning_rate": 2.653185717670978e-05,
-      "loss": 1.5424,
+      "epoch": 0.91,
+      "learning_rate": 1.8309807911040186e-05,
+      "loss": 1.4671,
       "step": 565
     },
     {
-      "epoch": 0.46,
-      "learning_rate": 2.6469992289275325e-05,
-      "loss": 1.5542,
+      "epoch": 0.92,
+      "learning_rate": 1.8113431236029078e-05,
+      "loss": 1.4486,
       "step": 570
     },
     {
-      "epoch": 0.47,
-      "learning_rate": 2.6407653985047665e-05,
-      "loss": 1.5571,
+      "epoch": 0.93,
+      "learning_rate": 1.7916495296447162e-05,
+      "loss": 1.483,
       "step": 575
     },
     {
-      "epoch": 0.47,
-      "learning_rate": 2.6344844837001508e-05,
-      "loss": 1.5838,
+      "epoch": 0.94,
+      "learning_rate": 1.771903546782883e-05,
+      "loss": 1.4896,
       "step": 580
     },
     {
-      "epoch": 0.47,
-      "learning_rate": 2.6281567437545347e-05,
-      "loss": 1.5626,
+      "epoch": 0.95,
+      "learning_rate": 1.7521087219814454e-05,
+      "loss": 1.5259,
       "step": 585
     },
     {
-      "epoch": 0.48,
-      "learning_rate": 2.6217824398414473e-05,
-      "loss": 1.5316,
+      "epoch": 0.95,
+      "learning_rate": 1.7322686109779032e-05,
+      "loss": 1.4845,
       "step": 590
     },
     {
-      "epoch": 0.48,
-      "learning_rate": 2.615361835056314e-05,
-      "loss": 1.6038,
+      "epoch": 0.96,
+      "learning_rate": 1.7123867776445e-05,
+      "loss": 1.4866,
       "step": 595
     },
     {
-      "epoch": 0.49,
-      "learning_rate": 2.6088951944056024e-05,
-      "loss": 1.6042,
+      "epoch": 0.97,
+      "learning_rate": 1.692466793348047e-05,
+      "loss": 1.4968,
       "step": 600
     },
     {
-      "epoch": 0.49,
-      "learning_rate": 2.6023827847958802e-05,
-      "loss": 1.5683,
+      "epoch": 0.98,
+      "learning_rate": 1.6725122363084004e-05,
+      "loss": 1.4582,
       "step": 605
     },
     {
-      "epoch": 0.49,
-      "learning_rate": 2.5958248750228018e-05,
-      "loss": 1.5166,
+      "epoch": 0.99,
+      "learning_rate": 1.6525266909557046e-05,
+      "loss": 1.4605,
       "step": 610
     },
     {
-      "epoch": 0.5,
-      "learning_rate": 2.5892217357600104e-05,
-      "loss": 1.5522,
+      "epoch": 1.0,
+      "learning_rate": 1.6325137472865262e-05,
+      "loss": 1.4391,
       "step": 615
     },
     {
-      "epoch": 0.5,
-      "learning_rate": 2.5825736395479708e-05,
-      "loss": 1.6122,
+      "epoch": 1.0,
+      "eval_loss": 3.6259799003601074,
+      "eval_runtime": 967.3528,
+      "eval_samples_per_second": 8.998,
+      "eval_steps_per_second": 2.249,
+      "step": 617
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 1.6124770002189804e-05,
+      "loss": 1.5992,
       "step": 620
     },
     {
-      "epoch": 0.51,
-      "learning_rate": 2.5758808607827153e-05,
-      "loss": 1.5722,
+      "epoch": 1.01,
+      "learning_rate": 1.5924200489469782e-05,
+      "loss": 1.1688,
       "step": 625
     },
     {
-      "epoch": 0.51,
-      "learning_rate": 2.5691436757045232e-05,
-      "loss": 1.5333,
+      "epoch": 1.02,
+      "learning_rate": 1.572346496293706e-05,
+      "loss": 1.1778,
       "step": 630
     },
     {
-      "epoch": 0.51,
-      "learning_rate": 2.5623623623865152e-05,
-      "loss": 1.563,
+      "epoch": 1.03,
+      "learning_rate": 1.5522599480644496e-05,
+      "loss": 1.1652,
       "step": 635
     },
     {
-      "epoch": 0.52,
-      "learning_rate": 2.5555372007231777e-05,
-      "loss": 1.5318,
+      "epoch": 1.04,
+      "learning_rate": 1.532164012398886e-05,
+      "loss": 1.1344,
       "step": 640
     },
     {
-      "epoch": 0.52,
-      "learning_rate": 2.548668472418811e-05,
-      "loss": 1.548,
+      "epoch": 1.05,
+      "learning_rate": 1.5120622991229545e-05,
+      "loss": 1.1474,
       "step": 645
     },
     {
-      "epoch": 0.53,
-      "learning_rate": 2.5417564609759005e-05,
-      "loss": 1.5749,
+      "epoch": 1.05,
+      "learning_rate": 1.4919584191004244e-05,
+      "loss": 1.1457,
       "step": 650
     },
     {
-      "epoch": 0.53,
-      "learning_rate": 2.5348014516834175e-05,
-      "loss": 1.6166,
+      "epoch": 1.06,
+      "learning_rate": 1.471855983584276e-05,
+      "loss": 1.1441,
       "step": 655
     },
     {
-      "epoch": 0.53,
-      "learning_rate": 2.5278037316050417e-05,
-      "loss": 1.61,
+      "epoch": 1.07,
+      "learning_rate": 1.4517586035680145e-05,
+      "loss": 1.1546,
       "step": 660
     },
     {
-      "epoch": 0.54,
-      "learning_rate": 2.5207635895673138e-05,
-      "loss": 1.4852,
+      "epoch": 1.08,
+      "learning_rate": 1.431669889137027e-05,
+      "loss": 1.1526,
       "step": 665
     },
     {
-      "epoch": 0.54,
-      "learning_rate": 2.513681316147715e-05,
-      "loss": 1.5476,
+      "epoch": 1.09,
+      "learning_rate": 1.4115934488201047e-05,
+      "loss": 1.1778,
       "step": 670
     },
     {
-      "epoch": 0.55,
-      "learning_rate": 2.506557203662673e-05,
-      "loss": 1.6291,
+      "epoch": 1.09,
+      "learning_rate": 1.3915328889412434e-05,
+      "loss": 1.1468,
       "step": 675
     },
     {
-      "epoch": 0.55,
-      "learning_rate": 2.4993915461554974e-05,
-      "loss": 1.5232,
+      "epoch": 1.1,
+      "learning_rate": 1.3714918129718418e-05,
+      "loss": 1.1367,
       "step": 680
     },
     {
-      "epoch": 0.55,
-      "learning_rate": 2.4921846393842414e-05,
-      "loss": 1.5827,
+      "epoch": 1.11,
+      "learning_rate": 1.3514738208834112e-05,
+      "loss": 1.1972,
       "step": 685
     },
     {
-      "epoch": 0.56,
-      "learning_rate": 2.484936780809497e-05,
-      "loss": 1.5019,
+      "epoch": 1.12,
+      "learning_rate": 1.331482508500912e-05,
+      "loss": 1.1701,
       "step": 690
     },
     {
-      "epoch": 0.56,
-      "learning_rate": 2.4776482695821154e-05,
-      "loss": 1.5315,
+      "epoch": 1.13,
+      "learning_rate": 1.31152146685684e-05,
+      "loss": 1.1911,
       "step": 695
     },
     {
-      "epoch": 0.57,
-      "learning_rate": 2.470319406530862e-05,
-      "loss": 1.5725,
+      "epoch": 1.13,
+      "learning_rate": 1.2915942815461677e-05,
+      "loss": 1.1758,
       "step": 700
     },
     {
-      "epoch": 0.57,
-      "learning_rate": 2.4629504941499984e-05,
-      "loss": 1.6135,
+      "epoch": 1.14,
+      "learning_rate": 1.2717045320822658e-05,
+      "loss": 1.1486,
       "step": 705
     },
     {
-      "epoch": 0.57,
-      "learning_rate": 2.4555418365867965e-05,
-      "loss": 1.556,
+      "epoch": 1.15,
+      "learning_rate": 1.2518557912539185e-05,
+      "loss": 1.1502,
       "step": 710
     },
     {
-      "epoch": 0.58,
-      "learning_rate": 2.4480937396289856e-05,
-      "loss": 1.5323,
+      "epoch": 1.16,
+      "learning_rate": 1.232051624483541e-05,
+      "loss": 1.1459,
       "step": 715
     },
     {
-      "epoch": 0.58,
-      "learning_rate": 2.4406065106921332e-05,
-      "loss": 1.5726,
+      "epoch": 1.17,
+      "learning_rate": 1.2122955891867278e-05,
+      "loss": 1.1546,
       "step": 720
     },
     {
-      "epoch": 0.59,
-      "learning_rate": 2.4330804588069536e-05,
-      "loss": 1.6077,
+      "epoch": 1.17,
+      "learning_rate": 1.1925912341332324e-05,
+      "loss": 1.165,
       "step": 725
     },
     {
-      "epoch": 0.59,
-      "learning_rate": 2.4255158946065542e-05,
-      "loss": 1.5959,
+      "epoch": 1.18,
+      "learning_rate": 1.1729420988095042e-05,
+      "loss": 1.1548,
       "step": 730
     },
     {
-      "epoch": 0.59,
-      "learning_rate": 2.4179131303136146e-05,
-      "loss": 1.5373,
+      "epoch": 1.19,
+      "learning_rate": 1.1533517127828926e-05,
+      "loss": 1.1454,
       "step": 735
     },
     {
-      "epoch": 0.6,
-      "learning_rate": 2.4102724797274994e-05,
-      "loss": 1.5624,
+      "epoch": 1.2,
+      "learning_rate": 1.1338235950676305e-05,
+      "loss": 1.19,
       "step": 740
     },
     {
-      "epoch": 0.6,
-      "learning_rate": 2.4025942582113067e-05,
-      "loss": 1.5431,
+      "epoch": 1.21,
+      "learning_rate": 1.1143612534927153e-05,
+      "loss": 1.1475,
       "step": 745
     },
     {
-      "epoch": 0.61,
-      "learning_rate": 2.3948787826788495e-05,
-      "loss": 1.5807,
+      "epoch": 1.22,
+      "learning_rate": 1.0949681840717997e-05,
+      "loss": 1.1754,
       "step": 750
     },
     {
-      "epoch": 0.61,
-      "learning_rate": 2.3871263715815802e-05,
-      "loss": 1.5659,
+      "epoch": 1.22,
+      "learning_rate": 1.0756478703752036e-05,
+      "loss": 1.1041,
       "step": 755
     },
     {
-      "epoch": 0.62,
-      "learning_rate": 2.3793373448954406e-05,
-      "loss": 1.5356,
+      "epoch": 1.23,
+      "learning_rate": 1.0564037829041609e-05,
+      "loss": 1.1465,
       "step": 760
     },
     {
-      "epoch": 0.62,
-      "learning_rate": 2.3715120241076602e-05,
-      "loss": 1.5181,
+      "epoch": 1.24,
+      "learning_rate": 1.037239378467416e-05,
+      "loss": 1.1704,
       "step": 765
     },
     {
-      "epoch": 0.62,
-      "learning_rate": 2.3636507322034844e-05,
-      "loss": 1.6218,
+      "epoch": 1.25,
+      "learning_rate": 1.0181580995602766e-05,
+      "loss": 1.1716,
       "step": 770
     },
     {
-      "epoch": 0.63,
-      "learning_rate": 2.355753793652844e-05,
-      "loss": 1.6171,
+      "epoch": 1.26,
+      "learning_rate": 9.991633737462405e-06,
+      "loss": 1.1902,
       "step": 775
     },
     {
-      "epoch": 0.63,
-      "learning_rate": 2.3478215343969623e-05,
-      "loss": 1.5306,
+      "epoch": 1.26,
+      "learning_rate": 9.802586130413045e-06,
+      "loss": 1.1342,
       "step": 780
     },
     {
-      "epoch": 0.64,
-      "learning_rate": 2.3398542818349042e-05,
-      "loss": 1.5147,
+      "epoch": 1.27,
+      "learning_rate": 9.614472133010623e-06,
+      "loss": 1.1848,
       "step": 785
     },
     {
-      "epoch": 0.64,
-      "learning_rate": 2.33185236481006e-05,
-      "loss": 1.5031,
+      "epoch": 1.28,
+      "learning_rate": 9.42732553610712e-06,
+      "loss": 1.1652,
       "step": 790
     },
     {
-      "epoch": 0.64,
-      "learning_rate": 2.323816113596575e-05,
-      "loss": 1.5471,
+      "epoch": 1.29,
+      "learning_rate": 9.241179956780689e-06,
+      "loss": 1.1832,
       "step": 795
     },
     {
-      "epoch": 0.65,
-      "learning_rate": 2.3157458598857164e-05,
-      "loss": 1.5904,
+      "epoch": 1.3,
+      "learning_rate": 9.056068832297041e-06,
+      "loss": 1.1935,
       "step": 800
     },
     {
-      "epoch": 0.65,
-      "learning_rate": 2.3076419367721834e-05,
-      "loss": 1.6055,
+      "epoch": 1.3,
+      "learning_rate": 8.872025414103135e-06,
+      "loss": 1.1714,
       "step": 805
     },
     {
-      "epoch": 0.66,
-      "learning_rate": 2.299504678740359e-05,
-      "loss": 1.6248,
+      "epoch": 1.31,
+      "learning_rate": 8.689082761854213e-06,
+      "loss": 1.1525,
       "step": 810
     },
     {
-      "epoch": 0.66,
-      "learning_rate": 2.2913344216505043e-05,
-      "loss": 1.5905,
+      "epoch": 1.32,
+      "learning_rate": 8.507273737475307e-06,
+      "loss": 1.1701,
       "step": 815
     },
     {
-      "epoch": 0.66,
-      "learning_rate": 2.283131502724896e-05,
-      "loss": 1.4956,
+      "epoch": 1.33,
+      "learning_rate": 8.326630999258286e-06,
+      "loss": 1.1761,
       "step": 820
     },
     {
-      "epoch": 0.67,
-      "learning_rate": 2.2748962605339066e-05,
-      "loss": 1.537,
+      "epoch": 1.34,
+      "learning_rate": 8.14718699599542e-06,
+      "loss": 1.1323,
       "step": 825
     },
     {
-      "epoch": 0.67,
-      "learning_rate": 2.266629034982033e-05,
-      "loss": 1.5609,
+      "epoch": 1.34,
+      "learning_rate": 7.968973961150653e-06,
+      "loss": 1.1768,
       "step": 830
     },
     {
-      "epoch": 0.68,
-      "learning_rate": 2.2583301672938648e-05,
-      "loss": 1.6342,
+      "epoch": 1.35,
+      "learning_rate": 7.792023907069486e-06,
+      "loss": 1.1487,
       "step": 835
     },
     {
-      "epoch": 0.68,
-      "learning_rate": 2.25e-05,
-      "loss": 1.6218,
+      "epoch": 1.36,
+      "learning_rate": 7.616368619228645e-06,
+      "loss": 1.1445,
       "step": 840
     },
     {
-      "epoch": 0.68,
-      "learning_rate": 2.24163887692291e-05,
-      "loss": 1.6035,
+      "epoch": 1.37,
+      "learning_rate": 7.442039650526419e-06,
+      "loss": 1.1348,
       "step": 845
     },
     {
-      "epoch": 0.69,
-      "learning_rate": 2.233247143162746e-05,
-      "loss": 1.4957,
+      "epoch": 1.38,
+      "learning_rate": 7.2690683156148705e-06,
+      "loss": 1.1478,
       "step": 850
     },
     {
-      "epoch": 0.69,
-      "learning_rate": 2.224825145083096e-05,
-      "loss": 1.4736,
+      "epoch": 1.39,
+      "learning_rate": 7.097485685274776e-06,
+      "loss": 1.1625,
       "step": 855
     },
     {
-      "epoch": 0.7,
-      "learning_rate": 2.216373230296689e-05,
-      "loss": 1.581,
+      "epoch": 1.39,
+      "learning_rate": 6.927322580834376e-06,
+      "loss": 1.1331,
       "step": 860
     },
     {
-      "epoch": 0.7,
-      "learning_rate": 2.2078917476510483e-05,
-      "loss": 1.6488,
+      "epoch": 1.4,
+      "learning_rate": 6.758609568632982e-06,
+      "loss": 1.1608,
       "step": 865
     },
     {
-      "epoch": 0.7,
-      "learning_rate": 2.1993810472140908e-05,
-      "loss": 1.5308,
+      "epoch": 1.41,
+      "learning_rate": 6.591376954530345e-06,
+      "loss": 1.1349,
       "step": 870
     },
     {
-      "epoch": 0.71,
-      "learning_rate": 2.190841480259681e-05,
-      "loss": 1.4948,
+      "epoch": 1.42,
+      "learning_rate": 6.4256547784628e-06,
+      "loss": 1.1569,
       "step": 875
     },
     {
-      "epoch": 0.71,
-      "learning_rate": 2.1822733992531294e-05,
-      "loss": 1.5796,
+      "epoch": 1.43,
+      "learning_rate": 6.261472809047244e-06,
+      "loss": 1.152,
       "step": 880
     },
     {
-      "epoch": 0.72,
-      "learning_rate": 2.1736771578366472e-05,
-      "loss": 1.6163,
+      "epoch": 1.43,
+      "learning_rate": 6.098860538233769e-06,
+      "loss": 1.1498,
       "step": 885
     },
     {
-      "epoch": 0.72,
-      "learning_rate": 2.1650531108147493e-05,
-      "loss": 1.5086,
+      "epoch": 1.44,
+      "learning_rate": 5.937847176008072e-06,
+      "loss": 1.1802,
       "step": 890
     },
     {
-      "epoch": 0.72,
-      "learning_rate": 2.1564016141396093e-05,
-      "loss": 1.564,
+      "epoch": 1.45,
+      "learning_rate": 5.778461645144438e-06,
+      "loss": 1.1717,
       "step": 895
     },
     {
-      "epoch": 0.73,
-      "learning_rate": 2.1477230248963675e-05,
-      "loss": 1.6207,
+      "epoch": 1.46,
+      "learning_rate": 5.6207325760103845e-06,
+      "loss": 1.152,
       "step": 900
     },
     {
-      "epoch": 0.73,
-      "learning_rate": 2.139017701288394e-05,
-      "loss": 1.5649,
+      "epoch": 1.47,
+      "learning_rate": 5.464688301423782e-06,
+      "loss": 1.1414,
       "step": 905
     },
     {
-      "epoch": 0.74,
-      "learning_rate": 2.1302860026225027e-05,
-      "loss": 1.5835,
+      "epoch": 1.47,
+      "learning_rate": 5.310356851563427e-06,
+      "loss": 1.154,
       "step": 910
     },
     {
-      "epoch": 0.74,
-      "learning_rate": 2.121528289294122e-05,
-      "loss": 1.5632,
+      "epoch": 1.48,
+      "learning_rate": 5.1577659489340255e-06,
+      "loss": 1.1353,
       "step": 915
     },
     {
-      "epoch": 0.74,
-      "learning_rate": 2.1127449227724186e-05,
-      "loss": 1.5381,
+      "epoch": 1.49,
+      "learning_rate": 5.00694300338638e-06,
+      "loss": 1.1783,
       "step": 920
     },
     {
-      "epoch": 0.75,
-      "learning_rate": 2.1039362655853796e-05,
-      "loss": 1.5546,
+      "epoch": 1.5,
+      "learning_rate": 4.857915107193783e-06,
+      "loss": 1.14,
       "step": 925
     },
     {
-      "epoch": 0.75,
-      "learning_rate": 2.0951026813048475e-05,
-      "loss": 1.4658,
+      "epoch": 1.51,
+      "learning_rate": 4.710709030185422e-06,
+      "loss": 1.1755,
       "step": 930
     },
     {
-      "epoch": 0.76,
-      "learning_rate": 2.0862445345315165e-05,
-      "loss": 1.5073,
+      "epoch": 1.51,
+      "learning_rate": 4.565351214937748e-06,
+      "loss": 1.1121,
       "step": 935
     },
     {
-      "epoch": 0.76,
-      "learning_rate": 2.0773621908798818e-05,
-      "loss": 1.5564,
+      "epoch": 1.52,
+      "learning_rate": 4.421867772024601e-06,
+      "loss": 1.171,
       "step": 940
     },
     {
-      "epoch": 0.76,
-      "learning_rate": 2.068456016963149e-05,
-      "loss": 1.6254,
+      "epoch": 1.53,
+      "learning_rate": 4.280284475326948e-06,
+      "loss": 1.1517,
       "step": 945
     },
     {
-      "epoch": 0.77,
-      "learning_rate": 2.0595263803781037e-05,
-      "loss": 1.5497,
+      "epoch": 1.54,
+      "learning_rate": 4.140626757403176e-06,
+      "loss": 1.1694,
       "step": 950
     },
     {
-      "epoch": 0.77,
-      "learning_rate": 2.050573649689938e-05,
-      "loss": 1.5791,
+      "epoch": 1.55,
+      "learning_rate": 4.002919704920607e-06,
+      "loss": 1.1465,
       "step": 955
     },
     {
-      "epoch": 0.78,
-      "learning_rate": 2.0415981944170405e-05,
-      "loss": 1.4777,
+      "epoch": 1.56,
+      "learning_rate": 3.8671880541492236e-06,
+      "loss": 1.179,
       "step": 960
     },
     {
-      "epoch": 0.78,
-      "learning_rate": 2.0326003850157408e-05,
-      "loss": 1.5316,
+      "epoch": 1.56,
+      "learning_rate": 3.7334561865182694e-06,
+      "loss": 1.1498,
       "step": 965
     },
     {
-      "epoch": 0.78,
-      "learning_rate": 2.0235805928650214e-05,
-      "loss": 1.5288,
+      "epoch": 1.57,
+      "learning_rate": 3.6017481242366503e-06,
+      "loss": 1.1438,
       "step": 970
     },
     {
-      "epoch": 0.79,
-      "learning_rate": 2.0145391902511905e-05,
-      "loss": 1.5519,
+      "epoch": 1.58,
+      "learning_rate": 3.472087525977823e-06,
+      "loss": 1.1647,
       "step": 975
     },
     {
-      "epoch": 0.79,
-      "learning_rate": 2.0054765503525136e-05,
-      "loss": 1.5359,
+      "epoch": 1.59,
+      "learning_rate": 3.3444976826299754e-06,
+      "loss": 1.1475,
       "step": 980
     },
     {
-      "epoch": 0.8,
-      "learning_rate": 1.9963930472238126e-05,
-      "loss": 1.5367,
+      "epoch": 1.6,
+      "learning_rate": 3.219001513112329e-06,
+      "loss": 1.15,
       "step": 985
     },
     {
-      "epoch": 0.8,
-      "learning_rate": 1.9872890557810258e-05,
-      "loss": 1.5804,
+      "epoch": 1.6,
+      "learning_rate": 3.0956215602581933e-06,
+      "loss": 1.1613,
       "step": 990
     },
     {
-      "epoch": 0.81,
-      "learning_rate": 1.978164951785733e-05,
-      "loss": 1.6531,
+      "epoch": 1.61,
+      "learning_rate": 2.974379986765622e-06,
+      "loss": 1.1672,
       "step": 995
     },
     {
-      "epoch": 0.81,
-      "learning_rate": 1.96902111182965e-05,
-      "loss": 1.6579,
+      "epoch": 1.62,
+      "learning_rate": 2.855298571216316e-06,
+      "loss": 1.1702,
       "step": 1000
     },
     {
-      "epoch": 0.81,
-      "learning_rate": 1.959857913319078e-05,
-      "loss": 1.6308,
+      "epoch": 1.63,
+      "learning_rate": 2.738398704163561e-06,
+      "loss": 1.1634,
       "step": 1005
     },
     {
-      "epoch": 0.82,
-      "learning_rate": 1.9506757344593345e-05,
-      "loss": 1.6248,
+      "epoch": 1.64,
+      "learning_rate": 2.6237013842898533e-06,
+      "loss": 1.1756,
       "step": 1010
     },
     {
-      "epoch": 0.82,
-      "learning_rate": 1.941474954239136e-05,
-      "loss": 1.5152,
+      "epoch": 1.64,
+      "learning_rate": 2.511227214634887e-06,
+      "loss": 1.1075,
       "step": 1015
-    },
-    {
-      "epoch": 0.83,
-      "learning_rate": 1.9322559524149603e-05,
-      "loss": 1.5966,
-      "step": 1020
-    },
-    {
-      "epoch": 0.83,
-      "learning_rate": 1.92301910949537e-05,
-      "loss": 1.5624,
-      "step": 1025
-    },
-    {
-      "epoch": 0.83,
-      "learning_rate": 1.9137648067253087e-05,
-      "loss": 1.6706,
-      "step": 1030
-    },
-    {
-      "epoch": 0.84,
-      "learning_rate": 1.9044934260703623e-05,
-      "loss": 1.566,
-      "step": 1035
-    },
-    {
-      "epoch": 0.84,
-      "learning_rate": 1.895205350200998e-05,
-      "loss": 1.508,
-      "step": 1040
-    },
-    {
-      "epoch": 0.85,
-      "learning_rate": 1.885900962476767e-05,
-      "loss": 1.6144,
-      "step": 1045
-    },
-    {
-      "epoch": 0.85,
-      "learning_rate": 1.8765806469304814e-05,
-      "loss": 1.6291,
-      "step": 1050
-    },
-    {
-      "epoch": 0.85,
-      "learning_rate": 1.8672447882523644e-05,
-      "loss": 1.6088,
-      "step": 1055
-    },
-    {
-      "epoch": 0.86,
-      "learning_rate": 1.8578937717741727e-05,
-      "loss": 1.5956,
-      "step": 1060
-    },
-    {
-      "epoch": 0.86,
-      "learning_rate": 1.8485279834532923e-05,
-      "loss": 1.6554,
-      "step": 1065
-    },
-    {
-      "epoch": 0.87,
-      "learning_rate": 1.839147809856807e-05,
-      "loss": 1.6333,
-      "step": 1070
-    },
-    {
-      "epoch": 0.87,
-      "learning_rate": 1.8297536381455434e-05,
-      "loss": 1.6016,
-      "step": 1075
-    },
-    {
-      "epoch": 0.87,
-      "learning_rate": 1.8203458560580934e-05,
-      "loss": 1.5765,
-      "step": 1080
-    },
-    {
-      "epoch": 0.88,
-      "learning_rate": 1.810924851894807e-05,
-      "loss": 1.4411,
-      "step": 1085
-    },
-    {
-      "epoch": 0.88,
-      "learning_rate": 1.801491014501768e-05,
-      "loss": 1.6064,
-      "step": 1090
-    },
-    {
-      "epoch": 0.89,
-      "learning_rate": 1.7920447332547423e-05,
-      "loss": 1.5165,
-      "step": 1095
-    },
-    {
-      "epoch": 0.89,
-      "learning_rate": 1.7825863980431106e-05,
-      "loss": 1.628,
-      "step": 1100
-    },
-    {
-      "epoch": 0.89,
-      "learning_rate": 1.773116399253772e-05,
-      "loss": 1.5792,
-      "step": 1105
-    },
-    {
-      "epoch": 0.9,
-      "learning_rate": 1.7636351277550324e-05,
-      "loss": 1.5908,
-      "step": 1110
-    },
-    {
-      "epoch": 0.9,
-      "learning_rate": 1.7541429748804722e-05,
-      "loss": 1.5477,
-      "step": 1115
-    },
-    {
-      "epoch": 0.91,
-      "learning_rate": 1.744640332412794e-05,
-      "loss": 1.6118,
-      "step": 1120
-    },
-    {
-      "epoch": 0.91,
-      "learning_rate": 1.7351275925676517e-05,
-      "loss": 1.6034,
-      "step": 1125
-    },
-    {
-      "epoch": 0.91,
-      "learning_rate": 1.725605147977461e-05,
-      "loss": 1.6053,
-      "step": 1130
-    },
-    {
-      "epoch": 0.92,
-      "learning_rate": 1.716073391675197e-05,
-      "loss": 1.5638,
-      "step": 1135
-    },
-    {
-      "epoch": 0.92,
-      "learning_rate": 1.7065327170781678e-05,
-      "loss": 1.5809,
-      "step": 1140
-    },
-    {
-      "epoch": 0.93,
-      "learning_rate": 1.69698351797178e-05,
-      "loss": 1.5822,
-      "step": 1145
-    },
-    {
-      "epoch": 0.93,
-      "learning_rate": 1.6874261884932842e-05,
-      "loss": 1.5981,
-      "step": 1150
-    },
-    {
-      "epoch": 0.93,
-      "learning_rate": 1.677861123115506e-05,
-      "loss": 1.5995,
-      "step": 1155
-    },
-    {
-      "epoch": 0.94,
-      "learning_rate": 1.668288716630566e-05,
-      "loss": 1.5323,
-      "step": 1160
-    },
-    {
-      "epoch": 0.94,
-      "learning_rate": 1.6587093641335844e-05,
-      "loss": 1.5981,
-      "step": 1165
-    },
-    {
-      "epoch": 0.95,
-      "learning_rate": 1.6491234610063732e-05,
-      "loss": 1.6448,
-      "step": 1170
-    },
-    {
-      "epoch": 0.95,
-      "learning_rate": 1.6395314029011184e-05,
-      "loss": 1.63,
-      "step": 1175
-    },
-    {
-      "epoch": 0.95,
-      "learning_rate": 1.6299335857240484e-05,
-      "loss": 1.6529,
-      "step": 1180
-    },
-    {
-      "epoch": 0.96,
-      "learning_rate": 1.6203304056190938e-05,
-      "loss": 1.6172,
-      "step": 1185
-    },
-    {
-      "epoch": 0.96,
-      "learning_rate": 1.6107222589515376e-05,
-      "loss": 1.5247,
-      "step": 1190
-    },
-    {
-      "epoch": 0.97,
-      "learning_rate": 1.601109542291652e-05,
-      "loss": 1.5932,
-      "step": 1195
-    },
-    {
-      "epoch": 0.97,
-      "learning_rate": 1.591492652398336e-05,
-      "loss": 1.6186,
-      "step": 1200
-    },
-    {
-      "epoch": 0.98,
-      "learning_rate": 1.5818719862027338e-05,
-      "loss": 1.6524,
-      "step": 1205
-    },
-    {
-      "epoch": 0.98,
-      "learning_rate": 1.572247940791856e-05,
-      "loss": 1.611,
-      "step": 1210
-    },
-    {
-      "epoch": 0.98,
-      "learning_rate": 1.5626209133921863e-05,
-      "loss": 1.5913,
-      "step": 1215
-    },
-    {
-      "epoch": 0.99,
-      "learning_rate": 1.5529913013532894e-05,
-      "loss": 1.7348,
-      "step": 1220
-    },
-    {
-      "epoch": 0.99,
-      "learning_rate": 1.543359502131409e-05,
-      "loss": 1.719,
-      "step": 1225
-    },
-    {
-      "epoch": 1.0,
-      "learning_rate": 1.533725913273064e-05,
-      "loss": 1.6274,
-      "step": 1230
-    },
-    {
-      "epoch": 1.0,
-      "learning_rate": 1.5240909323986384e-05,
-      "loss": 1.5266,
-      "step": 1235
-    },
-    {
-      "epoch": 1.0,
-      "eval_loss": 3.4220237731933594,
-      "eval_runtime": 954.9678,
-      "eval_samples_per_second": 9.114,
-      "eval_steps_per_second": 2.279,
-      "step": 1235
-    },
-    {
-      "epoch": 1.0,
-      "learning_rate": 1.5144549571859711e-05,
-      "loss": 1.3304,
-      "step": 1240
-    },
-    {
-      "epoch": 1.01,
-      "learning_rate": 1.5048183853539425e-05,
-      "loss": 1.0536,
-      "step": 1245
-    },
-    {
-      "epoch": 1.01,
-      "learning_rate": 1.4951816146460574e-05,
-      "loss": 1.0564,
-      "step": 1250
-    },
-    {
-      "epoch": 1.02,
-      "learning_rate": 1.4855450428140291e-05,
-      "loss": 1.0105,
-      "step": 1255
-    },
-    {
-      "epoch": 1.02,
-      "learning_rate": 1.4759090676013616e-05,
-      "loss": 1.0252,
-      "step": 1260
-    },
-    {
-      "epoch": 1.02,
-      "learning_rate": 1.4662740867269361e-05,
-      "loss": 1.0029,
-      "step": 1265
-    },
-    {
-      "epoch": 1.03,
-      "learning_rate": 1.4566404978685912e-05,
-      "loss": 0.9685,
-      "step": 1270
-    },
-    {
-      "epoch": 1.03,
-      "learning_rate": 1.447008698646711e-05,
-      "loss": 1.0095,
-      "step": 1275
-    },
-    {
-      "epoch": 1.04,
-      "learning_rate": 1.4373790866078143e-05,
-      "loss": 0.9759,
-      "step": 1280
-    },
-    {
-      "epoch": 1.04,
-      "learning_rate": 1.4277520592081442e-05,
-      "loss": 1.0217,
-      "step": 1285
-    },
-    {
-      "epoch": 1.04,
-      "learning_rate": 1.4181280137972666e-05,
-      "loss": 0.9472,
-      "step": 1290
-    },
-    {
-      "epoch": 1.05,
-      "learning_rate": 1.4085073476016642e-05,
-      "loss": 1.0432,
-      "step": 1295
-    },
-    {
-      "epoch": 1.05,
-      "learning_rate": 1.3988904577083481e-05,
-      "loss": 1.0355,
-      "step": 1300
-    },
-    {
-      "epoch": 1.06,
-      "learning_rate": 1.389277741048463e-05,
-      "loss": 0.9971,
-      "step": 1305
-    },
-    {
-      "epoch": 1.06,
-      "learning_rate": 1.3796695943809063e-05,
-      "loss": 0.9947,
-      "step": 1310
-    },
-    {
-      "epoch": 1.06,
-      "learning_rate": 1.3700664142759521e-05,
-      "loss": 1.0532,
-      "step": 1315
-    },
-    {
-      "epoch": 1.07,
-      "learning_rate": 1.3604685970988817e-05,
-      "loss": 0.9791,
-      "step": 1320
-    },
-    {
-      "epoch": 1.07,
-      "learning_rate": 1.3508765389936272e-05,
-      "loss": 1.0299,
-      "step": 1325
-    },
-    {
-      "epoch": 1.08,
-      "learning_rate": 1.341290635866416e-05,
-      "loss": 1.036,
-      "step": 1330
-    },
-    {
-      "epoch": 1.08,
-      "learning_rate": 1.3317112833694344e-05,
-      "loss": 1.0534,
-      "step": 1335
-    },
-    {
-      "epoch": 1.08,
-      "learning_rate": 1.3221388768844937e-05,
-      "loss": 1.0144,
-      "step": 1340
-    },
-    {
-      "epoch": 1.09,
-      "learning_rate": 1.3125738115067159e-05,
-      "loss": 1.0696,
-      "step": 1345
-    },
-    {
-      "epoch": 1.09,
-      "learning_rate": 1.30301648202822e-05,
-      "loss": 1.0651,
-      "step": 1350
-    },
-    {
-      "epoch": 1.1,
-      "learning_rate": 1.2934672829218327e-05,
-      "loss": 1.0598,
-      "step": 1355
-    },
-    {
-      "epoch": 1.1,
-      "learning_rate": 1.283926608324804e-05,
-      "loss": 1.0514,
-      "step": 1360
-    },
-    {
-      "epoch": 1.11,
-      "learning_rate": 1.2743948520225391e-05,
-      "loss": 0.9963,
-      "step": 1365
-    },
-    {
-      "epoch": 1.11,
-      "learning_rate": 1.2648724074323492e-05,
-      "loss": 1.0307,
-      "step": 1370
-    },
-    {
-      "epoch": 1.11,
-      "learning_rate": 1.255359667587206e-05,
-      "loss": 1.0354,
-      "step": 1375
-    },
-    {
-      "epoch": 1.12,
-      "learning_rate": 1.2458570251195279e-05,
-      "loss": 1.0965,
-      "step": 1380
-    },
-    {
-      "epoch": 1.12,
-      "learning_rate": 1.2363648722449679e-05,
-      "loss": 1.0484,
-      "step": 1385
-    },
-    {
-      "epoch": 1.13,
-      "learning_rate": 1.2268836007462284e-05,
-      "loss": 1.0421,
-      "step": 1390
-    },
-    {
-      "epoch": 1.13,
-      "learning_rate": 1.21741360195689e-05,
-      "loss": 1.0695,
-      "step": 1395
-    },
-    {
-      "epoch": 1.13,
-      "learning_rate": 1.207955266745258e-05,
-      "loss": 1.0323,
-      "step": 1400
-    },
-    {
-      "epoch": 1.14,
-      "learning_rate": 1.198508985498233e-05,
-      "loss": 1.0874,
-      "step": 1405
-    },
-    {
-      "epoch": 1.14,
-      "learning_rate": 1.1890751481051933e-05,
-      "loss": 1.0702,
-      "step": 1410
-    },
-    {
-      "epoch": 1.15,
-      "learning_rate": 1.179654143941907e-05,
-      "loss": 1.0211,
-      "step": 1415
-    },
-    {
-      "epoch": 1.15,
-      "learning_rate": 1.1702463618544562e-05,
-      "loss": 1.0252,
-      "step": 1420
-    },
-    {
-      "epoch": 1.15,
-      "learning_rate": 1.1608521901431932e-05,
-      "loss": 1.0588,
-      "step": 1425
-    },
-    {
-      "epoch": 1.16,
-      "learning_rate": 1.1514720165467076e-05,
-      "loss": 1.0644,
-      "step": 1430
-    },
-    {
-      "epoch": 1.16,
-      "learning_rate": 1.1421062282258276e-05,
-      "loss": 1.0831,
-      "step": 1435
-    },
-    {
-      "epoch": 1.17,
-      "learning_rate": 1.1327552117476363e-05,
-      "loss": 1.0262,
-      "step": 1440
-    },
-    {
-      "epoch": 1.17,
-      "learning_rate": 1.1234193530695189e-05,
-      "loss": 1.0775,
-      "step": 1445
-    },
-    {
-      "epoch": 1.17,
-      "learning_rate": 1.1140990375232336e-05,
-      "loss": 1.0476,
-      "step": 1450
-    },
-    {
-      "epoch": 1.18,
-      "learning_rate": 1.1047946497990016e-05,
-      "loss": 1.0373,
-      "step": 1455
-    },
-    {
-      "epoch": 1.18,
-      "learning_rate": 1.0955065739296378e-05,
-      "loss": 1.0734,
-      "step": 1460
     }
   ],
-  "max_steps": 2470,
+  "max_steps": 1234,
   "num_train_epochs": 2,
-  "total_flos": 3.470813417302917e+17,
+  "total_flos": 4.827958566725878e+17,
   "trial_name": null,
   "trial_params": null
 }

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b96a0dcac8585789fb3747446a827a33824ee43a375597e136dd02ef9963c321
-size 4335
+oid sha256:41b3235764a2f119a2586ca316b39c8e5986e318c5b3665781ee71759bb93cc0
+size 4271

zero_to_fp32.py

@@ -0,0 +1,484 @@
+#!/usr/bin/env python
+
+# This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example: python zero_to_fp32.py . pytorch_model.bin
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+from collections import OrderedDict
+
+# 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.
+import deepspeed
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION,
+                                            OPTIMIZER_STATE_DICT,
+                                            PARAM_SHAPES,
+                                            SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS,
+                                            ZERO_STAGE,
+                                            PARTITION_COUNT,
+                                            PARAM_SHAPES,
+                                            BUFFER_NAMES)
+
+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_optim_files(checkpoint_dir):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    optim_files = sorted(glob.glob(os.path.join(checkpoint_dir,
+                                                "*_optim_states.pt")),
+                         key=natural_keys)
+
+    if len(optim_files) == 0:
+        raise FileNotFoundError(
+            f"can't find '*_optim_states.pt' files in directory '{checkpoint_dir}'")
+
+    return optim_files
+
+
+def parse_model_state(file):
+    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]
+
+    ds_version = state_dict.get(DS_VERSION, None)
+
+    return buffers, param_shapes, ds_version
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+
+    total_files = len(files)
+    state_dicts = []
+    for f in files:
+        state_dicts.append(torch.load(f, map_location=device))
+
+    if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage == 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    if zero_stage == 2:
+        fp32_flat_groups = [
+            state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key]
+            for i in range(len(state_dicts))
+        ]
+    elif zero_stage == 3:
+        # if there is more than one param group, there will be multiple flattened tensors - one
+        # flattened tensor per group - for simplicity merge them into a single tensor
+        #
+        # XXX: could make the script more memory efficient for when there are multiple groups - it
+        # will require matching the sub-lists of param_shapes for each param group flattened tensor
+
+        fp32_flat_groups = [
+            torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key],
+                      0) for i in range(len(state_dicts))
+        ]
+
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(
+        f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_file = get_model_state_file(ds_checkpoint_dir, zero_stage)
+    buffers, param_shapes, ds_version = parse_model_state(model_file)
+    print(f'Parsing checkpoint created by deepspeed=={ds_version}')
+
+    if zero_stage == 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size,
+                                                          param_shapes,
+                                                          fp32_flat_groups,
+                                                          buffers)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size,
+                                                          param_shapes,
+                                                          fp32_flat_groups,
+                                                          buffers)
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size,
+                                               param_shapes,
+                                               fp32_flat_groups,
+                                               buffers):
+
+    # 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.")
+
+    state_dict = OrderedDict()
+
+    # buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel()
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(
+                    f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} "
+                )
+            state_dict[name] = full_single_fp32_vector.narrow(
+                0,
+                offset,
+                unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(
+                f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(
+        f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements"
+    )
+
+    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 _get_fp32_state_dict_from_zero3_checkpoint(world_size,
+                                               param_shapes,
+                                               fp32_flat_groups,
+                                               buffers):
+
+    # 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
+
+    avail_numel = fp32_flat_groups[0].numel() * world_size
+    # 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
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    state_dict = OrderedDict()
+
+    # buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    # 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"{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 fp32 state dict with {total_params} params {total_numel} elements"
+    )
+
+    return state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
+    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+    """
+
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+    print(f"Saving fp32 state dict to {output_file}")
+    torch.save(state_dict, output_file)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info(f"Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info(f"Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "checkpoint_dir",
+        type=str,
+        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help=
+        "path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)"
+    )
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)