End of training

README.md

@@ -4,7 +4,7 @@ base_model: BridgeTower/bridgetower-large-itm-mlm-itc
 tags:
 - generated_from_trainer
 datasets:
-- newyorker_caption_contest
+- jmhessel/newyorker_caption_contest
 model-index:
 - name: bridgetower
   results: []
@@ -15,7 +15,12 @@ should probably proofread and complete it, then remove this comment. -->
 
 # bridgetower
 
-This model is a fine-tuned version of [BridgeTower/bridgetower-large-itm-mlm-itc](https://huggingface.co/BridgeTower/bridgetower-large-itm-mlm-itc) on the newyorker_caption_contest dataset.
+This model is a fine-tuned version of [BridgeTower/bridgetower-large-itm-mlm-itc](https://huggingface.co/BridgeTower/bridgetower-large-itm-mlm-itc) on the jmhessel/newyorker_caption_contest matching dataset.
+It achieves the following results on the evaluation set:
+- Loss: 0.1360
+- Memory Allocated (gb): 51.27
+- Max Memory Allocated (gb): 57.18
+- Total Memory Available (gb): 94.62
 
 ## Model description
 

all_results.json

@@ -1,15 +1,15 @@
 {
     "epoch": 5.0,
-    "eval_loss": 0.12152472138404846,
-    "eval_runtime": 3.7753,
-    "eval_samples_per_second": 141.222,
-    "eval_steps_per_second": 8.826,
-    "max_memory_allocated (GB)": 60.52,
+    "eval_loss": 0.13596950471401215,
+    "eval_runtime": 5.0456,
+    "eval_samples_per_second": 129.562,
+    "eval_steps_per_second": 8.098,
+    "max_memory_allocated (GB)": 57.18,
     "memory_allocated (GB)": 51.27,
     "total_flos": 3.0598946525952e+16,
     "total_memory_available (GB)": 94.62,
-    "train_loss": 0.06085505417415074,
-    "train_runtime": 1020.8061,
-    "train_samples_per_second": 55.51,
-    "train_steps_per_second": 1.389
+    "train_loss": 0.06072675045655698,
+    "train_runtime": 1077.821,
+    "train_samples_per_second": 52.682,
+    "train_steps_per_second": 1.318
 }

runs/Oct15_14-09-47_workload-ai-workshop/events.out.tfevents.1729002622.workload-ai-workshop.5557.1

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

test_results.json

@@ -1,10 +1,10 @@
 {
     "epoch": 5.0,
-    "eval_loss": 0.12152472138404846,
-    "eval_runtime": 3.7753,
-    "eval_samples_per_second": 141.222,
-    "eval_steps_per_second": 8.826,
-    "max_memory_allocated (GB)": 60.52,
+    "eval_loss": 0.13596950471401215,
+    "eval_runtime": 5.0456,
+    "eval_samples_per_second": 129.562,
+    "eval_steps_per_second": 8.098,
+    "max_memory_allocated (GB)": 57.18,
     "memory_allocated (GB)": 51.27,
     "total_memory_available (GB)": 94.62
 }

tokenizer.json

@@ -1,21 +1,7 @@
 {
   "version": "1.0",
-  "truncation": {
-    "direction": "Right",
-    "max_length": 128,
-    "strategy": "LongestFirst",
-    "stride": 0
-  },
-  "padding": {
-    "strategy": {
-      "Fixed": 128
-    },
-    "direction": "Right",
-    "pad_to_multiple_of": null,
-    "pad_id": 1,
-    "pad_type_id": 0,
-    "pad_token": "<pad>"
-  },
+  "truncation": null,
+  "padding": null,
   "added_tokens": [
     {
       "id": 0,

train_results.json

@@ -1,11 +1,11 @@
 {
     "epoch": 5.0,
-    "max_memory_allocated (GB)": 60.52,
+    "max_memory_allocated (GB)": 57.18,
     "memory_allocated (GB)": 50.57,
     "total_flos": 3.0598946525952e+16,
     "total_memory_available (GB)": 94.62,
-    "train_loss": 0.06085505417415074,
-    "train_runtime": 1020.8061,
-    "train_samples_per_second": 55.51,
-    "train_steps_per_second": 1.389
+    "train_loss": 0.06072675045655698,
+    "train_runtime": 1077.821,
+    "train_samples_per_second": 52.682,
+    "train_steps_per_second": 1.318
 }

trainer_state.json

@@ -10,9 +10,9 @@
   "log_history": [
     {
       "epoch": 0.04081632653061224,
-      "grad_norm": 9.589848518371582,
+      "grad_norm": 16.731555938720703,
       "learning_rate": 9.918367346938776e-06,
-      "loss": 0.2612,
+      "loss": 0.2616,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 10,
@@ -20,9 +20,9 @@
     },
     {
       "epoch": 0.08163265306122448,
-      "grad_norm": 6.701302528381348,
+      "grad_norm": 18.564252853393555,
       "learning_rate": 9.836734693877552e-06,
-      "loss": 0.154,
+      "loss": 0.1554,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 20,
@@ -30,9 +30,9 @@
     },
     {
       "epoch": 0.12244897959183673,
-      "grad_norm": 5.337311267852783,
+      "grad_norm": 7.802953720092773,
       "learning_rate": 9.755102040816327e-06,
-      "loss": 0.1235,
+      "loss": 0.126,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 30,
@@ -40,9 +40,9 @@
     },
     {
       "epoch": 0.16326530612244897,
-      "grad_norm": 4.5042338371276855,
+      "grad_norm": 4.599184036254883,
       "learning_rate": 9.673469387755103e-06,
-      "loss": 0.1096,
+      "loss": 0.1103,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 40,
@@ -50,9 +50,9 @@
     },
     {
       "epoch": 0.20408163265306123,
-      "grad_norm": 4.461822032928467,
+      "grad_norm": 6.734092712402344,
       "learning_rate": 9.591836734693878e-06,
-      "loss": 0.1196,
+      "loss": 0.1217,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 50,
@@ -60,9 +60,9 @@
     },
     {
       "epoch": 0.24489795918367346,
-      "grad_norm": 2.2825701236724854,
+      "grad_norm": 1.7976917028427124,
       "learning_rate": 9.510204081632653e-06,
-      "loss": 0.0805,
+      "loss": 0.08,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 60,
@@ -70,9 +70,9 @@
     },
     {
       "epoch": 0.2857142857142857,
-      "grad_norm": 3.725268602371216,
+      "grad_norm": 4.0598835945129395,
       "learning_rate": 9.42857142857143e-06,
-      "loss": 0.1026,
+      "loss": 0.1023,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 70,
@@ -80,9 +80,9 @@
     },
     {
       "epoch": 0.32653061224489793,
-      "grad_norm": 1.707739233970642,
+      "grad_norm": 2.0350396633148193,
       "learning_rate": 9.346938775510204e-06,
-      "loss": 0.1111,
+      "loss": 0.1128,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 80,
@@ -90,9 +90,9 @@
     },
     {
       "epoch": 0.3673469387755102,
-      "grad_norm": 4.5863938331604,
+      "grad_norm": 3.785065174102783,
       "learning_rate": 9.26530612244898e-06,
-      "loss": 0.0856,
+      "loss": 0.0842,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 90,
@@ -100,9 +100,9 @@
     },
     {
       "epoch": 0.40816326530612246,
-      "grad_norm": 11.972647666931152,
+      "grad_norm": 1.9926950931549072,
       "learning_rate": 9.183673469387756e-06,
-      "loss": 0.0759,
+      "loss": 0.0728,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 100,
@@ -110,9 +110,9 @@
     },
     {
       "epoch": 0.4489795918367347,
-      "grad_norm": 4.550654888153076,
+      "grad_norm": 4.878537178039551,
       "learning_rate": 9.102040816326532e-06,
-      "loss": 0.0717,
+      "loss": 0.0692,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 110,
@@ -120,9 +120,9 @@
     },
     {
       "epoch": 0.4897959183673469,
-      "grad_norm": 4.418276786804199,
+      "grad_norm": 2.5495357513427734,
       "learning_rate": 9.020408163265307e-06,
-      "loss": 0.0717,
+      "loss": 0.0691,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 120,
@@ -130,9 +130,9 @@
     },
     {
       "epoch": 0.5306122448979592,
-      "grad_norm": 1.651443600654602,
+      "grad_norm": 3.0813372135162354,
       "learning_rate": 8.938775510204082e-06,
-      "loss": 0.0581,
+      "loss": 0.0597,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 130,
@@ -140,9 +140,9 @@
     },
     {
       "epoch": 0.5714285714285714,
-      "grad_norm": 1.5251814126968384,
+      "grad_norm": 1.2760021686553955,
       "learning_rate": 8.857142857142858e-06,
-      "loss": 0.0481,
+      "loss": 0.0507,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 140,
@@ -150,9 +150,9 @@
     },
     {
       "epoch": 0.6122448979591837,
-      "grad_norm": 1.7455183267593384,
+      "grad_norm": 2.0892932415008545,
       "learning_rate": 8.775510204081633e-06,
-      "loss": 0.0625,
+      "loss": 0.0634,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 150,
@@ -160,9 +160,9 @@
     },
     {
       "epoch": 0.6530612244897959,
-      "grad_norm": 1.7588891983032227,
+      "grad_norm": 2.0191688537597656,
       "learning_rate": 8.69387755102041e-06,
-      "loss": 0.0711,
+      "loss": 0.0718,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 160,
@@ -170,9 +170,9 @@
     },
     {
       "epoch": 0.6938775510204082,
-      "grad_norm": 2.7675328254699707,
+      "grad_norm": 2.4970450401306152,
       "learning_rate": 8.612244897959184e-06,
-      "loss": 0.0747,
+      "loss": 0.0753,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 170,
@@ -180,9 +180,9 @@
     },
     {
       "epoch": 0.7346938775510204,
-      "grad_norm": 1.781469464302063,
+      "grad_norm": 1.5428143739700317,
       "learning_rate": 8.530612244897961e-06,
-      "loss": 0.061,
+      "loss": 0.0612,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 180,
@@ -190,9 +190,9 @@
     },
     {
       "epoch": 0.7755102040816326,
-      "grad_norm": 2.3728435039520264,
+      "grad_norm": 3.6041452884674072,
       "learning_rate": 8.448979591836736e-06,
-      "loss": 0.0588,
+      "loss": 0.0573,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 190,
@@ -200,9 +200,9 @@
     },
     {
       "epoch": 0.8163265306122449,
-      "grad_norm": 0.8711996674537659,
+      "grad_norm": 1.0115529298782349,
       "learning_rate": 8.36734693877551e-06,
-      "loss": 0.062,
+      "loss": 0.0631,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 200,
@@ -210,9 +210,9 @@
     },
     {
       "epoch": 0.8571428571428571,
-      "grad_norm": 1.1986733675003052,
+      "grad_norm": 0.8029147982597351,
       "learning_rate": 8.285714285714287e-06,
-      "loss": 0.0627,
+      "loss": 0.0643,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 210,
@@ -220,9 +220,9 @@
     },
     {
       "epoch": 0.8979591836734694,
-      "grad_norm": 2.8968520164489746,
+      "grad_norm": 1.130996584892273,
       "learning_rate": 8.204081632653062e-06,
-      "loss": 0.0604,
+      "loss": 0.0608,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 220,
@@ -230,9 +230,9 @@
     },
     {
       "epoch": 0.9387755102040817,
-      "grad_norm": 0.8414793610572815,
+      "grad_norm": 0.7962849140167236,
       "learning_rate": 8.122448979591837e-06,
-      "loss": 0.0559,
+      "loss": 0.0535,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 230,
@@ -240,9 +240,9 @@
     },
     {
       "epoch": 0.9795918367346939,
-      "grad_norm": 0.7434167861938477,
+      "grad_norm": 2.3987386226654053,
       "learning_rate": 8.040816326530613e-06,
-      "loss": 0.0498,
+      "loss": 0.0499,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 240,
@@ -250,9 +250,9 @@
     },
     {
       "epoch": 1.0204081632653061,
-      "grad_norm": 0.8703041076660156,
+      "grad_norm": 1.084067463874817,
       "learning_rate": 7.959183673469388e-06,
-      "loss": 0.0618,
+      "loss": 0.0612,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 250,
@@ -260,9 +260,9 @@
     },
     {
       "epoch": 1.0612244897959184,
-      "grad_norm": 1.0856379270553589,
+      "grad_norm": 3.247530460357666,
       "learning_rate": 7.877551020408164e-06,
-      "loss": 0.056,
+      "loss": 0.055,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 260,
@@ -270,9 +270,9 @@
     },
     {
       "epoch": 1.1020408163265305,
-      "grad_norm": 0.8847401142120361,
+      "grad_norm": 1.5961194038391113,
       "learning_rate": 7.79591836734694e-06,
-      "loss": 0.0625,
+      "loss": 0.0632,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 270,
@@ -280,9 +280,9 @@
     },
     {
       "epoch": 1.1428571428571428,
-      "grad_norm": 1.5929882526397705,
+      "grad_norm": 0.32916340231895447,
       "learning_rate": 7.714285714285716e-06,
-      "loss": 0.0571,
+      "loss": 0.0565,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 280,
@@ -290,9 +290,9 @@
     },
     {
       "epoch": 1.183673469387755,
-      "grad_norm": 0.8007532954216003,
+      "grad_norm": 0.6009345054626465,
       "learning_rate": 7.63265306122449e-06,
-      "loss": 0.0511,
+      "loss": 0.0503,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 290,
@@ -300,9 +300,9 @@
     },
     {
       "epoch": 1.2244897959183674,
-      "grad_norm": 1.2002859115600586,
+      "grad_norm": 9.806236267089844,
       "learning_rate": 7.551020408163265e-06,
-      "loss": 0.065,
+      "loss": 0.0635,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 300,
@@ -310,7 +310,7 @@
     },
     {
       "epoch": 1.2653061224489797,
-      "grad_norm": 12.871713638305664,
+      "grad_norm": 1.231447696685791,
       "learning_rate": 7.469387755102041e-06,
       "loss": 0.0664,
       "max_memory_allocated (GB)": 57.18,
@@ -320,9 +320,9 @@
     },
     {
       "epoch": 1.306122448979592,
-      "grad_norm": 2.46173357963562,
+      "grad_norm": 1.016727328300476,
       "learning_rate": 7.387755102040817e-06,
-      "loss": 0.0495,
+      "loss": 0.0502,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 320,
@@ -330,9 +330,9 @@
     },
     {
       "epoch": 1.346938775510204,
-      "grad_norm": 0.860598087310791,
+      "grad_norm": 0.7703081965446472,
       "learning_rate": 7.306122448979592e-06,
-      "loss": 0.0603,
+      "loss": 0.0599,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 330,
@@ -340,9 +340,9 @@
     },
     {
       "epoch": 1.3877551020408163,
-      "grad_norm": 2.5583598613739014,
+      "grad_norm": 5.310096740722656,
       "learning_rate": 7.224489795918368e-06,
-      "loss": 0.0547,
+      "loss": 0.0541,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 340,
@@ -350,9 +350,9 @@
     },
     {
       "epoch": 1.4285714285714286,
-      "grad_norm": 0.37155964970588684,
+      "grad_norm": 1.1444560289382935,
       "learning_rate": 7.1428571428571436e-06,
-      "loss": 0.048,
+      "loss": 0.0493,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 350,
@@ -360,9 +360,9 @@
     },
     {
       "epoch": 1.469387755102041,
-      "grad_norm": 1.808316707611084,
+      "grad_norm": 3.1823084354400635,
       "learning_rate": 7.061224489795919e-06,
-      "loss": 0.0462,
+      "loss": 0.0459,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 360,
@@ -370,9 +370,9 @@
     },
     {
       "epoch": 1.510204081632653,
-      "grad_norm": 1.0183931589126587,
+      "grad_norm": 1.1924108266830444,
       "learning_rate": 6.979591836734695e-06,
-      "loss": 0.0594,
+      "loss": 0.058,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 370,
@@ -380,9 +380,9 @@
     },
     {
       "epoch": 1.5510204081632653,
-      "grad_norm": 0.5249583721160889,
+      "grad_norm": 3.677968740463257,
       "learning_rate": 6.8979591836734705e-06,
-      "loss": 0.0479,
+      "loss": 0.0468,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 380,
@@ -390,9 +390,9 @@
     },
     {
       "epoch": 1.5918367346938775,
-      "grad_norm": 1.1005572080612183,
+      "grad_norm": 1.0082957744598389,
       "learning_rate": 6.816326530612245e-06,
-      "loss": 0.0649,
+      "loss": 0.0652,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 390,
@@ -400,9 +400,9 @@
     },
     {
       "epoch": 1.6326530612244898,
-      "grad_norm": 0.6047573089599609,
+      "grad_norm": 0.6584922075271606,
       "learning_rate": 6.734693877551021e-06,
-      "loss": 0.0607,
+      "loss": 0.0605,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 400,
@@ -410,9 +410,9 @@
     },
     {
       "epoch": 1.6734693877551021,
-      "grad_norm": 0.7261654734611511,
+      "grad_norm": 1.2264763116836548,
       "learning_rate": 6.653061224489797e-06,
-      "loss": 0.0606,
+      "loss": 0.0599,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 410,
@@ -420,9 +420,9 @@
     },
     {
       "epoch": 1.7142857142857144,
-      "grad_norm": 0.848527193069458,
+      "grad_norm": 1.824320912361145,
       "learning_rate": 6.571428571428572e-06,
-      "loss": 0.0532,
+      "loss": 0.053,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 420,
@@ -430,9 +430,9 @@
     },
     {
       "epoch": 1.7551020408163265,
-      "grad_norm": 0.23483288288116455,
+      "grad_norm": 0.24405838549137115,
       "learning_rate": 6.489795918367348e-06,
-      "loss": 0.068,
+      "loss": 0.067,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 430,
@@ -440,9 +440,9 @@
     },
     {
       "epoch": 1.7959183673469388,
-      "grad_norm": 2.0767459869384766,
+      "grad_norm": 1.0410544872283936,
       "learning_rate": 6.408163265306124e-06,
-      "loss": 0.0617,
+      "loss": 0.06,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 440,
@@ -450,9 +450,9 @@
     },
     {
       "epoch": 1.836734693877551,
-      "grad_norm": 0.5654011368751526,
+      "grad_norm": 0.39765527844429016,
       "learning_rate": 6.326530612244899e-06,
-      "loss": 0.044,
+      "loss": 0.043,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 450,
@@ -460,9 +460,9 @@
     },
     {
       "epoch": 1.8775510204081631,
-      "grad_norm": 0.7382919788360596,
+      "grad_norm": 0.29981690645217896,
       "learning_rate": 6.244897959183675e-06,
-      "loss": 0.0537,
+      "loss": 0.0533,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 460,
@@ -470,9 +470,9 @@
     },
     {
       "epoch": 1.9183673469387754,
-      "grad_norm": 1.3547204732894897,
+      "grad_norm": 0.6634105443954468,
       "learning_rate": 6.163265306122449e-06,
-      "loss": 0.0432,
+      "loss": 0.0428,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 470,
@@ -480,9 +480,9 @@
     },
     {
       "epoch": 1.9591836734693877,
-      "grad_norm": 0.19681082665920258,
+      "grad_norm": 0.2060549110174179,
       "learning_rate": 6.0816326530612245e-06,
-      "loss": 0.0498,
+      "loss": 0.0505,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 480,
@@ -490,9 +490,9 @@
     },
     {
       "epoch": 2.0,
-      "grad_norm": 1.109737515449524,
+      "grad_norm": 1.0800402164459229,
       "learning_rate": 6e-06,
-      "loss": 0.0639,
+      "loss": 0.0647,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 490,
@@ -500,9 +500,9 @@
     },
     {
       "epoch": 2.0408163265306123,
-      "grad_norm": 0.5894625782966614,
+      "grad_norm": 0.6078592538833618,
       "learning_rate": 5.918367346938776e-06,
-      "loss": 0.0593,
+      "loss": 0.0596,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 500,
@@ -510,9 +510,9 @@
     },
     {
       "epoch": 2.0816326530612246,
-      "grad_norm": 0.7122555375099182,
+      "grad_norm": 0.6262472867965698,
       "learning_rate": 5.8367346938775515e-06,
-      "loss": 0.0498,
+      "loss": 0.0495,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 510,
@@ -520,9 +520,9 @@
     },
     {
       "epoch": 2.122448979591837,
-      "grad_norm": 0.8958902955055237,
+      "grad_norm": 1.4308090209960938,
       "learning_rate": 5.755102040816327e-06,
-      "loss": 0.0457,
+      "loss": 0.0461,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 520,
@@ -530,7 +530,7 @@
     },
     {
       "epoch": 2.163265306122449,
-      "grad_norm": 11.620415687561035,
+      "grad_norm": 0.8761769533157349,
       "learning_rate": 5.673469387755103e-06,
       "loss": 0.0626,
       "max_memory_allocated (GB)": 57.18,
@@ -540,9 +540,9 @@
     },
     {
       "epoch": 2.204081632653061,
-      "grad_norm": 0.3538230061531067,
+      "grad_norm": 0.6150110363960266,
       "learning_rate": 5.591836734693878e-06,
-      "loss": 0.0584,
+      "loss": 0.0576,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 540,
@@ -550,9 +550,9 @@
     },
     {
       "epoch": 2.2448979591836733,
-      "grad_norm": 1.5313146114349365,
+      "grad_norm": 9.280069351196289,
       "learning_rate": 5.510204081632653e-06,
-      "loss": 0.0627,
+      "loss": 0.0634,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 550,
@@ -560,9 +560,9 @@
     },
     {
       "epoch": 2.2857142857142856,
-      "grad_norm": 1.3519809246063232,
+      "grad_norm": 2.136209726333618,
       "learning_rate": 5.428571428571429e-06,
-      "loss": 0.0572,
+      "loss": 0.0579,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 560,
@@ -570,9 +570,9 @@
     },
     {
       "epoch": 2.326530612244898,
-      "grad_norm": 1.0263270139694214,
+      "grad_norm": 0.9547491073608398,
       "learning_rate": 5.3469387755102045e-06,
-      "loss": 0.0585,
+      "loss": 0.0589,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 570,
@@ -580,9 +580,9 @@
     },
     {
       "epoch": 2.36734693877551,
-      "grad_norm": 0.8926671147346497,
+      "grad_norm": 2.1748363971710205,
       "learning_rate": 5.26530612244898e-06,
-      "loss": 0.0673,
+      "loss": 0.0672,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 580,
@@ -590,9 +590,9 @@
     },
     {
       "epoch": 2.4081632653061225,
-      "grad_norm": 0.3185974955558777,
+      "grad_norm": 0.9653811454772949,
       "learning_rate": 5.183673469387756e-06,
-      "loss": 0.0537,
+      "loss": 0.0536,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 590,
@@ -600,9 +600,9 @@
     },
     {
       "epoch": 2.4489795918367347,
-      "grad_norm": 0.944624662399292,
+      "grad_norm": 0.7852123379707336,
       "learning_rate": 5.1020408163265315e-06,
-      "loss": 0.0442,
+      "loss": 0.0443,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 600,
@@ -610,9 +610,9 @@
     },
     {
       "epoch": 2.489795918367347,
-      "grad_norm": 0.32796111702919006,
+      "grad_norm": 0.7405697107315063,
       "learning_rate": 5.020408163265307e-06,
-      "loss": 0.0413,
+      "loss": 0.041,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 610,
@@ -620,9 +620,9 @@
     },
     {
       "epoch": 2.5306122448979593,
-      "grad_norm": 0.7929801940917969,
+      "grad_norm": 0.7166327834129333,
       "learning_rate": 4.938775510204082e-06,
-      "loss": 0.0428,
+      "loss": 0.0431,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 620,
@@ -630,9 +630,9 @@
     },
     {
       "epoch": 2.571428571428571,
-      "grad_norm": 0.910254955291748,
+      "grad_norm": 0.5675875544548035,
       "learning_rate": 4.857142857142858e-06,
-      "loss": 0.0813,
+      "loss": 0.0807,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 630,
@@ -640,9 +640,9 @@
     },
     {
       "epoch": 2.612244897959184,
-      "grad_norm": 1.101942539215088,
+      "grad_norm": 0.954290509223938,
       "learning_rate": 4.775510204081633e-06,
-      "loss": 0.0495,
+      "loss": 0.0492,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 640,
@@ -650,9 +650,9 @@
     },
     {
       "epoch": 2.6530612244897958,
-      "grad_norm": 0.7182526588439941,
+      "grad_norm": 0.9061315655708313,
       "learning_rate": 4.693877551020409e-06,
-      "loss": 0.0471,
+      "loss": 0.0472,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 650,
@@ -660,9 +660,9 @@
     },
     {
       "epoch": 2.693877551020408,
-      "grad_norm": 0.8068158626556396,
+      "grad_norm": 0.7333698868751526,
       "learning_rate": 4.612244897959184e-06,
-      "loss": 0.0469,
+      "loss": 0.0474,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 660,
@@ -670,9 +670,9 @@
     },
     {
       "epoch": 2.7346938775510203,
-      "grad_norm": 1.2375913858413696,
+      "grad_norm": 1.7820810079574585,
       "learning_rate": 4.530612244897959e-06,
-      "loss": 0.0857,
+      "loss": 0.0869,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 670,
@@ -680,9 +680,9 @@
     },
     {
       "epoch": 2.7755102040816326,
-      "grad_norm": 1.1524357795715332,
+      "grad_norm": 1.3184058666229248,
       "learning_rate": 4.448979591836735e-06,
-      "loss": 0.0488,
+      "loss": 0.0494,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 680,
@@ -690,9 +690,9 @@
     },
     {
       "epoch": 2.816326530612245,
-      "grad_norm": 0.3913586437702179,
+      "grad_norm": 0.4212433397769928,
       "learning_rate": 4.367346938775511e-06,
-      "loss": 0.0451,
+      "loss": 0.0445,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 690,
@@ -700,9 +700,9 @@
     },
     {
       "epoch": 2.857142857142857,
-      "grad_norm": 0.47935113310813904,
+      "grad_norm": 0.39547938108444214,
       "learning_rate": 4.2857142857142855e-06,
-      "loss": 0.0433,
+      "loss": 0.0429,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 700,
@@ -710,9 +710,9 @@
     },
     {
       "epoch": 2.8979591836734695,
-      "grad_norm": 0.8084143996238708,
+      "grad_norm": 1.0822113752365112,
       "learning_rate": 4.204081632653061e-06,
-      "loss": 0.0548,
+      "loss": 0.0549,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 710,
@@ -720,9 +720,9 @@
     },
     {
       "epoch": 2.938775510204082,
-      "grad_norm": 1.7315497398376465,
+      "grad_norm": 1.0838605165481567,
       "learning_rate": 4.122448979591837e-06,
-      "loss": 0.0587,
+      "loss": 0.0582,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 720,
@@ -730,9 +730,9 @@
     },
     {
       "epoch": 2.979591836734694,
-      "grad_norm": 0.20743349194526672,
+      "grad_norm": 0.2726249098777771,
       "learning_rate": 4.040816326530612e-06,
-      "loss": 0.0342,
+      "loss": 0.0341,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 730,
@@ -740,7 +740,7 @@
     },
     {
       "epoch": 3.020408163265306,
-      "grad_norm": 0.8024761080741882,
+      "grad_norm": 0.836703360080719,
       "learning_rate": 3.959183673469388e-06,
       "loss": 0.053,
       "max_memory_allocated (GB)": 57.18,
@@ -750,9 +750,9 @@
     },
     {
       "epoch": 3.061224489795918,
-      "grad_norm": 0.45326006412506104,
+      "grad_norm": 0.6878814697265625,
       "learning_rate": 3.877551020408164e-06,
-      "loss": 0.0619,
+      "loss": 0.0618,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 750,
@@ -760,7 +760,7 @@
     },
     {
       "epoch": 3.1020408163265305,
-      "grad_norm": 0.6953087449073792,
+      "grad_norm": 1.3712306022644043,
       "learning_rate": 3.795918367346939e-06,
       "loss": 0.0527,
       "max_memory_allocated (GB)": 57.18,
@@ -770,9 +770,9 @@
     },
     {
       "epoch": 3.142857142857143,
-      "grad_norm": 1.2290390729904175,
+      "grad_norm": 2.1615536212921143,
       "learning_rate": 3.7142857142857146e-06,
-      "loss": 0.0689,
+      "loss": 0.069,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 770,
@@ -780,9 +780,9 @@
     },
     {
       "epoch": 3.183673469387755,
-      "grad_norm": 0.6281890869140625,
+      "grad_norm": 0.45463302731513977,
       "learning_rate": 3.6326530612244903e-06,
-      "loss": 0.0647,
+      "loss": 0.064,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 780,
@@ -790,9 +790,9 @@
     },
     {
       "epoch": 3.2244897959183674,
-      "grad_norm": 0.3096281588077545,
+      "grad_norm": 0.40240758657455444,
       "learning_rate": 3.5510204081632655e-06,
-      "loss": 0.0522,
+      "loss": 0.0528,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 790,
@@ -800,9 +800,9 @@
     },
     {
       "epoch": 3.2653061224489797,
-      "grad_norm": 0.9390127062797546,
+      "grad_norm": 0.5408643484115601,
       "learning_rate": 3.469387755102041e-06,
-      "loss": 0.0432,
+      "loss": 0.0433,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 800,
@@ -810,9 +810,9 @@
     },
     {
       "epoch": 3.306122448979592,
-      "grad_norm": 0.87565016746521,
+      "grad_norm": 0.6606118679046631,
       "learning_rate": 3.3877551020408168e-06,
-      "loss": 0.0555,
+      "loss": 0.0548,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 810,
@@ -820,9 +820,9 @@
     },
     {
       "epoch": 3.3469387755102042,
-      "grad_norm": 1.0797837972640991,
+      "grad_norm": 0.89394211769104,
       "learning_rate": 3.3061224489795924e-06,
-      "loss": 0.0455,
+      "loss": 0.0448,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 820,
@@ -830,9 +830,9 @@
     },
     {
       "epoch": 3.387755102040816,
-      "grad_norm": 0.3658354878425598,
+      "grad_norm": 0.3227229416370392,
       "learning_rate": 3.2244897959183672e-06,
-      "loss": 0.0487,
+      "loss": 0.0491,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 830,
@@ -840,9 +840,9 @@
     },
     {
       "epoch": 3.4285714285714284,
-      "grad_norm": 0.4766336977481842,
+      "grad_norm": 0.928611159324646,
       "learning_rate": 3.142857142857143e-06,
-      "loss": 0.053,
+      "loss": 0.0523,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 840,
@@ -850,9 +850,9 @@
     },
     {
       "epoch": 3.4693877551020407,
-      "grad_norm": 0.49318933486938477,
+      "grad_norm": 0.6023752093315125,
       "learning_rate": 3.0612244897959185e-06,
-      "loss": 0.0812,
+      "loss": 0.0813,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 850,
@@ -860,7 +860,7 @@
     },
     {
       "epoch": 3.510204081632653,
-      "grad_norm": 1.3475311994552612,
+      "grad_norm": 1.1742885112762451,
       "learning_rate": 2.979591836734694e-06,
       "loss": 0.0451,
       "max_memory_allocated (GB)": 57.18,
@@ -870,9 +870,9 @@
     },
     {
       "epoch": 3.5510204081632653,
-      "grad_norm": 0.36763882637023926,
+      "grad_norm": 0.21970601379871368,
       "learning_rate": 2.8979591836734694e-06,
-      "loss": 0.0646,
+      "loss": 0.0644,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 870,
@@ -880,9 +880,9 @@
     },
     {
       "epoch": 3.5918367346938775,
-      "grad_norm": 3.085198402404785,
+      "grad_norm": 1.2270339727401733,
       "learning_rate": 2.816326530612245e-06,
-      "loss": 0.0439,
+      "loss": 0.044,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 880,
@@ -890,9 +890,9 @@
     },
     {
       "epoch": 3.63265306122449,
-      "grad_norm": 0.17229312658309937,
+      "grad_norm": 0.15171077847480774,
       "learning_rate": 2.7346938775510207e-06,
-      "loss": 0.0288,
+      "loss": 0.0291,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 890,
@@ -900,9 +900,9 @@
     },
     {
       "epoch": 3.673469387755102,
-      "grad_norm": 1.0760900974273682,
+      "grad_norm": 2.9102306365966797,
       "learning_rate": 2.6530612244897964e-06,
-      "loss": 0.0514,
+      "loss": 0.0516,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 900,
@@ -910,9 +910,9 @@
     },
     {
       "epoch": 3.7142857142857144,
-      "grad_norm": 0.45855164527893066,
+      "grad_norm": 0.6113564968109131,
       "learning_rate": 2.571428571428571e-06,
-      "loss": 0.0602,
+      "loss": 0.061,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 910,
@@ -920,9 +920,9 @@
     },
     {
       "epoch": 3.7551020408163263,
-      "grad_norm": 0.15575875341892242,
+      "grad_norm": 0.19307516515254974,
       "learning_rate": 2.489795918367347e-06,
-      "loss": 0.0543,
+      "loss": 0.0539,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 920,
@@ -930,9 +930,9 @@
     },
     {
       "epoch": 3.795918367346939,
-      "grad_norm": 0.779755175113678,
+      "grad_norm": 0.38657116889953613,
       "learning_rate": 2.4081632653061225e-06,
-      "loss": 0.0497,
+      "loss": 0.05,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 930,
@@ -940,9 +940,9 @@
     },
     {
       "epoch": 3.836734693877551,
-      "grad_norm": 0.7307060956954956,
+      "grad_norm": 0.9545940160751343,
       "learning_rate": 2.326530612244898e-06,
-      "loss": 0.0486,
+      "loss": 0.0481,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 940,
@@ -950,9 +950,9 @@
     },
     {
       "epoch": 3.877551020408163,
-      "grad_norm": 1.062565803527832,
+      "grad_norm": 1.3555389642715454,
       "learning_rate": 2.244897959183674e-06,
-      "loss": 0.0594,
+      "loss": 0.0596,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 950,
@@ -960,9 +960,9 @@
     },
     {
       "epoch": 3.9183673469387754,
-      "grad_norm": 0.3031039535999298,
+      "grad_norm": 0.3931931257247925,
       "learning_rate": 2.1632653061224495e-06,
-      "loss": 0.0497,
+      "loss": 0.0498,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 960,
@@ -970,9 +970,9 @@
     },
     {
       "epoch": 3.9591836734693877,
-      "grad_norm": 2.310593843460083,
+      "grad_norm": 0.6945517063140869,
       "learning_rate": 2.0816326530612247e-06,
-      "loss": 0.0746,
+      "loss": 0.0734,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 970,
@@ -980,265 +980,265 @@
     },
     {
       "epoch": 4.0,
-      "grad_norm": 0.6998704075813293,
+      "grad_norm": 1.0659183263778687,
       "learning_rate": 2.0000000000000003e-06,
-      "loss": 0.0703,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0707,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 980,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.040816326530612,
-      "grad_norm": 0.7492395639419556,
+      "grad_norm": 0.8017714023590088,
       "learning_rate": 1.9183673469387756e-06,
-      "loss": 0.0486,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0476,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 990,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.081632653061225,
-      "grad_norm": 0.7633445858955383,
+      "grad_norm": 1.1733373403549194,
       "learning_rate": 1.8367346938775512e-06,
-      "loss": 0.0625,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0617,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1000,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.122448979591836,
-      "grad_norm": 0.6911561489105225,
+      "grad_norm": 0.6048600077629089,
       "learning_rate": 1.7551020408163267e-06,
-      "loss": 0.0632,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0625,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1010,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.163265306122449,
-      "grad_norm": 0.33521902561187744,
+      "grad_norm": 1.096007227897644,
       "learning_rate": 1.6734693877551023e-06,
-      "loss": 0.0406,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0407,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1020,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.204081632653061,
-      "grad_norm": 0.7509037852287292,
+      "grad_norm": 4.377923965454102,
       "learning_rate": 1.5918367346938775e-06,
-      "loss": 0.0531,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0528,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1030,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.244897959183674,
-      "grad_norm": 0.5234070420265198,
+      "grad_norm": 0.7295175194740295,
       "learning_rate": 1.5102040816326532e-06,
-      "loss": 0.0396,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.04,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1040,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.285714285714286,
-      "grad_norm": 0.7997304797172546,
+      "grad_norm": 0.5129045248031616,
       "learning_rate": 1.4285714285714286e-06,
-      "loss": 0.05,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0493,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1050,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.326530612244898,
-      "grad_norm": 0.2255077213048935,
+      "grad_norm": 0.505799412727356,
       "learning_rate": 1.3469387755102043e-06,
-      "loss": 0.0457,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.046,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1060,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.36734693877551,
-      "grad_norm": 0.5182124376296997,
+      "grad_norm": 0.8460046052932739,
       "learning_rate": 1.2653061224489795e-06,
-      "loss": 0.0485,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0483,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1070,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.408163265306122,
-      "grad_norm": 0.35046374797821045,
+      "grad_norm": 0.5636487603187561,
       "learning_rate": 1.1836734693877552e-06,
-      "loss": 0.0519,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0523,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1080,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.448979591836735,
-      "grad_norm": 0.3923434615135193,
+      "grad_norm": 0.6893309354782104,
       "learning_rate": 1.1020408163265308e-06,
-      "loss": 0.0507,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0505,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1090,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.489795918367347,
-      "grad_norm": 0.23866137862205505,
+      "grad_norm": 0.6645925045013428,
       "learning_rate": 1.0204081632653063e-06,
-      "loss": 0.0362,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0361,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1100,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.530612244897959,
-      "grad_norm": 0.15117916464805603,
+      "grad_norm": 0.14709816873073578,
       "learning_rate": 9.387755102040817e-07,
-      "loss": 0.0464,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0458,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1110,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.571428571428571,
-      "grad_norm": 0.5993088483810425,
+      "grad_norm": 0.8794461488723755,
       "learning_rate": 8.571428571428572e-07,
-      "loss": 0.0404,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0414,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1120,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.612244897959184,
-      "grad_norm": 0.30265432596206665,
+      "grad_norm": 1.345910906791687,
       "learning_rate": 7.755102040816327e-07,
-      "loss": 0.0545,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0526,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1130,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.653061224489796,
-      "grad_norm": 0.6385183334350586,
+      "grad_norm": 0.5657418370246887,
       "learning_rate": 6.938775510204082e-07,
-      "loss": 0.0731,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0721,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1140,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.6938775510204085,
-      "grad_norm": 1.128566026687622,
+      "grad_norm": 2.0943238735198975,
       "learning_rate": 6.122448979591837e-07,
-      "loss": 0.0516,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0515,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1150,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.73469387755102,
-      "grad_norm": 1.1660116910934448,
+      "grad_norm": 3.1750009059906006,
       "learning_rate": 5.306122448979592e-07,
-      "loss": 0.0611,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0607,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1160,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.775510204081632,
-      "grad_norm": 0.5327439904212952,
+      "grad_norm": 0.8910292387008667,
       "learning_rate": 4.489795918367347e-07,
-      "loss": 0.0549,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.054,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1170,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.816326530612245,
-      "grad_norm": 0.8764423131942749,
+      "grad_norm": 0.9796111583709717,
       "learning_rate": 3.6734693877551025e-07,
-      "loss": 0.0441,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0432,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1180,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.857142857142857,
-      "grad_norm": 0.47835007309913635,
+      "grad_norm": 1.7865172624588013,
       "learning_rate": 2.8571428571428575e-07,
-      "loss": 0.0541,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0549,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1190,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.8979591836734695,
-      "grad_norm": 1.048047661781311,
+      "grad_norm": 0.7347144484519958,
       "learning_rate": 2.0408163265306121e-07,
-      "loss": 0.0731,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0711,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1200,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.938775510204081,
-      "grad_norm": 0.3101171851158142,
+      "grad_norm": 1.210486888885498,
       "learning_rate": 1.2244897959183673e-07,
-      "loss": 0.0648,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0652,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1210,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.979591836734694,
-      "grad_norm": 0.76802659034729,
+      "grad_norm": 1.2078640460968018,
       "learning_rate": 4.0816326530612253e-08,
-      "loss": 0.0418,
-      "max_memory_allocated (GB)": 60.52,
+      "loss": 0.0414,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1220,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 5.0,
-      "max_memory_allocated (GB)": 60.52,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1225,
       "total_flos": 3.0598946525952e+16,
       "total_memory_available (GB)": 94.62,
-      "train_loss": 0.06085505417415074,
-      "train_runtime": 1020.8061,
-      "train_samples_per_second": 55.51,
-      "train_steps_per_second": 1.389
+      "train_loss": 0.06072675045655698,
+      "train_runtime": 1077.821,
+      "train_samples_per_second": 52.682,
+      "train_steps_per_second": 1.318
     }
   ],
   "logging_steps": 10,

validation_results.json

@@ -1,10 +1,10 @@
 {
     "epoch": 5.0,
-    "eval_loss": 0.1272137612104416,
-    "eval_runtime": 30.5176,
-    "eval_samples_per_second": 31.06,
-    "eval_steps_per_second": 1.994,
-    "max_memory_allocated (GB)": 60.52,
+    "eval_loss": 0.12326910346746445,
+    "eval_runtime": 26.1444,
+    "eval_samples_per_second": 36.653,
+    "eval_steps_per_second": 2.352,
+    "max_memory_allocated (GB)": 57.18,
     "memory_allocated (GB)": 51.27,
     "total_memory_available (GB)": 94.62
 }