Prepare a MFU training automation script

[miladm]

PM would like to have a script like the following available to PyTorch/XLA customers.

https://github.com/google/maxtext/blob/main/MaxText/train.py#L75-L86

@alanwaketan do you think it's better to support it here or somewhere in torch_xla core. (in the absence of a side library)

[alanwaketan]

This is LLM specific... It doesn't make any sense to put it in the generic framework...

[alanwaketan]

Maybe somewhere in the tpu-pytorch model garden.