ColossalAI/tests/test_utils/test_checkpoint
..
test_checkpoint_1d.py
test_checkpoint_2d.py
test_checkpoint_2p5d.py
test_checkpoint_3d.py