Add resilience to TRL internal API reclassification (#4111)

* Add resilience to TRL internal API reclassification

TRL is moving toward v1.0 and will reclassify several
currently-importable symbols as internal with no stability
guarantees. This adds try/except cascading imports with local
fallbacks so Unsloth keeps working regardless of whether TRL
removes, moves, or restructures these symbols.

Changes:
- rl.py: Add try/except cascade for unwrap_model_for_generation
  with local contextmanager fallback. Wire sanitize_logprob from
  RL_REPLACEMENTS into the compiled trainer template (same pipeline
  as selective_log_softmax and other global functions). Add import
  math and import logging to the template header.
- rl_replacements.py: Remove inline import of sanitize_logprob
  from trl.scripts.vllm_serve in the regex replacement. The
  function is now a module-level global in the compiled file.
- tokenizer_utils.py: Wrap dynamic exec import with per-item
  fallback so a single removed symbol does not break the entire
  bulk import.

Depends on unslothai/unsloth-zoo#516.

Tested across all TRL versions from 0.22.2 through 0.29.0.dev0
(git main). Training losses and grad norms are bit-identical
to unpatched runs.

* [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

@@ -94,7 +94,40 @@ def vLLMSamplingParams(**kwargs):
 
 
 def PatchRL(FastLanguageModel):
-    from trl.models.utils import unwrap_model_for_generation
+    try:
+        from trl.models.utils import unwrap_model_for_generation
+    except ImportError:
+        try:
+            from trl.models import unwrap_model_for_generation
+        except ImportError:
+            # Local fallback -- TRL removed or moved this symbol
+            from contextlib import contextmanager as _cm
+
+            @_cm
+            def unwrap_model_for_generation(
+                model, accelerator, gather_deepspeed3_params = True
+            ):
+                unwrapped_model = accelerator.unwrap_model(model)
+                is_gc = getattr(unwrapped_model, "is_gradient_checkpointing", False)
+                if is_gc:
+                    unwrapped_model.gradient_checkpointing_disable()
+                if (
+                    getattr(accelerator, "state", None) is not None
+                    and getattr(accelerator.state, "deepspeed_plugin", None) is not None
+                    and accelerator.state.deepspeed_plugin.zero_stage == 3
+                ):
+                    if not gather_deepspeed3_params:
+                        yield accelerator.unwrap_model(model)
+                    else:
+                        import deepspeed
+
+                        with deepspeed.zero.GatheredParameters(model.parameters()):
+                            yield accelerator.unwrap_model(model)
+                else:
+                    yield unwrapped_model
+                if is_gc:
+                    unwrapped_model.gradient_checkpointing_enable()
+
     from contextlib import contextmanager
 
     @contextmanager
@@ -253,9 +286,12 @@ create_completion_attention_mask = RL_REPLACEMENTS["create_completion_attention_
 left_pack_padding = RL_REPLACEMENTS["left_pack_padding"]
 align_logprobs_with_mask = RL_REPLACEMENTS["align_logprobs_with_mask"]
 autotune_batch_and_chunks = RL_REPLACEMENTS["grpo_autotune_batch_and_chunks"]
+sanitize_logprob = RL_REPLACEMENTS["sanitize_logprob"]
 
 RLTrainer_replacement = '''
 import os
+import math
+import logging
 from typing import *
 from dataclasses import dataclass, field
 from packaging.version import Version
@@ -324,6 +360,7 @@ torch_compile_options = {{
 {left_pack_padding_code}
 {align_logprobs_with_mask_code}
 {autotune_batch_and_chunks_code}
+{sanitize_logprob_code}
 
 {RL_pre}
 
@@ -1228,6 +1265,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
     left_pack_padding_code = inspect.getsource(left_pack_padding)
     align_logprobs_with_mask_code = inspect.getsource(align_logprobs_with_mask)
     autotune_batch_and_chunks_code = inspect.getsource(autotune_batch_and_chunks)
+    sanitize_logprob_code = inspect.getsource(sanitize_logprob)
     # Get final source code
     RLTrainer_source = RLTrainer_replacement.format(
         RLTrainer_name = RLTrainer_name,
@@ -1256,6 +1294,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         autotune_batch_and_chunks_code = autotune_batch_and_chunks_code,
         left_pack_padding_code = left_pack_padding_code,
         align_logprobs_with_mask_code = align_logprobs_with_mask_code,
+        sanitize_logprob_code = sanitize_logprob_code,
     )
 
     if RLTrainer_name == "GRPOTrainer":

unsloth/models/rl_replacements.py

@@ -355,8 +355,9 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
         re.DOTALL | re.MULTILINE,
     )
 
+    # sanitize_logprob is injected as a module-level function via RLTrainer_replacement
+    # template in rl.py (from RL_REPLACEMENTS), so just reference it directly here.
     replacement_text = (
-        r"\1from trl.scripts.vllm_serve import sanitize_logprob\n"
         r"\1all_logprobs = [\n"
         r"\1    [sanitize_logprob(next(iter(logprob.values()))) for logprob in output.logprobs]\n"
         r"\1    for outputs in all_outputs\n"

unsloth/tokenizer_utils.py

@@ -1007,7 +1007,14 @@ def patch_sft_trainer_tokenizer():
             function = function.replace(replacer, check_text + replacer)
 
         x = [x for x in all_imports if x in function]
-        exec(f"from trl.trainer.sft_trainer import ({','.join(x)})", locals())
+        try:
+            exec(f"from trl.trainer.sft_trainer import ({','.join(x)})", locals())
+        except ImportError:
+            for _item in x:
+                try:
+                    exec(f"from trl.trainer.sft_trainer import {_item}", locals())
+                except ImportError:
+                    pass
         exec(function, locals(), globals())
         exec(
             f"trl.trainer.sft_trainer.SFTTrainer.{function_name} = {function_name}",