Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/recipe-usage.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions src/megatron/bridge/recipes/moonlight/h100/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
27 changes: 26 additions & 1 deletion src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
# =============================================================================
Expand Down Expand Up @@ -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",
Expand Down
4 changes: 4 additions & 0 deletions src/megatron/bridge/recipes/nemotronh/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand Down
4 changes: 4 additions & 0 deletions src/megatron/bridge/recipes/nemotronh/gb200/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand All @@ -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",
]
29 changes: 28 additions & 1 deletion src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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",
]
37 changes: 35 additions & 2 deletions src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_ultra.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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",
]
2 changes: 2 additions & 0 deletions src/megatron/bridge/recipes/qwen/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down
2 changes: 2 additions & 0 deletions src/megatron/bridge/recipes/qwen/gb200/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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",
Expand Down
23 changes: 22 additions & 1 deletion src/megatron/bridge/recipes/qwen/gb200/qwen35.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
12 changes: 12 additions & 0 deletions src/megatron/bridge/recipes/utils/environment_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
}
8 changes: 7 additions & 1 deletion src/megatron/bridge/training/initialize.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
22 changes: 22 additions & 0 deletions tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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

Expand Down Expand Up @@ -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
Loading
Loading