[Epic][Improvement] Setup and training loop reorganization

[JamesKunstle]

The main training entrypoint, setup functions, and training loop are wound fairly tightly, with functions doing multiple tasks at once. We should refactor to unwind these functions.