📍 Disable caching when grad checkpointing enable in GRPO (#2653)

* disable caching when grad checkpointing * style
huggingface · Jan 25, 2025 · 807046b · 807046b
1 parent 317d2d4
commit 807046b
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
@@ -172,6 +172,10 @@ def __init__(
                     "Invalid `torch_dtype` passed to `GRPOConfig`. Expected either 'auto' or a string representing "
                     f"a `torch.dtype` (e.g., 'float32'), but got {torch_dtype}."
                 )
+            # Disable caching if gradient checkpointing is enabled (not supported)
+            model_init_kwargs["use_cache"] = (
+                False if args.gradient_checkpointing else model_init_kwargs.get("use_cache")
+            )
             model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)
         else:
             model_id = model.config._name_or_path