diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 50fa566d1b6..bab6ed8276b 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/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 diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 8f0117c68ec..d54e6a4f58d 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/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) diff --git a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py index 08f17e54996..695b1af54c8 100644 --- a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py +++ b/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,17 +421,36 @@ 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": @@ -425,6 +458,17 @@ def configure_gtp_remat_from_recipe( 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=""): 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). diff --git a/megatron/core/tensor_parallel/gtp_api.py b/megatron/core/tensor_parallel/gtp_api.py index b49a5c02ded..778f4d979d5 100644 --- a/megatron/core/tensor_parallel/gtp_api.py +++ b/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", diff --git a/megatron/core/tensor_parallel/gtp_symm.py b/megatron/core/tensor_parallel/gtp_symm.py new file mode 100644 index 00000000000..c9b7ed6c4e3 --- /dev/null +++ b/megatron/core/tensor_parallel/gtp_symm.py @@ -0,0 +1,270 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""GTP symmetric-memory pools: the shared registration primitive. + +One ``ncclMemAlloc``-backed ``torch.cuda.MemPool`` per GTP process group, registered once +via ``backend.register_mem_pool(pool, symm=True)``. PyTorch's ProcessGroupNCCL segment hook +then auto-registers every later allocation made under ``gtp_mem_pool_ctx(group)``. With both +ends of a collective in such a pool on the same comm, NCCL selects its symmetric/NVLS kernels. + +This module owns only the pool + registration concern. Consumers ride on ``gtp_mem_pool_ctx``: +``GTPWeightCache`` for AG output buffers, ``RegisteredLifoPool`` for RS send buffers, and +``register_ddp_buffers_on_gtp_groups`` for the DDP param buffer (the AG input). + +Which params participate is decided at wrap time from ``--gtp-nccl-ub`` (dense) / +``--egtp-nccl-ub`` (expert) -- independent of ``--use-nccl-ub``, which covers the DP group -- +and stamped by ``_gtp_attach_attrs`` as ``gtp_smr`` (use this GTP pool) and +``param_needs_nccl_mem`` (back the DDP param buffer with ncclMemAlloc). Sites read both via +getattr; there is no module-level env gate here. + +Launcher must export ``NCCL_NVLS_ENABLE=1`` and +``TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK=0`` before ``init_process_group``. +""" + +import logging +import math +from collections import defaultdict +from contextlib import AbstractContextManager + +import torch +import torch.distributed as dist + +import megatron.core.nccl_allocator as nccl_allocator +from megatron.core.utils import log_single_rank + +logger = logging.getLogger(__name__) + +# group.group_name -> per-group MemPool (one pool per group, registered once). +_pools: "dict[str, torch.cuda.MemPool]" = {} +# group.group_name -> group for pools whose registration is live. A dict (not a set) so the +# shutdown deregister can reach each group; membership still gates re-registration. +_registered: "dict[str, object]" = {} +# group.group_name for groups whose NCCL comm has been warmed (lazy comms are created on the +# first collective). Shared by all GTP register paths so each group is warmed exactly once. +_warmed_groups: "set[str]" = set() + + +def get_gtp_pool(group: dist.ProcessGroup) -> torch.cuda.MemPool: + """Return the per-group ``ncclMemAlloc``-backed MemPool, creating it once.""" + name = group.group_name + pool = _pools.get(name) + if pool is None: + nccl_allocator.init() + pool = nccl_allocator.create_nccl_mem_pool(symmetric=True) + _pools[name] = pool + return pool + + +def _warmup_group_comm(group: dist.ProcessGroup) -> None: + """Force lazy NCCL comm creation for ``group`` once, so a subsequent register_mem_pool + sees an initialized communicator. Idempotent across all GTP register paths.""" + if group.group_name in _warmed_groups: + return + warmup = torch.zeros(1, device=torch.cuda.current_device()) + dist.all_reduce(warmup, group=group) + _warmed_groups.add(group.group_name) + + +def register_gtp_pool(group: dist.ProcessGroup) -> torch.cuda.MemPool: + """Create (if needed) and register the per-group pool on ``group``. Idempotent. + + Call once at model-construction time (before any CUDA-graph capture or + forward), because it issues a collective (comm warmup). After this, the + segment hook auto-registers future segments, so ``gtp_mem_pool_ctx`` itself + is capture-safe and collective-free. + """ + pool = get_gtp_pool(group) + if group.group_name in _registered: + return pool + _warmup_group_comm(group) + nccl_allocator.register_mem_pool(pool, group, symmetric=True) + _registered[group.group_name] = group + log_single_rank( + logger, + logging.INFO, + f"[MCORE][GTP] Registered GTP cache pool on group {group.group_name} " + f"(size={group.size()})", + ) + return pool + + +def gtp_mem_pool_ctx(group: dist.ProcessGroup) -> AbstractContextManager[None]: + """Context manager: allocations inside land in ``group``'s registered pool. + + Pure ``use_mem_pool`` -- no collective -- so it is safe inside CUDA-graph + capture. The pool must already be registered via ``register_gtp_pool`` for + new segments to be window-registered by the segment hook. + """ + return torch.cuda.use_mem_pool(get_gtp_pool(group)) + + +def is_gtp_pool_registered(group: dist.ProcessGroup | None) -> bool: + """True once ``register_gtp_pool`` has registered this group's symmetric pool. + + Full "use the symm pool for this group?" predicate: also rejects None and trivial + (size-1) groups, which are never registered. + """ + return group is not None and group.size() > 1 and group.group_name in _registered + + +class RegisteredLifoPool: + """Group-aware LIFO cache of reduce-scatter send buffers in the per-group symmetric pool. + + Holds the full-shape wgrad the bwd GEMM writes and then scatters. Fresh allocations go + through ``gtp_mem_pool_ctx`` so they are window-registered; freed buffers are recycled + (keyed by numel + dtype + group), keeping memory flat rather than one buffer per weight. + Storage is 1-D so one key serves any shape with that numel, and ``alloc`` returns a view + tagged with ``_gtp_symm_group`` for recycling. + + The free-list grows to peak RS concurrency during the eager warmup iterations, so under + CUDA-graph capture ``alloc`` only ever pops. A fresh allocation during capture would + re-register the pool mid-graph, so it raises instead. + """ + + def __init__(self) -> None: + # (numel, dtype, group_name) -> list of free 1-D buffers. + self._free: "dict[tuple, list]" = defaultdict(list) + + def alloc( + self, + shape: torch.Size | tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + group: dist.ProcessGroup, + ) -> torch.Tensor: + """Return a buffer of ``shape`` from ``group``'s free-list, allocating one if empty. + + Raises RuntimeError if a fresh allocation would be needed during CUDA-graph capture. + """ + numel = int(math.prod(shape)) + bucket = self._free[(numel, dtype, group.group_name)] + if bucket: + flat = bucket.pop() + else: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "[GTP] RegisteredLifoPool exhausted during CUDA-graph capture " + f"(group={group.group_name}, numel={numel}, dtype={dtype}). The " + "eager warmup did not pre-populate enough RS send buffers for " + "the reduce-scatter overlap depth -- run more warmup iters, or " + "the RS concurrency changed between warmup and capture." + ) + # Symm pool only when registered (gtp_smr implies it); else a plain buffer (non-symm). + if is_gtp_pool_registered(group): + with gtp_mem_pool_ctx(group): + flat = torch.empty(numel, dtype=dtype, device=device) + else: + flat = torch.empty(numel, dtype=dtype, device=device) + out = flat.view(shape) + out._gtp_symm_group = group # tag so callers can recycle via tag dispatch + return out + + def free(self, buf: torch.Tensor) -> None: + """Return ``buf`` to the free-list of the group it was allocated from. + + No-op for a buffer this pool did not allocate (no ``_gtp_symm_group`` tag). + """ + group = getattr(buf, "_gtp_symm_group", None) + if group is None: + return + self._free[(buf.numel(), buf.dtype, group.group_name)].append(buf.reshape(-1)) + + def clear(self) -> None: + """Drop every cached buffer. Called at teardown, before the pools they alias go away.""" + self._free.clear() + + +# Imported by generalized_tensor_parallelism. Lives here so deregister_gtp_pools can drop its +# buffers at teardown, alongside the pools they alias. +_gtp_wgrad_pool = RegisteredLifoPool() + + +def _ddp_buffers(ddp_module: torch.nn.Module) -> list: + """All param/grad buffers of a DDP-wrapped module (dense + expert-parallel).""" + return list(getattr(ddp_module, "buffers", [])) + list( + getattr(ddp_module, "expert_parallel_buffers", []) + ) + + +def _buffer_symm_groups(buf) -> list[dist.ProcessGroup]: + """GTP comm groups this buffer's pool must be (de)registered on: params with + ``param_needs_nccl_mem`` set and group size > 1. + + Shared by register and deregister so the two stay in step. Sorted by name because the + collective (de)register order must match across ranks or the comms cross-deadlock. + """ + if getattr(buf, "nccl_mem_pool", None) is None: + return [] + # Needs a param_data all-gather input: without distopt the pool backs only grad_data, which + # the GTP AG never reads, so registering it would cost a window for nothing. + if getattr(buf, "param_data", None) is None: + return [] + groups = {} + for param in buf.params: + if not getattr(param, "param_needs_nccl_mem", False): + continue + group = getattr(param, "group", None) + if group is not None and group.size() > 1: + groups.setdefault(group.group_name, group) + return [group for _, group in sorted(groups.items())] + + +def register_ddp_buffers_on_gtp_groups(ddp_module: torch.nn.Module) -> None: + """Register each DDP buffer's NCCL pool on the GTP group(s) its params opted into + via ``param_needs_nccl_mem``, so the DDP param buffer (the AG input) is in the + symmetric window. The GTP-owned cache/RS pool (AG/RS output) is registered separately + by ``configure_gtp_remat_from_recipe`` before construction. Call once per model chunk + *after* DDP construction (buffers/pools must exist). + + Always symmetric: ``--disable-symmetric-registration`` scopes to the DP-group + registration, not the GTP groups, which are opted into by --gtp-nccl-ub/--egtp-nccl-ub. + """ + for buf in _ddp_buffers(ddp_module): + for group in _buffer_symm_groups(buf): + # buf.nccl_mem_pool is non-None here (checked in _buffer_symm_groups). + _warmup_group_comm(group) + nccl_allocator.register_mem_pool(buf.nccl_mem_pool, group, symmetric=True) + log_single_rank( + logger, + logging.INFO, + f"[MCORE][GTP] Registered DDP param/grad pool on GTP group " + f"{group.group_name} (size={group.size()})", + ) + + +def deregister_ddp_buffers_from_gtp_groups(ddp_module: torch.nn.Module) -> None: + """Mirror of ``register_ddp_buffers_on_gtp_groups``: deregister each buffer's pool from + the same GTP group set. Call at graceful exit *before* the ProcessGroupNCCL destructor + -- window-registered handles left on a comm make its ncclCommDeregister abort + ("Could not find handle"). The DP group is the core buffer's concern, handled + separately by the training loop. + """ + for buf in _ddp_buffers(ddp_module): + for group in _buffer_symm_groups(buf): + # buf.nccl_mem_pool is non-None here (checked in _buffer_symm_groups). + nccl_allocator.deregister_mem_pool(buf.nccl_mem_pool, group) + log_single_rank( + logger, + logging.INFO, + f"[MCORE][GTP] Deregistered DDP param/grad pool from GTP group " + f"{group.group_name} (size={group.size()})", + ) + + +def deregister_gtp_pools() -> None: + """Deregister all GTP-owned symmetric pools. + + Must be called (collectively, on all ranks) before process-group teardown when GTP + symmetric memory was used -- leftover windows abort the ProcessGroupNCCL destructor. + Training shutdown does this; test fixtures enabling GTP NCCL-UB must do likewise. + No-op when nothing was registered. Also drops the recycled RS send buffers, which alias + the pools being torn down here. + """ + for name in sorted(_registered): + nccl_allocator.deregister_mem_pool(_pools[name], _registered[name]) + _registered.clear() + _pools.clear() + _gtp_wgrad_pool.clear() + _warmed_groups.clear() diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index fd30984b38a..df66a03856f 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3014,6 +3014,11 @@ def _add_distributed_args(parser): group.add_argument('--disable-symmetric-registration', action='store_true', dest='disable_symmetric_registration', default=False, help='Disable symmetric (window) registration for NCCL userbuffer registration.' 'This option will force to use conventional (local) userbuffer registration when use-nccl-ub is set.') + group.add_argument('--gtp-nccl-ub', action='store_true', dest='gtp_nccl_ub', + default=False, help='Register dense GTP param buffers with NCCL symmetric memory on the GTP ' + 'group, independent of --use-nccl-ub (which covers the DP group).') + group.add_argument('--egtp-nccl-ub', action='store_true', dest='egtp_nccl_ub', + default=False, help='Like --gtp-nccl-ub but for routed-expert (EGTP) groups.') group.add_argument('--fsdp-manual-registration', action='store_true', dest='fsdp_manual_registration', default=False, help='Manually register the FSDP communication buffers to NCCL user buffer.' 'This option is only effective when use-megatron-fsdp and use-nccl-ub is set.') diff --git a/megatron/training/training.py b/megatron/training/training.py index 57b8efe9390..2b672d27fc4 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2071,6 +2071,9 @@ def _build_model_wrapper(wrap_with_ddp: bool): fp8_recipe=getattr(args, 'fp8_recipe', None), fp8=getattr(args, 'fp8', None) is not None, calculate_per_token_loss=getattr(args, 'calculate_per_token_loss', False), + gtp_nccl_ub=getattr(args, 'gtp_nccl_ub', False), + egtp_nccl_ub=getattr(args, 'egtp_nccl_ub', False), + use_distributed_optimizer=getattr(args, 'use_distributed_optimizer', False), ) model = _build_model_wrapper(wrap_with_ddp) @@ -2089,6 +2092,17 @@ def _build_model_wrapper(wrap_with_ddp: bool): cuda_graph_impl=getattr(args, 'cuda_graph_impl', 'none'), ) + if is_gtp_remat_active(args) and ( + getattr(args, 'gtp_nccl_ub', False) or getattr(args, 'egtp_nccl_ub', False) + ): + from megatron.core.tensor_parallel.gtp_api import register_ddp_buffers_on_gtp_groups + + # Register DDP param buffers on GTP comm groups so both ends of the + # GTP all-gather are in the NCCL symmetric window (enables NVLS). + for model_chunk in model: + if isinstance(model_chunk, DDP): + register_ddp_buffers_on_gtp_groups(model_chunk) + if args.logits_save_dir is not None: from megatron.training.distillation import LogitsSaverHooks @@ -4137,10 +4151,27 @@ def trace_handler(p): # ncclCommDeregister on handles created by ncclCommWindowRegister, # causing "NCCL WARN Deregister: Could not find handle" and a crash. torch.distributed.barrier() + if is_gtp_remat_active(args) and ( + getattr(args, 'gtp_nccl_ub', False) or getattr(args, 'egtp_nccl_ub', False) + ): + from megatron.core.tensor_parallel.gtp_api import ( + deregister_ddp_buffers_from_gtp_groups, + deregister_gtp_pools, + ) + + for model_module in model: + if isinstance(model_module, DDP): + deregister_ddp_buffers_from_gtp_groups(model_module) + # Cache/RS pools are process-global (shared across chunks): deregister once, after + # the per-chunk buffer loop above. + deregister_gtp_pools() + # Deregister the DP-group registration only when it exists (--use-nccl-ub). With + # --gtp-nccl-ub alone the pool was registered on the GTP group (handled above), never on + # DP, so a DP deregister here would abort with "Could not find handle". 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: + if getattr(buf, 'nccl_ub', False) and getattr(buf, 'nccl_mem_pool', None) is not None: nccl_allocator.deregister_mem_pool(buf.nccl_mem_pool, buf.data_parallel_group) one_logger and one_logger.log_metrics( {'app_finish_time': one_logger_utils.get_timestamp_in_ms()}