@Jing-Bi Can you check if your problem is resolved in the pytorch nightly release? This might be happening earlier due to a barrier() call in init_process_group which was removed in https://github.com/pytorch/pytorch/pull/49419 10.