ColossalAI/colossalai/auto_parallel/pipeline_shard/__init__.py