From 53af4a1b3e78f2d1cac9401fe15baf8720fc2303 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 10 Apr 2026 07:58:15 -0700 Subject: [PATCH] Fix Gemma-4 GRPO catastrophic KL divergence with TRL 1.0.0+ (#4934) * Fix Gemma-4 GRPO catastrophic KL divergence with TRL 1.0.0+ Two compounding bugs caused Gemma-4 GRPO training to diverge with KL ~10^12 at step 1 against TRL 1.0.0+. Both fixes are runtime patches in the existing TRL/model patch flow and are no-ops for models and TRL versions that are not affected. Fix 1 (rl.py): replace trl.models.utils.disable_gradient_checkpointing with a no-op context manager. TRL 1.0.0+ wraps generation in `with torch.no_grad(), disable_gradient_checkpointing(self.model, ...):` purely to suppress a cosmetic PyTorch warning ("None of the inputs have requires_grad=True"). Inside torch.no_grad() the gradient checkpointing state has no functional effect on the forward pass. On context exit, TRL calls model.gradient_checkpointing_enable() which dispatches to HF's generic implementation and overwrites Unsloth's custom `use_gradient_checkpointing="unsloth"` wrapper, corrupting Gemma-4 forward numerics. Replacing the toggle with a no-op preserves Unsloth's custom GC wrapper across generation passes. The patch walks sys.modules dynamically to also rebind the symbol on every trl.* module that already imported it (grpo_trainer, dpo_trainer, rloo_trainer, dppo_trainer, gfpo_trainer, grpo_with_replay_buffer_trainer, and any future trainer module). Fix 2 (vision.py): inject `final_logit_softcapping` from `config.text_config` into the top-level `model.config` for multimodal models. Unsloth's GRPO trainer reads `getattr(model.config, "final_logit_softcapping", 0)` but for Gemma-4 the attribute lives only on the nested `Gemma4TextConfig`, so the lookup silently defaults to 0 instead of 30. Backwards compatibility: - trl 0.22.2: no `disable_gradient_checkpointing` symbol exists, the patch early-returns via `hasattr` guard. - trl 0.27.1: same broken pattern as 1.0.0, the noop replacement is correct. - trl 1.0.0+: end-to-end verified on `unsloth/gemma-4-E2B-it` GRPO with TRL 1.0.0 and transformers 5.5.0. Step 1 loss=2.46e-08, kl=2.92e-05 (machine zero) vs broken baseline loss=1.37e+06, kl=1.76e+09. - Llama / non-VLM text models: Fix 2 is a no-op (no `text_config`); Fix 1 is functionally identical (Unsloth's GC wrapper is preserved). - Qwen3-VL and other VLMs without final_logit_softcapping: Fix 2 is a no-op (text_config.final_logit_softcapping is None). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply loop 1 review fixes for PR #4934 - Move Fix 2 from vision.py to rl_replacements.py:858 and :1110 at the actual consumer sites. This avoids mutating model.config (which could leak into save_pretrained output) and covers text-only Gemma-4 paths that do not flow through FastBaseModel.from_pretrained. - Revert the vision.py injection block entirely. - Narrow the bare except blocks in patch_trl_disable_gradient_checkpointing from `except Exception:` to `(AttributeError, ImportError)` and `(AttributeError, TypeError)` to avoid masking unrelated bugs. - Add logger.warning_once when the noop patch is installed, matching patch_trl_openenv and patch_trl_vllm_generation convention. - Remove the dead per-module `_unsloth_noop_patched` sentinel check inside the sys.modules walk. The function-level early return already covers this case. - Move `import sys` and `from contextlib import contextmanager` to the module-level imports instead of inside the function body. - Rewrite the ordering comment in PatchFastRL to accurately describe why patch_trl_disable_gradient_checkpointing must run before patch_trl_rl_trainers. - Fix keyword default spacing to match surrounding rl.py style. End-to-end verified: Gemma-4-E2B GRPO on TRL 1.0.0 + transformers 5.5.0 step 1 loss=2.464e-08 kl=2.921e-05, all 5 steps succeed. * Apply loop 2 review fix for PR #4934 Extract the final_logit_softcapping fallback logic into a shared helper `_unsloth_get_final_logit_softcapping(config)` defined in rl_replacements.py and injected into the compiled cache via RL_PRE_ITEMS["grpo_trainer"]. Both call sites (`grpo_trainer__generate_and_score_completions` and `grpo_trainer_compute_loss`) now use the helper instead of inlining the same text_config fallback block twice. Verified: compiled cache file lists the helper at module scope and both consumer sites call it. Gemma-4-E2B GRPO step 1 loss=2.464e-08 kl=2.921e-05 (unchanged), all 5 steps pass. * Apply loop 3 review fix for PR #4934 Extend _unsloth_get_final_logit_softcapping to also fall back to config.get_text_config() for composite configs such as T5GemmaConfig where the text sub-config is not exposed via the text_config attribute but only via the get_text_config() method. Guard against (TypeError, ValueError) raised by ambiguous composite configs, and skip the self-referential case where get_text_config() returns self. This addresses the 6/7 reviewer consensus from the third review loop. Verified: - Helper returns 30.0 for Gemma-4, T5Gemma, and Gemma 1/2 configs. - Helper returns 0 for Llama, Qwen, Mistral, Cohere, Granite, and ambiguous configs raising ValueError. - Gemma-4-E2B GRPO step 1 loss=2.464e-08 kl=2.921e-05 (unchanged). - Llama-3.2-1B GRPO all 5 steps loss=0 kl=0 (no regression). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/rl.py | 87 +++++++++++++++++++++++++++++++ unsloth/models/rl_replacements.py | 35 ++++++++++--- 2 files changed, 116 insertions(+), 6 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 5651a7da41..f444f6bd37 100755 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -22,6 +22,8 @@ from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union 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 @@ -1947,6 +1949,83 @@ def patch_trl_rl_trainers(): 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 exists only to suppress a cosmetic PyTorch warning + # ("None of the inputs have requires_grad=True"). Inside torch.no_grad() + # the gradient checkpointing state has no functional effect on the + # forward pass. + # + # On exit, the context manager calls model.gradient_checkpointing_enable() + # which dispatches to HuggingFace's generic implementation and overwrites + # Unsloth's custom `use_gradient_checkpointing="unsloth"` wrapper. For + # Gemma-4 (and likely other models) this corrupts the forward numerics + # enough to make GRPO KL divergence explode to ~10^12 at step 1. + # + # Replacing the context manager with a no-op preserves Unsloth's custom + # gradient checkpointing wrapper across generation/inference passes. + # + # Backwards compatibility: + # - trl < 1.0.0 (no disable_gradient_checkpointing): early return. + # - trl >= 1.0.0: noop is functionally equivalent for forward + # correctness. The only loss is a cosmetic warning being emitted + # by PyTorch when use_reentrant=True (which is exactly the warning + # TRL added the toggle to suppress in the first place). + 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, so the noop applies even when the trainer module cached the + # original at import time. We walk sys.modules dynamically rather than + # hardcoding a list, so this picks up every trainer that does + # `from ...models.utils import disable_gradient_checkpointing` + # (grpo, dpo, rloo, dppo, gfpo, grpo_with_replay_buffer, and any future + # TRL trainer module). + 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 + + 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__}") @@ -1981,6 +2060,14 @@ def patch_trl_vllm_generation(): def PatchFastRL(algorithm = None, FastLanguageModel = None): if FastLanguageModel is not None: PatchRL(FastLanguageModel) + # Install the disable_gradient_checkpointing noop BEFORE + # patch_trl_rl_trainers. patch_trl_rl_trainers imports extra trl.* trainer + # submodules while generating the compiled cache; any new trl.* modules + # imported after the sys.modules walk would keep their original (broken) + # binding of disable_gradient_checkpointing. Running the noop install + # first ensures the canonical trl.models.utils symbol is already replaced + # before those submodules bind it. + patch_trl_disable_gradient_checkpointing() patch_trl_rl_trainers() patch_trl_openenv() patch_trl_vllm_generation() diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 2544afe82e..93a7f89bcb 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -855,9 +855,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_sizes_chunks = chunk_optional(image_sizes, B) temperature = self.temperature - logit_softcapping = getattr(model.config, "final_logit_softcapping", 0) - if logit_softcapping is None: - logit_softcapping = 0 + logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) logit_scale_multiply = getattr(model.config, "logit_scale", 0) if logit_scale_multiply is None: logit_scale_multiply = 0 @@ -1004,11 +1002,38 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entropies) + +def _unsloth_get_final_logit_softcapping(config): + """Return final_logit_softcapping for a model config, falling back to the + nested text sub-config for composite models. Handles both: + - Gemma-4-style configs where the attribute lives on ``config.text_config`` + - T5Gemma-style composite configs where the text sub-config is only + reachable via ``config.get_text_config()`` + Returns 0 if unset, matching the previous behaviour. + """ + softcap = getattr(config, "final_logit_softcapping", None) + if softcap is None: + text_cfg = getattr(config, "text_config", None) + if text_cfg is None: + get_text_config = getattr(config, "get_text_config", None) + if callable(get_text_config): + try: + text_cfg = get_text_config() + except (TypeError, ValueError): + text_cfg = None + if text_cfg is not None and text_cfg is not config: + softcap = getattr(text_cfg, "final_logit_softcapping", None) + return 0 if softcap is None else softcap + + grpo_compute_loss = RL_REPLACEMENTS["grpo_compute_loss"] grpo_compute_loss_slow = RL_REPLACEMENTS["grpo_compute_loss_slow"] UnslothEfficientGRPO = RL_REPLACEMENTS["UnslothEfficientGRPO"] grpo_accumulated_loss = RL_REPLACEMENTS["grpo_accumulated_loss"] grpo_update_SamplingParams = RL_REPLACEMENTS["grpo_update_SamplingParams"] +RL_PRE_ITEMS["grpo_trainer"].append( + inspect.getsource(_unsloth_get_final_logit_softcapping) +) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_compute_loss)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(UnslothEfficientGRPO)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_accumulated_loss)) @@ -1107,9 +1132,7 @@ def grpo_trainer_compute_loss(function_name, function): input_ids = input_ids[:, -logits_to_keep:] # Get logit softcapping and logit scale - logit_softcapping = getattr(model.config, "final_logit_softcapping", 0) # Gemma - if logit_softcapping is None: - logit_softcapping = 0 + logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) # Gemma logit_scale_multiply = getattr(model.config, "logit_scale", 0) # Cohere if logit_scale_multiply is None: logit_scale_multiply = 0