Move from an epoch-based optimization to a step-based optimization

[jpata]

This PR moves the training from an epoch-based formulation to a step-based formulation. The advantage is that very long trainings can be continued directly where they left off, as soon as they finish, rather than having to restart at the epoch boundary.

Technically, this is done by keeping track of the dataset and sampler state across multiple workers.

As validation, I trained a model (light blue) for a very long time (500k steps, ~20 days) on 8x AMD MI200X cards. It can be compared to the previous model freeze at #437 or https://huggingface.co/jpata/particleflow/tree/main/cms/v2.6.0pre1/pyg-cms_20250722_101813_274478 (dark blue). The cms_pf_qcd_nopu dataset size was also increased in this run from 5M events to 20M events. Overall, this corresponds to about 10 epochs over the full set of concatenated data, but the epochs between these two runs are not exactly the same. We see that the jet resolution metric (IQR) and the jet reco-gen matching fraction metric keep improving throughout the training.

The training loss is stable across ~20x 24-hour restarts. The regression loss fluctuates somewhat, and potentially a more stable solution could be better for optimization.

The validation loss is computed over a subset of the data (to avoid wasting a lot of time not training) at a regular interval and shows convergence.