From 83466b644fa08f5e9785458f34af76c1ca52a23b Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Mon, 20 Apr 2026 13:03:21 +0800 Subject: [PATCH 1/5] fix a post_init check Signed-off-by: Xiaowei Ren --- megatron/core/distributed/distributed_data_parallel_config.py | 3 ++- .../fsdp/src/megatron_fsdp/distributed_data_parallel_config.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 50878e149de..cc02ef88bc4 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -5,6 +5,7 @@ import torch +from ..utils import is_torch_min_version @dataclass class DistributedDataParallelConfig: @@ -213,7 +214,7 @@ def __post_init__(self): if self.reuse_grad_buf_for_mxfp8_param_ag: assert self.fp8_param_gather, "Reuse grad buffer only when keeping params in MXFP8." - if self.nccl_ub: + if self.nccl_ub and not is_torch_min_version("2.11.0a0"): if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','): raise ValueError( "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is currently not supported " diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py index a2feb99cb23..b23a2b80442 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py @@ -5,6 +5,7 @@ import torch +from megatron.core.utils import is_torch_min_version @dataclass class DistributedDataParallelConfig: @@ -155,7 +156,7 @@ def __post_init__(self): import os """Check the validity of the config.""" - if self.nccl_ub: + if self.nccl_ub and not is_torch_min_version("2.11.0a0"): if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','): raise ValueError( "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is currently not supported " From ebd98437064adab6b4d2b54007b34f510aeccbbb Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Tue, 28 Apr 2026 15:43:32 +0800 Subject: [PATCH 2/5] manually deregister mem pool Signed-off-by: Xiaowei Ren --- megatron/core/distributed/param_and_grad_buffer.py | 2 ++ megatron/training/training.py | 11 ++++++++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index dc3014d72d2..c2f4e55eb9c 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -988,6 +988,7 @@ def __init__( self.param_data = None self.grad_data = None self.extra_main_grads = [] + self.nccl_mem_pool = None if self.nccl_ub: # If nccl_ub is True, use nccl_allocator to allocate memory for param_data/grad_data. @@ -995,6 +996,7 @@ def __init__( pool = nccl_allocator.create_nccl_mem_pool( symmetric=not self.ddp_config.disable_symmetric_registration ) + self.nccl_mem_pool = pool mem_alloc_context = functools.partial( nccl_allocator.nccl_mem, pool, diff --git a/megatron/training/training.py b/megatron/training/training.py index b2f7e52f3cc..fc63f5364f0 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -133,7 +133,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): CallWrapper = type(None) -from megatron.core import mpu, tensor_parallel +from megatron.core import mpu, tensor_parallel, nccl_allocator from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( is_linear_attention_variant, ) @@ -3355,6 +3355,15 @@ def trace_handler(p): # If any exit conditions (signal handler, duration, iterations) have been reached, exit. if should_exit: + # Deregister NCCL user-buffer memory pools before exit. + # Without this, ProcessGroupNCCL's destructor calls abort() which uses + # ncclCommDeregister on handles created by ncclCommWindowRegister, + # causing "NCCL WARN Deregister: Could not find handle" and a crash. + for model_module in model: + if isinstance(model_module, DDP): + for buf in model_module.buffers + model_module.expert_parallel_buffers: + if getattr(buf, 'nccl_mem_pool', None) is not None: + nccl_allocator.deregister_nccl_mem_pool(buf.nccl_mem_pool, buf.data_parallel_group) wandb_writer = get_wandb_writer() if wandb_writer: wandb_writer.finish() From 708e22181624789808fd047e74e0b5f04ee26a31 Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Tue, 28 Apr 2026 15:45:36 +0800 Subject: [PATCH 3/5] fix a typo Signed-off-by: Xiaowei Ren --- megatron/training/training.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/training/training.py b/megatron/training/training.py index fc63f5364f0..e31a5be9493 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3363,7 +3363,7 @@ def trace_handler(p): if isinstance(model_module, DDP): for buf in model_module.buffers + model_module.expert_parallel_buffers: if getattr(buf, 'nccl_mem_pool', None) is not None: - nccl_allocator.deregister_nccl_mem_pool(buf.nccl_mem_pool, buf.data_parallel_group) + nccl_allocator.deregister_mem_pool(buf.nccl_mem_pool, buf.data_parallel_group) wandb_writer = get_wandb_writer() if wandb_writer: wandb_writer.finish() From b8af80b34a8673ff57d56254665a393190befa62 Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Tue, 28 Apr 2026 16:23:04 +0800 Subject: [PATCH 4/5] add a barrier before deregister nccl mem pool Signed-off-by: Xiaowei Ren --- megatron/training/training.py | 1 + 1 file changed, 1 insertion(+) diff --git a/megatron/training/training.py b/megatron/training/training.py index e31a5be9493..b31c5767eba 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3359,6 +3359,7 @@ def trace_handler(p): # Without this, ProcessGroupNCCL's destructor calls abort() which uses # ncclCommDeregister on handles created by ncclCommWindowRegister, # causing "NCCL WARN Deregister: Could not find handle" and a crash. + torch.distributed.barrier() for model_module in model: if isinstance(model_module, DDP): for buf in model_module.buffers + model_module.expert_parallel_buffers: From 90981ac2ca095995e60ffe77ba79eab541dfc993 Mon Sep 17 00:00:00 2001 From: Xiaowei Ren Date: Tue, 28 Apr 2026 16:58:13 +0800 Subject: [PATCH 5/5] auto formatting Signed-off-by: Xiaowei Ren --- megatron/core/distributed/distributed_data_parallel_config.py | 1 + .../fsdp/src/megatron_fsdp/distributed_data_parallel_config.py | 1 + 2 files changed, 2 insertions(+) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index cc02ef88bc4..ffffc1a4701 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -7,6 +7,7 @@ from ..utils import is_torch_min_version + @dataclass class DistributedDataParallelConfig: """Configuration for DistributedDataParallel.""" diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py index b23a2b80442..923efccb907 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py @@ -7,6 +7,7 @@ from megatron.core.utils import is_torch_min_version + @dataclass class DistributedDataParallelConfig: """