update examples and sphnix docs for the new api (#63)

docs/colossalai/colossalai.checkpointing.rst

-colossalai.checkpointing
-========================
-
-.. automodule:: colossalai.checkpointing
-   :members:

docs/colossalai/colossalai.communication.rst

 colossalai.communication
 ========================
 
-.. automodule:: colossalai.communication
-   :members:
-
-
 .. toctree::
    :maxdepth: 2
 
@@ -12,3 +8,7 @@ colossalai.communication
    colossalai.communication.p2p
    colossalai.communication.ring
    colossalai.communication.utils
+
+
+.. automodule:: colossalai.communication
+   :members:

docs/colossalai/colossalai.context.random.rst

 colossalai.context.random
 =========================
 
-.. automodule:: colossalai.context.random
-   :members:
-
-
 .. toctree::
    :maxdepth: 2
 
    colossalai.context.random.seed_manager
+
+
+.. automodule:: colossalai.context.random
+   :members:

docs/colossalai/colossalai.context.rst

 colossalai.context
 ==================
 
-.. automodule:: colossalai.context
-   :members:
-
 .. toctree::
    :maxdepth: 2
 
@@ -17,3 +14,7 @@ colossalai.context
    colossalai.context.config
    colossalai.context.parallel_context
    colossalai.context.parallel_mode
+
+
+.. automodule:: colossalai.context
+   :members:

docs/colossalai/colossalai.engine.amp.amp_type.rst

-colossalai.engine.amp.amp\_type
-===============================
-
-.. automodule:: colossalai.engine.amp.amp_type
-   :members:

docs/colossalai/colossalai.engine.amp.grad_scaler.rst

-colossalai.engine.amp.grad\_scaler
-==================================
-
-.. automodule:: colossalai.engine.amp.grad_scaler
-   :members:

docs/colossalai/colossalai.engine.amp.rst

-colossalai.engine.amp
-=====================
-
-.. automodule:: colossalai.engine.amp
-   :members:
-
-
-.. toctree::
-   :maxdepth: 2
-
-   colossalai.engine.amp.amp_type
-   colossalai.engine.amp.grad_scaler

docs/colossalai/colossalai.engine.rst

 colossalai.engine
 =================
 
-.. automodule:: colossalai.engine
-   :members:
-
 .. toctree::
    :maxdepth: 2
 
-   colossalai.engine.amp
    colossalai.engine.gradient_handler
    colossalai.engine.schedule
+
+
+.. automodule:: colossalai.engine
+   :members:

docs/colossalai/colossalai.logging.rst

 colossalai.logging
 ==================
 
-.. automodule:: colossalai.logging
-   :members:
-
-
 .. toctree::
    :maxdepth: 2
 
    colossalai.logging.logging
+
+
+.. automodule:: colossalai.logging
+   :members:

docs/colossalai/colossalai.nn.data.base_dataset.rst

-colossalai.nn.data.base\_dataset
-================================
-
-.. automodule:: colossalai.nn.data.base_dataset
-   :members:

docs/colossalai/colossalai.nn.data.cifar10_dataset.rst

-colossalai.nn.data.cifar10\_dataset
-===================================
-
-.. automodule:: colossalai.nn.data.cifar10_dataset
-   :members:

docs/colossalai/colossalai.nn.data.rst

-colossalai.nn.data
-==================
-
-.. automodule:: colossalai.nn.data
-   :members:
-
-.. toctree::
-   :maxdepth: 2
-
-   colossalai.nn.data.sampler
-
-
-.. toctree::
-   :maxdepth: 2
-
-   colossalai.nn.data.base_dataset
-   colossalai.nn.data.caltech101_dataset
-   colossalai.nn.data.cifar10_dataset

docs/colossalai/colossalai.nn.data.sampler.base_sampler.rst

-colossalai.nn.data.sampler.base\_sampler
-========================================
-
-.. automodule:: colossalai.nn.data.sampler.base_sampler
-   :members:

docs/colossalai/colossalai.nn.data.sampler.data_parallel_sampler.rst

-colossalai.nn.data.sampler.data\_parallel\_sampler
-==================================================
-
-.. automodule:: colossalai.nn.data.sampler.data_parallel_sampler
-   :members:

docs/colossalai/colossalai.nn.data.sampler.rst

-colossalai.nn.data.sampler
-==========================
-
-.. automodule:: colossalai.nn.data.sampler
-   :members:
-
-
-.. toctree::
-   :maxdepth: 2
-
-   colossalai.nn.data.sampler.base_sampler
-   colossalai.nn.data.sampler.data_parallel_sampler

docs/colossalai/colossalai.nn.data.caltech101_dataset.rst → docs/colossalai/colossalai.nn.layer.non_parallel_layers.rst

-colossalai.nn.data.caltech101\_dataset
+colossalai.nn.layer.non\_parallel\_layers
 ======================================
 
-.. automodule:: colossalai.nn.data.caltech101_dataset
+.. automodule:: colossalai.nn.layer.non_parallel_layers
    :members:

docs/colossalai/colossalai.nn.layer.parallel_1d.rst

 colossalai.nn.layer.parallel\_1d
 ================================
 
-.. automodule:: colossalai.nn.layer.parallel_1d
-   :members:
-
-
 .. toctree::
    :maxdepth: 2
 
    colossalai.nn.layer.parallel_1d.layers
+
+
+.. automodule:: colossalai.nn.layer.parallel_1d
+   :members:

docs/colossalai/colossalai.nn.layer.parallel_2d.rst

 colossalai.nn.layer.parallel\_2d
 ================================
 
-.. automodule:: colossalai.nn.layer.parallel_2d
-   :members:
-
-
 .. toctree::
    :maxdepth: 2
 
    colossalai.nn.layer.parallel_2d.layers
+
+
+.. automodule:: colossalai.nn.layer.parallel_2d
+   :members:

docs/colossalai/colossalai.nn.layer.parallel_2p5d.rst

 colossalai.nn.layer.parallel\_2p5d
 ==================================
 
-.. automodule:: colossalai.nn.layer.parallel_2p5d
-   :members:
-
-
 .. toctree::
    :maxdepth: 2
 
    colossalai.nn.layer.parallel_2p5d.layers
+
+
+.. automodule:: colossalai.nn.layer.parallel_2p5d
+   :members:

docs/colossalai/colossalai.nn.layer.parallel_3d.rst

 colossalai.nn.layer.parallel\_3d
 ================================
 
-.. automodule:: colossalai.nn.layer.parallel_3d
-   :members:
-
-
 .. toctree::
    :maxdepth: 2
 
    colossalai.nn.layer.parallel_3d.layers
+
+
+.. automodule:: colossalai.nn.layer.parallel_3d
+   :members: