Refactor: Move metrics logic to MetricLogger class

[SurbhiJainUSC]

Description

This PR initiates the refactoring of train.py. As part of refactoring effort, we will be:

• Consolidating all the metrics-related code into a dedicated MetricLogger class (PR)
• Moving some methods in train.py to a util module that can be re-used by other trainers (To be followed)

Tests

Ran MaxText/train.py with gcs_metrics=True and metrics_file="/tmp/maxtext_metrics.txt" configurations to ensure that the metrics were stored in a local file, GCS and TensorBoard.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[gemini-code-assist]

Important

The terms of service for this installation has not been accepted. Please ask the Organization owners to visit the Gemini Code Assist Admin Console to sign it.

[bvandermoon]

Is this clear no longer needed?

[SurbhiJainUSC]

This method is basically setting the global variables to None. I have moved those global variables to class variables instead. AFAIK, python's garbage collector will automatically clear the memory occupied by that class object when program exits. So, this method shouldn't be required.
@gobbleturk - LMK if I am missing anything here.

[gobbleturk]

I have seen issues when calling the main() of train.py mutiple times within the same python invocation, e.g. if you run a script like

  from train.py import main as train_main 
  train_main()
  train_main()

then the second call of train_main() will inherit the global variables (the metrics) of the end of the first unless they are explicitly cleared. I don't have much understanding of python garbage collection though. If we refactor train well enough perhaps this clear method is not needed.

[SurbhiJainUSC]

In this case, each invocation of train_main() will have its own metric logger object. All the metrics related data can only be accessed by the class object. So, IMO the second call of train_main() will not be able to access the metrics from the first train_main() call.

[gobbleturk]

LGTM, this is awesome thanks Surbhi!

I do think we still need the clear method as described in the comment though - but if you can run a test with two train_main calls in the same python script and see that the second run has fresh metrics perhaps it can be removed

[bvandermoon]

LGTM pending the outcome of the clear method discussion

[SurbhiJainUSC]

LGTM pending the outcome of the clear method discussion

Verified that two sequential train_main() function calls within the same script produced different sets of metrics.