[train][release] Attach a quick checkpoint when reporting metrics

[liulehui]

1. In Train V2, Free-floating metrics without corresponding checkpoint are no longer automatically saved. See context in https://docs.ray.io/en/master/train/user-guides/monitoring-logging.html#deprecated-reporting-free-floating-metrics.
2. By creating a quick checkpoint, we can leverage Train V2 by adding RAY_TRAIN_V2_ENABLED=1.
3. Since in V2, we have a 2s polling interval for train.report() and in this test we only care about final loss, move the metrics reporting part out of epoch loop, only report once at the end of training loop.
4. example run: https://buildkite.com/ray-project/release/builds/59517

Why are these changes needed?

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[gemini-code-assist]

Code Review

This pull request updates the torch benchmark to be compatible with Ray Train V2 by attaching a checkpoint when reporting metrics. However, the current implementation has a critical issue where train.report() is only called by the rank 0 worker. In a distributed setting, all workers must call train.report(), otherwise the job will likely hang. I've provided a suggestion to fix this by moving the train.report() call out of the rank-check conditional block, ensuring all workers participate in reporting.

[justinvyu]

Thanks!