From 393807643e7a8318b084cf80bbc13987d0bcb878 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Tue, 13 Oct 2026 10:04:38 -0700 Subject: [PATCH] fix(training): reuse Megatron-Core determinism setup under --deterministic (#6271) Signed-off-by: Zhiyu Li Co-authored-by: Claude Opus 5.5 Signed-off-by: NeMo Bot --- .../bridge/perf_recipes/deepseek/__init__.py | 1 + .../deepseek/gb300/deepseek_v3.py | 48 ++++++++ .../bridge/recipes/utils/determinism_utils.py | 15 +-- src/megatron/bridge/training/config.py | 21 ++-- .../recipes/test_gb300_ncclep_perf_recipes.py | 21 ++++ .../recipes/test_perf_recipe_environment.py | 5 +- .../recipes/utils/test_determinism_utils.py | 41 +++++++ .../performance/test_recipe_environment.py | 6 +- tests/unit_tests/training/test_config.py | 103 +++++++++++++----- 9 files changed, 214 insertions(+), 47 deletions(-) create mode 100644 tests/unit_tests/recipes/utils/test_determinism_utils.py diff --git a/src/megatron/bridge/perf_recipes/deepseek/__init__.py b/src/megatron/bridge/perf_recipes/deepseek/__init__.py index 5d9c852e22..a784910aff 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/__init__.py +++ b/src/megatron/bridge/perf_recipes/deepseek/__init__.py @@ -35,6 +35,7 @@ deepseek_v3_pretrain_256gpu_gb300_bf16_config, deepseek_v3_pretrain_256gpu_gb300_fp8cs_config, deepseek_v3_pretrain_256gpu_gb300_fp8mx_config, + deepseek_v3_pretrain_256gpu_gb300_fp8mx_deterministic_config, deepseek_v3_pretrain_256gpu_gb300_fp8mx_hsdp_config, deepseek_v3_pretrain_256gpu_gb300_fp8mx_large_scale_config, deepseek_v3_pretrain_256gpu_gb300_fp8mx_partial_cg_dev_config, diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v3.py index a5bf28ec0b..d6bf20b7c0 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v3.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v3.py @@ -26,6 +26,7 @@ set_deepseek_v3_pipeline_model_parallel_layout, ) from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS +from megatron.bridge.recipes.utils.determinism_utils import apply_determinism_overrides def _build_deepseek_v3_gb300_bf16() -> ConfigContainer: @@ -197,6 +198,53 @@ def deepseek_v3_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer: return cfg +def deepseek_v3_pretrain_256gpu_gb300_fp8mx_deterministic_config() -> ConfigContainer: + """DeepSeek V3 pretrain: 256× GB300, MXFP8, deterministic. + + Wraps :func:`deepseek_v3_pretrain_256gpu_gb300_fp8mx_config` and applies + :func:`~megatron.bridge.recipes.utils.determinism_utils.apply_determinism_overrides`. + """ + cfg = deepseek_v3_pretrain_256gpu_gb300_fp8mx_config() + # TODO: revert back to NCCL EP once the NCCL EP issue is fixed. + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_use_grouped_tensor = False + apply_determinism_overrides(cfg) + # Keep process settings next to the recipe so users can see the exact benchmark environment. + cfg.env_vars = { + **COMMON_PERF_ENV_VARS, + # ---- Determinism (same values as apply_determinism_overrides) ---- + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "NCCL_ALGO": "Ring", + "NCCL_EP_HT_EM_AG_SCAN_MODE": "1", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", + "MAMBA_DETERMINISTIC": "1", + "CAUSAL_CONV1D_DETERMINISTIC": "1", + # ---- Others ---- + # CUDA stream scheduling for this model and parallel layout. + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + # CUDA graph and allocator behavior for this recipe. + "NCCL_GRAPH_REGISTER": 0, + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True", + "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, + # NCCL user-buffer and launch settings. + "NCCL_NVLS_ENABLE": 0, + # HybridEP topology for the target system. + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + # Transformer Engine overlap settings for this model. + "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Use cuDNN LayerNorm for this measured baseline. + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, + } + return cfg + + def _build_deepseek_v3_gb300_nvfp4() -> ConfigContainer: """Shared HybridEP NVFP4 base for the DeepSeek V3 256-GPU GB300 recipe and its VRNVL72 alias.""" cfg = deepseek_v3_pretrain_config() diff --git a/src/megatron/bridge/recipes/utils/determinism_utils.py b/src/megatron/bridge/recipes/utils/determinism_utils.py index 9782d4c06f..32df65bf2b 100644 --- a/src/megatron/bridge/recipes/utils/determinism_utils.py +++ b/src/megatron/bridge/recipes/utils/determinism_utils.py @@ -14,6 +14,8 @@ """Config-level overrides for deterministic training.""" +from megatron.training.determinism import DETERMINISM_ENV_VAR_DEFAULTS + from megatron.bridge.training.config import ConfigContainer @@ -35,13 +37,12 @@ def apply_determinism_overrides(cfg: ConfigContainer) -> None: """ cfg.model.deterministic_mode = True cfg.model.cross_entropy_loss_fusion = False - cfg.env_vars.update( - { - "CUBLAS_WORKSPACE_CONFIG": ":4096:8", - "NCCL_ALGO": "Ring", - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, - } - ) + # Only the MoE aux-loss fusion is non-deterministic; fused TopK routing stays on. + cfg.model.moe_router_aux_loss_fusion = False + # Exported at launch so the values are in place before TE, cuBLAS and NCCL first read them. + cfg.env_vars.update(DETERMINISM_ENV_VAR_DEFAULTS) + # Pin the Mamba/causal-conv1d kernels explicitly instead of letting them follow torch's flag. + cfg.env_vars.update({"MAMBA_DETERMINISTIC": "1", "CAUSAL_CONV1D_DETERMINISTIC": "1"}) if cfg.comm_overlap is not None: cfg.comm_overlap.tp_comm_overlap = False diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index 8cf107fed6..193a67a3c2 100644 --- a/src/megatron/bridge/training/config.py +++ b/src/megatron/bridge/training/config.py @@ -18,6 +18,7 @@ import warnings from dataclasses import MISSING, dataclass, field, fields from pathlib import Path +from types import SimpleNamespace from typing import Any, Dict, Literal, Optional, Tuple, Union import torch @@ -41,6 +42,7 @@ from megatron.training.config import SchedulerConfig as MTrainSchedulerConfig from megatron.training.config import StragglerDetectionConfig as MTrainStragglerDetectionConfig from megatron.training.config import TrainingConfig as MTrainTrainingConfig +from megatron.training.determinism import apply_determinism_to_args from megatron.bridge.data.base import ( DataloaderConfig, @@ -1072,19 +1074,16 @@ def _validate_and_apply_deterministic_mode(self) -> None: if not getattr(self.model, "deterministic_mode", False): return - # Disallow cross-entropy loss fusion as it is not deterministic - assert not getattr(self.model, "cross_entropy_loss_fusion", False), ( - "Cross Entropy Fusion is currently not deterministic." - ) - - all_reduce_choices = ("Tree", "Ring", "CollnetDirect", "CollnetChain", "^NVLS") - assert os.getenv("NCCL_ALGO", -1) != -1 and os.getenv("NCCL_ALGO") in all_reduce_choices, ( - f"NCCL_ALGO must be one of {all_reduce_choices}." + # Megatron-Core's --deterministic-mode setup, so both frameworks enforce the same rules. + apply_determinism_to_args( + SimpleNamespace( + cross_entropy_loss_fusion=getattr(self.model, "cross_entropy_loss_fusion", False), + tp_comm_overlap=self.comm_overlap.tp_comm_overlap if self.comm_overlap is not None else False, + moe_router_fusion=getattr(self.model, "moe_router_fusion", False), + moe_router_aux_loss_fusion=getattr(self.model, "moe_router_aux_loss_fusion", None), + ) ) - # Enable deterministic algorithms in torch - torch.use_deterministic_algorithms(True) - def _validate_and_apply_megatron_fsdp_configs(self) -> None: """Validate and apply configuration required by the selected Megatron-FSDP version.""" # Set configs needed for Megatron-FSDP. diff --git a/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py index 64054808df..9fb2d0d503 100644 --- a/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py +++ b/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py @@ -18,6 +18,7 @@ import importlib import pytest +from megatron.training.determinism import DETERMINISM_ENV_VAR_DEFAULTS from megatron.bridge.perf_recipes.environment import HYBRID_EP_ENV_NAMES from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies @@ -26,6 +27,7 @@ pytestmark = pytest.mark.unit _NCCL_EP_ENV = {"NCCL_EP_HT_EM_PULL_PUSH": 1} +_SSM_DETERMINISM_ENV = {"MAMBA_DETERMINISTIC": "1", "CAUSAL_CONV1D_DETERMINISTIC": "1"} # (module under megatron.bridge.perf_recipes, recipe factory) for the GB300 recipes defaulted to NCCL EP. _GB300_NCCLEP_RECIPES = ( @@ -136,3 +138,22 @@ def test_vrnvl72_aliases_do_not_inherit_nccl_ep(module_name: str, factory_name: assert cfg.env_vars.keys() >= HYBRID_EP_ENV_NAMES else: assert cfg.model.moe_token_dispatcher_type == "alltoall" + + +def test_deepseek_v3_gb300_fp8mx_deterministic_is_base_plus_determinism_overrides() -> None: + module = "deepseek.gb300.deepseek_v3" + base = _build(module, "deepseek_v3_pretrain_256gpu_gb300_fp8mx_config") + cfg = _build(module, "deepseek_v3_pretrain_256gpu_gb300_fp8mx_deterministic_config") + + assert cfg.model.deterministic_mode is True + assert cfg.model.cross_entropy_loss_fusion is False + assert cfg.model.moe_router_aux_loss_fusion is False + assert cfg.model.moe_router_fusion == base.model.moe_router_fusion + # HybridEP with the legacy grouped MLP until the NCCL EP path is deterministic. + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_use_grouped_tensor is False + assert cfg.env_vars.keys() >= HYBRID_EP_ENV_NAMES + assert cfg.env_vars.keys().isdisjoint(_NCCL_EP_ENV) + expected_env = {name: value for name, value in base.env_vars.items() if name not in _NCCL_EP_ENV} + expected_env.update({name: cfg.env_vars[name] for name in HYBRID_EP_ENV_NAMES}) + assert cfg.env_vars == {**expected_env, **DETERMINISM_ENV_VAR_DEFAULTS, **_SSM_DETERMINISM_ENV} diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index ccf490fe5e..a9f75dbd17 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -217,7 +217,10 @@ def test_explicit_environment_invariants_across_all_flat_recipes(): assert environment.keys().isdisjoint(_DEEPSEEK_NON_BASELINE_ENV_NAMES) assert environment["NVTE_FWD_LAYERNORM_SM_MARGIN"] == 20 assert environment["NVTE_BWD_LAYERNORM_SM_MARGIN"] == 20 - if path.parent.name == "gb300" and path.stem == "deepseek_v3": + if function_name.endswith("_deterministic_config"): + # Bit-exact recipes pin deterministic TE kernels. + assert environment["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == "0" + elif path.parent.name == "gb300" and path.stem == "deepseek_v3": # The DeepSeek V3 GB300 recipes run the TE op-fuser grouped MLP, whose quantized scale # gradient is accumulated with atomics; TE raises when bit-exact execution is requested, # so these recipes must not pin NVTE_ALLOW_NONDETERMINISTIC_ALGO=0. diff --git a/tests/unit_tests/recipes/utils/test_determinism_utils.py b/tests/unit_tests/recipes/utils/test_determinism_utils.py new file mode 100644 index 0000000000..effa0e0c39 --- /dev/null +++ b/tests/unit_tests/recipes/utils/test_determinism_utils.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +from megatron.training.determinism import DETERMINISM_ENV_VAR_DEFAULTS + +from megatron.bridge.recipes.utils.determinism_utils import apply_determinism_overrides + + +def test_overrides_follow_megatron_core(): + """Only the aux-loss fusion is disabled; env is Megatron-Core's defaults plus the SSM pins.""" + cfg = SimpleNamespace( + model=SimpleNamespace(moe_router_fusion=True, moe_router_aux_loss_fusion=None), + env_vars={}, + comm_overlap=SimpleNamespace(tp_comm_overlap=True), + ) + + apply_determinism_overrides(cfg) + + assert cfg.model.deterministic_mode is True + assert cfg.model.cross_entropy_loss_fusion is False + assert cfg.model.moe_router_aux_loss_fusion is False + assert cfg.model.moe_router_fusion is True + assert cfg.comm_overlap.tp_comm_overlap is False + assert cfg.env_vars == { + **DETERMINISM_ENV_VAR_DEFAULTS, + "MAMBA_DETERMINISTIC": "1", + "CAUSAL_CONV1D_DETERMINISTIC": "1", + } diff --git a/tests/unit_tests/scripts/performance/test_recipe_environment.py b/tests/unit_tests/scripts/performance/test_recipe_environment.py index 87b0610f03..c84d20b3ad 100644 --- a/tests/unit_tests/scripts/performance/test_recipe_environment.py +++ b/tests/unit_tests/scripts/performance/test_recipe_environment.py @@ -738,7 +738,7 @@ def apply_argparse(config, args): assert effective_recipe.env_vars["USER_SELECTED"] == "custom" assert "REMOVE_ME" not in effective_recipe.env_vars assert effective_recipe.env_vars["CUBLAS_WORKSPACE_CONFIG"] == ":4096:8" - assert effective_recipe.env_vars["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == 0 + assert effective_recipe.env_vars["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == "0" def test_flat_environment_compatibility_preserves_explicit_hydra_values(): @@ -1222,7 +1222,9 @@ def test_runner_applies_determinism_before_environment_export(): "RECIPE_ENV": 1, "CUBLAS_WORKSPACE_CONFIG": ":4096:8", "NCCL_ALGO": "Ring", - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", + "MAMBA_DETERMINISTIC": "1", + "CAUSAL_CONV1D_DETERMINISTIC": "1", } diff --git a/tests/unit_tests/training/test_config.py b/tests/unit_tests/training/test_config.py index 3d65673f83..6c8d20e124 100644 --- a/tests/unit_tests/training/test_config.py +++ b/tests/unit_tests/training/test_config.py @@ -507,27 +507,80 @@ def test_cannot_set_blend_fields(self): ).finalize() +_DETERMINISM_ENV_VARS = ( + "NCCL_ALGO", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO", + "CUBLAS_WORKSPACE_CONFIG", + "MAMBA_DETERMINISTIC", + "CAUSAL_CONV1D_DETERMINISTIC", + "TRITON_CACHE_AUTOTUNING", +) + + +@pytest.fixture +def isolated_determinism_state(monkeypatch): + """Keep Megatron-Core's determinism setup from leaking env vars or torch global state.""" + import torch.utils.deterministic + + for name in _DETERMINISM_ENV_VARS: + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr( + torch.utils.deterministic, "fill_uninitialized_memory", torch.utils.deterministic.fill_uninitialized_memory + ) + calls = [] + monkeypatch.setattr(torch, "use_deterministic_algorithms", lambda flag, **_kwargs: calls.append(flag)) + return calls + + class TestConfigContainerValidation: - def test_deterministic_mode_disallows_ce_fusion(self, monkeypatch): + def test_deterministic_mode_disallows_ce_fusion(self, isolated_determinism_state): """Test that deterministic mode disallows cross-entropy loss fusion.""" gpt_model_cfg = create_test_gpt_config( deterministic_mode=True, cross_entropy_loss_fusion=True, ) - # Ensure NCCL_ALGO present but valid, so we fail on CE fusion - monkeypatch.setenv("NCCL_ALGO", "Tree") - container, og_ws, cfg_mod = create_test_config_container(world_size_override=1, model_config=gpt_model_cfg) try: - with pytest.raises(AssertionError, match="Cross Entropy Fusion is currently not deterministic"): + with pytest.raises(AssertionError, match="cross_entropy_loss_fusion=False"): container.validate() finally: restore_get_world_size_safe(og_ws, cfg_mod) - def test_deterministic_mode_requires_nccl_algo_and_sets_torch(self, monkeypatch): - """Test that deterministic mode requires NCCL_ALGO and sets torch.use_deterministic_algorithms.""" + def test_deterministic_mode_rejects_moe_aux_loss_fusion(self, isolated_determinism_state): + """Deterministic mode rejects the MoE aux-loss fusion but keeps fused TopK routing.""" + gpt_model_cfg = create_test_gpt_config(deterministic_mode=True, cross_entropy_loss_fusion=False) + gpt_model_cfg.moe_router_fusion = True + gpt_model_cfg.moe_router_aux_loss_fusion = True + + container, og_ws, cfg_mod = create_test_config_container(world_size_override=1, model_config=gpt_model_cfg) + + try: + with pytest.raises(AssertionError, match="moe_router_aux_loss_fusion=False"): + container._validate_and_apply_deterministic_mode() + + gpt_model_cfg.moe_router_aux_loss_fusion = False + container._validate_and_apply_deterministic_mode() + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + def test_deterministic_mode_rejects_tp_comm_overlap(self, isolated_determinism_state): + """TP communication overlap is not bit-exact and must be off.""" + gpt_model_cfg = create_test_gpt_config(deterministic_mode=True, cross_entropy_loss_fusion=False) + container, og_ws, cfg_mod = create_test_config_container(world_size_override=1, model_config=gpt_model_cfg) + container.comm_overlap = SimpleNamespace(tp_comm_overlap=True) + + try: + with pytest.raises(AssertionError, match="tp_comm_overlap=False"): + container._validate_and_apply_deterministic_mode() + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + def test_deterministic_mode_env_matches_megatron_core(self, monkeypatch, isolated_determinism_state): + """Env vars follow Megatron-Core's rules: defaults filled, non-deterministic values rejected.""" + import torch.utils.deterministic + gpt_model_cfg = create_test_gpt_config( deterministic_mode=True, cross_entropy_loss_fusion=False, @@ -537,27 +590,25 @@ def test_deterministic_mode_requires_nccl_algo_and_sets_torch(self, monkeypatch) container, og_ws, cfg_mod = create_test_config_container(world_size_override=1, model_config=gpt_model_cfg) try: - # Missing NCCL_ALGO - monkeypatch.delenv("NCCL_ALGO", raising=False) - with pytest.raises(AssertionError, match="NCCL_ALGO must be one of"): - container.validate() - - # Invalid NCCL_ALGO - monkeypatch.setenv("NCCL_ALGO", "AllReduce") - with pytest.raises(AssertionError, match="NCCL_ALGO must be one of"): - container.validate() + # Tree has no controllable reduction order and is rejected, as in Megatron-Core. + monkeypatch.setenv("NCCL_ALGO", "Tree") + with pytest.raises(AssertionError, match="NCCL_ALGO='Tree'"): + container._validate_and_apply_deterministic_mode() - # Valid NCCL_ALGO -> should pass and call torch deterministic monkeypatch.setenv("NCCL_ALGO", "Ring") - - called = {"det": False} - - def _mock_use_deterministic(flag): - called["det"] = flag - - with patch.object(torch, "use_deterministic_algorithms", side_effect=_mock_use_deterministic): - container.validate() - assert called["det"] is True + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1") + with pytest.raises(AssertionError, match="NVTE_ALLOW_NONDETERMINISTIC_ALGO='1'"): + container._validate_and_apply_deterministic_mode() + + # Unset values get the canonical defaults; torch deterministic mode is enabled. + monkeypatch.delenv("NCCL_ALGO") + monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO") + container._validate_and_apply_deterministic_mode() + assert os.environ["NCCL_ALGO"] == "Ring" + assert os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == "0" + assert os.environ["CUBLAS_WORKSPACE_CONFIG"] == ":4096:8" + assert isolated_determinism_state == [True] + assert torch.utils.deterministic.fill_uninitialized_memory is False finally: restore_get_world_size_safe(og_ws, cfg_mod)