Fix BucketBatchSampler cache alignment in DreamBooth scripts

[azolotenkov]

What does this PR do?

This PR fixes a bug where BucketBatchSampler reshuffled precomputed batches on each __iter__() call.

DreamBooth training scripts precompute latents and/or prompt embeddings and later consume them by dataloader step index, so changing batch order between the caching pass and the training pass can misalign the cached tensors with the current batch.

This PR removes the per-iteration random.shuffle(self.batches) call from BucketBatchSampler.__iter__() and instead shuffles the precomputed batch list once at sampler construction time, keeping the batch order fixed across epochs.

A future rework may still be needed to support true epoch-wise reshuffling, since batch membership is currently fixed once at sampler construction time.

Applied to:

• examples/dreambooth/train_dreambooth_lora_flux2.py
• examples/dreambooth/train_dreambooth_lora_flux2_img2img.py
• examples/dreambooth/train_dreambooth_lora_flux2_klein.py
• examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py
• examples/dreambooth/train_dreambooth_lora_z_image.py

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Training examples: @sayakpaul

[Copilot]

Pull request overview

Fixes cache misalignment in DreamBooth training examples by making BucketBatchSampler yield batches in a stable, precomputed order across __iter__() calls, preventing reshuffles that break step-indexed latent/prompt-embedding caches.

Changes:

• Removed random.shuffle(self.batches) from BucketBatchSampler.__iter__() in multiple DreamBooth example scripts.
• Added an explanatory comment describing why batch order must remain stable for cache alignment.

Reviewed changes

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

Show a summary per file

File	Description
examples/dreambooth/train_dreambooth_lora_flux2.py	Keep precomputed batch order stable in BucketBatchSampler.__iter__() to avoid cache misalignment.
examples/dreambooth/train_dreambooth_lora_flux2_img2img.py	Same sampler iteration stabilization for cache alignment.
examples/dreambooth/train_dreambooth_lora_flux2_klein.py	Same sampler iteration stabilization for cache alignment.
examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py	Same sampler iteration stabilization for cache alignment.
examples/dreambooth/train_dreambooth_lora_z_image.py	Same sampler iteration stabilization for cache alignment.

---

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

[azolotenkov]

Updated based on review: the sampler now shuffles precomputed batches once at construction time, while keeping iteration order fixed across epochs for cache alignment.

[linoytsaban]

good catch @azolotenkov!
I think it would make more sense if we apply this fix only when --cache_latents is enabled, since in other cases it's not issue. In light of this bug, ideally it would be best to support per-epoch reshuffle and latent caching but we can leave that for a follow up PR imo.
@sayakpaul wdyt?

[sayakpaul]

@linoytsaban yes, that would be the most appropriate!

[azolotenkov]

Updated. I kept the per-epoch reshuffle in the normal path, and only switch to a fixed precomputed batch order when the script relies on step-indexed caches. I scoped that to args.cache_latents or train_dataset.custom_instance_prompts, since custom prompt embedding caching has the same alignment requirement here. I also pulled that condition into a separate has_step_indexed_caches variable since it’s used both for the sampler behavior and for the caching path.
@linoytsaban wdyt?