huggingface / huggingface/diffusers

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

Open
#14,430 1 comment 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

### Describe the bug

The SD3 DreamBooth LoRA trainer creates the training dataloader with shuffle enabled.
When latent caching is enabled, the script traverses the shuffled dataloader once and stores the VAE latent distributions in a list. Training then starts another traversal of the same shuffled dataloader, but retrieves cached latents using only the current batch position.
Because the second traversal can produce a different sample order, a per-image caption in the current batch may be paired with the cached latent of another image. This problem is mostly hidden when every image uses the same instance prompt, but it can silently corrupt training when the caption column contains a different caption for each image.
Cached latents should be stored and retrieved using a stable dataset index or sample ID instead of the dataloader step. This would preserve shuffling while keeping each image, latent and caption correctly associated.
A similar correspondence issue was reported for the Qwen-Image trainer in issue #12124.

### Reproduction

The issue follows directly from the current data flow in train_dreambooth_lora_sd3.py.
The training dataloader is created with shuffle enabled. When latent caching is enabled, the first traversal appends each latent distribution to a list. Training starts a second traversal of the shuffled dataloader and reads the cached latent using latents_cache[step].
The first and second traversals are not guaranteed to have the same sample order. Therefore, the current per-image caption and latents_cache[step] are not guaranteed to belong to the same dataset sample.
This can be observed by printing the sample index during the latent caching traversal and again during the first training epoch.

### Logs

```shell
No runtime exception is raised. This is a silent image-caption correspondence issue that can affect the learned result without producing an error message.
```

### System Info

Diffusers: 0.40.0.dev0
Diffusers source commit: d6726f38a0c5ca6c06a8f227fb7bade3486ed98d
Python: 3.12
PyTorch: 2.8.0
CUDA: 12.8
Platform: Ubuntu 22.04
GPU: NVIDIA GeForce RTX 4090, 24GB
Training script: examples/dreambooth/train_dreambooth_lora_sd3.py
Model: stabilityai/stable-diffusion-3.5-medium
Dataset mode: local ImageFolder with per-image captions
Latent caching: enabled
Dataloader shuffling: enabled

### Who can help?

@sayakpaul @DN6

Contributor guide

Open the contributing guide

Research direction

Start in examples/dreambooth/train_dreambooth_lora_sd3.py and trace the latent-caching traversal alongside the first training epoch, focusing on how latents_cache[step] is matched to each sample. Reproduce with shuffled ImageFolder data and distinct per-image captions; done means cached latents remain paired with their original samples and captions across both traversals.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.