perf: selective activation checkpointing feature support

examples/configs/distillation_math.yaml

@@ -90,6 +90,17 @@ policy: &POLICY_BASE
         force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
         empty_unused_memory_level: 0
         activation_checkpointing: false
+        # recompute_granularity controls activation checkpointing depth.
+        # "full": recompute all activations (default, max memory savings).
+        # "selective": recompute only specific modules (see recompute_modules).
+        # Only takes effect when activation_checkpointing: true.
+        recompute_granularity: "full"
+        # Modules to selectively recompute when recompute_granularity="selective".
+        # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+        # Null defaults to ["core_attn"]. Full list and per-module constraints:
+        # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+        # Use ["moe"] for MoE models to recompute expert activations only.
+        recompute_modules: null
         converter_type: "Qwen3ForCausalLM"
         tensor_model_parallel_size: 2
         expert_tensor_parallel_size: 1

examples/configs/distillation_math_megatron.yaml

@@ -39,6 +39,17 @@ policy: &POLICY_BASE
         enabled: true
         empty_unused_memory_level: 0
         activation_checkpointing: false
+        # recompute_granularity controls activation checkpointing depth.
+        # "full": recompute all activations (default, max memory savings).
+        # "selective": recompute only specific modules (see recompute_modules).
+        # Only takes effect when activation_checkpointing: true.
+        recompute_granularity: "full"
+        # Modules to selectively recompute when recompute_granularity="selective".
+        # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+        # Null defaults to ["core_attn"]. Full list and per-module constraints:
+        # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+        # Use ["moe"] for MoE models to recompute expert activations only.
+        recompute_modules: null
         converter_type: "Qwen3ForCausalLM"
         tensor_model_parallel_size: 2
         expert_tensor_parallel_size: 1

examples/configs/dpo.yaml

@@ -129,6 +129,17 @@ policy:
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1
     activation_checkpointing: false
+    # recompute_granularity controls activation checkpointing depth.
+    # "full": recompute all activations (default, max memory savings).
+    # "selective": recompute only specific modules (see recompute_modules).
+    # Only takes effect when activation_checkpointing: true.
+    recompute_granularity: "full"
+    # Modules to selectively recompute when recompute_granularity="selective".
+    # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+    # Null defaults to ["core_attn"]. Full list and per-module constraints:
+    # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+    # Use ["moe"] for MoE models to recompute expert activations only.
+    recompute_modules: null
     tensor_model_parallel_size: 2
     expert_tensor_parallel_size: 1
     expert_model_parallel_size: 1

examples/configs/grpo_math_1B.yaml

@@ -134,6 +134,17 @@ policy:
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     activation_checkpointing: false
+    # recompute_granularity controls activation checkpointing depth.
+    # "full": recompute all activations (default, max memory savings).
+    # "selective": recompute only specific modules (see recompute_modules).
+    # Only takes effect when activation_checkpointing: true.
+    recompute_granularity: "full"
+    # Modules to selectively recompute when recompute_granularity="selective".
+    # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+    # Null defaults to ["core_attn"]. Full list and per-module constraints:
+    # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+    # Use ["moe"] for MoE models to recompute expert activations only.
+    recompute_modules: null
     converter_type: "Qwen2ForCausalLM"
     tensor_model_parallel_size: 1
     expert_tensor_parallel_size: 1

examples/configs/grpo_math_1B_megatron.yaml

@@ -86,6 +86,17 @@ policy:
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     activation_checkpointing: false
+    # recompute_granularity controls activation checkpointing depth.
+    # "full": recompute all activations (default, max memory savings).
+    # "selective": recompute only specific modules (see recompute_modules).
+    # Only takes effect when activation_checkpointing: true.
+    recompute_granularity: "full"
+    # Modules to selectively recompute when recompute_granularity="selective".
+    # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+    # Null defaults to ["core_attn"]. Full list and per-module constraints:
+    # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+    # Use ["moe"] for MoE models to recompute expert activations only.
+    recompute_modules: null
     converter_type: "Qwen2ForCausalLM"
     tensor_model_parallel_size: 1
     expert_tensor_parallel_size: 1

examples/configs/sft.yaml

@@ -110,6 +110,17 @@ policy:
     env_vars: {}
     empty_unused_memory_level: 1
     activation_checkpointing: false
+    # recompute_granularity controls activation checkpointing depth.
+    # "full": recompute all activations (default, max memory savings).
+    # "selective": recompute only specific modules (see recompute_modules).
+    # Only takes effect when activation_checkpointing: true.
+    recompute_granularity: "full"
+    # Modules to selectively recompute when recompute_granularity="selective".
+    # MCore options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+    # Null defaults to ["core_attn"]. Full list and per-module constraints:
+    # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+    # Use ["moe"] for MoE models to recompute expert activations only.
+    recompute_modules: null
     tensor_model_parallel_size: 1
     expert_tensor_parallel_size: 1
     expert_model_parallel_size: 1

nemo_rl/models/megatron/setup.py

@@ -669,9 +669,26 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None:
 
     # Activation checkpointing
     if config["megatron_cfg"]["activation_checkpointing"]:
-        model_cfg.recompute_granularity = "full"
-        model_cfg.recompute_method = "uniform"
-        model_cfg.recompute_num_layers = 1
+        granularity = config["megatron_cfg"].get("recompute_granularity", "full")
+        model_cfg.recompute_granularity = granularity
+        if granularity == "full":
+            model_cfg.recompute_method = "uniform"
+            model_cfg.recompute_num_layers = 1
+        elif granularity == "selective":
+            recompute_modules = config["megatron_cfg"].get("recompute_modules")
+            if recompute_modules is not None:
+                # NOTE: MCore validates recompute_modules in TransformerConfig.__post_init__,
+                # but that validation doesn't re-run after attribute assignment here.
+                # Valid values: core_attn, moe_act, layernorm, mla_up_proj, mlp, moe, shared_experts
+                # See: https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L1365
+                # Tracking: https://github.com/NVIDIA-NeMo/RL/issues/2291
+                model_cfg.recompute_modules = recompute_modules
+            # else: MCore defaults to ["core_attn"] when recompute_modules is None
+        else:
+            raise ValueError(
+                f"Invalid recompute_granularity: {granularity!r}. "
+                "Valid options are 'full' or 'selective'."
+            )
 
     # Activation function validation
     if not model_cfg.gated_linear_unit:

nemo_rl/models/policy/__init__.py

@@ -204,6 +204,16 @@ class MegatronConfig(TypedDict):
     # Setting to 0 is faster, but you are more likely to run out of GPU memory. In SFT/DPO, the default is 0.
     empty_unused_memory_level: int
     activation_checkpointing: bool
+    # Recompute granularity: "full" recomputes all activations, "selective" recomputes
+    # only specific modules (see recompute_modules). "selective" typically saves ~10-18GB
+    # for MoE models while retaining higher throughput than "full".
+    recompute_granularity: NotRequired[Literal["full", "selective"]]
+    # Modules to selectively recompute when recompute_granularity="selective".
+    # MCore valid options: ["core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", "shared_experts"].
+    # Defaults to ["core_attn"] when None. Full list and per-module constraints:
+    # https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L483
+    # when None. Use ["moe"] to recompute only expert activations (production-proven config).
+    recompute_modules: NotRequired[list[str] | None]
     tensor_model_parallel_size: int
     pipeline_model_parallel_size: int
     num_layers_in_first_pipeline_stage: int | None

tests/unit/models/megatron/test_megatron_setup.py

@@ -1005,6 +1005,93 @@ def test_fp8_param_warning(self):
         with pytest.warns(UserWarning, match="fp8_param=True sometimes causes NaN"):
             _apply_performance_config(model_cfg, config)
 
+    def test_recompute_granularity_full_explicit(self):
+        """granularity='full' sets uniform method with 1 layer."""
+        from nemo_rl.models.megatron.setup import _apply_performance_config
+
+        model_cfg = MagicMock()
+        model_cfg.gated_linear_unit = True
+        config = {
+            "megatron_cfg": {
+                "activation_checkpointing": True,
+                "recompute_granularity": "full",
+                "apply_rope_fusion": False,
+                "bias_activation_fusion": False,
+                "gradient_accumulation_fusion": False,
+            }
+        }
+
+        _apply_performance_config(model_cfg, config)
+
+        assert model_cfg.recompute_granularity == "full"
+        assert model_cfg.recompute_method == "uniform"
+        assert model_cfg.recompute_num_layers == 1
+
+    def test_recompute_granularity_selective_with_modules(self):
+        """granularity='selective' with explicit modules sets recompute_modules."""
+        from nemo_rl.models.megatron.setup import _apply_performance_config
+
+        model_cfg = MagicMock()
+        model_cfg.gated_linear_unit = True
+        modules = ["core_attn", "moe_act"]
+        config = {
+            "megatron_cfg": {
+                "activation_checkpointing": True,
+                "recompute_granularity": "selective",
+                "recompute_modules": modules,
+                "apply_rope_fusion": False,
+                "bias_activation_fusion": False,
+                "gradient_accumulation_fusion": False,
+            }
+        }
+
+        _apply_performance_config(model_cfg, config)
+
+        assert model_cfg.recompute_granularity == "selective"
+        assert model_cfg.recompute_modules == modules
+
+    def test_recompute_granularity_selective_without_modules_uses_mcore_default(self):
+        """granularity='selective' without recompute_modules leaves attr untouched (MCore default applies)."""
+        from nemo_rl.models.megatron.setup import _apply_performance_config
+
+        model_cfg = MagicMock(spec=["gated_linear_unit"])
+        model_cfg.gated_linear_unit = True
+        config = {
+            "megatron_cfg": {
+                "activation_checkpointing": True,
+                "recompute_granularity": "selective",
+                "apply_rope_fusion": False,
+                "bias_activation_fusion": False,
+                "gradient_accumulation_fusion": False,
+            }
+        }
+
+        _apply_performance_config(model_cfg, config)
+
+        assert model_cfg.recompute_granularity == "selective"
+        assert not hasattr(model_cfg, "recompute_modules")
+        assert not hasattr(model_cfg, "recompute_method")
+        assert not hasattr(model_cfg, "recompute_num_layers")
+
+    def test_recompute_granularity_invalid_raises(self):
+        """Invalid granularity raises ValueError with a helpful message."""
+        from nemo_rl.models.megatron.setup import _apply_performance_config
+
+        model_cfg = MagicMock()
+        model_cfg.gated_linear_unit = True
+        config = {
+            "megatron_cfg": {
+                "activation_checkpointing": True,
+                "recompute_granularity": "block",
+                "apply_rope_fusion": False,
+                "bias_activation_fusion": False,
+                "gradient_accumulation_fusion": False,
+            }
+        }
+
+        with pytest.raises(ValueError, match="Invalid recompute_granularity"):
+            _apply_performance_config(model_cfg, config)
+
 
 @pytest.mark.mcore
 class TestValidateOptimizerConfig:

tests/unit/reference_configs/distillation_math.yaml

@@ -90,6 +90,8 @@ policy: &POLICY_BASE
         force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
         empty_unused_memory_level: 0
         activation_checkpointing: false
+        recompute_granularity: "full"
+        recompute_modules: null
         converter_type: "Qwen3ForCausalLM"
         tensor_model_parallel_size: 2
         expert_tensor_parallel_size: 1

tests/unit/reference_configs/dpo.yaml

@@ -124,6 +124,8 @@ policy:
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1
     activation_checkpointing: false
+    recompute_granularity: "full"
+    recompute_modules: null
     tensor_model_parallel_size: 2
     expert_tensor_parallel_size: 1
     expert_model_parallel_size: 1

tests/unit/reference_configs/grpo_math_1B.yaml

@@ -135,6 +135,8 @@ policy:
     force_reconvert_from_hf: False # Set to True to force reconvert of the model from Hugging Face
     empty_unused_memory_level: 1  # 1 is the minimum recommendation for RL since we almost always need to offload before beginning generation. Setting to 0 is faster, but you are more likely to run out of GPU memory.
     activation_checkpointing: false
+    recompute_granularity: "full"
+    recompute_modules: null
     converter_type: "Qwen2ForCausalLM"
     tensor_model_parallel_size: 1
     expert_tensor_parallel_size: 1

tests/unit/reference_configs/sft.yaml

@@ -105,6 +105,8 @@ policy:
     env_vars: {}
     empty_unused_memory_level: 1
     activation_checkpointing: false
+    recompute_granularity: "full"
+    recompute_modules: null
     tensor_model_parallel_size: 1
     expert_tensor_parallel_size: 1
     expert_model_parallel_size: 1