Double buffer metrics In Train to avoid gap between steps

[rwitten]

Old gap (4.1 ms): https://screenshot.googleplex.com/3H8J9bAkWGAFfQL
New gap (.092 ms): https://screenshot.googleplex.com/BZzKRkh6dQ875VX

Unfortunately our reporting is a bit sus:

Per train step, total TFLOPs will be 170.71, split as 94.20% learnable weight flops and 5.80% attention flops
completed step: 0, seconds: 5.120, TFLOP/s/device: 33.343, loss: 12.620
To see full metrics 'tensorboard --logdir=gs://runner-maxtext-logs/fake/tensorboard/'
completed step: 1, seconds: 0.811, TFLOP/s/device: 210.381, loss: 12.603
completed step: 2, seconds: 0.609, TFLOP/s/device: 280.094, loss: 10.725
completed step: 3, seconds: 0.996, TFLOP/s/device: 171.314, loss: 9.784
completed step: 4, seconds: 1.207, TFLOP/s/device: 141.395, loss: 9.447
completed step: 5, seconds: 1.202, TFLOP/s/device: 141.962, loss: 9.180
completed step: 6, seconds: 1.204, TFLOP/s/device: 141.783, loss: 9.001
completed step: 7, seconds: 1.205, TFLOP/s/device: 141.693, loss: 8.896
completed step: 8, seconds: 1.205, TFLOP/s/device: 141.641, loss: 8.807
completed step: 9, seconds: 1.206, TFLOP/s/device: 141.548, loss: 8.725

I honestly don't see a fix for this reporting issue.

[gobbleturk]

Nit: Do these need to be global variables? Where else are they used? Are they global so they are not input/outputs of write_metrics?

[rwitten]

Not a nit! They need to be stateful so they can be "remembered" -- a pure function can't remember them. In the future as I hint we should encapsulate all of the metrics details into a class and that class can take care of this.

[ZhiyuLi-goog]

A tiny question, does this line change the dataloading? Previously, we load first batch before entering the loop and now inside the loop. I remember there was a PR relevant.