ColossalAI/examples/tutorial/sequence_parallel/model/__init__.py