unslothai--unsloth/unsloth/models/rl.py

# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

__all__ = [
    "PatchFastRL",
    "vLLMSamplingParams",
]

import torch
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
import importlib
import inspect
import os
import re
import sys
from contextlib import contextmanager
from unsloth_zoo.compiler import create_new_function
from unsloth_zoo.log import logger
from unsloth_zoo.logging_utils import PatchRLStatistics
from unsloth_zoo.rl_replacements import RL_REPLACEMENTS
from ..device_type import DEVICE_TYPE
from .rl_replacements import (
    RL_EXTRA_ARGS,
    RL_FUNCTIONS,
    RL_PRE_ITEMS,
    RL_CONFIG_CHANGES,
    RL_METRICS_CHANGES,
    RL_ADDITIONAL_FUNCTIONS,
)

torch_compile_options = {
    "epilogue_fusion": True,
    "max_autotune": False,  # Disable Triton mm kernels
    "shape_padding": True,
    "trace.enabled": False,
    "triton.cudagraphs": False,
}

# vLLM compatibility shim (TRL expects GuidedDecodingParams even if vLLM doesn't provide it)
try:
    import vllm.sampling_params as _unsloth_vllm_sp
    if not hasattr(_unsloth_vllm_sp, "GuidedDecodingParams"):

        class GuidedDecodingParams:
            def __init__(self, **kwargs):
                self.kwargs = kwargs

        _unsloth_vllm_sp.GuidedDecodingParams = GuidedDecodingParams
except Exception:
    pass

from trl import __version__ as trl_version_raw
from importlib.metadata import version as importlib_version
from unsloth_zoo.utils import Version

try:
    trl_version = Version(trl_version_raw)
except Exception:
    try:
        trl_version = Version(importlib_version("trl"))
    except Exception:
        trl_version = Version("0.0.0")

# Get PyTorch version for feature detection
try:
    torch_version = Version(torch.__version__.split("+")[0].split("a")[0].split("b")[0])
except Exception:
    torch_version = Version("0.0.0")

# Get transformers version for feature detection
try:
    from transformers import __version__ as _transformers_version_raw
    transformers_version = Version(_transformers_version_raw)
except Exception:
    transformers_version = Version("0.0.0")


def vLLMSamplingParams(**kwargs):
    from vllm import SamplingParams

    sampling_params = SamplingParams(**kwargs)
    sampling_params._set_kwargs = kwargs
    return sampling_params


def _maybe_prepare_vllm_for_resume(trainer):
    if not torch.cuda.is_available():
        return

    llm = getattr(trainer, "llm", None)
    if llm is None:
        llm = getattr(getattr(trainer, "model", None), "vllm_engine", None)
    if llm is None:
        return

    model_config = getattr(
        getattr(getattr(llm, "llm_engine", None), "vllm_config", None),
        "model_config",
        None,
    )
    if not getattr(model_config, "enable_sleep_mode", False):
        return

    try:
        llm.sleep(1)
    except Exception:
        pass

    import gc

    for _ in range(3):
        gc.collect()
        torch.cuda.empty_cache()


def _patch_resume_from_checkpoint_memory(trainer_class):
    original_train = getattr(trainer_class, "train", None)
    if original_train is None:
        return
    if getattr(original_train, "_unsloth_resume_guard", False):
        return

    def _unsloth_train_with_resume_guard(self, *args, **kwargs):
        resume_from_checkpoint = kwargs.get("resume_from_checkpoint", None)
        if resume_from_checkpoint is None:
            resume_from_checkpoint = kwargs.get("model_path", None)
        if resume_from_checkpoint is None and len(args) != 0:
            resume_from_checkpoint = args[0]

        if resume_from_checkpoint:
            _maybe_prepare_vllm_for_resume(self)
        return original_train(self, *args, **kwargs)

    _unsloth_train_with_resume_guard._unsloth_resume_guard = True
    trainer_class.train = _unsloth_train_with_resume_guard


def PatchRL(FastLanguageModel):
    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
    def unsloth_unwrap_model_for_generation(model, *args, **kwargs):
        # why: snapshot before TRL's unwrap context manager, which calls
        # gradient_checkpointing_disable() before yielding; preserve the actual
        # mode value (e.g. "unsloth") rather than collapsing it to a bool, so
        # the finally restore matches the caller's configured GC mode.
        use_gradient_checkpointing = next(
            (
                v
                for v in (getattr(m, "gradient_checkpointing", False) for m in model.modules())
                if v
            ),
            False,
        )
        with unwrap_model_for_generation(model, *args, **kwargs) as unwrapped_model:
            # Put the model in inference mode.
            FastLanguageModel.for_inference(model)

            # We must use .clone for Unsloth since we force inference_mode
            # Rather we should have used no_grad
            original_generate = unwrapped_model.generate

            def generate_with_clone(*args, **kwargs):
                out = original_generate(*args, **kwargs)
                if isinstance(out, torch.Tensor):
                    return out.clone()
                return out

            unwrapped_model.generate = generate_with_clone

            try:
                yield unwrapped_model
            finally:
                # Restore generate and return
                unwrapped_model.generate = original_generate
                FastLanguageModel.for_training(
                    model,
                    use_gradient_checkpointing = use_gradient_checkpointing,
                )

    from transformers import Trainer
    from transformers.trainer_pt_utils import nested_detach

    @torch.no_grad()
    def unsloth_prediction_step(self, model, inputs, prediction_loss_only, ignore_keys):
        """
        Perform an evaluation step on `model` using `inputs`.
        Subclass and override to inject custom behavior.
        Args:
            model (`nn.Module`):
                The model to evaluate.
            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
                The inputs and targets of the model.
                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument `labels`. Check your model's documentation for all accepted arguments.
            prediction_loss_only (`bool`):
                Whether or not to return the loss only.
            ignore_keys (`List[str]`, *optional*):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.
        Return:
            Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss,
            logits and labels (each being optional).
        """
        has_labels = (
            False
            if len(self.label_names) == 0
            else all(inputs.get(k) is not None for k in self.label_names)
        )
        # For CLIP-like models capable of returning loss values.
        # If `return_loss` is not specified or being `None` in `inputs`, we check if the default value of `return_loss`
        # is `True` in `model.forward`.
        return_loss = inputs.get("return_loss", None)
        if return_loss is None:
            return_loss = self.can_return_loss
        loss_without_labels = True if len(self.label_names) == 0 and return_loss else False

        inputs = self._prepare_inputs(inputs)
        if ignore_keys is None:
            if hasattr(self.model, "config"):
                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
            else:
                ignore_keys = []

        # labels may be popped when computing the loss (label smoothing for instance) so we grab them first.
        if has_labels or loss_without_labels:
            labels = nested_detach(tuple(inputs.get(name) for name in self.label_names))
            if len(labels) == 1:
                labels = labels[0]
        else:
            labels = None

        os.environ["UNSLOTH_RETURN_LOGITS"] = "1"
        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,
                        num_items_in_batch = num_items_in_batch,
                    )
                loss = loss.mean().detach()

                if isinstance(outputs, dict):
                    logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"])
                else:
                    logits = outputs[1:]
            else:
                loss = None
                with self.compute_loss_context_manager():
                    tokenized_output = self.processing_class(
                        inputs["prompt"],
                        padding = True,
                        truncation = True,
                        return_tensors = "pt",
                    ).to(model.device)
                    outputs = model(**tokenized_output)
                if isinstance(outputs, dict):
                    logits = tuple(v for k, v in outputs.items() if k not in ignore_keys)
                else:
                    logits = outputs
                # TODO: this needs to be fixed and made cleaner later.
                if self.args.past_index >= 0:
                    self._past = outputs[self.args.past_index - 1]
        os.environ["UNSLOTH_RETURN_LOGITS"] = "0"
        if prediction_loss_only:
            return (loss, None, None)

        logits = nested_detach(logits)
        if len(logits) == 1:
            logits = logits[0]

        return (loss, logits, labels)

    import trl.trainer

    trainers = dir(trl.trainer)
    trainers = [x for x in trainers if x.endswith("_trainer")]
    unwrap = "unwrap_model_for_generation"
    for trainer in trainers:
        try:
            current_trainer = getattr(trl.trainer, trainer)
        except:
            continue
        if hasattr(current_trainer, unwrap):
            try:
                setattr(current_trainer, unwrap, unsloth_unwrap_model_for_generation)
            except:
                continue
    Trainer.prediction_step = unsloth_prediction_step


grpo_selective_log_softmax = RL_REPLACEMENTS["grpo_selective_log_softmax"]
selective_log_softmax = RL_REPLACEMENTS["selective_log_softmax"]
calculate_pad_tokens_in_prompt = RL_REPLACEMENTS["calculate_pad_tokens_in_prompt"]
create_completion_attention_mask = RL_REPLACEMENTS["create_completion_attention_mask"]
left_pack_padding = RL_REPLACEMENTS["left_pack_padding"]
align_logprobs_with_mask = RL_REPLACEMENTS["align_logprobs_with_mask"]
align_completion_tool_mask = RL_REPLACEMENTS.get("align_completion_tool_mask")
if align_completion_tool_mask is None:

    def align_completion_tool_mask(
        tool_mask: torch.Tensor, completion_mask: torch.Tensor
    ) -> torch.Tensor:
        if tool_mask is None:
            return completion_mask
        raise RuntimeError(
            "env_mask/tool_mask GRPO requires an unsloth_zoo build whose "
            "grpo_accumulated_loss handles tool_mask. Please upgrade "
            "unsloth_zoo."
        )


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
import torch
import numpy as np
from contextlib import nullcontext
from torch.nn import functional as F
import inspect
from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling as TransformersDataCollatorForLanguageModeling
from transformers.training_args import ParallelMode
from unsloth_zoo.device_type import DEVICE_TYPE, device_synchronize

# Wrap trainer with padding to right and enable training mode
import functools
from types import MethodType
try:
    from unsloth_zoo.gradient_checkpointing import reset_unsloth_gradient_checkpointing_buffers
except:
    def reset_unsloth_gradient_checkpointing_buffers(): pass
# Canonical reset lives in unsloth.models._utils so the SFT auto-packing wrapper and the plain
# Trainer loop can import the same helper; fall back to a no-op only if it can't be imported.
try:
    from unsloth.models._utils import _unsloth_reset_stray_compile_cache
except Exception:
    def _unsloth_reset_stray_compile_cache(self): pass
def prepare_for_training_mode(f):
    @functools.wraps(f)
    def wrapper(self, *args, **kwargs):
        # Drop any torch.compile graph cache poisoned by a stray pre-train forward.
        try:
            _unsloth_reset_stray_compile_cache(self)
        except Exception:
            pass
        # Finish the previous W&B run if this is a subsequent train() call.
        # We do this at the START of train() (not the end) so that
        # evaluate() / log() still work after train() completes.
        # HF's WandbCallback.setup() will call wandb.init() for the new run.
        # See: https://github.com/unslothai/unsloth/issues/3954
        if getattr(self, '_unsloth_training_completed', False):
            try:
                import wandb
                if wandb.run is not None:
                    wandb.finish()
                    # Reset HF's WandbCallback so it calls wandb.init() for the new run
                    for cb in self.callback_handler.callbacks:
                        if type(cb).__name__ == 'WandbCallback':
                            cb._initialized = False
                            break
            except:
                pass
        # Enable training mode
        _was_training = None
        # Restore the GC mode the model was configured with at setup; fall back to
        # the training args only when it wasn't recorded (issue #4735). Use hasattr,
        # not a None sentinel, so a deliberately-recorded None is restored verbatim.
        _model = getattr(self, 'model', None)
        if hasattr(_model, '_unsloth_gradient_checkpointing'):
            use_gc = _model._unsloth_gradient_checkpointing
        else:
            use_gc = getattr(self.args, 'gradient_checkpointing', True)
        if hasattr(self, 'model') and hasattr(self.model, "training"):
            _was_training = self.model.training
        if hasattr(self, 'model') and hasattr(self.model, "for_training"):
            self.model.for_training(use_gradient_checkpointing=use_gc)
        output = f(self, *args, **kwargs)
        # Restore previous mode when possible
        if hasattr(self, 'model') and hasattr(self.model, "for_inference"):
            if _was_training is False:
                self.model.for_inference()
            elif _was_training is True and hasattr(self.model, "for_training"):
                self.model.for_training(use_gradient_checkpointing=use_gc)
        # Reset gradient checkpointing buffers to free memory while staying ready for next run
        try:
            reset_unsloth_gradient_checkpointing_buffers()
        except:
            pass
        # Mark that training completed so the next train() call can
        # finish this W&B run before starting a new one
        self._unsloth_training_completed = True
        return output
    return wrapper
pass

torch_compile_options = {{
    "epilogue_fusion"   : True,
    "max_autotune"      : False,
    "shape_padding"     : True,
    "trace.enabled"     : False,
    "triton.cudagraphs" : False,
}}

{grpo_selective_log_softmax_code}
{selective_log_softmax_code}
{calculate_pad_tokens_in_prompt_code}
{create_completion_attention_mask_code}
{left_pack_padding_code}
{align_logprobs_with_mask_code}
{align_completion_tool_mask_code}
{autotune_batch_and_chunks_code}
{sanitize_logprob_code}

{RL_pre}

@dataclass
class Unsloth{RLConfig_name}({RLConfig_name}):
    """
    {__RLConfig_doc__}
    """
    vllm_sampling_params: Optional[Any] = field(
        default = None,
        metadata = {{'help': 'vLLM SamplingParams'}},
    )
    unsloth_num_chunks : Optional[int] = field(
        default = -1,
        metadata = {{'help': 'Chunk size to reduce memory usage. -1 is most efficient.'}},
    )
    unsloth_logit_chunk_multiplier : Optional[int] = field(
            default = None,
            metadata = {{'help': 'Multiplier for chunked logit computations.'}},
        )
    unsloth_grpo_mini_batch : Optional[int] = field(
        default = None,
        metadata = {{'help': 'Mini batch size for GRPO hidden state accumulation. Default is None unless user defines it.'}},
    )
    {max_seq_length_pre}
    def __init__({RLConfig_arguments},
        vllm_sampling_params = None,
        unsloth_num_chunks = -1,
        unsloth_logit_chunk_multiplier = None,
        unsloth_grpo_mini_batch = None,
        {max_seq_length_call}
        **kwargs,
    ):
{RLConfig_extra_args}
        super().__init__({RLConfig_call_args}{RLConfig_kwargs})
        self.vllm_sampling_params = vllm_sampling_params
        self.unsloth_num_chunks = unsloth_num_chunks
        if unsloth_grpo_mini_batch is not None:
            if self.generation_batch_size >= unsloth_grpo_mini_batch:
                self.unsloth_grpo_mini_batch = unsloth_grpo_mini_batch
            else:
                raise ValueError(
                    f"Unsloth GRPO mini batch size needs to be less than or equal to the effective generation batch size, "
                    f"which is self.per_device_train_batch_size * gradient_accumulation_steps."
                )
        self.unsloth_logit_chunk_multiplier = unsloth_logit_chunk_multiplier
        {max_seq_length_post}
{RLConfig_post}
pass

{RLTrainer_extras}

class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}):
    """
    {__RLTrainer_doc__}
    """
    def __init__({RLTrainer_arguments},
        **kwargs
    ):
        if args is None: args = Unsloth{RLConfig_name}()
{RLTrainer_extra_args}
        # [TODO] Fix up DataParallel multiplying batch sizes
        # [TODO] DDP works, but DP seems to not work? [TODO]
        if getattr(args, "parallel_mode", None) == ParallelMode.NOT_DISTRIBUTED and args.n_gpu > 1:
            if getattr(args, "_n_gpu", 1) != 1:
                args._n_gpu = 1
        if "model" in locals() and hasattr(model, "for_training"):
            _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True)
            model.for_training(use_gradient_checkpointing=_use_gc)
        super().__init__({RLTrainer_call_args}{RLTrainer_kwargs})
        if "model" in locals() and hasattr(model, "for_inference"):
            model.for_inference()
{RLTrainer_post}
pass
'''


def _wrap_grpo_generate_and_score(trainer_cls):
    if not hasattr(trainer_cls, "_generate_and_score_completions"):
        return
    original = trainer_cls._generate_and_score_completions
    if getattr(original, "_unsloth_restore_training_wrapped", False):
        return

    def wrapped(self, *args, **kwargs):
        was_training = getattr(getattr(self, "model", None), "training", None)
        try:
            return original(self, *args, **kwargs)
        finally:
            if (
                was_training is False
                and hasattr(self, "model")
                and hasattr(self.model, "for_inference")
            ):
                try:
                    self.model.for_inference()
                except Exception:
                    pass

    wrapped._unsloth_restore_training_wrapped = True
    trainer_cls._generate_and_score_completions = wrapped


_UNSLOTH_RETURN_HIDDEN_STATES_SUPPORT_MARKER = "__UNSLOTH_SUPPORTS_RETURN_HIDDEN_STATES__"
_UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR = "_unsloth_grpo_hidden_states_forward_wrapped"
_UNSLOTH_GRPO_HIDDEN_STATES_WARNING_ATTR = "_unsloth_grpo_hidden_states_warning_issued"


def _grpo_hidden_states_wrap_target(model):
    if model is None:
        return None
    get_base_model = getattr(model, "get_base_model", None)
    if callable(get_base_model):
        base_model = get_base_model()
        if base_model is not None and base_model is not model:
            return base_model
    for attr in ("base_model", "model"):
        child = getattr(model, attr, None)
        if child is not None and child is not model and hasattr(child, "forward"):
            return child
    return model


def _model_supports_unsloth_return_hidden_states(model):
    target_model = _grpo_hidden_states_wrap_target(model)
    for candidate in (model, target_model):
        if candidate is None:
            continue
        if getattr(candidate, _UNSLOTH_RETURN_HIDDEN_STATES_SUPPORT_MARKER, False):
            return True
        if getattr(type(candidate), _UNSLOTH_RETURN_HIDDEN_STATES_SUPPORT_MARKER, False):
            return True
    return False


def _drop_forward_kwargs_consumed_positionally(forward_signature, args, kwargs):
    if len(args) == 0 or len(kwargs) == 0:
        return kwargs

    consumed_names = []
    for parameter in forward_signature.parameters.values():
        if parameter.kind == inspect.Parameter.VAR_POSITIONAL:
            break
        if parameter.kind in (
            inspect.Parameter.POSITIONAL_ONLY,
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
        ):
            consumed_names.append(parameter.name)
        if len(consumed_names) >= len(args):
            break

    if len(consumed_names) == 0:
        return kwargs

    kwargs = dict(kwargs)
    for name in consumed_names:
        kwargs.pop(name, None)
    return kwargs


def _get_num_logits_to_keep(forward_signature, args, kwargs):
    try:
        bound = forward_signature.bind_partial(*args, **kwargs)
        arguments = bound.arguments
        num_logits_to_keep = arguments.get("num_logits_to_keep", 0) or 0
        logits_to_keep = arguments.get("logits_to_keep", 0) or 0
        for parameter in forward_signature.parameters.values():
            if parameter.kind != inspect.Parameter.VAR_KEYWORD:
                continue
            extra_kwargs = arguments.get(parameter.name, {})
            num_logits_to_keep = max(
                num_logits_to_keep,
                extra_kwargs.get("num_logits_to_keep", 0) or 0,
            )
            logits_to_keep = max(
                logits_to_keep,
                extra_kwargs.get("logits_to_keep", 0) or 0,
            )
            break
        return max(num_logits_to_keep, logits_to_keep)
    except TypeError:
        logger.debug(
            "Unsloth: Could not bind forward arguments for GRPO hidden-state fallback.",
            exc_info = True,
        )

    num_logits_to_keep = kwargs.get("num_logits_to_keep", 0) or 0
    logits_to_keep = kwargs.get("logits_to_keep", 0) or 0
    return max(num_logits_to_keep, logits_to_keep)


def _warn_grpo_hidden_states_fallback_once(model, message):
    if getattr(model, _UNSLOTH_GRPO_HIDDEN_STATES_WARNING_ATTR, False):
        return
    setattr(model, _UNSLOTH_GRPO_HIDDEN_STATES_WARNING_ATTR, True)
    logger.warning(message)


def _replace_outputs_logits(outputs, hidden_states):
    if hasattr(outputs, "logits"):
        outputs.logits = hidden_states
        return outputs
    if isinstance(outputs, dict):
        outputs["logits"] = hidden_states
        return outputs
    if isinstance(outputs, tuple) and len(outputs) != 0:
        return (hidden_states,) + tuple(outputs[1:])
    raise TypeError(f"Unsupported output type for GRPO hidden-state fallback: {type(outputs)}")


def _install_grpo_hidden_states_forward_wrapper(model):
    if model is None or getattr(model, _UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR, False):
        return False
    if _model_supports_unsloth_return_hidden_states(model):
        return False

    target_model = _grpo_hidden_states_wrap_target(model)
    if getattr(target_model, _UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR, False):
        setattr(model, _UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR, True)
        return False

    original_forward = target_model.forward
    forward_signature = inspect.signature(original_forward)
    model_name = type(target_model).__name__

    def wrapped_forward(*args, **kwargs):
        if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") != "1":
            return original_forward(*args, **kwargs)

        forward_kwargs = _drop_forward_kwargs_consumed_positionally(forward_signature, args, kwargs)
        num_logits_to_keep = _get_num_logits_to_keep(forward_signature, args, forward_kwargs)
        forward_kwargs["output_hidden_states"] = True
        forward_kwargs["return_dict"] = True
        try:
            outputs = original_forward(*args, **forward_kwargs)
        except TypeError as error:
            if "output_hidden_states" not in str(error) and "return_dict" not in str(error):
                raise
            _warn_grpo_hidden_states_fallback_once(
                target_model,
                f"Unsloth: GRPO fallback could not request hidden states for unsupported model {model_name}; using logits directly.",
            )
            return original_forward(*args, **kwargs)

        hidden_states = getattr(outputs, "hidden_states", None)
        if hidden_states is None or len(hidden_states) == 0:
            _warn_grpo_hidden_states_fallback_once(
                target_model,
                f"Unsloth: GRPO fallback did not receive hidden states for unsupported model {model_name}; using logits directly.",
            )
            return outputs

        hidden_states = hidden_states[-1]
        if num_logits_to_keep != 0:
            hidden_states = hidden_states[:, -num_logits_to_keep:, :]
        return _replace_outputs_logits(outputs, hidden_states)

    wrapped_forward._unsloth_grpo_hidden_states_forward_wrapped = True
    target_model.forward = wrapped_forward
    setattr(target_model, _UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR, True)
    setattr(model, _UNSLOTH_GRPO_HIDDEN_STATES_WRAPPED_ATTR, True)
    return True


def _wrap_grpo_hidden_states_fallback(trainer_cls):
    original_init = trainer_cls.__init__
    if getattr(original_init, "_unsloth_grpo_hidden_states_init_wrapped", False):
        return

    def wrapped_init(self, *args, **kwargs):
        original_init(self, *args, **kwargs)
        _install_grpo_hidden_states_forward_wrapper(getattr(self, "model", None))
        _install_grpo_hidden_states_forward_wrapper(getattr(self, "ref_model", None))

    wrapped_init._unsloth_grpo_hidden_states_init_wrapped = True
    trainer_cls.__init__ = wrapped_init


def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
    # Defensive wrapper: matches patch_trl_rl_trainers()'s try/except so
    # direct callers don't see exceptions from the impl on TRL versions
    # that rename or move classes (e.g. TRL 1.x trl.experimental).
    try:
        return _patch_trl_rl_trainers_impl(trainer_file)
    except Exception as e:
        logger.info(
            f"Unsloth: Could not patch trl.trainer.{trainer_file}: " f"{type(e).__name__}: {e}"
        )
        return


def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"):
    # Patch for vLLM and Unsloth PEFT
    import trl
    import trl.trainer

    try:
        trainer = eval(f"trl.trainer.{trainer_file}")
    except Exception as error:
        logger.info(f"Unsloth: Could not import trl.trainer.{trainer_file}: {error}")
        return

    # Get SFTTrainer and SFTConfig names
    name = [
        x
        for x in dir(trainer)
        if x.endswith("Trainer")
        and x != "Trainer"
        and not x.startswith("_")
        and trainer_file.split("_")[0] in x.lower()
    ]
    config = [
        x
        for x in dir(trainer)
        if x.endswith("Config")
        and x != "Config"
        and not x.startswith("_")
        and trainer_file.split("_")[0] in x.lower()
    ]
    if len(name) != 1:
        logger.info(
            f"Unsloth: Could not find Trainer class in trl.trainer.{trainer_file}. Found: {name}"
        )
        return
    if len(config) != 1:
        # TRL 0.26+: Config may be in a separate *_config.py module
        config_module_name = trainer_file.replace("_trainer", "_config")
        try:
            config_mod = eval(f"trl.trainer.{config_module_name}")
            config = [
                x
                for x in dir(config_mod)
                if x.endswith("Config")
                and x != "Config"
                and not x.startswith("_")
                and trainer_file.split("_")[0] in x.lower()
            ]
        except Exception:
            pass
    if len(config) != 1 and len(name) == 1:
        # Thin wrapper fallback: walk the Trainer's MRO to find Config
        # in the real implementation module (e.g., trl.experimental.bco)
        try:
            _temp_cls = eval(f"trl.trainer.{trainer_file}.{name[0]}")
            for _parent in _temp_cls.__mro__[1:]:
                if _parent is object:
                    continue
                _parent_mod = inspect.getmodule(_parent)
                if _parent_mod is None or _parent_mod.__name__ == f"trl.trainer.{trainer_file}":
                    continue
                config = [
                    x
                    for x in dir(_parent_mod)
                    if x.endswith("Config")
                    and x != "Config"
                    and not x.startswith("_")
                    and trainer_file.split("_")[0] in x.lower()
                ]
                if len(config) == 1:
                    break
        except Exception:
            pass
    if len(config) != 1:
        logger.info(
            f"Unsloth: Could not find Config class in trl.trainer.{trainer_file}. Found: {config}"
        )
        return

    # Get SFTTrainer, SFTConfig
    RLTrainer_name = name[0]
    RLConfig_name = config[0]
    try:
        RLTrainer = eval(f"trl.trainer.{trainer_file}.{RLTrainer_name}")
    except Exception as e:
        logger.info(
            f"Unsloth: Could not load {RLTrainer_name} from trl.trainer.{trainer_file}: {e}"
        )
        return
    _config_resolved_module = None
    try:
        RLConfig = eval(f"trl.trainer.{trainer_file}.{RLConfig_name}")
    except Exception:
        # TRL 0.26+: Config may be in a separate *_config.py module
        try:
            config_module_name = trainer_file.replace("_trainer", "_config")
            RLConfig = eval(f"trl.trainer.{config_module_name}.{RLConfig_name}")
        except Exception:
            # Thin wrapper fallback: load Config from parent trainer's module
            _config_loaded = False
            try:
                _temp_cls = eval(f"trl.trainer.{trainer_file}.{name[0]}")
                for _parent in _temp_cls.__mro__[1:]:
                    if _parent is object:
                        continue
                    _parent_mod = inspect.getmodule(_parent)
                    if _parent_mod is None or _parent_mod.__name__ == f"trl.trainer.{trainer_file}":
                        continue
                    if hasattr(_parent_mod, RLConfig_name):
                        RLConfig = getattr(_parent_mod, RLConfig_name)
                        _config_resolved_module = _parent_mod
                        _config_loaded = True
                        break
            except Exception:
                pass
            if not _config_loaded:
                logger.info(f"Unsloth: Could not load {RLConfig_name}")
                return

    # Check name
    if RLTrainer.__name__.startswith("Unsloth"):
        print(f"Unsloth: {RLTrainer.__name__} is already patched.")
        return
    if RLConfig.__name__.startswith("Unsloth"):
        print(f"Unsloth: {RLConfig.__name__} is already patched.")
        return

    # TRL 0.26+: Resolve thin wrappers to their experimental parent class.
    # Thin wrappers are deprecation shims in trl.trainer that just forward
    # *args/**kwargs to the real implementation in trl.experimental.
    # Only resolve if a parent class actually lives in a trl.experimental module.
    _trainer_resolved_module = None
    try:
        _trainer_src = inspect.getsource(RLTrainer)
        _trainer_module = inspect.getmodule(RLTrainer)
        _trainer_module_src = inspect.getsource(_trainer_module) if _trainer_module else ""
        if "trl.experimental" in _trainer_src or "trl.experimental" in _trainer_module_src:
            for _parent in RLTrainer.__mro__[1:]:
                if _parent is object:
                    continue
                _parent_mod = inspect.getmodule(_parent)
                if _parent_mod is None:
                    continue
                # Only resolve to a parent that lives in trl.experimental
                if "trl.experimental" in _parent_mod.__name__:
                    RLTrainer = _parent
                    _trainer_resolved_module = _parent_mod
                    break
    except Exception:
        pass

    try:
        _config_src = inspect.getsource(RLConfig)
        _config_module = inspect.getmodule(RLConfig)
        _config_module_src = inspect.getsource(_config_module) if _config_module else ""
        if "trl.experimental" in _config_src or "trl.experimental" in _config_module_src:
            for _parent in RLConfig.__mro__[1:]:
                if _parent is object:
                    continue
                _parent_mod = inspect.getmodule(_parent)
                if _parent_mod is None:
                    continue
                # Only resolve to a parent that lives in trl.experimental
                if "trl.experimental" in _parent_mod.__name__:
                    RLConfig = _parent
                    break
    except Exception:
        pass

    # Get old source
    old_RLTrainer_source = inspect.getsource(RLTrainer)
    old_RLConfig_source = inspect.getsource(RLConfig)

    if _trainer_resolved_module is not None:
        all_imports = dir(_trainer_resolved_module)
    elif _config_resolved_module is not None:
        all_imports = dir(_config_resolved_module)
    else:
        all_imports = dir(trainer)
    # Fix _deprecate_arguments not getting imported so stop __ but not _
    imports = [x for x in all_imports if not x.startswith("__")]

    # Get default arguments
    EMPTY = inspect.Parameter.empty
    processed = []
    for RLobject in [RLTrainer, RLConfig]:
        parameters = inspect.signature(RLobject.__init__).parameters
        types = (
            bool,
            type(None),
            int,
            float,
            str,
        )
        arguments = ["self"]
        call_args = []
        for k, v in parameters.items():
            if k == "self":
                continue
            v = v.default
            if v == "\n":
                v = re.escape("\n")
            if v is EMPTY:
                arguments.append(k)
            elif type(v) is str:
                arguments.append(f"{k} = '{v}'")
            elif type(v) in types:
                arguments.append(f"{k} = {v}")
            else:
                continue
            call_args.append(f"{k} = {k}")
        arguments = f"\n{' ' * 8}" + f",\n{' ' * 8}".join(arguments)
        call_args = f"\n{' ' * 12}" + f",\n{' ' * 12}".join(call_args)
        processed.append(
            (
                arguments,
                call_args,
            )
        )

    # Process RLTrainer first
    arguments, call_args = processed[0]
    RLTrainer_post = ""

    # Add tokenizer if not seen
    if "tokenizer" not in parameters and "processing_class" in parameters:
        arguments += f",\n{' ' * 8}tokenizer = None"
        call_args = call_args.replace(
            "processing_class = processing_class",
            "processing_class = tokenizer if tokenizer is not None else processing_class",
        )

    # Edit bf16, fp16 by checking model's dtype/torch_dtype directly
    extra_args = ""
    if "args" in call_args and "model" in call_args:
        mixed_precision = (
            "use_bf16 = getattr(args, 'bf16', False)\n"
            "if type(use_bf16) is not bool: use_bf16 = False\n"
            "use_fp16 = getattr(args, 'fp16', False)\n"
            "if type(use_fp16) is not bool: use_fp16 = False\n"
            "force_float32 = False\n"
            # device-aware bf16 check (CUDA/XPU/HIP), so V100/T4 never pick bf16
            # but AMD/Intel are unaffected; fall back on older unsloth_zoo.
            "try:\n"
            "    from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported\n"
            "except Exception:\n"
            "    _bf16_supported = torch.cuda.is_bf16_supported\n"
            # FORCE_FLOAT32 models (Gemma3, gpt_oss, ...) cannot use float16. On a GPU without
            # bf16 (V100/T4) keep them in float32 so they never autocast to fp16. On a bf16 GPU,
            # full finetuning can still use bf16 autocast (master weights stay float32), which is
            # faster and uses less memory; LoRA/QLoRA keep float32 when forced.
            "full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1'\n"
            "if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' and not (full_finetuning and _bf16_supported()):\n"
            "    print('Unsloth: Switching to float32 training since model cannot work with float16')\n"
            "    force_float32 = True\n"
            "mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32')\n"
            "dtype = getattr(model.config, 'dtype', None) or getattr(model.config, 'torch_dtype', None)\n"
            "if dtype is None: dtype = model.get_input_embeddings().weight.dtype\n"
            "from unsloth_zoo.utils import _get_dtype\n"
            "dtype = _get_dtype(dtype)\n"
            "float16 = dtype == torch.float16\n"
            "bfloat16 = dtype == torch.bfloat16\n"
            "if full_finetuning:\n"
            "    if bfloat16 and use_fp16: use_fp16 = False\n"
            "    if float16 and use_bf16: use_bf16 = False\n"
            "if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n"
            "if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n"
            "if force_float32:\n"
            "    # Forced float32 training\n"
            "    args.fp16 = False\n"
            "    args.bf16 = False\n"
            "    os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'\n"
            "    if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'\n"
            "    # args.mixed_precision is a new argument which needs to be set now\n"
            "elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32':\n"
            "    # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16.\n"
            "    use_bf16_amp = (not float16) and _bf16_supported()\n"
            "    args.fp16 = not use_bf16_amp\n"
            "    args.bf16 = use_bf16_amp\n"
            "    os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16'\n"
            "    if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16'\n"
            "    # args.mixed_precision is a new argument which needs to be set now\n"
            "elif mixed_precision_dtype == 'bfloat16':\n"
            "    # Both False since bfloat16 full finetuning doesn't do any autocasting.\n"
            "    args.fp16 = False\n"
            "    args.bf16 = False\n"
            "    os.environ['ACCELERATE_MIXED_PRECISION'] = 'no'\n"
            "    if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'\n"
            "    # args.mixed_precision is a new argument which needs to be set now\n"
            "\n"
        )
        extra_args += mixed_precision

    # Check if per_device_eval_batch_size (default 8) bigger than bsz
    # Also use FP16 / BF16 evaluation
    if "args" in call_args:
        # Check eval_dataset first
        if "eval_dataset" in call_args:
            check_eval_dataset = (
                "if getattr(args, 'eval_dataset', None) is not None and "
                "getattr(args, 'eval_strategy', 'no') == 'no':\n"
                "    args.eval_strategy = 'steps'\n"
                "    if getattr(args, 'eval_steps', None) is None: args.eval_steps = 0.1\n"
            )
            extra_args += check_eval_dataset

        # Check if gradient accumulation bug fix is applied
        check_ga = (
            "ga_steps = getattr(args, 'gradient_accumulation_steps', None)\n"
            "if ga_steps is not None and ga_steps > 1:\n"
            "    from transformers import __version__ as transformers_version\n"
            "    if Version(transformers_version) <= Version('4.45.2'):\n"
            "        print('**** Unsloth: Please use our fixed gradient_accumulation_steps by updating transformers, TRL and Unsloth!\\n'\n"
            "              '`pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth transformers trl unsloth_zoo`')\n"
        )
        extra_args += check_ga

        eval_changes = (
            "if getattr(args, 'eval_strategy', 'no') != 'no':\n"
            "    eval_bsz = getattr(args, 'per_device_eval_batch_size', 8)\n"
            "    if eval_bsz == 8 and args.per_device_train_batch_size < eval_bsz: args.per_device_eval_batch_size = args.per_device_train_batch_size\n"
            "    if getattr(args, 'eval_accumulation_steps', None) is None and ga_steps is not None: args.eval_accumulation_steps = ga_steps\n"
            "fp16_full_eval = getattr(args, 'fp16_full_eval', False)\n"
            "if type(fp16_full_eval) is not bool: fp16_full_eval = False\n"
            "bf16_full_eval = getattr(args, 'bf16_full_eval', False)\n"
            "if type(bf16_full_eval) is not bool: bf16_full_eval = False\n"
            "if args.fp16 and bf16_full_eval: args.bf16_full_eval = False; args.fp16_full_eval = True\n"
            "if args.bf16 and fp16_full_eval: args.bf16_full_eval = True; args.fp16_full_eval = False\n"
            "if force_float32:\n"
            "    args.bf16_full_eval = False\n"
            "    args.fp16_full_eval = False\n"
            "elif os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32') == 'bfloat16':\n"
            "    args.bf16_full_eval = True\n"
            "    args.fp16_full_eval = False\n"
            "elif not bf16_full_eval and not fp16_full_eval:\n"
            "    args.bf16_full_eval = args.bf16\n"
            "    args.fp16_full_eval = args.fp16\n"
        )
        extra_args += eval_changes

    # Force logits to be produced if preprocess_logits_for_metrics or compute_metrics is used
    if "model" in call_args:
        logits_check = (
            "_output_logits = False\n"
            "if locals().get('compute_metrics', None) is not None: _output_logits = True\n"
            "if locals().get('preprocess_logits_for_metrics', None) is not None: _output_logits = True\n"
            "if _output_logits:\n"
            "    os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n"
        )
        extra_args += logits_check
        warnings_issued_check = (
            "if model is not None:\n"
            "    _warnings_issued = getattr(model, 'warnings_issued', None)\n"
            "    if _warnings_issued is None:\n"
            "        model.warnings_issued = {}\n"
            "    elif not isinstance(_warnings_issued, dict):\n"
            "        try:\n"
            "            model.warnings_issued = dict(_warnings_issued)\n"
            "        except Exception:\n"
            "            model.warnings_issued = {}\n"
        )
        extra_args += warnings_issued_check

    # Check max_seq_length
    if "model" in call_args:
        length_check = (
            "if 'max_seq_length' not in locals() and not hasattr(args, 'max_seq_length'):\n"
            "    pass\n"
            "else:\n"
            "    model_max_seq_length = getattr(model, 'max_seq_length', None)\n"
            "    args_max_seq_length  = getattr(args,  'max_seq_length', None)\n"
            "    if args_max_seq_length is None and model_max_seq_length is not None:\n"
            "        max_seq_length = model.max_seq_length\n"
            "        if hasattr(args, 'max_seq_length'): args.max_seq_length = max_seq_length\n"
            "    elif args_max_seq_length is not None and model_max_seq_length is not None:\n"
            "        if args_max_seq_length > model_max_seq_length:\n"
            "            print('Unsloth: You set `max_seq_length` as ' + str(args_max_seq_length) + ' but '\n"
            "                   'the maximum the model supports is ' + str(model_max_seq_length) + '. We shall reduce it.')\n"
            "            args.max_seq_length = model_max_seq_length\n"
        )
        extra_args += length_check

        # At this point max_seq_length might be set, but trl is moving to max_length
        if trainer_file == "sft_trainer":
            max_length_check = (
                "if 'max_length' not in locals() and not hasattr(args, 'max_length'):\n"
                "    pass\n"
                "else:\n"
                "    if hasattr(args, 'max_seq_length') and args.max_seq_length is not None and args.max_seq_length > 0:\n"
                "        if hasattr(args, 'max_length'):\n"
                "            args.max_length = args.max_seq_length\n"
                "            max_length = args.max_length\n"
                "    else:\n"
                "        model_max_length = getattr(model, 'max_seq_length', None)\n"
                "        if model_max_length is None: model_max_length = getattr(model, 'max_length', None)\n"
                "        if model_max_length is not None:\n"
                "            args.max_length = model_max_length\n"
                "            max_length = args.max_length\n"
                "        elif hasattr(args, 'max_length') and args.max_length is not None:\n"
                "            max_length = args.max_length\n"
                "            # if we are here, then we are in a weird case where max_length is set but max_seq_length is not set\n"
                "            setattr(model, 'max_seq_length', max_length)\n"
                "        else:\n"
                "            print('Unsloth: We did not find `max_seq_length` or `max_length` in the model or args. We will set it to 1024.')\n"
                "            args.max_length = 1024\n"
            )
            extra_args += max_length_check

    # Enable for training and move padding side of tokenizer to right
    if "model" in call_args:
        training_check = (
            "if model is not None and hasattr(model, 'for_training'):\n"
            "    _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True)\n"
            "    model.for_training(use_gradient_checkpointing=_use_gc)\n"
            "if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right'\n"
            "if 'processing_class' in locals():\n"
            "    if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right'\n"
            "    if hasattr(processing_class, 'tokenizer') and hasattr(processing_class.tokenizer, 'padding_side'): "
            "processing_class.tokenizer.padding_side = 'right'\n"
        )
        extra_args += training_check

    # Check data collator if it's correct!
    if "data_collator" in call_args and "train_dataset" in call_args:
        data_collator_check = (
            "__tokenizer = processing_class if 'processing_class' in locals() else tokenizer\n"
            "from unsloth_zoo.vision_utils import UnslothVisionDataCollator\n"
            "if not isinstance(data_collator, UnslothVisionDataCollator):\n"
            "    if isinstance(data_collator, DataCollatorForSeq2Seq) and 'labels' not in train_dataset.column_names:\n"
            "        data_collator = TransformersDataCollatorForLanguageModeling(\n"
            "            __tokenizer,\n"
            "            mlm = False,\n"
            "            mlm_probability = 0.0,\n"
            "            pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None),\n"
            "        )\n"
            "    elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling) and 'labels' in train_dataset.column_names:\n"
            "        data_collator = DataCollatorForSeq2Seq(\n"
            "            __tokenizer,\n"
            "            pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None),\n"
            "        )\n"
            "else:\n"
            "    if hasattr(args, 'remove_unused_columns'): args.remove_unused_columns = False\n"
            "    if hasattr(args, 'dataset_text_field'): args.dataset_text_field = ''\n"
            "    if hasattr(args, 'dataset_kwargs'): args.dataset_kwargs = {'skip_prepare_dataset': True}\n"
        )
        extra_args += data_collator_check

        # Also check if .pad exists -> if not, and is VLM, then change it!
        # Only swap LM/Seq2Seq collators; leave preference collators
        # (DPODataCollatorWithPadding etc.) alone so ORPO/DPO/CPO/KTO keep
        # their own prompt/chosen/rejected handling.
        pad_check = (
            "if not isinstance(data_collator, UnslothVisionDataCollator):\n"
            "    if not hasattr(__tokenizer, 'pad') and hasattr(__tokenizer, 'tokenizer'):\n"
            "        if isinstance(data_collator, DataCollatorForSeq2Seq):\n"
            "            data_collator = DataCollatorForSeq2Seq(\n"
            "                __tokenizer.tokenizer,\n"
            "                pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None),\n"
            "            )\n"
            "        elif isinstance(data_collator, TransformersDataCollatorForLanguageModeling):\n"
            "            data_collator = TransformersDataCollatorForLanguageModeling(\n"
            "                __tokenizer.tokenizer,\n"
            "                mlm = False,\n"
            "                mlm_probability = 0.0,\n"
            "                pad_to_multiple_of = getattr(args, 'pad_to_multiple_of', None),\n"
            "            )\n"
        )
        extra_args += pad_check

    # Check NEFTune
    if "model" in call_args:
        neftune_check = (
            "if hasattr(self, 'neftune_hook_handle'):\n"
            "    self.neftune_hook_handle.remove()\n"
            "    if hasattr(self, 'neftune_hook_handle'): del self.neftune_hook_handle\n"
            "if getattr(args, 'neftune_noise_alpha', None) is not None:\n"
            "    model.get_input_embeddings().neftune_noise_alpha = self.neftune_noise_alpha\n"
            "pass\n"
        )
        RLTrainer_post += neftune_check

    # Add accelerator scaler to model
    if "model" in call_args:
        accelerator_check = (
            "if hasattr(self, 'accelerator'):\n"
            "    scaler = self.accelerator.scaler\n"
            "    current_model = model\n"
            "    while hasattr(current_model, 'model'):\n"
            "        current_model.accelerator_scaler = scaler\n"
            "        current_model = current_model.model\n"
            "    current_model.accelerator_scaler = scaler\n"
            "pass\n"
        )
        RLTrainer_post += accelerator_check

    # Add enabling and disabling training modes
    if "model" in call_args:
        training_check = (
            "if hasattr(self, 'train'):\n"
            "    self.train = MethodType(prepare_for_training_mode(self.__class__.train), self)\n"
            "pass\n"
        )
        RLTrainer_post += training_check

    # Sync chat_template from processing_class to vLLM's tokenizer
    # This fixes base models that have custom chat templates applied after loading
    if "model" in call_args:
        vllm_chat_template_sync = (
            "if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'):\n"
            "    _vllm_tok = self.llm.get_tokenizer()\n"
            "    _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None)\n"
            "    if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None:\n"
            "        _vllm_tok.chat_template = _pc.chat_template\n"
            "pass\n"
        )
        RLTrainer_post += vllm_chat_template_sync

    # Edit optional metrics
    other_metrics_processor = ""
    if trainer_file in RL_METRICS_CHANGES:
        process_extra_args = RL_METRICS_CHANGES[trainer_file]
        for process_extra_arg in process_extra_args:
            other_metrics_processor += process_extra_arg(old_RLTrainer_source, old_RLConfig_source)

    # Add statistics as well!
    extra_args += (
        "other_metrics = []\n"
        f"{other_metrics_processor}\n"
        "from unsloth_zoo.logging_utils import PatchRLStatistics\n"
        f"PatchRLStatistics('{trainer_file}', other_metrics)\n"
    )

    # Patch optional args
    if trainer_file in RL_EXTRA_ARGS:
        process_extra_args = RL_EXTRA_ARGS[trainer_file]
        for process_extra_arg in process_extra_args:
            extra_args += process_extra_arg(call_args, extra_args)

    # Create RLTrainer args
    extra_args = extra_args.split("\n")
    extra_args = "\n".join(" " * 8 + x for x in extra_args)
    RLTrainer_post = RLTrainer_post.split("\n")
    RLTrainer_post = "\n".join(" " * 8 + x for x in RLTrainer_post)
    RLTrainer_arguments = arguments
    RLTrainer_extra_args = extra_args
    RLTrainer_call_args = call_args

    # Fix RLConfig next
    arguments, call_args = processed[1]
    extra_args = ""

    # Edit GA / bsz and weight_decay
    replacements = {
        "output_dir": None,
        "logging_nan_inf_filter": False,
        "per_device_train_batch_size": 4,
        "gradient_accumulation_steps": 2,
        # LoRA decays A and B toward 0 so effective W = W_init + (alpha/r) * B @ A is pulled toward W_init, not 0 as in full FT.
        # 0.001 keeps a small Frobenius prior |A|_F^2 + |B|_F^2 without measurably dragging the merged adapter back to base.
        "weight_decay": 0.001,
        "seed": 3407,
        "optim": "adamw_8bit",
        "learning_rate": 5e-05,
        "per_device_eval_batch_size": 4,
        "eval_accumulation_steps": 2,
        "torch_empty_cache_steps": 250,
        "logging_steps": 1,
        "max_seq_length": None,
        "num_generations": 8,
        # "steps_per_generation"          : 1, # Otherwise defaults to ga_steps which is wrong
        # "generation_batch_size"         : None, # Useless. If steps_per_generation set, generation_batch_size clashes
        "top_k": None,
        "vllm_mode": "colocate",
        "generation_kwargs": {},
        "bf16": False,
        "fp16": False,
        "report_to": "none",
        "include_tokens_per_second": False,
        "include_num_input_tokens_seen": False,
        "auto_find_batch_size": False,  # Auto /2 batch size - too many people complained so removing
        "dataloader_pin_memory": True,
        "padding_free": None,  # None = user didn't set it, allows auto-enable detection
        # Might fail so disable for now
        # "dataloader_persistent_workers" : True, # Keeps dataloader in RAM
        # "dataloader_prefetch_factor"    : 2,
        # "dataloader_num_workers"        : 2, # Default is 0 means 1
    }
    # warmup_ratio deprecated in transformers >= 5.0; warmup_steps accepts float
    if transformers_version >= Version("5.0.0"):
        replacements["warmup_steps"] = 0.1
    else:
        replacements["warmup_ratio"] = 0.1

    for k, v in replacements.items():
        x = f"{k}( = [^,\n]{{1,}})?,\n"
        y = f"'{v}'" if type(v) is str else f"{v}"
        y = f"{k} = {y},\n"
        arguments = re.sub(x, y, arguments)

    # Fix GRPO beta default as 0.001 TRL used to be 0.04, now 0.00!
    # https://github.com/huggingface/trl/pull/3516
    # https://verl.readthedocs.io/en/latest/examples/config.html
    if trainer_file == "grpo_trainer":
        replacements = {
            "loss_type": "bnpo",  # Default GRPO paper
            "beta": 0.001,  # Recommended as seen in verl
            "auto_find_batch_size": False,  # Cannot work on GRPO
            # [TODO] See https://fengyao.notion.site/off-policy-rl
            # https://github.com/huggingface/trl/pull/3867 (August 7th)
            "vllm_importance_sampling_correction": False,
            # TRL >= 1.7.0 enables the MoE router aux loss by default (0.001); the optimized
            # GRPO forward does not compute it, so default off. Opt in via router_aux_loss_coef > 0.
            "router_aux_loss_coef": 0.0,
        }
        for k, v in replacements.items():
            x = f"{k}( = [^,\n]{{1,}})?,\n"
            y = f"'{v}'" if type(v) is str else f"{v}"
            y = f"{k} = {y},\n"
            arguments = re.sub(x, y, arguments)

    # Warn on too large or too small learning rate
    if "learning_rate" in call_args:
        learning_rate_check = (
            "if learning_rate < 1e-7: print(f'Unsloth: Your learning rate of `{learning_rate}` is too small and less than 1e-7! "
            "Consider increasing it, otherwise gradient updates will be close to 0!')\n"
            "if learning_rate > 1: print(f'Unsloth: Your learning rate of `{learning_rate}` is way too larger > 1! "
            "Consider decreasing it to 1e-1, otherwise gradient updates will explode!')\n"
        )
        extra_args += learning_rate_check

    # Fix num_train_epochs = None causing TypeError in Trainer.__init__
    # Trainer does `args.num_train_epochs > 0` which fails when None
    if "num_train_epochs" in call_args:
        num_train_epochs_check = (
            "if num_train_epochs is None:\n"
            "    num_train_epochs = 3.0  # Default to 3 epochs if None, max_steps will override\n"
        )
        extra_args += num_train_epochs_check

    # Check if max_seq_length is NOT defined (max_length is now default)
    if "max_seq_length" not in call_args and "max_length" in call_args:
        max_seq_length_pre = """max_seq_length : Optional[int] = field(
        default = None,
        metadata = {'help': 'Maximum sequence length to truncate to.'},
    )"""
        max_seq_length_call = "max_seq_length = None,"
        max_seq_length_post = "self.max_seq_length = max_seq_length"
    else:
        max_seq_length_pre = ""
        max_seq_length_call = ""
        max_seq_length_post = ""

    # Add output_dir saving
    if "output_dir" in call_args:
        # Default checks
        saving_check = (
            "if output_dir is None and save_strategy == 'steps' and save_steps == 500:\n"
            "    output_dir = 'unsloth_training_checkpoints'\n"
            "    save_strategy = 'no'\n"
        )
        extra_args += saving_check

    # Edit dataset_num_proc
    if "dataset_num_proc" in call_args:
        num_proc_check = (
            "import multiprocessing as _mp\n"
            "if dataset_num_proc is None:\n"
            "    if _mp.get_start_method() != 'fork':\n"
            "        dataset_num_proc = None\n"
            "    else:\n"
            "        import psutil\n"
            "        dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)\n"
            "        memory_gb_left = psutil.virtual_memory().available / (1024**3)\n"
            "        if memory_gb_left <= 2: dataset_num_proc = 1\n"
            "        else: dataset_num_proc = min(dataset_num_proc, int(memory_gb_left))\n"
        )
        extra_args += num_proc_check

    # Add padding if flex attention is added
    if "pad_to_multiple_of" in call_args:
        pad_to_multiple_of = (
            "if os.environ.get('UNSLOTH_ENABLE_FLEX_ATTENTION', '0') == '1':\n"
            "    from unsloth_zoo.flex_attention import HAS_FLEX_ATTENTION\n"
            "    if HAS_FLEX_ATTENTION and pad_to_multiple_of is None:\n"
            "        from unsloth_zoo.flex_attention import FLEX_ATTENTION_BLOCK_SIZE\n"
            "        pad_to_multiple_of = FLEX_ATTENTION_BLOCK_SIZE\n"
            "\n"
        )
        extra_args += pad_to_multiple_of

    # Check for loss_type = dr_grpo and scale_rewards for GRPO
    if "loss_type" in call_args and "scale_rewards" in call_args:
        # See https://github.com/huggingface/trl/issues/3130#issuecomment-2746947835
        # DAPO uses per token loss so BNPO loss used
        check_dr_grpo = (
            "if loss_type.lower() == 'dr_grpo':\n"
            "    loss_type = 'dr_grpo'\n"
            "elif loss_type.lower() == 'dapo':\n"
            "    loss_type = 'dapo'\n"
            "if loss_type.lower() == 'dr_grpo':\n"
            "    if scale_rewards == None:\n"
            "        scale_rewards = True\n"
            "    elif scale_rewards == True:\n"
            "        print('Unsloth: The Dr GRPO paper recommends setting `scale_rewards` to False! Will override. Set it to `None` to force False.')\n"
            "        scale_rewards = False\n"
            "elif loss_type.lower() == 'dapo':\n"
            "    if mask_truncated_completions != True:\n"
            "        print('Unsloth: The DAPO paper recommends `mask_truncated_completions = True` - we will set it.')\n"
            "    if epsilon_high != 0.28:\n"
            "        print('Unsloth: The DAPO paper recommends `epsilon_high = 0.28` - we will set it.')\n"
            "    if beta != 0.0:\n"
            "        print(f'[WARNING] Unsloth: The DAPO paper recommends setting `beta = 0.0` to remove the KL term - You have set it to {beta}.')\n"
            "    mask_truncated_completions = True\n"
            "    epsilon_high = 0.28\n"
            "\n"
        )
        extra_args += check_dr_grpo

    # Check GRPO num_generations mismatch
    if (
        "per_device_train_batch_size" in call_args
        and "num_generations" in call_args
        and "steps_per_generation" in call_args
        and "generation_batch_size" in call_args
    ):
        # if world size is not set by accelerate or torchrun at this point it will be 1
        check_num_generations = (
            "if steps_per_generation is None and generation_batch_size is None:\n"
            "    ga = gradient_accumulation_steps\n"
            "    world_size = int(os.environ.get('WORLD_SIZE', '1'))\n"
            "    if (ga * world_size * per_device_train_batch_size) % num_generations != 0:\n"
            "        print('Unsloth: We now expect `per_device_train_batch_size` * `gradient_accumulation_steps` * `world_size` to be a multiple of `num_generations`.\\n"
            "We will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations))\n"
            "        per_device_train_batch_size = num_generations\n"
            "\n"
        )
        extra_args += check_num_generations
    elif "per_device_train_batch_size" in call_args and "num_generations" in call_args:
        if "steps_per_generation" not in call_args:
            print(f"Unsloth: Could not find `steps_per_generation` in {trainer_file}")
        if "generation_batch_size" not in call_args:
            print(f"Unsloth: Could not find `generation_batch_size` in {trainer_file}")

        check_num_generations = (
            "if (per_device_train_batch_size // num_generations) * num_generations != per_device_train_batch_size:\n"
            "    print('Unsloth: We now expect `per_device_train_batch_size` to be a multiple of `num_generations`.\\n"
            "We will change the batch size of ' + str(per_device_train_batch_size) + ' to the `num_generations` of ' + str(num_generations))\n"
            "    per_device_train_batch_size = num_generations\n"
            "\n"
        )
        extra_args += check_num_generations

    # Check temperature must not be <= 0. Also stop if >= 10
    if "temperature" in call_args:
        check_temperature = (
            "if temperature <= 0:\n"
            "    raise ValueError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')\n"
            "elif temperature >= 10:\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

    # Edit config with anything extra
    if trainer_file in RL_CONFIG_CHANGES:
        process_extra_args = RL_CONFIG_CHANGES[trainer_file]
        for process_extra_arg in process_extra_args:
            extra_args += process_extra_arg(old_RLTrainer_source, old_RLConfig_source)

    # Create RLConfig args
    extra_args = extra_args.split("\n")
    extra_args = "\n".join(" " * 8 + x for x in extra_args)
    RLConfig_arguments = arguments
    RLConfig_extra_args = extra_args
    RLConfig_call_args = call_args

    # TRL 0.27.0+ forces use_reentrant=False in gradient_checkpointing_kwargs.
    # Unsloth gradient checkpointing requires use_reentrant=True, so we remove
    # the setting after super().__init__() when it gets auto-applied.
    RLConfig_post = ""
    if trl_version >= Version("0.27.0"):
        RLConfig_post = (
            "        # Unsloth: Remove use_reentrant=False forced by TRL 0.27.0+\n"
            "        if getattr(self, 'gradient_checkpointing_kwargs', None) is not None:\n"
            "            if 'use_reentrant' in self.gradient_checkpointing_kwargs:\n"
            "                del self.gradient_checkpointing_kwargs['use_reentrant']\n"
        )

    # Patch vLLM and other functions
    RLTrainer_extras = patch_functions(
        RLTrainer, trainer_file, RLTrainer_name, all_imports, imports
    )
    if RLTrainer_extras is None:
        RLTrainer_extras = f"_Unsloth{RLTrainer_name} = {RLTrainer_name}"

    # Create full module
    exec(f"from trl.trainer import ({RLTrainer_name}, {RLConfig_name},)")
    __RLTrainer_doc__ = eval(f"trl.trainer.{RLTrainer_name}").__doc__
    if __RLTrainer_doc__ is None:
        __RLTrainer_doc__ = ""
    __RLConfig_doc__ = eval(f"trl.trainer.{RLConfig_name}").__doc__
    if __RLConfig_doc__ is None:
        __RLConfig_doc__ = ""

    # Get all pre-modules
    if trainer_file in RL_PRE_ITEMS:
        RL_pre = "\n".join(RL_PRE_ITEMS[trainer_file])
    else:
        RL_pre = ""

    # Check if SamplingParams is in there
    if "SamplingParams" in old_RLTrainer_source:
        RL_pre = RL_pre + "\n" + inspect.getsource(vLLMSamplingParams)

    # Selective log softmax and other functions
    selective_log_softmax_code = inspect.getsource(selective_log_softmax)
    grpo_selective_log_softmax_code = inspect.getsource(grpo_selective_log_softmax)
    calculate_pad_tokens_in_prompt_code = inspect.getsource(calculate_pad_tokens_in_prompt)
    create_completion_attention_mask_code = inspect.getsource(create_completion_attention_mask)
    left_pack_padding_code = inspect.getsource(left_pack_padding)
    align_logprobs_with_mask_code = inspect.getsource(align_logprobs_with_mask)
    align_completion_tool_mask_code = inspect.getsource(align_completion_tool_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,
        __RLTrainer_doc__ = __RLTrainer_doc__,
        RLTrainer_arguments = RLTrainer_arguments,
        RLTrainer_extra_args = RLTrainer_extra_args,
        RLTrainer_call_args = RLTrainer_call_args,
        RLTrainer_kwargs = ",**kwargs"[1 if RLTrainer_call_args.endswith(",") else 0 :],
        RLConfig_name = RLConfig_name,
        __RLConfig_doc__ = __RLConfig_doc__,
        RLConfig_arguments = RLConfig_arguments,
        RLConfig_extra_args = RLConfig_extra_args,
        RLConfig_call_args = RLConfig_call_args,
        RLConfig_kwargs = ",**kwargs"[1 if RLConfig_call_args.endswith(",") else 0 :],
        RLConfig_post = RLConfig_post,
        RLTrainer_extras = RLTrainer_extras,
        RLTrainer_post = RLTrainer_post,
        RL_pre = RL_pre,
        max_seq_length_pre = max_seq_length_pre,
        max_seq_length_call = max_seq_length_call,
        max_seq_length_post = max_seq_length_post,
        selective_log_softmax_code = selective_log_softmax_code,
        grpo_selective_log_softmax_code = grpo_selective_log_softmax_code,
        calculate_pad_tokens_in_prompt_code = calculate_pad_tokens_in_prompt_code,
        create_completion_attention_mask_code = create_completion_attention_mask_code,
        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,
        align_completion_tool_mask_code = align_completion_tool_mask_code,
        sanitize_logprob_code = sanitize_logprob_code,
    )

    if RLTrainer_name == "GRPOTrainer":
        # Base torch_compile_options shared by all device types
        base_options = """torch_compile_options = {
            "epilogue_fusion"   : True,
            "max_autotune"      : False,
            "shape_padding"     : True,
            "trace.enabled"     : False,"""

        # Generate torch_compile_options based on device type
        if DEVICE_TYPE == "cuda":
            # CUDA-specific options (added to base options)
            cuda_options = """
            "triton.enable_persistent_tma_matmul": torch.cuda.get_device_capability()[0] >= 9,"""
            # cutlass options were added in PyTorch 2.8.0
            if torch_version >= Version("2.8.0"):
                cuda_options += """
            "cuda.cutlass_epilogue_fusion_enabled": torch.cuda.get_device_capability()[0] >= 9,
            "cuda.cutlass_tma_only": torch.cuda.get_device_capability()[0] >= 9,"""
            cuda_options += """
            "cuda.compile_opt_level"              : "-O2",
            "cuda.enable_cuda_lto"                : True,
        }"""
            new_options = base_options + cuda_options
        else:
            # XPU, HIP, and other device types use base options only
            new_options = (
                base_options
                + """
        }"""
            )

        pattern = r"torch_compile_options\s*=\s*\{[^}]*\}"

        RLTrainer_source = re.sub(pattern, new_options, RLTrainer_source, flags = re.DOTALL)

        if trl_version >= Version("1.4.0"):
            # The `elif is_peft_model(model) and args.beta != 0.0:` ref-adapter block
            # was introduced in TRL 1.4.0 and is used through 1.7.x. Remove only that
            # block, anchored on the final ref_param copy so we do NOT also swallow the
            # following gradient-checkpointing enable_input_require_grads() block.
            peft_pattern = (
                r"\s*elif is_peft_model\(model\) and args\.beta != 0\.0:"
                r".*?"
                r"ref_param\.data\.copy_\(param\.data\)"
            )

            replacement_comment = (
                "\n        # PEFT initialization logic removed via script for trl >= 1.4.0\n"
            )

            RLTrainer_source = re.sub(
                peft_pattern, replacement_comment, RLTrainer_source, flags = re.DOTALL
            )

            if trl_version >= Version("1.7.0"):
                # router_aux_loss_coef / aux_loss_enabled were added in TRL 1.7.0. Unsloth's
                # optimized GRPO forward cannot compute the MoE router aux loss, so reject
                # explicit opt-in (router_aux_loss_coef > 0) at init rather than silently ignoring it.
                RLTrainer_source = RLTrainer_source.replace(
                    "self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0",
                    "self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0\n"
                    '        if self.aux_loss_enabled: raise NotImplementedError("Unsloth GRPO does not compute the MoE router auxiliary loss; set router_aux_loss_coef = 0 (the Unsloth default).")',
                )

        elif trl_version >= Version("0.27.0"):
            peft_pattern = (
                r"\s*if is_peft_available\(\) and is_peft_model\(model\) and args\.beta != 0\.0:"
                r".*?"
                r"param\.data = param\.data\.to\(torch\.bfloat16\)"
            )

            replacement_comment = (
                "\n        # PEFT initialization logic removed via script for trl >= 0.27.0\n"
            )

            RLTrainer_source = re.sub(
                peft_pattern, replacement_comment, RLTrainer_source, flags = re.DOTALL
            )

        elif trl_version >= Version("0.26.0"):
            peft_block_pattern = (
                r"\s*if is_peft_available\(\) and isinstance\(model, PeftModel\) and peft_config is not None:"
                r".*?"
                r"param\.data = param\.data\.to\(torch\.bfloat16\)"
            )

            RLTrainer_source = re.sub(
                peft_block_pattern,
                "\n        # TRL PEFT 0.26.0 initialization logic removed on unsloth side.\n",
                RLTrainer_source,
                flags = re.DOTALL,
            )

    # Remove TRL 0.26.0's unconditional bfloat16 cast of trainable params. It
    # hardcodes bfloat16 for QLoRA, ignoring the user's dtype and breaking
    # GradScaler with fp16=True. Unsloth already handles adapter dtype via
    # patch_model_and_tokenizer, so the block is unnecessary (and already a
    # no-op for GRPO, whose peft init block is removed above).
    RLTrainer_source = RLTrainer_source.replace(
        'if getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False):',
        "if False:",
    )
    # TRL >= 1.7.0 spells the same QLoRA bf16 cast as `if _is_quantized_model:`.
    RLTrainer_source = RLTrainer_source.replace(
        "if _is_quantized_model:",
        "if False:",
    )

    if RLTrainer_name == "SFTTrainer":
        original_text = (
            'self._signature_columns = ["input_ids", "attention_mask", "completion_mask"]'
        )
        new_text = (
            'self._signature_columns = ["input_ids", "attention_mask", "completion_mask","labels"]'
        )
        RLTrainer_source = RLTrainer_source.replace(original_text, new_text)

        # Do NOT override _is_vlm -- let TRL detect VLM models naturally
        # (forcing _is_vlm=False errors on vision datasets in TRL 0.27.1+).
        # But some notebooks pass a bare tokenizer as processing_class, so TRL
        # sets _is_vlm=False even for VLMs; add an architecture-based override
        # before the validation check.
        _vlm_check_original = (
            '        self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample\n'
            "        if self._is_vision_dataset and not self._is_vlm:"
        )
        _vlm_check_patched = (
            '        self._is_vision_dataset = "image" in dataset_sample or "images" in dataset_sample\n'
            "        # Unsloth: override _is_vlm for VLM models that pass a bare tokenizer\n"
            "        if not self._is_vlm and self._is_vision_dataset:\n"
            "            _m = model\n"
            '            if hasattr(_m, "model"): _m = _m.model\n'
            '            if hasattr(getattr(_m, "config", None), "vision_config") or \\\n'
            '               _m.__class__.__name__.endswith("ForConditionalGeneration"):\n'
            "                self._is_vlm = True\n"
            "        if self._is_vision_dataset and not self._is_vlm:"
        )
        if _vlm_check_original in RLTrainer_source:
            RLTrainer_source = RLTrainer_source.replace(_vlm_check_original, _vlm_check_patched)

        # Fix TRL 0.22.x: VLM models with text-only datasets. It checks _is_vlm
        # (model type), not _is_vision_dataset (added in 0.25.1+); with
        # _is_vlm=True the vision-only signature columns don't overlap tokenized
        # text columns. Fix: merge both column sets into the VLM branch. Extra
        # columns are ignored by _remove_unused_columns (raises only on zero match).
        _sig_vlm_old = 'self._signature_columns = ["messages", "prompt", "completion", "images"]'
        _sig_vlm_new = (
            'self._signature_columns = ["messages", "prompt", "completion", "images",'
            ' "input_ids", "labels", "attention_mask", "seq_lengths", "completion_mask", "assistant_masks"]'
        )
        RLTrainer_source = RLTrainer_source.replace(_sig_vlm_old, _sig_vlm_new)

        # Inject model reference before _prepare_dataset for dynamic
        # token_type_ids detection in sft_prepare_dataset
        _prep_pattern = r"([ \t]*)train_dataset = self\._prepare_dataset\("
        _prep_replacement = (
            r"\1self._unsloth_model_ref = model\n\1train_dataset = self._prepare_dataset("
        )
        RLTrainer_source = re.sub(_prep_pattern, _prep_replacement, RLTrainer_source, count = 1)

    # Silence TRL's noisy batch_size=1 + padding-free warning (handles both
    # the original "anihilate" typo and the corrected "annihilate" spelling)
    for _typo in ("anihilate", "annihilate"):
        _idx = RLTrainer_source.find(_typo)
        if _idx == -1:
            continue
        # Walk backwards to find "if args.per_device_train_batch_size"
        _block_start = RLTrainer_source.rfind("if args.per_device_train_batch_size == 1", 0, _idx)
        if _block_start == -1:
            continue
        # Walk backwards to the newline before the if
        _line_start = RLTrainer_source.rfind("\n", 0, _block_start)
        # Walk forwards past the closing paren to the end of the block
        _close = RLTrainer_source.find(")", _idx)
        if _close == -1:
            continue
        _block_end = RLTrainer_source.find("\n", _close)
        if _block_end == -1:
            continue
        RLTrainer_source = RLTrainer_source[:_line_start] + RLTrainer_source[_block_end:]
        break

    # Remove multiple doc strings
    if __RLConfig_doc__ != "" and RLTrainer_source.count(__RLTrainer_doc__) == 2:
        RLTrainer_source = RLTrainer_source.replace(__RLTrainer_doc__, "", 1)

    # Remove multiple newlines
    RLTrainer_source = re.sub(r"[\n]{3,}", "\n", RLTrainer_source)

    # Create new function
    _resolved_module = _trainer_resolved_module or _config_resolved_module
    _model_location = (
        _resolved_module.__name__ if _resolved_module is not None else f"trl.trainer.{trainer_file}"
    )
    created_module = create_new_function(
        f"Unsloth{RLTrainer_name}",
        RLTrainer_source,
        _model_location,
        imports,
        overwrite = False,
    )
    patched_trainer = getattr(created_module, f"Unsloth{RLTrainer_name}")
    if trainer_file == "grpo_trainer":
        _patch_resume_from_checkpoint_memory(patched_trainer)

    # Patch Trainer
    exec(
        f"trl.{RLTrainer_name} = created_module.Unsloth{RLTrainer_name}",
        locals(),
        globals(),
    )
    exec(
        f"trl.trainer.{RLTrainer_name} = created_module.Unsloth{RLTrainer_name}",
        locals(),
        globals(),
    )
    exec(
        f"trl.trainer.{trainer_file}.{RLTrainer_name} = created_module.Unsloth{RLTrainer_name}",
        locals(),
        globals(),
    )

    # Patch Config
    exec(
        f"trl.{RLConfig_name} = created_module.Unsloth{RLConfig_name}",
        locals(),
        globals(),
    )
    exec(
        f"trl.trainer.{RLConfig_name} = created_module.Unsloth{RLConfig_name}",
        locals(),
        globals(),
    )
    exec(
        f"trl.trainer.{trainer_file}.{RLConfig_name} = created_module.Unsloth{RLConfig_name}",
        locals(),
        globals(),
    )
    try:
        config_module_name = trainer_file.replace("_trainer", "_config")
        config_module = importlib.import_module(f"trl.trainer.{config_module_name}")
        if hasattr(config_module, RLConfig_name):
            setattr(
                config_module,
                RLConfig_name,
                getattr(created_module, f"Unsloth{RLConfig_name}"),
            )
    except Exception:
        pass

    if trainer_file == "grpo_trainer":
        try:
            _wrap_grpo_generate_and_score(getattr(created_module, f"Unsloth{RLTrainer_name}"))
        except Exception as e:
            logger.info(
                f"Unsloth: Could not wrap _generate_and_score_completions for {RLTrainer_name}: {e}"
            )
        try:
            _wrap_grpo_hidden_states_fallback(getattr(created_module, f"Unsloth{RLTrainer_name}"))
        except Exception as e:
            logger.info(
                f"Unsloth: Could not wrap GRPO hidden-state fallback for {RLTrainer_name}: {e}"
            )


def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports):
    init = inspect.getsource(RLTrainer.__init__)
    old_init = init

    # Remove brackets in comments since it interferes ie (...)
    comments = re.findall(r"\#[^\n]{1,}\n", init)
    bracketed_comments = [x for x in comments if "(" in x or ")" in x]
    # Replace with [...] instead
    for bracketed_comment in bracketed_comments:
        init = init.replace(
            bracketed_comment,
            bracketed_comment.replace("(", "[").replace(")", "]"),
        )

    # Remove peft_config
    init = init.replace("elif peft_config is None:", "elif False:")
    init = init.replace("elif peft_config is not None:", "elif False:")
    init = init.replace("if peft_config is None:", "if False:")
    init = init.replace("if peft_config is not None:", "if False:")
    init = init.replace("get_peft_model(model, peft_config)", "model")
    # New TRL 0.20.0
    init = init.replace(
        "if peft_config is not None or (is_peft_available() and isinstance(model, PeftModel)):",
        "if False:",
    )
    # New TRL 0.20.0
    init = init.replace("model = self._prepare_peft_model(model, peft_config, args)\n", "pass\n")
    # TRL 0.22.0+ uses prepare_peft_model as a standalone function
    init = init.replace("model = prepare_peft_model(model, peft_config, args)", "pass")

    # Skip add_adapter("ref") for reference model computation
    # Unsloth: We comment out the "ref" adapter creation because:
    # 1. We want to use the original BASE MODEL as the reference model, not the SFT/LoRA model
    # 2. PEFT doesn't allow multiple adapters when target_parameters is used (MoE models)
    # When "ref" is not in peft_config, GRPO/RLOO fallback uses disable_adapter()
    # which gives the base model logits - exactly what we want
    add_adapter_block_pattern = (
        r"([ \t]*)"  # Capture leading indentation
        r"if\s+is_peft_available\(\)\s+and\s+is_peft_model\(model\)\s+and\s+args\.beta\s*!=\s*0\.0\s*:"
        r"(.*?)"  # Match the entire block until ref_param.data.copy_
        r"ref_param\.data\.copy_\(param\.data\)"
    )

    def comment_out_block(match):
        """Comment out each line in the matched block, preserving indentation."""
        full_match = match.group(0)
        indent = match.group(1)
        lines = full_match.split("\n")
        commented_lines = []
        # Add explanation comment first
        commented_lines.append(
            f"{indent}# Unsloth: Commented out - use base model as reference, not SFT/LoRA model"
        )
        # Comment out each line - insert # after leading whitespace to preserve indentation
        for line in lines:
            if line.strip():
                stripped = line.lstrip()
                leading_ws = line[: len(line) - len(stripped)]
                commented_lines.append(f"{leading_ws}# {stripped}")
            else:
                commented_lines.append(line)
        return "\n".join(commented_lines)

    init = re.sub(add_adapter_block_pattern, comment_out_block, init, flags = re.DOTALL)

    # Set use_vllm if not set
    if "args.use_vllm" in init and "model" in init and "args" in init:
        # .*? matches first match. .+? matches final match.
        replacer = re.findall(
            r"def __init__\(.*?\).*?\:\n",
            init,
            flags = re.MULTILINE | re.DOTALL,
        )
        if len(replacer) != 0:
            replacer = replacer[0]
            vllm_setter = (
                "\n"
                + " " * 8
                + "if hasattr(model, 'vllm_engine') and hasattr(args, 'use_vllm'):\n"
                + " " * 12
                + "if (getattr(args, 'use_vllm', False) == False):\n"
                + " " * 16
                + "args.use_vllm = True\n"
            )
            # " " * 16 + "args.vllm_importance_sampling_correction = True\n" + \
            # " " * 16 + "args.vllm_importance_sampling_cap = 2.0\n"

            if "grpo" in trainer_file and trl_version >= Version("0.18.0"):
                # If model has vllm_engine, then use vllm in colocate mode. Donot wait for server
                vllm_setter += " " * 12 + "args.vllm_mode='colocate'\n"
                if trl_version >= Version("0.23.0"):
                    # Align TRL sleep mode with the engine's actual enable_sleep_mode
                    # (the vision standby gate may have disabled it); fall back to the
                    # standby env var when the engine cannot be introspected.
                    vllm_setter += (
                        " " * 12
                        + "_unsloth_esm = getattr(getattr(getattr(getattr(model.vllm_engine, 'llm_engine', None), 'vllm_config', None), 'model_config', None), 'enable_sleep_mode', None)\n"
                        + " " * 12
                        + "if (_unsloth_esm if _unsloth_esm is not None else os.environ.get('UNSLOTH_VLLM_STANDBY', '0') != '0'):\n"
                        + " " * 16
                        + "args.vllm_enable_sleep_mode=True\n"
                    )

            init = init.replace(replacer, replacer + vllm_setter)

    # breakpoint()

    vllm_part = re.findall(
        r"(\n[\s]{8}" r"if (self|args)\.use_vllm\:.*?" r"\n[\s]{8}" "else:\n)",
        init,
        flags = re.MULTILINE | re.DOTALL,
    )

    if len(vllm_part) == 1:
        vllm_part, args = vllm_part[0][0], vllm_part[0][1]
        # Strip all comments
        new_vllm_part = re.sub(
            r"^\s*\#[^\n]*\n?", "", vllm_part, flags = re.MULTILINE
        )  # to also remove whole comment line instead of just starting at #
        new_vllm_part = re.sub(
            r"\s*\#.*$", "", new_vllm_part, flags = re.MULTILINE
        )  # remove comments that occur after code

        # Get SamplingParams
        sampling_params = re.findall(
            r"\n[\s]{4,}(self\.[^\s]{1,}[\s]{0,}\=[\s]{0,}SamplingParams\(.+?\))",
            new_vllm_part,
            flags = re.MULTILINE | re.DOTALL,
        )

        if len(sampling_params) == 1:
            sampling_params = sampling_params[0]
            # Fix guided_decoding
            sampling_params = sampling_params.replace(
                "guided_decoding=guided_decoding,",
                "guided_decoding="
                'GuidedDecodingParams(backend="outlines", regex=args.vllm_guided_decoding_regex) '
                'if getattr(args, "vllm_guided_decoding_regex", None) is not None else None,',
            )
            # Replace with our vLLM engine when sharing weights
            sampling_params = (
                " " * 12
                + "if getattr(getattr(model, 'vllm_engine', None), 'shared_weights', False): "
                + "self.llm = model.vllm_engine; self._last_loaded_step = 0\n"
                + " " * 12
                + sampling_params
            )

            # count the indentation of last line of sampling_params.
            splitted_sampling_params = sampling_params.split("\n")
            if len(splitted_sampling_params) >= 2:
                last_line = splitted_sampling_params[-1]
                last_prev_line = splitted_sampling_params[-2]
                last_prev_indentation = len(last_prev_line) - len(last_prev_line.lstrip())
                last_indentation = len(last_line) - len(last_line.lstrip())

                # Add extra arguments to SamplingParams
                extra = "**getattr(getattr(args, 'vllm_sampling_params', vLLMSamplingParams()), '_set_kwargs', {})"
                # Backwards replace
                to_replace = (
                    ",\n"
                    + " " * last_prev_indentation
                    + extra
                    + ",\n"
                    + " " * last_indentation
                    + ")"
                )
                sampling_params = to_replace.join(sampling_params.rsplit(")", 1))
                # Strip multiple commas
                sampling_params = re.sub(r"[\,][\s]{0,}\,", ",", sampling_params)

                new_vllm_part = (
                    f"\n{' ' * 8}if {args}.use_vllm:\n{sampling_params}" f"\n{' ' * 8}else:\n"
                )

        if trl_version >= Version("0.18.0"):
            # Guard LLM init - use existing vLLM engine when sharing weights,
            # otherwise keep the original LLM() creation for sync/reload path
            vllm_llm_init_pattern = r"(?P<indent>[ \t]*)self\.llm\s*=\s*LLM\(.*?\)*\)\s*?\n(?!,)"

            def guard_llm_init(match):
                indent = match.group("indent")
                original = match.group(0)
                return (
                    f"{indent}if getattr(getattr(model, 'vllm_engine', None), 'shared_weights', False):\n"
                    f"{indent}    self.llm = model.vllm_engine\n"
                    f"{indent}else:\n"
                    f"{indent}    {original.lstrip()}"
                )

            new_vllm_part = re.sub(
                vllm_llm_init_pattern,
                guard_llm_init,
                new_vllm_part,
                flags = re.DOTALL,
            )

        init = init.replace(vllm_part, new_vllm_part)

    # Search for vLLM calling in all child functions
    functions = dir(RLTrainer)
    RLTrainer_source = inspect.getsource(RLTrainer)
    functions = [x for x in functions if f"def {x}" in RLTrainer_source]

    changed = {
        "__init__": (
            old_init,
            init,
        )
    }
    edit_functions = RL_FUNCTIONS.get(trainer_file, [])

    for function in functions:
        if not hasattr(RLTrainer, function):
            continue
        if function in changed:
            original_source, source = changed[function]
        else:
            fx = getattr(RLTrainer, function)
            try:
                source = inspect.getsource(fx)
            except:
                continue
            original_source = source

        # Check for function
        for edit_function in edit_functions:
            source = edit_function(function, source)

        """
        import torch
        X = torch.ones((2, 2048, 201088), dtype = torch.bfloat16, device = "cuda")
        X[torch.randperm(2, dtype = torch.int64, device = X.device)]

        will error out in torch 2.8 AcceleratorError: CUDA error: invalid configuration argument
        """
        source = re.sub(
            r"(\n[\s]{4,})generation_batch = shuffle_sequence_dict\(generation_batch\)\n",
            r"\n\1try: generation_batch = shuffle_sequence_dict(generation_batch)\n\1except: pass\n",
            source,
        )

        # llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
        source = re.sub(
            r"(\n[\s]{4,}).+?model_executor\.driver_worker.+?\n",
            r"\n\1pass\n",
            source,
        )

        # llm_model.load_weights(model.state_dict().items())
        source = re.sub(
            r"(\n[\s]{4,}).+?load_weights\(.+?\n",
            r"\n\1pass\n",
            source,
        )

        # .state_dict()
        source = re.sub(
            r"\.state_dict\(\)",
            r"",
            source,
        )

        # Replace self.llm.generate and self.llm.chat with lora_request (only when sharing weights)
        if "CUDA_VISIBLE_DEVICES" in os.environ:
            lora_name = (
                trainer_file
                + "_lora_model_' + "
                + "(os.environ.get('CUDA_VISIBLE_DEVICES', '0').replace(',',''))"
            )
        else:
            lora_name = trainer_file + "_lora_model'"
        source = re.sub(
            r"(self\.llm\.(?:generate|chat)\([^\)]{1,})\)",
            r"\1, lora_request = self.model.load_lora('"
            + lora_name
            + r", load_tensors = True)"
            + r" if getattr(self.llm, 'shared_weights', False)"
            + r" else None)",
            source,
        )
        # All these are to fix multiple commas before lora_request (in case the original code ends with something like ",)")
        # https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L1388 for eg has such an ending
        source = re.sub(r"\,[\s]{1,}\,[\s]{0,}lora_request", ", lora_request", source)
        source = re.sub(r"[\s]{1,}\,[\s]{0,}lora_request", ", lora_request", source)
        source = re.sub(r"[\,]{1,}[\s]{0,}lora_request", ", lora_request", source)
        # Prefer using unsloth's sampling params and fallback to trl's if not found
        # We'll enable this later separately when combining both this and GRPOConfig params
        # source = re.sub(
        #     r"sampling_params\s*=\s*sampling_params",
        #     r"sampling_params = getattr(self.args, 'vllm_sampling_params', sampling_params)",
        #     source
        # )
        # Fix later versions of SamplingParams via grpo_update_SamplingParams
        source = source.replace(
            "sampling_params = SamplingParams(**generation_kwargs)",
            "sampling_params = SamplingParams("
            "**grpo_update_SamplingParams("
            "SamplingParams, generation_kwargs, "
            "getattr(self.args, 'vllm_sampling_params', None)"
            ")"
            ")",
        )

        # Skip if no changes done
        if source == original_source:
            continue

        # Find all imports
        imports += [x for x in all_imports if not x.startswith("_") and x in source]

        changed[function] = (
            original_source,
            source,
        )

    # Import all functions
    imports = list(set(imports))

    # Patch all functions
    for function in changed:
        old, new = changed[function]
        RLTrainer_source = RLTrainer_source.replace(old, new)

    RLTrainer_source = RLTrainer_source.replace(
        f"class {RLTrainer_name}", f"class _Unsloth{RLTrainer_name}", 1
    )
    return RLTrainer_source


def patch_trl_rl_trainers():
    # Patch all TRL modules if they have vLLM or PEFT
    import trl.trainer

    all_trainers = dir(trl.trainer)
    all_trainers = [
        x for x in all_trainers if x.islower() and x.endswith("_trainer") and x != "base_trainer"
    ]
    for trainer in all_trainers:
        try:
            _patch_trl_rl_trainers(trainer)
        except Exception as e:
            logger.warning_once(f"Unsloth: Could not patch trl.trainer.{trainer}: {e}")
    return


def patch_trl_disable_gradient_checkpointing():
    # TRL 1.0.0+ wraps generation in:
    #   with torch.no_grad(), disable_gradient_checkpointing(self.model, ...):
    # The toggle only suppresses a cosmetic PyTorch warning; under no_grad it
    # has no functional effect. But on exit it calls
    # gradient_checkpointing_enable(), overwriting Unsloth's custom
    # "unsloth" wrapper -- for Gemma-4 this corrupts forward numerics and
    # blows GRPO KL divergence up to ~10^12 at step 1.
    #
    # Replacing the context manager with a no-op preserves Unsloth's wrapper.
    # trl < 1.0.0 (no disable_gradient_checkpointing): early return.
    # trl >= 1.0.0: noop is correct; only loss is the cosmetic warning.
    try:
        import trl.models.utils as _tmu
    except ImportError:
        return
    if not hasattr(_tmu, "disable_gradient_checkpointing"):
        return
    if getattr(
        _tmu.disable_gradient_checkpointing,
        "_unsloth_noop_patched",
        False,
    ):
        return

    @contextmanager
    def _noop_disable_gradient_checkpointing(model, gradient_checkpointing_kwargs = None):
        yield

    _noop_disable_gradient_checkpointing._unsloth_noop_patched = True

    _tmu.disable_gradient_checkpointing = _noop_disable_gradient_checkpointing

    # Also rebind any trl.* module that already imported the symbol by
    # reference (cached at import time). Walk sys.modules dynamically so this
    # catches every trainer doing
    # `from ...models.utils import disable_gradient_checkpointing`.
    for _mod_name, _mod in list(sys.modules.items()):
        if _mod is None or not _mod_name.startswith("trl."):
            continue
        try:
            _bound = getattr(_mod, "disable_gradient_checkpointing", None)
        except (AttributeError, ImportError):
            continue
        if _bound is None:
            continue
        try:
            setattr(
                _mod,
                "disable_gradient_checkpointing",
                _noop_disable_gradient_checkpointing,
            )
        except (AttributeError, TypeError):
            pass

    if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1":
        logger.warning_once(
            "Unsloth: Patched trl.models.utils.disable_gradient_checkpointing with "
            "a no-op to preserve Unsloth gradient checkpointing across TRL "
            "generation passes."
        )
    return


def patch_trl_openenv():
    for function in RL_ADDITIONAL_FUNCTIONS["openenv"]:
        logger.info(f"Unsloth: Patching trl openenv with function: {function.__name__}")
        function()  # Call the function to apply the patch
    return


def patch_trl_vllm_generation():
    # trl moved vllm stuff to trl/generation/vllm_generation.py
    # We need to min_p patch it to not instantiate another vLLM instance if we already have one with fast_inference
    # Find the instance of self.llm = LLM(..) (multiline) and wrap it around an if clause
    for function in RL_ADDITIONAL_FUNCTIONS["vllm_generation"]:
        logger.info(f"Unsloth: Patching trl VLLMGeneration with function: {function.__name__}")
        function()
    return


def patch_trl_vllm_generation():
    # trl moved vllm stuff to trl/generation/vllm_generation.py
    # We need to min_p patch it to not instantiate another vLLM instance if we already have one with fast_inference
    # Find the instance of self.llm = LLM(..) (multiline) and wrap it around an if clause
    for function in RL_ADDITIONAL_FUNCTIONS["vllm_generation"]:
        logger.info(f"Unsloth: Patching trl VLLMGeneration with function: {function.__name__}")
        function()
    return


def PatchFastRL(algorithm = None, FastLanguageModel = None):
    if FastLanguageModel is not None:
        PatchRL(FastLanguageModel)
    # Under UNSLOTH_ALLOW_CPU=1 (CPU-only CI), skip TRL trainer rewriting so
    # downstream `inspect.getsource(trl.SFTTrainer)` drift detectors see the
    # pristine upstream class, not the compiled Unsloth* wrappers.
    if os.environ.get("UNSLOTH_ALLOW_CPU", "0") == "1":
        return
    # Install the disable_gradient_checkpointing noop BEFORE
    # patch_trl_rl_trainers, which imports extra trl.* submodules; any module
    # imported after the sys.modules walk would keep the original broken
    # binding. Installing first ensures the canonical symbol is replaced before
    # those submodules bind it.
    patch_trl_disable_gradient_checkpointing()
    patch_trl_rl_trainers()
    patch_trl_openenv()
    patch_trl_vllm_generation()
    if type(algorithm) is str and algorithm.islower():
        PatchRLStatistics(algorithm)