aireenmei
commented
3 months ago - Switch to write_metrics (logging with buffer)
- add log_period. Usage suggestiong: logging will always trigger sync in jax, and break the overlapping of data input and computation. We can set log_period to a small value for debugging, but for submission, we should set to a big value, or even turn it off, only outputs at the end.
- cloud log for training 5000 steps with log_period 100: https://cloudlogging.app.goo.gl/i3a8d4LfTvzrsWqh7
- xprof: https://xprof.corp.google.com/trace_viewer/aireenmei-15432370690143073139?hosts=gke-tpu-fc96cd7e-vn72&host_index=0