diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index 446c53a371..d3a1e1cbbc 100644 --- a/src/megatron/bridge/training/config.py +++ b/src/megatron/bridge/training/config.py @@ -1405,6 +1405,7 @@ def validate(self) -> None: # Run validations _validate_and_sync_distributed_optimizer_settings(self) _validate_mixed_precision_consistency(self) + _validate_fine_grained_activation_offloading(self) if self.dist.use_megatron_fsdp and self.dist.use_torch_fsdp2: raise ValueError("Using use_megatron_fsdp and use_torch_fsdp2 at the same time is not supported.") @@ -1705,3 +1706,39 @@ def _validate_mixed_precision_consistency(config: ConfigContainer) -> None: "model is using fp32 precision (model.bf16=False, model.fp16=False) and " "use_precision_aware_optimizer=True." ) + + +def _validate_fine_grained_activation_offloading(config: ConfigContainer) -> None: + """Validate fine-grained activation offloading configuration. + + This function ensures that fine-grained activation offloading is only enabled + with compatible configurations (transformer_engine implementation) and that + necessary environment variables are set for newer TE versions. + + Args: + config: The configuration container to validate. + + Raises: + ValueError: If fine-grained activation offloading is enabled with incompatible settings. + """ + from megatron.core.utils import is_te_min_version + + model_cfg = config.model + + if not model_cfg.fine_grained_activation_offloading: + return + + # Fine-grained activation offloading requires transformer_engine implementation + if model_cfg.transformer_impl != "transformer_engine": + raise ValueError( + "Fine-grained activation offloading is only supported with transformer_engine implementation. " + f"Current transformer_impl: {model_cfg.transformer_impl}" + ) + + # For TE >= 2.10.0, NVTE_CPU_OFFLOAD_V1 must be set to avoid offloading weights + if is_te_min_version("2.10.0"): + if os.getenv("NVTE_CPU_OFFLOAD_V1", "0") != "1": + raise ValueError( + "For fine-grained activation offloading with TE >= 2.10.0, " + "NVTE_CPU_OFFLOAD_V1 environment variable should be set to 1 to avoid offloading weights." + ) diff --git a/src/megatron/bridge/training/setup.py b/src/megatron/bridge/training/setup.py index 85a0dd136c..42b70c07b4 100644 --- a/src/megatron/bridge/training/setup.py +++ b/src/megatron/bridge/training/setup.py @@ -139,6 +139,12 @@ def setup( restart_store=restart_store, ) + # Set CPU affinity for optimal host-device transfers when fine-grained activation offloading is enabled + if cfg.model.fine_grained_activation_offloading: + from megatron.core.pipeline_parallel.utils import set_ideal_affinity_for_current_gpu + + set_ideal_affinity_for_current_gpu() + timers = state.timers if cfg.logger.log_progress: diff --git a/tests/unit_tests/training/test_config.py b/tests/unit_tests/training/test_config.py index 9e412deb98..60842a83cf 100644 --- a/tests/unit_tests/training/test_config.py +++ b/tests/unit_tests/training/test_config.py @@ -1206,6 +1206,102 @@ def test_modelopt_requires_no_gradient_accumulation_fusion(self, monkeypatch): finally: restore_get_world_size_safe(og_ws, cfg_mod) + @patch("megatron.core.utils.is_te_min_version") + def test_fine_grained_activation_offloading_requires_transformer_engine(self, mock_is_te_min_version, monkeypatch): + """Test that fine_grained_activation_offloading requires transformer_engine implementation.""" + mock_is_te_min_version.return_value = False # Pretend TE < 2.10.0 + + gpt_model_cfg = create_test_gpt_config( + fine_grained_activation_offloading=True, + offload_modules=["attn_norm"], # Required when fine_grained_activation_offloading=True + transformer_impl="local", # Using local instead of transformer_engine + ) + container, og_ws, cfg_mod = create_test_config_container(world_size_override=1, model_config=gpt_model_cfg) + + try: + with pytest.raises( + ValueError, + match="Fine-grained activation offloading is only supported with transformer_engine implementation", + ): + container.validate() + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + @patch("megatron.core.utils.is_te_min_version") + def test_fine_grained_activation_offloading_with_transformer_engine_passes( + self, mock_is_te_min_version, monkeypatch + ): + """Test that fine_grained_activation_offloading passes with transformer_engine implementation.""" + mock_is_te_min_version.return_value = False # Pretend TE < 2.10.0 to skip env var check + + gpt_model_cfg = create_test_gpt_config( + fine_grained_activation_offloading=True, + offload_modules=["attn_norm"], # Required when fine_grained_activation_offloading=True + transformer_impl="transformer_engine", + ) + container, og_ws, cfg_mod = create_test_config_container(world_size_override=1, model_config=gpt_model_cfg) + + try: + container.validate() # Should pass without error + assert container.model.fine_grained_activation_offloading is True + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + @patch.dict("os.environ", {"NVTE_CPU_OFFLOAD_V1": "0"}) + @patch("megatron.core.utils.is_te_min_version") + def test_fine_grained_activation_offloading_te_2_10_requires_env_var(self, mock_is_te_min_version, monkeypatch): + """Test that fine_grained_activation_offloading with TE >= 2.10.0 requires NVTE_CPU_OFFLOAD_V1=1.""" + mock_is_te_min_version.return_value = True # Pretend TE >= 2.10.0 + + gpt_model_cfg = create_test_gpt_config( + fine_grained_activation_offloading=True, + offload_modules=["attn_norm"], # Required when fine_grained_activation_offloading=True + transformer_impl="transformer_engine", + ) + container, og_ws, cfg_mod = create_test_config_container(world_size_override=1, model_config=gpt_model_cfg) + + try: + with pytest.raises( + ValueError, + match="NVTE_CPU_OFFLOAD_V1 environment variable should be set to 1", + ): + container.validate() + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + @patch.dict("os.environ", {"NVTE_CPU_OFFLOAD_V1": "1"}) + @patch("megatron.core.utils.is_te_min_version") + def test_fine_grained_activation_offloading_te_2_10_with_env_var_passes(self, mock_is_te_min_version, monkeypatch): + """Test that fine_grained_activation_offloading with TE >= 2.10.0 and NVTE_CPU_OFFLOAD_V1=1 passes.""" + mock_is_te_min_version.return_value = True # Pretend TE >= 2.10.0 + + gpt_model_cfg = create_test_gpt_config( + fine_grained_activation_offloading=True, + offload_modules=["attn_norm"], # Required when fine_grained_activation_offloading=True + transformer_impl="transformer_engine", + ) + container, og_ws, cfg_mod = create_test_config_container(world_size_override=1, model_config=gpt_model_cfg) + + try: + container.validate() # Should pass without error + assert container.model.fine_grained_activation_offloading is True + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + def test_fine_grained_activation_offloading_disabled_skips_validation(self, monkeypatch): + """Test that validation is skipped when fine_grained_activation_offloading is disabled.""" + gpt_model_cfg = create_test_gpt_config( + fine_grained_activation_offloading=False, + transformer_impl="local", # Would fail if validation was run + ) + container, og_ws, cfg_mod = create_test_config_container(world_size_override=1, model_config=gpt_model_cfg) + + try: + container.validate() # Should pass without error since offloading is disabled + assert container.model.fine_grained_activation_offloading is False + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + class TestRerunConfigValidation: """