[gemini] improve compatibility and add static placement policy (#4479)

* [gemini] remove distributed-related part from colotensor (#4379)

* [gemini] remove process group dependency

* [gemini] remove tp part from colo tensor

* [gemini] patch inplace op

* [gemini] fix param op hook and update tests

* [test] remove useless tests

* [test] remove useless tests

* [misc] fix requirements

* [test] fix model zoo

* [test] fix model zoo

* [test] fix model zoo

* [test] fix model zoo

* [test] fix model zoo

* [misc] update requirements

* [gemini] refactor gemini optimizer and gemini ddp (#4398)

* [gemini] update optimizer interface

* [gemini] renaming gemini optimizer

* [gemini] refactor gemini ddp class

* [example] update gemini related example

* [example] update gemini related example

* [plugin] fix gemini plugin args

* [test] update gemini ckpt tests

* [gemini] fix checkpoint io

* [example] fix opt example requirements

* [example] fix opt example

* [example] fix opt example

* [example] fix opt example

* [gemini] add static placement policy (#4443)

* [gemini] add static placement policy

* [gemini] fix param offload

* [test] update gemini tests

* [plugin] update gemini plugin

* [plugin] update gemini plugin docstr

* [misc] fix flash attn requirement

* [test] fix gemini checkpoint io test

* [example] update resnet example result (#4457)

* [example] update bert example result (#4458)

* [doc] update gemini doc (#4468)

* [example] update gemini related examples (#4473)

* [example] update gpt example

* [example] update dreambooth example

* [example] update vit

* [example] update opt

* [example] update palm

* [example] update vit and opt benchmark

* [hotfix] fix bert in model zoo (#4480)

* [hotfix] fix bert in model zoo

* [test] remove chatglm gemini test

* [test] remove sam gemini test

* [test] remove vit gemini test

* [hotfix] fix opt tutorial example (#4497)

* [hotfix] fix opt tutorial example

* [hotfix] fix opt tutorial example

tests/test_zero/test_low_level/test_zero_init.py

-import pytest
-import torch
-import torch.distributed as dist
-import torch.nn as nn
-
-import colossalai
-from colossalai.tensor import ProcessGroup
-from colossalai.testing import spawn
-from colossalai.utils import get_current_device
-from colossalai.zero import ColoInitContext, LowLevelZeroOptimizer
-
-
-class MlpModel(nn.Module):
-
-    def __init__(self):
-        super(MlpModel, self).__init__()
-        self.linear1 = nn.Linear(128, 256)
-        self.linear2 = nn.Linear(256, 512)
-
-    def forward(self, x):
-        x = self.linear1(x)
-        x = self.linear2(x)
-        return x
-
-
-def exam_zero_init():
-    dp_2_tp_2_pg = ProcessGroup(dp_degree=2, tp_degree=2)
-    model1 = MlpModel().cuda()
-    with ColoInitContext(device=get_current_device(), default_pg=dp_2_tp_2_pg):
-        model2 = MlpModel()
-    optimizer1 = LowLevelZeroOptimizer(torch.optim.Adam(model1.parameters(), lr=1))
-    optimizer2 = LowLevelZeroOptimizer(torch.optim.Adam(model2.parameters(), lr=1))
-
-    assert optimizer1._local_rank == optimizer2._local_rank
-    assert optimizer1._world_size == optimizer2._world_size
-
-    mp_group1 = optimizer1.tp_pg
-    mp_group2 = optimizer2.tp_pg
-    assert dist.get_world_size(mp_group1) == dist.get_world_size(mp_group2)
-    assert dist.get_rank(mp_group1) == dist.get_rank(mp_group2)
-
-
-def run_dist(rank, world_size, port):
-    config_dict = dict(parallel=dict(data=2, tensor=dict(size=2, mode='1d')))
-    colossalai.launch(config=config_dict, rank=rank, world_size=world_size, port=port, host='localhost')
-    exam_zero_init()
-
-
-@pytest.mark.dist
-def test_zero_init():
-    spawn(run_dist, 4)
-
-
-if __name__ == '__main__':
-    test_zero_init()

tests/test_zero/test_low_level/test_zero_tp.py

@@ -85,6 +85,7 @@ def run_dist(rank, world_size, port):
     exam_zero_with_tp()
 
 
+@pytest.mark.skip('this will be rewritten by shardformer')
 @pytest.mark.dist
 @rerun_if_address_is_in_use()
 def test_zero_with_tp():