Skip to content

Fix latent/caption mismatch in SD3 dreambooth lora trainer with cache_latents + shuffle - #14431

Open
lcheng321 wants to merge 1 commit into
huggingface:mainfrom
lcheng321:fix/sd3-dreambooth-latent-cache-index-mismatch
Open

Fix latent/caption mismatch in SD3 dreambooth lora trainer with cache_latents + shuffle#14431
lcheng321 wants to merge 1 commit into
huggingface:mainfrom
lcheng321:fix/sd3-dreambooth-latent-cache-index-mismatch

Conversation

@lcheng321

Copy link
Copy Markdown
Contributor

Fixes #14430

Root cause

Latents were cached keyed by dataloader step, but the dataloader has shuffle=True and is traversed twice: once during caching, once during training. Each traversal produces a different sample order. latents_cache[step] returns whatever image happened to land at that step during caching, not the image in the current training batch. No exception is raised, the mismatch is silent.

Fix

Cache latents keyed by stable dataset index instead of step. Index is generated in DreamBoothDataset.__getitem__ and passed through collate_fn.

Also added a check that rejects --cache_latents combined with --with_prior_preservation, since class images are not indexed the same way and would hit the same bug through a different path.

Scope

This PR only covers train_dreambooth_lora_sd3.py. It does not touch the Qwen-Image trainer or #12124, which is a separate file with a similar but independent issue.

Testing

1. Isolated repro, no model or GPU required

Dataset index used directly as the tensor value, so any mismatch between cached value and actual batch value is visible by eye. Same dataloader, same seed, two traversals.

Before fix, 4/4 batches mismatched:

step=0: latents_cache[0]=[2.0, 0.0] vs actual pixel_values=[0.0, 3.0] MISMATCH
step=1: latents_cache[1]=[5.0, 3.0] vs actual pixel_values=[7.0, 1.0] MISMATCH
step=2: latents_cache[2]=[1.0, 4.0] vs actual pixel_values=[6.0, 2.0] MISMATCH
step=3: latents_cache[3]=[7.0, 6.0] vs actual pixel_values=[5.0, 4.0] MISMATCH

After fix, 4/4 batches matched:

step=0: indices=[0, 3] cached=[0.0, 3.0] actual=[0.0, 3.0] OK
step=1: indices=[7, 1] cached=[7.0, 1.0] actual=[7.0, 1.0] OK
step=2: indices=[6, 2] cached=[6.0, 2.0] actual=[6.0, 2.0] OK
step=3: indices=[5, 4] cached=[5.0, 4.0] actual=[5.0, 4.0] OK
all_match: True

2. Existing test suite

pytest examples/dreambooth/test_dreambooth_lora_sd3.py examples/dreambooth/test_dreambooth_sd3.py -v

11 passed in 344.51s

All existing tests pass, including test_dreambooth_lora_latent_caching.

@github-actions github-actions Bot added size/S PR with diff < 50 LOC examples fixes-issue labels Aug 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[SD3 DreamBooth] Cached latents can be paired with the wrong per-image captions when shuffle=True

1 participant