Conversation
…Config to BaseModel Second PR in the utils-layer config v1 -> v2 series (NVIDIA-NeMo#3916 migrated ClusterConfig; pattern from NVIDIA-NeMo#2518, alongside NVIDIA-NeMo#3197/NVIDIA-NeMo#3289). Schema: - Both classes become BaseModel(extra="allow"); NotRequired fields become Optional[...] = None. - Hidden call-site defaults are centralized on the class exactly as they behave today: save_optimizer=True, save_data_plane=False, save_consolidated=False, model_cache_dir="", model_repo_id="", is_peft=False, is_async=False. - model_save_format stays Optional[str] = None on purpose: the DTensor v1 worker rejects any non-None value, and the automodel layers apply their own "safetensors" fallback (CheckpointManager keeps that resolution, now explicit instead of a .get default). - load_replay_buffer stays tri-state (Optional[bool] = None): the async-GRPO restore path only opts out on an explicit False. Call sites (~85 across all algorithms, SC utils, lm_policy) move from key access / .get() to attribute access. Two deliberate dict boundaries are preserved with model_dump(): - policy_config["pretrained_checkpoint"] hand-off (PolicyConfig is still TypedDict-shaped; Megatron setup consumes dict keys), and - checkpointing_cfg forwarded to save_checkpoint workers (the automodel checkpoint layer consumes .items()/.get()). Tests: fixtures built via model_construct bypass validation, so checkpointing dict literals are wrapped in CheckpointingConfig.model_construct(**{...}) (missing required fields now raise AttributeError where they raised KeyError before); CheckpointManager unit tests validate their full config dicts via model_validate. No exemplar YAML or reference-config changes; class defaults mirror the documented/exemplar values, so tests/unit/test_config_v2.py is unaffected. Verified: py_compile on all touched files; standalone pydantic semantics checks (exemplar blocks round-trip, sft.yaml block keeps model_save_format=None, nested pretrained_checkpoint coerces and dumps back, extra keys preserved, model_construct default-fill); ruff format + check clean. Signed-off-by: Shubhankar Tripathy <reach2shubhankar@gmail.com>
Contributor
Author
|
/ok to test d99873c |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Second PR in the utils-layer config v1 → v2 series (after #3916 /
ClusterConfig; pattern from #2518, running alongside #3197 policy-config and #3289 generation-config series). MigratesCheckpointingConfigandPretrainedCheckpointConfig(nemo_rl/utils/checkpoint.py) fromTypedDicttoBaseModel(extra="allow").Schema decisions (each mirrors today's effective behavior — no user-visible default changes):
save_optimizer=True,save_data_plane=False,save_consolidated=False,model_cache_dir="",model_repo_id="",is_peft=False,is_async=False— this removes the.get(key, default)forbidden-pattern instances inCheckpointManager.__init__.model_save_formatstaysOptional[str] = None:Noneis load-bearing (the DTensor v1 worker rejects any non-None value); the automodel layers keep applying their own"safetensors"fallback, now written explicitly inCheckpointManagerinstead of hiding in a.getdefault.load_replay_bufferstays tri-state (Optional[bool] = None): the async-GRPO restore only opts out on an explicitFalse, so an omitted key must remain distinguishable fromtrue.save_optimizer/keep_top_kpreviously hit a latentKeyError(config["save_optimizer"]on aNotRequiredkey); they now get the documented defaults.Call sites: ~85 across all algorithms, SC utils, and
lm_policymove to attribute access. Two dict boundaries are deliberately preserved viamodel_dump()so neighboring migration lanes stay untouched:policy_config["pretrained_checkpoint"]hand-off (PolicyConfig is still TypedDict-shaped;megatron/setup.pyconsumes dict keys — @NolenLiang's lane), andcheckpointing_cfgforwarded tosave_checkpointworkers (the automodel checkpoint layer consumes.items()/.get()).Tests: fixtures built via
model_constructbypass validation, so checkpointing dict literals are wrapped inCheckpointingConfig.model_construct(**{...});CheckpointManagerunit tests validate their full config dicts viamodel_validate.No exemplar YAML or reference-config changes;
tests/unit/test_config_v2.pyis unaffected.Issues
Part of the TypedDict → BaseModel migration tracked by
tests/unit/test_config_v2.py(see also #1675 / #2102).Usage
No user-facing YAML change. Unknown/legacy keys under
checkpointing:continue to load; missing required keys (enabled,checkpoint_dir,metric_name,higher_is_better,save_period) now fail at config validation time with a clear pydantic error.Before your PR is "Ready for review"
model_save_format=None; nestedpretrained_checkpointcoerces and dumps back; extra keys preserved;model_constructdefault-fill)ruff format+checkclean (repo-pinned 0.9.9);py_compileon all touched files; zero3rdparty/changes (gitlinks verified againstmain)