diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index db3948562f5..4840658b041 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -753,6 +753,12 @@ def _does_param_require_new_bucket(param): group=self.data_parallel_group, symmetric=not self.ddp_config.disable_symmetric_registration, ) + # Since nccl communicator group is created lazily, we need to perform a warmup call to + # initialize NCCL comm buffers for this dp_group before doing buffer registration. + torch.distributed.barrier() + tmp_warmup_tensor = torch.zeros([1], device="cuda") + torch.distributed.all_reduce(tmp_warmup_tensor, group=self.data_parallel_group) + torch.distributed.barrier() else: # If nccl_ub is False, mem_alloc_context is nullcontext. mem_alloc_context = nullcontext