ColossalAI/examples/tutorial/auto_parallel/setup.py