ModuleNotFoundError: No module named 'colossalai.zero.init_ctx'

[humanintel]

Hi,

When I tried to run a sample training script with ColossalAI, I've got the following error message:

xxx@a930ba8955c6:~/colossal/data$ bash run.sh
WARNING:torch.distributed.run:

---

Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.

---

Traceback (most recent call last):
File "/home/xxx/colossal/ColossalAI-Examples/language/gpt/train_gpt.py", line 17, in
from colossalai.zero.init_ctx import ZeroInitContext
ModuleNotFoundError: No module named 'colossalai.zero.init_ctx'
Traceback (most recent call last):

I installed ColossalAI both from source and using pip, and don't know what might cause this. Any thoughts?
Thanks,

[yulingao]

same question

[digger-yu]

modify this file line 18 ,and try again
\ColossalAI\examples\language\gpt\titans\train_gpt.py

from colossalai.zero.legacy.init_ctx import ZeroInitContext