ColossalAI/tests/test_legacy/test_context/configs/parallel_2p5d_init.py