[Checkpoint] Separate implementation into generator (#105101)
Separates the non-reentrant AC implementation into a generator so that
other APIs such as composable checkpoint API can use the generator as pre and
post forward logic.
Differential Revision: [D47419387](https://our.internmc.facebook.com/intern/diff/D47419387/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/105101
Approved by: https://github.com/soulitzer
