From c7039b933430a200f4bcc16597e618896d6014b8 Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Fri, 30 Jan 2026 15:17:46 -0800 Subject: [PATCH 1/2] fix Signed-off-by: Youngeun Kwon --- megatron/core/distributed/param_and_grad_buffer.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index d34fdebaf75..4db3180a106 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -718,6 +718,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.tensor([0.0], 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 From 8a62df3f310e03d4bdf8b3239c8e879df81e5197 Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Fri, 30 Jan 2026 15:40:38 -0800 Subject: [PATCH 2/2] Address review's comment --- megatron/core/distributed/param_and_grad_buffer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 4db3180a106..b6bf9d0e95c 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -721,7 +721,7 @@ def _does_param_require_new_bucket(param): # 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.tensor([0.0], device="cuda") + tmp_warmup_tensor = torch.zeros([1], device="cuda") torch.distributed.all_reduce(tmp_warmup_tensor, group=self.data_parallel_group) torch.distributed.barrier() else: