diff --git a/colossalai/amp/torch_amp/_grad_scaler.py b/colossalai/amp/torch_amp/_grad_scaler.py
index d5e526d0c..de39b3e16 100644
--- a/colossalai/amp/torch_amp/_grad_scaler.py
+++ b/colossalai/amp/torch_amp/_grad_scaler.py
@@ -60,7 +60,7 @@ class GradScaler(object):
     * ``scaler.step(optimizer)`` safely unscales gradients and calls ``optimizer.step()``.
     * ``scaler.update()`` updates ``scaler``'s scale factor.
 
-    Example::
+    Example:
 
         # Creates a GradScaler once at the beginning of training.
         scaler = GradScaler()
diff --git a/colossalai/context/random/_helper.py b/colossalai/context/random/_helper.py
index 15fad47ef..973c4d9fa 100644
--- a/colossalai/context/random/_helper.py
+++ b/colossalai/context/random/_helper.py
@@ -102,10 +102,10 @@ def sync_states():
 def seed(parallel_mode: ParallelMode):
     """ A context for seed switch
 
-    Examples::
+    Examples:
 
-        with seed(ParallelMode.DATA):
-            output = F.dropout(input)
+        >>> with seed(ParallelMode.DATA):
+        >>>     output = F.dropout(input)
 
     Note:
         The parallel_mode should be concluded in ``ParallelMode``. More details about ``ParallelMode`` could be found
@@ -124,18 +124,18 @@ def with_seed(func, parallel_mode: ParallelMode):
     """
     A function wrapper which executes the function with a specified seed.
 
-    Examples::
+    Examples:
 
-        # use with decorator
-        @with_seed(ParallelMode.DATA)
-        def forward(input):
-            return F.dropout(input)
-        out = forward(input)
-        # OR use it inline
-        def forward(input):
-            return F.dropout(input)
-        wrapper_forward = with_seed(forward, ParallelMode.DATA)
-        out = wrapped_forward(input)
+        >>> # use with decorator
+        >>> @with_seed(ParallelMode.DATA)
+        >>> def forward(input):
+        >>>     return F.dropout(input)
+        >>> out = forward(input)
+        >>> # OR use it inline
+        >>> def forward(input):
+        >>>     return F.dropout(input)
+        >>> wrapper_forward = with_seed(forward, ParallelMode.DATA)
+        >>> out = wrapped_forward(input)
 
     Note:
         The parallel_mode should be concluded in ``ParallelMode``. More details about ``ParallelMode`` could be found
diff --git a/colossalai/fx/proxy.py b/colossalai/fx/proxy.py
index 464c3598f..c016b4db7 100644
--- a/colossalai/fx/proxy.py
+++ b/colossalai/fx/proxy.py
@@ -12,7 +12,8 @@ class ColoProxy(Proxy):
     ColoProxy is a proxy class which uses meta tensor to handle data-dependent control flow. The original torch.fx proxy
     cannot be used to infer the condition statement, with this proxy, torch.fx can still run even with if statements.
 
-    Usage:
+    Example::
+
         proxy = tracer.create_proxy(...)
         proxy.meta_data = torch.empty(4, 2, device='meta')
         print(len(proxy)) # expect output 4
diff --git a/colossalai/fx/tracer/tracer.py b/colossalai/fx/tracer/tracer.py
index 4b99f4154..bccdbf2ce 100644
--- a/colossalai/fx/tracer/tracer.py
+++ b/colossalai/fx/tracer/tracer.py
@@ -35,7 +35,8 @@ class ColoTracer(Tracer):
     ColoTracer is a symbolic tracer designed to support dynamic control flow by using meta tensors for the `colossalai.fx` module.
     This tracer is initialized in the same way as the original torch.fx.Tracer.
 
-    Usage:
+    Usage::
+
         class Model(nn.Module):
             def __init__(self):
                 super().__init__()
@@ -254,8 +255,8 @@ class ColoTracer(Tracer):
         non_meta_arg_names = sig_names - meta_arg_names
         for k, v in sig.parameters.items():
             if k in non_meta_arg_names and \
-                k not in concrete_args and \
-                v.default is not inspect.Parameter.empty:
+                    k not in concrete_args and \
+                    v.default is not inspect.Parameter.empty:
                 concrete_args[k] = v.default
 
         # get non concrete arg names
diff --git a/colossalai/gemini/gemini_mgr.py b/colossalai/gemini/gemini_mgr.py
index fc35f4c33..4717e6f24 100644
--- a/colossalai/gemini/gemini_mgr.py
+++ b/colossalai/gemini/gemini_mgr.py
@@ -13,6 +13,14 @@ class GeminiManager:
 
     PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
     https://arxiv.org/abs/2108.05818
+
+    Args:
+        placement_policy (str): Which device to place *held* tensors. It can be 'cpu', 'cuda' and 'auto'.
+            If it's 'cpu', parameters, gradients and optimizer states will be offloaded to CPU, which means min CUDA memory will be used.
+            If it's 'cuda', they won't be offloaded, which means max CUDA memory will be used.
+            If it's 'auto', they are moving dynamically based on CPU and CUDA memory usage. It will utilize heterogeneous memory space evenly and well.
+            Note that 'auto' policy can only work well when no other processes use CUDA during your training.
+        chunk_manager (ChunkManager): A ``ChunkManager`` instance.
     """
 
     def __init__(self, placement_policy: str, chunk_manager: ChunkManager) -> None:
diff --git a/colossalai/nn/parallel/data_parallel.py b/colossalai/nn/parallel/data_parallel.py
index 72d6d053f..ab5a9fea7 100644
--- a/colossalai/nn/parallel/data_parallel.py
+++ b/colossalai/nn/parallel/data_parallel.py
@@ -40,7 +40,7 @@ def _cast_float(args, dtype: torch.dtype):
 class ColoDDP(torch.nn.Module):
     """Distributed data parallel for ColoTensor. Nested ColoDDP is not supported now.
 
-    Example::
+    Example:
         >>> from colossalai.core import global_context as gpc
         >>> from colossalai.context import ParallelMode
         >>> model = torch.nn.Linear(20, 1)
@@ -148,7 +148,7 @@ class ColoDDP(torch.nn.Module):
         """Sets parameters to be ignored by DDP.
         This method must be called before initializing ColoDDP.
 
-        Example::
+        Example:
             >>> params_to_ignore = []
             >>> for p in module.parameters():
             >>>     if should_ignore(p):
@@ -174,7 +174,7 @@ class ZeroDDP(ColoDDP):
     We can configure chunk and gemini via ChunkManager and GeminiManager respectively.
     For more details, see the API reference of ``ChunkManager`` and ``GeminiManager``.
 
-    Example::
+    Example:
         >>> model = torch.nn.Linear(20, 1)
         >>> placement_policy = 'cuda'
         >>> chunk_size = ChunkManager.search_chunk_size(model, search_range, n_grids) if use_chunk else None
@@ -283,7 +283,7 @@ class ZeroDDP(ColoDDP):
             dict:
                 a dictionary containing a whole state of the module
 
-        Example::
+        Example:
 
             >>> module.state_dict().keys()
             ['bias', 'weight']
diff --git a/colossalai/tensor/op_wrapper.py b/colossalai/tensor/op_wrapper.py
index 1d8e847c2..1c00066f7 100644
--- a/colossalai/tensor/op_wrapper.py
+++ b/colossalai/tensor/op_wrapper.py
@@ -21,7 +21,7 @@ def colo_op_impl(func):
     by ``__torch_function__`` dispatch and has a ColoTensor as any of its
     parameters, the function provided will be invoked for that operator.
 
-    Example::
+    Example:
         >>> @colo_op_impl(torch.nn.functional.linear)
         >>> def my_custom_linear(types, args, kwargs, process_group):
         >>>   ....
diff --git a/colossalai/tensor/param_op_hook.py b/colossalai/tensor/param_op_hook.py
index a7937d79b..03cb090a6 100644
--- a/colossalai/tensor/param_op_hook.py
+++ b/colossalai/tensor/param_op_hook.py
@@ -41,7 +41,7 @@ class ParamOpHookManager:
     def use_hooks(*hooks: ParamOpHook):
         """Change the param op hooks you use. Nested calling is allowed.
 
-        Example::
+        Example:
             >>> with ParamOpHookManager.use_hooks(*hooks):
             >>>     do_something()
             >>>     with ParamOpHookManager.use_hooks():
diff --git a/colossalai/tensor/utils.py b/colossalai/tensor/utils.py
index 91cc3aeb9..d154ec7a7 100644
--- a/colossalai/tensor/utils.py
+++ b/colossalai/tensor/utils.py
@@ -26,7 +26,7 @@ def named_params_with_colotensor(
         (string, Union[Tensor, ColoTensor]): Tuple containing
             the name and parameter (or ColoTensor parameter)
 
-    Example::
+    Example:
 
         >>> model = torch.nn.Linear(*linear_size)
         >>> delattr(model.weight)
diff --git a/colossalai/zero/zero_optimizer.py b/colossalai/zero/zero_optimizer.py
index 3babc27b8..d93a6301d 100644
--- a/colossalai/zero/zero_optimizer.py
+++ b/colossalai/zero/zero_optimizer.py
@@ -16,6 +16,30 @@ class OptimState(Enum):
 
 
 class ZeroOptimizer(ColossalaiOptimizer):
+    """A wrapper for optimizer. ``ZeroDDP`` and ``ZeroOptimizer`` implement Zero Redundancy Optimizer (ZeRO state-3).
+
+    Note:
+        You must use ``ZeroDDP`` with ``ZeroOptimizer``.
+
+    Note:
+        Make sure you set ``placement_policy`` of ``GeminiManager`` to `"auto"`,
+        if you set ``gpu_margin_mem_ratio > 0``.
+
+    Args:
+        optim (Optimizer): An Optimizer instance.
+        module (ZeroDDP): A ``ZeroDDP`` instance.
+        gpu_margin_mem_ratio (float, optional): The ratio of GPU remaining memory (after the first forward-backward) 
+            which will be used when using hybrid CPU optimizer. 
+            This argument is meaningless when `placement_policy` of `GeminiManager` is not "auto".
+            Defaults to 0.0.
+        initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**32.
+        min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1.
+        growth_factor (float, optional): growth_factor used by DynamicGradScaler. Defaults to 2.
+        backoff_factor (float, optional): backoff_factor used by DynamicGradScaler. Defaults to 0.5.
+        growth_interval (float, optional): growth_interval used by DynamicGradScaler. Defaults to 1000.
+        hysteresis (float, optional): hysteresis used by DynamicGradScaler. Defaults to 2.
+        max_scale (int, optional): max_scale used by DynamicGradScaler. Defaults to 2**32.
+        """
 
     def __init__(self,
                  optim: Optimizer,
diff --git a/docs/colossalai/colossalai.amp.naive_amp.grad_scaler.rst b/docs/colossalai/colossalai.amp.naive_amp.grad_scaler.rst
new file mode 100644
index 000000000..12d477825
--- /dev/null
+++ b/docs/colossalai/colossalai.amp.naive_amp.grad_scaler.rst
@@ -0,0 +1,8 @@
+colossalai.amp.naive\_amp.grad\_scaler
+======================================
+
+.. automodule:: colossalai.amp.naive_amp.grad_scaler
+   :members:
+
+
+
diff --git a/docs/colossalai/colossalai.amp.naive_amp.rst b/docs/colossalai/colossalai.amp.naive_amp.rst
index 15917e174..fd364c053 100644
--- a/docs/colossalai/colossalai.amp.naive_amp.rst
+++ b/docs/colossalai/colossalai.amp.naive_amp.rst
@@ -4,6 +4,11 @@ colossalai.amp.naive\_amp
 .. automodule:: colossalai.amp.naive_amp
    :members:
 
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.amp.naive_amp.grad_scaler
+
 
 .. toctree::
    :maxdepth: 2
diff --git a/docs/colossalai/colossalai.builder.pipeline.rst b/docs/colossalai/colossalai.builder.pipeline.rst
deleted file mode 100644
index 7b8c960bb..000000000
--- a/docs/colossalai/colossalai.builder.pipeline.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.builder.pipeline
-===========================
-
-.. automodule:: colossalai.builder.pipeline
-   :members:
diff --git a/docs/colossalai/colossalai.builder.rst b/docs/colossalai/colossalai.builder.rst
index 60b8501c8..61163d7c1 100644
--- a/docs/colossalai/colossalai.builder.rst
+++ b/docs/colossalai/colossalai.builder.rst
@@ -9,4 +9,3 @@ colossalai.builder
    :maxdepth: 2
 
    colossalai.builder.builder
-   colossalai.builder.pipeline
diff --git a/docs/colossalai/colossalai.cli.benchmark.benchmark.rst b/docs/colossalai/colossalai.cli.benchmark.benchmark.rst
new file mode 100644
index 000000000..94a4170c8
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.benchmark.benchmark.rst
@@ -0,0 +1,5 @@
+colossalai.cli.benchmark.benchmark
+==================================
+
+.. automodule:: colossalai.cli.benchmark.benchmark
+   :members:
diff --git a/docs/colossalai/colossalai.cli.benchmark.models.rst b/docs/colossalai/colossalai.cli.benchmark.models.rst
new file mode 100644
index 000000000..4e6290288
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.benchmark.models.rst
@@ -0,0 +1,5 @@
+colossalai.cli.benchmark.models
+===============================
+
+.. automodule:: colossalai.cli.benchmark.models
+   :members:
diff --git a/docs/colossalai/colossalai.cli.benchmark.rst b/docs/colossalai/colossalai.cli.benchmark.rst
new file mode 100644
index 000000000..80fb43dde
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.benchmark.rst
@@ -0,0 +1,13 @@
+colossalai.cli.benchmark
+========================
+
+.. automodule:: colossalai.cli.benchmark
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.cli.benchmark.benchmark
+   colossalai.cli.benchmark.models
+   colossalai.cli.benchmark.utils
diff --git a/docs/colossalai/colossalai.cli.benchmark.utils.rst b/docs/colossalai/colossalai.cli.benchmark.utils.rst
new file mode 100644
index 000000000..12fbaf227
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.benchmark.utils.rst
@@ -0,0 +1,5 @@
+colossalai.cli.benchmark.utils
+==============================
+
+.. automodule:: colossalai.cli.benchmark.utils
+   :members:
diff --git a/docs/colossalai/colossalai.cli.check.check_installation.rst b/docs/colossalai/colossalai.cli.check.check_installation.rst
new file mode 100644
index 000000000..95b2d02ca
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.check.check_installation.rst
@@ -0,0 +1,5 @@
+colossalai.cli.check.check\_installation
+========================================
+
+.. automodule:: colossalai.cli.check.check_installation
+   :members:
diff --git a/docs/colossalai/colossalai.cli.check.rst b/docs/colossalai/colossalai.cli.check.rst
new file mode 100644
index 000000000..262ae7ad3
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.check.rst
@@ -0,0 +1,11 @@
+colossalai.cli.check
+====================
+
+.. automodule:: colossalai.cli.check
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.cli.check.check_installation
diff --git a/docs/colossalai/colossalai.cli.cli.rst b/docs/colossalai/colossalai.cli.cli.rst
new file mode 100644
index 000000000..8f83973d5
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.cli.rst
@@ -0,0 +1,5 @@
+colossalai.cli.cli
+==================
+
+.. automodule:: colossalai.cli.cli
+   :members:
diff --git a/docs/colossalai/colossalai.cli.launcher.hostinfo.rst b/docs/colossalai/colossalai.cli.launcher.hostinfo.rst
new file mode 100644
index 000000000..5bcd9dd8c
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.launcher.hostinfo.rst
@@ -0,0 +1,5 @@
+colossalai.cli.launcher.hostinfo
+================================
+
+.. automodule:: colossalai.cli.launcher.hostinfo
+   :members:
diff --git a/docs/colossalai/colossalai.cli.launcher.multinode_runner.rst b/docs/colossalai/colossalai.cli.launcher.multinode_runner.rst
new file mode 100644
index 000000000..223b0deac
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.launcher.multinode_runner.rst
@@ -0,0 +1,5 @@
+colossalai.cli.launcher.multinode\_runner
+=========================================
+
+.. automodule:: colossalai.cli.launcher.multinode_runner
+   :members:
diff --git a/docs/colossalai/colossalai.cli.launcher.rst b/docs/colossalai/colossalai.cli.launcher.rst
new file mode 100644
index 000000000..38bef61c7
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.launcher.rst
@@ -0,0 +1,13 @@
+colossalai.cli.launcher
+=======================
+
+.. automodule:: colossalai.cli.launcher
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.cli.launcher.hostinfo
+   colossalai.cli.launcher.multinode_runner
+   colossalai.cli.launcher.run
diff --git a/docs/colossalai/colossalai.cli.launcher.run.rst b/docs/colossalai/colossalai.cli.launcher.run.rst
new file mode 100644
index 000000000..8506fb9e3
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.launcher.run.rst
@@ -0,0 +1,5 @@
+colossalai.cli.launcher.run
+===========================
+
+.. automodule:: colossalai.cli.launcher.run
+   :members:
diff --git a/docs/colossalai/colossalai.cli.rst b/docs/colossalai/colossalai.cli.rst
new file mode 100644
index 000000000..8cc0dcb04
--- /dev/null
+++ b/docs/colossalai/colossalai.cli.rst
@@ -0,0 +1,18 @@
+colossalai.cli
+==============
+
+.. automodule:: colossalai.cli
+   :members:
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.cli.benchmark
+   colossalai.cli.check
+   colossalai.cli.launcher
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.cli.cli
diff --git a/docs/colossalai/colossalai.engine.gradient_accumulation.rst b/docs/colossalai/colossalai.engine.gradient_accumulation.rst
new file mode 100644
index 000000000..75fc0e9a2
--- /dev/null
+++ b/docs/colossalai/colossalai.engine.gradient_accumulation.rst
@@ -0,0 +1,5 @@
+colossalai.engine.gradient\_accumulation
+========================================
+
+.. automodule:: colossalai.engine.gradient_accumulation
+   :members:
diff --git a/docs/colossalai/colossalai.engine.gradient_handler.rst b/docs/colossalai/colossalai.engine.gradient_handler.rst
index d7d1633a6..27eb2b56a 100644
--- a/docs/colossalai/colossalai.engine.gradient_handler.rst
+++ b/docs/colossalai/colossalai.engine.gradient_handler.rst
@@ -3,3 +3,9 @@ colossalai.engine.gradient\_handler
 
 .. automodule:: colossalai.engine.gradient_handler
    :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.engine.gradient_handler.utils
diff --git a/docs/colossalai/colossalai.engine.gradient_handler.utils.rst b/docs/colossalai/colossalai.engine.gradient_handler.utils.rst
new file mode 100644
index 000000000..c8997e135
--- /dev/null
+++ b/docs/colossalai/colossalai.engine.gradient_handler.utils.rst
@@ -0,0 +1,5 @@
+colossalai.engine.gradient\_handler.utils
+=========================================
+
+.. automodule:: colossalai.engine.gradient_handler.utils
+   :members:
diff --git a/docs/colossalai/colossalai.engine.ophooks.zero_hook.rst b/docs/colossalai/colossalai.engine.ophooks.zero_hook.rst
deleted file mode 100644
index f7868dd3a..000000000
--- a/docs/colossalai/colossalai.engine.ophooks.zero_hook.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.gemini.ophooks.zero\_hook
-====================================
-
-.. automodule:: colossalai.gemini.ophooks.zero_hook
-   :members:
diff --git a/docs/colossalai/colossalai.engine.rst b/docs/colossalai/colossalai.engine.rst
index 740cb0334..3d194b706 100644
--- a/docs/colossalai/colossalai.engine.rst
+++ b/docs/colossalai/colossalai.engine.rst
@@ -7,6 +7,6 @@ colossalai.engine
 .. toctree::
    :maxdepth: 2
 
+   colossalai.engine.gradient_accumulation
    colossalai.engine.gradient_handler
-   colossalai.gemini.ophooks
    colossalai.engine.schedule
diff --git a/docs/colossalai/colossalai.fx.passes.adding_split_node_pass.rst b/docs/colossalai/colossalai.fx.passes.adding_split_node_pass.rst
new file mode 100644
index 000000000..6799fdc65
--- /dev/null
+++ b/docs/colossalai/colossalai.fx.passes.adding_split_node_pass.rst
@@ -0,0 +1,5 @@
+colossalai.fx.passes.adding\_split\_node\_pass
+==============================================
+
+.. automodule:: colossalai.fx.passes.adding_split_node_pass
+   :members:
diff --git a/docs/colossalai/colossalai.fx.passes.meta_info_prop.rst b/docs/colossalai/colossalai.fx.passes.meta_info_prop.rst
new file mode 100644
index 000000000..4e51732ce
--- /dev/null
+++ b/docs/colossalai/colossalai.fx.passes.meta_info_prop.rst
@@ -0,0 +1,5 @@
+colossalai.fx.passes.meta\_info\_prop
+=====================================
+
+.. automodule:: colossalai.fx.passes.meta_info_prop
+   :members:
diff --git a/docs/colossalai/colossalai.fx.passes.rst b/docs/colossalai/colossalai.fx.passes.rst
new file mode 100644
index 000000000..fac10b768
--- /dev/null
+++ b/docs/colossalai/colossalai.fx.passes.rst
@@ -0,0 +1,15 @@
+colossalai.fx.passes
+====================
+
+.. automodule:: colossalai.fx.passes
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.fx.passes.adding_split_node_pass
+   colossalai.fx.passes.meta_info_prop
+   colossalai.fx.passes.shard_1d_pass
+   colossalai.fx.passes.split_module
+   colossalai.fx.passes.utils
diff --git a/docs/colossalai/colossalai.fx.passes.shard_1d_pass.rst b/docs/colossalai/colossalai.fx.passes.shard_1d_pass.rst
new file mode 100644
index 000000000..0942e96d4
--- /dev/null
+++ b/docs/colossalai/colossalai.fx.passes.shard_1d_pass.rst
@@ -0,0 +1,5 @@
+colossalai.fx.passes.shard\_1d\_pass
+====================================
+
+.. automodule:: colossalai.fx.passes.shard_1d_pass
+   :members:
diff --git a/docs/colossalai/colossalai.fx.passes.split_module.rst b/docs/colossalai/colossalai.fx.passes.split_module.rst
new file mode 100644
index 000000000..9e5e58259
--- /dev/null
+++ b/docs/colossalai/colossalai.fx.passes.split_module.rst
@@ -0,0 +1,5 @@
+colossalai.fx.passes.split\_module
+==================================
+
+.. automodule:: colossalai.fx.passes.split_module
+   :members:
diff --git a/docs/colossalai/colossalai.fx.passes.utils.rst b/docs/colossalai/colossalai.fx.passes.utils.rst
new file mode 100644
index 000000000..4afd92563
--- /dev/null
+++ b/docs/colossalai/colossalai.fx.passes.utils.rst
@@ -0,0 +1,5 @@
+colossalai.fx.passes.utils
+==========================
+
+.. automodule:: colossalai.fx.passes.utils
+   :members:
diff --git a/docs/colossalai/colossalai.fx.proxy.rst b/docs/colossalai/colossalai.fx.proxy.rst
new file mode 100644
index 000000000..4b92da41c
--- /dev/null
+++ b/docs/colossalai/colossalai.fx.proxy.rst
@@ -0,0 +1,5 @@
+colossalai.fx.proxy
+===================
+
+.. automodule:: colossalai.fx.proxy
+   :members:
diff --git a/docs/colossalai/colossalai.fx.rst b/docs/colossalai/colossalai.fx.rst
new file mode 100644
index 000000000..778d642c3
--- /dev/null
+++ b/docs/colossalai/colossalai.fx.rst
@@ -0,0 +1,17 @@
+colossalai.fx
+=============
+
+.. automodule:: colossalai.fx
+   :members:
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.fx.passes
+   colossalai.fx.tracer
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.fx.proxy
diff --git a/docs/colossalai/colossalai.fx.tracer.rst b/docs/colossalai/colossalai.fx.tracer.rst
new file mode 100644
index 000000000..d2f743d67
--- /dev/null
+++ b/docs/colossalai/colossalai.fx.tracer.rst
@@ -0,0 +1,11 @@
+colossalai.fx.tracer
+====================
+
+.. automodule:: colossalai.fx.tracer
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.fx.tracer.tracer
diff --git a/docs/colossalai/colossalai.fx.tracer.tracer.rst b/docs/colossalai/colossalai.fx.tracer.tracer.rst
new file mode 100644
index 000000000..83b98bafd
--- /dev/null
+++ b/docs/colossalai/colossalai.fx.tracer.tracer.rst
@@ -0,0 +1,5 @@
+colossalai.fx.tracer.tracer
+===========================
+
+.. automodule:: colossalai.fx.tracer.tracer
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.chunk.rst b/docs/colossalai/colossalai.gemini.chunk.rst
new file mode 100644
index 000000000..9fe1c2b41
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.chunk.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.chunk
+=======================
+
+.. automodule:: colossalai.gemini.chunk
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.chunk_mgr.rst b/docs/colossalai/colossalai.gemini.chunk_mgr.rst
new file mode 100644
index 000000000..acb554faf
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.chunk_mgr.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.chunk\_mgr
+============================
+
+.. automodule:: colossalai.gemini.chunk_mgr
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.gemini_context.rst b/docs/colossalai/colossalai.gemini.gemini_context.rst
new file mode 100644
index 000000000..be4884062
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.gemini_context.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.gemini\_context
+=================================
+
+.. automodule:: colossalai.gemini.gemini_context
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.gemini_mgr.rst b/docs/colossalai/colossalai.gemini.gemini_mgr.rst
new file mode 100644
index 000000000..5d7f944f7
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.gemini_mgr.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.gemini\_mgr
+=============================
+
+.. automodule:: colossalai.gemini.gemini_mgr
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.memory_tracer.memory_monitor.rst b/docs/colossalai/colossalai.gemini.memory_tracer.memory_monitor.rst
new file mode 100644
index 000000000..e8088a609
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.memory_tracer.memory_monitor.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.memory\_tracer.memory\_monitor
+================================================
+
+.. automodule:: colossalai.gemini.memory_tracer.memory_monitor
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.memory_tracer.memstats_collector.rst b/docs/colossalai/colossalai.gemini.memory_tracer.memstats_collector.rst
new file mode 100644
index 000000000..e2682220c
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.memory_tracer.memstats_collector.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.memory\_tracer.memstats\_collector
+====================================================
+
+.. automodule:: colossalai.gemini.memory_tracer.memstats_collector
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.memory_tracer.model_data_memtracer.rst b/docs/colossalai/colossalai.gemini.memory_tracer.model_data_memtracer.rst
new file mode 100644
index 000000000..ccdfe6682
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.memory_tracer.model_data_memtracer.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.memory\_tracer.model\_data\_memtracer
+=======================================================
+
+.. automodule:: colossalai.gemini.memory_tracer.model_data_memtracer
+   :members:
diff --git a/docs/colossalai/colossalai.utils.memory_tracer.rst b/docs/colossalai/colossalai.gemini.memory_tracer.rst
similarity index 63%
rename from docs/colossalai/colossalai.utils.memory_tracer.rst
rename to docs/colossalai/colossalai.gemini.memory_tracer.rst
index 6adf0ee64..f3d9c4d76 100644
--- a/docs/colossalai/colossalai.utils.memory_tracer.rst
+++ b/docs/colossalai/colossalai.gemini.memory_tracer.rst
@@ -1,5 +1,5 @@
-colossalai.utils.memory\_tracer
-===============================
+colossalai.gemini.memory\_tracer
+================================
 
 .. automodule:: colossalai.gemini.memory_tracer
    :members:
@@ -8,6 +8,6 @@ colossalai.utils.memory\_tracer
 .. toctree::
    :maxdepth: 2
 
-   colossalai.gemini.memory_tracer.async_memtracer
+   colossalai.gemini.memory_tracer.memory_monitor
    colossalai.gemini.memory_tracer.memstats_collector
    colossalai.gemini.memory_tracer.model_data_memtracer
diff --git a/docs/colossalai/colossalai.engine.ophooks.rst b/docs/colossalai/colossalai.gemini.ophooks.rst
similarity index 78%
rename from docs/colossalai/colossalai.engine.ophooks.rst
rename to docs/colossalai/colossalai.gemini.ophooks.rst
index 0173aa2a4..af87ab568 100644
--- a/docs/colossalai/colossalai.engine.ophooks.rst
+++ b/docs/colossalai/colossalai.gemini.ophooks.rst
@@ -8,4 +8,4 @@ colossalai.gemini.ophooks
 .. toctree::
    :maxdepth: 2
 
-   colossalai.gemini.ophooks.zero_hook
+   colossalai.gemini.ophooks.utils
diff --git a/docs/colossalai/colossalai.gemini.ophooks.utils.rst b/docs/colossalai/colossalai.gemini.ophooks.utils.rst
new file mode 100644
index 000000000..5c5917047
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.ophooks.utils.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.ophooks.utils
+===============================
+
+.. automodule:: colossalai.gemini.ophooks.utils
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.paramhooks.rst b/docs/colossalai/colossalai.gemini.paramhooks.rst
new file mode 100644
index 000000000..28a823d4e
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.paramhooks.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.paramhooks
+============================
+
+.. automodule:: colossalai.gemini.paramhooks
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.placement_policy.rst b/docs/colossalai/colossalai.gemini.placement_policy.rst
new file mode 100644
index 000000000..9de0ed523
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.placement_policy.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.placement\_policy
+===================================
+
+.. automodule:: colossalai.gemini.placement_policy
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.rst b/docs/colossalai/colossalai.gemini.rst
new file mode 100644
index 000000000..4f6efe386
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.rst
@@ -0,0 +1,27 @@
+colossalai.gemini
+=================
+
+.. automodule:: colossalai.gemini
+   :members:
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.gemini.memory_tracer
+   colossalai.gemini.ophooks
+   colossalai.gemini.paramhooks
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.gemini.chunk
+   colossalai.gemini.chunk_mgr
+   colossalai.gemini.gemini_context
+   colossalai.gemini.gemini_mgr
+   colossalai.gemini.placement_policy
+   colossalai.gemini.stateful_tensor
+   colossalai.gemini.stateful_tensor_container
+   colossalai.gemini.stateful_tensor_mgr
+   colossalai.gemini.tensor_placement_policy
+   colossalai.gemini.tensor_utils
diff --git a/docs/colossalai/colossalai.gemini.stateful_tensor.rst b/docs/colossalai/colossalai.gemini.stateful_tensor.rst
new file mode 100644
index 000000000..02d526d1b
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.stateful_tensor.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.stateful\_tensor
+==================================
+
+.. automodule:: colossalai.gemini.stateful_tensor
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.stateful_tensor_container.rst b/docs/colossalai/colossalai.gemini.stateful_tensor_container.rst
new file mode 100644
index 000000000..be56c2aa8
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.stateful_tensor_container.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.stateful\_tensor\_container
+=============================================
+
+.. automodule:: colossalai.gemini.stateful_tensor_container
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.stateful_tensor_mgr.rst b/docs/colossalai/colossalai.gemini.stateful_tensor_mgr.rst
new file mode 100644
index 000000000..3456192bd
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.stateful_tensor_mgr.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.stateful\_tensor\_mgr
+=======================================
+
+.. automodule:: colossalai.gemini.stateful_tensor_mgr
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.tensor_placement_policy.rst b/docs/colossalai/colossalai.gemini.tensor_placement_policy.rst
new file mode 100644
index 000000000..81dcac339
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.tensor_placement_policy.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.tensor\_placement\_policy
+===========================================
+
+.. automodule:: colossalai.gemini.tensor_placement_policy
+   :members:
diff --git a/docs/colossalai/colossalai.gemini.tensor_utils.rst b/docs/colossalai/colossalai.gemini.tensor_utils.rst
new file mode 100644
index 000000000..385baf4b5
--- /dev/null
+++ b/docs/colossalai/colossalai.gemini.tensor_utils.rst
@@ -0,0 +1,5 @@
+colossalai.gemini.tensor\_utils
+===============================
+
+.. automodule:: colossalai.gemini.tensor_utils
+   :members:
diff --git a/docs/colossalai/colossalai.nn.graph.graph_node.rst b/docs/colossalai/colossalai.nn.graph.graph_node.rst
new file mode 100644
index 000000000..335ecfe62
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.graph.graph_node.rst
@@ -0,0 +1,5 @@
+colossalai.nn.graph.graph\_node
+===============================
+
+.. automodule:: colossalai.nn.graph.graph_node
+   :members:
diff --git a/docs/colossalai/colossalai.nn.graph.rst b/docs/colossalai/colossalai.nn.graph.rst
new file mode 100644
index 000000000..4510b3374
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.graph.rst
@@ -0,0 +1,12 @@
+colossalai.nn.graph
+===================
+
+.. automodule:: colossalai.nn.graph
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.nn.graph.graph_node
+   colossalai.nn.graph.utils
diff --git a/docs/colossalai/colossalai.nn.graph.utils.rst b/docs/colossalai/colossalai.nn.graph.utils.rst
new file mode 100644
index 000000000..866a93cd9
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.graph.utils.rst
@@ -0,0 +1,5 @@
+colossalai.nn.graph.utils
+=========================
+
+.. automodule:: colossalai.nn.graph.utils
+   :members:
diff --git a/docs/colossalai/colossalai.nn.layer.wrapper.lambda_wrapper.rst b/docs/colossalai/colossalai.nn.layer.wrapper.lambda_wrapper.rst
deleted file mode 100644
index f2ced6725..000000000
--- a/docs/colossalai/colossalai.nn.layer.wrapper.lambda_wrapper.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.nn.layer.wrapper.lambda\_wrapper
-===========================================
-
-.. automodule:: colossalai.nn.layer.wrapper.lambda_wrapper
-   :members:
diff --git a/docs/colossalai/colossalai.nn.layer.wrapper.rst b/docs/colossalai/colossalai.nn.layer.wrapper.rst
index 4e66651dc..761bf843a 100644
--- a/docs/colossalai/colossalai.nn.layer.wrapper.rst
+++ b/docs/colossalai/colossalai.nn.layer.wrapper.rst
@@ -8,5 +8,4 @@ colossalai.nn.layer.wrapper
 .. toctree::
    :maxdepth: 2
 
-   colossalai.nn.layer.wrapper.lambda_wrapper
    colossalai.nn.layer.wrapper.pipeline_wrapper
diff --git a/docs/colossalai/colossalai.nn.model.model_from_config.rst b/docs/colossalai/colossalai.nn.model.model_from_config.rst
deleted file mode 100644
index fadb5fd0f..000000000
--- a/docs/colossalai/colossalai.nn.model.model_from_config.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.nn.model.model\_from\_config
-=======================================
-
-.. automodule:: colossalai.nn.model.model_from_config
-   :members:
diff --git a/docs/colossalai/colossalai.nn.model.rst b/docs/colossalai/colossalai.nn.model.rst
deleted file mode 100644
index 5756e11cd..000000000
--- a/docs/colossalai/colossalai.nn.model.rst
+++ /dev/null
@@ -1,11 +0,0 @@
-colossalai.nn.model
-===================
-
-.. automodule:: colossalai.nn.model
-   :members:
-
-
-.. toctree::
-   :maxdepth: 2
-
-   colossalai.nn.model.model_from_config
diff --git a/docs/colossalai/colossalai.nn.optimizer.rst b/docs/colossalai/colossalai.nn.optimizer.rst
index d0063179d..ede9cc496 100644
--- a/docs/colossalai/colossalai.nn.optimizer.rst
+++ b/docs/colossalai/colossalai.nn.optimizer.rst
@@ -16,3 +16,4 @@ colossalai.nn.optimizer
    colossalai.nn.optimizer.hybrid_adam
    colossalai.nn.optimizer.lamb
    colossalai.nn.optimizer.lars
+   colossalai.nn.optimizer.utils
diff --git a/docs/colossalai/colossalai.nn.optimizer.utils.rst b/docs/colossalai/colossalai.nn.optimizer.utils.rst
new file mode 100644
index 000000000..9b2bc2f01
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.optimizer.utils.rst
@@ -0,0 +1,5 @@
+colossalai.nn.optimizer.utils
+=============================
+
+.. automodule:: colossalai.nn.optimizer.utils
+   :members:
diff --git a/docs/colossalai/colossalai.nn.parallel.data_parallel.rst b/docs/colossalai/colossalai.nn.parallel.data_parallel.rst
new file mode 100644
index 000000000..ba987c2ee
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.parallel.data_parallel.rst
@@ -0,0 +1,5 @@
+colossalai.nn.parallel.data\_parallel
+=====================================
+
+.. automodule:: colossalai.nn.parallel.data_parallel
+   :members:
diff --git a/docs/colossalai/colossalai.nn.parallel.layers.colo_module.rst b/docs/colossalai/colossalai.nn.parallel.layers.colo_module.rst
new file mode 100644
index 000000000..c80fff6d5
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.parallel.layers.colo_module.rst
@@ -0,0 +1,5 @@
+colossalai.nn.parallel.layers.colo\_module
+==========================================
+
+.. automodule:: colossalai.nn.parallel.layers.colo_module
+   :members:
diff --git a/docs/colossalai/colossalai.nn.parallel.layers.embedding.rst b/docs/colossalai/colossalai.nn.parallel.layers.embedding.rst
new file mode 100644
index 000000000..1e7ecc50f
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.parallel.layers.embedding.rst
@@ -0,0 +1,5 @@
+colossalai.nn.parallel.layers.embedding
+=======================================
+
+.. automodule:: colossalai.nn.parallel.layers.embedding
+   :members:
diff --git a/docs/colossalai/colossalai.nn.parallel.layers.linear.rst b/docs/colossalai/colossalai.nn.parallel.layers.linear.rst
new file mode 100644
index 000000000..bbc5e3257
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.parallel.layers.linear.rst
@@ -0,0 +1,5 @@
+colossalai.nn.parallel.layers.linear
+====================================
+
+.. automodule:: colossalai.nn.parallel.layers.linear
+   :members:
diff --git a/docs/colossalai/colossalai.nn.parallel.layers.module_utils.rst b/docs/colossalai/colossalai.nn.parallel.layers.module_utils.rst
new file mode 100644
index 000000000..5190ab403
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.parallel.layers.module_utils.rst
@@ -0,0 +1,5 @@
+colossalai.nn.parallel.layers.module\_utils
+===========================================
+
+.. automodule:: colossalai.nn.parallel.layers.module_utils
+   :members:
diff --git a/docs/colossalai/colossalai.nn.parallel.layers.rst b/docs/colossalai/colossalai.nn.parallel.layers.rst
new file mode 100644
index 000000000..782a206e8
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.parallel.layers.rst
@@ -0,0 +1,14 @@
+colossalai.nn.parallel.layers
+=============================
+
+.. automodule:: colossalai.nn.parallel.layers
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.nn.parallel.layers.colo_module
+   colossalai.nn.parallel.layers.embedding
+   colossalai.nn.parallel.layers.linear
+   colossalai.nn.parallel.layers.module_utils
diff --git a/docs/colossalai/colossalai.nn.parallel.reducer.rst b/docs/colossalai/colossalai.nn.parallel.reducer.rst
new file mode 100644
index 000000000..d80841f69
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.parallel.reducer.rst
@@ -0,0 +1,5 @@
+colossalai.nn.parallel.reducer
+==============================
+
+.. automodule:: colossalai.nn.parallel.reducer
+   :members:
diff --git a/docs/colossalai/colossalai.nn.parallel.rst b/docs/colossalai/colossalai.nn.parallel.rst
new file mode 100644
index 000000000..19e9d1eef
--- /dev/null
+++ b/docs/colossalai/colossalai.nn.parallel.rst
@@ -0,0 +1,17 @@
+colossalai.nn.parallel
+======================
+
+.. automodule:: colossalai.nn.parallel
+   :members:
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.nn.parallel.layers
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.nn.parallel.data_parallel
+   colossalai.nn.parallel.reducer
diff --git a/docs/colossalai/colossalai.nn.rst b/docs/colossalai/colossalai.nn.rst
index 32e5eae2f..7e683952f 100644
--- a/docs/colossalai/colossalai.nn.rst
+++ b/docs/colossalai/colossalai.nn.rst
@@ -7,12 +7,13 @@ colossalai.nn
 .. toctree::
    :maxdepth: 2
 
+   colossalai.nn.graph
    colossalai.nn.layer
    colossalai.nn.loss
    colossalai.nn.lr_scheduler
    colossalai.nn.metric
-   colossalai.nn.model
    colossalai.nn.optimizer
+   colossalai.nn.parallel
 
 
 .. toctree::
diff --git a/docs/colossalai/colossalai.pipeline.layer_sepc.rst b/docs/colossalai/colossalai.pipeline.layer_sepc.rst
new file mode 100644
index 000000000..0ff6a83c2
--- /dev/null
+++ b/docs/colossalai/colossalai.pipeline.layer_sepc.rst
@@ -0,0 +1,5 @@
+colossalai.pipeline.layer\_sepc
+===============================
+
+.. automodule:: colossalai.pipeline.layer_sepc
+   :members:
diff --git a/docs/colossalai/colossalai.pipeline.pipelinable.rst b/docs/colossalai/colossalai.pipeline.pipelinable.rst
new file mode 100644
index 000000000..5c2b02ba6
--- /dev/null
+++ b/docs/colossalai/colossalai.pipeline.pipelinable.rst
@@ -0,0 +1,5 @@
+colossalai.pipeline.pipelinable
+===============================
+
+.. automodule:: colossalai.pipeline.pipelinable
+   :members:
diff --git a/docs/colossalai/colossalai.pipeline.rst b/docs/colossalai/colossalai.pipeline.rst
new file mode 100644
index 000000000..adaebea2d
--- /dev/null
+++ b/docs/colossalai/colossalai.pipeline.rst
@@ -0,0 +1,13 @@
+colossalai.pipeline
+===================
+
+.. automodule:: colossalai.pipeline
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.pipeline.layer_sepc
+   colossalai.pipeline.pipelinable
+   colossalai.pipeline.utils
diff --git a/docs/colossalai/colossalai.pipeline.utils.rst b/docs/colossalai/colossalai.pipeline.utils.rst
new file mode 100644
index 000000000..a33bf42cf
--- /dev/null
+++ b/docs/colossalai/colossalai.pipeline.utils.rst
@@ -0,0 +1,5 @@
+colossalai.pipeline.utils
+=========================
+
+.. automodule:: colossalai.pipeline.utils
+   :members:
diff --git a/docs/colossalai/colossalai.rst b/docs/colossalai/colossalai.rst
index 15599f587..921f15a97 100644
--- a/docs/colossalai/colossalai.rst
+++ b/docs/colossalai/colossalai.rst
@@ -1,18 +1,26 @@
 colossalai
 ==========
 
+.. automodule:: colossalai
+   :members:
+
 .. toctree::
    :maxdepth: 2
 
    colossalai.amp
    colossalai.builder
+   colossalai.cli
    colossalai.communication
    colossalai.context
    colossalai.engine
+   colossalai.fx
+   colossalai.gemini
    colossalai.kernel
    colossalai.logging
    colossalai.nn
+   colossalai.pipeline
    colossalai.registry
+   colossalai.tensor
    colossalai.testing
    colossalai.trainer
    colossalai.utils
@@ -22,4 +30,7 @@ colossalai
 .. toctree::
    :maxdepth: 2
 
+   colossalai.constants
+   colossalai.core
+   colossalai.global_variables
    colossalai.initialize
diff --git a/docs/colossalai/colossalai.tensor.colo_parameter.rst b/docs/colossalai/colossalai.tensor.colo_parameter.rst
new file mode 100644
index 000000000..9b65029db
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.colo_parameter.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.colo\_parameter
+=================================
+
+.. automodule:: colossalai.tensor.colo_parameter
+   :members:
diff --git a/docs/colossalai/colossalai.tensor.colo_tensor.rst b/docs/colossalai/colossalai.tensor.colo_tensor.rst
new file mode 100644
index 000000000..9161ac22f
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.colo_tensor.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.colo\_tensor
+==============================
+
+.. automodule:: colossalai.tensor.colo_tensor
+   :members:
diff --git a/docs/colossalai/colossalai.tensor.compute_spec.rst b/docs/colossalai/colossalai.tensor.compute_spec.rst
new file mode 100644
index 000000000..e2d7235d9
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.compute_spec.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.compute\_spec
+===============================
+
+.. automodule:: colossalai.tensor.compute_spec
+   :members:
diff --git a/docs/colossalai/colossalai.tensor.const.rst b/docs/colossalai/colossalai.tensor.const.rst
new file mode 100644
index 000000000..a22a27893
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.const.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.const
+=======================
+
+.. automodule:: colossalai.tensor.const
+   :members:
diff --git a/docs/colossalai/colossalai.tensor.dist_spec_mgr.rst b/docs/colossalai/colossalai.tensor.dist_spec_mgr.rst
new file mode 100644
index 000000000..043cf2260
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.dist_spec_mgr.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.dist\_spec\_mgr
+=================================
+
+.. automodule:: colossalai.tensor.dist_spec_mgr
+   :members:
diff --git a/docs/colossalai/colossalai.tensor.distspec.rst b/docs/colossalai/colossalai.tensor.distspec.rst
new file mode 100644
index 000000000..2b4b0e5fa
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.distspec.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.distspec
+==========================
+
+.. automodule:: colossalai.tensor.distspec
+   :members:
diff --git a/docs/colossalai/colossalai.tensor.op_wrapper.rst b/docs/colossalai/colossalai.tensor.op_wrapper.rst
new file mode 100644
index 000000000..a246e0a6a
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.op_wrapper.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.op\_wrapper
+=============================
+
+.. automodule:: colossalai.tensor.op_wrapper
+   :members:
diff --git a/docs/colossalai/colossalai.tensor.param_op_hook.rst b/docs/colossalai/colossalai.tensor.param_op_hook.rst
new file mode 100644
index 000000000..475ada452
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.param_op_hook.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.param\_op\_hook
+=================================
+
+.. automodule:: colossalai.tensor.param_op_hook
+   :members:
diff --git a/docs/colossalai/colossalai.tensor.process_group.rst b/docs/colossalai/colossalai.tensor.process_group.rst
new file mode 100644
index 000000000..b71409e3b
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.process_group.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.process\_group
+================================
+
+.. automodule:: colossalai.tensor.process_group
+   :members:
diff --git a/docs/colossalai/colossalai.tensor.rst b/docs/colossalai/colossalai.tensor.rst
new file mode 100644
index 000000000..68e06552b
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.rst
@@ -0,0 +1,21 @@
+colossalai.tensor
+=================
+
+.. automodule:: colossalai.tensor
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.tensor.colo_parameter
+   colossalai.tensor.colo_tensor
+   colossalai.tensor.compute_spec
+   colossalai.tensor.const
+   colossalai.tensor.dist_spec_mgr
+   colossalai.tensor.distspec
+   colossalai.tensor.op_wrapper
+   colossalai.tensor.param_op_hook
+   colossalai.tensor.process_group
+   colossalai.tensor.tensor_spec
+   colossalai.tensor.utils
diff --git a/docs/colossalai/colossalai.tensor.tensor_spec.rst b/docs/colossalai/colossalai.tensor.tensor_spec.rst
new file mode 100644
index 000000000..7125b9cbc
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.tensor_spec.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.tensor\_spec
+==============================
+
+.. automodule:: colossalai.tensor.tensor_spec
+   :members:
diff --git a/docs/colossalai/colossalai.tensor.utils.rst b/docs/colossalai/colossalai.tensor.utils.rst
new file mode 100644
index 000000000..5d9bd1b03
--- /dev/null
+++ b/docs/colossalai/colossalai.tensor.utils.rst
@@ -0,0 +1,5 @@
+colossalai.tensor.utils
+=======================
+
+.. automodule:: colossalai.tensor.utils
+   :members:
diff --git a/docs/colossalai/colossalai.utils.checkpoint.module_checkpoint.rst b/docs/colossalai/colossalai.utils.checkpoint.module_checkpoint.rst
new file mode 100644
index 000000000..237ad380b
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.checkpoint.module_checkpoint.rst
@@ -0,0 +1,5 @@
+colossalai.utils.checkpoint.module\_checkpoint
+==============================================
+
+.. automodule:: colossalai.utils.checkpoint.module_checkpoint
+   :members:
diff --git a/docs/colossalai/colossalai.utils.checkpoint.rst b/docs/colossalai/colossalai.utils.checkpoint.rst
new file mode 100644
index 000000000..220c270f0
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.checkpoint.rst
@@ -0,0 +1,12 @@
+colossalai.utils.checkpoint
+===========================
+
+.. automodule:: colossalai.utils.checkpoint
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.utils.checkpoint.module_checkpoint
+   colossalai.utils.checkpoint.utils
diff --git a/docs/colossalai/colossalai.utils.checkpoint.utils.rst b/docs/colossalai/colossalai.utils.checkpoint.utils.rst
new file mode 100644
index 000000000..7fdeefd53
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.checkpoint.utils.rst
@@ -0,0 +1,5 @@
+colossalai.utils.checkpoint.utils
+=================================
+
+.. automodule:: colossalai.utils.checkpoint.utils
+   :members:
diff --git a/docs/colossalai/colossalai.utils.gradient_accumulation.rst b/docs/colossalai/colossalai.utils.gradient_accumulation.rst
deleted file mode 100644
index 6ad2ca3ae..000000000
--- a/docs/colossalai/colossalai.utils.gradient_accumulation.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.utils.gradient\_accumulation
-=======================================
-
-.. automodule:: colossalai.utils.gradient_accumulation
-   :members:
diff --git a/docs/colossalai/colossalai.utils.memory.rst b/docs/colossalai/colossalai.utils.memory.rst
new file mode 100644
index 000000000..67c5d6002
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.memory.rst
@@ -0,0 +1,5 @@
+colossalai.utils.memory
+=======================
+
+.. automodule:: colossalai.utils.memory
+   :members:
diff --git a/docs/colossalai/colossalai.utils.memory_tracer.async_memtracer.rst b/docs/colossalai/colossalai.utils.memory_tracer.async_memtracer.rst
deleted file mode 100644
index 2256c08d3..000000000
--- a/docs/colossalai/colossalai.utils.memory_tracer.async_memtracer.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.utils.memory\_tracer.async\_memtracer
-================================================
-
-.. automodule:: colossalai.gemini.memory_tracer.async_memtracer
-   :members:
diff --git a/docs/colossalai/colossalai.utils.memory_tracer.memstats_collector.rst b/docs/colossalai/colossalai.utils.memory_tracer.memstats_collector.rst
deleted file mode 100644
index 99a48cb8a..000000000
--- a/docs/colossalai/colossalai.utils.memory_tracer.memstats_collector.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.utils.memory\_tracer.memstats\_collector
-===================================================
-
-.. automodule:: colossalai.gemini.memory_tracer.memstats_collector
-   :members:
diff --git a/docs/colossalai/colossalai.utils.memory_tracer.model_data_memtracer.rst b/docs/colossalai/colossalai.utils.memory_tracer.model_data_memtracer.rst
deleted file mode 100644
index 8535a907d..000000000
--- a/docs/colossalai/colossalai.utils.memory_tracer.model_data_memtracer.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.utils.memory\_tracer.model\_data\_memtracer
-======================================================
-
-.. automodule:: colossalai.gemini.memory_tracer.model_data_memtracer
-   :members:
diff --git a/docs/colossalai/colossalai.utils.memory_utils.memory_monitor.rst b/docs/colossalai/colossalai.utils.memory_utils.memory_monitor.rst
deleted file mode 100644
index 4b6d805ae..000000000
--- a/docs/colossalai/colossalai.utils.memory_utils.memory_monitor.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.utils.memory\_utils.memory\_monitor
-==============================================
-
-.. automodule:: colossalai.utils.memory_utils.memory_monitor
-   :members:
diff --git a/docs/colossalai/colossalai.utils.memory_utils.rst b/docs/colossalai/colossalai.utils.memory_utils.rst
deleted file mode 100644
index 6777334e1..000000000
--- a/docs/colossalai/colossalai.utils.memory_utils.rst
+++ /dev/null
@@ -1,12 +0,0 @@
-colossalai.utils.memory\_utils
-==============================
-
-.. automodule:: colossalai.utils.memory_utils
-   :members:
-
-
-.. toctree::
-   :maxdepth: 2
-
-   colossalai.utils.memory_utils.memory_monitor
-   colossalai.utils.memory_utils.utils
diff --git a/docs/colossalai/colossalai.utils.memory_utils.utils.rst b/docs/colossalai/colossalai.utils.memory_utils.utils.rst
deleted file mode 100644
index a7bc1c607..000000000
--- a/docs/colossalai/colossalai.utils.memory_utils.utils.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.utils.memory\_utils.utils
-====================================
-
-.. automodule:: colossalai.utils.memory_utils.utils
-   :members:
diff --git a/docs/colossalai/colossalai.utils.model.colo_init_context.rst b/docs/colossalai/colossalai.utils.model.colo_init_context.rst
new file mode 100644
index 000000000..33ee44915
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.model.colo_init_context.rst
@@ -0,0 +1,5 @@
+colossalai.utils.model.colo\_init\_context
+==========================================
+
+.. automodule:: colossalai.utils.model.colo_init_context
+   :members:
diff --git a/docs/colossalai/colossalai.utils.model.lazy_init_context.rst b/docs/colossalai/colossalai.utils.model.lazy_init_context.rst
new file mode 100644
index 000000000..27c9a32c6
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.model.lazy_init_context.rst
@@ -0,0 +1,5 @@
+colossalai.utils.model.lazy\_init\_context
+==========================================
+
+.. automodule:: colossalai.utils.model.lazy_init_context
+   :members:
diff --git a/docs/colossalai/colossalai.utils.model.rst b/docs/colossalai/colossalai.utils.model.rst
new file mode 100644
index 000000000..9adfd1450
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.model.rst
@@ -0,0 +1,13 @@
+colossalai.utils.model
+======================
+
+.. automodule:: colossalai.utils.model
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.utils.model.colo_init_context
+   colossalai.utils.model.lazy_init_context
+   colossalai.utils.model.utils
diff --git a/docs/colossalai/colossalai.utils.model.utils.rst b/docs/colossalai/colossalai.utils.model.utils.rst
new file mode 100644
index 000000000..211106662
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.model.utils.rst
@@ -0,0 +1,5 @@
+colossalai.utils.model.utils
+============================
+
+.. automodule:: colossalai.utils.model.utils
+   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.comm_profiler.rst b/docs/colossalai/colossalai.utils.profiler.comm_profiler.rst
deleted file mode 100644
index 88532eff2..000000000
--- a/docs/colossalai/colossalai.utils.profiler.comm_profiler.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.utils.profiler.comm\_profiler
-========================================
-
-.. automodule:: colossalai.utils.profiler.comm_profiler
-   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.extention.rst b/docs/colossalai/colossalai.utils.profiler.extention.rst
new file mode 100644
index 000000000..5c8769261
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.profiler.extention.rst
@@ -0,0 +1,5 @@
+colossalai.utils.profiler.extention
+===================================
+
+.. automodule:: colossalai.utils.profiler.extention
+   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.legacy.comm_profiler.rst b/docs/colossalai/colossalai.utils.profiler.legacy.comm_profiler.rst
new file mode 100644
index 000000000..4329a3d60
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.profiler.legacy.comm_profiler.rst
@@ -0,0 +1,5 @@
+colossalai.utils.profiler.legacy.comm\_profiler
+===============================================
+
+.. automodule:: colossalai.utils.profiler.legacy.comm_profiler
+   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.legacy.mem_profiler.rst b/docs/colossalai/colossalai.utils.profiler.legacy.mem_profiler.rst
new file mode 100644
index 000000000..35c665c71
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.profiler.legacy.mem_profiler.rst
@@ -0,0 +1,5 @@
+colossalai.utils.profiler.legacy.mem\_profiler
+==============================================
+
+.. automodule:: colossalai.utils.profiler.legacy.mem_profiler
+   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.legacy.pcie_profiler.rst b/docs/colossalai/colossalai.utils.profiler.legacy.pcie_profiler.rst
new file mode 100644
index 000000000..7aa82b8f7
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.profiler.legacy.pcie_profiler.rst
@@ -0,0 +1,5 @@
+colossalai.utils.profiler.legacy.pcie\_profiler
+===============================================
+
+.. automodule:: colossalai.utils.profiler.legacy.pcie_profiler
+   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.legacy.prof_utils.rst b/docs/colossalai/colossalai.utils.profiler.legacy.prof_utils.rst
new file mode 100644
index 000000000..93af82b2f
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.profiler.legacy.prof_utils.rst
@@ -0,0 +1,5 @@
+colossalai.utils.profiler.legacy.prof\_utils
+============================================
+
+.. automodule:: colossalai.utils.profiler.legacy.prof_utils
+   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.legacy.rst b/docs/colossalai/colossalai.utils.profiler.legacy.rst
new file mode 100644
index 000000000..37fcebde5
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.profiler.legacy.rst
@@ -0,0 +1,14 @@
+colossalai.utils.profiler.legacy
+================================
+
+.. automodule:: colossalai.utils.profiler.legacy
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.utils.profiler.legacy.comm_profiler
+   colossalai.utils.profiler.legacy.mem_profiler
+   colossalai.utils.profiler.legacy.pcie_profiler
+   colossalai.utils.profiler.legacy.prof_utils
diff --git a/docs/colossalai/colossalai.utils.profiler.mem_profiler.rst b/docs/colossalai/colossalai.utils.profiler.mem_profiler.rst
deleted file mode 100644
index a68c17055..000000000
--- a/docs/colossalai/colossalai.utils.profiler.mem_profiler.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.utils.profiler.mem\_profiler
-=======================================
-
-.. automodule:: colossalai.utils.profiler.mem_profiler
-   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.pcie_profiler.rst b/docs/colossalai/colossalai.utils.profiler.pcie_profiler.rst
deleted file mode 100644
index fd2f4ea15..000000000
--- a/docs/colossalai/colossalai.utils.profiler.pcie_profiler.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.utils.profiler.pcie\_profiler
-========================================
-
-.. automodule:: colossalai.utils.profiler.pcie_profiler
-   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.prof_utils.rst b/docs/colossalai/colossalai.utils.profiler.prof_utils.rst
deleted file mode 100644
index 9b34a8505..000000000
--- a/docs/colossalai/colossalai.utils.profiler.prof_utils.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.utils.profiler.prof\_utils
-=====================================
-
-.. automodule:: colossalai.utils.profiler.prof_utils
-   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.profiler.rst b/docs/colossalai/colossalai.utils.profiler.profiler.rst
new file mode 100644
index 000000000..d35522837
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.profiler.profiler.rst
@@ -0,0 +1,5 @@
+colossalai.utils.profiler.profiler
+==================================
+
+.. automodule:: colossalai.utils.profiler.profiler
+   :members:
diff --git a/docs/colossalai/colossalai.utils.profiler.rst b/docs/colossalai/colossalai.utils.profiler.rst
index 2b49c8186..15681fcf2 100644
--- a/docs/colossalai/colossalai.utils.profiler.rst
+++ b/docs/colossalai/colossalai.utils.profiler.rst
@@ -4,11 +4,15 @@ colossalai.utils.profiler
 .. automodule:: colossalai.utils.profiler
    :members:
 
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.utils.profiler.legacy
+
 
 .. toctree::
    :maxdepth: 2
 
-   colossalai.utils.profiler.comm_profiler
-   colossalai.utils.profiler.mem_profiler
-   colossalai.utils.profiler.pcie_profiler
-   colossalai.utils.profiler.prof_utils
+   colossalai.utils.profiler.extention
+   colossalai.utils.profiler.profiler
+   colossalai.utils.profiler.stateful_tensor_mem_extention
diff --git a/docs/colossalai/colossalai.utils.profiler.stateful_tensor_mem_extention.rst b/docs/colossalai/colossalai.utils.profiler.stateful_tensor_mem_extention.rst
new file mode 100644
index 000000000..72a3fccec
--- /dev/null
+++ b/docs/colossalai/colossalai.utils.profiler.stateful_tensor_mem_extention.rst
@@ -0,0 +1,5 @@
+colossalai.utils.profiler.stateful\_tensor\_mem\_extention
+==========================================================
+
+.. automodule:: colossalai.utils.profiler.stateful_tensor_mem_extention
+   :members:
diff --git a/docs/colossalai/colossalai.utils.rst b/docs/colossalai/colossalai.utils.rst
index 975218511..8b232a12c 100644
--- a/docs/colossalai/colossalai.utils.rst
+++ b/docs/colossalai/colossalai.utils.rst
@@ -7,10 +7,9 @@ colossalai.utils
 .. toctree::
    :maxdepth: 2
 
+   colossalai.utils.checkpoint
    colossalai.utils.data_sampler
-   colossalai.utils.gradient_accumulation
-   colossalai.gemini.memory_tracer
-   colossalai.utils.memory_utils
+   colossalai.utils.model
    colossalai.utils.multi_tensor_apply
    colossalai.utils.profiler
    colossalai.utils.tensor_detector
@@ -23,5 +22,6 @@ colossalai.utils
    colossalai.utils.checkpointing
    colossalai.utils.common
    colossalai.utils.cuda
+   colossalai.utils.memory
    colossalai.utils.moe
    colossalai.utils.timer
diff --git a/docs/colossalai/colossalai.zero.rst b/docs/colossalai/colossalai.zero.rst
index 7a2bf227d..3bcaffd28 100644
--- a/docs/colossalai/colossalai.zero.rst
+++ b/docs/colossalai/colossalai.zero.rst
@@ -12,3 +12,10 @@ colossalai.zero
    colossalai.zero.sharded_model
    colossalai.zero.sharded_optim
    colossalai.zero.sharded_param
+   colossalai.zero.utils
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.zero.zero_optimizer
diff --git a/docs/colossalai/colossalai.zero.shard_utils.rst b/docs/colossalai/colossalai.zero.shard_utils.rst
index 0b7ee8683..580bfdab7 100644
--- a/docs/colossalai/colossalai.zero.shard_utils.rst
+++ b/docs/colossalai/colossalai.zero.shard_utils.rst
@@ -12,4 +12,3 @@ colossalai.zero.shard\_utils
    colossalai.zero.shard_utils.bucket_tensor_shard_strategy
    colossalai.zero.shard_utils.commons
    colossalai.zero.shard_utils.tensor_shard_strategy
-   colossalai.zero.shard_utils.tensor_utils
diff --git a/docs/colossalai/colossalai.zero.shard_utils.tensor_utils.rst b/docs/colossalai/colossalai.zero.shard_utils.tensor_utils.rst
deleted file mode 100644
index 84e66d514..000000000
--- a/docs/colossalai/colossalai.zero.shard_utils.tensor_utils.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.zero.shard\_utils.tensor\_utils
-==========================================
-
-.. automodule:: colossalai.zero.shard_utils.tensor_utils
-   :members:
diff --git a/docs/colossalai/colossalai.zero.sharded_param.rst b/docs/colossalai/colossalai.zero.sharded_param.rst
index f4a00cfd5..02e0fc6c2 100644
--- a/docs/colossalai/colossalai.zero.sharded_param.rst
+++ b/docs/colossalai/colossalai.zero.sharded_param.rst
@@ -10,4 +10,3 @@ colossalai.zero.sharded\_param
 
    colossalai.zero.sharded_param.sharded_param
    colossalai.zero.sharded_param.sharded_tensor
-   colossalai.zero.sharded_param.tensorful_state
diff --git a/docs/colossalai/colossalai.zero.sharded_param.tensorful_state.rst b/docs/colossalai/colossalai.zero.sharded_param.tensorful_state.rst
deleted file mode 100644
index 8dd3f7053..000000000
--- a/docs/colossalai/colossalai.zero.sharded_param.tensorful_state.rst
+++ /dev/null
@@ -1,5 +0,0 @@
-colossalai.zero.sharded\_param.tensorful\_state
-===============================================
-
-.. automodule:: colossalai.zero.sharded_param.tensorful_state
-   :members:
diff --git a/docs/colossalai/colossalai.zero.utils.rst b/docs/colossalai/colossalai.zero.utils.rst
new file mode 100644
index 000000000..15cf4d70d
--- /dev/null
+++ b/docs/colossalai/colossalai.zero.utils.rst
@@ -0,0 +1,12 @@
+colossalai.zero.utils
+=====================
+
+.. automodule:: colossalai.zero.utils
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.zero.utils.zero_hook
+   colossalai.zero.utils.zero_hook_v2
diff --git a/docs/colossalai/colossalai.zero.utils.zero_hook.rst b/docs/colossalai/colossalai.zero.utils.zero_hook.rst
new file mode 100644
index 000000000..424f466dd
--- /dev/null
+++ b/docs/colossalai/colossalai.zero.utils.zero_hook.rst
@@ -0,0 +1,5 @@
+colossalai.zero.utils.zero\_hook
+================================
+
+.. automodule:: colossalai.zero.utils.zero_hook
+   :members:
diff --git a/docs/colossalai/colossalai.zero.utils.zero_hook_v2.rst b/docs/colossalai/colossalai.zero.utils.zero_hook_v2.rst
new file mode 100644
index 000000000..6c9af62f1
--- /dev/null
+++ b/docs/colossalai/colossalai.zero.utils.zero_hook_v2.rst
@@ -0,0 +1,5 @@
+colossalai.zero.utils.zero\_hook\_v2
+====================================
+
+.. automodule:: colossalai.zero.utils.zero_hook_v2
+   :members:
diff --git a/docs/colossalai/colossalai.zero.zero_optimizer.rst b/docs/colossalai/colossalai.zero.zero_optimizer.rst
new file mode 100644
index 000000000..b945b081c
--- /dev/null
+++ b/docs/colossalai/colossalai.zero.zero_optimizer.rst
@@ -0,0 +1,5 @@
+colossalai.zero.zero\_optimizer
+===============================
+
+.. automodule:: colossalai.zero.zero_optimizer
+   :members: