diff --git a/docs/design-docs/training-backends.md b/docs/design-docs/training-backends.md index 4cb0eca8a9b..86b2d525642 100644 --- a/docs/design-docs/training-backends.md +++ b/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: diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml index 6270c7cf345..59f7616e0e1 100644 --- a/examples/configs/distillation_math.yaml +++ b/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). diff --git a/examples/configs/distillation_math_megatron.yaml b/examples/configs/distillation_math_megatron.yaml index 7fa821077ea..5e49a9c73ae 100644 --- a/examples/configs/distillation_math_megatron.yaml +++ b/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). diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml index e3e1bd9d76b..ec9fcfc338a 100755 --- a/examples/configs/dpo.yaml +++ b/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). diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 49627590547..c62856a23da 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/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). diff --git a/examples/configs/grpo_math_1B_megatron.yaml b/examples/configs/grpo_math_1B_megatron.yaml index c264ab32960..15700751845 100644 --- a/examples/configs/grpo_math_1B_megatron.yaml +++ b/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). diff --git a/examples/configs/ppo_math_1B.yaml b/examples/configs/ppo_math_1B.yaml index 366cf08cf31..97c5f9688de 100644 --- a/examples/configs/ppo_math_1B.yaml +++ b/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 diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml index b556b327ec9..d2f70c26a2d 100644 --- a/examples/configs/sft.yaml +++ b/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). diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index c1bbb4de879..bdf4d7849ac 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/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.""" diff --git a/nemo_rl/models/megatron/train.py b/nemo_rl/models/megatron/train.py index 751169107f8..d8ecb4f6590 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,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: diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 7d1ec11d46f..21fe14e5b9c 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/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] 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 c431e344ce8..e3f1c4c5c50 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -1258,6 +1258,161 @@ def test_fp8_configuration(self): assert model_cfg.fp8_recipe == "default" assert model_cfg.fp8_param is False + 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 + model_cfg.num_moe_experts = 8 + offload_modules = ["mlp_norm", "moe_act"] + config = { + "megatron_cfg": { + "activation_checkpointing": False, + "apply_rope_fusion": False, + "bias_activation_fusion": False, + "gradient_accumulation_fusion": False, + "use_fused_weighted_squared_relu": 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_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) + config = { + "megatron_cfg": { + "activation_checkpointing": False, + "apply_rope_fusion": False, + "bias_activation_fusion": False, + "gradient_accumulation_fusion": False, + "use_fused_weighted_squared_relu": False, + } + } + + _apply_performance_config(model_cfg, config) + + 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 == [] + + @pytest.mark.parametrize( + "offload_modules", + [[], None, "moe_act", 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, + "use_fused_weighted_squared_relu": 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 None → 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, + "use_fused_weighted_squared_relu": 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.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 diff --git a/tests/unit/models/megatron/test_train.py b/tests/unit/models/megatron/test_train.py index c13ffa36c86..130cbf9825b 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 @@ -769,6 +770,306 @@ 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: 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, + ) + + 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 + PipelineOffloadManager.OFFLOAD_MGR = empty_manager + 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}} + ) + + 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 empty_manager + assert empty_manager._is_warmup is True + assert empty_manager._cached_chunks_forward == [] + + 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 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.""" 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, diff --git a/tests/unit/reference_configs/distillation_math.yaml b/tests/unit/reference_configs/distillation_math.yaml index 7e0a60c7c07..5f5f2d8ddfb 100644 --- a/tests/unit/reference_configs/distillation_math.yaml +++ b/tests/unit/reference_configs/distillation_math.yaml @@ -97,6 +97,8 @@ policy: &POLICY_BASE cuda_graph_impl: "none" cuda_graph_modules: [] cuda_graph_warmup_steps: 3 + fine_grained_activation_offloading: false + offload_modules: null recompute_granularity: "full" recompute_modules: null tensor_model_parallel_size: 2 diff --git a/tests/unit/reference_configs/dpo.yaml b/tests/unit/reference_configs/dpo.yaml index e3293bd27d8..133624c0ea2 100755 --- a/tests/unit/reference_configs/dpo.yaml +++ b/tests/unit/reference_configs/dpo.yaml @@ -130,6 +130,8 @@ policy: cuda_graph_impl: "none" cuda_graph_modules: [] cuda_graph_warmup_steps: 3 + fine_grained_activation_offloading: false + offload_modules: null recompute_granularity: "full" recompute_modules: null tensor_model_parallel_size: 2 diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index f905aff1cc4..8e0b1892b52 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -188,6 +188,8 @@ policy: cuda_graph_impl: "none" cuda_graph_modules: [] cuda_graph_warmup_steps: 3 + fine_grained_activation_offloading: false + 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). 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 diff --git a/tests/unit/reference_configs/sft.yaml b/tests/unit/reference_configs/sft.yaml index 03ee9867bd9..9d787add94f 100644 --- a/tests/unit/reference_configs/sft.yaml +++ b/tests/unit/reference_configs/sft.yaml @@ -116,6 +116,8 @@ policy: cuda_graph_impl: "none" cuda_graph_modules: [] cuda_graph_warmup_steps: 3 + fine_grained_activation_offloading: false + offload_modules: null recompute_granularity: "full" recompute_modules: null tensor_model_parallel_size: 1