[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

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,

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'

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,

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'