Hey @JIE_LIU, which version of PyTorch are you using? If it is v1.6, I suspect it is  due to the DDP comm bucket reconstruction algorithm temporarily boost memory consumption, and then hit the OOM problem.
cc @Yanli_Zhao