feat(torchtitan): enable dynamic model parameter override via CLI (#254)

Author: Xiaoming-AMD

examples/megatron/prepare.py

@@ -191,7 +191,7 @@ def main():
 
     log_info(f"BACKEND_PATH {args.backend_path}")
     # primus_config = PrimusParser().parse(args)
-    primus_config = load_primus_config(args, unknown)
+    primus_config, _ = load_primus_config(args, unknown)
 
     primus_path = Path(args.primus_path).resolve()
     log_info(f"PRIMUS_PATH is set to: {primus_path}")

primus/core/launcher/parser.py

@@ -2,7 +2,7 @@
 import os
 from pathlib import Path
 from types import SimpleNamespace
-from typing import List
+from typing import Any, Dict, List, Tuple
 
 from primus.core.launcher.config import PrimusConfig
 from primus.core.utils import constant_vars, yaml_utils
@@ -142,6 +142,30 @@ def _check_keys_exist(ns: SimpleNamespace, overrides: dict, prefix=""):
             _check_keys_exist(attr_val, v, prefix=full_key)
 
 
+def _split_known_unknown(ns: SimpleNamespace, overrides: dict) -> Tuple[dict, dict]:
+    """
+    Split overrides into two dictionaries:
+      - known: keys that exist in the namespace
+      - unknown: keys not defined in the namespace
+    """
+    known, unknown = {}, {}
+    for k, v in overrides.items():
+        if hasattr(ns, k):
+            attr_val = getattr(ns, k)
+            if isinstance(v, dict) and isinstance(attr_val, SimpleNamespace):
+                sub_known, sub_unknown = _split_known_unknown(attr_val, v)
+                if sub_known:
+                    known[k] = sub_known
+                if sub_unknown:
+                    unknown[k] = sub_unknown
+            else:
+                known[k] = v
+        else:
+            unknown[k] = v
+            # print(f"[PrimusConfig] Unknown key '{k}' delegated to backend.")
+    return known, unknown
+
+
 def parse_args(extra_args_provider=None, ignore_unknown_args=False):
     args, unknown_args = _parse_args(extra_args_provider, ignore_unknown_args=True)
 
@@ -156,7 +180,7 @@ def parse_args(extra_args_provider=None, ignore_unknown_args=False):
     return primus_config
 
 
-def load_primus_config(args: argparse.Namespace, overrides: List[str]) -> PrimusConfig:
+def load_primus_config(args: argparse.Namespace, overrides: List[str]) -> Tuple[Any, Dict[str, Any]]:
     """
     Build the Primus configuration with optional command-line overrides.
 
@@ -177,10 +201,19 @@ def load_primus_config(args: argparse.Namespace, overrides: List[str]) -> Primus
 
     # 3 Apply overrides to pre_trainer module config
     pre_trainer_cfg = primus_config.get_module_config("pre_trainer")
-    _check_keys_exist(pre_trainer_cfg, override_ns)
-    _deep_merge_namespace(pre_trainer_cfg, override_ns)
+    # _check_keys_exist(pre_trainer_cfg, override_ns)
+    # _deep_merge_namespace(pre_trainer_cfg, override_ns)
 
-    return primus_config
+    # return primus_config
+    known_overrides, unknown_overrides = _split_known_unknown(pre_trainer_cfg, override_ns)
+
+    if known_overrides:
+        _deep_merge_namespace(pre_trainer_cfg, known_overrides)
+
+    if unknown_overrides:
+        print(f"[PrimusConfig] Detected unknown override keys: {list(unknown_overrides.keys())}")
+
+    return primus_config, unknown_overrides
 
 
 class PrimusParser(object):

primus/modules/trainer/megatron/pre_trainer.py

@@ -118,6 +118,15 @@ def pop(cls):
 class MegatronPretrainTrainer(MegatronTrainer):
     def __init__(self, *args, **kwargs):
         kwargs["module_name"] = "pre_trainer"
+
+        # Explicitly reject unknown extra_args
+        extra_args = kwargs.pop("extra_args", None)
+        if extra_args:
+            raise ValueError(
+                f"[MegatronPretrainTrainer] Unexpected extra_args detected: {extra_args}. "
+                f"Megatron backend does not support unregistered config keys."
+            )
+
         super().__init__(*args, **kwargs)
 
     def get_batch(self, data_iterator):

primus/modules/trainer/torchtitan/pre_trainer.py

@@ -13,6 +13,7 @@
 
 class TorchTitanPretrainTrainer(BaseModule):
     def __init__(self, *args, **kwargs):
+        extra_args = kwargs.pop("extra_args", None)
         super().__init__(*args, **kwargs)
 
         # important: make sure patch torchtitan logger first
@@ -26,6 +27,7 @@ def __init__(self, *args, **kwargs):
         cfg_dict = nested_namespace_to_dict(pre_trainer_cfg)
 
         self.patch_torchtitan_embedding_amp(cfg_dict["primus_turbo"]["enable_embedding_autocast"])
+        self.patch_titan_train_spec(pre_trainer_cfg.model.name, pre_trainer_cfg.model.flavor, extra_args)
 
         # ensure checkpoint patch applied before import torchtitan
         # background: consolidate_safetensors_files_on_every_rank is a new DCP
@@ -508,3 +510,89 @@ def new_init(self, *args, **kwargs):
         primus_logger.info(
             "[PrimusPatch][AMP] nn.Embedding.__init__ patched for AMP/mixed precision alignment."
         )
+
+    def patch_titan_train_spec(self, model_name: str, flavor: str, model_overrides: Dict[str, Any]):
+        """
+        Monkey patch torchtitan.train_spec.get_train_spec to override model args dynamically.
+        All override keys MUST start with "model." (e.g., {"model.n_layers": 8}).
+        """
+        from primus.core.utils.logger import _logger as primus_logger
+
+        if not model_overrides:
+            primus_logger.info("[PrimusPatch][ModelOverride] No model_overrides provided, skip patch.")
+            return
+
+        primus_logger.info(f"[PrimusPatch][ModelOverride] Applying model_overrides: {model_overrides}")
+
+        # --- flatten nested form {"model": {"n_layers": 4}} → {"model.n_layers": 4}
+        flat_overrides = {}
+        for k, v in model_overrides.items():
+            if k == "model" and isinstance(v, dict):
+                for subk, subv in v.items():
+                    flat_overrides[f"model.{subk}"] = subv
+            else:
+                flat_overrides[k] = v
+        model_overrides = flat_overrides
+
+        # Enforce `model.` prefix strictly
+        bad_keys = [k for k in model_overrides.keys() if not k.startswith("model.")]
+        if bad_keys:
+            raise ValueError(
+                # f"[PrimusPatch][ModelOverride] Unsupported override keys (must start with 'model.'): {bad_keys}"
+                f"[PrimusPatch][ModelOverride] Invalid override keys detected: {bad_keys}. "
+                "These parameters belong to the model configuration and must be specified "
+                "with the 'model.' prefix (e.g., 'model.n_layers', 'model.dim')."
+            )
+
+        primus_logger.info(
+            f"[PrimusPatch][ModelOverride] model_overrides provided for '{model_name}' (flavor={flavor}): {model_overrides}"
+        )
+
+        import torchtitan.protocols.train_spec as train_spec_module
+
+        orig_get_train_spec = train_spec_module.get_train_spec
+
+        def patched_get_train_spec(name: str):
+            spec = orig_get_train_spec(name)
+            if name != model_name:
+                return spec  # only patch targeted model
+
+            assert hasattr(
+                spec, "model_args"
+            ), f"[PrimusPatch][ModelOverride] train_spec for '{name}' missing model_args"
+            model_args_root = spec.model_args
+            assert isinstance(
+                model_args_root, dict
+            ), f"[PrimusPatch][ModelOverride] train_spec.model_args must be dict, got {type(model_args_root)}"
+
+            if flavor not in model_args_root:
+                raise KeyError(
+                    f"[PrimusPatch][ModelOverride] flavor '{flavor}' not found in model_args for '{name}'. "
+                    f"Available flavors: {list(model_args_root.keys())}"
+                )
+
+            target_args = model_args_root[flavor]
+            assert is_dataclass(
+                target_args
+            ), f"[PrimusPatch][ModelOverride] Expected dataclass model_args, got {type(target_args)}"
+
+            before = asdict(target_args)
+            for k, v in model_overrides.items():
+                field_name = k[len("model.") :]
+                if not hasattr(target_args, field_name):
+                    raise AttributeError(
+                        f"[PrimusPatch][ModelOverride] '{type(target_args).__name__}' has no field '{field_name}'"
+                    )
+                setattr(target_args, field_name, v)
+
+            primus_logger.info(
+                f"[PrimusPatch][ModelOverride] Patched dataclass model_args['{flavor}'] "
+                f"for '{name}' with {model_overrides} (before={before})"
+            )
+            return spec
+
+        # Apply the patch globally
+        train_spec_module.get_train_spec = patched_get_train_spec
+        primus_logger.info(
+            f"[PrimusPatch][ModelOverride] get_train_spec for '{model_name}' successfully monkey patched (flavor={flavor})."
+        )

primus/pretrain.py

@@ -98,7 +98,7 @@ def setup_env(data_path: str):
         print(f"[Primus CLI] HF_HOME already set: {hf_home}")
 
 
-def launch_pretrain_trainer(primus_cfg: PrimusConfig):
+def launch_pretrain_trainer(primus_cfg: PrimusConfig, extra_args=None):
     """
     Launch the training using the Primus trainer.
 
@@ -126,6 +126,7 @@ def launch_pretrain_trainer(primus_cfg: PrimusConfig):
         module_world_size=world_size,
         module_master_addr=master_addr,
         module_master_port=master_port,
+        extra_args=extra_args,
     )
 
     # Launch training
@@ -150,7 +151,7 @@ def launch_pretrain_from_cli(args, overrides):
 
     setup_env(data_path=args.data_path)
 
-    primus_cfg = load_primus_config(args, overrides)
+    primus_cfg, unknown_overrides = load_primus_config(args, overrides)
 
     # Export merged config if requested
     if args.export_config:
@@ -160,7 +161,7 @@ def launch_pretrain_from_cli(args, overrides):
     framework = primus_cfg.get_module_config("pre_trainer").framework
     setup_backend_path(framework=framework, backend_path=args.backend_path, verbose=True)
 
-    launch_pretrain_trainer(primus_cfg=primus_cfg)
+    launch_pretrain_trainer(primus_cfg=primus_cfg, extra_args=unknown_overrides)
 
 
 if __name__ == "__main__":