ColossalAI/colossalai/legacy/pipeline/__init__.py