improve vram use w gradient checkpointing (#1167) [skip ci]

axolotl-ai-cloud · Jan 23, 2024 · 802f966 · 802f966
1 parent b8e5603
commit 802f966
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/axolotl/utils/config.py b/src/axolotl/utils/config.py
@@ -159,6 +159,13 @@ def normalize_config(cfg):
     if isinstance(cfg.pretraining_dataset, dict):
         cfg.pretraining_dataset = [cfg.pretraining_dataset]
 
+    if (
+        cfg.gradient_checkpointing
+        and cfg.unfrozen_parameters is None
+        and cfg.gradient_checkpointing_kwargs is None
+    ):
+        cfg.gradient_checkpointing_kwargs = {"use_reentrant": True}
+
     log_gpu_memory_usage(LOG, "baseline", cfg.device)