ColossalAI/colossalai/legacy/nn/layer/parallel_sequence/__init__.py