From 8afad77a357d073f0e8e224affb22bfbe261561e Mon Sep 17 00:00:00 2001
From: Mashiro <57566630+HAOCHENYE@users.noreply.github.com>
Date: Fri, 30 Jun 2023 11:15:18 +0800
Subject: [PATCH] [Enhance] Update fsdp vit-huge and vit-large config (#1675)

* Update fsdp vit-huge and vit-large config

* Update fsdp vit-huge and vit-large config

* rename
---
 ....py => vit-huge-p14_8xb128-ds-coslr-50e_in1k.py} |  2 +-
 .../vit-huge-p14_8xb128-fsdp-coslr-50e_in1k.py      | 13 +++++++++++++
 ...py => vit-large-p16_8xb128-ds-coslr-50e_in1k.py} |  2 +-
 .../vit-large-p16_8xb128-fsdp-coslr-50e_in1k.py     | 13 +++++++++++++
 4 files changed, 28 insertions(+), 2 deletions(-)
 rename configs/mae/benchmarks/{vit-huge-p14_8xb128-ds-zero3-coslr-50e_in1k.py => vit-huge-p14_8xb128-ds-coslr-50e_in1k.py} (97%)
 create mode 100644 configs/mae/benchmarks/vit-huge-p14_8xb128-fsdp-coslr-50e_in1k.py
 rename configs/mae/benchmarks/{vit-large-p16_8xb128-ds-zero3-coslr-50e_in1k.py => vit-large-p16_8xb128-ds-coslr-50e_in1k.py} (97%)
 create mode 100644 configs/mae/benchmarks/vit-large-p16_8xb128-fsdp-coslr-50e_in1k.py

diff --git a/configs/mae/benchmarks/vit-huge-p14_8xb128-ds-zero3-coslr-50e_in1k.py b/configs/mae/benchmarks/vit-huge-p14_8xb128-ds-coslr-50e_in1k.py
similarity index 97%
rename from configs/mae/benchmarks/vit-huge-p14_8xb128-ds-zero3-coslr-50e_in1k.py
rename to configs/mae/benchmarks/vit-huge-p14_8xb128-ds-coslr-50e_in1k.py
index 92f6960f..474da1a6 100644
--- a/configs/mae/benchmarks/vit-huge-p14_8xb128-ds-zero3-coslr-50e_in1k.py
+++ b/configs/mae/benchmarks/vit-huge-p14_8xb128-ds-coslr-50e_in1k.py
@@ -18,7 +18,7 @@ strategy = dict(
     ),
     inputs_to_half=['inputs'],
     zero_optimization=dict(
-        stage=3,
+        stage=1,
         allgather_partitions=True,
         reduce_scatter=True,
         allgather_bucket_size=50000000,
diff --git a/configs/mae/benchmarks/vit-huge-p14_8xb128-fsdp-coslr-50e_in1k.py b/configs/mae/benchmarks/vit-huge-p14_8xb128-fsdp-coslr-50e_in1k.py
new file mode 100644
index 00000000..5f8dfb76
--- /dev/null
+++ b/configs/mae/benchmarks/vit-huge-p14_8xb128-fsdp-coslr-50e_in1k.py
@@ -0,0 +1,13 @@
+_base_ = ['./vit-huge-p14_8xb128-coslr-50e_in1k.py']
+
+strategy = dict(
+    type='FSDPStrategy',
+    model_wrapper=dict(
+        auto_wrap_policy=dict(
+            type='torch.distributed.fsdp.wrap.size_based_auto_wrap_policy',
+            min_num_params=1e7)))
+
+optim_wrapper = dict(type='AmpOptimWrapper')
+
+# runner which supports strategies
+runner_type = 'FlexibleRunner'
diff --git a/configs/mae/benchmarks/vit-large-p16_8xb128-ds-zero3-coslr-50e_in1k.py b/configs/mae/benchmarks/vit-large-p16_8xb128-ds-coslr-50e_in1k.py
similarity index 97%
rename from configs/mae/benchmarks/vit-large-p16_8xb128-ds-zero3-coslr-50e_in1k.py
rename to configs/mae/benchmarks/vit-large-p16_8xb128-ds-coslr-50e_in1k.py
index 2c457a09..12d4acc8 100644
--- a/configs/mae/benchmarks/vit-large-p16_8xb128-ds-zero3-coslr-50e_in1k.py
+++ b/configs/mae/benchmarks/vit-large-p16_8xb128-ds-coslr-50e_in1k.py
@@ -18,7 +18,7 @@ strategy = dict(
     ),
     inputs_to_half=['inputs'],
     zero_optimization=dict(
-        stage=3,
+        stage=1,
         allgather_partitions=True,
         reduce_scatter=True,
         allgather_bucket_size=50000000,
diff --git a/configs/mae/benchmarks/vit-large-p16_8xb128-fsdp-coslr-50e_in1k.py b/configs/mae/benchmarks/vit-large-p16_8xb128-fsdp-coslr-50e_in1k.py
new file mode 100644
index 00000000..3a8a6740
--- /dev/null
+++ b/configs/mae/benchmarks/vit-large-p16_8xb128-fsdp-coslr-50e_in1k.py
@@ -0,0 +1,13 @@
+_base_ = ['./vit-large-p16_8xb128-coslr-50e_in1k.py']
+
+strategy = dict(
+    type='FSDPStrategy',
+    model_wrapper=dict(
+        auto_wrap_policy=dict(
+            type='torch.distributed.fsdp.wrap.size_based_auto_wrap_policy',
+            min_num_params=1e7)))
+
+optim_wrapper = dict(type='AmpOptimWrapper')
+
+# runner which supports strategies
+runner_type = 'FlexibleRunner'