[sync] Fine-grained activation offloading

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."
+            )

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:

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:
     """