diff --git a/tests/unit_tests/test_override.py b/tests/unit_tests/test_override.py index d7cb5a2d67..4f6881121d 100644 --- a/tests/unit_tests/test_override.py +++ b/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=`` -- 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) diff --git a/torchtitan/config/manager.py b/torchtitan/config/manager.py index f253d23b63..2adebd0173 100644 --- a/torchtitan/config/manager.py +++ b/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=`` (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() diff --git a/torchtitan/config/override.py b/torchtitan/config/override.py index c781e3e74c..00d86eeb34 100644 --- a/torchtitan/config/override.py +++ b/torchtitan/config/override.py @@ -32,9 +32,10 @@ from __future__ import annotations import importlib +import json from dataclasses import dataclass, field, fields, is_dataclass from fnmatch import fnmatch -from typing import cast, TYPE_CHECKING, TypeVar +from typing import Any, cast, TYPE_CHECKING, TypeVar from torchtitan.tools.logging import logger @@ -46,36 +47,97 @@ # Precise return type for ``derive``: the constructed ``target_config`` class. _ConfigT = TypeVar("_ConfigT", bound="Configurable.Config") +# One ``override.imports`` entry: a bare module path, or a ``(module_path, +# kwargs)`` tuple whose ``kwargs`` are forwarded to that module's override +# factories. Kept as a runtime object (not just an annotation) so the CLI parser +# can register a tyro rule keyed on ``list[OverrideImport]``. +OverrideImport = str | tuple[str, dict[str, Any]] + @dataclass(kw_only=True, slots=True) class OverrideConfig: - imports: list[str] = field(default_factory=list) + imports: list[OverrideImport] = field(default_factory=list) """ - Python module paths to import for override registration. + Modules to import for override registration, each optionally carrying kwargs. + + An entry is either: + + - a module path string, e.g. ``"torchtitan.overrides.fused_swiglu"``; or + - a ``(module_path, kwargs)`` tuple, e.g. + ``("my_pkg.triton_rope", {"block_size": 256})``. The ``kwargs`` are passed + to every override factory registered by that module when it is applied, so + two config trees can share one override module yet configure it + differently (e.g. the RL trainer and generator activate one HybridEP + dispatch module with opposite capacity factors). Each module is imported once at startup, triggering any ``@override`` decorators it defines. Only overrides registered by these modules (or their - submodules) are applied — not the entire global registry. + submodules) are applied -- not the entire global registry. - Example: ``["torchtitan.overrides.fused_swiglu", "vendor_x.kernels"]`` + On the CLI, ``--override.imports`` takes space- or comma-separated module + paths; attach kwargs to a module with ``module=`` (see + :func:`parse_cli_imports`), e.g. + ``--override.imports 'my_pkg.triton_rope={"block_size": 256}'``. See ``torchtitan/overrides/README.md`` for details. """ +def parse_cli_imports(tokens: list[str]) -> list[OverrideImport]: + """Parse ``--override.imports`` CLI tokens into ``imports`` entries. + + Each token is one of: + + - ``module=`` -- a module path whose kwargs follow the name + directly (one token), e.g. ``my_pkg.triton_rope={"block_size": 256}``; or + - a plain module path, or a comma-separated group of them (``a.b,c.d``), + for entries without kwargs. + + So ``["fused_swiglu", 'my_pkg.triton_rope={"block_size": 256}']`` parses to + ``["fused_swiglu", ("my_pkg.triton_rope", {"block_size": 256})]``. A token + with kwargs must be a single shell token, so quote it. + """ + entries: list[OverrideImport] = [] + for token in tokens: + module, sep, raw_kwargs = token.partition("=") + if sep: + # module=: kwargs attached directly to this module. A module + # path never contains '=', so the first '=' is always the separator + # and any '=' inside the JSON is preserved. + kwargs = json.loads(raw_kwargs) + if not isinstance(kwargs, dict): + raise ValueError( + f"override.imports kwargs for '{module}' must be a JSON " + f"object, got {raw_kwargs!r}." + ) + entries.append((module, kwargs)) + else: + entries.extend(part for part in token.split(",") if part) + return entries + + +def format_cli_imports(entries: list[OverrideImport]) -> list[str]: + """Serialize ``imports`` entries back to CLI tokens (inverse of the parse).""" + return [ + entry if isinstance(entry, str) else f"{entry[0]}={json.dumps(entry[1])}" + for entry in entries + ] + + @dataclass class Override: """A single registered override. - ``factory`` constructs the replacement config from the matched config. - ``fqns`` selects which matched nodes the override claims (see - :func:`override`). ``exact`` restricts the match to exactly ``target_cls`` - rather than subclasses. + ``factory`` constructs the replacement config from the matched config, plus + any keyword arguments supplied by the ``override.imports`` entry that + activated this override. ``fqns`` selects which matched nodes the override + claims (see :func:`override`). ``exact`` restricts the match to exactly + ``target_cls`` rather than subclasses. """ name: str target_cls: type[Configurable.Config] - factory: Callable[[Configurable.Config], Configurable.Config] + factory: Callable[..., Configurable.Config] fqns: list[str] | None description: str origin_module: str @@ -120,7 +182,12 @@ class (and its subclasses) are candidates; ``fqns`` narrows them. ``target``. The default ``False`` preserves subclass matching, useful when a replacement is valid for the full target contract. - The decorated function takes the matched config and returns its replacement. + The decorated function takes the matched config as its first positional + argument and returns its replacement. It may also declare keyword parameters + (or ``**kwargs``); those are filled from the ``(module_path, kwargs)`` entry + that activated the override, letting one override module be configured per + actor (e.g. trainer vs. generator). A kwarg the factory does not accept + raises ``TypeError`` at apply time (normal Python argument binding). Example — fuse MoE only on the later layers:: @@ -226,36 +293,67 @@ def triton_rope(cfg: RoPE.Config) -> TritonRoPE.Config: return cast(_ConfigT, target_cls(**kwargs)) -def _resolve_active(imports: list[str]) -> list[Override]: - """Return overrides registered by the listed import modules (or submodules). +def _resolve_active( + entries: list[tuple[str, dict[str, Any]]] +) -> list[tuple[Override, dict[str, Any]]]: + """Pair each active override with the kwargs of the entry that activated it. Provenance is the module where the factory is defined (``fn.__module__``), matched against each listed import by exact name or submodule prefix. This keeps application strictly limited to the user's ``override.imports`` even when other code paths have registered overrides into the global table. + + When several listed modules are prefixes of one override's origin (e.g. + ``"pkg"`` and ``"pkg.sub"``), the most specific (longest) one supplies its + kwargs. An entry that carries kwargs but activates no override is a mistake + (wrong module path, or its kwargs were shadowed by a more specific entry), + so it raises rather than silently dropping the kwargs. """ - prefixes = tuple(imports) - active = [] + resolved: list[tuple[Override, dict[str, Any]]] = [] + used = [False] * len(entries) for ov in _REGISTRY.values(): origin = ov.origin_module - if any(origin == p or origin.startswith(p + ".") for p in prefixes): - active.append(ov) - return active + candidates = [ + i + for i, (mod, _) in enumerate(entries) + if origin == mod or origin.startswith(mod + ".") + ] + if not candidates: + continue + best = max(candidates, key=lambda i: len(entries[i][0])) + used[best] = True + resolved.append((ov, entries[best][1])) + + for i, (mod, kwargs) in enumerate(entries): + if kwargs and not used[i]: + raise ValueError( + f"override.imports entry '{mod}' passed kwargs {sorted(kwargs)} " + "but activated no override. Check the module path and that it " + "registers an @override (and is not shadowed by a more specific " + "listed module)." + ) + return resolved @dataclass class _Claim: - """One (override, matched node) pair, collected before any mutation.""" + """One (override, matched node) pair, collected before any mutation. + + ``kwargs`` are the values from the ``override.imports`` entry that activated + the override, forwarded to its factory when the claim is applied. + """ ov: Override fqn: str cfg: Configurable.Config parent: object | None attr: str | int | None + kwargs: dict[str, Any] = field(default_factory=dict) def _collect_claims( - active: list[Override], config_root: Configurable.Config + resolved: list[tuple[Override, dict[str, Any]]], + config_root: Configurable.Config, ) -> list[_Claim]: """Traverse the original tree and gather every node each override claims. @@ -263,12 +361,21 @@ def _collect_claims( are never re-traversed, so one override cannot affect another's matches. """ claims: list[_Claim] = [] - for ov in active: + for ov, kwargs in resolved: for fqn, cfg, parent, attr in config_root.traverse(ov.target_cls): if ov.exact and type(cfg) is not ov.target_cls: continue if ov.matches(fqn): - claims.append(_Claim(ov=ov, fqn=fqn, cfg=cfg, parent=parent, attr=attr)) + claims.append( + _Claim( + ov=ov, + fqn=fqn, + cfg=cfg, + parent=parent, + attr=attr, + kwargs=kwargs, + ) + ) return claims @@ -320,7 +427,12 @@ def apply_overrides( Returns a list of human-readable log lines describing each replacement. """ - for mod_path in override_config.imports: + # Each entry is a bare module path or a (module_path, kwargs) tuple; the + # kwargs are forwarded to that module's factories (see OverrideConfig.imports). + entries: list[tuple[str, dict[str, Any]]] = [ + (e, {}) if isinstance(e, str) else (e[0], e[1]) for e in override_config.imports + ] + for mod_path, _ in entries: try: importlib.import_module(mod_path) except ImportError as e: @@ -328,13 +440,15 @@ def apply_overrides( f"Failed to import override module '{mod_path}': {e}" ) from e - active = _resolve_active(override_config.imports) - claims = _collect_claims(active, config_root) + resolved = _resolve_active(entries) + claims = _collect_claims(resolved, config_root) _check_node_conflicts(claims) replacements: list[str] = [] for c in claims: - new_cfg = c.ov.factory(c.cfg) + # ``**{}`` for a bare (no-kwargs) entry is just ``factory(cfg)``; a kwarg + # the factory does not accept raises TypeError here, naming the factory. + new_cfg = c.ov.factory(c.cfg, **c.kwargs) if isinstance(c.parent, list) and isinstance(c.attr, int): c.parent[c.attr] = new_cfg elif isinstance(c.attr, str): @@ -344,9 +458,14 @@ def apply_overrides( f"Override '{c.ov.name}' claims root config '{c.fqn}', " "which cannot be replaced in place." ) + kwargs_note = ( + " with " + ", ".join(f"{k}={v!r}" for k, v in c.kwargs.items()) + if c.kwargs + else "" + ) replacements.append( f"[Override] {c.ov.name}: {c.fqn} " - f"{type(c.cfg).__qualname__} -> {type(new_cfg).__qualname__}" + f"{type(c.cfg).__qualname__} -> {type(new_cfg).__qualname__}{kwargs_note}" ) if replacements: diff --git a/torchtitan/distributed/deepep/hybridep_override.py b/torchtitan/distributed/deepep/hybridep_override.py new file mode 100644 index 0000000000..8cf85d3112 --- /dev/null +++ b/torchtitan/distributed/deepep/hybridep_override.py @@ -0,0 +1,50 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Override: set the HybridEP MoE capacity factor per actor. + +The RL trainer and generator share one ``model_spec`` but need opposite HybridEP +dispatch modes: the trainer runs eagerly and backprops, so it needs the blocking, +dropless path (capacity factor ``None``); the generator captures a CUDA graph and +needs the static, host-sync-free non-blocking path (a float capacity factor). + +Rather than hardcode that split in the trainer, an actor activates this one module +and passes the value it wants as a per-entry ``capacity_factor`` kwarg on its own +``override``:: + + # generator.override -- non-blocking / cudagraph-safe dispatch: + OverrideConfig(imports=[( + "torchtitan.distributed.deepep.hybridep_override", + {"capacity_factor": 0.0325}, + )]) + +The ``capacity_factor`` kwarg sets ``HybridEPTokenDispatcher.Config``'s +``non_blocking_capacity_factor`` (``None`` = blocking; a float in (0, 1] = the +non-blocking capacity factor; see that config for exact semantics). The blocking +path is the config default, so the trainer needs no override; only the actor that +wants the non-blocking path activates this. ``capacity_factor`` is required (this +override exists to set it). +""" + +from __future__ import annotations + +import dataclasses + +from torchtitan.config import override +from torchtitan.models.common.token_dispatcher import HybridEPTokenDispatcher + + +@override( + "hybridep_override", + target=HybridEPTokenDispatcher.Config, + description="Set the HybridEP non-blocking capacity factor per actor.", +) +def hybridep_override( + cfg: HybridEPTokenDispatcher.Config, + *, + capacity_factor: float | None, +) -> HybridEPTokenDispatcher.Config: + return dataclasses.replace(cfg, non_blocking_capacity_factor=capacity_factor) diff --git a/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py b/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py index 6dabb18565..159b007965 100644 --- a/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py +++ b/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py @@ -640,6 +640,48 @@ def rl_grpo_qwen3_moe_debug_deepep() -> Controller.Config: return config +def rl_grpo_qwen3_moe_debug_hybridep() -> Controller.Config: + """Debug MoE config on the HybridEP backend with a cudagraph-capturable generator + (8 GPUs: 4 gen + 4 train). + + Same EP/TP/DP layout as ``rl_grpo_qwen3_moe_debug_varlen``, but the MoE uses the + HybridEP comm backend. The trainer and generator share one ``model_spec`` yet need + opposite HybridEP dispatch modes: + + - Trainer: runs eagerly and backprops, so it needs the blocking, dropless path + (capacity factor ``None``). That is the shared model_spec's default, so the + trainer needs no override. + - Generator: captures a CUDA graph, so it needs the static, host-sync-free + non-blocking path (a float capacity factor). The generator-only ``hybridep_override`` + sets it via a per-entry ``capacity_factor`` kwarg, so the two actors + diverge from one shared spec without a hardcoded per-actor branch. (This mirrors + how the ``deepep`` recipe above flips only the generator's dispatch via + ``deepep_inference``.) + + ``capacity_factor`` in (0, 1] tunes the non-blocking sizing: 1.0 is worst-case / + dropless; lower saves memory but may drop tokens. + """ + config = rl_grpo_qwen3_moe_debug_varlen() + config.model_spec = model_registry( + "debugmodel_moe", attn_backend="varlen", moe_comm_backend="hybridep" + ) + # Generator-only: switch HybridEP dispatch to the non-blocking, cudagraph-safe path + # by passing a float capacity_factor as a kwarg to the shared override module. The + # trainer keeps the shared spec's default (None -> blocking, dropless, backward-able). + config.generator.override = OverrideConfig( + imports=[ + ( + "torchtitan.distributed.deepep.hybridep_override", + {"capacity_factor": 0.0325}, + ) + ] + ) + config.generator.cudagraph = VLLMCudagraphConfig( + enable=True, mode="FULL_AND_PIECEWISE" + ) + return config + + def rl_grpo_qwen3_moe_debug_varlen_batch_invariant() -> Controller.Config: """Batch-invariant MoE EP config for bitwise parity testing (8 GPUs). diff --git a/torchtitan/overrides/README.md b/torchtitan/overrides/README.md index 49508355fa..c031b96512 100644 --- a/torchtitan/overrides/README.md +++ b/torchtitan/overrides/README.md @@ -110,14 +110,19 @@ class TritonRoPE(RoPE): @override("triton_rope", target=RoPE.Config, description="Triton rotary embedding, ~2x faster on A100+") -def triton_rope(cfg: RoPE.Config) -> TritonRoPE.Config: +def triton_rope(cfg: RoPE.Config, *, block_size: int = 128) -> TritonRoPE.Config: # `derive` copies every field shared with RoPE.Config; the factory states - # only its deltas. See "Deriving the replacement config" below. - return derive(cfg, TritonRoPE.Config, block_size=128) + # only its deltas. See "Deriving the replacement config" below. `block_size` + # defaults to 128 but can be tuned per config tree via a per-entry kwarg + # (see "Optional per-entry kwargs" below). + return derive(cfg, TritonRoPE.Config, block_size=block_size) ``` -The factory receives the matched config and returns its replacement. *Which* -instances it applies to is controlled declaratively by `fqns` (see +The factory receives the matched config as its first positional argument and +returns its replacement. It may declare keyword parameters (like `block_size` +above) that callers fill via a `(module_path, kwargs)` import entry; a bare +module-path entry uses their defaults. *Which* instances the factory applies to +is controlled declaratively by `fqns` (see [per-instance targeting](#per-instance-targeting)), not by inspecting fields inside the factory. @@ -150,20 +155,54 @@ the matched node is being replaced and detached. `derive` is most valuable when the replacement subclasses the target (the common drop-in case), where new core fields are inherited and thus carried over for free. -**Configuration is code.** An override takes no CLI flags (it runs after CLI -parsing), so the override module *is* its configuration: bake sensible defaults -into the factory (above, `block_size=128`). A user who wants different values -imports their own module that registers a factory with their deltas — e.g. -`derive(cfg, TritonRoPE.Config, block_size=256)`. We deliberately don't add a -CLI-settable `kwargs` dict (untyped, error-prone) or a typed override-`Config` -parsed at config-construction time (which would re-couple the config registry to -the override package, defeating the no-touch goal). +**Configuration is code.** By default the override module *is* its +configuration: the factory bakes in sensible defaults (above, `block_size=128`), +so a bare `override.imports` entry needs no other input. + +**Optional per-entry kwargs.** When two config trees need the *same* override +module configured *differently*, a factory can expose keyword parameters and each +`override.imports` entry can supply them as a `(module_path, kwargs)` tuple +instead of a bare string; the `kwargs` are forwarded to that module's factories. +For `triton_rope` above, two trees can pick different block sizes without a +second module: + +```python +# one tree tunes block_size=128 (the default), another 256: +OverrideConfig(imports=["my_pkg.triton_rope"]) # -> 128 +OverrideConfig(imports=[("my_pkg.triton_rope", {"block_size": 256})]) +``` + +(The RL trainer and generator use this to activate one HybridEP dispatch module +with opposite `capacity_factor` values — blocking `None` for the trainer, a float +for the cudagraph-capturing generator — instead of two modules or a hardcoded +per-actor branch.) + +The factory declares the keyword parameters it accepts (or `**kwargs`); a kwarg +it does not accept raises rather than silently no-op'ing. On the CLI, attach +kwargs to the module name as `module=` (quote it as a single shell +token; `my_pkg.triton_rope` is a placeholder for your own module): + +```bash +torchtitan_train --module llama3 --config llama3_8b \ + --override.imports 'my_pkg.triton_rope={"block_size": 256}' +``` + +JSON keeps the values typed (numbers, `null`, strings, nesting), so the same +entry works from Python or the CLI. We still avoid a typed override-`Config` +parsed at config-construction time, which would re-couple the config registry to +the override package and defeat the no-touch goal. ### Activation ```bash +# One or more modules, space- or comma-separated: torchtitan_train --module llama3 --config llama3_8b \ --override.imports torchtitan.overrides.fused_swiglu + +# A module with per-entry kwargs -- attached to the name as module=, +# quoted as one shell token (my_pkg.triton_rope is a placeholder for your module): +torchtitan_train --module llama3 --config llama3_8b \ + --override.imports 'my_pkg.triton_rope={"block_size": 256}' ``` ### Application @@ -459,7 +498,6 @@ for the full recipe. recipe from "Custom kernels and `torch.compile`". `helion` is an optional dependency, so the module imports without it and falls back to the PyTorch RoPE when it (or CUDA) is unavailable; it is checkpoint-compatible with stock. - The `TritonRoPE` snippets above are illustrative — no `triton_rope.py` is shipped — but RoPE is a fully valid override target (`helion_rope.py` is a real one): each attention module owns a `rope` submodule (`RoPE.Config`), so a custom @@ -474,7 +512,7 @@ RoPE is an ordinary component override. | `torchtitan/protocols/model_spec.py` | `ModelSpec.traverse` exposes the nested model config to the traversal. | | `torchtitan/trainer.py` | Holds the `override` config field; applies overrides after `update_from_config`, before builds. | | `torchtitan/overrides/` | In-repo example implementations (`fused_swiglu.py`, `helion_rope.py`). | -| `tests/unit_tests/test_override.py` | Unit tests: registration, provenance, FQN / exact targeting, per-node conflicts, `derive`. | +| `tests/unit_tests/test_override.py` | Unit tests: registration, provenance, FQN / exact targeting, per-node conflicts, per-entry kwargs, `derive`. | Overriding a component requires no changes to any model's `config_registry.py` or `__init__.py` — that is the point of the design. @@ -482,7 +520,10 @@ or `__init__.py` — that is the point of the design. ## At a Glance - **Activation.** Explicit opt-in via `override.imports`; only overrides - registered by the listed modules (provenance-checked) are applied. + registered by the listed modules (provenance-checked) are applied. An entry + may carry kwargs (a `(module_path, kwargs)` tuple in Python, or + `module=` on the CLI) to configure the same module differently per + config tree. - **Conflicts.** Per-node, not per-class: overrides may share a target class as long as their claimed nodes are disjoint; same-node or ancestor/descendant claims error.