NVIDIA-NeMo / NVIDIA-NeMo/RL

refactor: leverage MCore TransformerConfig.__post_init__ validation for config fields set in _apply_performance_config

Open
#2,291 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Context

When `_apply_performance_config()` (and similar `_apply_*` functions in `nemo_rl/models/megatron/setup.py`) set fields on `model_cfg`, they do so via plain attribute assignment **after** the `TransformerConfig` dataclass has already been constructed and `__post_init__` has run. This means MCore's built-in validation (e.g., checking `recompute_modules` against the [allowed set](https://github.com/NVIDIA/Megatron-LM/blob/d30c3ae5469fe3f6a64d4fd2e63b6e7f7844ea81/megatron/core/transformer/transformer_config.py#L1363-L1389)) is bypassed for user-supplied values.

### Call chain

1. [`MegatronPolicyWorker.__init__`](https://github.com/NVIDIA-NeMo/RL/blob/6621d76218009f065e4622e03a3f34b3398768fe/nemo_rl/models/policy/workers/megatron_policy_worker.py#L150) → `validate_and_set_config()`
2. [`validate_and_set_config()`](https://github.com/NVIDIA-NeMo/RL/blob/6621d76218009f065e4622e03a3f34b3398768fe/nemo_rl/models/megatron/setup.py#L223) → `setup_model_config()`
3. [`setup_model_config()`](https://github.com/NVIDIA-NeMo/RL/blob/6621d76218009f065e4622e03a3f34b3398768fe/nemo_rl/models/megatron/setup.py#L308) calls `ConfigContainer.from_yaml()` — `__post_init__` runs here
4. [`_apply_performance_config()`](https://github.com/NVIDIA-NeMo/RL/blob/6621d76218009f065e4622e03a3f34b3398768fe/nemo_rl/models/megatron/setup.py#L341) mutates `model_cfg` via attribute assignment — `__post_init__` does **not** re-run

### Problem

A typo in a user config (e.g., `recompute_modules: ["mopee"]`) is silently accepted. MCore's runtime checks match module names via string comparison, so an unrecognized name is a silent no-op — the user thinks they're saving memory but nothing is actually recomputed.

### Proposal

Consider refactoring the config initialization flow so that MCore's `__post_init__` validation runs **after** all `_apply_*` mutations are complete. Options include:

- Constructing the `TransformerConfig` from a merged dict (user overrides applied first, then construct once)
- Re-running `__post_init__()` after all mutations (need to verify safety — it also sets defaults)
- Extracting MCore's validation into a callable utility and invoking it post-mutation

This would eliminate the need to duplicate validation logic in NeMo-RL and ensure any future MCore validation additions are automatically picked up.

### Related

Surfaced during review of PR #2280 (selective activation checkpointing).

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.