{"record":{"id":"2010036d7de86d21","repo":"huggingface/pytorch-image-models","slug":"rank-self-rank-number-of-indices-for-this-rank","errorCode":null,"errorMessage":"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.","messagePattern":"Rank (.+?): Number of indices for this rank \\((.+?)\\) does not match expected padded samples per rank \\((.+?)\\)\\. Epoch generation might be inconsistent\\.","errorType":"console","errorClass":null,"httpStatus":null,"severity":"warning","filePath":"timm/data/naflex_dataset.py","lineNumber":436,"sourceCode":"            padded_total_len = self._padded_samples_per_rank * self.world_size\n            if padded_total_len > total_len:\n                pad_size = padded_total_len - total_len\n                # Repeat initial elements from the *shuffled* list for padding\n                indices_for_ranks = all_indices_shuffled + all_indices_shuffled[:pad_size]\n            # Ensure length matches expectation\n            if len(indices_for_ranks) != padded_total_len:\n                 raise RuntimeError(f\"Internal Error: Padded index list length {len(indices_for_ranks)} does not match expected {padded_total_len}\")\n\n        # 3. Select indices for the current rank\n        if self.distributed and self.world_size > 1:\n            indices_this_rank = indices_for_ranks[self.rank::self.world_size]\n        else: # Non-distributed or world_size=1\n            indices_this_rank = indices_for_ranks\n\n        # Sanity check length\n        if len(indices_this_rank) != self._padded_samples_per_rank:\n             # This might happen if canonical schedule generation had warnings/issues\n             warnings.warn(\n                 f\"Rank {self.rank}: Number of indices for this rank ({len(indices_this_rank)}) \"\n                 f\"does not match expected padded samples per rank ({self._padded_samples_per_rank}). \"\n                 f\"Epoch generation might be inconsistent.\"\n              )\n             # Adjust expected samples? Or truncate/pad indices? Let's proceed but warn.\n             # Using min() prevents IndexError later if indices are fewer than expected.\n             effective_samples_this_rank = min(len(indices_this_rank), self._padded_samples_per_rank)\n             indices_this_rank = indices_this_rank[:effective_samples_this_rank]\n\n        else:\n             effective_samples_this_rank = self._padded_samples_per_rank\n\n        # 4. Shuffle the order of the canonical batch schedule for this epoch\n        if self.shuffle:\n            schedule_perm = torch.randperm(self._num_batches_per_rank, generator=g).tolist()\n            shuffled_schedule = [self._canonical_batch_schedule[i] for i in schedule_perm]\n        else:\n            shuffled_schedule = list(self._canonical_batch_schedule) # Keep original order","sourceCodeStart":418,"sourceCodeEnd":454,"githubUrl":"https://github.com/huggingface/pytorch-image-models/blob/9a5261e31b3b5128526eb2658333b4c0a54464ae/timm/data/naflex_dataset.py#L418-L454","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Fix the root cause: make schedule constraints evenly cover padded samples per rank (align dataset size with world_size and batch constraints)","Ensure rank/world_size passed to NaFlexDataset match the actual initialized process group","Pad the dataset to a multiple of world_size * per-rank schedule granularity"],"exampleFix":"# before\nNaFlexDataset(ds, ..., world_size=8)  # run launched with 4 GPUs\n# after\nws = torch.distributed.get_world_size()\nNaFlexDataset(ds, ..., rank=dist.get_rank(), world_size=ws)","handlingStrategy":"validation","validationCode":"import torch.distributed as dist\\nassert world_size == (dist.get_world_size() if dist.is_initialized() else 1)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive rank/world_size from the live process group, never hardcode","Keep dataset length fixed across epochs or re-instantiate the dataset"],"tags":["naflex","distributed","epoch-schedule","timm"],"backgroundTag":"uneven-data-sharding","analyzedSha":"9a5261e31b3b5128526eb2658333b4c0a54464ae","analyzedAt":"2026-08-27T02:34:25.417Z","schemaVersion":2},"datasetVersion":"2026-08-27T03:17:27.898Z"}