Fix Gemma-4 GRPO catastrophic KL divergence with TRL 1.0.0+

[danielhanchen]

Summary

Fixes Gemma-4 GRPO training diverging with KL ~10^12 at step 1 against TRL 1.0.0+, by adding two runtime patches to the existing TRL/model patch flow. Both patches are no-ops for models and TRL versions that are not affected.

The bugs

Bug 1 (primary): TRL disable_gradient_checkpointing overwrites Unsloth's custom GC wrapper

TRL 1.0.0+ wraps generation in:

with torch.no_grad(), disable_gradient_checkpointing(self.model, self.args.gradient_checkpointing_kwargs):

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 context exit, TRL 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 the training-step forward diverge from the reference forward, producing KL ~10^12 at step 1.

Bug 2 (secondary): final_logit_softcapping lookup misses for multimodal Gemma-4

UnslothGRPOTrainer reads getattr(model.config, "final_logit_softcapping", 0). For Gemma4ForConditionalGeneration the attribute lives only on the nested Gemma4TextConfig, so the lookup silently defaults to 0 instead of 30. Both ref and policy paths hit the same bug for LoRA so KL cancels, but full fine-tuning with a separate ref_model produces numerically incorrect logps.

The fix

Fix 1: unsloth/models/rl.py - new patch_trl_disable_gradient_checkpointing()

Replaces trl.models.utils.disable_gradient_checkpointing with a no-op context manager. The patch dynamically walks sys.modules for any trl.* module that already imported the symbol by reference and rebinds it, so it picks up:

• trl.trainer.grpo_trainer
• trl.trainer.dpo_trainer
• trl.trainer.rloo_trainer
• trl.experimental.dppo.dppo_trainer
• trl.experimental.gfpo.gfpo_trainer
• trl.experimental.grpo_with_replay_buffer.grpo_with_replay_buffer_trainer
• and any future TRL trainer module

The patch is wired into PatchFastRL BEFORE patch_trl_rl_trainers so the compiled cache picks up the noop at its from trl.trainer.grpo_trainer import disable_gradient_checkpointing binding time.

Fix 2: unsloth/models/vision.py - inject final_logit_softcapping from text_config

In FastBaseModel.from_pretrained, after the model is loaded, lifts final_logit_softcapping from config.text_config (or config.get_text_config()) to the top-level model.config if and only if the top-level config does not already expose it. Skips silently for models that already have it or do not use softcapping.

Backwards compatibility

TRL version	Behavior
0.22.2	disable_gradient_checkpointing symbol does not exist. The hasattr guard early-returns. Verified by installing trl 0.22.2 in a clean venv and inspecting the symbol.
0.27.1	Same broken positional-arg pattern as 1.0.0. The noop replacement applies. Verified by installing trl 0.27.1 and running the patch logic against it.
1.0.0+	End-to-end verified on unsloth/gemma-4-E2B-it GRPO with TRL 1.0.0 and transformers 5.5.0.

Model	Behavior
Gemma-4 (E2B/E4B)	Fix 1 fixes the KL blow-up. Fix 2 lifts softcap=30 to the top-level config.
Llama / Qwen / text models	Fix 1 is functionally identical (Unsloth's GC wrapper is preserved). Fix 2 is a no-op (no text_config).
Qwen3-VL and other VLMs without softcapping	Fix 2 is a no-op (text_config.final_logit_softcapping is None).

Test plan

• Gemma-4-E2B GRPO with TRL 1.0.0 + transformers 5.5.0: train 5 steps, expect step 1 KL near zero
• Llama-3.2-1B GRPO with the patches applied: train 5 steps, expect identical numerics to baseline
• Verify trl 0.22.2 hits the hasattr early-return path
• Verify trl 0.27.1 patch logic replaces the symbol and the noop context manager works
• Verify the dynamic sys.modules walker covers grpo_trainer, dpo_trainer, rloo_trainer, dppo_trainer, gfpo_trainer, grpo_with_replay_buffer_trainer
• Confirm Gemma-4-E2B final_logit_softcapping is 30.0 after FastModel.from_pretrained and survives get_peft_model

Empirical numbers

Run	Step 1 loss	Step 1 KL
Gemma-4-E2B without fix	1.37e+06	1.76e+09
Gemma-4-E2B with fix	2.46e-08	2.92e-05
Llama-3.2-1B with fix	0	0

[gemini-code-assist]

Code Review

This pull request introduces a patch to disable TRL's disable_gradient_checkpointing context manager with a no-op, preventing numerical issues with Unsloth's custom gradient checkpointing, particularly for models like Gemma-4. This patch is applied early in PatchFastRL. Additionally, logic is added to from_pretrained for vision models to ensure final_logit_softcapping is correctly exposed at the top-level model.config for multimodal models, which is essential for the GRPO trainer. The review comments suggest improving error handling by logging exceptions instead of silently passing them in the newly added code blocks, adhering to the rule against broad, silent exception handlers.

[gemini-code-assist]

Silently passing on exceptions here can hide potential issues where the monkey-patching fails. As per the rule 'Avoid using broad, silent exception handlers like except Exception: pass. Instead, log the exception, even if at a debug level, to aid in future debugging.', it's better to log these exceptions, even at a debug level, to help with future troubleshooting. This will make the patching process more transparent without crashing on unexpected errors.

Suggested change

	setattr(
	_mod,
	except Exception as e:
	logger.debug(f"Unsloth: Could not patch disable_gradient_checkpointing on {_mod_name}: {e}")

References

1. Avoid using broad, silent exception handlers like except Exception: pass. Instead, log the exception, even if at a debug level, to aid in future debugging.

[gemini-code-assist]

This block uses multiple nested try...except Exception: pass statements, which can make the code harder to debug and understand. As per the rule 'Avoid using broad, silent exception handlers like except Exception: pass. Instead, log the exception, even if at a debug level, to aid in future debugging.', it's better to flatten the logic and log exceptions when they occur. This will improve maintainability and help diagnose issues if this patching logic fails for certain model configurations.

            try:
                _top_config = model.config
                if getattr(_top_config, "final_logit_softcapping", None) is None:
                    _softcap = None
                    _text_cfg = getattr(_top_config, "text_config", None)
                    if _text_cfg is not None:
                        _softcap = getattr(
                            _text_cfg, "final_logit_softcapping", None,
                        )
                    if _softcap is None:
                        _get_text = getattr(_top_config, "get_text_config", None)
                        if callable(_get_text):
                            try:
                                _softcap = getattr(
                                    _get_text(),
                                    "final_logit_softcapping",
                                    None,
                                )
                            except Exception as e:
                                logger.debug(f"Unsloth: Failed to get text_config for softcapping: {e}")
                    if _softcap is not None:
                        try:
                            setattr(_top_config, "final_logit_softcapping", _softcap)
                        except Exception as e:
                            logger.warning(f"Unsloth: Could not set final_logit_softcapping on model config: {e}")
            except Exception as e:
                logger.warning(f"Unsloth: An unexpected error occurred during final_logit_softcapping patch: {e}")

References

1. Avoid using broad, silent exception handlers like except Exception: pass. Instead, log the exception, even if at a debug level, to aid in future debugging.

[Datta0]

Yeah I too noticed trl messing up with gradient checkpointing. This seems like a necessary change