[SPMD] Add optimizer states and steps to the return (#98579)
This will correctly functionalize the optimizer. Otherwise, there are orphand copy_.
Differential Revision: [D44761512](https://our.internmc.facebook.com/intern/diff/D44761512/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/98579
Approved by: https://github.com/mrshenli
