From 72c2792139c95a80c1f714cfd2fa0f6d7c86d388 Mon Sep 17 00:00:00 2001
From: bubbliiiing <3323290568@qq.com>
Date: Mon, 9 Jun 2025 09:42:32 +0000
Subject: [PATCH] Update Training

---
 scripts/wan2.1/train.py                  | 5 +++--
 scripts/wan2.1/train_lora.py             | 5 +++--
 scripts/wan2.1_fun/train.py              | 5 +++--
 scripts/wan2.1_fun/train_control.py      | 5 +++--
 scripts/wan2.1_fun/train_control_lora.py | 5 +++--
 scripts/wan2.1_fun/train_lora.py         | 5 +++--
 6 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/scripts/wan2.1/train.py b/scripts/wan2.1/train.py
index 2b9227f..4d3a827 100755
--- a/scripts/wan2.1/train.py
+++ b/scripts/wan2.1/train.py
@@ -759,8 +759,9 @@ def main():
         log_with=args.report_to,
         project_config=accelerator_project_config,
     )
-    deepspeed_plugin = accelerator.state.deepspeed_plugin
-    fsdp_plugin = accelerator.state.fsdp_plugin
+
+    deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None
+    fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None
     if deepspeed_plugin is not None:
         zero_stage = int(deepspeed_plugin.zero_stage)
         fsdp_stage = 0
diff --git a/scripts/wan2.1/train_lora.py b/scripts/wan2.1/train_lora.py
index 42ca41c..50a1883 100755
--- a/scripts/wan2.1/train_lora.py
+++ b/scripts/wan2.1/train_lora.py
@@ -758,8 +758,9 @@ def main():
         log_with=args.report_to,
         project_config=accelerator_project_config,
     )
-    deepspeed_plugin = accelerator.state.deepspeed_plugin
-    fsdp_plugin = accelerator.state.fsdp_plugin
+
+    deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None
+    fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None
     if deepspeed_plugin is not None:
         zero_stage = int(deepspeed_plugin.zero_stage)
         fsdp_stage = 0
diff --git a/scripts/wan2.1_fun/train.py b/scripts/wan2.1_fun/train.py
index 7743942..1863178 100755
--- a/scripts/wan2.1_fun/train.py
+++ b/scripts/wan2.1_fun/train.py
@@ -729,8 +729,9 @@ def main():
         log_with=args.report_to,
         project_config=accelerator_project_config,
     )
-    deepspeed_plugin = accelerator.state.deepspeed_plugin
-    fsdp_plugin = accelerator.state.fsdp_plugin
+
+    deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None
+    fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None
     if deepspeed_plugin is not None:
         zero_stage = int(deepspeed_plugin.zero_stage)
         fsdp_stage = 0
diff --git a/scripts/wan2.1_fun/train_control.py b/scripts/wan2.1_fun/train_control.py
index 9886c14..23d01c6 100755
--- a/scripts/wan2.1_fun/train_control.py
+++ b/scripts/wan2.1_fun/train_control.py
@@ -661,8 +661,9 @@ def main():
         log_with=args.report_to,
         project_config=accelerator_project_config,
     )
-    deepspeed_plugin = accelerator.state.deepspeed_plugin
-    fsdp_plugin = accelerator.state.fsdp_plugin
+
+    deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None
+    fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None
     if deepspeed_plugin is not None:
         zero_stage = int(deepspeed_plugin.zero_stage)
         fsdp_stage = 0
diff --git a/scripts/wan2.1_fun/train_control_lora.py b/scripts/wan2.1_fun/train_control_lora.py
index 9e85399..9a79d25 100755
--- a/scripts/wan2.1_fun/train_control_lora.py
+++ b/scripts/wan2.1_fun/train_control_lora.py
@@ -659,8 +659,9 @@ def main():
         log_with=args.report_to,
         project_config=accelerator_project_config,
     )
-    deepspeed_plugin = accelerator.state.deepspeed_plugin
-    fsdp_plugin = accelerator.state.fsdp_plugin
+
+    deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None
+    fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None
     if deepspeed_plugin is not None:
         zero_stage = int(deepspeed_plugin.zero_stage)
         fsdp_stage = 0
diff --git a/scripts/wan2.1_fun/train_lora.py b/scripts/wan2.1_fun/train_lora.py
index 090f14f..e64e224 100755
--- a/scripts/wan2.1_fun/train_lora.py
+++ b/scripts/wan2.1_fun/train_lora.py
@@ -721,8 +721,9 @@ def main():
         log_with=args.report_to,
         project_config=accelerator_project_config,
     )
-    deepspeed_plugin = accelerator.state.deepspeed_plugin
-    fsdp_plugin = accelerator.state.fsdp_plugin
+
+    deepspeed_plugin = accelerator.state.deepspeed_plugin if hasattr(accelerator.state, "deepspeed_plugin") else None
+    fsdp_plugin = accelerator.state.fsdp_plugin if hasattr(accelerator.state, "fsdp_plugin") else None
     if deepspeed_plugin is not None:
         zero_stage = int(deepspeed_plugin.zero_stage)
         fsdp_stage = 0