Mlperf/4.1 grain

[aireenmei]

• Use a custom grain wheel (Need to rebuild dependency image, see changes in requirements.txt and setup.sh)
• Add grain support for per_device_batch_size<1 while eval_per_device_batch_size>1
• Run on v5p-128 using xpk (per_device_batch_size=0.5, eval_per_device_batch_size=1), using the script in this PR. command: bash MaxText/configs/v5e/mlperf-grain-tpu.sh PER_DEVICE_BATCH_SIZE=0.5 GRAIN_WORKER_COUNT=1 ICI_TENSOR=16 PLATFORM=gke Results: https://cloudlogging.app.goo.gl/5BpmD9iCEm7uVGNx5
• The padding batch for eval doesn't work well. Suggest to first test with eval off (eval_interval=-1), then try eval on. Eval may not terminate but hang on some slice sizes. I'll continue working on that.

[aireenmei]

The padding batch issue for eval is fixed now. To use this feature, set eval_steps to the just enough to finish the eval data (the data iter will generate empty examples until eval_steps met). For mlperf eval data, set eval_steps = math.ceil(5700/global_batch_size) and check if the output total_weights==11590004. Add more eval_steps if needed for total_weights to reach that value.