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 <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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
 

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

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)