ColossalAI/colossalai/cli/benchmark/models.py