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

@@ -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

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

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