Disable unsafe rng flag

[eltsai]

We are running into an OOM error:

jax.errors.JaxRuntimeError: INTERNAL: RET_CHECK failure (platforms/xla/service/jellyfish/bounds_check.cc:708) allocation_size_words <= std::numeric_limits<int32_t>::max() 4180377600 %s1 = inlined_call_operand.hbm [shape: u32[64,75600,13824], index: 1, kind: output, shape index: {}] /* operand 1 / %5539 = dma.vmem_to_hbm [thread:$0] (%p7303), /vmem=/%s5537, /size_in_granules=/%s7302, /hbm=/%s5532, /dst_syncflagno=/%s5515 /.

if we set jax_default_prng_impl as unsafe_rbg and use --xla_tpu_spmd_rng_bit_generator_unsafe=true flag, it breaks the sharding of the mlp block in the transformer and caused OOM error.

This PR brings back the reverted optimizations and disable the unsafe rng flag.

The run with the flag disabled run successfully: log

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid