From 867fe636a36612fefde0ecf4f342741556a43a1d Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Fri, 22 May 2026 18:07:25 +0530 Subject: [PATCH] Respect GC for GRPO (#5269) * Respect GC for GRPO * Preserve gradient_checkpointing across post-generate training-mode restores Two sibling generation paths put the model into inference mode and then unconditionally restored training with the for_training default, which re-enabled gradient checkpointing even when the caller had it disabled: - unsloth/models/rl.py: unsloth_unwrap_model_for_generation, installed onto every TRL *_trainer module that exposes unwrap_model_for_generation. - unsloth/models/llama.py: unsloth_fast_generate, bound onto model.generate. Snapshot the active gradient_checkpointing state from the model modules before for_inference clears it, then thread the snapshot through the matching for_training call. Same one-line restore semantics already used by prepare_for_training_mode and the GRPO replacement at rl_replacements.py. The for_training(...) call on each line is preserved; only the kwarg is added. The pre-existing post-generate guards (the conditional restore in unsloth_fast_generate and the finally restore in unsloth_unwrap_model_for_generation) continue to run unchanged. * Snapshot pre-disable, preserve unsloth smart-GC mode across generation restores Two follow-ups to the post-generate gradient_checkpointing restore: 1. unsloth/models/rl.py: TRL's _unwrap_model_for_generation calls unwrapped_model.gradient_checkpointing_disable() before yielding (trl/models/utils.py:124-127 in 0.22.2, 0.27.1, and 1.3.0). The previous snapshot was taken inside the with-block and therefore read the post-disable state, restoring for_training with use_gradient_checkpointing=False even when the caller had it on. Move the snapshot above the with-block so it observes the caller's pre-disable configuration. 2. unsloth/models/{rl.py,llama.py}: any(getattr(m, "gradient_checkpointing")) collapses Unsloth's smart-GC mode value "unsloth" (a documented loader default at unsloth/models/_utils.py:212 and unsloth/models/llama.py 2824/3314, loader.py:248/854) into a plain True. After generation, the restore would silently downgrade "unsloth" smart GC to standard HF GC. Replace any() with a value-preserving next((v for ... if v), False) so the actual mode value survives the round-trip. The for_training(...) calls on each line are preserved; only the snapshot expression and its position change. The pre-existing post-generate restore guards continue to run unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/llama.py | 18 +++++++++++++++++- unsloth/models/rl.py | 19 ++++++++++++++++++- unsloth/models/rl_replacements.py | 2 +- 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 20b515c711..b2e6a45aa0 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2071,6 +2071,19 @@ def unsloth_fast_generate( ): # If the model starts out in training mode, restore training mode after generation restore_training_mode = self.training + # why: snapshot the actual GC mode value (e.g. "unsloth") before for_inference + # clears it, so the post-generate restore preserves the caller's configured GC + # mode rather than collapsing it to a plain bool. + use_gradient_checkpointing = next( + ( + v + for v in ( + getattr(m, "gradient_checkpointing", False) for m in self.modules() + ) + if v + ), + False, + ) FastLlamaModel.for_inference(self) @@ -2156,7 +2169,10 @@ def unsloth_fast_generate( # pass if restore_training_mode: - FastLlamaModel.for_training(self) + FastLlamaModel.for_training( + self, + use_gradient_checkpointing = use_gradient_checkpointing, + ) return output diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 8521b72602..7b7c3ac1a4 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -186,6 +186,20 @@ def PatchRL(FastLanguageModel): @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) @@ -207,7 +221,10 @@ def PatchRL(FastLanguageModel): finally: # Restore generate and return unwrapped_model.generate = original_generate - FastLanguageModel.for_training(model) + FastLanguageModel.for_training( + model, + use_gradient_checkpointing = use_gradient_checkpointing, + ) from transformers import Trainer from transformers.trainer_pt_utils import nested_detach diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 53ef25f95c..0f9a324d5b 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -723,7 +723,7 @@ def grpo_trainer__generate_and_score_completions(function_name, function): # Left pad prompt before calculation old and ref hidden states left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) max_left_pad = torch.max(left_pad_tokens_per_prompt).item() - self.model.for_training()""" + self.model.for_training(use_gradient_checkpointing=getattr(self.args, 'gradient_checkpointing', True))""" function = function.replace(line_to_replace, replacement_lines)