运行2D张量并行示例报错

[chinoll]

在运行2D张量并行的示例的时候报错了。
报错如下

  File "/home/xxx/anaconda3/lib/python3.8/site-packages/colossalai/context/parallel_context.py", line 134, in get_local_rank
    return self._local_ranks[parallel_mode]
KeyError: <ParallelMode.PARALLEL_2D_COL: '2d_col'>
    x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_2D_COL)]
  File "/home/xxx/anaconda3/lib/python3.8/site-packages/colossalai/context/parallel_context.py", line 134, in get_local_rank
    return self._local_ranks[parallel_mode]
KeyError: <ParallelMode.PARALLEL_2D_COL: '2d_col'>

命令行参数

python -m torch.distributed.launch --nproc_per_node 4 --master_addr localhost --master_port 29500  test.py

配置和示例中的一样，请问该怎么解决呢

[FrankLeeeee]

Hi， 请问你的ColossalAI版本是多少呢？

[chinoll]

colossalAI版本：0.1.3
驱动信息
NVIDIA-SMI 460.32.03 Driver Version: 460.32.03 CUDA Version: 11.2

[kurisusnowdeng]

@chinoll hi, 2d并行的初始化正确吗？可以用以下代码查看

from colossalai.core import global_context as gpc
from colossalai.context import ParallelMode
assert gpc.is_initialized(ParallelMode.PARALLEL_2D_COL)
assert gpc.is_initialized(ParallelMode.PARALLEL_2D_ROW)

[FrankLeeeee]

有点奇怪，能告知一下你的torch和CUDA版本吗？

[chinoll]

有点奇怪，能告知一下你的torch和CUDA版本吗？

torch version:1.10.2
NVIDIA-SMI 460.32.03 Driver Version: 460.32.03 CUDA Version: 11.2

[FrankLeeeee]

好的，我试试看能不能复现

[kurisusnowdeng]

@kurisusnowdeng 设置CUDA_LAUNCH_BLOCKING=1后报错消失了

这个应该是pytorch distributed run在多卡下退出时本身可能出现的问题，可以尝试在代码结尾加上gpc.destroy()，但也不一定能完全解决

[chinoll]

@kurisusnowdeng 设置CUDA_LAUNCH_BLOCKING=1后报错消失了

这个应该是pytorch distributed run在多卡下退出时本身可能出现的问题，可以尝试在代码结尾加上gpc.destroy()，但也不一定能完全解决

thanks，解决了