fix(data): reject a rank the epoch-rollover re-split left with no samples

[hjinnkim]

Summary

The bucket re-split that runs at every epoch rollover can leave a rank with no samples, and nothing
checks for it. Startup refuses the identical dataset state.

Problem

_handle_bucket_operations always ends by calling _handle_config_versioning (factory.py:3359),
which closes with a check that raises ValueError: Dataset produced no usable samples when the split
left this rank's backend with nothing usable (factory.py:3437-3476).

Trainer._epoch_rollover splits a second time (trainer.py:5536) and has no equivalent follow-up.
The only diagnostic on that path is the warning at the end of split_buckets_between_processes
(base.py:945-952), and it is gated on should_log() — global rank 0 — while the rank that ends up
empty is not rank 0.

Measured with 16 images across 8 ranks, train_batch_size=1, gradient_accumulation_steps=1,
repeats=1, --max_train_steps=100, oversubscription off:

path	dataset state	outcome
startup	buckets {1.0: 8, 1.5: 8}	all 8 ranks start
rollover	re-bucketed to {1.0: 4, 1.5: 4}	ranks 4-7 hold zero samples. No error and no warning, on any rank
startup	the same {1.0: 4, 1.5: 4}	ranks 4-7 fail with Dataset produced no usable samples

No data has to be deleted to get there. crop_aspect="random" is what puts a backend on this code
path in the first place, and re-bucketing the same 16 images is enough.

Reproduction condition. The pre-split validation inside split_buckets_between_processes
(base.py:787-870) does run at rollover, so what is missing is specifically a post-split check.
Staying silent requires repeats >= train_batch_size * gradient_accumulation_steps; with the default
repeats=0 every rank is caught earlier by Dataset configuration will produce zero usable batches.

What happens afterwards. Not nothing. With prefetch off the emptied rank raises from
MultiAspectSampler._reset_buckets (sampler.py:408) at the first fetch of the new epoch — an
unattributed No images found in the dataset from a rank that started fine, some distance from the
split that caused it. Reading the prefetch path suggests that with --dataloader_prefetch the same
exception is swallowed into a DEBUG record (batch_fetcher.py:83-85); I have not verified that
variant on hardware.

Resolution

Ten lines directly after the re-split in _epoch_rollover: if this rank's buckets hold no samples at
all, raise, naming the backend and the epoch it was about to enter.

Nothing else moves. factory.py is untouched, and the check does not import from or call into the
startup path, so the two guards stay independent — that is deliberate, and it is why the message is
written out here rather than shared. The only reason to couple them would be to keep one predicate,
and this PR does not want the startup predicate.

Raising rather than logging: the trainer already aborts mid-run by raising — a None batch
(trainer.py:6481) and an unrecognised gradient-clipping method (trainer.py:6694) raise
ValueError, a non-finite loss (trainer.py:6602) raises RuntimeError. A rank whose schedule is
empty cannot contribute a gradient for the epoch it is entering.

The message opens with Dataset produced no usable samples, the same phrase startup uses, so an
operator grepping logs or an existing alert finds both. The rest of it is rollover-specific: which
epoch, and the causes that actually apply here — per-epoch re-bucketing under crop_aspect=random,
or samples filtered out of the cache during the previous epoch.

The remedies it suggests are the ones that move a physical shard boundary.
--allow_dataset_oversubscription turns padding on (trainer.py:5538), which raises every non-empty
bucket's shard to ceil(len / dp_size) and so guarantees at least one sample per rank; fewer GPUs
lowers the divisor; more samples raises the dividend. It deliberately does not repeat startup's
advice to lower batch_size / gradient_accumulation_steps or raise repeats: none of those reach
divmod(len(trimmed_images), effective_dp_size). trim_limit is always >= len(images)
(base.py:903-906), so trimmed_images is the full bucket regardless of effective batch size, and
repeats only multiplies the sampler's logical passes over a shard that is already empty.

This guard is deliberately narrower than the startup one

It asks a literal question — did this rank end the split with zero samples — and nothing more.

The startup check uses len(metadata_backend) == 0, which counts buckets able to form a complete
local batch (discovery.py:767-780). Reusing that here would also abort a rank that still holds
samples but too few to fill one train_batch_size. Such a rank keeps training today on recycled
samples, and it keeps doing so after this PR. Aborting a run that is currently making progress is a
larger intervention than the defect requires, so the narrow predicate is the conservative choice:
the only state that newly aborts is the one that is already unusable.

The consequence is that this does not make the rollover fully agree with startup. A sub-batch shard
is still accepted here and refused there. That is a smaller divergence than the one being fixed, and
closing it means changing what startup considers fatal — a separate decision, not this PR's.

The startup check's missing_instance_dir escape hatch (factory.py:3437-3440) is not carried over.
It exists to tolerate a backend whose instance directory is not a local path — for type: aws it is
the S3 key prefix (factory.py:3032-3033), and os.path.exists() on that is always False. That is a
question about configuration, and it has already been answered by the time an epoch has finished
training. Whether the samples live on local disk or in a bucket has no bearing on whether this rank's
shard came back empty. Leaving it out also means the rollover guard actually covers remote backends,
where the startup one silently does not.

What this does not do

The guard is rank-local, exactly as startup's is: any rank that still has data proceeds. It is not
only the partial case — a re-bucketing that empties the dataset globally reaches this check on every
rank, because the pre-split validation skips empty buckets (base.py:790-791) and so does not catch
it first. Whether the ranks that survive a partial case then block on a collective is not something
I can demonstrate — the reproduction is an 8-rank
simulation inside a single process — so no claim is made about hangs or NCCL timeouts.

It also leaves the rollover's padding policy alone.
apply_padding=(self.config.overrode_max_train_steps or self.config.allow_dataset_oversubscription)
at trainer.py:5538 is a correct carry-forward of not args.max_train_steps, and
test_epoch_rollover_reuses_initial_bucket_padding_policy already pins it.

Tests

tests/test_trainer.py

• test_epoch_rollover_rejects_a_rank_the_resplit_left_empty — recorded RED on main
  (AssertionError: ValueError not raised). It also asserts the rejection lands after the re-split
  and before rebuild_cache(), so a bad schedule is never handed downstream.
• test_epoch_rollover_keeps_a_rank_that_cannot_fill_a_batch — the other direction. One sample
  against batch_size=4, so len(metadata_backend) reads 0 and the startup predicate would fire.
  This rank must keep going. Passes before and after, and fails if the predicate is ever widened.

test_epoch_rollover_reuses_initial_bucket_padding_policy gains one fixture line: its MagicMock
backend now carries an aspect_ratio_bucket_indices dict, because a bare MagicMock iterates empty
and would read as a zero-sample shard. Its assertion is unchanged.

[Copilot]

Pull request overview

This PR adds a post-split validation during Trainer._epoch_rollover to fail fast when the per-epoch re-bucketing + re-splitting leaves a given rank with zero samples, aligning rollover behavior more closely with startup-time dataset validation and avoiding downstream “empty schedule” failures.

Changes:

• Add a rank-local guard after split_buckets_between_processes() during epoch rollover to raise ValueError when the local shard contains zero samples.
• Extend unit tests to cover both directions: reject a truly empty shard, but allow shards that have samples even if they can’t form a full local batch.
• Adjust an existing epoch-rollover padding-policy test fixture to provide non-empty aspect_ratio_bucket_indices (so it doesn’t trip the new guard).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
simpletuner/helpers/training/trainer.py	Adds a post re-split “zero samples on this rank” failure guard during epoch rollover.
tests/test_trainer.py	Adds targeted unit tests for the new rollover guard and updates an existing test fixture to remain valid.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.