Use snapshot to prevent warmup from affecting training and refactor warmup

[michaelmckinsey1]

• Fixes issue where warmup was affecting training by using snapshot to restore model state to that before warmup.
• Fixes issue where validation was not being warmed up properly.
• Refactor shared logic

With this configuration, runtime of the first epoch should now be similar to subsequent epochs.

[PatrickRMiles]

The snapshot is a cool approach, but I wonder if we can do this simpler by just creating another copy of the model and using that separate copy for warmup. We'd create two separate models the same way in worker.py, warmup on one, then train on the other. This would achieve what we want -- preventing warmup from impacting the model state before training -- without needing this snapshot logic. Worth testing at least before we merge this

[michaelmckinsey1]

The snapshot is a cool approach, but I wonder if we can do this simpler by just creating another copy of the model and using that separate copy for warmup. We'd create two separate models the same way in worker.py, warmup on one, then train on the other. This would achieve what we want -- preventing warmup from impacting the model state before training -- without needing this snapshot logic. Worth testing at least before we merge this

I did test that copying the model using copy.deepcopy would address the issue with model state. However, the other components like the optimizer, scheduler, and checkpoint manager would not be reverted. Additionally, I'm not sure what the impact would be for creating a new model object with a different address, as I think the pytorch components are tied to the original model object.

[PatrickRMiles]

We've decided to keep the snapshot approach. This looks good!