[feature] new zero implementation (#1623)

hpcaitech · Sep 24, 2022 · 5be118f · 5be118f
1 parent f921733
commit 5be118f
Show file tree

Hide file tree

Showing 29 changed files with 918 additions and 1,540 deletions.
diff --git a/colossalai/gemini/__init__.py b/colossalai/gemini/__init__.py
@@ -1,10 +1,6 @@
-from .chunk import TensorInfo, Chunk, TensorState
-from .chunk_mgr import ChunkManager
+from .chunk import TensorInfo, TensorState
 from .stateful_tensor_mgr import StatefulTensorMgr
 from .tensor_placement_policy import TensorPlacementPolicyFactory
 from .gemini_mgr import GeminiManager
 
-__all__ = [
-    'StatefulTensorMgr', 'TensorPlacementPolicyFactory', 'GeminiManager', 'ChunkManager', 'TensorInfo', 'Chunk',
-    'TensorState'
-]
+__all__ = ['StatefulTensorMgr', 'TensorPlacementPolicyFactory', 'GeminiManager', 'TensorInfo', 'TensorState']
diff --git a/colossalai/gemini/chunk.py b/colossalai/gemini/chunk.py
diff --git a/colossalai/gemini/chunk/__init__.py b/colossalai/gemini/chunk/__init__.py
@@ -0,0 +1,3 @@
+from .chunk import TensorState, TensorInfo, ChunkFullError, Chunk
+from .manager import ChunkManager
+from .search_utils import clasify_params, search_chunk_configuration