From 287c1e5e8e6622c5e3858f10ad7514f7521c62b4 Mon Sep 17 00:00:00 2001 From: hjinnkim Date: Tue, 4 Aug 2026 15:18:29 +0900 Subject: [PATCH 1/2] fix(data): reject a rank the epoch-rollover re-split left with no samples --- simpletuner/helpers/training/trainer.py | 11 ++++ tests/test_trainer.py | 68 +++++++++++++++++++++++++ 2 files changed, 79 insertions(+) diff --git a/simpletuner/helpers/training/trainer.py b/simpletuner/helpers/training/trainer.py index 42d01b29c..46d7d5943 100644 --- a/simpletuner/helpers/training/trainer.py +++ b/simpletuner/helpers/training/trainer.py @@ -5537,6 +5537,17 @@ def _epoch_rollover(self, epoch): gradient_accumulation_steps=self.config.gradient_accumulation_steps, apply_padding=(self.config.overrode_max_train_steps or self.config.allow_dataset_oversubscription), ) + if sum(len(b) for b in backend["metadata_backend"].aspect_ratio_bucket_indices.values()) == 0: + raise ValueError( + f"(id={backend_id}) Dataset produced no usable samples. The epoch rollover" + f" re-split left this rank with zero samples, so epoch {epoch} would train" + f" against an empty schedule here.\n" + f"This usually means the per-epoch re-bucketing (crop_aspect=random) produced" + f" buckets too small to divide across the data-parallel ranks, or that samples" + f" were filtered out of the cache during the previous epoch.\n" + f"Enable --allow_dataset_oversubscription so short buckets are padded across" + f" ranks, use fewer GPUs, or add more samples to the dataset." + ) # we have to rebuild the VAE cache if it exists. if "vaecache" in backend: logger.info("Rebuilding VAE cache..") diff --git a/tests/test_trainer.py b/tests/test_trainer.py index b9aea9736..221e87d7f 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -10,6 +10,7 @@ import time import types import unittest +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch @@ -1862,6 +1863,7 @@ def test_epoch_rollover_reuses_initial_bucket_padding_policy(self): allow_dataset_oversubscription=allow_oversubscription, ) metadata_backend = MagicMock(read_only=True) + metadata_backend.aspect_ratio_bucket_indices = {"1.0": ["image-0.jpg"]} backends = {"train": {"metadata_backend": metadata_backend}} backend_config = { "crop": True, @@ -1886,6 +1888,72 @@ def test_epoch_rollover_reuses_initial_bucket_padding_policy(self): apply_padding=expected, ) + def _rollover_fixture(self, bucket_contents, usable_batches): + trainer = object.__new__(Trainer) + trainer.state = {"first_epoch": 1, "current_epoch": 1} + trainer.accelerator = MagicMock(is_main_process=True) + trainer.extra_lr_scheduler_kwargs = {} + trainer.get_steps_per_epoch_for_epoch = MagicMock(return_value=100) + trainer.config = SimpleNamespace( + num_train_epochs=5, + aspect_bucket_disable_rebuild=False, + lr_scheduler="constant", + num_update_steps_per_epoch=100, + gradient_accumulation_steps=1, + overrode_max_train_steps=False, + allow_dataset_oversubscription=False, + ) + metadata_backend = MagicMock(read_only=True, batch_size=4) + metadata_backend.aspect_ratio_bucket_indices = bucket_contents + metadata_backend.__len__.return_value = usable_batches + vaecache = MagicMock() + backends = {"train": {"metadata_backend": metadata_backend, "vaecache": vaecache}} + return trainer, metadata_backend, vaecache, backends + + @contextmanager + def _rollover_patches(self, backends): + with ( + patch("simpletuner.helpers.training.trainer.StateTracker.set_epoch"), + patch( + "simpletuner.helpers.training.trainer.StateTracker.get_data_backends", + return_value=backends, + ), + patch( + "simpletuner.helpers.training.trainer.StateTracker.get_data_backend_config", + return_value={"crop": True, "crop_aspect": "random"}, + ), + ): + yield + + def test_epoch_rollover_rejects_a_rank_the_resplit_left_empty(self): + trainer, metadata_backend, vaecache, backends = self._rollover_fixture( + bucket_contents={"1.0": [], "1.5": []}, usable_batches=0 + ) + + with self._rollover_patches(backends): + with self.assertRaises(ValueError) as context: + trainer._epoch_rollover(2) + + self.assertIn("Dataset produced no usable samples", str(context.exception)) + # The rejection must land after the re-split and before anything downstream + # consumes the new schedule. + metadata_backend.split_buckets_between_processes.assert_called_once() + vaecache.rebuild_cache.assert_not_called() + + def test_epoch_rollover_keeps_a_rank_that_cannot_fill_a_batch(self): + # One sample against batch_size=4, so the startup guard's len() would read 0 here. + # This guard asks only whether the shard is empty, and this rank keeps training on + # recycled samples exactly as it does today. + trainer, metadata_backend, vaecache, backends = self._rollover_fixture( + bucket_contents={"1.0": ["image-0.jpg"]}, usable_batches=0 + ) + + with self._rollover_patches(backends): + trainer._epoch_rollover(2) + + metadata_backend.split_buckets_between_processes.assert_called_once() + vaecache.rebuild_cache.assert_called_once() + @patch( "simpletuner.helpers.training.trainer.Trainer.parse_arguments", return_value=Mock(), From 67900276ce77105e21cb50ad3343643bc19d9386 Mon Sep 17 00:00:00 2001 From: bagheera <59658056+bghira@users.noreply.github.com> Date: Tue, 4 Aug 2026 08:39:12 -0600 Subject: [PATCH 2/2] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- simpletuner/helpers/training/trainer.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/simpletuner/helpers/training/trainer.py b/simpletuner/helpers/training/trainer.py index 46d7d5943..37e74a088 100644 --- a/simpletuner/helpers/training/trainer.py +++ b/simpletuner/helpers/training/trainer.py @@ -5537,12 +5537,15 @@ def _epoch_rollover(self, epoch): gradient_accumulation_steps=self.config.gradient_accumulation_steps, apply_padding=(self.config.overrode_max_train_steps or self.config.allow_dataset_oversubscription), ) - if sum(len(b) for b in backend["metadata_backend"].aspect_ratio_bucket_indices.values()) == 0: + local_sample_count = sum( + len(bucket) for bucket in backend["metadata_backend"].aspect_ratio_bucket_indices.values() + ) + if local_sample_count == 0: raise ValueError( f"(id={backend_id}) Dataset produced no usable samples. The epoch rollover" f" re-split left this rank with zero samples, so epoch {epoch} would train" f" against an empty schedule here.\n" - f"This usually means the per-epoch re-bucketing (crop_aspect=random) produced" + f"This usually means per-epoch re-bucketing (e.g., crop_aspect=random) produced" f" buckets too small to divide across the data-parallel ranks, or that samples" f" were filtered out of the cache during the previous epoch.\n" f"Enable --allow_dataset_oversubscription so short buckets are padded across"