huggingface/pytorch-image-models · warning

Rank {self.rank}: Number of indices for this rank ({len(indi

Error message

Rank {self.rank}: Number of indices for this rank ({len(indices_this_rank)}) does not match expected padded samples per rank ({self._padded_samples_per_rank}). Epoch generation might be inconsistent.

What it means

NaFlexDataset._prepare_epoch_batches verifies that the number of indices assigned to this rank equals _padded_samples_per_rank. A mismatch means the canonical schedule didn't schedule exactly padded count — typically because schedule generation hit the constraint warnings (266/267) — and downstream code will proceed with min() clamping, possibly yielding fewer batches than the DataLoader expects.

Source

Thrown at timm/data/naflex_dataset.py:436

            padded_total_len = self._padded_samples_per_rank * self.world_size
            if padded_total_len > total_len:
                pad_size = padded_total_len - total_len
                # Repeat initial elements from the *shuffled* list for padding
                indices_for_ranks = all_indices_shuffled + all_indices_shuffled[:pad_size]
            # Ensure length matches expectation
            if len(indices_for_ranks) != padded_total_len:
                 raise RuntimeError(f"Internal Error: Padded index list length {len(indices_for_ranks)} does not match expected {padded_total_len}")

        # 3. Select indices for the current rank
        if self.distributed and self.world_size > 1:
            indices_this_rank = indices_for_ranks[self.rank::self.world_size]
        else: # Non-distributed or world_size=1
            indices_this_rank = indices_for_ranks

        # Sanity check length
        if len(indices_this_rank) != self._padded_samples_per_rank:
             # This might happen if canonical schedule generation had warnings/issues
             warnings.warn(
                 f"Rank {self.rank}: Number of indices for this rank ({len(indices_this_rank)}) "
                 f"does not match expected padded samples per rank ({self._padded_samples_per_rank}). "
                 f"Epoch generation might be inconsistent."
              )
             # Adjust expected samples? Or truncate/pad indices? Let's proceed but warn.
             # Using min() prevents IndexError later if indices are fewer than expected.
             effective_samples_this_rank = min(len(indices_this_rank), self._padded_samples_per_rank)
             indices_this_rank = indices_this_rank[:effective_samples_this_rank]

        else:
             effective_samples_this_rank = self._padded_samples_per_rank

        # 4. Shuffle the order of the canonical batch schedule for this epoch
        if self.shuffle:
            schedule_perm = torch.randperm(self._num_batches_per_rank, generator=g).tolist()
            shuffled_schedule = [self._canonical_batch_schedule[i] for i in schedule_perm]
        else:
            shuffled_schedule = list(self._canonical_batch_schedule) # Keep original order

View on GitHub (pinned to 9a5261e31b)

Solutions

1. Fix the root cause: make schedule constraints evenly cover padded samples per rank (align dataset size with world_size and batch constraints)
2. Ensure rank/world_size passed to NaFlexDataset match the actual initialized process group
3. Pad the dataset to a multiple of world_size * per-rank schedule granularity

Example fix

# before
NaFlexDataset(ds, ..., world_size=8)  # run launched with 4 GPUs
# after
ws = torch.distributed.get_world_size()
NaFlexDataset(ds, ..., rank=dist.get_rank(), world_size=ws)

import torch.distributed as dist\nassert world_size == (dist.get_world_size() if dist.is_initialized() else 1)

When it happens

Trigger: Distributed training where the schedule leaves samples unscheduled (see warning 267), causing len(indices_this_rank) != padded count; rank/world_size misconfiguration between dataset creation and process group.

Common situations: Uneven dataset size across ranks; changing world size between runs while reusing cached schedules. Also observed in the persistent-worker shared-epoch test path, where it signals schedule/indices drift.

Related errors

• Rank {self.rank}: Canonical schedule accounts for {total_sch
• Rank {self.rank}: Ran out of samples ({idx_pos}/{effective_s
• Rank {self.rank}: Assigned {scheduled_samples_count} samples
• Error processing sample index {idx}. Error: {e}. Skipping sa
• Calculated batch size <= 0 (seq_len={seq_len}, remaining={re

AI-assisted analysis of huggingface/pytorch-image-models@9a5261e31b (2026-08-27). Data as JSON: /api/errors/2010036d7de86d21. Report an issue: GitHub.