Fix per-rank sampler state resume path

[hjinnkim]

Fix per-rank sampler state resume path

Summary

• Resume rebuilt the per-rank sampler state filename independently from the save hook.
• Nonzero ranks saved through SaveHookManager.training_state_path, e.g. training_state-rank1.json, but sampler resume looked for training_state-1.json before backend state path mangling.
• Load sampler state from the hook-owned self.model_hooks.training_state_path instead.

Changes

• simpletuner/helpers/training/trainer.py: use SaveHookManager.training_state_path when loading sampler state on checkpoint resume.
• tests/test_trainer.py: add a regression test for nonzero-rank sampler resume loading training_state-rank1.json.

Notes

• This intentionally leaves completed-partition sampler rollover behavior unchanged. A fully consumed sampler state can be the epoch-end signal used to coordinate ranks before the next training step.
• Model, optimizer, and scheduler resume behavior is unchanged.

Verification

/Users/kash/src/SimpleTuner/.venv/bin/python -m unittest tests.test_trainer.TestTrainer.test_init_resume_checkpoint tests.test_trainer.TestTrainer.test_init_resume_checkpoint_loads_sampler_from_hook_training_state_path -v

[Copilot]

Pull request overview

This PR fixes sampler resume behavior in distributed training by (1) ensuring each rank loads the exact per-rank training state filename selected by SaveHookManager, and (2) preventing a fully-consumed per-rank sampler partition from resuming as an immediately-exhausted epoch.

Changes:

• Use self.model_hooks.training_state_path when loading sampler state on resume, avoiding duplicated/incorrect per-rank filename construction.
• When restoring seen_images, detect a completed local partition and reset buckets without raising a live exhaustion signal (while preserving partial progress unchanged).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
simpletuner/helpers/training/trainer.py	Align sampler resume state path with SaveHookManager.training_state_path (per-rank filename consistency).
simpletuner/helpers/multiaspect/sampler.py	Reset completed epoch state on resume to avoid immediate exhaustion for fully-consumed partitions.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[bghira]

i split the filename fix from the sampler change and removed the sampler part. i'd like for that to be submitted separately with some demonstration that it doesn't cause a barrier wait condition. the thing is, the MultiDatasetExhausted signal is really the final say on the sampler resetting its buckets, and when running ddp, an all-reduce is done across nodes/GPUs to determine if any rank had an exhaustion signal. if they do, then they all roll to the next epoch so that things are kept in sync. can you please explore that for the next time the sampler change is submitted? thank you for your pull requests and improvements you've provided so far.