make sure to enable fsdp activation checkpointing

axolotl-ai-cloud · Mar 30, 2024 · c4ad7a8 · c4ad7a8
1 parent ad47dbf
commit c4ad7a8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py
@@ -302,6 +302,8 @@ def calc_sample_packing_eff_est(estimates: List[float]):
 
 def setup_fsdp_envs(cfg):
     os.environ["ACCELERATE_USE_FSDP"] = "true"
+    if cfg.fsdp_config.fsdp_activation_checkpointing:
+        os.environ["FSDP_ACTIVATION_CHECKPOINTING"] = "true"
     if cfg.fsdp_config.fsdp_offload_params:
         os.environ["FSDP_OFFLOAD_PARAMS"] = "true"
     if cfg.fsdp_config.fsdp_sync_module_states: