From 6b5da2ea0f96af59321993fff45ffce1fe842731 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 27 Mar 2026 03:06:59 -0700 Subject: [PATCH] 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> --- unsloth/models/rl.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 581244e4d3..5651a7da41 100755 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -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()