set use_reentrant=False

2024-01-19 23:29:54 +08:00 · 2024-01-19 23:29:54 +08:00 · 0ff9a1fb4f
parent 12043aab9c
commit 0ff9a1fb4f
1 changed files with 1 additions and 2 deletions
--- a/src/llmtuner/model/patcher.py
+++ b/src/llmtuner/model/patcher.py
@ -205,8 +205,7 @@ def _prepare_model_for_training(
        if not getattr(model, "supports_gradient_checkpointing", False):
            logger.warning("Current model does not support gradient checkpointing.")
        else:
-            model.enable_input_require_grads()
-            model.gradient_checkpointing_enable()
+            model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
            model.config.use_cache = False # turn off when gradient checkpointing is enabled
            logger.info("Gradient checkpointing enabled.")