Skip to content

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

Description

@terrykong

Context

During review of #2420 (TensorRT-LLM rollout backend), it was suggested 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).

The author pointed out that the base GenerationConfig (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):

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

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions