Hi @kurisusnowdeng , I have reviewed the code and found some problems.

colossalai/engine/schedule/_base_schedule.py

1. split_batch only splits the batch dimension. Where can the user to split other dimensions? e.g. user needs to split the last dimension if using 2D
2. The current implementation assumes that the batch data consists of an input and an ouput and does not support multiple inputs and ouputs like those in BERT. If the user wants to use BERT inputs and outputs, they have to build a new schedule class and overrides the load_batch method. I think this might cause inconvenience for the user.
3. @ver217 is updating the load_batch function as well, I think you might talk to him to avoid conflicts.

colossalai/nn/layer/parallel_2p5d/_operation.py

1. In line 123 and 125, we need to call contiguous() on reshaped tensors

colossalai/nn/layer/parallel_1d/layers.py

1. In line 382, conditional_context is still used. I would suggest to use a simple if statement as suggested in the last comment. Do look out for similar cases too.

colossalai/nn/layer/parallel_1d/layers.py

1. In line 48, if gather_output is True for Linear1D_Col, set_parallel_input(True) should be False.