nextrng not checkpointed, consider using fold_in(config.seed, step)

[reinerp]

Looks like the nextrng value is not saved to the checkpoint:

maxtext/MaxText/train.py

Line 285 in bab4ee7

state, metrics, nextrng = p_train_step(

It doesn't matter in most cases since most people don't use dropout or stochastic rounding. But in cases where it does matter, it's cleaner to generate the RNG for a training step using jax.random.fold_in(config.seed, state.step). This way, no checkpointing is required, and there's also some other side advantages listed in https://twitter.com/cgarciae88/status/1615022554992738315.

[rwitten]

We will follow the suggestion here.

[rwitten]

done!