Fix missing num_items_in_batch in unsloth_prediction_step (#4616)

* Fix missing num_items_in_batch in unsloth_prediction_step

unsloth_prediction_step calls compute_loss without num_items_in_batch
during evaluation. This causes _unsloth_pre_compute_loss to see
num_items_in_batch=None, which triggers a spurious warning for every
model when gradient_accumulation_steps > 1:

  "Unsloth: Not an error, but {model} does not accept num_items_in_batch.
   Using gradient accumulation will be very slightly less accurate."

The standard transformers prediction_step computes num_items_in_batch
via _get_num_items_in_batch before passing it to compute_loss. This
patch does the same in unsloth_prediction_step.

Tested on Llama-3.2-1B-Instruct and Olmo-3-7B-Instruct with
gradient_accumulation_steps=3 and eval_steps=3. Warning is gone and
eval loss is computed correctly for both.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Guard _get_num_items_in_batch for older transformers versions

_get_num_items_in_batch was added in transformers 4.46. Wrap the call
in try/except so older versions fall back to num_items_in_batch=None,
which preserves the original behavior of not passing it.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-03-27 03:06:59 -07:00 committed by GitHub
commit 6b5da2ea0f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -273,8 +273,17 @@ def PatchRL(FastLanguageModel):
with torch.no_grad():
if has_labels or loss_without_labels:
with self.compute_loss_context_manager():
try:
num_items_in_batch = self._get_num_items_in_batch(
[inputs], self.args.device
)
except (AttributeError, TypeError):
num_items_in_batch = None
loss, outputs = self.compute_loss(
model, inputs, return_outputs = True
model,
inputs,
return_outputs = True,
num_items_in_batch = num_items_in_batch,
)
loss = loss.mean().detach()