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