[override] Allow passing kwargs in override from both CLI and config#3894
[override] Allow passing kwargs in override from both CLI and config#3894wwwjn wants to merge 1 commit into
Conversation
| name: str | ||
| target_cls: type[Configurable.Config] | ||
| factory: Callable[[Configurable.Config], Configurable.Config] | ||
| factory: Callable[..., Configurable.Config] |
There was a problem hiding this comment.
can we improve this? At least the first arg should be config
| """Decorator to register an override factory. | ||
|
|
||
| Args: | ||
| name: Unique identifier for this override (e.g. ``"triton_rope"``). |
There was a problem hiding this comment.
The wrapped function itself (import path + function name) should already be "unique identifier". Can we deprecate this name field?
| ``"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. |
There was a problem hiding this comment.
Oh, I don't like this ambiguity.
It's caused by that imports requires a module. Do you think we can let imports contain the exact (path to) the function? If so we don't need the special treatment here.
E.g. instead of imports = ["torchtitan.distributed.deepep.hybridep_override"] can we do imports = ["torchtitan.distributed.deepep.hybridep_override.hybridep_override"]
There was a problem hiding this comment.
seems not a good location. Maybe put in overrides folder as you are doing for deepep inference override?
Why
The override mechanism (override.imports) could only activate a module by name — the
factory had no way to receive parameters. The documented workaround was
"configuration is code": if you wanted different values, you wrote a second override
module.
That breaks down when two config trees that share one model_spec need the same
override configured differently. The motivating case is RL: the trainer and
generator share one model spec but need opposite HybridEP MoE dispatch modes:
(capacity_factor=None)
path (a float capacity_factor)
With name-only overrides, the only options were (a) two near-identical modules, or
(b) a hardcoded if hybridep: capacity_factor = None branch in the trainer's
_build_model (what PR #3871 had to do). Both are avoidable if the override can just
take a parameter
What
An override.imports entry can now be either a bare module path or a (module_path,
kwargs) tuple; the kwargs are forwarded to that module's factory:
generator.override = OverrideConfig(imports=[(
"torchtitan.distributed.deepep.hybridep_override", {"capacity_factor": 0.0325},
)]) # trainer needs no override — capacity_factor=None is the shared spec's default
--override.imports 'my_pkg.triton_rope={"block_size": 256}'
raises a normal TypeError; an entry whose kwargs activate no override raises (no
silent no-op).
HybridEPTokenDispatcher.Config.non_blocking_capacity_factor from a capacity_factor
kwarg) + RL recipe rl_grpo_qwen3_moe_debug_hybridep — trainer uses the default
blocking path, generator activates the override with a float, no hardcoded per-actor
branch.