From 5247c5004e4df75ab0ab47c6883fe402a54a879d Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Fri, 18 Aug 2023 20:28:09 -0400
Subject: [PATCH] set env for FSDP offload params (#433)

---
 src/axolotl/utils/trainer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py
index fdf67509b..fa0f0e384 100644
--- a/src/axolotl/utils/trainer.py
+++ b/src/axolotl/utils/trainer.py
@@ -355,6 +355,8 @@ def calculate_total_num_steps(cfg, train_dataset, tokenizer):
 
 def setup_fsdp_envs(cfg):
     os.environ["ACCELERATE_USE_FSDP"] = "true"
+    if cfg.fsdp_config.fsdp_offload_params:
+        os.environ["FSDP_OFFLOAD_PARAMS"] = "true"
     if cfg.fsdp_config.fsdp_sync_module_states:
         os.environ["FSDP_SYNC_MODULE_STATES"] = "true"
     if cfg.fsdp_config.fsdp_state_dict_type: