diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index ab13fe0647..736a75cf65 100644 --- a/src/megatron/bridge/training/config.py +++ b/src/megatron/bridge/training/config.py @@ -20,9 +20,11 @@ from pathlib import Path from typing import Any, Literal, Optional, Tuple, Union +import torch from megatron.core.datasets.gpt_dataset import GPTDatasetConfig as MCoreGPTDatasetConfig from megatron.core.distributed import DistributedDataParallelConfig as MCoreDistributedDataParallelConfig from megatron.core.optimizer import OptimizerConfig as MCoreOptimizerConfig +from megatron.core.transformer.enums import AttnBackend from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs from megatron.bridge.models import GPTModelProvider, T5ModelProvider @@ -1084,6 +1086,32 @@ def set_data_parallel_size(self) -> None: if self.comm_overlap is not None: self.comm_overlap.data_parallel_size = self.data_parallel_size + def _validate_and_apply_deterministic_mode(self) -> None: + """Apply and validate deterministic mode requirements. + + This enforces restrictions and settings that must hold when + the model is configured to run in deterministic mode. + """ + if not getattr(self.model, "deterministic_mode", False): + return + + # Disallow flash attention when running deterministically + if getattr(self.model, "attention_backend", None) == AttnBackend.flash: + raise AssertionError("Flash attention can not be used in deterministic mode.") + + # 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}." + ) + + # Enable deterministic algorithms in torch + torch.use_deterministic_algorithms(True) + def _sync_and_validate_external_cuda_graph(self) -> None: """Sync necessary configs for external CUDA Graphs and and validates it.""" @@ -1149,6 +1177,9 @@ def validate(self) -> None: if self.comm_overlap is not None: self.comm_overlap.data_parallel_size = self.data_parallel_size + # Deterministic mode validations and settings + self._validate_and_apply_deterministic_mode() + # Run validations _validate_and_sync_distributed_optimizer_settings(self) diff --git a/tests/unit_tests/training/test_config.py b/tests/unit_tests/training/test_config.py index 83c4ee8c3c..6b77c670b8 100644 --- a/tests/unit_tests/training/test_config.py +++ b/tests/unit_tests/training/test_config.py @@ -383,7 +383,66 @@ def test_cannot_set_blend_fields(self): class TestConfigContainerValidation: - """Tests for the `validate` method of the `ConfigContainer` class.""" + def test_deterministic_mode_disallows_flash_and_ce_fusion(self, monkeypatch): + """Test that deterministic mode disallows flash attention and cross-entropy loss fusion.""" + from megatron.core.transformer.enums import AttnBackend + + gpt_model_cfg = create_test_gpt_config( + deterministic_mode=True, + attention_backend=AttnBackend.flash, + cross_entropy_loss_fusion=True, + ) + + # Ensure NCCL_ALGO present but valid, so we fail earlier on flash/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="Flash attention can not be used in deterministic mode"): + container.validate() + + # Fix attention, still CE fusion should fail + container.model.attention_backend = AttnBackend.local + with pytest.raises(AssertionError, match="Cross Entropy Fusion is currently not deterministic"): + 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.""" + gpt_model_cfg = create_test_gpt_config( + deterministic_mode=True, + cross_entropy_loss_fusion=False, + transformer_impl="transformer_engine", + ) + + 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() + + # 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 + finally: + restore_get_world_size_safe(og_ws, cfg_mod) @pytest.mark.parametrize( "world_size, expect_assertion_error",