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.

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]
 

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.