Support TP-compatible Torch AMP and Update trainer API (#27)

* Add gradient accumulation, fix lr scheduler

* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)

* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes

* fixed trainer

* Revert "fixed trainer"

This reverts commit 2e0b0b7699.

* improved consistency between trainer, engine and schedule (#23)

Co-authored-by: 1SAA <c2h214748@gmail.com>

Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>

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

@@ -0,0 +1,5 @@
+colossalai.engine.amp.amp\_type
+===============================
+
+.. automodule:: colossalai.engine.amp.amp_type
+   :members:

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

@@ -0,0 +1,5 @@
+colossalai.engine.amp.grad\_scaler
+==================================
+
+.. automodule:: colossalai.engine.amp.grad_scaler
+   :members:

docs/colossalai/colossalai.engine.amp.rst

@@ -0,0 +1,12 @@
+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.amp_type.rst

@@ -1,5 +0,0 @@
-colossalai.engine.amp\_type
-===========================
-
-.. automodule:: colossalai.engine.amp_type
-   :members:

docs/colossalai/colossalai.engine.rst

@@ -7,11 +7,6 @@ colossalai.engine
 .. toctree::
    :maxdepth: 2
 
+   colossalai.engine.amp
    colossalai.engine.gradient_handler
    colossalai.engine.schedule
-
-
-.. toctree::
-   :maxdepth: 2
-
-   colossalai.engine.amp_type

docs/colossalai/colossalai.rst

@@ -21,7 +21,6 @@ colossalai
 .. toctree::
    :maxdepth: 2
 
-   colossalai.checkpointing
    colossalai.constants
    colossalai.core
    colossalai.initialize

docs/colossalai/colossalai.utils.checkpointing.rst

@@ -0,0 +1,5 @@
+colossalai.utils.checkpointing
+==============================
+
+.. automodule:: colossalai.utils.checkpointing
+   :members:

docs/colossalai/colossalai.utils.rst

@@ -9,6 +9,7 @@ colossalai.utils
    :maxdepth: 2
 
    colossalai.utils.activation_checkpoint
+   colossalai.utils.checkpointing
    colossalai.utils.common
    colossalai.utils.cuda
    colossalai.utils.memory