[BUG]: ModuleNotFoundError: No module named 'colossalai.utils.model.colo_init_context' #3535

🐛 Describe the bug

CUDA_VISIBLE_DEVICES=6 python train.py
Traceback (most recent call last):
File "train.py", line 13, in
from colossalai.utils.model.colo_init_context import ColoInitContext
ModuleNotFoundError: No module named 'colossalai.utils.model.colo_init_context'

Environment

absl-py 1.4.0
accelerate 0.17.1
aiohttp 3.8.4
aiosignal 1.3.1
albumentations 1.3.0
altair 4.2.2
antlr4-python3-runtime 4.8
anyio 3.6.2
anykeystore 0.2
apex 0.9.10.dev0
appdirs 1.4.4
arrow 1.2.3
asttokens 2.2.1
async-timeout 4.0.2
attrs 22.2.0
backcall 0.2.0
backports.zoneinfo 0.2.1
bcrypt 4.0.1
beautifulsoup4 4.11.2
bitsandbytes 0.37.1
black 23.1.0
blessed 1.20.0
blinker 1.5
boto3 1.26.94
botocore 1.29.94
braceexpand 0.1.7
cachetools 5.3.0
certifi 2022.12.7
cffi 1.15.1
cfgv 3.3.1
charset-normalizer 3.1.0
click 8.1.3
cmake 3.26.0
colossalai 0.2.8
contexttimer 0.3.3
contourpy 1.0.7
coverage 7.2.2
croniter 1.3.8
cryptacular 1.6.2
cryptography 39.0.2
cycler 0.11.0
Cython 0.29.33
datasets 2.10.1
dateutils 0.6.12
decorator 5.1.1
deepdiff 6.3.0
defusedxml 0.7.1
diffusers 0.14.0
dill 0.3.6
distlib 0.3.6
distro 1.8.0
DLLogger 1.0.0
dnspython 2.3.0
docker 6.0.1
docstring-parser 0.8.1
einops 0.3.0
email-validator 1.3.1
entrypoints 0.4
exceptiongroup 1.1.1
executing 1.2.0
fabric 3.0.0
fastapi 0.88.0
fbgemm-gpu 0.2.0
ffmpy 0.3.0
filelock 3.10.0
fire 0.5.0
Flask 2.0.3
fonttools 4.39.2
frozenlist 1.3.3
fsspec 2023.3.0
ftfy 6.1.1
future 0.18.3
gitdb 4.0.10
GitPython 3.1.31
google-auth 2.16.2
google-auth-oauthlib 0.4.6
gradio 3.11.0
greenlet 2.0.2
grpcio 1.51.3
h11 0.12.0
h5py 3.8.0
httpcore 0.15.0
httptools 0.5.0
httpx 0.23.3
huggingface-hub 0.13.2
hupper 1.11
hypothesis 6.70.0
identify 2.5.21
idna 3.4
imageio 2.9.0
imageio-ffmpeg 0.4.2
importlib-metadata 6.1.0
importlib-resources 5.12.0
iniconfig 2.0.0
inquirer 3.1.3
invoke 2.0.0
iopath 0.1.10
ipdb 0.13.13
ipython 8.11.0
itsdangerous 2.1.2
jedi 0.18.2
Jinja2 3.1.2
jmespath 1.0.1
joblib 1.2.0
jsonschema 4.17.3
kiwisolver 1.4.4
lazy_loader 0.1
libcst 0.4.9
lightning 1.9.0
lightning-cloud 0.5.31
lightning-utilities 0.8.0
linkify-it-py 2.0.0
lit 15.0.7
Markdown 3.4.1
markdown-it-py 2.2.0
MarkupSafe 2.1.2
matplotlib 3.7.1
matplotlib-inline 0.1.6
mdit-py-plugins 0.3.5
mdurl 0.1.2
moreorless 0.4.0
mpi4py 3.1.4
mpmath 1.3.0
multidict 6.0.4
multiprocess 0.70.14
mypy-extensions 1.0.0
networkx 3.0
ninja 1.11.1
nltk 3.8.1
nodeenv 1.7.0
numpy 1.24.2
nvidia-cublas-cu11 11.10.3.66
nvidia-cuda-cupti-cu11 11.7.101
nvidia-cuda-nvrtc-cu11 11.7.99
nvidia-cuda-runtime-cu11 11.7.99
nvidia-cudnn-cu11 8.5.0.96
nvidia-cufft-cu11 10.9.0.58
nvidia-curand-cu11 10.2.10.91
nvidia-cusolver-cu11 11.4.0.1
nvidia-cusparse-cu11 11.7.4.91
nvidia-nccl-cu11 2.14.3
nvidia-nvtx-cu11 11.7.91
oauthlib 3.2.2
omegaconf 2.1.1
open-clip-torch 2.7.0
opencv-python 4.6.0.66
opencv-python-headless 4.7.0.72
ordered-set 4.1.0
orjson 3.8.7
packaging 22.0
pandas 1.5.3
paramiko 3.1.0
parso 0.8.3
PasteDeploy 3.0.1
pathspec 0.11.1
pbkdf2 1.3
peft 0.3.0.dev0 /mnt1/wcp/llm-adapters/llm-adapters/peft
pexpect 4.8.0
pickleshare 0.7.5
Pillow 9.4.0
pip 23.0.1
pkgutil_resolve_name 1.3.10
plaster 1.1.2
plaster-pastedeploy 1.0.1
platformdirs 3.1.1
pluggy 1.0.0
portalocker 2.7.0
pre-commit 3.2.0
prefetch-generator 1.0.3
progressbar 2.5
prompt-toolkit 3.0.38
protobuf 3.20.3
psutil 5.9.4
ptyprocess 0.7.0
pudb 2019.2
pure-eval 0.2.2
pyarrow 11.0.0
pyasn1 0.4.8
pyasn1-modules 0.2.8
pycparser 2.21
pycryptodome 3.17
pydantic 1.10.6
pydeck 0.8.0
pyDeprecate 0.3.2
pydub 0.25.1
Pygments 2.14.0
PyJWT 2.6.0
Pympler 1.0.1
PyNaCl 1.5.0
pyparsing 3.0.9
pyramid 2.0.1
pyramid-mailer 0.15.1
pyre-extensions 0.0.27
pyrsistent 0.19.3
pytest 7.2.2
pytest-cov 4.0.0
python-dateutil 2.8.2
python-dotenv 1.0.0
python-editor 1.0.4
python-multipart 0.0.6
python3-openid 3.2.0
pytz 2022.7.1
pytz-deprecation-shim 0.1.0.post0
PyWavelets 1.4.1
PyYAML 6.0
qudida 0.0.4
readchar 4.0.5
regex 2022.10.31
repoze.sendmail 4.4.1
requests 2.28.2
requests-oauthlib 1.3.1
responses 0.18.0
rfc3986 1.5.0
rich 13.3.2
rsa 4.9
s3transfer 0.6.0
safetensors 0.3.0
scikit-build 0.16.7
scikit-image 0.20.0
scikit-learn 1.2.2
scipy 1.9.1
semver 2.13.0
sentencepiece 0.1.97
setuptools 65.6.3
six 1.16.0
smmap 5.0.0
sniffio 1.3.0
sortedcontainers 2.4.0
soupsieve 2.4
SQLAlchemy 1.4.47
stack-data 0.6.2
starlette 0.22.0
starsessions 1.3.0
stdlibs 2022.10.9
streamlit 1.20.0
sympy 1.11.1
tabulate 0.9.0
tensorboard 2.12.0
tensorboard-data-server 0.7.0
tensorboard-plugin-wit 1.8.1
termcolor 2.2.0
test-tube 0.7.5
threadpoolctl 3.1.0
tifffile 2023.3.15
timm 0.6.12
titans 0.0.7
tokenize-rt 5.0.0
tokenizers 0.13.2
toml 0.10.2
tomli 2.0.1
toolz 0.12.0
torch 2.0.0
torchaudio 2.0.1
torchmetrics 0.7.0
torchrec 0.2.0
torchvision 0.15.1
torchx-nightly 2023.3.19
tornado 6.2
tqdm 4.65.0
trailrunner 1.3.0
traitlets 5.9.0
transaction 3.1.0
transformers 4.28.0.dev0
translationstring 1.4
triton 2.0.0
typing_extensions 4.5.0
typing-inspect 0.8.0
tzdata 2022.7
tzlocal 4.3
uc-micro-py 1.0.1
ujson 5.7.0
urllib3 1.26.15
urwid 2.1.2
usort 1.0.5
uvicorn 0.21.1
uvloop 0.17.0
validators 0.20.0
velruse 1.1.1
venusian 3.0.0
virtualenv 20.21.0
watchdog 2.3.1
watchfiles 0.18.1
wcwidth 0.2.6
webdataset 0.2.5
WebOb 1.8.7
websocket-client 1.5.1
websockets 10.4
Werkzeug 2.2.3
wheel 0.38.4
WTForms 3.0.1
wtforms-recaptcha 0.3.2
xxhash 3.2.0
yarl 1.8.2
zipp 3.15.0
zope.deprecation 4.4.0
zope.interface 6.0
zope.sqlalchemy 2.0