From 00fe9a40c06968b54b4e0bb35e1747045e1a48a2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 25 Feb 2026 06:34:21 -0800 Subject: [PATCH] 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 | 41 ++++++++++++++++++++++++++++++- unsloth/models/rl_replacements.py | 3 ++- unsloth/tokenizer_utils.py | 9 ++++++- 3 files changed, 50 insertions(+), 3 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 181e9479df..fbb32b2c50 100755 --- a/unsloth/models/rl.py +++ b/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": diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 27f00f10f3..b3a55440f9 100755 --- a/unsloth/models/rl_replacements.py +++ b/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" diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 1c107bd84f..c445879df7 100644 --- a/unsloth/tokenizer_utils.py +++ b/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}",