How to Properly Resume Training w/ Respect to Schedulers?

Not sure if it's an error but I'm having the issue where resuming the model training using the --resume command seems to load everything up successfully, but the training steps are reset back to 0, as opposed to picking up from the last epoch it was trained on.

Ex: If I stop training after 2/10 epochs, or on the second state, and then resume training, it'll load up the state folder successfully but will start back at 0/10 epochs, which I would assume would result in 12 epochs being trained instead of 10(?)

I would assume the fix in this example would be to decrease the max epochs so that it's 8 instead of 10, resulting in 10 max epochs being trained instead of 12, but it seems like this would completely mess up the schedulers and timesteps. Is this intentional, or is there some other command I need to include or an existing command I need to modify alongside the --resume command?

Also, finetuning Flux just for reference.

Start by tracing the --resume path for Flux finetuning and how it restores training steps, epochs, schedulers, and timesteps. Reproduce the reported 2-of-10-epochs scenario, then verify whether resumed training preserves the prior progress and scheduler state without requiring manual changes to max epochs.

Written by the indexing model from the issue text.