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