From 1db830b328f84e472bbf497651d2fc27ac14d147 Mon Sep 17 00:00:00 2001 From: Vu Dinh Date: Fri, 29 May 2026 10:27:33 -0400 Subject: [PATCH 1/3] [fix] Preserve original AutoConfig in MegatronWorker.init_configs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a Megatron-Bridge SFT run was launched with overrides like `++transformer_config_kwargs.mtp_num_layers=0`, the exported HF checkpoint's `config.json` carried the runtime override (`num_nextn_predict_layers=0`) instead of the on-disk source-of-truth value (e.g. 4 for Nemotron-H 120B), making the saved directory silently inconsistent with the reference repo. Root cause: `update_model_config` mutated the input `AutoConfig` in place via `setattr` (recursing into sub-configs the same way). `MegatronWorker.init_configs` then assigned that already-mutated object to `self.strategy.hf_config`, and `MegatronStrategy .save_hf_configs` published it via `model_config.save_pretrained (work_dir)`. Fix: - `skyrl/train/utils/utils.py`: `update_model_config` now returns a `deepcopy` of the input with overrides applied. The recursive setattr logic is preserved unchanged in a private `_apply_overrides_in_place` helper used internally on the copy. - `skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py`: `init_configs` keeps the un-mutated `AutoConfig.from_pretrained (...)` result as `hf_config_original` and assigns it to `self.strategy.hf_config`, while the Megatron provider/runtime path continues to use the override-applied copy returned by `update_model_config`. Verification: - pre-commit hooks (ruff, black, gitleaks): pass on both files. - Standalone unit-style smoke test (recreating the function bodies in isolation, since the module pulls in heavy runtime deps): asserts the original config is unchanged, the returned copy carries top-level overrides, and nested-dict overrides also do not leak back into the input. Includes a Bug B-style repro: `update_model_config(hf_cfg, {"num_nextn_predict_layers": 0})` leaves `hf_cfg.num_nextn_predict_layers` at the original value. - No existing test in `tests/` references `update_model_config`, `MegatronWorker.init_configs`, or `save_hf_configs`, so this path was previously uncovered. Backward compatibility: the only call site of `update_model_config` is `MegatronWorker.init_configs`, which is updated in this commit to consume the returned value. The recursion-via-setattr semantics are identical to before — only the target object changed (a deep-copied config instead of the caller's input). Signed-off-by: Vu Dinh --- .../workers/megatron/megatron_worker.py | 9 ++++++--- skyrl/train/utils/utils.py | 20 +++++++++++++++++-- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 58061e9c3c..0c9544b9f8 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -331,7 +331,7 @@ def init_configs( Initialize the Megatron-Bridge bridge and provider objects + hf_config and tokenizer """ tokenizer = get_tokenizer(model_path, trust_remote_code=True) - hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + hf_config_original = AutoConfig.from_pretrained(model_path, trust_remote_code=True) override_config_kwargs = { "bos_token_id": tokenizer.bos_token_id, @@ -339,7 +339,7 @@ def init_configs( "pad_token_id": tokenizer.pad_token_id, } override_config_kwargs.update(model_config_kwargs.get("model_config", {})) - update_model_config(hf_config, override_config_kwargs=override_config_kwargs) + hf_config = update_model_config(hf_config_original, override_config_kwargs=override_config_kwargs) transformer_config_kwargs = ( transformer_config_kwargs @@ -402,7 +402,10 @@ def init_configs( self.provider = provider self.bridge = bridge - self.strategy.hf_config = hf_config + # strategy.hf_config is the on-disk source-of-truth used by + # save_hf_configs and must NOT carry runtime overrides like + # mtp_num_layers=0; assign the un-mutated AutoConfig here. + self.strategy.hf_config = hf_config_original self.tokenizer = tokenizer self.enable_router_replay = megatron_config.moe_enable_routing_replay diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index 4ac09c3eb3..0811859bea 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -6,6 +6,7 @@ import socket import sys import time +from copy import deepcopy from datetime import datetime from pathlib import Path @@ -976,15 +977,30 @@ def peer_access_supported(max_num_gpus_per_node: int): def update_model_config(module_config, override_config_kwargs): - """Update the module config with the override_config_kwargs. + """Return a copy of ``module_config`` with ``override_config_kwargs`` applied. + + The returned config is a deep copy, so the caller's input is left + unmodified. Nested dict values in ``override_config_kwargs`` recurse into + the corresponding sub-config attribute (which is already part of the deep + copy, so the recursion mutates the copy in place). Args: module_config: The module config from Huggingface Transformers. override_config_kwargs: The kwargs to override the module config. + + Returns: + A new module config with the overrides applied. """ + new_config = deepcopy(module_config) + _apply_overrides_in_place(new_config, override_config_kwargs) + return new_config + + +def _apply_overrides_in_place(module_config, override_config_kwargs): + """Apply override kwargs to ``module_config`` in place (used for sub-configs).""" for key, val in override_config_kwargs.items(): if isinstance(val, dict): - update_model_config(getattr(module_config, key), val) + _apply_overrides_in_place(getattr(module_config, key), val) else: setattr(module_config, key, val) From 3a46f971324b490dc241009220d35f6f6baa059b Mon Sep 17 00:00:00 2001 From: Vu Dinh Date: Fri, 29 May 2026 21:31:02 -0400 Subject: [PATCH 2/3] [test] Add CPU unit tests for update_model_config non-mutating contract Locks in three properties of update_model_config: - the caller's input config is not mutated at the top level, - the returned copy carries top-level overrides, - nested-dict overrides do not leak back into the caller's input. Includes the Bug B repro: update_model_config(cfg, {"num_nextn_predict_layers": 0}) leaves the input's num_nextn_predict_layers at its original value. Runs on cpu_skyrl_train.yaml (uv run --extra skyrl-train --extra dev pytest tests/train/...). Signed-off-by: Vu Dinh --- tests/train/utils/test_utils.py | 38 +++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 tests/train/utils/test_utils.py diff --git a/tests/train/utils/test_utils.py b/tests/train/utils/test_utils.py new file mode 100644 index 0000000000..5546d7f6b9 --- /dev/null +++ b/tests/train/utils/test_utils.py @@ -0,0 +1,38 @@ +""" +uv run --isolated --extra dev pytest tests/train/utils/test_utils.py +""" + +from copy import deepcopy +from types import SimpleNamespace + +from skyrl.train.utils.utils import update_model_config + + +def _make_config(): + """Build a config-like object with a nested sub-config attribute.""" + sub = SimpleNamespace(mtp_num_layers=4, hidden_size=128) + return SimpleNamespace(num_nextn_predict_layers=4, sub_config=sub) + + +class TestUpdateModelConfigNonMutating: + """Lock in the non-mutating contract of ``update_model_config``.""" + + def test_input_is_not_mutated_at_top_level(self): + cfg = _make_config() + before = deepcopy(cfg) + update_model_config(cfg, {"num_nextn_predict_layers": 0}) + assert cfg.num_nextn_predict_layers == before.num_nextn_predict_layers + assert cfg.sub_config.mtp_num_layers == before.sub_config.mtp_num_layers + + def test_returned_copy_carries_top_level_overrides(self): + cfg = _make_config() + new_cfg = update_model_config(cfg, {"num_nextn_predict_layers": 0}) + assert new_cfg.num_nextn_predict_layers == 0 + assert new_cfg is not cfg + + def test_nested_overrides_do_not_leak_back_into_input(self): + cfg = _make_config() + new_cfg = update_model_config(cfg, {"sub_config": {"mtp_num_layers": 0}}) + assert new_cfg.sub_config.mtp_num_layers == 0 + assert cfg.sub_config.mtp_num_layers == 4 + assert new_cfg.sub_config is not cfg.sub_config From db83caff7bb7cdfc07f821660d9059f93aeab6be Mon Sep 17 00:00:00 2001 From: Vu Dinh Date: Fri, 29 May 2026 21:36:17 -0400 Subject: [PATCH 3/3] [test] Rename test_utils.py to test_update_model_config.py Avoids a pytest module-name collision with tests/train/generators/test_utils.py (no __init__.py in either directory; pytest's rootdir-based collection treats both as the same top-level module name and aborts collection). The new basename is also more descriptive of what the file actually exercises. Signed-off-by: Vu Dinh --- .../train/utils/{test_utils.py => test_update_model_config.py} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename tests/train/utils/{test_utils.py => test_update_model_config.py} (94%) diff --git a/tests/train/utils/test_utils.py b/tests/train/utils/test_update_model_config.py similarity index 94% rename from tests/train/utils/test_utils.py rename to tests/train/utils/test_update_model_config.py index 5546d7f6b9..4e8591d2db 100644 --- a/tests/train/utils/test_utils.py +++ b/tests/train/utils/test_update_model_config.py @@ -1,5 +1,5 @@ """ -uv run --isolated --extra dev pytest tests/train/utils/test_utils.py +uv run --isolated --extra dev pytest tests/train/utils/test_update_model_config.py """ from copy import deepcopy