make sure to set gradient checkpointing too

axolotl-ai-cloud · Jan 25, 2024 · 44595f6 · 44595f6
1 parent 239000f
commit 44595f6
Showing 1 changed file with 11 additions and 7 deletions.
diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py
@@ -1015,14 +1015,18 @@ def build_training_arguments(self, total_num_steps):
             training_args_kwargs[
                 "dataloader_prefetch_factor"
             ] = self.cfg.dataloader_prefetch_factor
-        if self.cfg.gradient_checkpointing_kwargs is not None:
+        if self.cfg.gradient_checkpointing:
             training_args_kwargs[
-                "gradient_checkpointing_kwargs"
-            ] = self.cfg.gradient_checkpointing_kwargs
-        else:
-            training_args_kwargs["gradient_checkpointing_kwargs"] = {
-                "use_reentrant": False
-            }
+                "gradient_checkpointing"
+            ] = self.cfg.gradient_checkpointing
+            if self.cfg.gradient_checkpointing_kwargs is not None:
+                training_args_kwargs[
+                    "gradient_checkpointing_kwargs"
+                ] = self.cfg.gradient_checkpointing_kwargs
+            else:
+                training_args_kwargs["gradient_checkpointing_kwargs"] = {
+                    "use_reentrant": False
+                }
 
         training_args = TrainingArguments(
             per_device_train_batch_size=self.cfg.micro_batch_size,