python --version  # Make sure your use the Python 3.10
python -m pip install --upgrade pip

# Prepare Colab Env
pip install -r colab_env.txt

# Install t5x
git clone --branch=main https://github.com/Orion-Zheng/t5x
python -m pip install ./t5x

# Install ColossalAI
git clone --branch my_openmoe https://github.com/Orion-Zheng/ColossalAI.git
pip install ./ColossalAI
python -m pip install -r ./ColossalAI/examples/language/openmoe/requirements.txt