[SPMD] Move some functions to IterGraphModule.setup() (#99076)
Since users will have to call these steps before calling `setup()`, moving these steps to `setup()` can reduce the API usage complexity.
Differential Revision: [D44973726](https://our.internmc.facebook.com/intern/diff/D44973726/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/99076
Approved by: https://github.com/lessw2020
