Frees unused opt_states in pmap model initialization
Frees unused opt_states in pmap model initialization

If we decide to re-compute the opt_states after the model variables are updated, we can free it before re-computing to regain some memory, as sometimes the opt_states can be very large.

PiperOrigin-RevId: 482900021

google/paxml,google/paxml