Hi, @aireenmei

Awesome and thank you for the fix!

Do we expect to unblock eval_per_device_batch_size < 1. with this PR?
And do you happen to have some test results?

@ZhiyuLi-goog It should support eval_per_device_batch_size < 1. But let me actually run the tests to make sure. Will get back to you.

@khatwanimohit @ZhiyuLi-goog
I made more chnages. Now it supports eval_per_device_batch_size < 1. Could both of you review again?

Cases tested on v4-128:

• M_EXPANSION_FACTOR_REAL_DATA=2, per_device_batch_size=per_eval_device_batch_size=8.0
  https://console.cloud.google.com/kubernetes/service/us-central2/v4-128-bodaborg-us-central2-b/default/aireen-20241011-12-10/details?project=cloud-tpu-multipod-dev
• per_device_batch_size=eval_per_device_batch_size=0.5 ici_tensor_parallelism=4, expansion_factor=-1: https://console.cloud.google.com/kubernetes/service/us-central2/v4-128-bodaborg-us-central2-b/default/aireen-20241010-18-28/details?project=cloud-tpu-multipod-dev
• per_device_batch_size=1.0 eval_per_device_batch_size=0.5, ici_tensor_parallelism=4, expansion_factor=-1: https://console.cloud.google.com/kubernetes/service/us-central2/v4-128-bodaborg-us-central2-b/default/aireen-20241010-19-02/details?project=cloud-tpu-multipod-dev
• per_device_batch_size=0.5 eval_per_device_batch_size=1, ici_tensor_parallelism=4, expansion_factor=-1: https://console.cloud.google.com/kubernetes/service/us-central2/v4-128-bodaborg-us-central2-b/default/aireen-20241010-18-36/details?project=cloud-tpu-multipod-dev

@khatwanimohit @ZhiyuLi-goog I made more chnages. Now it supports eval_per_device_batch_size < 1. Could both of you review again?

Cases tested on v4-128:

• M_EXPANSION_FACTOR_REAL_DATA=2, per_device_batch_size=per_eval_device_batch_size=8.0
  https://console.cloud.google.com/kubernetes/service/us-central2/v4-128-bodaborg-us-central2-b/default/aireen-20241011-12-10/details?project=cloud-tpu-multipod-dev
• per_device_batch_size=eval_per_device_batch_size=0.5 ici_tensor_parallelism=4, expansion_factor=-1: https://console.cloud.google.com/kubernetes/service/us-central2/v4-128-bodaborg-us-central2-b/default/aireen-20241010-18-28/details?project=cloud-tpu-multipod-dev
• per_device_batch_size=1.0 eval_per_device_batch_size=0.5, ici_tensor_parallelism=4, expansion_factor=-1: https://console.cloud.google.com/kubernetes/service/us-central2/v4-128-bodaborg-us-central2-b/default/aireen-20241010-19-02/details?project=cloud-tpu-multipod-dev
• per_device_batch_size=0.5 eval_per_device_batch_size=1, ici_tensor_parallelism=4, expansion_factor=-1: https://console.cloud.google.com/kubernetes/service/us-central2/v4-128-bodaborg-us-central2-b/default/aireen-20241010-18-36/details?project=cloud-tpu-multipod-dev

Thank you @aireenmei for those experiments!

Look great to me!

@ZhiyuLi-goog I reliazed _tfds_data_processing_c4_mlperf.py also needs update. Please review my latest push^.