ColossalAI/colossalai/legacy/nn/parallel/__init__.py