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.1284
+- 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.13596950471401215,
-    "eval_runtime": 5.0456,
-    "eval_samples_per_second": 129.562,
-    "eval_steps_per_second": 8.098,
+    "eval_loss": 0.12840959429740906,
+    "eval_runtime": 7.7898,
+    "eval_samples_per_second": 66.741,
+    "eval_steps_per_second": 4.171,
     "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.06072675045655698,
-    "train_runtime": 1077.821,
-    "train_samples_per_second": 52.682,
-    "train_steps_per_second": 1.318
+    "train_loss": 0.06098026679486644,
+    "train_runtime": 1192.2443,
+    "train_samples_per_second": 46.607,
+    "train_steps_per_second": 1.166
 }

runs/Oct16_16-24-55_workload-ai-workshop/events.out.tfevents.1729097262.workload-ai-workshop.7719.1

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

test_results.json

@@ -1,9 +1,9 @@
 {
     "epoch": 5.0,
-    "eval_loss": 0.13596950471401215,
-    "eval_runtime": 5.0456,
-    "eval_samples_per_second": 129.562,
-    "eval_steps_per_second": 8.098,
+    "eval_loss": 0.12840959429740906,
+    "eval_runtime": 7.7898,
+    "eval_samples_per_second": 66.741,
+    "eval_steps_per_second": 4.171,
     "max_memory_allocated (GB)": 57.18,
     "memory_allocated (GB)": 51.27,
     "total_memory_available (GB)": 94.62

train_results.json

@@ -4,8 +4,8 @@
     "memory_allocated (GB)": 50.57,
     "total_flos": 3.0598946525952e+16,
     "total_memory_available (GB)": 94.62,
-    "train_loss": 0.06072675045655698,
-    "train_runtime": 1077.821,
-    "train_samples_per_second": 52.682,
-    "train_steps_per_second": 1.318
+    "train_loss": 0.06098026679486644,
+    "train_runtime": 1192.2443,
+    "train_samples_per_second": 46.607,
+    "train_steps_per_second": 1.166
 }

trainer_state.json

@@ -20,9 +20,9 @@
     },
     {
       "epoch": 0.08163265306122448,
-      "grad_norm": 18.564252853393555,
+      "grad_norm": 7.052234649658203,
       "learning_rate": 9.836734693877552e-06,
-      "loss": 0.1554,
+      "loss": 0.1555,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 20,
@@ -30,9 +30,9 @@
     },
     {
       "epoch": 0.12244897959183673,
-      "grad_norm": 7.802953720092773,
+      "grad_norm": 6.5298075675964355,
       "learning_rate": 9.755102040816327e-06,
-      "loss": 0.126,
+      "loss": 0.1251,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 30,
@@ -40,9 +40,9 @@
     },
     {
       "epoch": 0.16326530612244897,
-      "grad_norm": 4.599184036254883,
+      "grad_norm": 4.405805587768555,
       "learning_rate": 9.673469387755103e-06,
-      "loss": 0.1103,
+      "loss": 0.1102,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 40,
@@ -50,9 +50,9 @@
     },
     {
       "epoch": 0.20408163265306123,
-      "grad_norm": 6.734092712402344,
+      "grad_norm": 4.870044708251953,
       "learning_rate": 9.591836734693878e-06,
-      "loss": 0.1217,
+      "loss": 0.1232,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 50,
@@ -60,9 +60,9 @@
     },
     {
       "epoch": 0.24489795918367346,
-      "grad_norm": 1.7976917028427124,
+      "grad_norm": 1.6433866024017334,
       "learning_rate": 9.510204081632653e-06,
-      "loss": 0.08,
+      "loss": 0.0797,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 60,
@@ -70,9 +70,9 @@
     },
     {
       "epoch": 0.2857142857142857,
-      "grad_norm": 4.0598835945129395,
+      "grad_norm": 4.2432074546813965,
       "learning_rate": 9.42857142857143e-06,
-      "loss": 0.1023,
+      "loss": 0.1031,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 70,
@@ -80,9 +80,9 @@
     },
     {
       "epoch": 0.32653061224489793,
-      "grad_norm": 2.0350396633148193,
+      "grad_norm": 2.0352487564086914,
       "learning_rate": 9.346938775510204e-06,
-      "loss": 0.1128,
+      "loss": 0.1115,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 80,
@@ -90,9 +90,9 @@
     },
     {
       "epoch": 0.3673469387755102,
-      "grad_norm": 3.785065174102783,
+      "grad_norm": 4.201560020446777,
       "learning_rate": 9.26530612244898e-06,
-      "loss": 0.0842,
+      "loss": 0.0817,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 90,
@@ -100,9 +100,9 @@
     },
     {
       "epoch": 0.40816326530612246,
-      "grad_norm": 1.9926950931549072,
+      "grad_norm": 3.618368625640869,
       "learning_rate": 9.183673469387756e-06,
-      "loss": 0.0728,
+      "loss": 0.0768,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 100,
@@ -110,9 +110,9 @@
     },
     {
       "epoch": 0.4489795918367347,
-      "grad_norm": 4.878537178039551,
+      "grad_norm": 4.793916702270508,
       "learning_rate": 9.102040816326532e-06,
-      "loss": 0.0692,
+      "loss": 0.071,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 110,
@@ -120,9 +120,9 @@
     },
     {
       "epoch": 0.4897959183673469,
-      "grad_norm": 2.5495357513427734,
+      "grad_norm": 2.3223495483398438,
       "learning_rate": 9.020408163265307e-06,
-      "loss": 0.0691,
+      "loss": 0.0707,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 120,
@@ -130,9 +130,9 @@
     },
     {
       "epoch": 0.5306122448979592,
-      "grad_norm": 3.0813372135162354,
+      "grad_norm": 3.5389153957366943,
       "learning_rate": 8.938775510204082e-06,
-      "loss": 0.0597,
+      "loss": 0.0599,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 130,
@@ -140,9 +140,9 @@
     },
     {
       "epoch": 0.5714285714285714,
-      "grad_norm": 1.2760021686553955,
+      "grad_norm": 1.586653232574463,
       "learning_rate": 8.857142857142858e-06,
-      "loss": 0.0507,
+      "loss": 0.0491,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 140,
@@ -150,9 +150,9 @@
     },
     {
       "epoch": 0.6122448979591837,
-      "grad_norm": 2.0892932415008545,
+      "grad_norm": 1.5236841440200806,
       "learning_rate": 8.775510204081633e-06,
-      "loss": 0.0634,
+      "loss": 0.0632,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 150,
@@ -160,9 +160,9 @@
     },
     {
       "epoch": 0.6530612244897959,
-      "grad_norm": 2.0191688537597656,
+      "grad_norm": 2.752020835876465,
       "learning_rate": 8.69387755102041e-06,
-      "loss": 0.0718,
+      "loss": 0.0722,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 160,
@@ -170,9 +170,9 @@
     },
     {
       "epoch": 0.6938775510204082,
-      "grad_norm": 2.4970450401306152,
+      "grad_norm": 7.606927394866943,
       "learning_rate": 8.612244897959184e-06,
-      "loss": 0.0753,
+      "loss": 0.0756,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 170,
@@ -180,9 +180,9 @@
     },
     {
       "epoch": 0.7346938775510204,
-      "grad_norm": 1.5428143739700317,
+      "grad_norm": 1.5622702836990356,
       "learning_rate": 8.530612244897961e-06,
-      "loss": 0.0612,
+      "loss": 0.0617,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 180,
@@ -190,9 +190,9 @@
     },
     {
       "epoch": 0.7755102040816326,
-      "grad_norm": 3.6041452884674072,
+      "grad_norm": 0.9614956378936768,
       "learning_rate": 8.448979591836736e-06,
-      "loss": 0.0573,
+      "loss": 0.0572,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 190,
@@ -200,9 +200,9 @@
     },
     {
       "epoch": 0.8163265306122449,
-      "grad_norm": 1.0115529298782349,
+      "grad_norm": 0.7814755439758301,
       "learning_rate": 8.36734693877551e-06,
-      "loss": 0.0631,
+      "loss": 0.0636,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 200,
@@ -210,9 +210,9 @@
     },
     {
       "epoch": 0.8571428571428571,
-      "grad_norm": 0.8029147982597351,
+      "grad_norm": 1.352851390838623,
       "learning_rate": 8.285714285714287e-06,
-      "loss": 0.0643,
+      "loss": 0.0648,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 210,
@@ -220,9 +220,9 @@
     },
     {
       "epoch": 0.8979591836734694,
-      "grad_norm": 1.130996584892273,
+      "grad_norm": 1.6814969778060913,
       "learning_rate": 8.204081632653062e-06,
-      "loss": 0.0608,
+      "loss": 0.0604,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 220,
@@ -230,9 +230,9 @@
     },
     {
       "epoch": 0.9387755102040817,
-      "grad_norm": 0.7962849140167236,
+      "grad_norm": 0.859993040561676,
       "learning_rate": 8.122448979591837e-06,
-      "loss": 0.0535,
+      "loss": 0.0549,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 230,
@@ -240,9 +240,9 @@
     },
     {
       "epoch": 0.9795918367346939,
-      "grad_norm": 2.3987386226654053,
+      "grad_norm": 0.6439819931983948,
       "learning_rate": 8.040816326530613e-06,
-      "loss": 0.0499,
+      "loss": 0.0493,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 240,
@@ -250,9 +250,9 @@
     },
     {
       "epoch": 1.0204081632653061,
-      "grad_norm": 1.084067463874817,
+      "grad_norm": 0.8465150594711304,
       "learning_rate": 7.959183673469388e-06,
-      "loss": 0.0612,
+      "loss": 0.0624,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 250,
@@ -260,9 +260,9 @@
     },
     {
       "epoch": 1.0612244897959184,
-      "grad_norm": 3.247530460357666,
+      "grad_norm": 1.0257333517074585,
       "learning_rate": 7.877551020408164e-06,
-      "loss": 0.055,
+      "loss": 0.056,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 260,
@@ -270,9 +270,9 @@
     },
     {
       "epoch": 1.1020408163265305,
-      "grad_norm": 1.5961194038391113,
+      "grad_norm": 2.619938850402832,
       "learning_rate": 7.79591836734694e-06,
-      "loss": 0.0632,
+      "loss": 0.0648,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 270,
@@ -280,9 +280,9 @@
     },
     {
       "epoch": 1.1428571428571428,
-      "grad_norm": 0.32916340231895447,
+      "grad_norm": 0.4946042001247406,
       "learning_rate": 7.714285714285716e-06,
-      "loss": 0.0565,
+      "loss": 0.0586,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 280,
@@ -290,9 +290,9 @@
     },
     {
       "epoch": 1.183673469387755,
-      "grad_norm": 0.6009345054626465,
+      "grad_norm": 1.0154733657836914,
       "learning_rate": 7.63265306122449e-06,
-      "loss": 0.0503,
+      "loss": 0.0505,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 290,
@@ -300,9 +300,9 @@
     },
     {
       "epoch": 1.2244897959183674,
-      "grad_norm": 9.806236267089844,
+      "grad_norm": 1.0347952842712402,
       "learning_rate": 7.551020408163265e-06,
-      "loss": 0.0635,
+      "loss": 0.0646,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 300,
@@ -310,9 +310,9 @@
     },
     {
       "epoch": 1.2653061224489797,
-      "grad_norm": 1.231447696685791,
+      "grad_norm": 0.7844366431236267,
       "learning_rate": 7.469387755102041e-06,
-      "loss": 0.0664,
+      "loss": 0.0676,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 310,
@@ -320,9 +320,9 @@
     },
     {
       "epoch": 1.306122448979592,
-      "grad_norm": 1.016727328300476,
+      "grad_norm": 1.1971337795257568,
       "learning_rate": 7.387755102040817e-06,
-      "loss": 0.0502,
+      "loss": 0.0499,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 320,
@@ -330,9 +330,9 @@
     },
     {
       "epoch": 1.346938775510204,
-      "grad_norm": 0.7703081965446472,
+      "grad_norm": 0.6674404740333557,
       "learning_rate": 7.306122448979592e-06,
-      "loss": 0.0599,
+      "loss": 0.0602,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 330,
@@ -340,9 +340,9 @@
     },
     {
       "epoch": 1.3877551020408163,
-      "grad_norm": 5.310096740722656,
+      "grad_norm": 1.511208415031433,
       "learning_rate": 7.224489795918368e-06,
-      "loss": 0.0541,
+      "loss": 0.0547,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 340,
@@ -350,9 +350,9 @@
     },
     {
       "epoch": 1.4285714285714286,
-      "grad_norm": 1.1444560289382935,
+      "grad_norm": 0.5328841209411621,
       "learning_rate": 7.1428571428571436e-06,
-      "loss": 0.0493,
+      "loss": 0.0486,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 350,
@@ -360,9 +360,9 @@
     },
     {
       "epoch": 1.469387755102041,
-      "grad_norm": 3.1823084354400635,
+      "grad_norm": 1.464439034461975,
       "learning_rate": 7.061224489795919e-06,
-      "loss": 0.0459,
+      "loss": 0.0464,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 360,
@@ -370,9 +370,9 @@
     },
     {
       "epoch": 1.510204081632653,
-      "grad_norm": 1.1924108266830444,
+      "grad_norm": 0.834863543510437,
       "learning_rate": 6.979591836734695e-06,
-      "loss": 0.058,
+      "loss": 0.0591,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 370,
@@ -380,9 +380,9 @@
     },
     {
       "epoch": 1.5510204081632653,
-      "grad_norm": 3.677968740463257,
+      "grad_norm": 0.5399609208106995,
       "learning_rate": 6.8979591836734705e-06,
-      "loss": 0.0468,
+      "loss": 0.0464,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 380,
@@ -390,9 +390,9 @@
     },
     {
       "epoch": 1.5918367346938775,
-      "grad_norm": 1.0082957744598389,
+      "grad_norm": 0.8577661514282227,
       "learning_rate": 6.816326530612245e-06,
-      "loss": 0.0652,
+      "loss": 0.0654,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 390,
@@ -400,9 +400,9 @@
     },
     {
       "epoch": 1.6326530612244898,
-      "grad_norm": 0.6584922075271606,
+      "grad_norm": 0.5057955384254456,
       "learning_rate": 6.734693877551021e-06,
-      "loss": 0.0605,
+      "loss": 0.0609,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 400,
@@ -410,9 +410,9 @@
     },
     {
       "epoch": 1.6734693877551021,
-      "grad_norm": 1.2264763116836548,
+      "grad_norm": 0.9135333895683289,
       "learning_rate": 6.653061224489797e-06,
-      "loss": 0.0599,
+      "loss": 0.0607,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 410,
@@ -420,9 +420,9 @@
     },
     {
       "epoch": 1.7142857142857144,
-      "grad_norm": 1.824320912361145,
+      "grad_norm": 2.9697179794311523,
       "learning_rate": 6.571428571428572e-06,
-      "loss": 0.053,
+      "loss": 0.054,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 420,
@@ -430,9 +430,9 @@
     },
     {
       "epoch": 1.7551020408163265,
-      "grad_norm": 0.24405838549137115,
+      "grad_norm": 0.3473312556743622,
       "learning_rate": 6.489795918367348e-06,
-      "loss": 0.067,
+      "loss": 0.0685,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 430,
@@ -440,9 +440,9 @@
     },
     {
       "epoch": 1.7959183673469388,
-      "grad_norm": 1.0410544872283936,
+      "grad_norm": 1.4528335332870483,
       "learning_rate": 6.408163265306124e-06,
-      "loss": 0.06,
+      "loss": 0.0611,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 440,
@@ -450,9 +450,9 @@
     },
     {
       "epoch": 1.836734693877551,
-      "grad_norm": 0.39765527844429016,
+      "grad_norm": 0.48578280210494995,
       "learning_rate": 6.326530612244899e-06,
-      "loss": 0.043,
+      "loss": 0.0438,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 450,
@@ -460,9 +460,9 @@
     },
     {
       "epoch": 1.8775510204081631,
-      "grad_norm": 0.29981690645217896,
+      "grad_norm": 0.3472760021686554,
       "learning_rate": 6.244897959183675e-06,
-      "loss": 0.0533,
+      "loss": 0.0544,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 460,
@@ -470,9 +470,9 @@
     },
     {
       "epoch": 1.9183673469387754,
-      "grad_norm": 0.6634105443954468,
+      "grad_norm": 1.0984327793121338,
       "learning_rate": 6.163265306122449e-06,
-      "loss": 0.0428,
+      "loss": 0.0438,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 470,
@@ -480,9 +480,9 @@
     },
     {
       "epoch": 1.9591836734693877,
-      "grad_norm": 0.2060549110174179,
+      "grad_norm": 0.20147933065891266,
       "learning_rate": 6.0816326530612245e-06,
-      "loss": 0.0505,
+      "loss": 0.0518,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 480,
@@ -490,9 +490,9 @@
     },
     {
       "epoch": 2.0,
-      "grad_norm": 1.0800402164459229,
+      "grad_norm": 1.1583309173583984,
       "learning_rate": 6e-06,
-      "loss": 0.0647,
+      "loss": 0.0637,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 490,
@@ -500,7 +500,7 @@
     },
     {
       "epoch": 2.0408163265306123,
-      "grad_norm": 0.6078592538833618,
+      "grad_norm": 0.6601622104644775,
       "learning_rate": 5.918367346938776e-06,
       "loss": 0.0596,
       "max_memory_allocated (GB)": 57.18,
@@ -510,9 +510,9 @@
     },
     {
       "epoch": 2.0816326530612246,
-      "grad_norm": 0.6262472867965698,
+      "grad_norm": 0.5227305293083191,
       "learning_rate": 5.8367346938775515e-06,
-      "loss": 0.0495,
+      "loss": 0.0493,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 510,
@@ -520,7 +520,7 @@
     },
     {
       "epoch": 2.122448979591837,
-      "grad_norm": 1.4308090209960938,
+      "grad_norm": 0.8996191620826721,
       "learning_rate": 5.755102040816327e-06,
       "loss": 0.0461,
       "max_memory_allocated (GB)": 57.18,
@@ -530,9 +530,9 @@
     },
     {
       "epoch": 2.163265306122449,
-      "grad_norm": 0.8761769533157349,
+      "grad_norm": 1.0684189796447754,
       "learning_rate": 5.673469387755103e-06,
-      "loss": 0.0626,
+      "loss": 0.0629,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 530,
@@ -540,9 +540,9 @@
     },
     {
       "epoch": 2.204081632653061,
-      "grad_norm": 0.6150110363960266,
+      "grad_norm": 0.5558530688285828,
       "learning_rate": 5.591836734693878e-06,
-      "loss": 0.0576,
+      "loss": 0.0581,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 540,
@@ -550,9 +550,9 @@
     },
     {
       "epoch": 2.2448979591836733,
-      "grad_norm": 9.280069351196289,
+      "grad_norm": 1.1996757984161377,
       "learning_rate": 5.510204081632653e-06,
-      "loss": 0.0634,
+      "loss": 0.0626,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 550,
@@ -560,9 +560,9 @@
     },
     {
       "epoch": 2.2857142857142856,
-      "grad_norm": 2.136209726333618,
+      "grad_norm": 1.2928632497787476,
       "learning_rate": 5.428571428571429e-06,
-      "loss": 0.0579,
+      "loss": 0.0575,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 560,
@@ -570,9 +570,9 @@
     },
     {
       "epoch": 2.326530612244898,
-      "grad_norm": 0.9547491073608398,
+      "grad_norm": 0.7934871912002563,
       "learning_rate": 5.3469387755102045e-06,
-      "loss": 0.0589,
+      "loss": 0.0577,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 570,
@@ -580,9 +580,9 @@
     },
     {
       "epoch": 2.36734693877551,
-      "grad_norm": 2.1748363971710205,
+      "grad_norm": 3.946485757827759,
       "learning_rate": 5.26530612244898e-06,
-      "loss": 0.0672,
+      "loss": 0.0663,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 580,
@@ -590,9 +590,9 @@
     },
     {
       "epoch": 2.4081632653061225,
-      "grad_norm": 0.9653811454772949,
+      "grad_norm": 0.43567588925361633,
       "learning_rate": 5.183673469387756e-06,
-      "loss": 0.0536,
+      "loss": 0.0539,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 590,
@@ -600,9 +600,9 @@
     },
     {
       "epoch": 2.4489795918367347,
-      "grad_norm": 0.7852123379707336,
+      "grad_norm": 0.5725533962249756,
       "learning_rate": 5.1020408163265315e-06,
-      "loss": 0.0443,
+      "loss": 0.0438,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 600,
@@ -610,7 +610,7 @@
     },
     {
       "epoch": 2.489795918367347,
-      "grad_norm": 0.7405697107315063,
+      "grad_norm": 0.44328320026397705,
       "learning_rate": 5.020408163265307e-06,
       "loss": 0.041,
       "max_memory_allocated (GB)": 57.18,
@@ -620,9 +620,9 @@
     },
     {
       "epoch": 2.5306122448979593,
-      "grad_norm": 0.7166327834129333,
+      "grad_norm": 1.338100790977478,
       "learning_rate": 4.938775510204082e-06,
-      "loss": 0.0431,
+      "loss": 0.0424,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 620,
@@ -630,9 +630,9 @@
     },
     {
       "epoch": 2.571428571428571,
-      "grad_norm": 0.5675875544548035,
+      "grad_norm": 0.92643803358078,
       "learning_rate": 4.857142857142858e-06,
-      "loss": 0.0807,
+      "loss": 0.0811,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 630,
@@ -640,7 +640,7 @@
     },
     {
       "epoch": 2.612244897959184,
-      "grad_norm": 0.954290509223938,
+      "grad_norm": 1.1147398948669434,
       "learning_rate": 4.775510204081633e-06,
       "loss": 0.0492,
       "max_memory_allocated (GB)": 57.18,
@@ -650,9 +650,9 @@
     },
     {
       "epoch": 2.6530612244897958,
-      "grad_norm": 0.9061315655708313,
+      "grad_norm": 0.6104307174682617,
       "learning_rate": 4.693877551020409e-06,
-      "loss": 0.0472,
+      "loss": 0.0468,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 650,
@@ -660,9 +660,9 @@
     },
     {
       "epoch": 2.693877551020408,
-      "grad_norm": 0.7333698868751526,
+      "grad_norm": 0.9826134443283081,
       "learning_rate": 4.612244897959184e-06,
-      "loss": 0.0474,
+      "loss": 0.0471,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 660,
@@ -670,9 +670,9 @@
     },
     {
       "epoch": 2.7346938775510203,
-      "grad_norm": 1.7820810079574585,
+      "grad_norm": 0.7680672407150269,
       "learning_rate": 4.530612244897959e-06,
-      "loss": 0.0869,
+      "loss": 0.0858,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 670,
@@ -680,9 +680,9 @@
     },
     {
       "epoch": 2.7755102040816326,
-      "grad_norm": 1.3184058666229248,
+      "grad_norm": 0.9682340025901794,
       "learning_rate": 4.448979591836735e-06,
-      "loss": 0.0494,
+      "loss": 0.0484,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 680,
@@ -690,9 +690,9 @@
     },
     {
       "epoch": 2.816326530612245,
-      "grad_norm": 0.4212433397769928,
+      "grad_norm": 0.37712323665618896,
       "learning_rate": 4.367346938775511e-06,
-      "loss": 0.0445,
+      "loss": 0.0443,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 690,
@@ -700,9 +700,9 @@
     },
     {
       "epoch": 2.857142857142857,
-      "grad_norm": 0.39547938108444214,
+      "grad_norm": 0.34970754384994507,
       "learning_rate": 4.2857142857142855e-06,
-      "loss": 0.0429,
+      "loss": 0.0434,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 700,
@@ -710,9 +710,9 @@
     },
     {
       "epoch": 2.8979591836734695,
-      "grad_norm": 1.0822113752365112,
+      "grad_norm": 0.9949877262115479,
       "learning_rate": 4.204081632653061e-06,
-      "loss": 0.0549,
+      "loss": 0.0553,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 710,
@@ -720,9 +720,9 @@
     },
     {
       "epoch": 2.938775510204082,
-      "grad_norm": 1.0838605165481567,
+      "grad_norm": 1.4436949491500854,
       "learning_rate": 4.122448979591837e-06,
-      "loss": 0.0582,
+      "loss": 0.0583,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 720,
@@ -730,9 +730,9 @@
     },
     {
       "epoch": 2.979591836734694,
-      "grad_norm": 0.2726249098777771,
+      "grad_norm": 0.1619979739189148,
       "learning_rate": 4.040816326530612e-06,
-      "loss": 0.0341,
+      "loss": 0.0336,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 730,
@@ -740,9 +740,9 @@
     },
     {
       "epoch": 3.020408163265306,
-      "grad_norm": 0.836703360080719,
+      "grad_norm": 1.2799049615859985,
       "learning_rate": 3.959183673469388e-06,
-      "loss": 0.053,
+      "loss": 0.0536,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 740,
@@ -750,9 +750,9 @@
     },
     {
       "epoch": 3.061224489795918,
-      "grad_norm": 0.6878814697265625,
+      "grad_norm": 0.5613189935684204,
       "learning_rate": 3.877551020408164e-06,
-      "loss": 0.0618,
+      "loss": 0.062,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 750,
@@ -760,7 +760,7 @@
     },
     {
       "epoch": 3.1020408163265305,
-      "grad_norm": 1.3712306022644043,
+      "grad_norm": 0.827383279800415,
       "learning_rate": 3.795918367346939e-06,
       "loss": 0.0527,
       "max_memory_allocated (GB)": 57.18,
@@ -770,9 +770,9 @@
     },
     {
       "epoch": 3.142857142857143,
-      "grad_norm": 2.1615536212921143,
+      "grad_norm": 0.6983201503753662,
       "learning_rate": 3.7142857142857146e-06,
-      "loss": 0.069,
+      "loss": 0.0691,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 770,
@@ -780,9 +780,9 @@
     },
     {
       "epoch": 3.183673469387755,
-      "grad_norm": 0.45463302731513977,
+      "grad_norm": 1.0466923713684082,
       "learning_rate": 3.6326530612244903e-06,
-      "loss": 0.064,
+      "loss": 0.0644,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 780,
@@ -790,9 +790,9 @@
     },
     {
       "epoch": 3.2244897959183674,
-      "grad_norm": 0.40240758657455444,
+      "grad_norm": 0.3068871796131134,
       "learning_rate": 3.5510204081632655e-06,
-      "loss": 0.0528,
+      "loss": 0.0524,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 790,
@@ -800,9 +800,9 @@
     },
     {
       "epoch": 3.2653061224489797,
-      "grad_norm": 0.5408643484115601,
+      "grad_norm": 0.40160393714904785,
       "learning_rate": 3.469387755102041e-06,
-      "loss": 0.0433,
+      "loss": 0.0434,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 800,
@@ -810,9 +810,9 @@
     },
     {
       "epoch": 3.306122448979592,
-      "grad_norm": 0.6606118679046631,
+      "grad_norm": 0.880214512348175,
       "learning_rate": 3.3877551020408168e-06,
-      "loss": 0.0548,
+      "loss": 0.056,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 810,
@@ -820,9 +820,9 @@
     },
     {
       "epoch": 3.3469387755102042,
-      "grad_norm": 0.89394211769104,
+      "grad_norm": 0.9539953470230103,
       "learning_rate": 3.3061224489795924e-06,
-      "loss": 0.0448,
+      "loss": 0.0464,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 820,
@@ -830,9 +830,9 @@
     },
     {
       "epoch": 3.387755102040816,
-      "grad_norm": 0.3227229416370392,
+      "grad_norm": 0.24522298574447632,
       "learning_rate": 3.2244897959183672e-06,
-      "loss": 0.0491,
+      "loss": 0.0485,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 830,
@@ -840,9 +840,9 @@
     },
     {
       "epoch": 3.4285714285714284,
-      "grad_norm": 0.928611159324646,
+      "grad_norm": 0.4946345388889313,
       "learning_rate": 3.142857142857143e-06,
-      "loss": 0.0523,
+      "loss": 0.0527,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 840,
@@ -850,7 +850,7 @@
     },
     {
       "epoch": 3.4693877551020407,
-      "grad_norm": 0.6023752093315125,
+      "grad_norm": 0.4724675416946411,
       "learning_rate": 3.0612244897959185e-06,
       "loss": 0.0813,
       "max_memory_allocated (GB)": 57.18,
@@ -860,9 +860,9 @@
     },
     {
       "epoch": 3.510204081632653,
-      "grad_norm": 1.1742885112762451,
+      "grad_norm": 0.9907402396202087,
       "learning_rate": 2.979591836734694e-06,
-      "loss": 0.0451,
+      "loss": 0.0447,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 860,
@@ -870,9 +870,9 @@
     },
     {
       "epoch": 3.5510204081632653,
-      "grad_norm": 0.21970601379871368,
+      "grad_norm": 0.19696560502052307,
       "learning_rate": 2.8979591836734694e-06,
-      "loss": 0.0644,
+      "loss": 0.0635,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 870,
@@ -880,9 +880,9 @@
     },
     {
       "epoch": 3.5918367346938775,
-      "grad_norm": 1.2270339727401733,
+      "grad_norm": 0.7972800135612488,
       "learning_rate": 2.816326530612245e-06,
-      "loss": 0.044,
+      "loss": 0.0438,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 880,
@@ -890,9 +890,9 @@
     },
     {
       "epoch": 3.63265306122449,
-      "grad_norm": 0.15171077847480774,
+      "grad_norm": 0.21193134784698486,
       "learning_rate": 2.7346938775510207e-06,
-      "loss": 0.0291,
+      "loss": 0.029,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 890,
@@ -900,9 +900,9 @@
     },
     {
       "epoch": 3.673469387755102,
-      "grad_norm": 2.9102306365966797,
+      "grad_norm": 0.6128103137016296,
       "learning_rate": 2.6530612244897964e-06,
-      "loss": 0.0516,
+      "loss": 0.0514,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 900,
@@ -910,7 +910,7 @@
     },
     {
       "epoch": 3.7142857142857144,
-      "grad_norm": 0.6113564968109131,
+      "grad_norm": 0.8112168312072754,
       "learning_rate": 2.571428571428571e-06,
       "loss": 0.061,
       "max_memory_allocated (GB)": 57.18,
@@ -920,9 +920,9 @@
     },
     {
       "epoch": 3.7551020408163263,
-      "grad_norm": 0.19307516515254974,
+      "grad_norm": 0.18730562925338745,
       "learning_rate": 2.489795918367347e-06,
-      "loss": 0.0539,
+      "loss": 0.0546,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 920,
@@ -930,9 +930,9 @@
     },
     {
       "epoch": 3.795918367346939,
-      "grad_norm": 0.38657116889953613,
+      "grad_norm": 0.3866801857948303,
       "learning_rate": 2.4081632653061225e-06,
-      "loss": 0.05,
+      "loss": 0.0501,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 930,
@@ -940,9 +940,9 @@
     },
     {
       "epoch": 3.836734693877551,
-      "grad_norm": 0.9545940160751343,
+      "grad_norm": 0.8816384077072144,
       "learning_rate": 2.326530612244898e-06,
-      "loss": 0.0481,
+      "loss": 0.0489,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 940,
@@ -950,9 +950,9 @@
     },
     {
       "epoch": 3.877551020408163,
-      "grad_norm": 1.3555389642715454,
+      "grad_norm": 0.5572797656059265,
       "learning_rate": 2.244897959183674e-06,
-      "loss": 0.0596,
+      "loss": 0.0599,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 950,
@@ -960,9 +960,9 @@
     },
     {
       "epoch": 3.9183673469387754,
-      "grad_norm": 0.3931931257247925,
+      "grad_norm": 0.38238489627838135,
       "learning_rate": 2.1632653061224495e-06,
-      "loss": 0.0498,
+      "loss": 0.0497,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 960,
@@ -970,9 +970,9 @@
     },
     {
       "epoch": 3.9591836734693877,
-      "grad_norm": 0.6945517063140869,
+      "grad_norm": 0.6144959926605225,
       "learning_rate": 2.0816326530612247e-06,
-      "loss": 0.0734,
+      "loss": 0.0741,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 970,
@@ -980,9 +980,9 @@
     },
     {
       "epoch": 4.0,
-      "grad_norm": 1.0659183263778687,
+      "grad_norm": 0.6087101697921753,
       "learning_rate": 2.0000000000000003e-06,
-      "loss": 0.0707,
+      "loss": 0.0703,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 980,
@@ -990,9 +990,9 @@
     },
     {
       "epoch": 4.040816326530612,
-      "grad_norm": 0.8017714023590088,
+      "grad_norm": 0.5187469720840454,
       "learning_rate": 1.9183673469387756e-06,
-      "loss": 0.0476,
+      "loss": 0.0482,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 990,
@@ -1000,9 +1000,9 @@
     },
     {
       "epoch": 4.081632653061225,
-      "grad_norm": 1.1733373403549194,
+      "grad_norm": 1.248850703239441,
       "learning_rate": 1.8367346938775512e-06,
-      "loss": 0.0617,
+      "loss": 0.0631,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1000,
@@ -1010,9 +1010,9 @@
     },
     {
       "epoch": 4.122448979591836,
-      "grad_norm": 0.6048600077629089,
+      "grad_norm": 0.5806276798248291,
       "learning_rate": 1.7551020408163267e-06,
-      "loss": 0.0625,
+      "loss": 0.0629,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1010,
@@ -1020,7 +1020,7 @@
     },
     {
       "epoch": 4.163265306122449,
-      "grad_norm": 1.096007227897644,
+      "grad_norm": 0.3565673828125,
       "learning_rate": 1.6734693877551023e-06,
       "loss": 0.0407,
       "max_memory_allocated (GB)": 57.18,
@@ -1030,9 +1030,9 @@
     },
     {
       "epoch": 4.204081632653061,
-      "grad_norm": 4.377923965454102,
+      "grad_norm": 0.6948438882827759,
       "learning_rate": 1.5918367346938775e-06,
-      "loss": 0.0528,
+      "loss": 0.053,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1030,
@@ -1040,9 +1040,9 @@
     },
     {
       "epoch": 4.244897959183674,
-      "grad_norm": 0.7295175194740295,
+      "grad_norm": 0.5245764851570129,
       "learning_rate": 1.5102040816326532e-06,
-      "loss": 0.04,
+      "loss": 0.0399,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1040,
@@ -1050,9 +1050,9 @@
     },
     {
       "epoch": 4.285714285714286,
-      "grad_norm": 0.5129045248031616,
+      "grad_norm": 0.7932385802268982,
       "learning_rate": 1.4285714285714286e-06,
-      "loss": 0.0493,
+      "loss": 0.0502,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1050,
@@ -1060,7 +1060,7 @@
     },
     {
       "epoch": 4.326530612244898,
-      "grad_norm": 0.505799412727356,
+      "grad_norm": 0.30140048265457153,
       "learning_rate": 1.3469387755102043e-06,
       "loss": 0.046,
       "max_memory_allocated (GB)": 57.18,
@@ -1070,9 +1070,9 @@
     },
     {
       "epoch": 4.36734693877551,
-      "grad_norm": 0.8460046052932739,
+      "grad_norm": 0.570467472076416,
       "learning_rate": 1.2653061224489795e-06,
-      "loss": 0.0483,
+      "loss": 0.0487,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1070,
@@ -1080,9 +1080,9 @@
     },
     {
       "epoch": 4.408163265306122,
-      "grad_norm": 0.5636487603187561,
+      "grad_norm": 0.43690067529678345,
       "learning_rate": 1.1836734693877552e-06,
-      "loss": 0.0523,
+      "loss": 0.0521,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1080,
@@ -1090,9 +1090,9 @@
     },
     {
       "epoch": 4.448979591836735,
-      "grad_norm": 0.6893309354782104,
+      "grad_norm": 0.5298590660095215,
       "learning_rate": 1.1020408163265308e-06,
-      "loss": 0.0505,
+      "loss": 0.0506,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1090,
@@ -1100,9 +1100,9 @@
     },
     {
       "epoch": 4.489795918367347,
-      "grad_norm": 0.6645925045013428,
+      "grad_norm": 0.2310735136270523,
       "learning_rate": 1.0204081632653063e-06,
-      "loss": 0.0361,
+      "loss": 0.036,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1100,
@@ -1110,9 +1110,9 @@
     },
     {
       "epoch": 4.530612244897959,
-      "grad_norm": 0.14709816873073578,
+      "grad_norm": 0.13128583133220673,
       "learning_rate": 9.387755102040817e-07,
-      "loss": 0.0458,
+      "loss": 0.0463,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1110,
@@ -1120,9 +1120,9 @@
     },
     {
       "epoch": 4.571428571428571,
-      "grad_norm": 0.8794461488723755,
+      "grad_norm": 0.7682464122772217,
       "learning_rate": 8.571428571428572e-07,
-      "loss": 0.0414,
+      "loss": 0.0403,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1120,
@@ -1130,9 +1130,9 @@
     },
     {
       "epoch": 4.612244897959184,
-      "grad_norm": 1.345910906791687,
+      "grad_norm": 0.6608971953392029,
       "learning_rate": 7.755102040816327e-07,
-      "loss": 0.0526,
+      "loss": 0.0543,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1130,
@@ -1140,9 +1140,9 @@
     },
     {
       "epoch": 4.653061224489796,
-      "grad_norm": 0.5657418370246887,
+      "grad_norm": 0.8803687691688538,
       "learning_rate": 6.938775510204082e-07,
-      "loss": 0.0721,
+      "loss": 0.0728,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1140,
@@ -1150,9 +1150,9 @@
     },
     {
       "epoch": 4.6938775510204085,
-      "grad_norm": 2.0943238735198975,
+      "grad_norm": 4.121662139892578,
       "learning_rate": 6.122448979591837e-07,
-      "loss": 0.0515,
+      "loss": 0.0514,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1150,
@@ -1160,9 +1160,9 @@
     },
     {
       "epoch": 4.73469387755102,
-      "grad_norm": 3.1750009059906006,
+      "grad_norm": 0.7500938773155212,
       "learning_rate": 5.306122448979592e-07,
-      "loss": 0.0607,
+      "loss": 0.0612,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1160,
@@ -1170,9 +1170,9 @@
     },
     {
       "epoch": 4.775510204081632,
-      "grad_norm": 0.8910292387008667,
+      "grad_norm": 0.6001973748207092,
       "learning_rate": 4.489795918367347e-07,
-      "loss": 0.054,
+      "loss": 0.0549,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1170,
@@ -1180,9 +1180,9 @@
     },
     {
       "epoch": 4.816326530612245,
-      "grad_norm": 0.9796111583709717,
+      "grad_norm": 0.7522645592689514,
       "learning_rate": 3.6734693877551025e-07,
-      "loss": 0.0432,
+      "loss": 0.0445,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1180,
@@ -1190,9 +1190,9 @@
     },
     {
       "epoch": 4.857142857142857,
-      "grad_norm": 1.7865172624588013,
+      "grad_norm": 0.6640497446060181,
       "learning_rate": 2.8571428571428575e-07,
-      "loss": 0.0549,
+      "loss": 0.0542,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1190,
@@ -1200,9 +1200,9 @@
     },
     {
       "epoch": 4.8979591836734695,
-      "grad_norm": 0.7347144484519958,
+      "grad_norm": 0.8031227588653564,
       "learning_rate": 2.0408163265306121e-07,
-      "loss": 0.0711,
+      "loss": 0.0728,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1200,
@@ -1210,9 +1210,9 @@
     },
     {
       "epoch": 4.938775510204081,
-      "grad_norm": 1.210486888885498,
+      "grad_norm": 0.39187708497047424,
       "learning_rate": 1.2244897959183673e-07,
-      "loss": 0.0652,
+      "loss": 0.065,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1210,
@@ -1220,9 +1220,9 @@
     },
     {
       "epoch": 4.979591836734694,
-      "grad_norm": 1.2078640460968018,
+      "grad_norm": 3.809382915496826,
       "learning_rate": 4.0816326530612253e-08,
-      "loss": 0.0414,
+      "loss": 0.0417,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1220,
@@ -1235,10 +1235,10 @@
       "step": 1225,
       "total_flos": 3.0598946525952e+16,
       "total_memory_available (GB)": 94.62,
-      "train_loss": 0.06072675045655698,
-      "train_runtime": 1077.821,
-      "train_samples_per_second": 52.682,
-      "train_steps_per_second": 1.318
+      "train_loss": 0.06098026679486644,
+      "train_runtime": 1192.2443,
+      "train_samples_per_second": 46.607,
+      "train_steps_per_second": 1.166
     }
   ],
   "logging_steps": 10,

validation_results.json

@@ -1,9 +1,9 @@
 {
     "epoch": 5.0,
-    "eval_loss": 0.12326910346746445,
-    "eval_runtime": 26.1444,
-    "eval_samples_per_second": 36.653,
-    "eval_steps_per_second": 2.352,
+    "eval_loss": 0.12305960804224014,
+    "eval_runtime": 31.7562,
+    "eval_samples_per_second": 32.832,
+    "eval_steps_per_second": 2.107,
     "max_memory_allocated (GB)": 57.18,
     "memory_allocated (GB)": 51.27,
     "total_memory_available (GB)": 94.62