From 982ae7bbebc1bef9c24dae857c794ac82cd75981 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 31 Dec 2025 21:35:48 -0800 Subject: [PATCH] Fix correctness bugs in rl.py, rl_replacements.py, and vision.py (#3811) * Fix correctness bugs in rl.py, rl_replacements.py, and vision.py 1. rl_replacements.py (lines 864, 870): Fixed undefined `nanmin`/`nanmax` functions by using `.nan_to_num(nan=inf/-inf).min()/.max()` pattern. PyTorch doesn't have torch.nanmin/nanmax, so we replace NaN values before computing min/max. 2. vision.py (line 150): Fixed bug where code checked for "input" key but then accessed kwargs["input_ids"] instead of kwargs["input"]. 3. vision.py (line 159): Fixed bug where literal string "key" was used instead of the variable `key` when accessing kwargs. 4. rl.py (lines 903, 905): Fixed non-existent `MathError` exception by replacing with `ValueError`. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/rl.py | 4 ++-- unsloth/models/rl_replacements.py | 10 ++++++++-- unsloth/models/vision.py | 4 ++-- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 03f2c44701..e1c43b8b85 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -900,9 +900,9 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): if "temperature" in call_args: check_temperature = ( "if temperature <= 0:\n" - " raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')\n" + " raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')\n" "elif temperature >= 10:\n" - " raise MathError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.')\n" + " raise ValueError('Unsloth: Please set a positive non-zero temperature less than 10, since sampling will be quite erratic.')\n" "\n" ) extra_args += check_temperature diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 3dfeea6871..5e079335ae 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -861,13 +861,19 @@ def grpo_trainer_compute_loss(function_name, function): else torch.tensor(0.0, device = self.model.device) ) self._metrics[mode]["sampling/importance_sampling_ratio/min"].append( - nanmin(self.accelerator.gather(min_importance_sampling_ratio)).item() + self.accelerator.gather(min_importance_sampling_ratio) + .nan_to_num(nan = float("inf")) + .min() + .item() ) self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append( self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item() ) self._metrics[mode]["sampling/importance_sampling_ratio/max"].append( - nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item() + self.accelerator.gather(max_importance_sampling_ratio) + .nan_to_num(nan = float("-inf")) + .max() + .item() ) return loss diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 36cfbf0b17..c909f963b9 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -147,7 +147,7 @@ def unsloth_base_fast_generate( elif "input_ids" in kwargs: input_ids = kwargs["input_ids"] elif "input" in kwargs: - input_ids = kwargs["input_ids"] + input_ids = kwargs["input"] elif "input_features" in kwargs: input_ids = kwargs["input_features"] elif "input_embeds" in kwargs: @@ -156,7 +156,7 @@ def unsloth_base_fast_generate( input_ids = kwargs["inputs"] else: key = next(iter(kwargs.keys())) - if type(kwargs["key"]) is not torch.Tensor: + if type(kwargs[key]) is not torch.Tensor: raise TypeError("Unsloth: You need to pass in input_ids to .generate!") input_ids = kwargs[key] assert type(input_ids) is torch.Tensor