[Dev] Add zero-CTA all-gather policy for Megatron-FSDP

docs/user-guide/features/megatron_fsdp.md

@@ -172,6 +172,9 @@ unset CUDA_DEVICE_MAX_CONNECTIONS
 --use-nccl-ub
 --fsdp-double-buffer
 --fsdp-manual-registration
+
+# Request NCCL zero-CTA/copy-engine collectives for eligible Megatron-FSDP parameter all-gather buffers. Requires NCCL user buffers and symmetric registration.
+--megatron-fsdp-zero-sm-all-gather
 ```
 
 ### 🤖 Megatron-Core
@@ -591,6 +594,7 @@ Megatron-FSDP sharding and communication buffers support mixed-precision, such t
 | Optimization | Description | `Megatron-Core` Config | `fully_shard` Config |
 |--------------|-------------|----------------------|----------------------|
 | **NCCL User Buffers** | Allocate and register Megatron-FSDP communication buffers with NCCL, which enables zero-`COPY`, high-precision reduction, copy-engine collectives, and symmetric kernels. Uses double buffering. | `--use-nccl-ub` | `nccl_ub=True` |
+| **Zero-SM All-Gather** | Request NCCL zero-CTA/copy-engine policy on dedicated Megatron-FSDP parameter all-gather groups. NCCL uses copy engines for eligible symmetrically registered buffers and falls back otherwise. | `--megatron-fsdp-zero-sm-all-gather` | `megatron_fsdp_zero_sm_all_gather=True` |
 | **NCCL Manual Registration** | Instead of registering NCCL user buffers on first allocation, batch registration of all communication buffers at the end of the initial training step. Reduces registration latency. | `--fsdp-manual-registration` | N/A (Megatron-Core Only) |
 | **Disable Symmetric Registration** | Disable symmetric registration with NCCL. Optional, as symmetric registration failure defaults to normal registration. | `--disable-symmetric-registration` | `disable_symmetric_registration=True` |
 
@@ -605,4 +609,6 @@ NCCL (`v2.27+`) supports symmetric allocation or registration for communicators
 - **Copy-Engine (CE) Collectives**: Instead of using SMs (or CTAs) for common non-computational collectives like AG in Megatron-FSDP, copy engines are instead used to perform all-gather collectives, dedicating SM resources to compute and reduction during FSDP. Requires NCCL `v2.28+`.
 - **High-Precision Reduction**: When training large models, high-precision gradient reduction and accumulation is desired for accuracy and convergence, but communicating FP32 gradients is expensive. With symmetric registration, FP32 accumulators enable gradients to be reduced in FP32 but communicated in BF16, which decreases gradient RS communication latency while maintaining high accuracy during training. Megatron-FSDP supports FP32 main gradient accumulation but BF16 gradient communication, customizable through `megatron_fsdp.MixedPrecisionPolicy`.
 
+`--megatron-fsdp-zero-sm-all-gather` sets `NCCL_CTA_POLICY_ZERO` only on the dedicated dense and expert parameter all-gather communicators. Zero-CTA requires CUDA driver `12.5+` and NCCL `v2.28+` within supported NVLink domains; network transport requires NCCL `v2.30.6+`. The policy is best-effort: buffers that are not symmetrically registered, including FSDP buckets that fall back outside the fixed double-buffer pool, use NCCL's kernel path. The zero-CTA policy prioritizes SM availability and may increase isolated collective latency, so evaluate end-to-end compute/communication overlap rather than collective latency alone.
+
 These optimizations significantly reduce SM resource contention for overlapped compute and communication kernels in FSDP. Symmetric registration, allocation, and pooling is also supported in PyTorch: [`torch.distributed._symmetric_memory`](https://docs.pytorch.org/docs/stable/symmetric_memory.html).

megatron/core/dist_checkpointing/strategies/nvrx.py

@@ -17,6 +17,9 @@
 
 def has_nvrx_async_support() -> bool:
     """Checks whether the NVRx async checkpointing symbols Megatron uses are importable."""
+    if not is_nvrx_min_version():
+        return False
+
     try:
         core = import_module("nvidia_resiliency_ext.checkpointing.async_ckpt.core")
         cached_metadata_reader = import_module(
@@ -41,9 +44,6 @@ def has_nvrx_async_support() -> bool:
         getattr(state_dict_saver, "save_state_dict_async_finalize", None),
         getattr(state_dict_saver, "save_state_dict_async_plan", None),
     )
-    assert (
-        is_nvrx_min_version()
-    ), f"Minimum required nvidia-resiliency-ext package version is {NVRX_MIN_VERSION}."
 
     return all(symbol is not None for symbol in required_symbols) and hasattr(
         filesystem_async, "_results_queue"
@@ -71,9 +71,7 @@ def make_nvrx_async_request(
 def is_nvrx_min_version(version: str = NVRX_MIN_VERSION) -> bool:
     """Check if minimum version of `NVRx` is installed."""
     if not HAVE_PACKAGING:
-        raise ImportError(
-            "packaging is not installed. Please install it with `pip install packaging`."
-        )
+        return False
 
     try:
         import nvidia_resiliency_ext as nvrx

megatron/core/dist_checkpointing/strategies/torch.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 """ Strategies using PyTorch distributed.checkpoint as an underlying format. """
 import inspect
@@ -50,7 +50,12 @@
 )
 from .async_utils import AsyncRequest
 from .checkpointable import CheckpointableShardedTensor, LocalShardsContainer
-from .nvrx import has_nvrx_async_support, make_nvrx_async_request
+from .nvrx import (
+    NVRX_MIN_VERSION,
+    has_nvrx_async_support,
+    is_nvrx_min_version,
+    make_nvrx_async_request,
+)
 
 if TYPE_CHECKING:
     from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncRequest as NVRxAsyncRequest
@@ -1041,6 +1046,12 @@ def remove_sharded_tensors(self, checkpoint_dir: str, key_prefix: str):
 def get_async_strategy(async_strategy: str = "nvrx", module: str = None) -> tuple:
     """Returns async strategy and related async imported modules"""
     if async_strategy == "nvrx":
+        if not is_nvrx_min_version():
+            raise ModuleNotFoundError(
+                "A compatible `nvidia-resiliency-ext` installation is required for "
+                f'`async_strategy="nvrx"`; minimum version is {NVRX_MIN_VERSION}. '
+                'Please install it or set `async_strategy` to `mcore`.'
+            )
         try:
             # nvrx async imports
             from nvidia_resiliency_ext.checkpointing.async_ckpt.cached_metadata_filesystem_reader import (  # pylint: disable=line-too-long

megatron/core/distributed/distributed_data_parallel_config.py

@@ -177,6 +177,13 @@ class DistributedDataParallelConfig:
       when nccl_ub is set.
     """
 
+    megatron_fsdp_zero_sm_all_gather: bool = False
+    """If true, request NCCL zero-CTA for Megatron-FSDP parameter all-gather
+      copy-engine collectives. Requires Megatron-FSDP, NCCL user-buffer registration,
+      symmetric registration, and dedicated all-gather process groups. NCCL uses zero
+      CTA whenever the collective and its buffers are eligible.
+    """
+
     fsdp_manual_registration: bool = False
     """If true, manually register the FSDP communication buffers to NCCL user buffer.
       This option is only effective when use_megatron_fsdp and nccl_ub is set.
@@ -281,6 +288,19 @@ def __post_init__(self):
                     "with nccl_ub due to compatibility issue with torch.cuda.MemPool API."
                 )
 
+        if self.megatron_fsdp_zero_sm_all_gather:
+            if not self.use_megatron_fsdp:
+                raise ValueError(
+                    "megatron_fsdp_zero_sm_all_gather is only supported with Megatron-FSDP."
+                )
+            if not self.nccl_ub:
+                raise ValueError(
+                    "megatron_fsdp_zero_sm_all_gather requires NCCL user-buffer registration."
+                )
+            if self.disable_symmetric_registration:
+                raise ValueError(
+                    "megatron_fsdp_zero_sm_all_gather requires symmetric NCCL registration."
+                )
         if len(self.param_name_patterns_for_fp32_local_accumulation) > 0:
             assert not self.grad_reduce_in_fp32, (
                 "Only need to explicitly specify param_name patterns for FP32 local accumulation "

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

@@ -114,6 +114,13 @@ class DistributedDataParallelConfig:
       when nccl_ub is set.
     """
 
+    megatron_fsdp_zero_sm_all_gather: bool = False
+    """If true, request NCCL zero-CTA for Megatron-FSDP parameter all-gather
+      copy-engine collectives. Requires NCCL user-buffer registration, symmetric
+      registration, and dedicated all-gather process groups. NCCL uses zero CTA whenever
+      the collective and its buffers are eligible.
+    """
+
     fsdp_manual_registration: bool = False
     """If true, manually register the FSDP communication buffers to NCCL user buffer.
       This option is only effective when use_megatron_fsdp and nccl_ub is set.
@@ -203,3 +210,13 @@ def __post_init__(self):
                     "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is currently not supported "
                     "with nccl_ub due to compatibility issue with torch.cuda.MemPool API."
                 )
+
+        if self.megatron_fsdp_zero_sm_all_gather:
+            if not self.nccl_ub:
+                raise ValueError(
+                    "megatron_fsdp_zero_sm_all_gather requires NCCL user-buffer registration."
+                )
+            if self.disable_symmetric_registration:
+                raise ValueError(
+                    "megatron_fsdp_zero_sm_all_gather requires symmetric NCCL registration."
+                )

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

@@ -1736,6 +1736,19 @@ def __init__(
         )
         self.ubr_groups = None
         self.already_registered = False
+        if self.ddp_config.megatron_fsdp_zero_sm_all_gather:
+            assert self.ddp_config.nccl_ub, (
+                "Megatron-FSDP zero-SM all-gather requires NCCL user-buffer registration. "
+                "Please set nccl_ub=True."
+            )
+            assert not self.ddp_config.disable_symmetric_registration, (
+                "Megatron-FSDP zero-SM all-gather requires symmetric NCCL registration. "
+                "Please leave disable_symmetric_registration=False."
+            )
+            assert NCCL_ALLOCATOR == "MCORE", (
+                "Megatron-FSDP zero-SM all-gather requires megatron.core.nccl_allocator "
+                "for symmetric NCCL memory allocation."
+            )
         # User buffer registration related settings
         if self.ddp_config.nccl_ub:
             assert nccl_allocator is not None, (

megatron/core/parallel_state.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 """Model and data parallel groups."""
 
@@ -146,6 +146,68 @@
 _global_process_group_list = None
 
 
+def _get_nccl_cta_policy_value(cta_policy):
+    """Map a NCCL CTA policy config value to ProcessGroupNCCL's enum value."""
+    if isinstance(cta_policy, int):
+        return cta_policy
+    if not isinstance(cta_policy, str):
+        raise RuntimeError(
+            f"cta_policy ({cta_policy}) must be an int or one of "
+            "'default', 'efficiency', or 'zero'."
+        )
+
+    policy_name = cta_policy.lower()
+    pg_nccl = torch.distributed.ProcessGroupNCCL
+    cta_policies = {
+        "default": "NCCL_CTA_POLICY_DEFAULT",
+        "efficiency": "NCCL_CTA_POLICY_EFFICIENCY",
+        "zero": "NCCL_CTA_POLICY_ZERO",
+    }
+    if policy_name not in cta_policies:
+        raise RuntimeError(
+            f"cta_policy ({cta_policy}) is not supported. "
+            "Accepted values: 'default', 'efficiency', 'zero', or an integer NCCL policy value."
+        )
+    attr_name = cta_policies[policy_name]
+    if not hasattr(pg_nccl, attr_name):
+        raise RuntimeError(
+            f"cta_policy={cta_policy!r} requires torch.distributed.ProcessGroupNCCL.{attr_name}, "
+            "which is not available in this PyTorch build."
+        )
+    return getattr(pg_nccl, attr_name)
+
+
+def _load_nccl_comm_cfgs(nccl_communicator_config_path):
+    """Load NCCL communicator configs from YAML, returning an empty dict when unset."""
+    if nccl_communicator_config_path is None:
+        return {}
+
+    try:
+        import yaml
+    except ImportError:
+        raise RuntimeError(
+            "Cannot import `yaml`. Setting custom nccl communicator configs "
+            "requires the yaml package."
+        )
+
+    with open(nccl_communicator_config_path, "r") as stream:
+        return yaml.safe_load(stream) or {}
+
+
+def _apply_high_priority_stream_groups(nccl_comm_cfgs, high_priority_stream_groups):
+    """Apply high-priority stream overrides to a NCCL communicator config map."""
+    high_priority_stream_groups = high_priority_stream_groups or []
+    for pg_name in high_priority_stream_groups:
+        overwrite_nccl_comm_cfgs(nccl_comm_cfgs, pg_name, ("is_high_priority_stream", True))
+
+
+def _apply_zero_sm_all_gather_group_options(nccl_comm_cfgs, for_expert_parallelism):
+    """Configure FSDP all-gather groups to use NCCL's zero-CTA policy."""
+    overwrite_nccl_comm_cfgs(nccl_comm_cfgs, "dp_cp_ag", ("cta_policy", "zero"))
+    if for_expert_parallelism:
+        overwrite_nccl_comm_cfgs(nccl_comm_cfgs, "ep_dp_ag", ("cta_policy", "zero"))
+
+
 def get_nccl_options(pg_name, nccl_comm_cfgs):
     """Set the NCCL process group options.
 
@@ -168,6 +230,10 @@ def get_nccl_options(pg_name, nccl_comm_cfgs):
             nccl_options.config.max_ctas = nccl_comm_cfgs[pg_name]["max_ctas"]
         if "min_ctas" in nccl_comm_cfgs[pg_name]:
             nccl_options.config.min_ctas = nccl_comm_cfgs[pg_name]["min_ctas"]
+        if "cta_policy" in nccl_comm_cfgs[pg_name]:
+            nccl_options.config.cta_policy = _get_nccl_cta_policy_value(
+                nccl_comm_cfgs[pg_name]["cta_policy"]
+            )
         if "net_name" in nccl_comm_cfgs[pg_name]:
             nccl_options.config.net_name = nccl_comm_cfgs[pg_name]["net_name"]
             # verify net_name value
@@ -181,6 +247,19 @@ def get_nccl_options(pg_name, nccl_comm_cfgs):
         return None
 
 
+def _get_nccl_options_with_fallback(primary_pg_name, fallback_pg_name, nccl_comm_cfgs):
+    """Get NCCL options for a process group, inheriting base-group config when present."""
+    if primary_pg_name in nccl_comm_cfgs and fallback_pg_name in nccl_comm_cfgs:
+        merged_cfgs = {
+            primary_pg_name: {**nccl_comm_cfgs[fallback_pg_name], **nccl_comm_cfgs[primary_pg_name]}
+        }
+        return get_nccl_options(primary_pg_name, merged_cfgs)
+
+    return get_nccl_options(primary_pg_name, nccl_comm_cfgs) or get_nccl_options(
+        fallback_pg_name, nccl_comm_cfgs
+    )
+
+
 def update_pg_timeout(
     timeout: timedelta, pg: Optional[torch._C._distributed_c10d.ProcessGroup] = None
 ):
@@ -641,8 +720,8 @@ def initialize_model_parallel(
 
         nccl_communicator_config_path (str, default = None):
             Path to the yaml file of NCCL communicator configurations.
-            `min_ctas`, `max_ctas`, and `cga_cluster_size` can be set
-            for each communicator.
+            `min_ctas`, `max_ctas`, `cga_cluster_size`, and `cta_policy`
+            can be set for each communicator.
 
         distributed_timeout_minutes (int, default = 30): Timeout, in
             minutes,for operations executed against distributed
@@ -748,23 +827,8 @@ def initialize_model_parallel(
 
     rank = torch.distributed.get_rank()
 
-    nccl_comm_cfgs = {}
-    if nccl_communicator_config_path is not None:
-        try:
-            import yaml
-        except ImportError:
-            raise RuntimeError(
-                "Cannot import `yaml`. Setting custom nccl communicator configs "
-                "requires the yaml package."
-            )
-
-        with open(nccl_communicator_config_path, "r") as stream:
-            nccl_comm_cfgs = yaml.safe_load(stream)
-
-    # Set is_high_priority_stream flag to the nccl_comm_cfgs if it is in high_priority_stream_groups
-    high_priority_stream_groups = high_priority_stream_groups or []
-    for pg_name in high_priority_stream_groups:
-        overwrite_nccl_comm_cfgs(nccl_comm_cfgs, pg_name, ("is_high_priority_stream", True))
+    nccl_comm_cfgs = _load_nccl_comm_cfgs(nccl_communicator_config_path)
+    _apply_high_priority_stream_groups(nccl_comm_cfgs, high_priority_stream_groups)
 
     decoder_rank_generator = RankGenerator(
         tp=tensor_model_parallel_size,
@@ -1353,7 +1417,14 @@ def initialize_model_parallel(
     _set_global_memory_buffer()
 
 
-def create_all_gather_groups(for_expert_parallelism=False, timeout=None, nccl_comm_cfgs=None):
+def create_all_gather_groups(
+    for_expert_parallelism=False,
+    timeout=None,
+    nccl_comm_cfgs=None,
+    nccl_communicator_config_path=None,
+    high_priority_stream_groups=None,
+    zero_sm_all_gather=False,
+):
     """
     Helper function to create all-gather process groups for AG/RS overlap.
 
@@ -1364,6 +1435,11 @@ def create_all_gather_groups(for_expert_parallelism=False, timeout=None, nccl_co
         for_expert_parallelism (bool): If True, also creates AG group for expert parameters.
         timeout (timedelta): Timeout for distributed collectives.
         nccl_comm_cfgs (dict): NCCL communicator configurations.
+        nccl_communicator_config_path (str): Path to NCCL communicator YAML.
+        high_priority_stream_groups (List[str]): Communicator groups that should use
+            high priority streams.
+        zero_sm_all_gather (bool): If true, request NCCL's zero-CTA policy on the
+            dedicated all-gather groups. NCCL falls back for ineligible buffers.
 
     Returns:
         tuple: (dp_cp_ag_group, expt_dp_ag_group) where expt_dp_ag_group is None
@@ -1386,6 +1462,12 @@ def create_all_gather_groups(for_expert_parallelism=False, timeout=None, nccl_co
             "Call initialize_model_parallel() first."
         )
 
+    if nccl_comm_cfgs is None:
+        nccl_comm_cfgs = _load_nccl_comm_cfgs(nccl_communicator_config_path)
+    _apply_high_priority_stream_groups(nccl_comm_cfgs, high_priority_stream_groups)
+    if zero_sm_all_gather:
+        _apply_zero_sm_all_gather_group_options(nccl_comm_cfgs, for_expert_parallelism)
+
     rank = torch.distributed.get_rank()
     pp_size = get_pipeline_model_parallel_world_size()
     cp_size = get_context_parallel_world_size()
@@ -1403,7 +1485,7 @@ def create_all_gather_groups(for_expert_parallelism=False, timeout=None, nccl_co
         group_with_cp_ag = create_group(
             ranks_with_cp,
             timeout=timeout,
-            pg_options=get_nccl_options('dp_cp', nccl_comm_cfgs or {}),
+            pg_options=_get_nccl_options_with_fallback('dp_cp_ag', 'dp_cp', nccl_comm_cfgs),
             group_desc='DATA_PARALLEL_GROUP_WITH_CP_AG',
         )
         if rank in ranks_with_cp:
@@ -1429,7 +1511,7 @@ def create_all_gather_groups(for_expert_parallelism=False, timeout=None, nccl_co
             expert_dp_ag = create_group(
                 expert_dp_ranks,
                 timeout=timeout,
-                pg_options=get_nccl_options("ep_dp", nccl_comm_cfgs or {}),
+                pg_options=_get_nccl_options_with_fallback("ep_dp_ag", "ep_dp", nccl_comm_cfgs),
                 group_desc='EXPERT_DATA_PARALLEL_GROUP_AG',
             )
             if rank in expert_dp_ranks: