colossalai.nn

This part contains different colossalai layers for constructing your model. You can easily use them as the way of using layers in torch.nn.

Now colossalai support layer types below: Linear, Classifier, Embedding, PatchEmbedding, LayerNorm, Dropout for different parallelisms.

This part contains different loss functions for different parallelisms.

This part contains different learning rate schedules to control your learning rate in training process

This part contains different metric to measure performance of your model.

This part contains some colossalai optimizers

This part contains different methods to initialize weights.