From 64af21bcb2f68386cdfa3bd1f289781306bc510c Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Fri, 11 Aug 2023 08:37:33 -0400
Subject: [PATCH] set env vars trainer needs for FSDP

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

diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py
index 3eed07b4c..5e9fa0355 100644
--- a/src/axolotl/utils/trainer.py
+++ b/src/axolotl/utils/trainer.py
@@ -353,7 +353,17 @@ def calculate_total_num_steps(cfg, train_dataset, tokenizer):
     return total_num_steps
 
 
+def setup_fsdp_envs(cfg):
+    os.environ["ACCELERATE_USE_FSDP"] = "true"
+    if cfg.fsdp_config.fsdp_sync_module_states:
+        os.environ["FSDP_SYNC_MODULE_STATES"] = "true"
+    if cfg.fsdp_config.fsdp_state_dict_type:
+        os.environ["FSDP_STATE_DICT_TYPE"] = cfg.fsdp_config.fsdp_state_dict_type
+
+
 def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps):
+    if cfg.fsdp:
+        setup_fsdp_envs(cfg)
     warmup_steps = (
         cfg.warmup_steps
         if cfg.warmup_steps is not None