From 7ad95258c297d586fa54a0a43868eebe151a832b Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Sun, 28 Jun 2026 11:43:13 +0000 Subject: [PATCH 1/2] perf: add zero-sm all-gather for Megatron-FSDP Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- docs/user-guide/features/megatron_fsdp.md | 6 + .../distributed_data_parallel_config.py | 20 +++ .../distributed_data_parallel_config.py | 17 +++ .../megatron_fsdp/param_and_grad_buffer.py | 13 ++ megatron/core/parallel_state.py | 128 ++++++++++++++---- megatron/training/arguments.py | 26 ++++ megatron/training/config/common_config.py | 2 +- megatron/training/training.py | 8 +- .../test_mcore_fully_sharded_data_parallel.py | 62 ++++++++- .../distributed/test_param_and_grad_buffer.py | 83 +++++++++++- tests/unit_tests/test_parallel_state.py | 50 ++++++- 11 files changed, 387 insertions(+), 28 deletions(-) diff --git a/docs/user-guide/features/megatron_fsdp.md b/docs/user-guide/features/megatron_fsdp.md index 36fcc68893c..2e6a59e7a46 100644 --- a/docs/user-guide/features/megatron_fsdp.md +++ b/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). diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 56ec9e89539..6ecb4c2b13a 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/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 " diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py index 938e17a5b3f..9a2856283ac 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py +++ b/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." + ) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py index 4a0b854abde..72b328909ff 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py +++ b/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, ( diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py index 863b5d55d9d..f51d09dab70 100644 --- a/megatron/core/parallel_state.py +++ b/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: diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index f082674935c..a91c4667b36 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1304,6 +1304,23 @@ def validate_args(args, defaults={}): "--megatron-fsdp-cache-param-bucket-views requires " "--use-megatron-fsdp." ) + if args.megatron_fsdp_zero_sm_all_gather: + assert ( + args.use_megatron_fsdp + ), "Megatron-FSDP zero-SM all-gather requires --use-megatron-fsdp." + assert args.nccl_ub, "Megatron-FSDP zero-SM all-gather requires --use-nccl-ub." + assert not args.disable_symmetric_registration, ( + "Megatron-FSDP zero-SM all-gather requires symmetric NCCL registration. " + "Do not set --disable-symmetric-registration." + ) + if not args.create_all_gather_group: + args.create_all_gather_group = True + warn_rank_0( + 'Megatron-FSDP zero-SM all-gather requires a dedicated all-gather ' + 'process group. Enabling --create-all-gather-group.', + args.rank, + ) + if args.nccl_ub and args.use_megatron_fsdp: # In Megatron-LM, required implementation for manual registration is already provided. # So we enable the manual registration by default when nccl-ub and use_megatron_fsdp is set. @@ -4121,6 +4138,15 @@ def _add_distributed_args(parser): 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( + '--megatron-fsdp-zero-sm-all-gather', + action='store_true', + dest='megatron_fsdp_zero_sm_all_gather', + default=False, + help='Request NCCL zero-CTA/copy-engine collectives for eligible Megatron-FSDP ' + 'parameter all-gather buffers. Requires --use-megatron-fsdp, --use-nccl-ub, ' + 'symmetric registration, and dedicated all-gather process groups.', + ) group.add_argument( '--fsdp-manual-registration', action='store_true', diff --git a/megatron/training/config/common_config.py b/megatron/training/config/common_config.py index 2107816bd85..f42b95d0d33 100644 --- a/megatron/training/config/common_config.py +++ b/megatron/training/config/common_config.py @@ -103,7 +103,7 @@ class DistributedInitConfig: nccl_communicator_config_path: str | None = None """Path to the yaml file with NCCL communicator configurations. The number of min/max thread groups and thread group cluster size of each communicator can be configured by setting - `min_ctas`, `max_ctas`, and `cga_cluster_size`.""" + `min_ctas`, `max_ctas`, `cga_cluster_size`, and `cta_policy`.""" use_tp_pp_dp_mapping: bool = False """If set, distributed ranks initialize order is changed from tp-cp-ep-dp-pp to tp-cp-ep-pp-dp. diff --git a/megatron/training/training.py b/megatron/training/training.py index e4075078f0e..ca06c7312b5 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2058,12 +2058,18 @@ def get_model( else None ) dp_cp_ag, expt_dp_ag = create_all_gather_groups( - for_expert_parallelism=(args.expert_model_parallel_size > 1), timeout=timeout + for_expert_parallelism=(args.expert_model_parallel_size > 1), + timeout=timeout, + nccl_communicator_config_path=getattr(args, "nccl_communicator_config_path", None), + high_priority_stream_groups=getattr(args, "high_priority_stream_groups", None), + zero_sm_all_gather=getattr(args, "megatron_fsdp_zero_sm_all_gather", False), ) pg_collection.dp_cp_ag = dp_cp_ag pg_collection.expt_dp_ag = expt_dp_ag print_rank_0("> created all-gather process groups for AG/RS overlap") + if getattr(args, "megatron_fsdp_zero_sm_all_gather", False): + print_rank_0("> zero-SM all-gather enabled for Megatron-FSDP AG groups") if expt_dp_ag is not None: print_rank_0("> including expert parallelism AG group") diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py index 4888c60c4c3..587d08150ff 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_mcore_fully_sharded_data_parallel.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import copy import gc import random @@ -72,6 +72,66 @@ def forward(self, x): return x +def _make_zero_sm_all_gather_args(**extra_overrides): + import argparse + + from megatron.training.arguments import add_megatron_arguments + + parser = argparse.ArgumentParser(allow_abbrev=False) + add_megatron_arguments(parser) + defaults = { + "num_layers": 2, + "hidden_size": 128, + "num_attention_heads": 4, + "seq_length": 256, + "max_position_embeddings": 256, + "micro_batch_size": 1, + "global_batch_size": 8, + "train_iters": 1, + "lr": 1e-4, + "mock_data": True, + "tokenizer_type": "NullTokenizer", + "vocab_size": 256, + "bf16": True, + "use_megatron_fsdp": True, + "ckpt_format": "fsdp_dtensor", + "nccl_ub": True, + "megatron_fsdp_zero_sm_all_gather": True, + "data_parallel_sharding_strategy": "optim_grads_params", + "check_for_nan_in_loss_and_grad": False, + "eval_iters": 0, + "eval_interval": 1, + } + defaults.update(extra_overrides) + parser.set_defaults(**defaults) + args = parser.parse_args([]) + args._is_global_batch_size_explicitly_specified = args.global_batch_size is not None + args.rank = 0 + args.world_size = 1 + return args + + +def test_zero_sm_all_gather_validation_enables_ag_group(monkeypatch): + from megatron.training.arguments import validate_args + + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + args = _make_zero_sm_all_gather_args(create_all_gather_group=False) + + validate_args(args) + + assert args.create_all_gather_group + + +def test_zero_sm_all_gather_validation_requires_nccl_ub(monkeypatch): + from megatron.training.arguments import validate_args + + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + args = _make_zero_sm_all_gather_args(nccl_ub=False) + + with pytest.raises(AssertionError, match="requires --use-nccl-ub"): + validate_args(args) + + def setup_seed(seed): random.seed(seed) # Set Python's built-in random seed np.random.seed(seed) # Set NumPy's random seed diff --git a/tests/unit_tests/distributed/test_param_and_grad_buffer.py b/tests/unit_tests/distributed/test_param_and_grad_buffer.py index ac1fdfe2ed6..e56515ed2d5 100644 --- a/tests/unit_tests/distributed/test_param_and_grad_buffer.py +++ b/tests/unit_tests/distributed/test_param_and_grad_buffer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import contextlib import math @@ -10,12 +10,93 @@ from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig +from megatron.core.distributed.fsdp.src.megatron_fsdp.distributed_data_parallel_config import ( + DistributedDataParallelConfig as StandaloneFSDPDistributedDataParallelConfig, +) from megatron.core.distributed.param_and_grad_buffer import _ParamAndGradBuffer, partition_buckets from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer from megatron.core.transformer import TransformerConfig from tests.unit_tests.test_utilities import TestModel, Utils +def test_megatron_fsdp_zero_sm_all_gather_config_valid(monkeypatch): + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + ddp_config = DistributedDataParallelConfig( + use_megatron_fsdp=True, nccl_ub=True, megatron_fsdp_zero_sm_all_gather=True + ) + + assert ddp_config.megatron_fsdp_zero_sm_all_gather + + +def test_standalone_fsdp_zero_sm_all_gather_config_valid(monkeypatch): + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + ddp_config = StandaloneFSDPDistributedDataParallelConfig( + nccl_ub=True, megatron_fsdp_zero_sm_all_gather=True + ) + + assert ddp_config.megatron_fsdp_zero_sm_all_gather + + +@pytest.mark.parametrize( + "config_cls, kwargs", + [ + (DistributedDataParallelConfig, {"use_megatron_fsdp": True}), + (StandaloneFSDPDistributedDataParallelConfig, {}), + ], +) +def test_megatron_fsdp_zero_sm_all_gather_preserves_allocator_fallback( + config_cls, kwargs, monkeypatch +): + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + ddp_config = config_cls( + nccl_ub=True, + megatron_fsdp_zero_sm_all_gather=True, + fsdp_db_use_persist_buf_on_alloc_fail=False, + **kwargs, + ) + + assert not ddp_config.fsdp_db_use_persist_buf_on_alloc_fail + + +def test_megatron_fsdp_zero_sm_all_gather_config_requires_megatron_fsdp(monkeypatch): + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + with pytest.raises(ValueError, match="only supported with Megatron-FSDP"): + DistributedDataParallelConfig(nccl_ub=True, megatron_fsdp_zero_sm_all_gather=True) + + +@pytest.mark.parametrize( + "config_cls, kwargs", + [ + (DistributedDataParallelConfig, {"use_megatron_fsdp": True}), + (StandaloneFSDPDistributedDataParallelConfig, {}), + ], +) +def test_megatron_fsdp_zero_sm_all_gather_config_requires_nccl_ub(config_cls, kwargs, monkeypatch): + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + with pytest.raises(ValueError, match="requires NCCL user-buffer registration"): + config_cls(megatron_fsdp_zero_sm_all_gather=True, **kwargs) + + +@pytest.mark.parametrize( + "config_cls, kwargs", + [ + (DistributedDataParallelConfig, {"use_megatron_fsdp": True}), + (StandaloneFSDPDistributedDataParallelConfig, {}), + ], +) +def test_megatron_fsdp_zero_sm_all_gather_config_requires_symmetric_registration( + config_cls, kwargs, monkeypatch +): + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + with pytest.raises(ValueError, match="requires symmetric NCCL registration"): + config_cls( + nccl_ub=True, + disable_symmetric_registration=True, + megatron_fsdp_zero_sm_all_gather=True, + **kwargs, + ) + + class TestModelWithExperts(torch.nn.Module): """Model with both dense and expert-parallel parameters. diff --git a/tests/unit_tests/test_parallel_state.py b/tests/unit_tests/test_parallel_state.py index 64c3a5b36a6..698f377e029 100644 --- a/tests/unit_tests/test_parallel_state.py +++ b/tests/unit_tests/test_parallel_state.py @@ -1,6 +1,7 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from math import log2 +from types import SimpleNamespace import pytest import torch @@ -14,6 +15,53 @@ test_parallel_order = ['tp-cp-ep-dp-pp', 'tp-cp-pp-ep-dp'] +class _FakeNCCLOptions: + def __init__(self, is_high_priority_stream=False): + self.is_high_priority_stream = is_high_priority_stream + self.config = SimpleNamespace() + + +@pytest.fixture +def fake_process_group_nccl(monkeypatch): + monkeypatch.setattr( + torch.distributed, + "ProcessGroupNCCL", + SimpleNamespace( + Options=_FakeNCCLOptions, + NCCL_CTA_POLICY_DEFAULT=0, + NCCL_CTA_POLICY_EFFICIENCY=1, + NCCL_CTA_POLICY_ZERO=2, + ), + raising=False, + ) + + +def test_get_nccl_options_sets_cta_policy(fake_process_group_nccl): + options = ps.get_nccl_options( + "dp_cp_ag", + {"dp_cp_ag": {"is_high_priority_stream": True, "min_ctas": 2, "cta_policy": "zero"}}, + ) + + assert options.is_high_priority_stream + assert options.config.min_ctas == 2 + assert options.config.cta_policy == 2 + + +def test_get_nccl_options_with_fallback_merges_base_options(fake_process_group_nccl): + options = ps._get_nccl_options_with_fallback( + "dp_cp_ag", + "dp_cp", + { + "dp_cp": {"is_high_priority_stream": True, "min_ctas": 1}, + "dp_cp_ag": {"min_ctas": 2, "cta_policy": "zero"}, + }, + ) + + assert options.is_high_priority_stream + assert options.config.min_ctas == 2 + assert options.config.cta_policy == 2 + + @pytest.mark.parametrize('order', test_parallel_order) @pytest.mark.flaky @pytest.mark.flaky_in_dev From 77ad57ca0a7fa345d97db7df68c8670c8ca9ac9f Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Sun, 28 Jun 2026 12:05:46 +0000 Subject: [PATCH 2/2] fix: tolerate old optional NVRX package Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- .../core/dist_checkpointing/strategies/nvrx.py | 10 ++++------ .../dist_checkpointing/strategies/torch.py | 15 +++++++++++++-- .../dist_checkpointing/test_async_save.py | 18 +++++++++++++----- 3 files changed, 30 insertions(+), 13 deletions(-) diff --git a/megatron/core/dist_checkpointing/strategies/nvrx.py b/megatron/core/dist_checkpointing/strategies/nvrx.py index 1d609e1186f..afe947cc608 100644 --- a/megatron/core/dist_checkpointing/strategies/nvrx.py +++ b/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 diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py index 31782acb851..af382908a4f 100644 --- a/megatron/core/dist_checkpointing/strategies/torch.py +++ b/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 diff --git a/tests/unit_tests/dist_checkpointing/test_async_save.py b/tests/unit_tests/dist_checkpointing/test_async_save.py index 3c5f37c0133..f66a1d4f4ff 100644 --- a/tests/unit_tests/dist_checkpointing/test_async_save.py +++ b/tests/unit_tests/dist_checkpointing/test_async_save.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import sys from unittest import mock @@ -129,7 +129,7 @@ def test_get_async_strategy_missing_nvrx_cached_metadata_reader(self): class TestHasNvrxAsyncSupport: - """Tests for has_nvrx_async_support, focusing on the minimum-version assertion.""" + """Tests for has_nvrx_async_support, focusing on the minimum-version gate.""" def _fake_modules(self): """MagicMock modules that satisfy every symbol and hasattr check in has_nvrx_async_support.""" @@ -150,7 +150,7 @@ def test_version_check_passes(self): assert has_nvrx_async_support() is True def test_version_check_fails(self): - """Raises AssertionError when all NVRx symbols are present but version is too old.""" + """Returns False when all NVRx symbols are present but version is too old.""" with ( mock.patch( 'megatron.core.dist_checkpointing.strategies.nvrx.import_module', @@ -161,5 +161,13 @@ def test_version_check_fails(self): return_value=False, ), ): - with pytest.raises(AssertionError, match="Minimum required nvidia-resiliency-ext"): - has_nvrx_async_support() + assert has_nvrx_async_support() is False + + def test_get_async_strategy_rejects_old_nvrx(self): + """Explicit NVRx async still fails clearly when the installed version is too old.""" + with mock.patch( + 'megatron.core.dist_checkpointing.strategies.torch.is_nvrx_min_version', + return_value=False, + ): + with pytest.raises(ModuleNotFoundError, match="minimum version"): + get_async_strategy("nvrx")