Worker-class overrides are split across three mechanisms instead of one config field
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
### Summary
Three separate mechanisms decide which Ray actor class a policy or generation worker group runs, and they don't compose — the later one just overwrites the earlier one's result. I'd like them unified onto the config field `worker_extension_cls_fqn`.
### Current state
| # | Mechanism | Where it decides | Added by |
|---|---|---|---|
| 1 | `quant_cfg` | [`resolve_policy_worker_cls`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/models/policy/utils.py#L120) / [`resolve_generation_worker_cls`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/models/generation/vllm/utils.py#L637), via the hard-coded [`POLICY_WORKER_OVERRIDES`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/models/policy/utils.py#L113) / [`GENERATION_WORKER_OVERRIDES`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/models/generation/vllm/utils.py#L631) maps | pre-existing |
| 2 | `policy.worker_extension_cls_fqn` / `policy.generation.worker_extension_cls_fqn` | overwrite the resolver's result, after the fact | #3809 |
| 3 | `Policy(worker_extension_cls_fqn=...)` constructor arg | same overwrite site ([`lm_policy.py:240-244`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/models/policy/lm_policy.py#L240)) | pre-existing; only in-tree caller is [`single_update.py:105`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/research/template_project/single_update.py#L105) |
### Why it's worth consolidating
- (1) and (2) both replace the resolved class, so they're mutually exclusive by a hand-written check that is duplicated in two files — and on the policy side it guards only the config channel, not the constructor one.
- (3) duplicates (2) with its own precedence rule plus a reconciliation `raise`, so the same FQN can be accepted or rejected depending on which channel delivered it.
- Both resolvers already carry the same TODO asking for exactly this ([`policy/utils.py:109-112`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/models/policy/utils.py#L109), [`vllm/utils.py:627-630`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/models/generation/vllm/utils.py#L627)): *"Replace this hard-coded map with a generic plugin-registration hook … so core has no knowledge of ModelOpt-specific worker classes."*
### Proposed direction
Make the config field the only thing that decides the worker class, and turn `quant_cfg` into a *constraint* on it rather than a second substitution mechanism.
1. Reconcile `worker_extension_cls_fqn` against `quant_cfg` in one place, with three cases:
- **not set** — fill it in from [`POLICY_WORKER_OVERRIDES`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/models/policy/utils.py#L113) / [`GENERATION_WORKER_OVERRIDES`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/models/generation/vllm/utils.py#L631) and warn, so the recipe still runs and the user is told which class they got;
- **set to the matching quant worker** — proceed;
- **set to anything else** — raise, naming the FQN that was expected.
The two maps stop being substitution tables applied behind the user's back and become the source for that fill-and-validate step, so every worker-class decision ends up expressed in the config field.
2. Deprecate the `Policy` constructor argument and migrate `single_update.py` to set `policy.worker_extension_cls_fqn` in [`research/template_project/configs/grpo_math_1B.yaml`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/research/template_project/configs/grpo_math_1B.yaml) instead.
This removes the footgun the current inline comment describes — that a config author cannot know which worker `quant_cfg` resolved to — because the recipe either states the class or is told what it got.
### Migration notes
- The quant workers are already in [`MODELOPT_ACTOR_REGISTRY`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/modelopt/registry.py#L37), which is merged into `ACTOR_ENVIRONMENT_REGISTRY` on import, so naming one from YAML resolves its venv today. No registration work is needed for this migration.
- Not a breaking change: the recipes that set only `quant_cfg` (11 files today, most of them setting it under both `policy` and `policy.generation`) keep working and just warn. Naming the FQN explicitly in those recipes is a follow-up that silences the warning.
- The expected FQN depends on backend, `dtensor_cfg._v2` and `vllm_cfg.async_engine` — which is exactly why the unset case should auto-fill rather than hard-require: someone who flips `async_engine` should not have to remember to change the FQN as well. Same reason the error in the third case should name the FQN it expected instead of just rejecting.
- [`teacher_worker_group.py:179-184`](https://github.com/NVIDIA-NeMo/RL/blob/a366bc8cffec730080354415d74c2c12462ea028/nemo_rl/models/policy/teacher_worker_group.py#L179) drops a student-side `quant_cfg` to run the teacher unquantized; it would also need to clear an explicitly-set `worker_extension_cls_fqn`, or the teacher gets a quant worker with no quant config.
Contributor guide
Assessment
This issue has not been assessed yet.