cp: fix(training): reuse Megatron-Core determinism setup under --deterministic (6271) into r0.7.0

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,

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()

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

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.

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}

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.

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",
+    }

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",
     }
 
 

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)