NVIDIA-NeMo / NVIDIA-NeMo/Automodel
Configure checkpoint saving behavior through YAML
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 960
- Forks
- 316
- Avg merge
- 3d 20h
- Merged PRs (30d)
- 143
Description
Is your feature request related to a problem? Please describe.
Currently, checkpoint saving behavior is defined by the base recipe. In particular, checkpoints include optimizer state by default, and users need to create or modify a custom recipe script if they want to disable saving the optimizer state.
This makes it unnecessarily difficult to configure checkpoint behavior for common use cases such as fine-tuning or inference-oriented training, where only the model weights are needed.
It would be useful to expose checkpoint-saving options through the YAML recipe so users can change this behavior without writing a custom Python recipe.
Describe the solution you'd like
I'd like to be able to configure checkpoint saving behavior directly in the YAML recipe.
For example, something along the lines of:
checkpoint:
save_optimizer: false
or an equivalent configuration supported by the existing checkpointing implementation.
With this configuration, the training script would save only the model state (and other required metadata) without saving optimizer state.
The goal is to make checkpoint behavior configurable without requiring users to create a custom recipe script.
Describe alternatives you've considered
One current alternative is to create a custom recipe/script that modifies the checkpoint configuration or saving logic. However, this adds unnecessary Python code for what is essentially a configuration-level option.
Another option is to keep the current behavior and manually remove optimizer state after training, but this still requires additional processing and wastes storage during training.
Additional context
This would be particularly useful for fine-tuning workflows where the optimizer state is not needed after training and can significantly increase checkpoint size ( especially with large model ).
Ideally, the existing checkpoint-saving implementation could expose these options through the YAML configuration while preserving the current defaults for backward compatibility.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by tracing the existing checkpoint-saving implementation and the YAML recipe parsing entry point to find how optimizer state is currently selected. Done means a YAML option controls optimizer-state saving, current defaults remain unchanged, and coverage verifies both enabled and disabled configurations.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100