Tracing the HuggingFace Trainer Dataloader Path — From `trainer.train()` to a `next(epoch_iterator)`

All permalinks in this post point to huggingface/transformers tag v4.57.1 and huggingface/accelerate tag v1.13.0. Line numbers were verified against those tags.

Why this post

Trainer is the most-used entry point in modern Transformer fine-tuning — TRL’s SFTTrainer, DPOTrainer, and GRPOTrainer all subclass it. The default dataloader path “just works”, which means most people never look at it. Then you hit:

• “Why does my IterableDataset behave differently?”
• “Where does LengthGroupedSampler get plugged in — and why does it never run for me?”
• “What is num_items_in_batch and why does the loss change when I forward it?”
• “Why is one rank always slower than the others?”

Each answer lives in 5-10 lines of trainer.py. This post traces the path from trainer.train() down to a single next(epoch_iterator) call, with permalinks.

High-level call graph

trainer.train()
  └─ _inner_training_loop()                          [trainer.py:2353]
       ├─ train_dataloader = self.get_train_dataloader()
       │      └─ self._get_dataloader(...)            [trainer.py:1082]
       │           ├─ sampler = self._get_train_sampler(dataset)  [trainer.py:1053]
       │           ├─ DataLoader(dataset, sampler=..., collate_fn=self.data_collator, ...)
       │           └─ self.accelerator.prepare(...)
       │                  → wraps into DataLoaderShard / DataLoaderDispatcher
       │                                              [accelerate/data_loader.py:1006]
       ├─ epoch_dataloader = train_dataloader
       │  if hasattr(epoch_dataloader, "set_epoch"): epoch_dataloader.set_epoch(epoch)
       ├─ epoch_iterator = iter(epoch_dataloader)
       └─ for _ in range(total_updates):
              batch_samples, num_items_in_batch = self.get_batch_samples(...)
                                                    [trainer.py:5644]
              for inputs in batch_samples:
                  loss = self.training_step(model, inputs, num_items_in_batch)

That’s the spine. Now zoom in to each box.

1. get_train_dataloader — a thin wrapper

The public entry point is intentionally trivial — just delegates to _get_dataloader with the training-specific args:

trainer.py:1128

def get_train_dataloader(self) -> DataLoader:
    if self.train_dataset is None:
        raise ValueError("Trainer: training requires a train_dataset.")

    return self._get_dataloader(
        dataset=self.train_dataset,
        description="Training",
        batch_size=self._train_batch_size,
        sampler_fn=self._get_train_sampler,
        is_training=True,
    )

If you want custom dataloader behavior, override this method in a subclass — that’s the contract Trainer exposes. The body of _get_dataloader is where the real wiring happens.

2. _get_dataloader — the actual construction

trainer.py:1082

def _get_dataloader(
    self,
    dataset: Dataset,
    description: str,
    batch_size: int,
    sampler_fn: Optional[Callable[[Dataset], torch.utils.data.Sampler]] = None,
    is_training: bool = False,
    dataloader_key: Optional[str] = None,
) -> DataLoader:
    data_collator = self.data_collator
    if is_datasets_available() and isinstance(dataset, datasets.Dataset):
        dataset = self._remove_unused_columns(dataset, description=description)
    else:
        data_collator = self._get_collator_with_removed_columns(self.data_collator, description=description)

    dataloader_params = {
        "batch_size": batch_size,
        "collate_fn": data_collator,
        "num_workers": self.args.dataloader_num_workers,
        "pin_memory": self.args.dataloader_pin_memory,
        "persistent_workers": self.args.dataloader_persistent_workers,
    }

    if not isinstance(dataset, torch.utils.data.IterableDataset):
        if sampler_fn is not None:
            dataloader_params["sampler"] = sampler_fn(dataset)
        dataloader_params["drop_last"] = self.args.dataloader_drop_last
        dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
        if is_training:
            dataloader_params["worker_init_fn"] = partial(
                seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index
            )

    dataloader = self.accelerator.prepare(DataLoader(dataset, **dataloader_params))
    ...
    return dataloader

Three things worth pointing out:

2a. The isinstance(dataset, IterableDataset) branch

This is the single most important behavioral switch in the Trainer’s dataloader path. If your dataset is an IterableDataset:

• sampler is not set (no sampler_fn is called at all)
• drop_last is not set
• prefetch_factor is not set
• worker_init_fn is not set

The reasoning is that an IterableDataset knows how to shard itself (via get_worker_info() in its __iter__), so PyTorch’s sampler machinery is bypassed. That also means: group_by_length=True silently does nothing for an IterableDataset. No warning, no error — _get_train_sampler is just never called. If you’ve ever set group_by_length=True and seen no effect, this is why.

It also means prefetch_factor is not configurable for an IterableDataset via this code path — it falls back to PyTorch’s default of 2.

2b. _remove_unused_columns

For a HuggingFace datasets.Dataset, the trainer strips columns the model’s forward() doesn’t accept (trainer.py:1094). For everything else (raw torch.utils.data.Dataset, IterableDataset, custom subclasses), it instead wraps the collator to strip unknown keys — see _get_collator_with_removed_columns. So if you have an unexpected key making it through to the model, check whether you’re on the dataset path or the collator path.

2c. self.accelerator.prepare(...) — the seam between Trainer and Accelerate

The plain DataLoader gets handed to accelerator.prepare(...). This is where DDP-related transforms happen:

• For DDP, the dataset gets wrapped in a DataLoaderShard (each rank only sees its own shard).
• For dispatcher-mode multi-node, it becomes a DataLoaderDispatcher (one rank reads, scatters to others).
• For XLA/TPU, additional wrapping.

The wrapping logic lives in accelerate/data_loader.py:1006 (prepare_data_loader). The wrapper classes are DataLoaderShard and DataLoaderDispatcher.

A subtle consequence: when you call len(epoch_dataloader) inside _inner_training_loop, you might be calling __len__ on the wrapped dataloader, which can return a different value than the unwrapped one (e.g., len // world_size for DDP sharding).

3. _get_train_sampler — sampler selection

trainer.py:1053

def _get_train_sampler(self, train_dataset: Optional[Dataset] = None) -> Optional[torch.utils.data.Sampler]:
    if train_dataset is None:
        train_dataset = self.train_dataset
    if train_dataset is None or not has_length(train_dataset):
        return None

    # Build the sampler.
    if self.args.group_by_length:
        ...
        return LengthGroupedSampler(
            self.args.train_batch_size * self.args.gradient_accumulation_steps,
            dataset=train_dataset,
            lengths=lengths,
            model_input_name=model_input_name,
        )

    else:
        return RandomSampler(train_dataset)

The early-return on not has_length(dataset) is the second place IterableDataset support diverges — even if _get_train_sampler somehow got called for an IterableDataset, it would return None (no sampler). Some IterableDatasets implement __len__ (e.g., HuggingFace IterableDataset for finite splits), but most streaming ones don’t.

LengthGroupedSampler itself lives in trainer_pt_utils.py:615. Its __iter__ calls get_length_grouped_indices, which:

1. randomly permutes the full index list
2. chunks it into “mega-batches” of mega_batch_mult * batch_size (default mega_batch_mult=50)
3. sorts within each mega-batch by length descending
4. emits the indices flat, with the largest sequence placed in the very first batch (so OOM happens fast, not in epoch 3)

So it doesn’t globally sort by length — it does local sorting within mega-batches to preserve some randomness. Cute trade-off.

4. The collator — collate_fn decides batch shape

The trainer passes self.data_collator as the collate_fn to DataLoader. For text training, the collator decides:

• Whether to pad each sample to longest-in-batch (padded mode), or
• Whether to concatenate samples into one flat row with position_ids resetting per sample (padding-free mode, used by DataCollatorForLanguageModeling(padding_free=True))

TRL’s SFTTrainer instantiates DataCollatorForLanguageModeling and sets padding_free=True automatically if you set packing_strategy="bfd". See trl/trainer/sft_trainer.py:899 for the auto-coupling logic, and get_position_ids_from_packed_seq_lengths for how the resetting position_ids are constructed.

This is the seam where “packing” and “padding-free” become indistinguishable from the model’s perspective: whether you call the upstream concatenation “packing” or not, the collator produces a (1, sum_of_lengths) tensor with position_ids that resets, and FlashAttention-varlen derives cu_seqlens from those resets.

5. The training loop — _inner_training_loop

trainer.py:2353. The relevant slice:

for epoch in range(epochs_trained, num_train_epochs):
    epoch_dataloader = train_dataloader
    if hasattr(epoch_dataloader, "set_epoch"):
        epoch_dataloader.set_epoch(epoch)
    ...
    epoch_iterator = iter(epoch_dataloader)
    ...
    for _ in range(total_updates):
        update_step += 1
        num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder
        batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches, args.device)
        ...
        for i, inputs in enumerate(batch_samples):
            ...
            loss = self.training_step(model, inputs, num_items_in_batch)

Two things to notice:

5a. set_epoch is opportunistic

The trainer calls epoch_dataloader.set_epoch(epoch) only if the attribute exists. For DistributedSampler this is required to get fresh shuffling each epoch; for IterableDataset wrappers it depends on whether the underlying class implements it. If you have a custom IterableDataset that doesn’t expose set_epoch, your sampling stays static across epochs — a footgun.

5b. get_batch_samples prefetches gradient_accumulation_steps batches at once

This is the key insight people miss. trainer.py:5644:

def get_batch_samples(
    self, epoch_iterator: Iterator, num_batches: int, device: torch.device
) -> tuple[list, Optional[Union[torch.Tensor, int]]]:
    batch_samples = []
    for _ in range(num_batches):
        try:
            batch_samples.append(next(epoch_iterator))
        except StopIteration:
            break

    num_items_in_batch = self._get_num_items_in_batch(batch_samples, device)
    return batch_samples, num_items_in_batch

For each update step, the trainer pulls gradient_accumulation_steps batches eagerly, then computes num_items_in_batch across them, then loops over them for the forward+backward+accumulate-grad pass. This means num_items_in_batch reflects the summed completion-token count across the entire accumulation window, not just one batch.

6. num_items_in_batch — the loss-weighting correction nobody reads about

trainer.py:5593:

def _get_num_items_in_batch(self, batch_samples: list, device: torch.device) -> Optional[Union[torch.Tensor, int]]:
    num_items_in_batch = None
    count_num_items_in_batch = (
        len(batch_samples) > 0
        and "labels" in batch_samples[0]
        and (
            self.model_accepts_loss_kwargs
            or self.compute_loss_func is not None
        )
    )
    if count_num_items_in_batch:
        try:
            num_items_in_batch = sum((batch["labels"].ne(-100)).sum() for batch in batch_samples)
        except (TypeError, AttributeError):
            pass

    if num_items_in_batch is not None:
        if self.args.average_tokens_across_devices and self.args.world_size >= 1:
            num_items_in_batch = self.accelerator.gather(num_items_in_batch.to(device)).sum()
        ...
    return num_items_in_batch

Two gates: the model’s forward must accept **loss_kwargs (or you must supply compute_loss_func), and labels must be present. If those don’t hold, num_items_in_batch is None and the model uses its default per-rank token-mean loss — which under DDP gives the biased gradient I’ve written about elsewhere: ranks with fewer counted tokens contribute disproportionately per-token.

When num_items_in_batch is passed through, the model’s forward uses it to scale the loss so the DDP average produces the correct globally-token-weighted gradient. The average_tokens_across_devices flag controls whether to all-reduce the count across ranks (you want this on for unbiased DDP loss; default is off).

For modern HF model forwards (Llama, Qwen, etc.), **loss_kwargs is supported and the model internally does:

loss = nn.functional.cross_entropy(logits, labels, reduction="sum") / num_items_in_batch

instead of the default reduction="mean". The two are equivalent on a single rank but differ when averaged across DDP ranks with different token counts.

Custom compute_loss overrides need to forward num_items_in_batch to keep this correction working — see trainer.py:training_step for the call site. If you wrap or replace compute_loss, this is the kwarg to plumb through.

7. The IterableDataset case — what you don’t get for free

Summarizing what’s missing when your dataset is an IterableDataset (vs a map-style Dataset):

feature	map-style	IterableDataset
RandomSampler shuffle	✓	dataset’s own __iter__ decides
LengthGroupedSampler	✓ via group_by_length=True	silently disabled
DistributedSampler sharding	✓ via accelerator.prepare	dataset must shard itself in __iter__
prefetch_factor config	✓	PyTorch default (2), no Trainer knob
worker_init_fn (seed_worker)	✓	not set
drop_last config	✓	not set
__len__ for epoch budgeting	✓	falls back to args.max_steps

For streaming data (e.g., Avro from HDFS, Parquet from S3) this is usually fine — your dataset’s __iter__ already implements per-worker sharding via torch.utils.data.get_worker_info(), and you don’t want the trainer trying to add a sampler on top. But if you’re surprised that group_by_length is a no-op, this is why.

8. The “what actually happens at each step” cheat sheet

For someone reading a trainer.train() for the first time, here’s the simplified per-step flow:

At step N (post-warmup, normal training):

1. epoch_iterator    →  produces 1 batch
                        = (1, sum_of_lengths) for padding_free, or
                          (B, max_in_batch) for padded
                        Workers are reading the NEXT batch in parallel,
                        thanks to DataLoader prefetch.

2. get_batch_samples(epoch_iterator, num_batches=grad_accum_steps)
                     →  collects grad_accum_steps batches into a list
                     →  computes num_items_in_batch = total non-ignored
                        label tokens across all those batches
                     →  optionally all-reduces num_items_in_batch
                        if average_tokens_across_devices=True

3. for each batch in the accumulation list:
   a. inputs → model.forward(**inputs, **loss_kwargs)
      where loss_kwargs may include num_items_in_batch
   b. loss = (raw CE summed over tokens) / num_items_in_batch
      (or default mean if num_items_in_batch is None — biased under DDP)
   c. loss.backward()  → accumulates grad

4. After all grad_accum_steps batches:
   a. NCCL all-reduce of gradients (DDP)
   b. optimizer.step()
   c. optimizer.zero_grad()
   d. lr_scheduler.step()

Most of the “where does X get called?” questions reduce to one of these 4 boxes.

9. Useful permalinks index

For when you’re debugging at 2am:

Topic	File	Line
get_train_dataloader	trainer.py	1128
_get_dataloader (construction + accelerator.prepare)	trainer.py	1082
_get_train_sampler (LengthGrouped vs Random)	trainer.py	1053
_inner_training_loop (epoch loop body)	trainer.py	2353
get_batch_samples (grad-accum prefetch)	trainer.py	5644
_get_num_items_in_batch (DDP loss correction count)	trainer.py	5593
LengthGroupedSampler class	trainer_pt_utils.py	615
get_length_grouped_indices (mega-batch sort)	trainer_pt_utils.py	580
prepare_data_loader (accelerator-side wrap)	accelerate/data_loader.py	1006
DataLoaderShard (DDP per-rank wrap)	accelerate/data_loader.py	502
DataLoaderDispatcher (rank-0 reads + scatter)	accelerate/data_loader.py	714
skip_first_batches (used for checkpoint resume)	accelerate/data_loader.py	1385

10. Takeaways

1. The Trainer’s dataloader is just torch.utils.data.DataLoader + an Accelerate wrapper. No magic. If you can read _get_dataloader, you can read the whole path.
2. IterableDataset bypasses the sampler machinery. group_by_length, prefetch_factor, and worker_init_fn are no-ops. Your dataset’s __iter__ owns sharding and shuffling.
3. accelerator.prepare(...) is where DDP sharding gets injected. It returns a wrapped dataloader; len() may return a different value than the unwrapped one.
4. Gradient accumulation prefetches its full window of batches up-front in get_batch_samples, then computes num_items_in_batch across them all.
5. num_items_in_batch is the loss-weighting bias fix for DDP. If you write a custom compute_loss, forward this kwarg — otherwise your model falls back to per-rank token-mean and your gradient becomes biased toward whichever rank had fewer counted tokens that step.

If you have a custom Trainer subclass and any of the above surprises you when you read it, drop me a note — would love to hear the workload.