diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py index dc995fda8e..3fd4a137f1 100644 --- 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: