Propagate RL runtime config updates to module configs

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.

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:

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",
         [

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):