[GTP] NCCL Symmetric Memory Registration Support

megatron/core/distributed/param_and_grad_buffer.py

@@ -1145,6 +1145,16 @@ def __init__(
             tmp_warmup_tensor = torch.zeros([1], device="cuda")
             torch.distributed.all_reduce(tmp_warmup_tensor, group=self.data_parallel_group)
             torch.distributed.barrier()
+        elif self.ddp_config.use_distributed_optimizer and any(
+            getattr(p, 'param_needs_nccl_mem', False) for p in self.params
+        ):
+            # Params opted into ncclMemAlloc backing so the caller can window-register this
+            # buffer on other comm groups; it registers and warms them itself, so this branch
+            # only allocates. Distopt-gated: without it there is no param_data to back.
+            nccl_allocator.init()
+            pool = nccl_allocator.create_nccl_mem_pool(symmetric=True)
+            self.nccl_mem_pool = pool
+            mem_alloc_context = functools.partial(torch.cuda.use_mem_pool, pool)
         else:
             # If nccl_ub is False, mem_alloc_context is nullcontext.
             mem_alloc_context = nullcontext

megatron/core/extensions/transformer_engine.py

@@ -498,7 +498,15 @@ def _init_gtp_remat_context(
         rng_via_kwarg=rng_via_kwarg,
         out_split_size=out_split_size,
     )
-    yield out_features
+    # Route super().__init__'s pre-sharded weight (native FP8/MXFP8 storage under
+    # --fp8-param-gather) into the registered symmetric pool -> NVLS-eligible GTP all-gather input.
+    from megatron.core.tensor_parallel.gtp_api import gtp_mem_pool_ctx, is_gtp_pool_registered
+
+    if is_gtp_pool_registered(gtp_remat_group):
+        with gtp_mem_pool_ctx(gtp_remat_group):
+            yield out_features
+    else:
+        yield out_features
     _gtp_attach_post_init(module, gtp_ctx, is_grouped=is_grouped)
 
 

megatron/core/tensor_parallel/generalized_tensor_parallelism.py

@@ -36,6 +36,13 @@
 import torch
 from packaging.version import Version
 
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.tensor_parallel.gtp_symm import (
+    _gtp_wgrad_pool,
+    gtp_mem_pool_ctx,
+    is_gtp_pool_registered,
+    register_gtp_pool,
+)
 from megatron.core.utils import log_single_rank
 
 logger = logging.getLogger(__name__)
@@ -384,6 +391,13 @@ class GTPRematConfig:
     # DDP's 1/replicate scaling to yield the full (replicate x gtp) mean.
     calculate_per_token_loss: bool = False
 
+    # Back the AG output cache, RS send buffers and (under distopt) the DDP param buffer with
+    # ncclMemAlloc pools registered on the GTP / EGTP group. Independent of --use-nccl-ub (DP).
+    gtp_nccl_ub: bool = False
+    egtp_nccl_ub: bool = False
+    # Without distopt the slice-path shard keeps its own storage, so it goes to the symm pool.
+    use_distributed_optimizer: bool = False
+
 
 GTP_CONFIG = GTPRematConfig()
 
@@ -407,24 +421,54 @@ def tag_gtp_params_with_names(model):
             param._debug_name = name
 
 
+def _initialize_gtp_symmetric_group(group):
+    """Register a GTP/EGTP group's symmetric pool (idempotent; no-op for absent/trivial groups)."""
+    if group is not None and group.size() > 1:
+        register_gtp_pool(group)
+
+
 def configure_gtp_remat_from_recipe(
-    *, fp4=False, fp8_recipe=None, fp8=False, calculate_per_token_loss=False
+    *,
+    fp4=False,
+    fp8_recipe=None,
+    fp8=False,
+    calculate_per_token_loss=False,
+    gtp_nccl_ub=False,
+    egtp_nccl_ub=False,
+    use_distributed_optimizer=False,
 ):
     """
-    Configure GTP weight-remat (padding + loss reduction) from the quantization recipe.
+    Configure GTP weight-remat (padding + loss reduction + symm mem) from the training recipe.
     Must be called once BEFORE model construction.
     """
     # gtp_remat grad reduction SUMs (not means) the gtp_remat axis under per-token-loss.
     # check_param_states=False: GTP buffer reuse (notably under CUDA-graph capture) trips the
     # param-state debug asserts, so keep them off for GTP runs.
-    update_gtp_config(calculate_per_token_loss=calculate_per_token_loss, check_param_states=False)
+    update_gtp_config(
+        calculate_per_token_loss=calculate_per_token_loss,
+        check_param_states=False,
+        gtp_nccl_ub=gtp_nccl_ub,
+        egtp_nccl_ub=egtp_nccl_ub,
+        use_distributed_optimizer=use_distributed_optimizer,
+    )
     if fp4:
         update_gtp_config(pad_for_alignment=16)
     elif fp8_recipe == "mxfp8":
         update_gtp_config(pad_for_alignment=32)
     elif fp8:
         update_gtp_config(pad_for_alignment=16)
 
+    # Register the dense-GTP and EGTP pools centrally (pre-construction, before any forward),
+    # so all modules -- including TE pre-sharded ones -- share the registered per-group pools.
+    if gtp_nccl_ub or egtp_nccl_ub:
+        pg_collection = ProcessGroupCollection.use_mpu_process_groups(
+            required_pgs=["gtp_remat", "expt_gtp_remat"]
+        )
+        if gtp_nccl_ub:
+            _initialize_gtp_symmetric_group(pg_collection.gtp_remat)
+        if egtp_nccl_ub:
+            _initialize_gtp_symmetric_group(pg_collection.expt_gtp_remat)
+
     if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
         logger.info("> GTP_remat enabled. %s", GTP_CONFIG)
 
@@ -499,7 +543,13 @@ def _gtp_slice_one_param(param, gtp_remat_group, *, name="<unnamed>"):
 
     shard_size = tensor.shape[0] // gtp_remat_size
     shard = tensor[gtp_rank * shard_size : (gtp_rank + 1) * shard_size]
-    gtp_shard = GTPShardedParam(shard.clone())
+    # No distopt: the shard keeps this storage, so it must be the symm AG input. With distopt it
+    # is flattened into param_data (which carries symmetry) and then freed -- skip the register.
+    if is_gtp_pool_registered(gtp_remat_group) and not GTP_CONFIG.use_distributed_optimizer:
+        with gtp_mem_pool_ctx(gtp_remat_group):
+            gtp_shard = GTPShardedParam(shard.clone())
+    else:
+        gtp_shard = GTPShardedParam(shard.clone())
     gtp_shard.pad_length = pad_length
     # Preserve the source weight's TP attributes (dropped when wrapping into GTPShardedParam),
     # so param_is_not_tensor_parallel_duplicate still classifies it without GTP-specific code.
@@ -528,6 +578,18 @@ def _gtp_attach_attrs(gtp_shard, gtp_remat_group, *, is_grouped=False, expert_id
         gtp_shard.chain_id = GTPChain.UNGRAPHED.value
     gtp_shard.group = gtp_remat_group
     gtp_shard.gtp_remat_size = gtp_remat_group.size()
+
+    # gtp_smr: use per-group ncclMemAlloc pool for AG/RS buffers.
+    # param_needs_nccl_mem: signal DDP to allocate param_data from ncclMemAlloc (AG input).
+    # Expert vs dense by group identity check
+    _expert_group = ProcessGroupCollection.use_mpu_process_groups(
+        required_pgs=["expt_gtp_remat"]
+    ).expt_gtp_remat
+    _is_smr_enabled = (
+        GTP_CONFIG.egtp_nccl_ub if gtp_remat_group is _expert_group else GTP_CONFIG.gtp_nccl_ub
+    )
+    gtp_shard.gtp_smr = _is_smr_enabled
+    gtp_shard.param_needs_nccl_mem = _is_smr_enabled
     global _GTP_PARAMS
     _GTP_PARAMS.append(gtp_shard)
 
@@ -820,6 +882,9 @@ def _init_gtp_runtime_attrs(obj):
     obj._wgrad_rs_handle = None
     obj.rs_event = torch.cuda.Event(external=True)
     obj._rs_ticket = None
+    # Persistent symm RS send buffer (padded shape, window-registered via LIFO pool).
+    # None for non-symm params or before the first RS.
+    obj._wgrad_padded_buf = None
     # Padding
     obj.pad_length = 0
     # Debug
@@ -1443,8 +1508,37 @@ def batched_all_gather_and_prefetch(self, **kwargs):
         return self.all_gather_and_prefetch(**kwargs)
 
     def get_wgrad_tensor(self):
-        """Pool-allocate a wgrad scratch tensor of unsharded shape for the bwd GEMM."""
-        return _wgrad_pool_get(self._unsharded_shape, self.main_grad.dtype, self.device)
+        """Pool-allocate a wgrad scratch tensor of unsharded shape for the bwd GEMM.
+
+        Symm-RS path: return the [:unsharded] view of a persistent padded send buffer
+        from the LIFO pool. The bwd GEMM writes into the view; the RS sends the full
+        padded buffer directly (already window-registered), skipping F.pad.
+        """
+        if not getattr(self, "gtp_smr", False):
+            return _wgrad_pool_get(self._unsharded_shape, self.main_grad.dtype, self.device)
+        if self._wgrad_padded_buf is None:
+            self._wgrad_padded_buf = _gtp_wgrad_pool.alloc(
+                self._unsharded_shape_padded, self.main_grad.dtype, self.device, self.group
+            )
+            # The GEMM writes only the [:unsharded] rows but the RS sends the whole buffer;
+            # zero the pad tail so recycled LIFO storage (keyed by numel, not shape) can't
+            # scatter stale rows into the last rank's main_grad.
+            if self.pad_length:
+                self._wgrad_padded_buf[self._unsharded_shape[0] :].zero_()
+        return self._wgrad_padded_buf[: self._unsharded_shape[0]]
+
+    def _release_wgrad_sendbufs(self):
+        """Return this chain's symm RS send buffers to the LIFO after RS completes -- the release
+        counterpart of get_wgrad_tensor. Non-symm chains never allocated one, so skip without
+        touching self._weights. The LIFO holds buffers at stable addresses, so the same set
+        recycles across params (UNGRAPHED and GRAPHED) rather than pinning one per weight.
+        """
+        if not getattr(self, "gtp_smr", False):
+            return
+        for w in self._weights:
+            if w._wgrad_padded_buf is not None:
+                _gtp_wgrad_pool.free(w._wgrad_padded_buf)
+                w._wgrad_padded_buf = None
 
     def register_grad_accum_hook(self, grad_accum_node, hook):
         """Register a DDP backward hook to call after the wgrad RS finalize.
@@ -1511,6 +1605,7 @@ def _wait_reduce_scatter(self, finalize_grad=False):
                 for buf in self._wgrad_input_bufs:
                     _wgrad_pool_put(buf)
             self._wgrad_input_bufs = None
+        self._release_wgrad_sendbufs()
 
     def _prescale_wgrads_for_mean_rs(self, wgrads):
         """Pre-scale wgrad by 1/gtp_remat so the SUM reduce-scatter yields the gtp_remat mean.
@@ -1539,7 +1634,16 @@ def _reduce_scatter(self, wgrads, async_op, nvtx_label=None):
             for w in self._weights:
                 w._set_rs_state(new_rs_state)
 
-        if self.pad_length > 0:
+        # Symm-RS: if the bwd GEMM wrote into the registered padded send buffer (confirmed by
+        # storage aliasing), send it directly -- F.pad would make an unregistered copy.
+        use_persistent_wgrad = self.pad_length > 0 and all(
+            getattr(w, "_wgrad_padded_buf", None) is not None
+            and w._wgrad_padded_buf.untyped_storage().data_ptr() == g.untyped_storage().data_ptr()
+            for w, g in zip(self._weights, wgrads)
+        )
+        if use_persistent_wgrad:
+            wgrads = [w._wgrad_padded_buf for w in self._weights]
+        elif self.pad_length > 0:
             wgrads = [torch.nn.functional.pad(w, (0, 0, 0, self.pad_length)) for w in wgrads]
 
         if async_op:
@@ -1630,6 +1734,7 @@ def wgrad_reduce_scatter(self, wgrad, nvtx_label=None):
             if poolable:
                 for buf in wgrads:
                     _wgrad_pool_put(buf)
+            self._release_wgrad_sendbufs()
             ret = result if batched else result[0]
 
         # Wait for last reduce scatter if it was async
@@ -1802,8 +1907,17 @@ def _allocate_buffer(
         else:
             out_shape = param._unsharded_shape_padded
 
-        # Route GRAPHED-chain buffers into the CG mempool at creation (see _graphed_alloc).
-        with _graphed_alloc(getattr(param, "chain_id", GTPChain.UNGRAPHED.value)):
+        # Only the AG output needs the symmetric window (store-multicast); the RS output is a
+        # local write -> _graphed_alloc. The registered check keeps custom/unregistered groups
+        # (e.g. in tests) on the plain path.
+        symm = getattr(param, 'gtp_smr', False) and not reduce_scatter
+        group = getattr(param, "group", None)
+        if symm and is_gtp_pool_registered(group):
+            alloc_ctx = gtp_mem_pool_ctx(group)
+        else:
+            alloc_ctx = _graphed_alloc(getattr(param, "chain_id", GTPChain.UNGRAPHED.value))
+
+        with alloc_ctx:
             if not isinstance(dtype, torch.dtype):
                 # Use the gather quantizer copy: mutating the param's own quantizer usage
                 # would corrupt the optimizer's quantize_ update direction (frozen weights).

megatron/core/tensor_parallel/gtp_api.py

@@ -10,6 +10,16 @@
 so no core module uses GTP symbols without TE.
 """
 
+# Symmetric-memory helpers, outside the HAVE_TE guard on purpose: gtp_symm has no TE dependency
+# and shutdown must be able to deregister pools even where the TE-backed surface is unavailable.
+from megatron.core.tensor_parallel.gtp_symm import (
+    deregister_ddp_buffers_from_gtp_groups,
+    deregister_gtp_pools,
+    gtp_mem_pool_ctx,
+    is_gtp_pool_registered,
+    register_ddp_buffers_on_gtp_groups,
+)
+
 try:
     from megatron.core.tensor_parallel.generalized_tensor_parallelism import (
         HAVE_TE,
@@ -40,6 +50,11 @@
 
 __all__ = [
     "HAVE_GTP",
+    "deregister_ddp_buffers_from_gtp_groups",
+    "deregister_gtp_pools",
+    "gtp_mem_pool_ctx",
+    "is_gtp_pool_registered",
+    "register_ddp_buffers_on_gtp_groups",
     "GTPChain",
     "GTPEmbeddingWeight",
     "attach_gtp_to_presharded_module",