NCCL UB fix: reduce memory cost and correctly deregister NCCL mem pool

megatron/core/distributed/distributed_data_parallel_config.py

@@ -5,6 +5,8 @@
 
 import torch
 
+from ..utils import is_torch_min_version
+
 
 @dataclass
 class DistributedDataParallelConfig:
@@ -213,7 +215,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 "

megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py

@@ -5,6 +5,8 @@
 
 import torch
 
+from megatron.core.utils import is_torch_min_version
+
 
 @dataclass
 class DistributedDataParallelConfig:
@@ -155,7 +157,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 "

megatron/core/distributed/param_and_grad_buffer.py

@@ -998,13 +998,15 @@ 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.
             nccl_allocator.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,

megatron/training/training.py

@@ -128,7 +128,7 @@ def set_startup_timestamps(program_start=None, main_entry=None):
 except ImportError:
     has_nvidia_modelopt = False
 
-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,
 )
@@ -3412,6 +3412,16 @@ 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.
+        torch.distributed.barrier()
+        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_mem_pool(buf.nccl_mem_pool, buf.data_parallel_group)
         wandb_writer = get_wandb_writer()
         if wandb_writer:
             wandb_writer.finish()