From de681477206ba929aecfca9a376ee6ab451c5124 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 10 Aug 2026 21:30:29 +0000 Subject: [PATCH] Propagate RL runtime config updates Signed-off-by: Philip Petrakian --- megatron/core/transformer/utils.py | 38 +++++++- megatron/rl/rl_utils.py | 44 ++++++--- tests/unit_tests/rl/test_rl_utils.py | 104 ++++++++++++++++++--- tests/unit_tests/transformer/test_utils.py | 44 +++++++++ 4 files changed, 201 insertions(+), 29 deletions(-) diff --git a/megatron/core/transformer/utils.py b/megatron/core/transformer/utils.py index aee4e961b9e..9983f2f6dc0 100644 --- a/megatron/core/transformer/utils.py +++ b/megatron/core/transformer/utils.py @@ -1,6 +1,7 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. """Utilities for transformer layers.""" + import gc import logging from operator import itemgetter @@ -218,7 +219,7 @@ def make_sharded_object_for_checkpoint( def _get_extra_state_offsets( - sharded_offsets: Iterable[Tuple[int, int, int]] + sharded_offsets: Iterable[Tuple[int, int, int]], ) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: """Turns ShardedTensor offsets into offsets suitable for ShardedObject.""" if sharded_offsets: @@ -300,6 +301,41 @@ def sharded_state_dict_default( _sequence_parallel_attr_cache = None +def set_model_config_attribute(model: Any, attribute: str, value: Any) -> None: + """Set a config attribute on a model and all distinct child-module configs. + + Some models give individual layers separate config objects. Runtime model-wide + toggles must update those configs just as they did when every layer shared the + model's root config. + + Args: + model: Model whose configs should be updated. + attribute: Config attribute to set. + value: Value to assign. The same value object is assigned to every config. + """ + root_config = model.config + setattr(root_config, attribute, value) + updated_config_ids = {id(root_config)} + + module_root = model + visited_wrapper_ids = set() + while not isinstance(module_root, torch.nn.Module) or not hasattr(module_root, "_modules"): + visited_wrapper_ids.add(id(module_root)) + module_root = getattr(module_root, "module", None) + if module_root is None or id(module_root) in visited_wrapper_ids: + return + + for module in module_root.modules(): + config = getattr(module, "config", None) + if ( + config is not None + and id(config) not in updated_config_ids + and hasattr(config, attribute) + ): + setattr(config, attribute, value) + updated_config_ids.add(id(config)) + + def _init_sequence_parallel_cache(model, exclude_modules): """ Initialize the cache of modules with sequence parallel attributes. diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 0b834def67b..262c366110f 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -48,7 +48,11 @@ is_batch_invariant_mode_enabled, ) from megatron.core.transformer.enums import CudaGraphModule -from megatron.core.transformer.utils import toggle_cuda_graphs, transition_moe_cudagraphs +from megatron.core.transformer.utils import ( + set_model_config_attribute, + toggle_cuda_graphs, + transition_moe_cudagraphs, +) from megatron.core.utils import ( get_asyncio_loop, get_attr_wrapped_model, @@ -805,7 +809,7 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa # This is a hack to fix megatron's behaviour when flash-decode affects the training code flow. flash_decode = model.config.flash_decode - model.config.flash_decode = False + set_model_config_attribute(model, "flash_decode", False) fp32_output = not (args.fp16 or args.bf16) with torch.no_grad() if no_grad else nullcontext(): logits_or_hidden_states = model( @@ -816,7 +820,7 @@ def get_logprobs(model, tokens, position_ids, no_grad=False, sequence_packing=Fa runtime_gather_output=True, fp32_output=fp32_output, ) - model.config.flash_decode = flash_decode + set_model_config_attribute(model, "flash_decode", flash_decode) pg_collection = get_attr_wrapped_model(model, "pg_collection") pp_group = pg_collection.pp @@ -2200,9 +2204,11 @@ def megatron_rl_inference_mode( # Use local CUDA graphs during rollout inference. An empty module list preserves # full-layer capture when the configured inference scope is layer. - model[0].config.cuda_graph_modules = [] - model[0].config.cuda_graph_impl = "local" - model[0].config.inference_cuda_graph_scope = args.inference_cuda_graph_scope + set_model_config_attribute(model[0], "cuda_graph_modules", []) + set_model_config_attribute(model[0], "cuda_graph_impl", "local") + set_model_config_attribute( + model[0], "inference_cuda_graph_scope", args.inference_cuda_graph_scope + ) # If we get a lower precision wrapper, we go one object deeper. lang_module = model[0].module.module if hasattr(model[0].module, "module") else model[0].module @@ -2260,17 +2266,25 @@ def megatron_rl_inference_mode( # Restore cudagraph scope for training. # MoE partial capture requires specific scopes that aren't user-facing. - model[0].config.cuda_graph_impl = args.cuda_graph_impl - model[0].config.inference_cuda_graph_scope = args.inference_cuda_graph_scope + set_model_config_attribute(model[0], "cuda_graph_impl", args.cuda_graph_impl) + set_model_config_attribute( + model[0], "inference_cuda_graph_scope", args.inference_cuda_graph_scope + ) if args.num_experts is not None: - model[0].config.cuda_graph_modules = [ - CudaGraphModule.mamba, - CudaGraphModule.attn, - CudaGraphModule.moe_router, - CudaGraphModule.moe_preprocess, - ] + set_model_config_attribute( + model[0], + "cuda_graph_modules", + [ + CudaGraphModule.mamba, + CudaGraphModule.attn, + CudaGraphModule.moe_router, + CudaGraphModule.moe_preprocess, + ], + ) else: - model[0].config.cuda_graph_modules = copy.copy(args.cuda_graph_modules) + set_model_config_attribute( + model[0], "cuda_graph_modules", copy.copy(args.cuda_graph_modules) + ) # Switch MoE layers to partial CUDA graph capture for training if args.rl_training_cuda_graphs and args.num_experts is not None: diff --git a/tests/unit_tests/rl/test_rl_utils.py b/tests/unit_tests/rl/test_rl_utils.py index c37d5ec00e1..72b65689ff5 100644 --- a/tests/unit_tests/rl/test_rl_utils.py +++ b/tests/unit_tests/rl/test_rl_utils.py @@ -108,16 +108,37 @@ def make_token_rollout(trajectory, logprobs, generation_mask=None, reward=1.0, p ) -class DummyLangModule: +class DummyConfigModule(torch.nn.Module): def __init__(self, config): + super().__init__() + self.config = config + + +class DummyLogprobsModel(torch.nn.Module): + def __init__(self, config, layer_config): + super().__init__() + self.config = config + self.layer = DummyConfigModule(layer_config) + self.pg_collection = SimpleNamespace(pp=object()) + self.config_values_during_forward = None + + def forward(self, tokens, position_ids, attention_mask, **kwargs): + del position_ids, attention_mask, kwargs + self.config_values_during_forward = ( + self.config.flash_decode, + self.layer.config.flash_decode, + ) + return torch.ones((tokens.shape[0], tokens.shape[1], VOCAB)) + + +class DummyLangModule(torch.nn.Module): + def __init__(self, config): + super().__init__() self.config = config self.rotary_pos_emb = None self.eval = MagicMock() self.train = MagicMock() - def modules(self): - return iter(()) - class DummyMoELayer: def __init__(self, use_partial_cudagraphs): @@ -346,17 +367,33 @@ def _toggle(lang_module, set_to): return MagicMock(side_effect=_toggle) - def test_megatron_rl_inference_mode_restores_training_cuda_graph_state(self, monkeypatch): + @pytest.mark.parametrize( + "share_config", + [pytest.param(True, id="shared-config"), pytest.param(False, id="distinct-config")], + ) + @pytest.mark.parametrize("num_experts", [None, 8], ids=["dense", "moe"]) + def test_megatron_rl_inference_mode_restores_training_cuda_graph_state( + self, monkeypatch, share_config, num_experts + ): config = SimpleNamespace( cuda_graph_impl="none", cuda_graph_modules=[CudaGraphModule.attn], inference_cuda_graph_scope=InferenceCudaGraphScope.none, ) - lang_module = DummyLangModule(config) + layer_config = ( + config + if share_config + else SimpleNamespace( + cuda_graph_impl="none", + cuda_graph_modules=[CudaGraphModule.attn], + inference_cuda_graph_scope=InferenceCudaGraphScope.none, + ) + ) + lang_module = DummyLangModule(layer_config) model = [SimpleNamespace(config=config, module=lang_module)] args = SimpleNamespace( rl_training_cuda_graphs=False, - num_experts=None, + num_experts=num_experts, curr_iteration=11, cuda_graph_impl="local", cuda_graph_modules=[CudaGraphModule.attn], @@ -368,20 +405,61 @@ def test_megatron_rl_inference_mode_restores_training_cuda_graph_state(self, mon with rl_utils.megatron_rl_inference_mode(model, MagicMock(), "local", False) as result: assert result is interface - assert config.cuda_graph_impl == "local" - assert config.cuda_graph_modules == [] - assert config.inference_cuda_graph_scope == InferenceCudaGraphScope.block + for current_config in (config, layer_config): + assert current_config.cuda_graph_impl == "local" + assert current_config.cuda_graph_modules == [] + assert current_config.inference_cuda_graph_scope == InferenceCudaGraphScope.block assert toggle_cuda_graphs.call_args_list == [ call(lang_module, "local"), call(lang_module, "none"), ] - assert config.cuda_graph_impl == "local" - assert config.cuda_graph_modules == [CudaGraphModule.attn] - assert config.inference_cuda_graph_scope == InferenceCudaGraphScope.block + expected_modules = ( + [ + CudaGraphModule.mamba, + CudaGraphModule.attn, + CudaGraphModule.moe_router, + CudaGraphModule.moe_preprocess, + ] + if num_experts is not None + else [CudaGraphModule.attn] + ) + for current_config in (config, layer_config): + assert current_config.cuda_graph_impl == "local" + assert current_config.cuda_graph_modules == expected_modules + assert current_config.inference_cuda_graph_scope == InferenceCudaGraphScope.block lang_module.eval.assert_called_once() lang_module.train.assert_called_once() + @pytest.mark.parametrize( + "share_config", + [pytest.param(True, id="shared-config"), pytest.param(False, id="distinct-config")], + ) + def test_get_logprobs_updates_all_model_configs(self, monkeypatch, share_config): + config = SimpleNamespace(flash_decode=True) + layer_config = config if share_config else SimpleNamespace(flash_decode=True) + model = DummyLogprobsModel(config, layer_config) + monkeypatch.setattr(rl_utils, "get_args", lambda: SimpleNamespace(fp16=False, bf16=False)) + monkeypatch.setattr( + rl_utils, "get_nvtx_range", lambda: (lambda *args, **kwargs: nullcontext()) + ) + monkeypatch.setattr( + rl_utils, "get_attr_wrapped_model", lambda model, name: getattr(model, name) + ) + monkeypatch.setattr(rl_utils, "is_pp_last_stage", lambda _group: False) + + output = rl_utils.get_logprobs( + model, + torch.ones((1, 2), dtype=torch.long), + position_ids=None, + packed_seq_params=object(), + ) + + assert output.shape == (1, 2, VOCAB) + assert model.config_values_during_forward == (False, False) + assert config.flash_decode is True + assert layer_config.flash_decode is True + @pytest.mark.parametrize( "initialize_model_parallel", [ diff --git a/tests/unit_tests/transformer/test_utils.py b/tests/unit_tests/transformer/test_utils.py index 481a6bf5d61..9bbc97eb045 100644 --- a/tests/unit_tests/transformer/test_utils.py +++ b/tests/unit_tests/transformer/test_utils.py @@ -2,6 +2,7 @@ import inspect import os +from types import SimpleNamespace import pytest import torch @@ -13,11 +14,54 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.utils import ( is_layer_window_attention, + set_model_config_attribute, set_model_to_sequence_parallel, ) from tests.unit_tests.test_utilities import Utils +class _TrackingConfig: + def __init__(self, value): + self._runtime_value = value + self.update_count = 0 + + @property + def runtime_value(self): + return self._runtime_value + + @runtime_value.setter + def runtime_value(self, value): + self._runtime_value = value + self.update_count += 1 + + +class _ConfigModule(torch.nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + + +def test_set_model_config_attribute_updates_distinct_configs_once(): + root_config = _TrackingConfig("original") + child_config = _TrackingConfig("original") + unsupported_config = SimpleNamespace() + + module = _ConfigModule(root_config) + module.first_child = _ConfigModule(child_config) + module.second_child = _ConfigModule(child_config) + module.unsupported_child = _ConfigModule(unsupported_config) + model = SimpleNamespace(config=root_config, module=SimpleNamespace(module=module)) + new_value = object() + + set_model_config_attribute(model, "runtime_value", new_value) + + assert root_config.runtime_value is new_value + assert child_config.runtime_value is new_value + assert root_config.update_count == 1 + assert child_config.update_count == 1 + assert not hasattr(unsupported_config, "runtime_value") + + class TestGPTModel: def setup_method(self, method):