[zero] reorganize zero/gemini folder structure (#3424)

* [zero] refactor low-level zero folder structure

* [zero] fix legacy zero import path

* [zero] fix legacy zero import path

* [zero] remove useless import

* [zero] refactor gemini folder structure

* [zero] refactor gemini folder structure

* [zero] refactor legacy zero import path

* [zero] refactor gemini folder structure

* [zero] refactor gemini folder structure

* [zero] refactor gemini folder structure

* [zero] refactor legacy zero import path

* [zero] fix test import path

* [zero] fix test

* [zero] fix circular import

* [zero] update import

tests/test_optimizer/test_cpu_adam.py

@@ -56,7 +56,7 @@ def test_cpu_adam(adamw, step, p_dtype, g_dtype):
     eps = 1e-8
     weight_decay = 0
 
-    for i in range(1024):
+    for i in range(3):
         p_data = torch.rand(64, dtype=p_dtype)
         p_data_copy = p_data.clone().float()
         p_grad = torch.rand(64, dtype=g_dtype)

tests/test_optimizer/test_fused_adam_kernel.py

@@ -54,7 +54,7 @@ def test_adam(adamw, step, p_dtype, g_dtype):
 
     count = 0
 
-    for i in range(1024):
+    for i in range(3):
         p = torch.rand(64, dtype=p_dtype).cuda()
         p_copy = p.clone().float()
         g = torch.rand(p.shape, dtype=g_dtype).cuda()

tests/test_optimizer/test_hybrid_adam.py

@@ -1,12 +1,12 @@
 import torch
 import torch.nn as nn
-from torch.optim.adam import Adam
 from torch.optim import AdamW
+from torch.optim.adam import Adam
 
 from colossalai.nn.optimizer.hybrid_adam import HybridAdam
 from colossalai.testing import parameterize
 
-RE = 1024
+RE = 3
 
 
 @parameterize('adamw', [False, True])