perf: support fine-grained activation offloading

docs/design-docs/training-backends.md

@@ -57,6 +57,51 @@ 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
+    cuda_graph_impl: transformer_engine
+    env_vars:
+      NVTE_CPU_OFFLOAD_V1: "1"
+    fine_grained_activation_offloading: true
+    offload_modules: ["core_attn", "attn_proj"]
+```
+
+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`.
+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:
 

examples/configs/distillation_math.yaml

@@ -98,6 +98,18 @@ policy: &POLICY_BASE
         cuda_graph_impl: "none"
         cuda_graph_modules: []
         cuda_graph_warmup_steps: 3
+        # Offload specific module activations to CPU. Works for both dense and MoE
+        # 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.
+        # 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).
         # "selective": recompute only specific modules (see recompute_modules).

examples/configs/distillation_math_megatron.yaml

@@ -44,6 +44,18 @@ policy: &POLICY_BASE
         cuda_graph_impl: "none"
         cuda_graph_modules: []
         cuda_graph_warmup_steps: 3
+        # Offload specific module activations to CPU. Works for both dense and MoE
+        # 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.
+        # 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).
         # "selective": recompute only specific modules (see recompute_modules).

examples/configs/dpo.yaml

@@ -136,6 +136,18 @@ policy:
     cuda_graph_impl: "none"
     cuda_graph_modules: []
     cuda_graph_warmup_steps: 3
+    # Offload specific module activations to CPU. Works for both dense and MoE
+    # 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.
+    # 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).
     # "selective": recompute only specific modules (see recompute_modules).

examples/configs/grpo_math_1B.yaml

@@ -190,6 +190,18 @@ policy:
     cuda_graph_impl: "none"
     cuda_graph_modules: []
     cuda_graph_warmup_steps: 3
+    # Offload specific module activations to CPU. Works for both dense and MoE
+    # 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.
+    # 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).
     # "selective": recompute only specific modules (see recompute_modules).

examples/configs/grpo_math_1B_megatron.yaml

@@ -92,6 +92,18 @@ policy:
     cuda_graph_impl: "none"
     cuda_graph_modules: []
     cuda_graph_warmup_steps: 3
+    # Offload specific module activations to CPU. Works for both dense and MoE
+    # 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.
+    # 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).
     # "selective": recompute only specific modules (see recompute_modules).

examples/configs/ppo_math_1B.yaml

@@ -128,6 +128,15 @@ policy:
     enabled: false
     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. 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
     expert_tensor_parallel_size: 1
@@ -317,6 +326,15 @@ value:
     enabled: false
     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. 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
     expert_model_parallel_size: 1

examples/configs/sft.yaml

@@ -122,6 +122,18 @@ policy:
     cuda_graph_impl: "none"
     cuda_graph_modules: []
     cuda_graph_warmup_steps: 3
+    # Offload specific module activations to CPU. Works for both dense and MoE
+    # 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.
+    # 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).
     # "selective": recompute only specific modules (see recompute_modules).

nemo_rl/models/megatron/setup.py

@@ -1125,6 +1125,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_cfg = config["megatron_cfg"]
+    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 = []
+    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
+
 
 def _validate_optimizer_config(config: PolicyConfig) -> None:
     """Validate optimizer configuration."""

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,10 @@
     get_tensor_model_parallel_rank,
 )
 from megatron.core.pipeline_parallel import get_forward_backward_func
-from megatron.core.utils import StragglerDetector
+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 (
     TrainingSamplingParams,
@@ -70,6 +73,49 @@
 ]
 
 
+@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))
+
+    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:
+        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(
     model: GPTModel,
     data_dict: BatchedDataDict[Any],
@@ -366,20 +412,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:

nemo_rl/models/policy/__init__.py

@@ -414,6 +414,20 @@ class MegatronConfig(TypedDict):
     moe_pad_experts_for_cuda_graph_inference: NotRequired[bool]
     # 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 both dense and MoE models. Different from
+    # optimizer_cpu_offload which offloads optimizer states.
+    # 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). 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.
     use_gloo_process_groups: NotRequired[bool]

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 = []