Migrate GenerationConfig and backend generation configs (vllm/sglang/trtllm) from TypedDict to pydantic BaseModel
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Context
During review of #2420 (TensorRT-LLM rollout backend), [it was suggested](https://github.com/NVIDIA-NeMo/RL/pull/2420#discussion_r3580489863) that the new `TrtllmConfig` / `TrtllmSpecificArgs` in `nemo_rl/models/generation/trtllm/config.py` be written as `pydantic.BaseModel` per the v2 config convention (see #2325 and [design-and-philosophy.md — Configuration Schema](https://github.com/NVIDIA-NeMo/RL/blob/main/docs/design-docs/design-and-philosophy.md#configuration-schema-basemodel-dataclass-and-typeddict)).
The author [pointed out](https://github.com/NVIDIA-NeMo/RL/pull/2420#discussion_r3585558675) that the base `GenerationConfig` ([`nemo_rl/models/generation/interfaces.py`](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/generation/interfaces.py)) is still a `TypedDict`, so the backend configs that extend it (`VllmConfig`, `SglangConfig`, and now `TrtllmConfig`) can't individually migrate cleanly. This was accepted as a non-blocker for #2420, to be tracked as follow-up work.
## Task
Migrate the generation config hierarchy from `TypedDict` (v1) to `pydantic.BaseModel` (v2):
- [ ] `GenerationConfig` (base) in `nemo_rl/models/generation/interfaces.py`
- [ ] `VllmConfig` / `VllmSpecificArgs`
- [ ] `SglangConfig` / `SglangSpecificArgs`
- [ ] `TrtllmConfig` / `TrtllmSpecificArgs` (added in #2420)
Per the v2 convention: declare with `extra="allow"`, move defaults onto the BaseModel fields (exemplar YAMLs remain as documentation), and keep `tests/unit/test_config_v2.py` / `tests/unit/reference_configs/` green.
## References
- Migration pattern: #2325
- Originating discussion: https://github.com/NVIDIA-NeMo/RL/pull/2420#discussion_r3585558675
Contributor guide
Assessment
This issue has not been assessed yet.