ColossalAI/colossalai/legacy/utils/checkpoint/__init__.py