Update Training Code

scripts/wan2.1/train.py

@@ -773,8 +773,12 @@ def main():
         zero_stage = 0
         if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD:
             fsdp_stage = 3
+        elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2.
+            fsdp_stage = 3
         elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP:
             fsdp_stage = 2
+        else:
+            fsdp_stage = 0
         print(f"Using FSDP stage: {fsdp_stage}")
 
         args.use_fsdp = True

scripts/wan2.1/train_lora.py

@@ -772,8 +772,12 @@ def main():
         zero_stage = 0
         if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD:
             fsdp_stage = 3
+        elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2.
+            fsdp_stage = 3
         elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP:
             fsdp_stage = 2
+        else:
+            fsdp_stage = 0
         print(f"Using FSDP stage: {fsdp_stage}")
 
         args.use_fsdp = True

scripts/wan2.1_fun/train.py

@@ -743,8 +743,12 @@ def main():
         zero_stage = 0
         if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD:
             fsdp_stage = 3
+        elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2.
+            fsdp_stage = 3
         elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP:
             fsdp_stage = 2
+        else:
+            fsdp_stage = 0
         print(f"Using FSDP stage: {fsdp_stage}")
 
         args.use_fsdp = True

scripts/wan2.1_fun/train_control.py

@@ -675,8 +675,12 @@ def main():
         zero_stage = 0
         if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD:
             fsdp_stage = 3
+        elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2.
+            fsdp_stage = 3
         elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP:
             fsdp_stage = 2
+        else:
+            fsdp_stage = 0
         print(f"Using FSDP stage: {fsdp_stage}")
 
         args.use_fsdp = True

scripts/wan2.1_fun/train_control_lora.py

@@ -673,8 +673,12 @@ def main():
         zero_stage = 0
         if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD:
             fsdp_stage = 3
+        elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2.
+            fsdp_stage = 3
         elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP:
             fsdp_stage = 2
+        else:
+            fsdp_stage = 0
         print(f"Using FSDP stage: {fsdp_stage}")
 
         args.use_fsdp = True

scripts/wan2.1_fun/train_lora.py

@@ -735,8 +735,12 @@ def main():
         zero_stage = 0
         if fsdp_plugin.sharding_strategy is ShardingStrategy.FULL_SHARD:
             fsdp_stage = 3
+        elif fsdp_plugin.sharding_strategy is None: # The fsdp_plugin.sharding_strategy is None in FSDP 2.
+            fsdp_stage = 3
         elif fsdp_plugin.sharding_strategy is ShardingStrategy.SHARD_GRAD_OP:
             fsdp_stage = 2
+        else:
+            fsdp_stage = 0
         print(f"Using FSDP stage: {fsdp_stage}")
 
         args.use_fsdp = True