update documentation

README.md

@@ -1,6 +1,8 @@
 # ColossalAI
 
-An integrated large-scale model training system with efficient parallelization techniques
+An integrated large-scale model training system with efficient parallelization techniques.
+
+arXiv: [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://arxiv.org/abs/2110.14883)
 
 ## Installation
 

colossalai/nn/optimizer/lamb.py

@@ -12,6 +12,7 @@ from colossalai.registry import OPTIMIZERS
 class Lamb(Optimizer):
     r"""Implements Lamb algorithm.
     It has been proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
+
     Arguments:
         params (iterable): iterable of parameters to optimize or dicts defining
             parameter groups
@@ -23,7 +24,8 @@ class Lamb(Optimizer):
         weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
         adam (bool, optional): always use trust ratio = 1, which turns this into
             Adam. Useful for comparison purposes.
-    .. _Large Batch Optimization for Deep Learning: Training BERT in 76 minutes:
+
+    .. _Large Batch Optimization for Deep Learning\: Training BERT in 76 minutes:
         https://arxiv.org/abs/1904.00962
     """
 

colossalai/nn/optimizer/lars.py

@@ -12,6 +12,7 @@ from colossalai.registry import OPTIMIZERS
 class Lars(Optimizer):
     r"""Implements the LARS optimizer from `"Large batch training of convolutional networks"
     <https://arxiv.org/pdf/1708.03888.pdf>`_.
+
     Args:
         params (iterable): iterable of parameters to optimize or dicts defining
             parameter groups
@@ -35,7 +36,8 @@ class Lars(Optimizer):
         if momentum < 0.0:
             raise ValueError("Invalid momentum value: {}".format(momentum))
         if weight_decay < 0.0:
-            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+            raise ValueError(
+                "Invalid weight_decay value: {}".format(weight_decay))
         if eeta <= 0 or eeta > 1:
             raise ValueError("Invalid eeta value: {}".format(eeta))
         if epsilon < 0:
@@ -48,6 +50,7 @@ class Lars(Optimizer):
     @torch.no_grad()
     def step(self, closure=None):
         """Performs a single optimization step.
+
         Arguments:
             closure (callable, optional): A closure that reevaluates the model
                 and returns the loss.

docs/colossalai/colossalai.nn.data.prefetcher.rst

@@ -1,5 +0,0 @@
-colossalai.nn.data.prefetcher
-=============================
-
-.. automodule:: colossalai.nn.data.prefetcher
-   :members:

docs/colossalai/colossalai.nn.data.rst

@@ -16,5 +16,3 @@ colossalai.nn.data
    colossalai.nn.data.base_dataset
    colossalai.nn.data.caltech101_dataset
    colossalai.nn.data.cifar10_dataset
-   colossalai.nn.data.prefetcher
-   colossalai.nn.data.wiki_dataset

docs/colossalai/colossalai.nn.data.wiki_dataset.rst

@@ -1,5 +0,0 @@
-colossalai.nn.data.wiki\_dataset
-================================
-
-.. automodule:: colossalai.nn.data.wiki_dataset
-   :members:

docs/colossalai/colossalai.nn.model.bert.bert.rst

@@ -1,5 +0,0 @@
-colossalai.nn.model.bert.bert
-=============================
-
-.. automodule:: colossalai.nn.model.bert.bert
-   :members:

docs/colossalai/colossalai.nn.model.bert.rst

@@ -1,11 +0,0 @@
-colossalai.nn.model.bert
-========================
-
-.. automodule:: colossalai.nn.model.bert
-   :members:
-
-
-.. toctree::
-   :maxdepth: 2
-
-   colossalai.nn.model.bert.bert

docs/colossalai/colossalai.nn.model.rst

@@ -7,7 +7,6 @@ colossalai.nn.model
 .. toctree::
    :maxdepth: 2
 
-   colossalai.nn.model.bert
    colossalai.nn.model.vanilla_resnet
    colossalai.nn.model.vision_transformer
 

docs/colossalai/colossalai.nn.optimizer.lars.rst

@@ -0,0 +1,5 @@
+colossalai.nn.optimizer.lars
+============================
+
+.. automodule:: colossalai.nn.optimizer.lars
+   :members:

docs/colossalai/colossalai.nn.optimizer.rst

@@ -13,6 +13,7 @@ colossalai.nn.optimizer
    colossalai.nn.optimizer.fused_lamb
    colossalai.nn.optimizer.fused_sgd
    colossalai.nn.optimizer.lamb
+   colossalai.nn.optimizer.lars
    colossalai.nn.optimizer.loss_scaler
    colossalai.nn.optimizer.zero_redundancy_optimizer_level_1
    colossalai.nn.optimizer.zero_redundancy_optimizer_level_2