[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.

[cdoern]

seems related to #433

I'd like to work on some of this if possible

[cdoern]

assigning to myself for now

[github-actions]

This issue has been automatically marked as stale because it has not had activity within 90 days. It will be automatically closed if no further activity occurs within 30 days.