Skip to content

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

Open
wwwjn wants to merge 1 commit into
mainfrom
override-flexible
Open

[override] Allow passing kwargs in override from both CLI and config#3894
wwwjn wants to merge 1 commit into
mainfrom
override-flexible

Conversation

@wwwjn

@wwwjn wwwjn commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

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:

  • trainer runs eagerly + backprops → needs the blocking, dropless path
    (capacity_factor=None)
  • generator captures a CUDA graph → needs the static, host-sync-free non-blocking
    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:

  1. Python config — same module, different value per actor:
    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
  2. CLI — kwargs as JSON attached to the name:
    --override.imports 'my_pkg.triton_rope={"block_size": 256}'
  • Factory declares the keyword params it accepts (or **kwargs); an unknown kwarg
    raises a normal TypeError; an entry whose kwargs activate no override raises (no
    silent no-op).
  • Ships the concrete consumer: hybridep_override (sets
    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.

@wwwjn wwwjn requested review from fegin, tianyu-l and wconstab as code owners July 9, 2026 22:26
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 9, 2026
name: str
target_cls: type[Configurable.Config]
factory: Callable[[Configurable.Config], Configurable.Config]
factory: Callable[..., Configurable.Config]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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"``).

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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"]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

seems not a good location. Maybe put in overrides folder as you are doing for deepep inference override?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/rl ciflow/8gpu CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants