diff --git a/docs/recipe-usage.md b/docs/recipe-usage.md index e01df5343e5..9989544e962 100644 --- a/docs/recipe-usage.md +++ b/docs/recipe-usage.md @@ -326,6 +326,8 @@ Flat performance and hardware-specific model recipes keep their baseline environ Every flat performance recipe builder assigns `cfg.env_vars` directly before returning its config. `COMMON_PERF_ENV_VARS` contains only `TORCH_NCCL_HIGH_PRIORITY=1`, the one value shared by every workload; all GPU-, topology-, precision-, CUDA graph-, and model-specific values are written inline with comments in the owning recipe. This is intentionally verbose so a recipe user can see the exact benchmark environment without following a decorator or global derivation rules. +Deterministic recipes additionally spread `DETERMINISM_ENV_VARS` (from `megatron.bridge.recipes.utils.environment_utils`), a third shared constant, next to the two common-default ones, holding the process settings required for bit-exact reproducibility, such as `CUBLAS_WORKSPACE_CONFIG` and `NCCL_ALGO`. It is the only other constant a recipe may spread, and it is spread after `apply_determinism_overrides` because the recipe reassigns `cfg.env_vars`. + Hardware recipe builders follow the same explicit pattern. `COMMON_RECIPE_ENV_VARS` contains only the seven settings that are identical across every supported recipe; each of the 245 supported builders assigns `cfg.env_vars` directly before returning and writes any HybridEP or model-specific additions inline. There is no recipe decorator or generic launcher-side environment recomposition. Because the compatibility catalog currently reuses H100 builders on other targets, the recipe entry point adapts HybridEP topology after overrides and preserves legacy coupling for an explicit argparse `--nccl_ub`; explicit `env_vars` remain final. The nemo-ci executors under `scripts/performance` do not provide fallback values for recipe-owned process settings. Cluster fabric and orchestration settings do not belong in recipes. Keep credentials, cache/offline policy, scheduler metadata, W&B state, and cluster-specific NCCL network tuning in the launcher or executor. In particular, settings such as `NCCL_NET_GDR_LEVEL`, `NCCL_NET_GDR_C2C`, and `NCCL_IB_QPS_PER_CONNECTION` must remain deployment-owned. 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 d6bf20b7c0e..fca7653e4ea 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v3.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v3.py @@ -27,6 +27,7 @@ ) from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.recipes.utils.determinism_utils import apply_determinism_overrides +from megatron.bridge.recipes.utils.environment_utils import DETERMINISM_ENV_VARS def _build_deepseek_v3_gb300_bf16() -> ConfigContainer: @@ -212,13 +213,7 @@ def deepseek_v3_pretrain_256gpu_gb300_fp8mx_deterministic_config() -> ConfigCont # 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", + **DETERMINISM_ENV_VARS, # ---- Others ---- # CUDA stream scheduling for this model and parallel layout. "CUDA_DEVICE_MAX_CONNECTIONS": 32, diff --git a/src/megatron/bridge/recipes/moonlight/h100/__init__.py b/src/megatron/bridge/recipes/moonlight/h100/__init__.py index adae78bda55..4dc91d8eee3 100644 --- a/src/megatron/bridge/recipes/moonlight/h100/__init__.py +++ b/src/megatron/bridge/recipes/moonlight/h100/__init__.py @@ -19,6 +19,7 @@ "moonlight_16b_peft_2gpu_h100_bf16_config", "moonlight_16b_peft_4gpu_h100_bf16_config", "moonlight_16b_pretrain_16gpu_h100_bf16_config", + "moonlight_16b_pretrain_16gpu_h100_bf16_deterministic_config", "moonlight_16b_pretrain_8gpu_h100_bf16_config", "moonlight_16b_sft_8gpu_h100_bf16_8k_config", "moonlight_16b_sft_8gpu_h100_bf16_config", diff --git a/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py b/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py index ec0279aef4a..9dbf15d2b11 100644 --- a/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py +++ b/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py @@ -20,7 +20,8 @@ from megatron.bridge.peft.base import PEFT from megatron.bridge.recipes.common import _peft_common, _pretrain_common, _sft_common from megatron.bridge.recipes.utils.dataset_utils import default_peft_config -from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.recipes.utils.determinism_utils import apply_determinism_overrides +from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS, DETERMINISM_ENV_VARS from megatron.bridge.training.comm_overlap import CommOverlapConfig from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.mixed_precision import MixedPrecisionConfig @@ -391,6 +392,29 @@ def moonlight_16b_pretrain_16gpu_h100_bf16_config() -> ConfigContainer: return cfg +def moonlight_16b_pretrain_16gpu_h100_bf16_deterministic_config() -> ConfigContainer: + """Return a deterministic bounded-convergence pre-training config for Moonlight-16B. + + Wraps :func:`moonlight_16b_pretrain_16gpu_h100_bf16_config` and applies + :func:`~megatron.bridge.recipes.utils.determinism_utils.apply_determinism_overrides`. + """ + cfg = moonlight_16b_pretrain_16gpu_h100_bf16_config() + apply_determinism_overrides(cfg) + # Keep the complete process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + **DETERMINISM_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 8, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "USE_MNNVL": 0, + } + return cfg + + # ============================================================================= # SFT Config # ============================================================================= @@ -912,6 +936,7 @@ def moonlight_16b_peft_4gpu_h100_bf16_config( "moonlight_16b_peft_2gpu_h100_bf16_config", "moonlight_16b_peft_4gpu_h100_bf16_config", "moonlight_16b_pretrain_16gpu_h100_bf16_config", + "moonlight_16b_pretrain_16gpu_h100_bf16_deterministic_config", "moonlight_16b_pretrain_8gpu_h100_bf16_config", "moonlight_16b_sft_8gpu_h100_bf16_8k_config", "moonlight_16b_sft_8gpu_h100_bf16_config", diff --git a/src/megatron/bridge/recipes/nemotronh/__init__.py b/src/megatron/bridge/recipes/nemotronh/__init__.py index 19ad5a10a6a..f9ca938a063 100644 --- a/src/megatron/bridge/recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/__init__.py @@ -19,8 +19,10 @@ nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config, nemotron_3_nano_gb200_pretrain_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, + nemotron_3_nano_pretrain_8gpu_gb200_bf16_deterministic_config, nemotron_3_super_pretrain_64gpu_gb200_bf16_config, nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config, + nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_deterministic_config, ) from megatron.bridge.recipes.nemotronh.nemotron_3_nano import ( nemotron_3_5_lightning_peft_config, @@ -65,6 +67,7 @@ "nemotron_3_nano_peft_config", "nemotron_3_nano_gb200_pretrain_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", + "nemotron_3_nano_pretrain_8gpu_gb200_bf16_deterministic_config", # Nemotron 3 Nano 4B model "nemotron_3_nano_4b_pretrain_config", "nemotron_3_nano_4b_sft_config", @@ -77,6 +80,7 @@ "nemotron_3_super_peft_config", # Nemotron 3 Ultra models "nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config", + "nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_deterministic_config", "nemotron_3_ultra_pretrain_config", "nemotron_3_ultra_sft_openmathinstruct2_packed_config", "nemotron_3_ultra_peft_openmathinstruct2_packed_config", diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py index dfe61779988..edbcfc2fcfd 100644 --- a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py @@ -18,12 +18,14 @@ nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config, nemotron_3_nano_gb200_pretrain_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, + nemotron_3_nano_pretrain_8gpu_gb200_bf16_deterministic_config, ) from megatron.bridge.recipes.nemotronh.gb200.nemotron_3_super import ( nemotron_3_super_pretrain_64gpu_gb200_bf16_config, ) from megatron.bridge.recipes.nemotronh.gb200.nemotron_3_ultra import ( nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config, + nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_deterministic_config, ) @@ -33,6 +35,8 @@ "nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config", "nemotron_3_nano_gb200_pretrain_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", + "nemotron_3_nano_pretrain_8gpu_gb200_bf16_deterministic_config", "nemotron_3_super_pretrain_64gpu_gb200_bf16_config", "nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config", + "nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_deterministic_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py index 78283e737b3..01504bbd116 100644 --- a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py +++ b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py @@ -22,7 +22,8 @@ nemotron_3_5_lightning_pretrain_config, nemotron_3_5_lightning_sft_openmathinstruct2_packed_config, ) -from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.recipes.utils.determinism_utils import apply_determinism_overrides +from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS, DETERMINISM_ENV_VARS from megatron.bridge.training.comm_overlap import CommOverlapConfig from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.mixed_precision import get_mixed_precision_config @@ -157,6 +158,31 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: return cfg +def nemotron_3_nano_pretrain_8gpu_gb200_bf16_deterministic_config() -> ConfigContainer: + """Return a deterministic Nemotron 3 Nano BF16 pretraining config for eight GB200 GPUs. + + Wraps :func:`nemotron_3_nano_pretrain_8gpu_gb200_bf16_config` and applies + :func:`~megatron.bridge.recipes.utils.determinism_utils.apply_determinism_overrides`. + """ + cfg = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() + apply_determinism_overrides(cfg) + # Keep the exact process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + **DETERMINISM_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, + } + return cfg + + def nemotron_3_5_lightning_pretrain_8k_config() -> ConfigContainer: """Return the Nemotron 3.5 Lightning BF16 8K pretraining config.""" cfg = nemotron_3_5_lightning_pretrain_config() @@ -250,4 +276,5 @@ def nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config() -> ConfigCo "nemotron_3_5_lightning_sft_openmathinstruct2_packed_tp1_config", "nemotron_3_nano_gb200_pretrain_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", + "nemotron_3_nano_pretrain_8gpu_gb200_bf16_deterministic_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_ultra.py b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_ultra.py index 48cc894161f..197651282aa 100644 --- a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_ultra.py +++ b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_ultra.py @@ -20,7 +20,8 @@ NEMOTRON_3_ULTRA_PRETRAIN_SEQ_LENGTH, _nemotron_3_ultra_large_scale_bf16_config, ) -from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.recipes.utils.determinism_utils import apply_determinism_overrides +from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS, DETERMINISM_ENV_VARS from megatron.bridge.training.config import ConfigContainer @@ -100,4 +101,36 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config() -> ConfigContainer return _nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config() -__all__ = ["nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config"] +def nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_deterministic_config() -> ConfigContainer: + """Return a deterministic Nemotron 3 Ultra BF16 pretrain config for 256 GB200 GPUs. + + Wraps :func:`nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config` and applies + :func:`~megatron.bridge.recipes.utils.determinism_utils.apply_determinism_overrides` + and spreads ``DETERMINISM_ENV_VARS`` into the recipe's static process environment. + + Returns: + GB200 BF16 deterministic pretraining configuration. + """ + cfg = nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config() + apply_determinism_overrides(cfg) + # Keep the exact process environment visible on the recipe. + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + **DETERMINISM_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 16, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_CPU_OFFLOAD_V1": 1, + "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + } + return cfg + + +__all__ = [ + "nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config", + "nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_deterministic_config", +] diff --git a/src/megatron/bridge/recipes/qwen/__init__.py b/src/megatron/bridge/recipes/qwen/__init__.py index 9481d406909..07b5c774d4c 100644 --- a/src/megatron/bridge/recipes/qwen/__init__.py +++ b/src/megatron/bridge/recipes/qwen/__init__.py @@ -20,6 +20,7 @@ # Qwen3.5 GB200 models from .gb200.qwen35 import ( qwen35_text_9b_pretrain_8gpu_gb200_bf16_config, + qwen35_text_9b_pretrain_8gpu_gb200_bf16_deterministic_config, qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config, qwen35_text_35b_a3b_sft_8gpu_gb200_bf16_dynamic_cp_config, qwen35_text_35b_a3b_sft_long_context_16gpu_gb200_fp8mx_config, @@ -175,6 +176,7 @@ # Qwen3.5 text models "qwen35_text_9b_pretrain_config", "qwen35_text_9b_pretrain_8gpu_gb200_bf16_config", + "qwen35_text_9b_pretrain_8gpu_gb200_bf16_deterministic_config", "qwen35_text_35b_a3b_pretrain_config", "qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config", "qwen35_text_35b_a3b_sft_8gpu_gb200_bf16_dynamic_cp_config", diff --git a/src/megatron/bridge/recipes/qwen/gb200/__init__.py b/src/megatron/bridge/recipes/qwen/gb200/__init__.py index 5799317f1d7..cb009b0029d 100644 --- a/src/megatron/bridge/recipes/qwen/gb200/__init__.py +++ b/src/megatron/bridge/recipes/qwen/gb200/__init__.py @@ -17,6 +17,7 @@ ) from megatron.bridge.recipes.qwen.gb200.qwen35 import ( qwen35_text_9b_pretrain_8gpu_gb200_bf16_config, + qwen35_text_9b_pretrain_8gpu_gb200_bf16_deterministic_config, qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config, qwen35_text_35b_a3b_sft_8gpu_gb200_bf16_dynamic_cp_config, qwen35_text_35b_a3b_sft_long_context_16gpu_gb200_fp8mx_config, @@ -26,6 +27,7 @@ __all__ = [ "qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config", "qwen35_text_9b_pretrain_8gpu_gb200_bf16_config", + "qwen35_text_9b_pretrain_8gpu_gb200_bf16_deterministic_config", "qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config", "qwen35_text_35b_a3b_sft_8gpu_gb200_bf16_dynamic_cp_config", "qwen35_text_35b_a3b_sft_long_context_16gpu_gb200_fp8mx_config", diff --git a/src/megatron/bridge/recipes/qwen/gb200/qwen35.py b/src/megatron/bridge/recipes/qwen/gb200/qwen35.py index 26a6396785a..997be339a14 100644 --- a/src/megatron/bridge/recipes/qwen/gb200/qwen35.py +++ b/src/megatron/bridge/recipes/qwen/gb200/qwen35.py @@ -24,7 +24,8 @@ from megatron.bridge import AutoBridge from megatron.bridge.recipes.common import _pretrain_common, _sft_common from megatron.bridge.recipes.utils.dataset_utils import default_coderforge_config -from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.recipes.utils.determinism_utils import apply_determinism_overrides +from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS, DETERMINISM_ENV_VARS from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing from megatron.bridge.training.comm_overlap import CommOverlapConfig from megatron.bridge.training.config import ConfigContainer @@ -220,6 +221,26 @@ def qwen35_text_9b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: return cfg +def qwen35_text_9b_pretrain_8gpu_gb200_bf16_deterministic_config() -> ConfigContainer: + """Return a deterministic text-only Qwen3.5-9B pretraining config for eight GB200 GPUs. + + Wraps :func:`qwen35_text_9b_pretrain_8gpu_gb200_bf16_config` and applies + :func:`~megatron.bridge.recipes.utils.determinism_utils.apply_determinism_overrides`. + """ + cfg = qwen35_text_9b_pretrain_8gpu_gb200_bf16_config() + apply_determinism_overrides(cfg) + + # Standalone Qwen3.5 override, not part of apply_determinism_overrides: Megatron-Core rejects + # deterministic_mode with Gated DeltaNet unless the torch kernel backend is used. + cfg.model.gdn_kernel_backend = "torch" + + # Keep the exact process environment visible on the recipe. + cfg.env_vars = { + **DETERMINISM_ENV_VARS, + } + return cfg + + def qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: """Return a text-only Qwen3.5-35B-A3B pretraining config for eight GB200 GPUs.""" cfg = _pretrain_common() diff --git a/src/megatron/bridge/recipes/utils/environment_utils.py b/src/megatron/bridge/recipes/utils/environment_utils.py index d44381d4b0d..164f11371a6 100644 --- a/src/megatron/bridge/recipes/utils/environment_utils.py +++ b/src/megatron/bridge/recipes/utils/environment_utils.py @@ -31,3 +31,15 @@ "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, "TORCH_NCCL_HIGH_PRIORITY": 1, } + +# Values required for bit-exact reproducibility. Identical to what +# apply_determinism_overrides exports; spread into a deterministic recipe's +# static env_vars dict, which is reassigned after that call. +DETERMINISM_ENV_VARS: dict[str, str | int | float | bool] = { + "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", +} diff --git a/src/megatron/bridge/training/initialize.py b/src/megatron/bridge/training/initialize.py index 22b1fe63dbb..bfd8ab4d2a7 100644 --- a/src/megatron/bridge/training/initialize.py +++ b/src/megatron/bridge/training/initialize.py @@ -24,7 +24,7 @@ import torch import torch.distributed import torch.nn.functional as F -from megatron.core import parallel_state, tensor_parallel +from megatron.core import parallel_state, tensor_parallel, tuning from megatron.core.datasets.utils import compile_helpers from megatron.core.fusions.fused_bias_dropout import bias_dropout_add_fused_train from megatron.core.fusions.fused_bias_gelu import bias_gelu @@ -126,6 +126,12 @@ def initialize_megatron( # init rerun global state init_rerun_state(rerun_state_machine_config) + # Pin Triton autotune configs before any kernel runs, as Megatron-Core's training initializer does under + # --deterministic-mode. Timed autotuning can pick different tilings between ranks and runs, which changes the + # floating-point reduction order. Runs without deterministic_mode keep Triton's timed autotuning. + if getattr(model_config, "deterministic_mode", False): + tuning.install(deterministic=True) + # torch.distributed initialization result = torch_dist_init( model_config=model_config, diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py index 570504259fd..6951528ac8d 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py @@ -32,6 +32,10 @@ import torch from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.recipes.nemotronh.gb200.nemotron_3_nano import ( + nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, + nemotron_3_nano_pretrain_8gpu_gb200_bf16_deterministic_config, +) from megatron.bridge.recipes.nemotronh.h100 import nemotron_3_nano as recipe_module from megatron.bridge.recipes.nemotronh.nemotron_3_nano import ( nemotron_3_5_lightning_peft_config, @@ -41,6 +45,7 @@ nemotron_3_nano_pretrain_config, nemotron_3_nano_sft_config, ) +from megatron.bridge.recipes.utils.environment_utils import DETERMINISM_ENV_VARS from megatron.bridge.training.config import ConfigContainer from megatron.bridge.utils.cuda_graph import cuda_graph_module_names @@ -628,3 +633,20 @@ def test_moe_model_configuration(self, recipe_fn): assert config.model.moe_shared_expert_intermediate_size == 3712 assert config.model.moe_router_topk == 6 assert config.model.moe_router_topk_scaling_factor == 2.5 + + +@pytest.mark.unit +def test_gb200_8gpu_deterministic_pretrain_applies_determinism(): + """The deterministic wrapper keeps the GB200 layout and adds determinism settings and env vars.""" + base = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() + cfg = nemotron_3_nano_pretrain_8gpu_gb200_bf16_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.expert_model_parallel_size == base.model.expert_model_parallel_size + assert cfg.train.global_batch_size == base.train.global_batch_size + for key, value in DETERMINISM_ENV_VARS.items(): + assert cfg.env_vars[key] == value + assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 + assert cfg.env_vars["USE_MNNVL"] == 1 diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py index 26162a5ed8e..a73b9981cc6 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py @@ -27,6 +27,7 @@ ) from megatron.bridge.recipes.nemotronh.gb200.nemotron_3_ultra import ( nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config, + nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_deterministic_config, ) from megatron.bridge.recipes.nemotronh.h100.nemotron_3_ultra import ( _split_hybrid_layer_pattern, @@ -39,6 +40,7 @@ nemotron_3_ultra_pretrain_config, nemotron_3_ultra_sft_openmathinstruct2_packed_config, ) +from megatron.bridge.recipes.utils.environment_utils import DETERMINISM_ENV_VARS from megatron.bridge.utils.cuda_graph import cuda_graph_module_names from tests.unit_tests.recipes.recipe_test_utils import ( NEMOTRON_3_ULTRA_HYBRID_LAYER_PATTERN, @@ -169,6 +171,23 @@ def test_gb200_ep16_pretrain_preserves_functional_policy() -> None: assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 +@pytest.mark.unit +def test_gb200_ep16_deterministic_pretrain_applies_determinism() -> None: + assert recipes.nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_deterministic_config is ( + nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_deterministic_config + ) + + base = nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config() + cfg = nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_deterministic_config() + + assert cfg.model.deterministic_mode is True + assert cfg.model.cross_entropy_loss_fusion is False + assert cfg.model.expert_model_parallel_size == base.model.expert_model_parallel_size + for key, value in DETERMINISM_ENV_VARS.items(): + assert cfg.env_vars[key] == value + assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 + + @pytest.mark.unit def test_vrnvl72_perf_recipe_uses_nvl72_ultra_topology() -> None: """VRNVL72 Ultra preserves the GB300 execution layout with explicit NVL72 settings.""" diff --git a/tests/unit_tests/recipes/qwen/test_qwen35_dynamic_cp_recipe.py b/tests/unit_tests/recipes/qwen/test_qwen35_dynamic_cp_recipe.py index 567894a6bf7..ba43c52757d 100644 --- a/tests/unit_tests/recipes/qwen/test_qwen35_dynamic_cp_recipe.py +++ b/tests/unit_tests/recipes/qwen/test_qwen35_dynamic_cp_recipe.py @@ -18,6 +18,7 @@ from megatron.bridge.data.builders import GPTSFTDatasetConfig from megatron.bridge.recipes.qwen.gb200 import qwen35 +from megatron.bridge.recipes.utils.environment_utils import DETERMINISM_ENV_VARS from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies @@ -56,3 +57,15 @@ def test_dynamic_cp_sft_recipe_uses_real_data_and_dataset_driven_packing(): assert cfg.dataset.num_workers == 16 assert cfg.optimizer.use_precision_aware_optimizer is True assert cfg.validation.eval_iters == 0 + + +def test_text_9b_pretrain_deterministic_applies_determinism(): + base = qwen35.qwen35_text_9b_pretrain_8gpu_gb200_bf16_config() + cfg = qwen35.qwen35_text_9b_pretrain_8gpu_gb200_bf16_deterministic_config() + + assert cfg.model.deterministic_mode is True + assert cfg.model.cross_entropy_loss_fusion is False + assert cfg.model.gdn_kernel_backend == "torch" + assert cfg.model.tensor_model_parallel_size == base.model.tensor_model_parallel_size == 1 + assert cfg.train.global_batch_size == base.train.global_batch_size + assert cfg.env_vars == DETERMINISM_ENV_VARS diff --git a/tests/unit_tests/recipes/test_moonlight_recipes.py b/tests/unit_tests/recipes/test_moonlight_recipes.py index 2161918c532..aab225e70c4 100644 --- a/tests/unit_tests/recipes/test_moonlight_recipes.py +++ b/tests/unit_tests/recipes/test_moonlight_recipes.py @@ -329,6 +329,27 @@ def _assert_finetuning_optimizer_contract(cfg, *, precision_aware: bool): assert cfg.ddp.grad_reduce_in_fp32 is reduce_in_fp32 +def test_moonlight_16b_pretrain_deterministic_applies_determinism(monkeypatch: pytest.MonkeyPatch): + """The deterministic wrapper keeps the 16-GPU layout and sets the determinism env.""" + from megatron.bridge.recipes.moonlight import h100 as moonlight_h100 + from megatron.bridge.recipes.utils.environment_utils import DETERMINISM_ENV_VARS + + base_fn = moonlight_h100.moonlight_16b_pretrain_16gpu_h100_bf16_config + deterministic_fn = moonlight_h100.moonlight_16b_pretrain_16gpu_h100_bf16_deterministic_config + patch_recipe_module_global(monkeypatch, base_fn, "AutoBridge", _FakeBridge) + + base = base_fn() + cfg = deterministic_fn() + + assert cfg.model.deterministic_mode is True + assert cfg.model.cross_entropy_loss_fusion is False + assert cfg.model.expert_model_parallel_size == base.model.expert_model_parallel_size == 8 + assert cfg.train.global_batch_size == base.train.global_batch_size + for key, value in DETERMINISM_ENV_VARS.items(): + assert cfg.env_vars[key] == value + assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 8 + + def test_moonlight_16b_pretrain_convergence_contract(monkeypatch: pytest.MonkeyPatch): """Test the exact 16-GPU pretrain convergence contract.""" from megatron.bridge.recipes.moonlight import moonlight_16b_pretrain_config diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index 84d877f5408..96cbb2809b6 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -23,6 +23,7 @@ from megatron.bridge.perf_recipes._common import _benchmark_common from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS +from megatron.bridge.recipes.utils.environment_utils import DETERMINISM_ENV_VARS _CANONICAL_RECIPE_NAME = re.compile( @@ -102,8 +103,11 @@ def _explicit_environment(path: Path, function_name: str) -> dict[str, str | int common_expansions = 0 for key, value in zip(mapping.keys, mapping.values): if key is None: - assert isinstance(value, ast.Name) and value.id == "COMMON_PERF_ENV_VARS" - common_expansions += 1 + assert isinstance(value, ast.Name) and value.id in {"COMMON_PERF_ENV_VARS", "DETERMINISM_ENV_VARS"} + if value.id == "COMMON_PERF_ENV_VARS": + common_expansions += 1 + else: + result.update(DETERMINISM_ENV_VARS) continue if isinstance(value, ast.Name): assert value.id in local_constants diff --git a/tests/unit_tests/recipes/test_recipe_environment_defaults.py b/tests/unit_tests/recipes/test_recipe_environment_defaults.py index 1dc64e04e08..6ba0e545f30 100644 --- a/tests/unit_tests/recipes/test_recipe_environment_defaults.py +++ b/tests/unit_tests/recipes/test_recipe_environment_defaults.py @@ -20,7 +20,7 @@ import pytest -from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS +from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS, DETERMINISM_ENV_VARS _CANONICAL_RECIPE_NAME = re.compile( @@ -67,9 +67,12 @@ def _explicit_environment(path: Path, function_name: str) -> dict[str, str | int common_expansions = 0 for index, (key, value) in enumerate(zip(mapping.keys, mapping.values)): if key is None: - assert index == 0 - assert isinstance(value, ast.Name) and value.id == "COMMON_RECIPE_ENV_VARS" - common_expansions += 1 + assert isinstance(value, ast.Name) and value.id in {"COMMON_RECIPE_ENV_VARS", "DETERMINISM_ENV_VARS"} + if value.id == "COMMON_RECIPE_ENV_VARS": + assert index == 0 + common_expansions += 1 + else: + result.update(DETERMINISM_ENV_VARS) continue result[ast.literal_eval(key)] = ast.literal_eval(value) assert common_expansions == 1 diff --git a/tests/unit_tests/training/test_initialize.py b/tests/unit_tests/training/test_initialize.py index bba5e865f49..a6c0e58f232 100644 --- a/tests/unit_tests/training/test_initialize.py +++ b/tests/unit_tests/training/test_initialize.py @@ -28,6 +28,7 @@ from megatron.bridge.training.initialize import ( _initialize_tp_communicators, _setup_flight_recorder_env, + initialize_megatron, set_jit_fusion_options, ) @@ -561,3 +562,33 @@ def test_torch_dist_init_resets_default_logging_ranks(lazy): assert get_default_log_ranks() == (0,) finally: set_default_log_ranks(original) + + +@pytest.mark.unit +@pytest.mark.parametrize("deterministic_mode", [True, False]) +def test_initialize_megatron_pins_triton_autotune_only_in_deterministic_mode(deterministic_mode: bool) -> None: + """Timed Triton autotuning is pinned before any kernel runs, only when deterministic_mode is on.""" + cfg = Mock() + cfg.model.deterministic_mode = deterministic_mode + cfg.inprocess_restart = None + cfg.profiling = None + + with ( + patch("megatron.bridge.training.initialize.init_num_microbatches_calculator"), + patch("megatron.bridge.training.initialize.init_rerun_state"), + patch("megatron.bridge.training.initialize.torch_dist_init", return_value=None) as torch_dist_init, + patch("megatron.bridge.training.initialize.tuning") as tuning, + patch("megatron.bridge.training.initialize.torch.distributed.is_initialized", return_value=False), + ): + calls = Mock() + calls.attach_mock(tuning, "tuning") + calls.attach_mock(torch_dist_init, "torch_dist_init") + initialize_megatron(cfg, allow_no_cuda=True) + + called = [name for name, _, _ in calls.mock_calls] + if deterministic_mode: + tuning.install.assert_called_once_with(deterministic=True) + assert called.index("tuning.install") < called.index("torch_dist_init") + else: + tuning.install.assert_not_called() + torch_dist_init.assert_called_once()