[override] Allow passing kwargs in override from both CLI and config

tests/unit_tests/test_override.py

@@ -358,6 +358,82 @@ def fmt_swap(cfg: ComponentA.Config) -> ComponentB.Config:
             self.assertIn("->", line)
 
 
+class TestOverrideKwargs(unittest.TestCase):
+    """``override.imports`` entries may carry kwargs forwarded to the factory."""
+
+    def setUp(self):
+        clear_overrides()
+
+    def tearDown(self):
+        clear_overrides()
+
+    def test_same_module_different_kwargs_per_actor(self):
+        # The motivating case: two config trees share one override module but
+        # pass different kwargs (e.g. RL trainer vs. generator capacity factor).
+        @override("per_actor", target=ComponentA.Config, fqns=["child"])
+        def per_actor(cfg: ComponentA.Config, *, extra: int) -> ComponentB.Config:
+            return ComponentB.Config(dim=cfg.dim, extra=extra)
+
+        trainer_cfg = ParentComponent.Config()
+        generator_cfg = ParentComponent.Config()
+        apply_overrides(OverrideConfig(imports=[(__name__, {"extra": 1})]), trainer_cfg)
+        apply_overrides(
+            OverrideConfig(imports=[(__name__, {"extra": 2})]), generator_cfg
+        )
+
+        self.assertEqual(trainer_cfg.child.extra, 1)
+        self.assertEqual(generator_cfg.child.extra, 2)
+
+    def test_bare_string_entry_calls_factory_without_kwargs(self):
+        # A bare-string entry keeps the pre-kwargs contract: no kwargs passed.
+        @override("no_kw", target=ComponentA.Config, fqns=["child"])
+        def no_kw(cfg: ComponentA.Config, *, extra: int = 9) -> ComponentB.Config:
+            return ComponentB.Config(dim=cfg.dim, extra=extra)
+
+        parent_cfg = ParentComponent.Config()
+        apply_overrides(OverrideConfig(imports=[__name__]), parent_cfg)
+        self.assertEqual(parent_cfg.child.extra, 9)
+
+    def test_unknown_kwarg_raises(self):
+        @override("strict_kw", target=ComponentA.Config)
+        def strict_kw(cfg: ComponentA.Config, *, extra: int) -> ComponentB.Config:
+            return ComponentB.Config(dim=cfg.dim, extra=extra)
+
+        # A kwarg the factory does not accept is a plain TypeError from the call.
+        with self.assertRaisesRegex(TypeError, "unexpected keyword argument"):
+            apply_overrides(
+                OverrideConfig(imports=[(__name__, {"typo": 1})]),
+                ParentComponent.Config(),
+            )
+
+    def test_kwargs_for_no_matching_override_raises(self):
+        # kwargs that activate no override is a misconfiguration, not a silent
+        # no-op (the registry is empty here after setUp's clear_overrides).
+        with self.assertRaisesRegex(ValueError, "activated no override"):
+            apply_overrides(
+                OverrideConfig(imports=[("torchtitan.config.override", {"x": 1})]),
+                ParentComponent.Config(),
+            )
+
+    def test_parse_cli_imports(self):
+        from torchtitan.config.override import parse_cli_imports
+
+        # Plain modules (comma- or space-separated) and modules whose kwargs are
+        # attached to the name as ``module=<json>`` -- the CLI grammar backing
+        # ``--override.imports``.
+        self.assertEqual(parse_cli_imports(["a.b,c.d"]), ["a.b", "c.d"])
+        self.assertEqual(
+            parse_cli_imports(['mod={"block_size": 256, "flag": null}']),
+            [("mod", {"block_size": 256, "flag": None})],
+        )
+        self.assertEqual(
+            parse_cli_imports(["plain", 'mod={"x": 1}']),
+            ["plain", ("mod", {"x": 1})],
+        )
+        with self.assertRaisesRegex(ValueError, "must be a JSON object"):
+            parse_cli_imports(["mod=123"])
+
+
 class TestDerive(unittest.TestCase):
     def test_copies_shared_and_applies_deltas(self):
         src = ComponentA.Config(dim=99)

torchtitan/config/manager.py

@@ -243,6 +243,42 @@ def list_str_rule(type_info: tyro.constructors.PrimitiveTypeInfo):
                 str_from_instance=lambda instance: [",".join(instance)],
             )
 
+        @registry.primitive_rule
+        def override_imports_rule(type_info: tyro.constructors.PrimitiveTypeInfo):
+            """Parse ``--override.imports`` module paths, each with optional kwargs.
+
+            Needed because ``OverrideConfig.imports`` is a ``str | tuple`` union,
+            which tyro cannot render as a CLI flag on its own. Tokens are space-
+            or comma-separated module paths; attach kwargs to a module with
+            ``module=<json>`` (see ``parse_cli_imports``), e.g.
+            ``--override.imports 'my_pkg.triton_rope={"block_size": 256}'``.
+            """
+            from torchtitan.config.override import (
+                format_cli_imports,
+                OverrideImport,
+                parse_cli_imports,
+            )
+
+            if type_info.type != list[OverrideImport]:
+                return None
+            return tyro.constructors.PrimitiveConstructorSpec(
+                nargs="*",
+                metavar="MODULE[='{\"k\": v}'] ...",
+                instance_from_str=lambda args: parse_cli_imports(args),
+                is_instance=lambda instance: isinstance(instance, list)
+                and all(
+                    isinstance(i, str)
+                    or (
+                        isinstance(i, tuple)
+                        and len(i) == 2
+                        and isinstance(i[0], str)
+                        and isinstance(i[1], dict)
+                    )
+                    for i in instance
+                ),
+                str_from_instance=lambda instance: format_cli_imports(instance),
+            )
+
 
 # Initialize the custom registry for tyro
 custom_registry = tyro.constructors.ConstructorRegistry()