Fix deterministic distributed bucket splitting

[hjinnkim]

Fix deterministic distributed bucket splitting

Summary

• Process-local RNG state could give ranks different permutations of the same bucket before splitting.
• When each rank slices a different permutation, distributed training can duplicate some samples and omit others in an epoch.
• Use a shared shuffle seed and a dataset/bucket-specific RNG so every rank slices the same bucket order.

Changes

• simpletuner/helpers/metadata/backends/base.py: reuse the configured seed, or generate one on the main process and broadcast it; shuffle a bucket copy with a dataset-and-bucket-specific RNG before rank slicing.
• simpletuner/helpers/metadata/backends/base.py: avoid a second full-list copy when the trim limit keeps the complete bucket.
• AGENTS.md: document that checkpoint resume is not assumed to support topology or dataloader-setting changes between save and resume.

Notes

• This intentionally does not restore saved sampler partitions or add compatibility handling for changed DDP/FSDP/CP topology, world size, batch sizing, gradient accumulation, repeats, bucket splitting, shuffling, or dataloader configuration. Those changes may produce unexpected behavior when resuming and need explicit support before being treated as safe.

Verification

/Users/kash/src/SimpleTuner/.venv/bin/python -m unittest tests.test_factory_edge_cases.TestFactoryEdgeCases.test_early_validation_sufficient_repeats tests.test_factory_edge_cases.TestFactoryEdgeCases.test_eval_dataset_ignores_gradient_accumulation tests.test_factory_edge_cases.TestFactoryEdgeCases.test_oversubscription_auto_adjustment -v

Also verified with a focused 8-rank bucket split probe: 1,024 assigned, 1,024 unique, 0 missing, 0 duplicated entries.

[Copilot]

Pull request overview

This PR fixes correctness issues in distributed multi-aspect bucket splitting and sampler resume by ensuring all ranks slice the same deterministic bucket permutation and by restoring sampler partitions only when the checkpoint’s recorded DP/CP topology matches the current run.

Changes:

• Make bucket shuffling deterministic across ranks by deriving a per-dataset/per-bucket RNG seed (broadcast from main when needed) and shuffling a bucket copy prior to rank slicing.
• Avoid an unnecessary full-list copy when the trim limit keeps the entire bucket.
• Persist and validate sampler partition topology on checkpoint save/load; reject legacy distributed partitions without topology metadata and reject topology mismatches before applying any progress.

Reviewed changes

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

File	Description
simpletuner/helpers/multiaspect/sampler.py	Save/validate partition topology and restore saved per-rank partitions only when topology matches.
simpletuner/helpers/metadata/backends/base.py	Deterministic, rank-consistent bucket shuffling + minor trimming optimization to reduce copying.

---

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

[bghira]

i removed the sampler restore changes from your pull request because they're unexpectedly complex and fragile. i added guidance to AGENTS.md so that future sessions are unlikely to go down the path of resolving bugs that centre around unsupported use like altering training hardware layout or config.json (in many cases, at least). this is because changing these things can often cause unexpected behaviour anyway. it's difficult to verify at every level of the stack that adjusting something mid-training won't cause a problem.