fix: make sure to pass kwargs to enable checkpoint

axolotl-ai-cloud · Mar 16, 2024 · a914cb3 · a914cb3
1 parent 34211ff
commit a914cb3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -888,7 +888,7 @@ def load_and_quantize_parallel(name_param, model, **kwargs):
 
     if cfg.adapter in ["lora", "qlora"]:
         if cfg.gradient_checkpointing:
-            model.gradient_checkpointing_enable()
+            model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=cfg.gradient_checkpointing_kwargs)
         if (
             cfg.load_in_8bit or cfg.load_in_4bit
         ) and not skip_prepare_model_for_kbit_training: