Fix num_items_in_batch over-counting for causal LM losses

[qgallouedec]

The bug

Trainer._get_num_items_in_batch counts labels at every position. But for causal LM the loss shifts labels — position 0 is never a prediction target. The denominator is too large by num_rows per micro-batch, systematically under-scaling every causal LM loss and gradient.

labels (trainer counts these):           shift_labels = labels[..., 1:]
                                         (what ForCausalLMLoss uses)
┌────┬────┬────┬────┬────┐                    ┌────┬────┬────┬────┐
│ t0 │ t1 │ t2 │ t3 │ t4 │  count = 5         │ t1 │ t2 │ t3 │ t4 │  count = 4
└────┴────┴────┴────┴────┘                    └────┴────┴────┴────┘
  ↑
  position 0 is dropped by the shift — over-count = 1 per row

shift_labels = labels[..., 1:]           # ← numerator: 4 CE terms
loss = sum_ce / num_items_in_batch       # ← denominator: counted 5

loss_type is set on every PreTrainedModel (see modeling_utils.py). Only causal LM loss types are touched; ForMaskedLM, classification, etc. are unaffected.

Reported causal LM loss becomes slightly larger (correctly so, the denominator was too big). Shift is num_rows / total_tokens per step. Gradient magnitudes scale by the same factor.

Tests

All TrainerGradientAccumulationTest tests pass (4/4). A padded vs padding-free reproducer now matches with Δgrad_norm = 0.

How it surfaced

TRL's invariance suite compared SFT with padding_free=False vs padding_free=True. Loss curves almost matched, but grad_norm drifted with a systematic +0.17 bias over 50 steps. The padding-free collator masks labels[position_ids == 0] = -100, which incidentally matches the post-shift count, so it exposed the padded path's over-count.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[vasqu]

I think this is too simple. I think if we were to go that way we should inspect the actual loss function, e.g. CsmForConditionalGeneration would fail here, no?

But, my bigger issue is that this covers the most simple usecase where we only prepare the tokenized input and pass that but what if we were to use a data collator that properly prepares the shifted labels? We would now count 1 too much

But definitely needed to fix in general!

[vasqu]

Also definitely need a test that covers this edge case

[SunMarc]

Thanks, left some minor comments but happy to merge it in general. Can you add a small test to check that we are calculating correctly the num_item_per_batch in case we have ForCausalLM ?

[SunMarc]

as @vasqu pointed, maybe we can have the following so that it is a bit more robust.

from transformers.loss.loss_utils import LOSS_MAPPING, ForCausalLMLoss
    self._loss_shifts_labels = (
        LOSS_MAPPING.get(getattr(model_to_inspect, "loss_type", None)) is ForCausalLMLoss
    )

[SunMarc]

Maybe to fix @vasqu point, we can also take into account the case where shift_labels is prepared by the user ?

  labels_for_count = [
      batch["shift_labels"] if "shift_labels" in batch
      else batch["labels"][..., 1:] if self._loss_shifts_labels
      else batch["labels"]
      for batch in batch_samples
  ]

[SunMarc]

Thanks !

[vasqu]

Seems like we have 1 new failure https://app.circleci.com/pipelines/github/huggingface/transformers/175911/workflows/b1032ab4-41a0-44f1-b1cd-072bad362f4d/jobs/2326963

Thanks tho, overall LGTM as well

[qgallouedec]

Seems like we have 1 new failure

0a3d375 should fix it

[github-actions]

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46204&sha=0a3d37