ColossalAI/examples/language/mixtral/model_utils.py