ERROR: Accessing retired flag 'jax_enable_async_collective_offload'

[LeoXinhaoLee]

Hi, thank you so much for releasing this wonderful codebase. When I'm trying to run pretrain_llama_7b on some v3-tpu pod, I got this error:

ERROR: Accessing retired flag 'jax_enable_async_collective_offload'

It seems related to the flag specified before launching the job:

export LIBTPU_INIT_ARGS='--xla_jf_spmd_threshold_for_windowed_einsum_mib=0 \
--xla_tpu_spmd_threshold_for_allgather_cse=10000 \
--xla_tpu_spmd_rewrite_einsum_with_reshape=true \
--xla_enable_async_all_gather=true \
--jax_enable_async_collective_offload=true \
--xla_tpu_enable_latency_hiding_scheduler=true TPU_MEGACORE=MEGACORE_DENSE'

I am wondering if these flags are necessary and if some could cause the error? Thank you very much for your time and help!

[s-smits]

Same problem here.

[sharjeel103]

Did you find the solution. I have the same error