From feee53ef4e35418b352a7aaefcdcf21ed35e48cd Mon Sep 17 00:00:00 2001 From: sna Date: Fri, 17 Apr 2026 00:08:44 -0700 Subject: [PATCH 01/24] Add fine-grained activation offloading for Megatron policy Exposes Megatron-Core fine_grained_activation_offloading and offload_modules through PolicyConfig so training can offload specific submodule activations (moe_act, core_attn, qkv_linear, mlp_norm, attn_norm) to CPU. Works for both dense and MoE models. Validation of module names is left to Megatron. Signed-off-by: sna --- nemo_rl/models/megatron/setup.py | 18 ++++++++++++++++++ nemo_rl/models/policy/__init__.py | 8 ++++++++ 2 files changed, 26 insertions(+) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index fc5c6c44fa4..4bdc1eb0a27 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -514,6 +514,24 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: "Refer to https://github.com/NVIDIA-NeMo/RL/issues/1164 for latest updates with this issue." ) + # Fine-grained activation offloading moves specified submodule activations + # to CPU. Works for both dense and MoE models; the user picks which + # submodules to offload via offload_modules. Megatron owns the list of + # valid module names and their per-model-type compatibility, so we only + # require a non-empty list here and let Megatron validate the contents. + fine_grained_activation_offloading = config["megatron_cfg"].get( + "fine_grained_activation_offloading", False + ) + if fine_grained_activation_offloading: + offload_modules = config["megatron_cfg"].get("offload_modules", []) + if not offload_modules: + raise ValueError( + "offload_modules must be a non-empty list when " + "fine_grained_activation_offloading is True." + ) + model_cfg.fine_grained_activation_offloading = True + model_cfg.offload_modules = offload_modules + def _validate_optimizer_config(config: PolicyConfig) -> None: """Validate optimizer configuration.""" diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index ec4c9e66bbb..28d9a8bbc7b 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -236,6 +236,14 @@ class MegatronConfig(TypedDict): moe_token_dispatcher_type: str # Can be used only with 'alltoall' token dispatcher moe_shared_expert_overlap: bool + # Offload specific module activations to CPU to reduce peak GPU memory. + # Works with MoE models (offloads MoE expert activations). Different from + # optimizer_cpu_offload which offloads optimizer states. + fine_grained_activation_offloading: NotRequired[bool] + # Modules to offload when fine_grained_activation_offloading is True. + # Defaults to ["moe_act"] if not specified. Valid values include: + # "moe_act", "core_attn", "qkv_linear", "mlp_norm", "attn_norm". + offload_modules: NotRequired[list[str]] peft: NotRequired[MegatronPeftConfig | MegatronPeftConfigDisabled] optimizer: MegatronOptimizerConfig scheduler: MegatronSchedulerConfig From 0b530c4d4dcff294584b963fd453f2142d1b5c3f Mon Sep 17 00:00:00 2001 From: Seonjin Date: Thu, 23 Apr 2026 14:31:47 -0700 Subject: [PATCH 02/24] Update nemo_rl/models/policy/__init__.py Co-authored-by: Terry Kong Signed-off-by: Seonjin --- nemo_rl/models/policy/__init__.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 28d9a8bbc7b..fa59cd94e69 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -237,12 +237,15 @@ class MegatronConfig(TypedDict): # Can be used only with 'alltoall' token dispatcher moe_shared_expert_overlap: bool # Offload specific module activations to CPU to reduce peak GPU memory. - # Works with MoE models (offloads MoE expert activations). Different from + # Works with both dense and MoE models. Different from # optimizer_cpu_offload which offloads optimizer states. + # Requires transformer_engine implementation. fine_grained_activation_offloading: NotRequired[bool] # Modules to offload when fine_grained_activation_offloading is True. - # Defaults to ["moe_act"] if not specified. Valid values include: - # "moe_act", "core_attn", "qkv_linear", "mlp_norm", "attn_norm". + # Required (no default). Valid values: + # "attn_norm", "qkv_linear", "core_attn", "attn_proj", "mlp_norm", + # "expert_fc1", "moe_act". Note: "attn_proj" requires "core_attn". + # See: https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L1440-L1448 offload_modules: NotRequired[list[str]] peft: NotRequired[MegatronPeftConfig | MegatronPeftConfigDisabled] optimizer: MegatronOptimizerConfig From d5df80afa833702a1a96f808275463d70e43d427 Mon Sep 17 00:00:00 2001 From: Seonjin Date: Thu, 23 Apr 2026 14:32:14 -0700 Subject: [PATCH 03/24] Update nemo_rl/models/megatron/setup.py Co-authored-by: Terry Kong Signed-off-by: Seonjin --- nemo_rl/models/megatron/setup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 4bdc1eb0a27..caef9c10d99 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -524,11 +524,12 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: ) if fine_grained_activation_offloading: offload_modules = config["megatron_cfg"].get("offload_modules", []) - if not offload_modules: + if not isinstance(offload_modules, list) or not offload_modules: raise ValueError( "offload_modules must be a non-empty list when " "fine_grained_activation_offloading is True." ) + ) model_cfg.fine_grained_activation_offloading = True model_cfg.offload_modules = offload_modules From e23798739aace8e26989726f05e4902d78ea16b4 Mon Sep 17 00:00:00 2001 From: sna Date: Thu, 14 May 2026 14:16:44 -0700 Subject: [PATCH 04/24] fix: remove stray paren in setup.py raising ValueError A stray closing parenthesis after the raise ValueError block caused a SyntaxError, blocking the ruff/ruff-format pre-commit hooks in CI. Signed-off-by: sna --- nemo_rl/models/megatron/setup.py | 1 - 1 file changed, 1 deletion(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 685e8aa3e74..841ac8763b6 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -687,7 +687,6 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: "offload_modules must be a non-empty list when " "fine_grained_activation_offloading is True." ) - ) model_cfg.fine_grained_activation_offloading = True model_cfg.offload_modules = offload_modules From 06b4d4abcfd3eebe1356827bfeccba66c7bed8c4 Mon Sep 17 00:00:00 2001 From: sna Date: Fri, 15 May 2026 11:52:45 -0700 Subject: [PATCH 05/24] fix: pin NeMo Gym docs URL to v0.2.1 (latest 404) Signed-off-by: sna --- docs/design-docs/nemo-gym-integration.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/design-docs/nemo-gym-integration.md b/docs/design-docs/nemo-gym-integration.md index 33e324547bf..0263d36fef8 100644 --- a/docs/design-docs/nemo-gym-integration.md +++ b/docs/design-docs/nemo-gym-integration.md @@ -181,7 +181,7 @@ sequenceDiagram GRPO->>Policy: Compute loss and train ``` -> **NeMo Gym server types** (see [Core Components](https://docs.nvidia.com/nemo/gym/latest/about/concepts/core-components.html)): +> **NeMo Gym server types** (see [Core Components](https://docs.nvidia.com/nemo/gym/v0.2.1/about/concepts/core-components/)): > - **Agent Server**: Orchestrates the rollout loop > - **Model Server**: HTTP proxy to vLLM; translates Responses API ↔ Chat Completions > - **Resource Server**: Provides tools and rewards From 522521770628acb6ff9d31407d8c92fdb825261b Mon Sep 17 00:00:00 2001 From: sna Date: Sat, 16 May 2026 15:23:14 -0700 Subject: [PATCH 06/24] test: add unit tests for fine_grained_activation_offloading branch Covers _apply_performance_config offload-modules dispatch: - happy path: True + non-empty list sets both attrs - disabled: defaults skip the branch (no attrs touched) - invalid offload_modules ([], None, str, int) all raise ValueError - missing offload_modules key raises ValueError Lifts patch coverage above codecov target. Signed-off-by: sna --- .../models/megatron/test_megatron_setup.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 1eaa3a12477..c2dce64c315 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -794,6 +794,99 @@ 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_fine_grained_activation_offloading_enabled(self): + """Test happy path: enabled with non-empty offload_modules list.""" + from nemo_rl.models.megatron.setup import _apply_performance_config + + model_cfg = MagicMock() + model_cfg.gated_linear_unit = True + offload_modules = ["mlp", "moe_act"] + config = { + "megatron_cfg": { + "activation_checkpointing": False, + "apply_rope_fusion": False, + "bias_activation_fusion": False, + "gradient_accumulation_fusion": False, + "fine_grained_activation_offloading": True, + "offload_modules": offload_modules, + } + } + + _apply_performance_config(model_cfg, config) + + assert model_cfg.fine_grained_activation_offloading is True + assert model_cfg.offload_modules == offload_modules + + def test_fine_grained_activation_offloading_disabled_skips(self): + """When flag is False (default), no offload attrs should be set.""" + 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": False, + "apply_rope_fusion": False, + "bias_activation_fusion": False, + "gradient_accumulation_fusion": False, + } + } + + _apply_performance_config(model_cfg, config) + + assert not hasattr(model_cfg, "fine_grained_activation_offloading") + assert not hasattr(model_cfg, "offload_modules") + + @pytest.mark.parametrize( + "offload_modules", + [[], None, "mlp", 42], + ids=["empty_list", "none", "string", "int"], + ) + def test_fine_grained_activation_offloading_invalid_modules_raises( + self, offload_modules + ): + """offload_modules must be a non-empty list when feature is enabled.""" + from nemo_rl.models.megatron.setup import _apply_performance_config + + model_cfg = MagicMock() + model_cfg.gated_linear_unit = True + config = { + "megatron_cfg": { + "activation_checkpointing": False, + "apply_rope_fusion": False, + "bias_activation_fusion": False, + "gradient_accumulation_fusion": False, + "fine_grained_activation_offloading": True, + "offload_modules": offload_modules, + } + } + + with pytest.raises( + ValueError, match="offload_modules must be a non-empty list" + ): + _apply_performance_config(model_cfg, config) + + def test_fine_grained_activation_offloading_missing_modules_raises(self): + """When enabled but offload_modules key is absent, defaults to [] → raises.""" + from nemo_rl.models.megatron.setup import _apply_performance_config + + model_cfg = MagicMock() + model_cfg.gated_linear_unit = True + config = { + "megatron_cfg": { + "activation_checkpointing": False, + "apply_rope_fusion": False, + "bias_activation_fusion": False, + "gradient_accumulation_fusion": False, + "fine_grained_activation_offloading": True, + } + } + + with pytest.raises( + ValueError, match="offload_modules must be a non-empty list" + ): + _apply_performance_config(model_cfg, config) + @pytest.mark.mcore class TestValidateOptimizerConfig: From 211e31a37dc576837437582842774ad7b30ee172 Mon Sep 17 00:00:00 2001 From: Seonjin Date: Mon, 18 May 2026 22:14:36 -0700 Subject: [PATCH 07/24] Update nemo_rl/models/megatron/setup.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Seonjin --- nemo_rl/models/megatron/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 841ac8763b6..18571850483 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -681,7 +681,7 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: "fine_grained_activation_offloading", False ) if fine_grained_activation_offloading: - offload_modules = config["megatron_cfg"].get("offload_modules", []) + offload_modules = config["megatron_cfg"].get("offload_modules") if not isinstance(offload_modules, list) or not offload_modules: raise ValueError( "offload_modules must be a non-empty list when " From 502d2ddf1aeecfd2a7ecad6e325f158143941d98 Mon Sep 17 00:00:00 2001 From: Seonjin Date: Mon, 18 May 2026 22:14:48 -0700 Subject: [PATCH 08/24] Update nemo_rl/models/megatron/setup.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Seonjin --- nemo_rl/models/megatron/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 18571850483..b141215874f 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -678,7 +678,7 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: # valid module names and their per-model-type compatibility, so we only # require a non-empty list here and let Megatron validate the contents. fine_grained_activation_offloading = config["megatron_cfg"].get( - "fine_grained_activation_offloading", False + "fine_grained_activation_offloading" ) if fine_grained_activation_offloading: offload_modules = config["megatron_cfg"].get("offload_modules") From da947f801689a9e13e3d885599731b111d6f5e72 Mon Sep 17 00:00:00 2001 From: Seonjin Date: Mon, 18 May 2026 22:38:18 -0700 Subject: [PATCH 09/24] Update tests/unit/models/megatron/test_megatron_setup.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Seonjin --- tests/unit/models/megatron/test_megatron_setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index c2dce64c315..ff081ba68c7 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -864,7 +864,7 @@ def test_fine_grained_activation_offloading_invalid_modules_raises( with pytest.raises( ValueError, match="offload_modules must be a non-empty list" ): - _apply_performance_config(model_cfg, config) + """When enabled but offload_modules key is absent, defaults to None → raises.""" def test_fine_grained_activation_offloading_missing_modules_raises(self): """When enabled but offload_modules key is absent, defaults to [] → raises.""" From 254c7d097de8646f0b0671b22004218e96488371 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Tue, 19 May 2026 14:04:01 -0700 Subject: [PATCH 10/24] Fix syntax error from main merge in offload test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The merge resolution left an empty 'with pytest.raises(...)' block — restore the _apply_performance_config(model_cfg, config) call inside it. Signed-off-by: seonjinn --- tests/unit/models/megatron/test_megatron_setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 4ec60bda13d..617428fd116 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -1075,7 +1075,7 @@ def test_fine_grained_activation_offloading_invalid_modules_raises( with pytest.raises( ValueError, match="offload_modules must be a non-empty list" ): - """When enabled but offload_modules key is absent, defaults to None → raises.""" + _apply_performance_config(model_cfg, config) def test_fine_grained_activation_offloading_missing_modules_raises(self): """When enabled but offload_modules key is absent, defaults to [] → raises.""" From ed86517d14ffc7d3266eb5394570c079d383b3a3 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 23 May 2026 00:10:47 -0700 Subject: [PATCH 11/24] Address review comments on fine_grained_activation_offloading - Add fine_grained_activation_offloading and offload_modules to all megatron-capable exemplar configs (grpo_math_1B, grpo_math_1B_megatron, sft, dpo, distillation_math, distillation_math_megatron). - Sync tests/unit/reference_configs/*.yaml with the new keys so test_reference_configs_up_to_date passes. - Trim setup.py block comment to a single line per reviewer feedback. - Fix test docstring to reflect that .get() defaults to None. Signed-off-by: seonjinn --- examples/configs/distillation_math.yaml | 9 +++++++++ examples/configs/distillation_math_megatron.yaml | 9 +++++++++ examples/configs/dpo.yaml | 9 +++++++++ examples/configs/grpo_math_1B.yaml | 9 +++++++++ examples/configs/grpo_math_1B_megatron.yaml | 9 +++++++++ examples/configs/sft.yaml | 9 +++++++++ nemo_rl/models/megatron/setup.py | 6 +----- tests/unit/models/megatron/test_megatron_setup.py | 2 +- tests/unit/reference_configs/distillation_math.yaml | 2 ++ tests/unit/reference_configs/dpo.yaml | 2 ++ tests/unit/reference_configs/grpo_math_1B.yaml | 2 ++ tests/unit/reference_configs/sft.yaml | 2 ++ 12 files changed, 64 insertions(+), 6 deletions(-) diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml index ab88186661e..a1e66944673 100644 --- a/examples/configs/distillation_math.yaml +++ b/examples/configs/distillation_math.yaml @@ -90,6 +90,15 @@ 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 + # Offload specific module activations to CPU. Works for both dense and MoE + # models. Requires transformer_engine. Different from optimizer_cpu_offload + # which offloads optimizer states. + fine_grained_activation_offloading: false + # Modules to offload when fine_grained_activation_offloading is true. + # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", + # "mlp_norm", "expert_fc1", "moe_act"]. "attn_proj" requires "core_attn". + # See: https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L1440-L1448 + offload_modules: null converter_type: "Qwen3ForCausalLM" tensor_model_parallel_size: 2 expert_tensor_parallel_size: 1 diff --git a/examples/configs/distillation_math_megatron.yaml b/examples/configs/distillation_math_megatron.yaml index 5c0eb71a909..d2e56aab6fa 100644 --- a/examples/configs/distillation_math_megatron.yaml +++ b/examples/configs/distillation_math_megatron.yaml @@ -39,6 +39,15 @@ policy: &POLICY_BASE enabled: true empty_unused_memory_level: 0 activation_checkpointing: false + # Offload specific module activations to CPU. Works for both dense and MoE + # models. Requires transformer_engine. Different from optimizer_cpu_offload + # which offloads optimizer states. + fine_grained_activation_offloading: false + # Modules to offload when fine_grained_activation_offloading is true. + # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", + # "mlp_norm", "expert_fc1", "moe_act"]. "attn_proj" requires "core_attn". + # See: https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L1440-L1448 + offload_modules: null converter_type: "Qwen3ForCausalLM" tensor_model_parallel_size: 2 expert_tensor_parallel_size: 1 diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml index 17a49669e0e..d68bcbcd40d 100755 --- a/examples/configs/dpo.yaml +++ b/examples/configs/dpo.yaml @@ -129,6 +129,15 @@ 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 + # Offload specific module activations to CPU. Works for both dense and MoE + # models. Requires transformer_engine. Different from optimizer_cpu_offload + # which offloads optimizer states. + fine_grained_activation_offloading: false + # Modules to offload when fine_grained_activation_offloading is true. + # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", + # "mlp_norm", "expert_fc1", "moe_act"]. "attn_proj" requires "core_attn". + # See: https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L1440-L1448 + offload_modules: null tensor_model_parallel_size: 2 expert_tensor_parallel_size: 1 expert_model_parallel_size: 1 diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 4e2b8241f2b..66e3b61d4ed 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -134,6 +134,15 @@ 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 + # Offload specific module activations to CPU. Works for both dense and MoE + # models. Requires transformer_engine. Different from optimizer_cpu_offload + # which offloads optimizer states. + fine_grained_activation_offloading: false + # Modules to offload when fine_grained_activation_offloading is true. + # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", + # "mlp_norm", "expert_fc1", "moe_act"]. "attn_proj" requires "core_attn". + # See: https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L1440-L1448 + offload_modules: null converter_type: "Qwen2ForCausalLM" tensor_model_parallel_size: 1 expert_tensor_parallel_size: 1 diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index 084d6621304..68a7f5ea5c4 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/examples/configs/grpo_math_1B_megatron.yaml @@ -86,6 +86,15 @@ 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 + # Offload specific module activations to CPU. Works for both dense and MoE + # models. Requires transformer_engine. Different from optimizer_cpu_offload + # which offloads optimizer states. + fine_grained_activation_offloading: false + # Modules to offload when fine_grained_activation_offloading is true. + # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", + # "mlp_norm", "expert_fc1", "moe_act"]. "attn_proj" requires "core_attn". + # See: https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L1440-L1448 + offload_modules: null converter_type: "Qwen2ForCausalLM" tensor_model_parallel_size: 1 expert_tensor_parallel_size: 1 diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml index cf02bdfc74e..1ff1e18debc 100644 --- a/examples/configs/sft.yaml +++ b/examples/configs/sft.yaml @@ -110,6 +110,15 @@ policy: env_vars: {} empty_unused_memory_level: 1 activation_checkpointing: false + # Offload specific module activations to CPU. Works for both dense and MoE + # models. Requires transformer_engine. Different from optimizer_cpu_offload + # which offloads optimizer states. + fine_grained_activation_offloading: false + # Modules to offload when fine_grained_activation_offloading is true. + # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", + # "mlp_norm", "expert_fc1", "moe_act"]. "attn_proj" requires "core_attn". + # See: https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L1440-L1448 + offload_modules: null tensor_model_parallel_size: 1 expert_tensor_parallel_size: 1 expert_model_parallel_size: 1 diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 5e6bfe6d9f9..d3e69c3ac23 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -718,11 +718,7 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: "Refer to https://github.com/NVIDIA-NeMo/RL/issues/1164 for latest updates with this issue." ) - # Fine-grained activation offloading moves specified submodule activations - # to CPU. Works for both dense and MoE models; the user picks which - # submodules to offload via offload_modules. Megatron owns the list of - # valid module names and their per-model-type compatibility, so we only - # require a non-empty list here and let Megatron validate the contents. + # Megatron validates module names and per-model-type compatibility. fine_grained_activation_offloading = config["megatron_cfg"].get( "fine_grained_activation_offloading" ) diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 617428fd116..5ee070c74d7 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -1078,7 +1078,7 @@ def test_fine_grained_activation_offloading_invalid_modules_raises( _apply_performance_config(model_cfg, config) def test_fine_grained_activation_offloading_missing_modules_raises(self): - """When enabled but offload_modules key is absent, defaults to [] → raises.""" + """When enabled but offload_modules key is absent, defaults to None → raises.""" from nemo_rl.models.megatron.setup import _apply_performance_config model_cfg = MagicMock() diff --git a/tests/unit/reference_configs/distillation_math.yaml b/tests/unit/reference_configs/distillation_math.yaml index ab88186661e..1ff9917b518 100644 --- a/tests/unit/reference_configs/distillation_math.yaml +++ b/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 + fine_grained_activation_offloading: false + offload_modules: null converter_type: "Qwen3ForCausalLM" tensor_model_parallel_size: 2 expert_tensor_parallel_size: 1 diff --git a/tests/unit/reference_configs/dpo.yaml b/tests/unit/reference_configs/dpo.yaml index 415512addb0..b2a930dc020 100755 --- a/tests/unit/reference_configs/dpo.yaml +++ b/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 + fine_grained_activation_offloading: false + offload_modules: null tensor_model_parallel_size: 2 expert_tensor_parallel_size: 1 expert_model_parallel_size: 1 diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index b3bf195c1c0..2167e0b1fa7 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/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 + fine_grained_activation_offloading: false + offload_modules: null converter_type: "Qwen2ForCausalLM" tensor_model_parallel_size: 1 expert_tensor_parallel_size: 1 diff --git a/tests/unit/reference_configs/sft.yaml b/tests/unit/reference_configs/sft.yaml index 416895b6350..6f03e7cbd63 100644 --- a/tests/unit/reference_configs/sft.yaml +++ b/tests/unit/reference_configs/sft.yaml @@ -105,6 +105,8 @@ policy: env_vars: {} empty_unused_memory_level: 1 activation_checkpointing: false + fine_grained_activation_offloading: false + offload_modules: null tensor_model_parallel_size: 1 expert_tensor_parallel_size: 1 expert_model_parallel_size: 1 From d89ac84382a652147cc9deb6090f96c8518df80c Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 23 May 2026 00:13:02 -0700 Subject: [PATCH 12/24] Address remaining terrykong review comments - Add NVTE_CPU_OFFLOAD_V1=1 note (TE >= 2.10.0) to TypedDict comment in policy/__init__.py so users see the env requirement up front rather than via a late Megatron-Bridge validation error. - Document the NUMA affinity gap in megatron/setup.py: Megatron-Bridge's standalone path calls set_ideal_affinity_for_current_gpu() when this feature is on; NeMo-RL does not, so this comment points users who care about offload bandwidth at the external workaround. Signed-off-by: seonjinn --- nemo_rl/models/megatron/setup.py | 4 ++++ nemo_rl/models/policy/__init__.py | 4 +++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index d3e69c3ac23..29266a5b4e0 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -719,6 +719,10 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: ) # Megatron validates module names and per-model-type compatibility. + # Note: Megatron-Bridge's standalone training path also sets NUMA-aware + # CPU affinity via set_ideal_affinity_for_current_gpu() when this is on, + # which improves PCIe/DRAM throughput. NeMo-RL does not call it; users + # who need maximum offload bandwidth may want to set affinity externally. fine_grained_activation_offloading = config["megatron_cfg"].get( "fine_grained_activation_offloading" ) diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 3daabb7d455..64c49dd0c0a 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -240,7 +240,9 @@ class MegatronConfig(TypedDict): # Offload specific module activations to CPU to reduce peak GPU memory. # Works with both dense and MoE models. Different from # optimizer_cpu_offload which offloads optimizer states. - # Requires transformer_engine implementation. + # Requires transformer_engine. For TE >= 2.10.0 also requires + # NVTE_CPU_OFFLOAD_V1=1 in the environment (validated by + # Megatron-Bridge at runtime). fine_grained_activation_offloading: NotRequired[bool] # Modules to offload when fine_grained_activation_offloading is True. # Required (no default). Valid values: From aaece7b131786e62671ddae8e1c4598776357e16 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 23 May 2026 00:39:08 -0700 Subject: [PATCH 13/24] Allow offload_modules to be None in MegatronConfig TypedDict NotRequired[list[str]] caused pydantic to reject `offload_modules: null` in YAML, breaking L1 run_vlm_grpo (and any recipe loading an exemplar megatron config when the feature is off). Wrap with Optional so the exemplar default `null` is accepted; the runtime validation in setup.py already raises if the feature is on with a non-list / empty value. Signed-off-by: seonjinn --- nemo_rl/models/policy/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 64c49dd0c0a..b336703c86f 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Literal, NotRequired, TypedDict, Union +from typing import Any, Literal, NotRequired, Optional, TypedDict, Union from nemo_rl.models.generation.interfaces import GenerationConfig from nemo_rl.utils.checkpoint import PretrainedCheckpointConfig @@ -249,7 +249,7 @@ class MegatronConfig(TypedDict): # "attn_norm", "qkv_linear", "core_attn", "attn_proj", "mlp_norm", # "expert_fc1", "moe_act". Note: "attn_proj" requires "core_attn". # See: https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L1440-L1448 - offload_modules: NotRequired[list[str]] + offload_modules: NotRequired[Optional[list[str]]] # Enable grouped GEMM for MoE experts via CUTLASS. Significant throughput # gain when multiple experts are assigned per rank (num_local_experts > 1). # Requires TE >= 1.11.0 for FP8 and Ampere (sm_80) or newer. From 7051195fe0c827efd481b0a0cc470eabc80c16be Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 16:11:11 -0700 Subject: [PATCH 14/24] fix(megatron): honor activation offload overrides Signed-off-by: seonjinn --- nemo_rl/models/megatron/setup.py | 32 ++++++++-- .../models/megatron/test_megatron_setup.py | 58 +++++++++++++++++++ 2 files changed, 84 insertions(+), 6 deletions(-) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index ae767f7d07f..33252a5dc8c 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1103,17 +1103,37 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: except KeyError as e: raise KeyError(f"Missing key in fp8_cfg: {e}") - # Megatron validates module names and per-model-type compatibility. - fine_grained_activation_offloading = config["megatron_cfg"].get( - "fine_grained_activation_offloading" - ) - if fine_grained_activation_offloading: - offload_modules = config["megatron_cfg"].get("offload_modules") + megatron_cfg = config["megatron_cfg"] + if "fine_grained_activation_offloading" in megatron_cfg: + fine_grained_activation_offloading = megatron_cfg[ + "fine_grained_activation_offloading" + ] + else: + fine_grained_activation_offloading = None + + if fine_grained_activation_offloading is False: + # Preserve the legacy exemplar's disabled/null semantics and clear any + # enabled state carried by a provider or checkpoint. + model_cfg.fine_grained_activation_offloading = False + model_cfg.offload_modules = None + elif fine_grained_activation_offloading: + offload_modules = megatron_cfg.get("offload_modules") if not isinstance(offload_modules, list) or not offload_modules: raise ValueError( "offload_modules must be a non-empty list when " "fine_grained_activation_offloading is True." ) + moe_only_modules = {"expert_fc1", "moe_act", "fused_group_mlp"} + invalid_dense_modules = moe_only_modules.intersection(offload_modules) + if ( + invalid_dense_modules + and getattr(model_cfg, "num_moe_experts", None) is None + ): + raise ValueError( + "A MoE-only offload module requires a MoE model " + "(num_moe_experts must not be None): " + f"{sorted(invalid_dense_modules)}." + ) model_cfg.fine_grained_activation_offloading = True model_cfg.offload_modules = offload_modules diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 56fb0892791..fd11e1c2742 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -1247,6 +1247,44 @@ def test_fine_grained_activation_offloading_disabled_skips(self): assert not hasattr(model_cfg, "fine_grained_activation_offloading") assert not hasattr(model_cfg, "offload_modules") + def test_missing_offloading_flag_preserves_provider_values(self): + """An omitted setting does not overwrite the provider's offload configuration.""" + from nemo_rl.models.megatron.setup import _apply_performance_config + + offload_modules = ["core_attn"] + model_cfg = SimpleNamespace( + gated_linear_unit=True, + fine_grained_activation_offloading=True, + offload_modules=offload_modules, + ) + + _apply_performance_config(model_cfg, self._config()) + + assert model_cfg.fine_grained_activation_offloading is True + assert model_cfg.offload_modules == offload_modules + + def test_explicitly_disabled_offloading_clears_provider_values(self): + """An explicit false overrides enabled provider values from a checkpoint.""" + from nemo_rl.models.megatron.setup import _apply_performance_config + + config = self._config() + config["megatron_cfg"].update( + { + "fine_grained_activation_offloading": False, + "offload_modules": None, + } + ) + model_cfg = SimpleNamespace( + gated_linear_unit=True, + fine_grained_activation_offloading=True, + offload_modules=["core_attn"], + ) + + _apply_performance_config(model_cfg, config) + + assert model_cfg.fine_grained_activation_offloading is False + assert model_cfg.offload_modules is None + @pytest.mark.parametrize( "offload_modules", [[], None, "moe_act", 42], @@ -1299,6 +1337,26 @@ def test_fine_grained_activation_offloading_missing_modules_raises(self): ): _apply_performance_config(model_cfg, config) + @pytest.mark.parametrize( + "offload_module", + ["expert_fc1", "moe_act", "fused_group_mlp"], + ) + def test_moe_only_offload_module_rejected_for_dense_model(self, offload_module): + """MoE-only offload modules cannot silently no-op for dense models.""" + from nemo_rl.models.megatron.setup import _apply_performance_config + + config = self._config() + config["megatron_cfg"].update( + { + "fine_grained_activation_offloading": True, + "offload_modules": [offload_module], + } + ) + model_cfg = SimpleNamespace(gated_linear_unit=True, num_moe_experts=None) + + with pytest.raises(ValueError, match="requires a MoE model"): + _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 From 602e7b0a1c9b68b4261d59b8ad40322dc8ed05c6 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 16:12:18 -0700 Subject: [PATCH 15/24] docs: clarify Megatron activation CPU offload Signed-off-by: seonjinn --- docs/design-docs/training-backends.md | 36 +++++++++++++++++++ examples/configs/ppo_math_1B.yaml | 8 +++++ .../ppo_math_1B_megatron.yaml | 4 +++ 3 files changed, 48 insertions(+) diff --git a/docs/design-docs/training-backends.md b/docs/design-docs/training-backends.md index 4cb0eca8a9b..d3cce048d56 100644 --- a/docs/design-docs/training-backends.md +++ b/docs/design-docs/training-backends.md @@ -57,6 +57,42 @@ dictionaries follow the hierarchy of nested model-config objects. NeMo RL-specific settings such as optimizer, scheduler, checkpointing, and environment variables remain under their existing `megatron_cfg` sections. +#### Fine-grained activation CPU offload + +Fine-grained activation offloading asynchronously moves selected module-input +activations to CPU between forward and backward passes to reduce peak GPU +memory. It is distinct from `optimizer_cpu_offload`, which moves optimizer +states rather than activations. The following dense-model configuration is +runnable with the Megatron backend; `core_attn` and `attn_proj` are appropriate +for a dense Qwen model, and `attn_proj` must be paired with `core_attn`. + +```yaml +policy: + megatron_cfg: + enabled: true + env_vars: + NVTE_CPU_OFFLOAD_V1: "1" + fine_grained_activation_offloading: true + offload_modules: ["core_attn", "attn_proj"] +``` + +Supported module names are `attn_norm`, `qkv_linear`, `core_attn`, +`attn_proj`, `mlp_norm`, `expert_fc1`, `moe_act`, and `fused_group_mlp`. +The last three are MoE-specific. `fused_group_mlp` requires the Transformer +Engine op fuser and cannot be combined with `expert_fc1` or `moe_act`. + +Activation checkpointing is not blanket-incompatible with fine-grained +activation offload. However, selective recomputation of the whole MoE module +(`recompute_modules: ["moe"]`) conflicts with MoE-internal offload modules +(`expert_fc1`, `moe_act`, or `fused_group_mlp`), and layer-level +`cpu_offloading` conflicts with fine-grained activation offload. Megatron +Bridge/Megatron-Core setup validation owns the exact compatibility checks for +the pinned versions. + +Offloading saves GPU memory but adds CPU transfer and synchronization work, so +benchmark the memory and throughput tradeoff for the target model, sequence +length, hardware, and parallelism configuration. + ### DTensor Backend To enable DTensor (FSDP2) training: diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 366cf08cf31..78fabe5c9c9 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -128,6 +128,10 @@ policy: enabled: false empty_unused_memory_level: 1 activation_checkpointing: false + # Module-level activation CPU offload; see docs/design-docs/training-backends.md. + fine_grained_activation_offloading: false + # Modules to offload when enabled; null disables module selection. + offload_modules: null converter_type: "Qwen2ForCausalLM" tensor_model_parallel_size: 1 expert_tensor_parallel_size: 1 @@ -317,6 +321,10 @@ value: enabled: false empty_unused_memory_level: 1 activation_checkpointing: false + # Module-level activation CPU offload; see docs/design-docs/training-backends.md. + fine_grained_activation_offloading: false + # Modules to offload when enabled; null disables module selection. + offload_modules: null tensor_model_parallel_size: 1 expert_tensor_parallel_size: 1 expert_model_parallel_size: 1 diff --git a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml index 6af4416c2f5..9858962d779 100644 --- a/tests/unit/reference_configs/ppo_math_1B_megatron.yaml +++ b/tests/unit/reference_configs/ppo_math_1B_megatron.yaml @@ -130,6 +130,8 @@ policy: ckpt_assume_constant_structure: false empty_unused_memory_level: 1 activation_checkpointing: false + fine_grained_activation_offloading: false + offload_modules: null converter_type: "Qwen2ForCausalLM" tensor_model_parallel_size: 1 expert_tensor_parallel_size: 1 @@ -306,6 +308,8 @@ value: ckpt_assume_constant_structure: false empty_unused_memory_level: 1 activation_checkpointing: false + fine_grained_activation_offloading: false + offload_modules: null tensor_model_parallel_size: 1 expert_tensor_parallel_size: 1 expert_model_parallel_size: 1 From 5ddbde866f2cc4367d07578d2a631eca77bf60f1 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 16:15:52 -0700 Subject: [PATCH 16/24] docs: correct activation offload graph requirement Signed-off-by: seonjinn --- docs/design-docs/training-backends.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/design-docs/training-backends.md b/docs/design-docs/training-backends.md index d3cce048d56..f614d71d682 100644 --- a/docs/design-docs/training-backends.md +++ b/docs/design-docs/training-backends.md @@ -70,12 +70,16 @@ for a dense Qwen model, and `attn_proj` must be paired with `core_attn`. policy: megatron_cfg: enabled: true + cuda_graph_impl: transformer_engine env_vars: NVTE_CPU_OFFLOAD_V1: "1" fine_grained_activation_offloading: true offload_modules: ["core_attn", "attn_proj"] ``` +Pinned MCore requires the `transformer_engine` CUDA-graph implementation for +this dense module pair; `local` CUDA graphs support only partial MoE offload. + Supported module names are `attn_norm`, `qkv_linear`, `core_attn`, `attn_proj`, `mlp_norm`, `expert_fc1`, `moe_act`, and `fused_group_mlp`. The last three are MoE-specific. `fused_group_mlp` requires the Transformer From 2f39df66d6fd0a0b1b53cf472eb1599c6c05dfce Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 16:17:41 -0700 Subject: [PATCH 17/24] test(megatron): model activation offload happy path Signed-off-by: seonjinn --- tests/unit/models/megatron/test_megatron_setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index fd11e1c2742..4a2ec20b66a 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -1209,6 +1209,7 @@ def test_fine_grained_activation_offloading_enabled(self): model_cfg = MagicMock() model_cfg.gated_linear_unit = True + model_cfg.num_moe_experts = 8 offload_modules = ["mlp_norm", "moe_act"] config = { "megatron_cfg": { From bf3407fdd6dfdf82eff76886836f10e4ff1d4002 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 18:16:36 -0700 Subject: [PATCH 18/24] test(megatron): preserve activation offload warmup Signed-off-by: seonjinn --- tests/unit/models/megatron/test_train.py | 100 +++++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py index c13ffa36c86..5c515af415c 100644 --- a/tests/unit/models/megatron/test_train.py +++ b/tests/unit/models/megatron/test_train.py @@ -769,6 +769,106 @@ def test_megatron_forward_backward_forward_only(self, mock_get_fb): call_kwargs = mock_fb_func.call_args[1] assert call_kwargs["forward_only"] is True + @patch("nemo_rl.models.megatron.train.get_forward_backward_func") + def test_forward_only_preserves_activation_offload_warmup(self, mock_get_fb): + """Forward-only RL stages must not consume activation-offload warmup.""" + from nemo_rl.models.megatron.train import ( + LossPostProcessor, + megatron_forward_backward, + ) + + model_config = SimpleNamespace(fine_grained_activation_offloading=True) + model = SimpleNamespace(config=model_config) + + def run_forward_only(**kwargs): + assert kwargs["forward_only"] is True + assert model_config.fine_grained_activation_offloading is False + return {"logprobs": torch.tensor(0.0)} + + mock_get_fb.return_value = run_forward_only + post_processor = LossPostProcessor( + loss_fn=MagicMock(), cfg={"sequence_packing": {"enabled": False}} + ) + + megatron_forward_backward( + model=model, + data_iterator=iter([]), + num_microbatches=1, + seq_length=64, + mbs=1, + post_processing_fn=post_processor, + forward_only=True, + ) + + assert model_config.fine_grained_activation_offloading is True + + @patch("nemo_rl.models.megatron.train.get_forward_backward_func") + def test_forward_only_restores_activation_offload_after_error(self, mock_get_fb): + """A failed forward-only stage must restore activation offload for training.""" + from nemo_rl.models.megatron.train import ( + LossPostProcessor, + megatron_forward_backward, + ) + + model_config = SimpleNamespace(fine_grained_activation_offloading=True) + model = SimpleNamespace(config=model_config) + + def fail_forward_only(**kwargs): + assert kwargs["forward_only"] is True + assert model_config.fine_grained_activation_offloading is False + raise RuntimeError("forward-only failure") + + mock_get_fb.return_value = fail_forward_only + post_processor = LossPostProcessor( + loss_fn=MagicMock(), cfg={"sequence_packing": {"enabled": False}} + ) + + with pytest.raises(RuntimeError, match="forward-only failure"): + megatron_forward_backward( + model=model, + data_iterator=iter([]), + num_microbatches=1, + seq_length=64, + mbs=1, + post_processing_fn=post_processor, + forward_only=True, + ) + + assert model_config.fine_grained_activation_offloading is True + + @patch("nemo_rl.models.megatron.train.get_forward_backward_func") + def test_training_keeps_activation_offload_enabled(self, mock_get_fb): + """Training must retain activation offload so MCore can warm up and run it.""" + from nemo_rl.models.megatron.train import ( + LossPostProcessor, + megatron_forward_backward, + ) + + model_config = SimpleNamespace(fine_grained_activation_offloading=True) + model = SimpleNamespace(config=model_config) + + def run_training(**kwargs): + assert kwargs["forward_only"] is False + assert model_config.fine_grained_activation_offloading is True + return {"loss": torch.tensor(0.5)} + + mock_get_fb.return_value = run_training + post_processor = LossPostProcessor( + loss_fn=MagicMock(), cfg={"sequence_packing": {"enabled": False}} + ) + + megatron_forward_backward( + model=model, + data_iterator=iter([]), + num_microbatches=1, + seq_length=64, + mbs=1, + post_processing_fn=post_processor, + forward_only=False, + ) + + assert model_config.fine_grained_activation_offloading is True + class TestLossPostProcessor: """Tests for LossPostProcessor class.""" From 3a3bebdca537fd44910f04d90d6e67134baf085f Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 18:24:22 -0700 Subject: [PATCH 19/24] fix(megatron): preserve activation offload warmup Signed-off-by: seonjinn --- nemo_rl/models/megatron/train.py | 64 ++++++++++++++++++++++++-------- 1 file changed, 48 insertions(+), 16 deletions(-) diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index beb02a89ad0..2320904bb8d 100644 --- a/nemo_rl/models/megatron/train.py +++ b/nemo_rl/models/megatron/train.py @@ -13,7 +13,7 @@ # limitations under the License. from collections import defaultdict -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from functools import partial from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union @@ -27,7 +27,7 @@ get_tensor_model_parallel_rank, ) from megatron.core.pipeline_parallel import get_forward_backward_func -from megatron.core.utils import StragglerDetector +from megatron.core.utils import StragglerDetector, get_model_config from nemo_rl.algorithms.logits_sampling_utils import ( TrainingSamplingParams, @@ -70,6 +70,37 @@ ] +@contextmanager +def _suspend_activation_offload_for_forward_only( + model: Union[GPTModel, List[GPTModel]], forward_only: bool +) -> Iterator[None]: + """Keep inference-only RL phases from consuming MCore's training warmup.""" + if not forward_only: + yield + return + + model_chunks = model if isinstance(model, list) else [model] + original_values: List[Tuple[Any, bool]] = [] + seen_configs: set[int] = set() + for model_chunk in model_chunks: + model_config = get_model_config(model_chunk) + if id(model_config) in seen_configs: + continue + seen_configs.add(id(model_config)) + original_value = bool( + getattr(model_config, "fine_grained_activation_offloading", False) + ) + if original_value: + original_values.append((model_config, original_value)) + model_config.fine_grained_activation_offloading = False + + try: + yield + finally: + for model_config, original_value in original_values: + model_config.fine_grained_activation_offloading = original_value + + def model_forward( model: GPTModel, data_dict: BatchedDataDict[Any], @@ -357,20 +388,21 @@ def megatron_forward_backward( forward_backward_func = get_forward_backward_func() if use_router_replay: clear_router_replay(model) - try: - return forward_backward_func( - forward_step_func=forward_step, - data_iterator=data_iterator, - model=model, - num_microbatches=num_microbatches, - seq_length=seq_length, - micro_batch_size=mbs, - decoder_seq_length=seq_length, - forward_only=forward_only, - ) - finally: - if use_router_replay: - clear_router_replay(model) + with _suspend_activation_offload_for_forward_only(model, forward_only): + try: + return forward_backward_func( + forward_step_func=forward_step, + data_iterator=data_iterator, + model=model, + num_microbatches=num_microbatches, + seq_length=seq_length, + micro_batch_size=mbs, + decoder_seq_length=seq_length, + forward_only=forward_only, + ) + finally: + if use_router_replay: + clear_router_replay(model) class LossPostProcessor: From 9c79efea40a6f1f5ff6b87b355cc1c5ec1f08898 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 18:38:48 -0700 Subject: [PATCH 20/24] test(megatron): cover offload manager lifecycle Signed-off-by: seonjinn --- tests/unit/models/megatron/test_train.py | 217 +++++++++++++++++++++-- 1 file changed, 206 insertions(+), 11 deletions(-) diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py index 5c515af415c..1b011a5e2b6 100644 --- a/tests/unit/models/megatron/test_train.py +++ b/tests/unit/models/megatron/test_train.py @@ -24,6 +24,7 @@ from contextlib import nullcontext from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock, patch import pytest @@ -770,8 +771,14 @@ def test_megatron_forward_backward_forward_only(self, mock_get_fb): assert call_kwargs["forward_only"] is True @patch("nemo_rl.models.megatron.train.get_forward_backward_func") - def test_forward_only_preserves_activation_offload_warmup(self, mock_get_fb): + def test_forward_only_preserves_activation_offload_warmup( + self, mock_get_fb: MagicMock + ) -> None: """Forward-only RL stages must not consume activation-offload warmup.""" + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + from nemo_rl.models.megatron.train import ( LossPostProcessor, megatron_forward_backward, @@ -780,9 +787,10 @@ def test_forward_only_preserves_activation_offload_warmup(self, mock_get_fb): model_config = SimpleNamespace(fine_grained_activation_offloading=True) model = SimpleNamespace(config=model_config) - def run_forward_only(**kwargs): + def run_forward_only(**kwargs: Any) -> dict[str, torch.Tensor]: assert kwargs["forward_only"] is True assert model_config.fine_grained_activation_offloading is False + assert PipelineOffloadManager.OFFLOAD_MGR is None return {"logprobs": torch.tensor(0.0)} mock_get_fb.return_value = run_forward_only @@ -790,18 +798,205 @@ def run_forward_only(**kwargs): loss_fn=MagicMock(), cfg={"sequence_packing": {"enabled": False}} ) - megatron_forward_backward( - model=model, - data_iterator=iter([]), - num_microbatches=1, - seq_length=64, - mbs=1, - post_processing_fn=post_processor, - forward_only=True, - ) + with patch.object(PipelineOffloadManager, "OFFLOAD_MGR", None): + megatron_forward_backward( + model=model, + data_iterator=iter([]), + num_microbatches=1, + seq_length=64, + mbs=1, + post_processing_fn=post_processor, + forward_only=True, + ) + + assert PipelineOffloadManager.OFFLOAD_MGR is None assert model_config.fine_grained_activation_offloading is True + @patch("nemo_rl.models.megatron.train.get_forward_backward_func") + def test_forward_only_does_not_consume_warm_manager_chunks( + self, mock_get_fb: MagicMock + ) -> None: + """Logprob stages must not advance chunks cached by a prior training step.""" + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + from nemo_rl.models.megatron.train import ( + LossPostProcessor, + megatron_forward_backward, + ) + + class StatefulOffloadManager: + def __init__(self) -> None: + self.do_offload = True + self.cached_forward_index = 0 + + def disable_offload(self) -> None: + self.do_offload = False + + def enable_offload(self) -> None: + self.do_offload = True + + def consume_cached_chunk(self) -> None: + if self.do_offload: + self.cached_forward_index += 1 + + def reset(self) -> None: + self.cached_forward_index = 0 + + manager = StatefulOffloadManager() + model_config = SimpleNamespace(fine_grained_activation_offloading=True) + model = SimpleNamespace(config=model_config) + observed_phases: list[tuple[bool, int]] = [] + + def run_schedule(**kwargs: Any) -> dict[str, torch.Tensor]: + manager.consume_cached_chunk() + observed_phases.append( + (kwargs["forward_only"], manager.cached_forward_index) + ) + if not kwargs["forward_only"]: + manager.reset() + return {} + + mock_get_fb.return_value = run_schedule + post_processor = LossPostProcessor( + loss_fn=MagicMock(), cfg={"sequence_packing": {"enabled": False}} + ) + + with patch.object(PipelineOffloadManager, "OFFLOAD_MGR", manager): + for forward_only in (True, False, True, False): + megatron_forward_backward( + model=model, + data_iterator=iter([]), + num_microbatches=1, + seq_length=64, + mbs=1, + post_processing_fn=post_processor, + forward_only=forward_only, + ) + + assert manager.do_offload is True + + assert observed_phases == [(True, 0), (False, 1), (True, 0), (False, 1)] + + @patch("nemo_rl.models.megatron.train.get_forward_backward_func") + def test_forward_only_preserves_disabled_manager_state( + self, mock_get_fb: MagicMock + ) -> None: + """Nested callers that disabled offload must remain disabled.""" + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + from nemo_rl.models.megatron.train import ( + LossPostProcessor, + megatron_forward_backward, + ) + + manager = MagicMock() + manager.do_offload = False + model_config = SimpleNamespace(fine_grained_activation_offloading=True) + + def run_forward_only(**kwargs: Any) -> dict[str, torch.Tensor]: + assert kwargs["forward_only"] is True + assert manager.do_offload is False + return {} + + mock_get_fb.return_value = run_forward_only + post_processor = LossPostProcessor( + loss_fn=MagicMock(), cfg={"sequence_packing": {"enabled": False}} + ) + + with patch.object(PipelineOffloadManager, "OFFLOAD_MGR", manager): + megatron_forward_backward( + model=SimpleNamespace(config=model_config), + data_iterator=iter([]), + num_microbatches=1, + seq_length=64, + mbs=1, + post_processing_fn=post_processor, + forward_only=True, + ) + + manager.disable_offload.assert_not_called() + manager.enable_offload.assert_not_called() + assert manager.do_offload is False + + @patch("nemo_rl.models.megatron.train.get_forward_backward_func") + def test_forward_only_restores_vpp_configs(self, mock_get_fb: MagicMock) -> None: + """Shared and distinct VPP configs are suspended and restored atomically.""" + from nemo_rl.models.megatron.train import ( + LossPostProcessor, + megatron_forward_backward, + ) + + shared_config = SimpleNamespace(fine_grained_activation_offloading=True) + distinct_config = SimpleNamespace(fine_grained_activation_offloading=True) + model = [MagicMock(), MagicMock(), MagicMock()] + + def run_forward_only(**kwargs: Any) -> dict[str, torch.Tensor]: + assert shared_config.fine_grained_activation_offloading is False + assert distinct_config.fine_grained_activation_offloading is False + return {} + + mock_get_fb.return_value = run_forward_only + post_processor = LossPostProcessor( + loss_fn=MagicMock(), cfg={"sequence_packing": {"enabled": False}} + ) + + with patch( + "nemo_rl.models.megatron.train.get_model_config", + side_effect=[shared_config, shared_config, distinct_config], + ): + megatron_forward_backward( + model=model, + data_iterator=iter([]), + num_microbatches=1, + seq_length=64, + mbs=1, + post_processing_fn=post_processor, + forward_only=True, + ) + + assert shared_config.fine_grained_activation_offloading is True + assert distinct_config.fine_grained_activation_offloading is True + + @patch("nemo_rl.models.megatron.train.get_forward_backward_func") + def test_forward_only_config_discovery_is_atomic( + self, mock_get_fb: MagicMock + ) -> None: + """A later VPP wrapper error must not leave earlier configs disabled.""" + from nemo_rl.models.megatron.train import ( + LossPostProcessor, + megatron_forward_backward, + ) + + first_config = SimpleNamespace(fine_grained_activation_offloading=True) + post_processor = LossPostProcessor( + loss_fn=MagicMock(), cfg={"sequence_packing": {"enabled": False}} + ) + + with ( + patch( + "nemo_rl.models.megatron.train.get_model_config", + side_effect=[first_config, RuntimeError("invalid VPP wrapper")], + ), + pytest.raises(RuntimeError, match="invalid VPP wrapper"), + ): + megatron_forward_backward( + model=[MagicMock(), MagicMock()], + data_iterator=iter([]), + num_microbatches=1, + seq_length=64, + mbs=1, + post_processing_fn=post_processor, + forward_only=True, + ) + + mock_get_fb.return_value.assert_not_called() + assert first_config.fine_grained_activation_offloading is True + @patch("nemo_rl.models.megatron.train.get_forward_backward_func") def test_forward_only_restores_activation_offload_after_error(self, mock_get_fb): """A failed forward-only stage must restore activation offload for training.""" From aeeb3d6f5d5848893c12827c874747b2e6ad3551 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 18:43:32 -0700 Subject: [PATCH 21/24] fix(megatron): suspend activation offload manager Signed-off-by: seonjinn --- nemo_rl/models/megatron/train.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index 2320904bb8d..6302b31b28a 100644 --- a/nemo_rl/models/megatron/train.py +++ b/nemo_rl/models/megatron/train.py @@ -27,6 +27,9 @@ get_tensor_model_parallel_rank, ) from megatron.core.pipeline_parallel import get_forward_backward_func +from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, +) from megatron.core.utils import StragglerDetector, get_model_config from nemo_rl.algorithms.logits_sampling_utils import ( @@ -92,13 +95,25 @@ def _suspend_activation_offload_for_forward_only( ) if original_value: original_values.append((model_config, original_value)) - model_config.fine_grained_activation_offloading = False + + offload_manager = PipelineOffloadManager.OFFLOAD_MGR + suspend_manager = bool( + original_values and offload_manager is not None and offload_manager.do_offload + ) try: + for model_config, _ in original_values: + model_config.fine_grained_activation_offloading = False + if suspend_manager and offload_manager is not None: + offload_manager.disable_offload() yield finally: - for model_config, original_value in original_values: - model_config.fine_grained_activation_offloading = original_value + try: + if suspend_manager and offload_manager is not None: + offload_manager.enable_offload() + finally: + for model_config, original_value in original_values: + model_config.fine_grained_activation_offloading = original_value def model_forward( From 01398467224921c058a70702cb4a8285eb98fc71 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Wed, 12 Aug 2026 18:48:46 -0700 Subject: [PATCH 22/24] test(megatron): assert offload warmup state Signed-off-by: seonjinn --- tests/unit/models/megatron/test_train.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py index 1b011a5e2b6..130cbf9825b 100644 --- a/tests/unit/models/megatron/test_train.py +++ b/tests/unit/models/megatron/test_train.py @@ -786,11 +786,15 @@ def test_forward_only_preserves_activation_offload_warmup( model_config = SimpleNamespace(fine_grained_activation_offloading=True) model = SimpleNamespace(config=model_config) + empty_manager = SimpleNamespace( + _is_warmup=True, + _cached_chunks_forward=[], + ) def run_forward_only(**kwargs: Any) -> dict[str, torch.Tensor]: assert kwargs["forward_only"] is True assert model_config.fine_grained_activation_offloading is False - assert PipelineOffloadManager.OFFLOAD_MGR is None + PipelineOffloadManager.OFFLOAD_MGR = empty_manager return {"logprobs": torch.tensor(0.0)} mock_get_fb.return_value = run_forward_only @@ -809,7 +813,9 @@ def run_forward_only(**kwargs: Any) -> dict[str, torch.Tensor]: forward_only=True, ) - assert PipelineOffloadManager.OFFLOAD_MGR is None + assert PipelineOffloadManager.OFFLOAD_MGR is empty_manager + assert empty_manager._is_warmup is True + assert empty_manager._cached_chunks_forward == [] assert model_config.fine_grained_activation_offloading is True From b9ed160df89e55411ac3d52e4f01d83d04fdd7a8 Mon Sep 17 00:00:00 2001 From: seonjinn Date: Fri, 14 Aug 2026 23:34:08 -0700 Subject: [PATCH 23/24] fix(megatron): address activation offload review feedback Signed-off-by: seonjinn --- docs/design-docs/training-backends.md | 9 ++- examples/configs/distillation_math.yaml | 5 +- .../configs/distillation_math_megatron.yaml | 5 +- examples/configs/dpo.yaml | 5 +- examples/configs/grpo_math_1B.yaml | 5 +- examples/configs/grpo_math_1B_megatron.yaml | 5 +- examples/configs/ppo_math_1B.yaml | 2 + examples/configs/sft.yaml | 5 +- nemo_rl/models/megatron/setup.py | 11 ++- nemo_rl/models/megatron/train.py | 4 +- .../value/workers/megatron_value_worker.py | 22 +++--- .../models/megatron/test_megatron_setup.py | 6 +- .../value/test_megatron_value_worker.py | 72 +++++++++++++++++++ 13 files changed, 120 insertions(+), 36 deletions(-) diff --git a/docs/design-docs/training-backends.md b/docs/design-docs/training-backends.md index f614d71d682..86b2d525642 100644 --- a/docs/design-docs/training-backends.md +++ b/docs/design-docs/training-backends.md @@ -77,8 +77,13 @@ policy: offload_modules: ["core_attn", "attn_proj"] ``` -Pinned MCore requires the `transformer_engine` CUDA-graph implementation for -this dense module pair; `local` CUDA graphs support only partial MoE offload. +Activation offloading requires the Transformer Engine model implementation. +CUDA graphs are optional; this example was validated with Transformer Engine +CUDA graphs. If graphs are enabled, pinned MCore permits `transformer_engine` +or `full_iteration` for this dense module pair, but only the former is validated +here. `local` CUDA graphs support only partial MoE offload (`expert_fc1`, +`moe_act`, and `fused_group_mlp`). With the default `cuda_graph_impl: none`, no +graph-specific restriction applies. Supported module names are `attn_norm`, `qkv_linear`, `core_attn`, `attn_proj`, `mlp_norm`, `expert_fc1`, `moe_act`, and `fused_group_mlp`. diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml index 7f045497756..fdde83698fa 100644 --- a/examples/configs/distillation_math.yaml +++ b/examples/configs/distillation_math.yaml @@ -95,8 +95,9 @@ policy: &POLICY_BASE empty_unused_memory_level: 0 activation_checkpointing: false # Offload specific module activations to CPU. Works for both dense and MoE - # models. Requires transformer_engine. Different from optimizer_cpu_offload - # which offloads optimizer states. + # models and requires transformer_engine. Different from optimizer_cpu_offload, + # which offloads optimizer states. On TE >= 2.10, + # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", diff --git a/examples/configs/distillation_math_megatron.yaml b/examples/configs/distillation_math_megatron.yaml index c9625d3487f..04374fd10b7 100644 --- a/examples/configs/distillation_math_megatron.yaml +++ b/examples/configs/distillation_math_megatron.yaml @@ -41,8 +41,9 @@ policy: &POLICY_BASE empty_unused_memory_level: 0 activation_checkpointing: false # Offload specific module activations to CPU. Works for both dense and MoE - # models. Requires transformer_engine. Different from optimizer_cpu_offload - # which offloads optimizer states. + # models and requires transformer_engine. Different from optimizer_cpu_offload, + # which offloads optimizer states. On TE >= 2.10, + # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml index 06a77e06172..f7dcec9af49 100755 --- a/examples/configs/dpo.yaml +++ b/examples/configs/dpo.yaml @@ -133,8 +133,9 @@ policy: empty_unused_memory_level: 1 activation_checkpointing: false # Offload specific module activations to CPU. Works for both dense and MoE - # models. Requires transformer_engine. Different from optimizer_cpu_offload - # which offloads optimizer states. + # models and requires transformer_engine. Different from optimizer_cpu_offload, + # which offloads optimizer states. On TE >= 2.10, + # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 1e3bd4e80a3..a23a7ffd3b9 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -187,8 +187,9 @@ policy: 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 # Offload specific module activations to CPU. Works for both dense and MoE - # models. Requires transformer_engine. Different from optimizer_cpu_offload - # which offloads optimizer states. + # models and requires transformer_engine. Different from optimizer_cpu_offload, + # which offloads optimizer states. On TE >= 2.10, + # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index 8e3af788706..aa81d255e30 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/examples/configs/grpo_math_1B_megatron.yaml @@ -89,8 +89,9 @@ policy: 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 # Offload specific module activations to CPU. Works for both dense and MoE - # models. Requires transformer_engine. Different from optimizer_cpu_offload - # which offloads optimizer states. + # models and requires transformer_engine. Different from optimizer_cpu_offload, + # which offloads optimizer states. On TE >= 2.10, + # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 78fabe5c9c9..8e9b231cb1a 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -129,6 +129,7 @@ policy: empty_unused_memory_level: 1 activation_checkpointing: false # Module-level activation CPU offload; see docs/design-docs/training-backends.md. + # On TE >= 2.10 also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when enabled; null disables module selection. offload_modules: null @@ -322,6 +323,7 @@ value: empty_unused_memory_level: 1 activation_checkpointing: false # Module-level activation CPU offload; see docs/design-docs/training-backends.md. + # On TE >= 2.10 also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when enabled; null disables module selection. offload_modules: null diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml index 14938be890e..afe01caf413 100644 --- a/examples/configs/sft.yaml +++ b/examples/configs/sft.yaml @@ -119,8 +119,9 @@ policy: empty_unused_memory_level: 1 activation_checkpointing: false # Offload specific module activations to CPU. Works for both dense and MoE - # models. Requires transformer_engine. Different from optimizer_cpu_offload - # which offloads optimizer states. + # models and requires transformer_engine. Different from optimizer_cpu_offload, + # which offloads optimizer states. On TE >= 2.10, + # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 33252a5dc8c..ac040413066 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1104,18 +1104,15 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: raise KeyError(f"Missing key in fp8_cfg: {e}") megatron_cfg = config["megatron_cfg"] - if "fine_grained_activation_offloading" in megatron_cfg: - fine_grained_activation_offloading = megatron_cfg[ - "fine_grained_activation_offloading" - ] - else: - fine_grained_activation_offloading = None + fine_grained_activation_offloading = megatron_cfg.get( + "fine_grained_activation_offloading" + ) if fine_grained_activation_offloading is False: # Preserve the legacy exemplar's disabled/null semantics and clear any # enabled state carried by a provider or checkpoint. model_cfg.fine_grained_activation_offloading = False - model_cfg.offload_modules = None + model_cfg.offload_modules = [] elif fine_grained_activation_offloading: offload_modules = megatron_cfg.get("offload_modules") if not isinstance(offload_modules, list) or not offload_modules: diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index 6302b31b28a..7a860815a39 100644 --- a/nemo_rl/models/megatron/train.py +++ b/nemo_rl/models/megatron/train.py @@ -74,7 +74,7 @@ @contextmanager -def _suspend_activation_offload_for_forward_only( +def suspend_activation_offload_for_forward_only( model: Union[GPTModel, List[GPTModel]], forward_only: bool ) -> Iterator[None]: """Keep inference-only RL phases from consuming MCore's training warmup.""" @@ -403,7 +403,7 @@ def megatron_forward_backward( forward_backward_func = get_forward_backward_func() if use_router_replay: clear_router_replay(model) - with _suspend_activation_offload_for_forward_only(model, forward_only): + with suspend_activation_offload_for_forward_only(model, forward_only): try: return forward_backward_func( forward_step_func=forward_step, diff --git a/nemo_rl/models/value/workers/megatron_value_worker.py b/nemo_rl/models/value/workers/megatron_value_worker.py index fc23e48cc98..a4bf63b41d1 100644 --- a/nemo_rl/models/value/workers/megatron_value_worker.py +++ b/nemo_rl/models/value/workers/megatron_value_worker.py @@ -75,6 +75,7 @@ from nemo_rl.models.megatron.train import ( LossPostProcessor, megatron_forward_backward, + suspend_activation_offload_for_forward_only, ) from nemo_rl.models.policy.utils import get_runtime_env_for_policy_worker from nemo_rl.models.policy.workers.base_policy_worker import AbstractPolicyWorker @@ -730,16 +731,17 @@ def collection_fn(output_tensor): return output_tensor, collection_fn forward_backward_func = get_forward_backward_func() - list_of_values = forward_backward_func( - forward_step_func=forward_step_fn, - data_iterator=mb_iterator, - model=self.model, - num_microbatches=num_microbatches, - seq_length=padded_seq_length, - micro_batch_size=micro_batch_size_actual, - decoder_seq_length=padded_seq_length, - forward_only=True, - ) + with suspend_activation_offload_for_forward_only(self.model, True): + list_of_values = forward_backward_func( + forward_step_func=forward_step_fn, + data_iterator=mb_iterator, + model=self.model, + num_microbatches=num_microbatches, + seq_length=padded_seq_length, + micro_batch_size=micro_batch_size_actual, + decoder_seq_length=padded_seq_length, + forward_only=True, + ) if is_pipeline_last_stage(ignore_virtual=True): all_values_padded = [] diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 4a2ec20b66a..0142c46a554 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -1228,8 +1228,8 @@ def test_fine_grained_activation_offloading_enabled(self): assert model_cfg.fine_grained_activation_offloading is True assert model_cfg.offload_modules == offload_modules - def test_fine_grained_activation_offloading_disabled_skips(self): - """When flag is False (default), no offload attrs should be set.""" + def test_absent_offloading_flag_leaves_attrs_unset(self): + """When the key is absent and the provider has no offload attrs, none are added.""" from nemo_rl.models.megatron.setup import _apply_performance_config model_cfg = SimpleNamespace(gated_linear_unit=True) @@ -1284,7 +1284,7 @@ def test_explicitly_disabled_offloading_clears_provider_values(self): _apply_performance_config(model_cfg, config) assert model_cfg.fine_grained_activation_offloading is False - assert model_cfg.offload_modules is None + assert model_cfg.offload_modules == [] @pytest.mark.parametrize( "offload_modules", diff --git a/tests/unit/models/value/test_megatron_value_worker.py b/tests/unit/models/value/test_megatron_value_worker.py index 3adf97fe858..8c0efb183f5 100644 --- a/tests/unit/models/value/test_megatron_value_worker.py +++ b/tests/unit/models/value/test_megatron_value_worker.py @@ -30,6 +30,9 @@ import os from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import patch import pytest import ray @@ -49,6 +52,75 @@ pytestmark = pytest.mark.mcore +def test_get_values_suspends_activation_offload() -> None: + """Value inference must preserve the activation-offload training warmup.""" + from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + PipelineOffloadManager, + ) + + from nemo_rl.models.value.workers.megatron_value_worker import ( + MegatronValueWorkerImpl, + ) + + class OffloadManager: + def __init__(self) -> None: + self.do_offload = True + + def disable_offload(self) -> None: + self.do_offload = False + + def enable_offload(self) -> None: + self.do_offload = True + + manager = OffloadManager() + model_config = SimpleNamespace(fine_grained_activation_offloading=True) + model = SimpleNamespace(config=model_config, eval=lambda: None) + worker = SimpleNamespace( + cfg={"train_micro_batch_size": 1}, + model=model, + _policy_like_cfg={}, + mcore_state=SimpleNamespace(straggler_timer=None), + ) + observed_states: list[tuple[bool, bool]] = [] + + def run_forward_only(**kwargs: Any) -> list[dict[str, torch.Tensor]]: + assert kwargs["forward_only"] is True + observed_states.append( + (model_config.fine_grained_activation_offloading, manager.do_offload) + ) + return [{"values": torch.tensor([[1.0]])}] + + with ( + patch.object(PipelineOffloadManager, "OFFLOAD_MGR", manager), + patch( + "nemo_rl.models.value.workers.megatron_value_worker.get_microbatch_iterator", + return_value=(iter([]), 1, 1, 1, 1), + ), + patch( + "nemo_rl.models.value.workers.megatron_value_worker.get_forward_backward_func", + return_value=run_forward_only, + ), + patch( + "nemo_rl.models.value.workers.megatron_value_worker.get_pipeline_model_parallel_group", + return_value=None, + ), + patch( + "nemo_rl.models.value.workers.megatron_value_worker.is_pipeline_last_stage", + return_value=True, + ), + patch("nemo_rl.models.value.workers.megatron_value_worker.broadcast_tensor"), + patch("torch.distributed.get_rank", return_value=0), + patch("torch.cuda.nvtx.range_push"), + patch("torch.cuda.nvtx.range_pop"), + ): + result = MegatronValueWorkerImpl.get_values(worker, BatchedDataDict({})) + + torch.testing.assert_close(result["values"], torch.tensor([[1.0]])) + assert observed_states == [(False, False)] + assert model_config.fine_grained_activation_offloading is True + assert manager.do_offload is True + + def _create_value_test_config( model_name: str, tp: int = 1, From 6a4f1fc96cefb27312c3814f96672b172d71d67b Mon Sep 17 00:00:00 2001 From: seonjinn Date: Sat, 15 Aug 2026 00:08:03 -0700 Subject: [PATCH 24/24] docs(megatron): future-proof offload module references Signed-off-by: seonjinn --- examples/configs/distillation_math.yaml | 9 +++++---- examples/configs/distillation_math_megatron.yaml | 9 +++++---- examples/configs/dpo.yaml | 9 +++++---- examples/configs/grpo_math_1B.yaml | 9 +++++---- examples/configs/grpo_math_1B_megatron.yaml | 9 +++++---- examples/configs/ppo_math_1B.yaml | 12 ++++++++++-- examples/configs/sft.yaml | 9 +++++---- nemo_rl/models/policy/__init__.py | 10 +++++----- 8 files changed, 45 insertions(+), 31 deletions(-) diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml index 4df3dee2bd5..59f7616e0e1 100644 --- a/examples/configs/distillation_math.yaml +++ b/examples/configs/distillation_math.yaml @@ -104,10 +104,11 @@ policy: &POLICY_BASE # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. - # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", - # "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp"]. - # "attn_proj" requires "core_attn". - # See: https://github.com/NVIDIA/Megatron-LM/blob/d12f6c8c9aff51e166d872fd70151687a8e3f375/megatron/core/transformer/transformer_config.py#L1234-L1245 + # Common examples: ["core_attn", "attn_proj", "expert_fc1", "moe_act"]. + # Supported names depend on the pinned Megatron-LM version and are + # validated by MCore. "attn_proj" requires "core_attn". See the latest + # upstream module reference: + # https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#offloadable-modules offload_modules: null # recompute_granularity controls activation checkpointing depth. # "full": recompute all activations (default, max memory savings). diff --git a/examples/configs/distillation_math_megatron.yaml b/examples/configs/distillation_math_megatron.yaml index 61e64bbeea8..5e49a9c73ae 100644 --- a/examples/configs/distillation_math_megatron.yaml +++ b/examples/configs/distillation_math_megatron.yaml @@ -50,10 +50,11 @@ policy: &POLICY_BASE # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. - # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", - # "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp"]. - # "attn_proj" requires "core_attn". - # See: https://github.com/NVIDIA/Megatron-LM/blob/d12f6c8c9aff51e166d872fd70151687a8e3f375/megatron/core/transformer/transformer_config.py#L1234-L1245 + # Common examples: ["core_attn", "attn_proj", "expert_fc1", "moe_act"]. + # Supported names depend on the pinned Megatron-LM version and are + # validated by MCore. "attn_proj" requires "core_attn". See the latest + # upstream module reference: + # https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#offloadable-modules offload_modules: null # recompute_granularity controls activation checkpointing depth. # "full": recompute all activations (default, max memory savings). diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml index 530d63172bf..ec9fcfc338a 100755 --- a/examples/configs/dpo.yaml +++ b/examples/configs/dpo.yaml @@ -142,10 +142,11 @@ policy: # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. - # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", - # "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp"]. - # "attn_proj" requires "core_attn". - # See: https://github.com/NVIDIA/Megatron-LM/blob/d12f6c8c9aff51e166d872fd70151687a8e3f375/megatron/core/transformer/transformer_config.py#L1234-L1245 + # Common examples: ["core_attn", "attn_proj", "expert_fc1", "moe_act"]. + # Supported names depend on the pinned Megatron-LM version and are + # validated by MCore. "attn_proj" requires "core_attn". See the latest + # upstream module reference: + # https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#offloadable-modules offload_modules: null # recompute_granularity controls activation checkpointing depth. # "full": recompute all activations (default, max memory savings). diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 83419a06e7e..c62856a23da 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -196,10 +196,11 @@ policy: # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. - # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", - # "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp"]. - # "attn_proj" requires "core_attn". - # See: https://github.com/NVIDIA/Megatron-LM/blob/d12f6c8c9aff51e166d872fd70151687a8e3f375/megatron/core/transformer/transformer_config.py#L1234-L1245 + # Common examples: ["core_attn", "attn_proj", "expert_fc1", "moe_act"]. + # Supported names depend on the pinned Megatron-LM version and are + # validated by MCore. "attn_proj" requires "core_attn". See the latest + # upstream module reference: + # https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#offloadable-modules offload_modules: null # recompute_granularity controls activation checkpointing depth. # "full": recompute all activations (default, max memory savings). diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index 93f0d8385e7..15700751845 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/examples/configs/grpo_math_1B_megatron.yaml @@ -98,10 +98,11 @@ policy: # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. - # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", - # "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp"]. - # "attn_proj" requires "core_attn". - # See: https://github.com/NVIDIA/Megatron-LM/blob/d12f6c8c9aff51e166d872fd70151687a8e3f375/megatron/core/transformer/transformer_config.py#L1234-L1245 + # Common examples: ["core_attn", "attn_proj", "expert_fc1", "moe_act"]. + # Supported names depend on the pinned Megatron-LM version and are + # validated by MCore. "attn_proj" requires "core_attn". See the latest + # upstream module reference: + # https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#offloadable-modules offload_modules: null # recompute_granularity controls activation checkpointing depth. # "full": recompute all activations (default, max memory savings). diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 8e9b231cb1a..97c5f9688de 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/examples/configs/ppo_math_1B.yaml @@ -131,7 +131,11 @@ policy: # Module-level activation CPU offload; see docs/design-docs/training-backends.md. # On TE >= 2.10 also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false - # Modules to offload when enabled; null disables module selection. + # Modules to offload when enabled; null disables module selection. Common + # examples: ["core_attn", "attn_proj", "expert_fc1", "moe_act"]. Supported + # names depend on the pinned Megatron-LM version and are validated by MCore. + # "attn_proj" requires "core_attn". See the latest upstream module reference: + # https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#offloadable-modules offload_modules: null converter_type: "Qwen2ForCausalLM" tensor_model_parallel_size: 1 @@ -325,7 +329,11 @@ value: # Module-level activation CPU offload; see docs/design-docs/training-backends.md. # On TE >= 2.10 also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false - # Modules to offload when enabled; null disables module selection. + # Modules to offload when enabled; null disables module selection. Common + # examples: ["core_attn", "attn_proj", "expert_fc1", "moe_act"]. Supported + # names depend on the pinned Megatron-LM version and are validated by MCore. + # "attn_proj" requires "core_attn". See the latest upstream module reference: + # https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#offloadable-modules offload_modules: null tensor_model_parallel_size: 1 expert_tensor_parallel_size: 1 diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml index 1c26c5c2930..d2f70c26a2d 100644 --- a/examples/configs/sft.yaml +++ b/examples/configs/sft.yaml @@ -128,10 +128,11 @@ policy: # also set megatron_cfg.env_vars.NVTE_CPU_OFFLOAD_V1: "1". fine_grained_activation_offloading: false # Modules to offload when fine_grained_activation_offloading is true. - # Valid options: ["attn_norm", "qkv_linear", "core_attn", "attn_proj", - # "mlp_norm", "expert_fc1", "moe_act", "fused_group_mlp"]. - # "attn_proj" requires "core_attn". - # See: https://github.com/NVIDIA/Megatron-LM/blob/d12f6c8c9aff51e166d872fd70151687a8e3f375/megatron/core/transformer/transformer_config.py#L1234-L1245 + # Common examples: ["core_attn", "attn_proj", "expert_fc1", "moe_act"]. + # Supported names depend on the pinned Megatron-LM version and are + # validated by MCore. "attn_proj" requires "core_attn". See the latest + # upstream module reference: + # https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#offloadable-modules offload_modules: null # recompute_granularity controls activation checkpointing depth. # "full": recompute all activations (default, max memory savings). diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 7138f837e16..21fe14e5b9c 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -422,11 +422,11 @@ class MegatronConfig(TypedDict): # Megatron-Bridge at runtime). fine_grained_activation_offloading: NotRequired[bool] # Modules to offload when fine_grained_activation_offloading is True. - # Required (no default). Valid values: - # "attn_norm", "qkv_linear", "core_attn", "attn_proj", "mlp_norm", - # "expert_fc1", "moe_act", "fused_group_mlp". Note: "attn_proj" - # requires "core_attn". - # See: https://github.com/NVIDIA/Megatron-LM/blob/d12f6c8c9aff51e166d872fd70151687a8e3f375/megatron/core/transformer/transformer_config.py#L1234-L1245 + # Required (no default). Common examples: "core_attn", "attn_proj", + # "expert_fc1", and "moe_act". Supported names depend on the pinned + # Megatron-LM version and are validated by MCore. "attn_proj" requires + # "core_attn". See the latest upstream module reference: + # https://github.com/NVIDIA/Megatron-LM/blob/main/docs/user-guide/features/fine_grained_activation_offloading.md#offloadable-modules offload_modules: NotRequired[list[str] | None] # Create gloo process groups during Megatron distributed init. # Omitted: use the Megatron Bridge default.