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.1327
+- 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.12054020911455154,
-    "eval_runtime": 14.1795,
-    "eval_samples_per_second": 70.009,
-    "eval_steps_per_second": 4.376,
-    "max_memory_allocated (GB)": 63.75,
+    "eval_loss": 0.1326669603586197,
+    "eval_runtime": 7.2012,
+    "eval_samples_per_second": 81.55,
+    "eval_steps_per_second": 5.097,
+    "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.06080986156755564,
-    "train_runtime": 1168.6251,
-    "train_samples_per_second": 48.37,
-    "train_steps_per_second": 1.21
+    "train_loss": 0.06081954878203723,
+    "train_runtime": 1056.8939,
+    "train_samples_per_second": 53.754,
+    "train_steps_per_second": 1.345
 }

runs/Oct16_18-58-21_workload-ai-workshop/events.out.tfevents.1729106325.workload-ai-workshop.14402.1

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

test_results.json

@@ -1,10 +1,10 @@
 {
     "epoch": 5.0,
-    "eval_loss": 0.12054020911455154,
-    "eval_runtime": 14.1795,
-    "eval_samples_per_second": 70.009,
-    "eval_steps_per_second": 4.376,
-    "max_memory_allocated (GB)": 63.75,
+    "eval_loss": 0.1326669603586197,
+    "eval_runtime": 7.2012,
+    "eval_samples_per_second": 81.55,
+    "eval_steps_per_second": 5.097,
+    "max_memory_allocated (GB)": 57.18,
     "memory_allocated (GB)": 51.27,
     "total_memory_available (GB)": 94.62
 }

train_results.json

@@ -1,11 +1,11 @@
 {
     "epoch": 5.0,
-    "max_memory_allocated (GB)": 63.75,
+    "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.06080986156755564,
-    "train_runtime": 1168.6251,
-    "train_samples_per_second": 48.37,
-    "train_steps_per_second": 1.21
+    "train_loss": 0.06081954878203723,
+    "train_runtime": 1056.8939,
+    "train_samples_per_second": 53.754,
+    "train_steps_per_second": 1.345
 }

trainer_state.json

@@ -20,9 +20,9 @@
     },
     {
       "epoch": 0.08163265306122448,
-      "grad_norm": 8.167236328125,
+      "grad_norm": 6.701302528381348,
       "learning_rate": 9.836734693877552e-06,
-      "loss": 0.1542,
+      "loss": 0.154,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 20,
@@ -30,9 +30,9 @@
     },
     {
       "epoch": 0.12244897959183673,
-      "grad_norm": 5.958656311035156,
+      "grad_norm": 5.337311267852783,
       "learning_rate": 9.755102040816327e-06,
-      "loss": 0.1236,
+      "loss": 0.1235,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 30,
@@ -40,9 +40,9 @@
     },
     {
       "epoch": 0.16326530612244897,
-      "grad_norm": 4.233438968658447,
+      "grad_norm": 4.59727144241333,
       "learning_rate": 9.673469387755103e-06,
-      "loss": 0.1105,
+      "loss": 0.1094,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 40,
@@ -50,9 +50,9 @@
     },
     {
       "epoch": 0.20408163265306123,
-      "grad_norm": 4.4948201179504395,
+      "grad_norm": 4.675155162811279,
       "learning_rate": 9.591836734693878e-06,
-      "loss": 0.1195,
+      "loss": 0.1191,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 50,
@@ -60,9 +60,9 @@
     },
     {
       "epoch": 0.24489795918367346,
-      "grad_norm": 1.827812910079956,
+      "grad_norm": 2.7298097610473633,
       "learning_rate": 9.510204081632653e-06,
-      "loss": 0.0797,
+      "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.889023780822754,
+      "grad_norm": 3.581591844558716,
       "learning_rate": 9.42857142857143e-06,
-      "loss": 0.0989,
+      "loss": 0.1015,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 70,
@@ -80,9 +80,9 @@
     },
     {
       "epoch": 0.32653061224489793,
-      "grad_norm": 1.9397954940795898,
+      "grad_norm": 1.5557887554168701,
       "learning_rate": 9.346938775510204e-06,
-      "loss": 0.1102,
+      "loss": 0.1117,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 80,
@@ -90,9 +90,9 @@
     },
     {
       "epoch": 0.3673469387755102,
-      "grad_norm": 3.2782671451568604,
+      "grad_norm": 8.605263710021973,
       "learning_rate": 9.26530612244898e-06,
-      "loss": 0.0824,
+      "loss": 0.0872,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 90,
@@ -100,9 +100,9 @@
     },
     {
       "epoch": 0.40816326530612246,
-      "grad_norm": 5.711858749389648,
+      "grad_norm": 3.075289249420166,
       "learning_rate": 9.183673469387756e-06,
-      "loss": 0.0762,
+      "loss": 0.076,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 100,
@@ -110,7 +110,7 @@
     },
     {
       "epoch": 0.4489795918367347,
-      "grad_norm": 3.395564317703247,
+      "grad_norm": 3.8531429767608643,
       "learning_rate": 9.102040816326532e-06,
       "loss": 0.0726,
       "max_memory_allocated (GB)": 57.18,
@@ -120,9 +120,9 @@
     },
     {
       "epoch": 0.4897959183673469,
-      "grad_norm": 2.2739310264587402,
+      "grad_norm": 3.202249765396118,
       "learning_rate": 9.020408163265307e-06,
-      "loss": 0.0705,
+      "loss": 0.0708,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 120,
@@ -130,9 +130,9 @@
     },
     {
       "epoch": 0.5306122448979592,
-      "grad_norm": 2.418794870376587,
+      "grad_norm": 1.818903923034668,
       "learning_rate": 8.938775510204082e-06,
-      "loss": 0.0595,
+      "loss": 0.059,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 130,
@@ -140,9 +140,9 @@
     },
     {
       "epoch": 0.5714285714285714,
-      "grad_norm": 2.2896311283111572,
+      "grad_norm": 1.1054496765136719,
       "learning_rate": 8.857142857142858e-06,
-      "loss": 0.0498,
+      "loss": 0.0487,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 140,
@@ -150,7 +150,7 @@
     },
     {
       "epoch": 0.6122448979591837,
-      "grad_norm": 1.7899913787841797,
+      "grad_norm": 2.9758386611938477,
       "learning_rate": 8.775510204081633e-06,
       "loss": 0.0629,
       "max_memory_allocated (GB)": 57.18,
@@ -160,9 +160,9 @@
     },
     {
       "epoch": 0.6530612244897959,
-      "grad_norm": 1.9983731508255005,
+      "grad_norm": 1.993240237236023,
       "learning_rate": 8.69387755102041e-06,
-      "loss": 0.071,
+      "loss": 0.0713,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 160,
@@ -170,7 +170,7 @@
     },
     {
       "epoch": 0.6938775510204082,
-      "grad_norm": 1.7236266136169434,
+      "grad_norm": 3.4001684188842773,
       "learning_rate": 8.612244897959184e-06,
       "loss": 0.0748,
       "max_memory_allocated (GB)": 57.18,
@@ -180,9 +180,9 @@
     },
     {
       "epoch": 0.7346938775510204,
-      "grad_norm": 1.1180106401443481,
+      "grad_norm": 1.881607174873352,
       "learning_rate": 8.530612244897961e-06,
-      "loss": 0.0599,
+      "loss": 0.0603,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 180,
@@ -190,9 +190,9 @@
     },
     {
       "epoch": 0.7755102040816326,
-      "grad_norm": 2.3057782649993896,
+      "grad_norm": 3.133998155593872,
       "learning_rate": 8.448979591836736e-06,
-      "loss": 0.0588,
+      "loss": 0.0592,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 190,
@@ -200,9 +200,9 @@
     },
     {
       "epoch": 0.8163265306122449,
-      "grad_norm": 0.8334403038024902,
+      "grad_norm": 0.934773325920105,
       "learning_rate": 8.36734693877551e-06,
-      "loss": 0.0622,
+      "loss": 0.0626,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 200,
@@ -210,9 +210,9 @@
     },
     {
       "epoch": 0.8571428571428571,
-      "grad_norm": 1.0548275709152222,
+      "grad_norm": 0.8929204940795898,
       "learning_rate": 8.285714285714287e-06,
-      "loss": 0.0623,
+      "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.3381606340408325,
+      "grad_norm": 1.0954375267028809,
       "learning_rate": 8.204081632653062e-06,
-      "loss": 0.0605,
+      "loss": 0.0596,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 220,
@@ -230,9 +230,9 @@
     },
     {
       "epoch": 0.9387755102040817,
-      "grad_norm": 0.809412956237793,
+      "grad_norm": 0.8736925721168518,
       "learning_rate": 8.122448979591837e-06,
-      "loss": 0.0566,
+      "loss": 0.0567,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 230,
@@ -240,9 +240,9 @@
     },
     {
       "epoch": 0.9795918367346939,
-      "grad_norm": 0.7182928323745728,
+      "grad_norm": 0.6469990611076355,
       "learning_rate": 8.040816326530613e-06,
-      "loss": 0.0496,
+      "loss": 0.0501,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 240,
@@ -250,9 +250,9 @@
     },
     {
       "epoch": 1.0204081632653061,
-      "grad_norm": 1.1081018447875977,
+      "grad_norm": 0.7664241790771484,
       "learning_rate": 7.959183673469388e-06,
-      "loss": 0.0601,
+      "loss": 0.06,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 250,
@@ -260,9 +260,9 @@
     },
     {
       "epoch": 1.0612244897959184,
-      "grad_norm": 3.9478495121002197,
+      "grad_norm": 3.3073790073394775,
       "learning_rate": 7.877551020408164e-06,
-      "loss": 0.0541,
+      "loss": 0.0538,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 260,
@@ -270,9 +270,9 @@
     },
     {
       "epoch": 1.1020408163265305,
-      "grad_norm": 0.8079606294631958,
+      "grad_norm": 0.7618402242660522,
       "learning_rate": 7.79591836734694e-06,
-      "loss": 0.063,
+      "loss": 0.0621,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 270,
@@ -280,9 +280,9 @@
     },
     {
       "epoch": 1.1428571428571428,
-      "grad_norm": 1.4108704328536987,
+      "grad_norm": 2.667276620864868,
       "learning_rate": 7.714285714285716e-06,
-      "loss": 0.0579,
+      "loss": 0.0578,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 280,
@@ -290,9 +290,9 @@
     },
     {
       "epoch": 1.183673469387755,
-      "grad_norm": 1.544438123703003,
+      "grad_norm": 0.7617210149765015,
       "learning_rate": 7.63265306122449e-06,
-      "loss": 0.0518,
+      "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": 1.0270023345947266,
+      "grad_norm": 1.4268369674682617,
       "learning_rate": 7.551020408163265e-06,
-      "loss": 0.0658,
+      "loss": 0.0638,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 300,
@@ -310,9 +310,9 @@
     },
     {
       "epoch": 1.2653061224489797,
-      "grad_norm": 1.0347421169281006,
+      "grad_norm": 1.0776606798171997,
       "learning_rate": 7.469387755102041e-06,
-      "loss": 0.0667,
+      "loss": 0.0668,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 310,
@@ -320,9 +320,9 @@
     },
     {
       "epoch": 1.306122448979592,
-      "grad_norm": 1.5818060636520386,
+      "grad_norm": 1.2839634418487549,
       "learning_rate": 7.387755102040817e-06,
-      "loss": 0.0502,
+      "loss": 0.049,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 320,
@@ -330,9 +330,9 @@
     },
     {
       "epoch": 1.346938775510204,
-      "grad_norm": 1.2030609846115112,
+      "grad_norm": 0.6254063844680786,
       "learning_rate": 7.306122448979592e-06,
-      "loss": 0.0616,
+      "loss": 0.06,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 330,
@@ -340,9 +340,9 @@
     },
     {
       "epoch": 1.3877551020408163,
-      "grad_norm": 1.2504222393035889,
+      "grad_norm": 1.5291829109191895,
       "learning_rate": 7.224489795918368e-06,
-      "loss": 0.0543,
+      "loss": 0.055,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 340,
@@ -350,7 +350,7 @@
     },
     {
       "epoch": 1.4285714285714286,
-      "grad_norm": 0.8420921564102173,
+      "grad_norm": 0.6570836305618286,
       "learning_rate": 7.1428571428571436e-06,
       "loss": 0.0488,
       "max_memory_allocated (GB)": 57.18,
@@ -360,9 +360,9 @@
     },
     {
       "epoch": 1.469387755102041,
-      "grad_norm": 1.517096757888794,
+      "grad_norm": 4.310382843017578,
       "learning_rate": 7.061224489795919e-06,
-      "loss": 0.0467,
+      "loss": 0.0464,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 360,
@@ -370,7 +370,7 @@
     },
     {
       "epoch": 1.510204081632653,
-      "grad_norm": 1.4490768909454346,
+      "grad_norm": 1.3656598329544067,
       "learning_rate": 6.979591836734695e-06,
       "loss": 0.0585,
       "max_memory_allocated (GB)": 57.18,
@@ -380,9 +380,9 @@
     },
     {
       "epoch": 1.5510204081632653,
-      "grad_norm": 0.5360353589057922,
+      "grad_norm": 0.9290764927864075,
       "learning_rate": 6.8979591836734705e-06,
-      "loss": 0.046,
+      "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": 0.8193866610527039,
+      "grad_norm": 1.746748447418213,
       "learning_rate": 6.816326530612245e-06,
-      "loss": 0.0657,
+      "loss": 0.0648,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 390,
@@ -400,9 +400,9 @@
     },
     {
       "epoch": 1.6326530612244898,
-      "grad_norm": 0.5883302092552185,
+      "grad_norm": 0.5461153388023376,
       "learning_rate": 6.734693877551021e-06,
-      "loss": 0.0609,
+      "loss": 0.0613,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 400,
@@ -410,9 +410,9 @@
     },
     {
       "epoch": 1.6734693877551021,
-      "grad_norm": 0.6720415949821472,
+      "grad_norm": 0.4624268710613251,
       "learning_rate": 6.653061224489797e-06,
-      "loss": 0.0603,
+      "loss": 0.0604,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 410,
@@ -420,9 +420,9 @@
     },
     {
       "epoch": 1.7142857142857144,
-      "grad_norm": 1.368994951248169,
+      "grad_norm": 4.135178089141846,
       "learning_rate": 6.571428571428572e-06,
-      "loss": 0.0528,
+      "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.25535887479782104,
+      "grad_norm": 0.19531233608722687,
       "learning_rate": 6.489795918367348e-06,
-      "loss": 0.0674,
+      "loss": 0.0675,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 430,
@@ -440,9 +440,9 @@
     },
     {
       "epoch": 1.7959183673469388,
-      "grad_norm": 0.8839388489723206,
+      "grad_norm": 1.667291522026062,
       "learning_rate": 6.408163265306124e-06,
-      "loss": 0.0612,
+      "loss": 0.0606,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 440,
@@ -450,9 +450,9 @@
     },
     {
       "epoch": 1.836734693877551,
-      "grad_norm": 0.4119959771633148,
+      "grad_norm": 0.4277726709842682,
       "learning_rate": 6.326530612244899e-06,
-      "loss": 0.0435,
+      "loss": 0.044,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 450,
@@ -460,9 +460,9 @@
     },
     {
       "epoch": 1.8775510204081631,
-      "grad_norm": 0.45877301692962646,
+      "grad_norm": 0.3809329569339752,
       "learning_rate": 6.244897959183675e-06,
-      "loss": 0.0543,
+      "loss": 0.0537,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 460,
@@ -470,9 +470,9 @@
     },
     {
       "epoch": 1.9183673469387754,
-      "grad_norm": 0.9594807624816895,
+      "grad_norm": 1.0573110580444336,
       "learning_rate": 6.163265306122449e-06,
-      "loss": 0.0433,
+      "loss": 0.0427,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 470,
@@ -480,9 +480,9 @@
     },
     {
       "epoch": 1.9591836734693877,
-      "grad_norm": 0.168818861246109,
+      "grad_norm": 0.20000681281089783,
       "learning_rate": 6.0816326530612245e-06,
-      "loss": 0.0497,
+      "loss": 0.0502,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 480,
@@ -490,9 +490,9 @@
     },
     {
       "epoch": 2.0,
-      "grad_norm": 1.1468336582183838,
+      "grad_norm": 0.9971123337745667,
       "learning_rate": 6e-06,
-      "loss": 0.0651,
+      "loss": 0.0636,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 490,
@@ -500,7 +500,7 @@
     },
     {
       "epoch": 2.0408163265306123,
-      "grad_norm": 0.6001573204994202,
+      "grad_norm": 0.673917293548584,
       "learning_rate": 5.918367346938776e-06,
       "loss": 0.0591,
       "max_memory_allocated (GB)": 57.18,
@@ -510,9 +510,9 @@
     },
     {
       "epoch": 2.0816326530612246,
-      "grad_norm": 0.8006247878074646,
+      "grad_norm": 0.6252917647361755,
       "learning_rate": 5.8367346938775515e-06,
-      "loss": 0.0489,
+      "loss": 0.0485,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 510,
@@ -520,9 +520,9 @@
     },
     {
       "epoch": 2.122448979591837,
-      "grad_norm": 0.6839123964309692,
+      "grad_norm": 1.9926584959030151,
       "learning_rate": 5.755102040816327e-06,
-      "loss": 0.0456,
+      "loss": 0.0461,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 520,
@@ -530,9 +530,9 @@
     },
     {
       "epoch": 2.163265306122449,
-      "grad_norm": 1.020135521888733,
+      "grad_norm": 1.1171493530273438,
       "learning_rate": 5.673469387755103e-06,
-      "loss": 0.0627,
+      "loss": 0.063,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 530,
@@ -540,9 +540,9 @@
     },
     {
       "epoch": 2.204081632653061,
-      "grad_norm": 0.7156575322151184,
+      "grad_norm": 0.5303067564964294,
       "learning_rate": 5.591836734693878e-06,
-      "loss": 0.0584,
+      "loss": 0.0594,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 540,
@@ -550,9 +550,9 @@
     },
     {
       "epoch": 2.2448979591836733,
-      "grad_norm": 1.195730447769165,
+      "grad_norm": 1.5552964210510254,
       "learning_rate": 5.510204081632653e-06,
-      "loss": 0.0628,
+      "loss": 0.0632,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 550,
@@ -560,9 +560,9 @@
     },
     {
       "epoch": 2.2857142857142856,
-      "grad_norm": 0.8824738264083862,
+      "grad_norm": 7.580259799957275,
       "learning_rate": 5.428571428571429e-06,
-      "loss": 0.0572,
+      "loss": 0.058,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 560,
@@ -570,7 +570,7 @@
     },
     {
       "epoch": 2.326530612244898,
-      "grad_norm": 1.3354676961898804,
+      "grad_norm": 0.8162185549736023,
       "learning_rate": 5.3469387755102045e-06,
       "loss": 0.0585,
       "max_memory_allocated (GB)": 57.18,
@@ -580,9 +580,9 @@
     },
     {
       "epoch": 2.36734693877551,
-      "grad_norm": 1.9097795486450195,
+      "grad_norm": 3.4555113315582275,
       "learning_rate": 5.26530612244898e-06,
-      "loss": 0.0668,
+      "loss": 0.067,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 580,
@@ -590,9 +590,9 @@
     },
     {
       "epoch": 2.4081632653061225,
-      "grad_norm": 0.8386860489845276,
+      "grad_norm": 0.45223644375801086,
       "learning_rate": 5.183673469387756e-06,
-      "loss": 0.0542,
+      "loss": 0.0537,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 590,
@@ -600,9 +600,9 @@
     },
     {
       "epoch": 2.4489795918367347,
-      "grad_norm": 1.2374165058135986,
+      "grad_norm": 0.9515520930290222,
       "learning_rate": 5.1020408163265315e-06,
-      "loss": 0.0436,
+      "loss": 0.0435,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 600,
@@ -610,9 +610,9 @@
     },
     {
       "epoch": 2.489795918367347,
-      "grad_norm": 0.3719439208507538,
+      "grad_norm": 0.5981607437133789,
       "learning_rate": 5.020408163265307e-06,
-      "loss": 0.0408,
+      "loss": 0.0414,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 610,
@@ -620,9 +620,9 @@
     },
     {
       "epoch": 2.5306122448979593,
-      "grad_norm": 0.905327558517456,
+      "grad_norm": 0.8998698592185974,
       "learning_rate": 4.938775510204082e-06,
-      "loss": 0.0429,
+      "loss": 0.0432,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 620,
@@ -630,9 +630,9 @@
     },
     {
       "epoch": 2.571428571428571,
-      "grad_norm": 0.628597617149353,
+      "grad_norm": 0.5227445960044861,
       "learning_rate": 4.857142857142858e-06,
-      "loss": 0.0812,
+      "loss": 0.0813,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 630,
@@ -640,9 +640,9 @@
     },
     {
       "epoch": 2.612244897959184,
-      "grad_norm": 1.3098090887069702,
+      "grad_norm": 1.117499828338623,
       "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.5635781288146973,
+      "grad_norm": 0.6399075984954834,
       "learning_rate": 4.693877551020409e-06,
-      "loss": 0.0466,
+      "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.6197735071182251,
+      "grad_norm": 3.8376593589782715,
       "learning_rate": 4.612244897959184e-06,
-      "loss": 0.0474,
+      "loss": 0.0473,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 660,
@@ -670,9 +670,9 @@
     },
     {
       "epoch": 2.7346938775510203,
-      "grad_norm": 0.6390748620033264,
+      "grad_norm": 0.6460541486740112,
       "learning_rate": 4.530612244897959e-06,
-      "loss": 0.0863,
+      "loss": 0.0872,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 670,
@@ -680,9 +680,9 @@
     },
     {
       "epoch": 2.7755102040816326,
-      "grad_norm": 1.6307971477508545,
+      "grad_norm": 1.897635579109192,
       "learning_rate": 4.448979591836735e-06,
-      "loss": 0.0489,
+      "loss": 0.0487,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 680,
@@ -690,9 +690,9 @@
     },
     {
       "epoch": 2.816326530612245,
-      "grad_norm": 0.3477366864681244,
+      "grad_norm": 0.3259105682373047,
       "learning_rate": 4.367346938775511e-06,
-      "loss": 0.0448,
+      "loss": 0.0449,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 690,
@@ -700,9 +700,9 @@
     },
     {
       "epoch": 2.857142857142857,
-      "grad_norm": 0.41136085987091064,
+      "grad_norm": 0.3581603467464447,
       "learning_rate": 4.2857142857142855e-06,
-      "loss": 0.0431,
+      "loss": 0.0435,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 700,
@@ -710,9 +710,9 @@
     },
     {
       "epoch": 2.8979591836734695,
-      "grad_norm": 1.1029525995254517,
+      "grad_norm": 0.8618249297142029,
       "learning_rate": 4.204081632653061e-06,
-      "loss": 0.0551,
+      "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": 0.8994241952896118,
+      "grad_norm": 1.4629887342453003,
       "learning_rate": 4.122448979591837e-06,
-      "loss": 0.0581,
+      "loss": 0.0584,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 720,
@@ -730,7 +730,7 @@
     },
     {
       "epoch": 2.979591836734694,
-      "grad_norm": 0.1889757364988327,
+      "grad_norm": 0.21725626289844513,
       "learning_rate": 4.040816326530612e-06,
       "loss": 0.034,
       "max_memory_allocated (GB)": 57.18,
@@ -740,9 +740,9 @@
     },
     {
       "epoch": 3.020408163265306,
-      "grad_norm": 1.7815334796905518,
+      "grad_norm": 1.257843017578125,
       "learning_rate": 3.959183673469388e-06,
-      "loss": 0.0536,
+      "loss": 0.054,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 740,
@@ -750,9 +750,9 @@
     },
     {
       "epoch": 3.061224489795918,
-      "grad_norm": 0.4372510612010956,
+      "grad_norm": 0.7538784742355347,
       "learning_rate": 3.877551020408164e-06,
-      "loss": 0.0617,
+      "loss": 0.0611,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 750,
@@ -760,9 +760,9 @@
     },
     {
       "epoch": 3.1020408163265305,
-      "grad_norm": 5.120749473571777,
+      "grad_norm": 0.6883149147033691,
       "learning_rate": 3.795918367346939e-06,
-      "loss": 0.0518,
+      "loss": 0.0523,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 760,
@@ -770,9 +770,9 @@
     },
     {
       "epoch": 3.142857142857143,
-      "grad_norm": 6.453648090362549,
+      "grad_norm": 0.6788195967674255,
       "learning_rate": 3.7142857142857146e-06,
-      "loss": 0.069,
+      "loss": 0.0685,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 770,
@@ -780,9 +780,9 @@
     },
     {
       "epoch": 3.183673469387755,
-      "grad_norm": 0.6512885093688965,
+      "grad_norm": 0.5011599063873291,
       "learning_rate": 3.6326530612244903e-06,
-      "loss": 0.0649,
+      "loss": 0.0643,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 780,
@@ -790,9 +790,9 @@
     },
     {
       "epoch": 3.2244897959183674,
-      "grad_norm": 0.3266737759113312,
+      "grad_norm": 0.4473969042301178,
       "learning_rate": 3.5510204081632655e-06,
-      "loss": 0.0523,
+      "loss": 0.0525,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 790,
@@ -800,9 +800,9 @@
     },
     {
       "epoch": 3.2653061224489797,
-      "grad_norm": 0.4506176710128784,
+      "grad_norm": 0.5276008248329163,
       "learning_rate": 3.469387755102041e-06,
-      "loss": 0.0432,
+      "loss": 0.0422,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 800,
@@ -810,9 +810,9 @@
     },
     {
       "epoch": 3.306122448979592,
-      "grad_norm": 0.8929914236068726,
+      "grad_norm": 0.8036576509475708,
       "learning_rate": 3.3877551020408168e-06,
-      "loss": 0.0554,
+      "loss": 0.0553,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 810,
@@ -820,9 +820,9 @@
     },
     {
       "epoch": 3.3469387755102042,
-      "grad_norm": 0.7046924233436584,
+      "grad_norm": 0.7645472884178162,
       "learning_rate": 3.3061224489795924e-06,
-      "loss": 0.0453,
+      "loss": 0.0457,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 820,
@@ -830,9 +830,9 @@
     },
     {
       "epoch": 3.387755102040816,
-      "grad_norm": 0.29230576753616333,
+      "grad_norm": 0.27847281098365784,
       "learning_rate": 3.2244897959183672e-06,
-      "loss": 0.0491,
+      "loss": 0.0492,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 830,
@@ -840,7 +840,7 @@
     },
     {
       "epoch": 3.4285714285714284,
-      "grad_norm": 0.4533096253871918,
+      "grad_norm": 0.3790993094444275,
       "learning_rate": 3.142857142857143e-06,
       "loss": 0.0529,
       "max_memory_allocated (GB)": 57.18,
@@ -850,9 +850,9 @@
     },
     {
       "epoch": 3.4693877551020407,
-      "grad_norm": 0.5383632183074951,
+      "grad_norm": 0.4730982184410095,
       "learning_rate": 3.0612244897959185e-06,
-      "loss": 0.0823,
+      "loss": 0.0828,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 850,
@@ -860,9 +860,9 @@
     },
     {
       "epoch": 3.510204081632653,
-      "grad_norm": 2.8597779273986816,
+      "grad_norm": 1.0951735973358154,
       "learning_rate": 2.979591836734694e-06,
-      "loss": 0.0456,
+      "loss": 0.0453,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 860,
@@ -870,9 +870,9 @@
     },
     {
       "epoch": 3.5510204081632653,
-      "grad_norm": 0.26686975359916687,
+      "grad_norm": 0.46890848875045776,
       "learning_rate": 2.8979591836734694e-06,
-      "loss": 0.064,
+      "loss": 0.0643,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 870,
@@ -880,9 +880,9 @@
     },
     {
       "epoch": 3.5918367346938775,
-      "grad_norm": 0.7789614796638489,
+      "grad_norm": 0.9014816880226135,
       "learning_rate": 2.816326530612245e-06,
-      "loss": 0.0441,
+      "loss": 0.0439,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 880,
@@ -890,9 +890,9 @@
     },
     {
       "epoch": 3.63265306122449,
-      "grad_norm": 0.21532948315143585,
+      "grad_norm": 0.1719776690006256,
       "learning_rate": 2.7346938775510207e-06,
-      "loss": 0.0291,
+      "loss": 0.0288,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 890,
@@ -900,9 +900,9 @@
     },
     {
       "epoch": 3.673469387755102,
-      "grad_norm": 0.741765558719635,
+      "grad_norm": 0.6494420170783997,
       "learning_rate": 2.6530612244897964e-06,
-      "loss": 0.0512,
+      "loss": 0.0515,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 900,
@@ -910,9 +910,9 @@
     },
     {
       "epoch": 3.7142857142857144,
-      "grad_norm": 0.6416855454444885,
+      "grad_norm": 0.4721522629261017,
       "learning_rate": 2.571428571428571e-06,
-      "loss": 0.0606,
+      "loss": 0.0604,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 910,
@@ -920,9 +920,9 @@
     },
     {
       "epoch": 3.7551020408163263,
-      "grad_norm": 0.14841973781585693,
+      "grad_norm": 0.13098274171352386,
       "learning_rate": 2.489795918367347e-06,
-      "loss": 0.0542,
+      "loss": 0.055,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 920,
@@ -930,9 +930,9 @@
     },
     {
       "epoch": 3.795918367346939,
-      "grad_norm": 0.4417996108531952,
+      "grad_norm": 0.39990562200546265,
       "learning_rate": 2.4081632653061225e-06,
-      "loss": 0.0498,
+      "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.9759775400161743,
+      "grad_norm": 1.020967721939087,
       "learning_rate": 2.326530612244898e-06,
-      "loss": 0.0491,
+      "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.020371913909912,
+      "grad_norm": 0.6601265072822571,
       "learning_rate": 2.244897959183674e-06,
-      "loss": 0.0597,
+      "loss": 0.0594,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 950,
@@ -960,9 +960,9 @@
     },
     {
       "epoch": 3.9183673469387754,
-      "grad_norm": 0.3064863085746765,
+      "grad_norm": 0.34864309430122375,
       "learning_rate": 2.1632653061224495e-06,
-      "loss": 0.0499,
+      "loss": 0.0495,
       "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 960,
@@ -970,7 +970,7 @@
     },
     {
       "epoch": 3.9591836734693877,
-      "grad_norm": 0.7580925226211548,
+      "grad_norm": 1.2881830930709839,
       "learning_rate": 2.0816326530612247e-06,
       "loss": 0.0742,
       "max_memory_allocated (GB)": 57.18,
@@ -980,265 +980,265 @@
     },
     {
       "epoch": 4.0,
-      "grad_norm": 0.6833075881004333,
+      "grad_norm": 0.7139947414398193,
       "learning_rate": 2.0000000000000003e-06,
-      "loss": 0.0708,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0698,
+      "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.5641142725944519,
+      "grad_norm": 0.7273740172386169,
       "learning_rate": 1.9183673469387756e-06,
-      "loss": 0.0481,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0483,
+      "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.8568029403686523,
+      "grad_norm": 0.811747133731842,
       "learning_rate": 1.8367346938775512e-06,
-      "loss": 0.0626,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0623,
+      "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.5912718772888184,
+      "grad_norm": 0.595486044883728,
       "learning_rate": 1.7551020408163267e-06,
-      "loss": 0.0628,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0627,
+      "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.3173392713069916,
+      "grad_norm": 0.40334606170654297,
       "learning_rate": 1.6734693877551023e-06,
-      "loss": 0.0402,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0404,
+      "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.8902315497398376,
+      "grad_norm": 0.5992821455001831,
       "learning_rate": 1.5918367346938775e-06,
-      "loss": 0.0536,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0529,
+      "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.5009722113609314,
+      "grad_norm": 0.5139011740684509,
       "learning_rate": 1.5102040816326532e-06,
-      "loss": 0.0399,
-      "max_memory_allocated (GB)": 63.75,
+      "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": 1.8656221628189087,
+      "grad_norm": 0.35282936692237854,
       "learning_rate": 1.4285714285714286e-06,
-      "loss": 0.0499,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0501,
+      "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.4257819950580597,
+      "grad_norm": 0.32772839069366455,
       "learning_rate": 1.3469387755102043e-06,
-      "loss": 0.0459,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0456,
+      "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.5823583006858826,
+      "grad_norm": 0.6208099126815796,
       "learning_rate": 1.2653061224489795e-06,
-      "loss": 0.0488,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0486,
+      "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.40693071484565735,
+      "grad_norm": 0.7009238004684448,
       "learning_rate": 1.1836734693877552e-06,
-      "loss": 0.0525,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0521,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1080,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.448979591836735,
-      "grad_norm": 1.8890392780303955,
+      "grad_norm": 0.3578231930732727,
       "learning_rate": 1.1020408163265308e-06,
-      "loss": 0.0512,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0509,
+      "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.3593562841415405,
+      "grad_norm": 0.30579662322998047,
       "learning_rate": 1.0204081632653063e-06,
-      "loss": 0.0364,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0365,
+      "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.1553877741098404,
+      "grad_norm": 0.10309538245201111,
       "learning_rate": 9.387755102040817e-07,
-      "loss": 0.0465,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0462,
+      "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.6775248050689697,
+      "grad_norm": 0.6286255121231079,
       "learning_rate": 8.571428571428572e-07,
-      "loss": 0.0406,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.04,
+      "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.5735678672790527,
+      "grad_norm": 0.6702510714530945,
       "learning_rate": 7.755102040816327e-07,
-      "loss": 0.0539,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0537,
+      "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.7891528606414795,
+      "grad_norm": 7.299179553985596,
       "learning_rate": 6.938775510204082e-07,
-      "loss": 0.0732,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0725,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1140,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.6938775510204085,
-      "grad_norm": 0.7845800518989563,
+      "grad_norm": 0.8366469144821167,
       "learning_rate": 6.122448979591837e-07,
       "loss": 0.0515,
-      "max_memory_allocated (GB)": 63.75,
+      "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.0361818075180054,
+      "grad_norm": 0.7851714491844177,
       "learning_rate": 5.306122448979592e-07,
-      "loss": 0.0608,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.061,
+      "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.42603600025177,
+      "grad_norm": 0.5544641613960266,
       "learning_rate": 4.489795918367347e-07,
-      "loss": 0.0549,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0552,
+      "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.6727630496025085,
+      "grad_norm": 0.6136904358863831,
       "learning_rate": 3.6734693877551025e-07,
-      "loss": 0.0441,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.044,
+      "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.8499141335487366,
+      "grad_norm": 0.957547128200531,
       "learning_rate": 2.8571428571428575e-07,
-      "loss": 0.0544,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0545,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1190,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.8979591836734695,
-      "grad_norm": 0.7604736685752869,
+      "grad_norm": 0.7660542726516724,
       "learning_rate": 2.0408163265306121e-07,
-      "loss": 0.0728,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0725,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1200,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.938775510204081,
-      "grad_norm": 1.0298157930374146,
+      "grad_norm": 1.0952503681182861,
       "learning_rate": 1.2244897959183673e-07,
-      "loss": 0.0654,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0648,
+      "max_memory_allocated (GB)": 57.18,
       "memory_allocated (GB)": 50.57,
       "step": 1210,
       "total_memory_available (GB)": 94.62
     },
     {
       "epoch": 4.979591836734694,
-      "grad_norm": 1.1890877485275269,
+      "grad_norm": 1.1168698072433472,
       "learning_rate": 4.0816326530612253e-08,
-      "loss": 0.0413,
-      "max_memory_allocated (GB)": 63.75,
+      "loss": 0.0415,
+      "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)": 63.75,
+      "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.06080986156755564,
-      "train_runtime": 1168.6251,
-      "train_samples_per_second": 48.37,
-      "train_steps_per_second": 1.21
+      "train_loss": 0.06081954878203723,
+      "train_runtime": 1056.8939,
+      "train_samples_per_second": 53.754,
+      "train_steps_per_second": 1.345
     }
   ],
   "logging_steps": 10,

validation_results.json

@@ -1,10 +1,10 @@
 {
     "epoch": 5.0,
-    "eval_loss": 0.13143940269947052,
-    "eval_runtime": 46.706,
-    "eval_samples_per_second": 27.822,
-    "eval_steps_per_second": 1.786,
-    "max_memory_allocated (GB)": 63.75,
+    "eval_loss": 0.12077265232801437,
+    "eval_runtime": 29.4902,
+    "eval_samples_per_second": 49.245,
+    "eval_steps_per_second": 3.161,
+    "max_memory_allocated (GB)": 57.18,
     "memory_allocated (GB)": 51.27,
     "total_memory_available (GB)": 94.62
 }