Feature/zero (#279)

* add zero1 (#209)

* add zero1

* add test zero1

* update zero stage 1 develop (#212)

* Implement naive zero3 (#240)

* naive zero3 works well

* add zero3 param manager

* add TODOs in comments

* add gather full param ctx

* fix sub module streams

* add offload

* fix bugs of hook and add unit tests

* fix bugs of hook and add unit tests (#252)

* add gather full param ctx

* fix sub module streams

* add offload

* fix bugs of hook and add unit tests

* polish code and add state dict hook

* fix bug

* update unit test

* refactor reconstructed zero code

* clip_grad support zero3 and add unit test

* add unit test for Zero3ParameterManager

* [WIP] initialize the shard param class

* [WIP] Yet another sharded model implementation (#274)

* [WIP] initialize the shard param class

* [WIP] Yes another implementation of shardModel. Using a better hook method.

* torch.concat -> torch.cat

* fix test_zero_level_1.py::test_zero_level_1 unitest

* remove deepspeed implementation and refactor for the reconstructed zero module

* polish zero dp unittests

Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: Frank Lee <somerlee.9@gmail.com>

colossalai/engine/schedule/_pipeline_schedule.py

@@ -12,8 +12,7 @@ from colossalai.core import global_context as gpc
 from colossalai.logging import get_dist_logger
 from colossalai.utils import switch_virtual_pipeline_parallel_rank
 from colossalai.utils.cuda import get_current_device
-from colossalai.zero import (ZeroRedundancyOptimizer_Level_2,
-                             ZeroRedundancyOptimizer_Level_3)
+from colossalai.zero import ShardedOptimizer, ShardedModel
 
 from ._base_schedule import BaseSchedule
 
@@ -91,9 +90,10 @@ class PipelineSchedule(BaseSchedule):
         return self._move_to_device(data), self._move_to_device(label)
 
     def pre_processing(self, engine):
-        if isinstance(engine.optimizer, (ZeroRedundancyOptimizer_Level_2, ZeroRedundancyOptimizer_Level_3)):
+        # TODO: remove this after testing new zero with pipeline parallelism
+        if isinstance(engine.optimizer, ShardedOptimizer) or isinstance(engine.model, ShardedModel):
             raise TypeError(
-                "Pipeline schedule is currently not compatible with ZeRO Level 2 and Level 3"
+                "Pipeline schedule is currently not compatible with ZeRO"
             )
         model = engine.model
         if isinstance(model, NaiveAMPModel):