NVIDIA-NeMo / NVIDIA-NeMo/RL

Migrate GenerationConfig and backend generation configs (vllm/sglang/trtllm) from TypedDict to pydantic BaseModel

Open
#3,289 0 comments 0 reactions 1 assignee Claimed by @shuyixiong View on GitHub
Feature
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.