From d2e6192eb2cd75b103eea2108735982315cf2c47 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 01:50:23 +0000 Subject: [PATCH] Honor the configured SDXL LoRA batch size on datasets smaller than the batch The SDXL trainer drew min(train_batch_size, len(pairs)) indices, so a dataset with fewer images than the batch trained at a smaller effective batch than configured while the scheduler and samples-per-second still assumed the full batch. The shared PermutationBatchSampler already refills across permutation cycles to return exactly k indices, and the DiT trainer calls it with the full batch size, so drop the clamp and pass train_batch_size through for parity and to honor the configured batch. --- .../core/training/diffusion_lora_trainer.py | 8 +++++++- .../backend/tests/test_diffusion_training.py | 18 ++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 879b88860f..b4349a0b61 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -486,7 +486,13 @@ def run_diffusion_lora_training( index_sampler = PermutationBatchSampler(len(pairs), rng) def _next_batch() -> tuple[list[int], list[str], list[str]]: - idx = index_sampler.next_batch(min(cfg.train_batch_size, len(pairs))) + # Draw the full configured batch, not min(batch, n): PermutationBatchSampler refills + # across permutation cycles so a dataset smaller than train_batch_size still yields + # exactly train_batch_size indices (the DiT trainer calls next_batch the same way). + # Clamping to len(pairs) would silently train a tiny dataset at a smaller effective + # batch than configured while the scheduler and samples-per-second still assume the + # full batch. + idx = index_sampler.next_batch(cfg.train_batch_size) chosen = [pairs[i] for i in idx] return idx, [c[0] for c in chosen], [c[1] for c in chosen] diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index c6b817470c..fb947e6e14 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -617,6 +617,24 @@ def test_permutation_sampler_covers_dataset_once_per_cycle(): assert set(batch) == {0, 1, 2, 3} +def test_permutation_sampler_honors_batch_on_tiny_dataset(): + # Regression for the SDXL LoRA trainer clamp: a dataset smaller than train_batch_size must + # still yield exactly train_batch_size indices (refilled across cycles), so a tiny dataset + # trains at the configured/reported effective batch instead of silently shrinking it -- the + # contract the SDXL _next_batch now relies on (matching the DiT trainer, which never clamps). + import random + + from core.training.diffusion_train_common import PermutationBatchSampler + + sampler = PermutationBatchSampler(2, random.Random(0)) # 2-image dataset + batch = sampler.next_batch(8) # train_batch_size = 8, not clamped to 2 + assert len(batch) == 8 + assert set(batch) == {0, 1} + # A whole-multiple batch draws each image equally, so the effective gradient matches the + # configured batch rather than a shrunk one. + assert batch.count(0) == 4 and batch.count(1) == 4 + + def test_route_start_accepts_zero_max_grad_norm(client): # 0 is the documented "disable clipping" value (the trainer skips clip_grad_norm_); # the request model must not reject it.