[Dev] Update symmetric registration interface to sync-up with upstream pytorch change

megatron/core/distributed/fsdp/src/README.md

@@ -35,10 +35,14 @@ Megatron-FSDP can provide up to 25% speed up and 23% memory savings compared to
 - **Advanced Bucketing**: Data-type aware bucketing system to minimize the overhead of collective operations
 - **Buffer Management**: Zero copy communication is achieved by reorganizing the storage of parameters and main grad with `ParamAndGradBuffer` class
 - **Communication Overlapping**: Improved communication overlap of paramter all-gather and gradient reduce-scatter
-- **User-Buffer-Registration NCCL communication**: Offload NCCL collective communication to NVL/IB Sharp to reduce GPU SM usage for communication
 - **FP8 Mixed Precision with Transformer Engine**: Compatibility with Transformer Engine enables efficient FP8 mixed precision training
 - **Gradient accumulate fusion support with Transformer Engine**: Remove the explicit gradient copy to the communication buffer in backwards pass
 
+### Advanced Collective Communication
+- **SM Usage Reduction with SHARP**: FSDP's `All-Gather` (AG) and `Reduce-Scatter` (RS) collectives are designed to overlap with compute kernels. However, standard NCCL communication kernels can consume a significant number of GPU SMs (e.g., 16-32 SMs), "stealing" resources from compute (GEMM) kernels and reducing overall TFLOPS.
+- **In-Switch Processing**: We leverage **SHARP** (Scalable Hierarchical Aggregation and Reduction Protocol) to offload these collective operations. SHARP performs aggregation and reduction computations directly on the network switches (InfiniBand or NVLink Switch) instead of on the GPU SMs. This dramatically reduces the SM consumption for communication to **1-6 SM** freeing up GPU resources for compute. It also provides lower communication latency, especially in large, scaled-out workloads.
+- **Symmetric Optimizations for MNNVL**: We support **symmetric-based optimizations**, introduced in NCCL v2.27, which enable switch offloading for **Multi-Node NVLink (MNNVL)** systems such as GB200/GB300. This allows the same SM-saving benefits over the high-bandwidth NVLink fabric itself.
+- **Hierarchical Collectives**: When an FSDP sharding domain spans both NVLink and InfiniBand, the library utilizes **hierarchical SHARP collectives** (e.g., NVL-SHARP + IB-SHARP) to optimize the communication path across the entire system topology.
 <!-- ## 📊 Performance  -->
 
 ## 📦 Installation
@@ -207,6 +211,9 @@ optimizer.load_state_dict(ckpt_state_dict["optimizer"])
 - `nccl_ub` will allocate and register the NCCL userbuffer for param and grad buffers. This option enables an SM-efficient NCCL algorithm that could improve the performance of overlapped computations. This flag will be much more effective when used together with SHARP if the FSDP communication includes both NVL and IB domains. Enabling this option will cause additional memory overhead due to the requirement to enable the `fsdp_double_buffer` option.
     - **Only effective when using Megatron-LM.**
     - Defaults to `False`.
+    - By default we try to use NCCL window (symmetric) registration if it is available. If not it falls back to conventional local registraion.
+- `disable_symmetric_registration` will disable NCCL window (i.e. symmetric) registraion when using `nccl_ub`. 
+    - Dafaults to `False`.
 - `fsdp_double_buffer` will use persistently allocated double buffers for temporarily-defined memory needed in `MegatronFSDP` communications. Having persistent double buffers may increase peak VRAM utilization, but is required to register NCCL user buffers (`nccl_ub=True`) for `MegatronFSDP`. Currently, this is only supported for simple repetitive model structures such as GPT.
     - **Only effective when using Megatron-LM.**
     - Defaults to `False`. Automatically overridden to `True` when `nccl_ub` is enabled.

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

@@ -76,13 +76,19 @@
 except Exception:
     HAVE_TE = False
 
+NCCL_ALLOCATOR = None
+
 try:
     # Try to import the MCore NCCL nccl_allocator first.
     # If it fails, try to import the APEX NCCL nccl_allocator.
     import megatron.core.nccl_allocator as nccl_allocator
+
+    NCCL_ALLOCATOR = "MCORE"
 except ImportError:
     try:
         import apex.contrib.nccl_allocator as nccl_allocator
+
+        NCCL_ALLOCATOR = "APEX"
     except ImportError:
         nccl_allocator = None
 
@@ -94,8 +100,8 @@ def _p_assert(cond: Any, s: str, raise_assertion_error: bool = True) -> None:
     message ``s`` since otherwise, it is swallowed.
     """
     if not cond:
-        print(s)
-        traceback.print_stack()
+        logger.error(s)
+        logger.error(''.join(traceback.format_stack()))
         if raise_assertion_error:
             raise AssertionError(s)
 
@@ -205,7 +211,7 @@ def __exit__(self, *args):
         for group in self.groups[1:]:
             backend = group._get_backend(torch.device("cuda", torch.cuda.current_device()))
             if torch.distributed.get_rank() == 0:
-                print(
+                logger.info(
                     f"[MultiGroupUBRAllocator] Registering mem pool to group {group}, "
                     f"group.group_desc:{group.group_desc}"
                 )
@@ -1612,7 +1618,9 @@ def __init__(
         # If using nccl_ub, it returns a function that registers buffers to the NCCL memory pool
         # Buffer is registered to data_parallel_group and expert_data_parallel_group if it exists
         # In the case of not using nccl_ub, it returns a nullcontext
-        self.mem_alloc_context = self.get_mem_alloc_context(groups=self.ubr_groups)
+        self.mem_alloc_context = self.get_mem_alloc_context(
+            groups=self.ubr_groups, symmetric=not self.ddp_config.disable_symmetric_registration
+        )
 
         # Mark FP8 params. If TransformerEngine is not installed, we can skip this.
         meta_device_init_fp8_params = {}
@@ -1640,7 +1648,7 @@ def __init__(
 
         self._log_parameter_groups()
 
-    def get_mem_alloc_context(self, groups=None):
+    def get_mem_alloc_context(self, groups=None, symmetric=True):
         """
         Get the memory allocation context for the parameter and gradient buffers.
         """
@@ -1653,22 +1661,43 @@ def get_mem_alloc_context(self, groups=None):
             if groups is None:
                 # data parallel group is a default group for user buffer registration
                 groups = [self.dist_index.get_fsdp_group(is_expert_parallel=False)]
-            if len(groups) == 1:
-                # register buffers to the default group directly using apex memory allocator
-                mem_alloc_context = functools.partial(
-                    nccl_allocator.nccl_mem, NCCL_MEMORY_POOL, group=groups[0]
-                )
-            else:
-                if hasattr(nccl_allocator, "MultiGroupMemPoolAllocator"):
-                    # Case of MCore NCCL allocator
+
+            if NCCL_ALLOCATOR == "MCORE":
+                if len(groups) == 1:
+                    # register buffers to the default group directly using nccl memory allocator
                     mem_alloc_context = functools.partial(
-                        nccl_allocator.MultiGroupMemPoolAllocator, NCCL_MEMORY_POOL, groups=groups
+                        nccl_allocator.nccl_mem,
+                        NCCL_MEMORY_POOL,
+                        group=groups[0],
+                        symmetric=symmetric,
                     )
                 else:
-                    # Case of APEX NCCL allocator.
+                    mem_alloc_context = functools.partial(
+                        nccl_allocator.MultiGroupMemPoolAllocator,
+                        NCCL_MEMORY_POOL,
+                        groups=groups,
+                        symmetric=symmetric,
+                    )
+            elif NCCL_ALLOCATOR == "APEX":
+                if symmetric:
+                    logging.warning(
+                        "Symmetric registration is not supported for APEX NCCL allocator."
+                        "falling back to non-symmetric registration. "
+                        "Please use Megatron Core NCCL allocator for symmetric registration."
+                    )
+
+                if len(groups) == 1:
+                    # register buffers to the default group directly using nccl memory allocator
+                    mem_alloc_context = functools.partial(
+                        nccl_allocator.nccl_mem, NCCL_MEMORY_POOL, group=groups[0]
+                    )
+                else:
+                    # Supports multiple groups registration for APEX NCCL allocator.
                     mem_alloc_context = functools.partial(
                         MultiGroupUBRAllocator, NCCL_MEMORY_POOL, groups=groups
                     )
+            else:
+                raise ValueError(f"Invalid NCCL allocator: {NCCL_ALLOCATOR}")
             return mem_alloc_context
         else:
             return nullcontext

megatron/core/distributed/param_and_grad_buffer.py

@@ -685,7 +685,10 @@ def _does_param_require_new_bucket(param):
                 symmetric=not self.ddp_config.disable_symmetric_registration
             )
             mem_alloc_context = functools.partial(
-                nccl_allocator.nccl_mem, pool, group=self.data_parallel_group
+                nccl_allocator.nccl_mem,
+                pool,
+                group=self.data_parallel_group,
+                symmetric=not self.ddp_config.disable_symmetric_registration,
             )
         else:
             # If nccl_ub is False, mem_alloc_context is nullcontext.

megatron/core/nccl_allocator.py

@@ -2,6 +2,7 @@
 import logging
 import os
 from contextlib import nullcontext
+from functools import lru_cache
 
 import torch
 
@@ -94,6 +95,7 @@ def _build_nccl_allocator():
     _allocator = nccl_allocator.get_nccl_allocator()
 
 
+@lru_cache(maxsize=None)
 def get_func_args(func):
     """
     Get the argument names of a function.
@@ -122,15 +124,17 @@ def create_nccl_mem_pool(symmetric=None):  # symmetric: bool | None = None -> to
         _pool = torch.cuda.MemPool(_allocator)
     else:
         if 'symmetric' in get_func_args(torch.cuda.MemPool):
+            # The PyTorch version >= 2.9.0a0 and before PyTorch PR #161238,
+            # The symmetric knob should passed to the MemPool constructor.
+            # Since PyTorch PR #161238 symmetric knob is now in registration function.
             _pool = torch.cuda.MemPool(_allocator, symmetric=symmetric)
         elif 'symm_mem' in get_func_args(torch.cuda.MemPool):
             # This path handles argument name divergence between
             # nvidia pytorch and the official pytorch.
             _pool = torch.cuda.MemPool(_allocator, symm_mem=symmetric)
         else:
-            raise ValueError(
-                "symmetric setting with torch.cuda.MemPool requires " "higher PyTorch version"
-            )
+            # This path handles the case where the symmetric knob is in the registration function.
+            _pool = torch.cuda.MemPool(_allocator)
     return _pool
 
 
@@ -149,7 +153,7 @@ def init() -> None:
     # Disables the use of the tensor register allocator hook
     os.environ["TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK"] = "0"
     _build_nccl_allocator()
-    print(f"[MCORE][NCCL_ALLOCATOR] Initialized NCCL Allocator")
+    logging.info(f"[MCORE][NCCL_ALLOCATOR] Initialized NCCL Allocator")
 
 
 # Preserve the original APEX NCCL allocator interface for backward compatibility
@@ -158,11 +162,12 @@ class nccl_mem:
     An NCCL memory allocator, which inherits APEX nccl_allocator implementation.
     """
 
-    def __init__(self, pool, enabled=True, device=None, group=None):
+    def __init__(self, pool, enabled=True, device=None, group=None, symmetric=True):
         self.device = None
         self.group = None
         self.mem_context = None
         self.pool = pool
+        self.symmetric = symmetric
 
         if enabled:
             if device is None:
@@ -185,26 +190,41 @@ def __init__(self, pool, enabled=True, device=None, group=None):
     def __enter__(self):
         self.mem_context.__enter__()
         if self.group is not None:
-            backend = self.group._get_backend(self.device)
-            try:
-                # Deregister first to avoid duplicate registration of previously
-                # registered memory.
-                backend.deregister_mem_pool(self.pool)
-            except RuntimeError:
-                desc = getattr(self.group, "group_desc", None)
-                print(
-                    f"[MCORE][NCCL_ALLOCATOR] Failed to deregister mem pool from"
-                    f"{repr(self.group)}({desc}) group!!"
-                )
+            # If the pool is not empty, deregister the pool from the group.
+            if self.pool.snapshot():
+                backend = self.group._get_backend(self.device)
+                try:
+                    # Deregister first to avoid duplicate registration of previously
+                    # registered memory.
+                    backend.deregister_mem_pool(self.pool)
+                except RuntimeError:
+                    desc = getattr(self.group, "group_desc", None)
+                    logging.warning(
+                        f"[MCORE][NCCL_ALLOCATOR] Failed to deregister mem pool from"
+                        f"{repr(self.group)}({desc}) group!!"
+                    )
 
     def __exit__(self, *args):
         if self.group is not None:
             backend = self.group._get_backend(self.device)
             try:
-                backend.register_mem_pool(self.pool)
+                # Prefer attempting symmetric registration first; fall back if unsupported.
+                if self.symmetric:
+                    try:
+                        # Since PyTorch PR #161238 symmetric knob is now in registration function.
+                        backend.register_mem_pool(self.pool, symm=self.symmetric)
+                    except TypeError:
+                        # Older PyTorch/APIs without 'symm' keyword.
+                        logging.warning(
+                            f"[MCORE][NCCL_ALLOCATOR] Failed in symmetric registration."
+                            f"Falling back to non-symmetric registration!!"
+                        )
+                        backend.register_mem_pool(self.pool)
+                else:
+                    backend.register_mem_pool(self.pool)
             except RuntimeError:
                 desc = getattr(self.group, "group_desc", None)
-                print(
+                logging.warning(
                     f"[MCORE][NCCL_ALLOCATOR] Failed to register mem pool to"
                     f"{repr(self.group)}({desc}) group!!"
                 )
@@ -238,11 +258,12 @@ class MultiGroupMemPoolAllocator:
     """
 
     def __init__(
-        self, pool, groups
+        self, pool, groups, symmetric=True
     ):  # pool: torch.cuda.MemPool, groups: List[torch.distributed.ProcessGroup]
         self.pool = pool
         self.groups = groups
         self.mem_context = torch.cuda.use_mem_pool(self.pool)
+        self.symmetric = symmetric
 
         assert isinstance(self.pool, torch.cuda.MemPool), "pool must be a torch.cuda.MemPool"
         assert isinstance(self.groups, list), "groups must be a list"
@@ -252,28 +273,43 @@ def __init__(
 
     def __enter__(self):
         self.mem_context.__enter__()
-        for group in self.groups:
-            backend = group._get_backend(torch.device("cuda", torch.cuda.current_device()))
-            try:
-                # Since the registration is done in mempool granularity, we need to deregister
-                # the tensors in the mempool and re-register the mempool including the newly created
-                # tensors after the context is exited.
-                backend.deregister_mem_pool(self.pool)
-            except RuntimeError:
-                desc = getattr(group, "group_desc", None)
-                print(
-                    f"[MCORE][MultiGroupMemPoolAllocator] Failed to deregister mem pool from"
-                    f"{repr(group)}({desc}) group!!"
-                )
+        # If the pool is not empty, deregister the pool from all the groups.
+        if self.pool.snapshot():
+            for group in self.groups:
+                backend = group._get_backend(torch.device("cuda", torch.cuda.current_device()))
+                try:
+                    # Since the registration is done in mempool granularity, we need to deregister
+                    # the tensors in the mempool and re-register the mempool including
+                    # the newly created tensors after the context is exited.
+                    backend.deregister_mem_pool(self.pool)
+                except RuntimeError:
+                    desc = getattr(group, "group_desc", None)
+                    logging.warning(
+                        f"[MCORE][MultiGroupMemPoolAllocator] Failed to deregister mem pool from"
+                        f"{repr(group)}({desc}) group!!"
+                    )
 
     def __exit__(self, *args):
         for group in self.groups:
             backend = group._get_backend(torch.device("cuda", torch.cuda.current_device()))
             try:
-                backend.register_mem_pool(self.pool)
+                # Prefer attempting symmetric registration first; fall back if unsupported.
+                if self.symmetric:
+                    try:
+                        # Since PyTorch PR #161238 symmetric knob is now in registration function.
+                        backend.register_mem_pool(self.pool, symm=self.symmetric)
+                    except TypeError:
+                        # Older PyTorch/APIs without 'symm' keyword.
+                        logging.warning(
+                            f"[MCORE][MultiGroupMemPoolAllocator] Failed in symmetric registration."
+                            f"Falling back to non-symmetric registration!!"
+                        )
+                        backend.register_mem_pool(self.pool)
+                else:
+                    backend.register_mem_pool(self.pool)
             except RuntimeError:
                 desc = getattr(group, "group_desc", None)
-                print(
+                logging.warning(
                     f"[MCORE][MultiGroupMemPoolAllocator] Failed to register mem pool to"
                     f"{repr(group)}({desc}) group!!"
                 )