NVIDIA-NeMo / NVIDIA-NeMo/RL

Use `dataclass` instead of `TypedDict` for config

Open
#1,675 4 comments 1 reaction 1 assignee Claimed by @yuki-97 View on GitHub
community-request enhancement t-dataclass waiting-on-maintainers x-tencent
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

Currently, nemo-rl uses `TypedDict` to host the config options:
https://github.com/NVIDIA-NeMo/RL/blob/b238e41a6bad66324635ea6a65ef8bd8aeb8165e/nemo_rl/algorithms/grpo.py#L115-L177

We can consider to use `dataclass` to host these options like:
```python
@dataclass
class GRPOConfig:
num_prompts_per_step: int
num_generations_per_prompt: int
max_num_epochs: int
max_num_steps: int
# ...
```

It has several benifits, like, more friendly for type checking for `pyrefly`/`mypy`, and ease of development, e.g., you can directly access `config.logger.log_dir` with type hints in IDE/editor, instead of typing `config["logger"]["log_dir"]`. We may also add some runtime validation logic in the `__post_init__` dunder for dataclass.

The dataclass could be converted from omegaconf using `to_object` method.

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.