3d parallel with megatorn-gpt2

hi, deepspeed.
i'm trying to reproduce 3d-parallel results with megatron-gpt2.
i use PipeModelDataParallelTopology. and the topology setting is:

Using topology: {ProcessCoord(pipe=0, data=0, model=0): 0, ProcessCoord(pipe=0, data=0, model=1): 1, ProcessCoord(pipe=0, data=1, model=0): 2, ProcessCoord(pipe=0, data=1, model=1): 3, ProcessCoord(pipe=1, data=0, model=0): 4, ProcessCoord(pipe=1, data=0, model=1): 5, ProcessCoord(pipe=1, data=1, model=0): 6, ProcessCoord(pipe=1, data=1, model=1): 7}

however, my program hangs. it seems that the pipeline parallel is waiting for recving gradient of attention mask, but it has been poped from this line.

Another issue is, should we use PartitionedTensor with megatron-styled parallelization? i assume the output tensor for each BertLayer should be replicated across devices and thus no need to do PartitionedTensor, am i corect?

Thanks for the help.

Start by reading deepspeed/runtime/pipe/engine.py around line 828 and reviewing the reported PipeModelDataParallelTopology for the Megatron-GPT2 reproduction. Reproduce the hang with the described 3D-parallel setup, then determine why the attention-mask gradient is removed and whether PartitionedTensor is required; done means the behavior and expected tensor handling are established.