NVIDIA-NeMo / NVIDIA-NeMo/Automodel
Clean up and take ownership of HF/DCP checkpoint storage
Open
@yuhezhang-ai is already working on this.
Since Aug 25, 2026.
automodel
ckpt
enhancement
Performance
- Dominant language
- Python
- Stars
- 960
- Forks
- 316
- Avg merge
- 3d 20h
- Merged PRs (30d)
- 143
Description
Tracks the HF/DCP storage ownership cleanup identified under NVIDIA-NeMo/Automodel#3576.
Problem
components/checkpoint/_backports mixes historical compatibility code with maintained AutoModel-specific storage behavior. Three files are verified unused on current main: default_planner.py, planner_helpers.py, and _version.py (about 697 lines total). The remaining active implementation is blanket-excluded from coverage and type checking.
Scope
- Remove the three verified-unused files in a focused behavior-neutral PR.
- After the current checkpoint stack lands, move maintained implementation out of
_backportsinto a first-class internal package. - Replace blanket coverage and type-check exclusions with normal validation or narrow documented compatibility exceptions.
- Centralize duplicated safetensors header and dtype handling where ownership becomes clear.
Acceptance criteria
- No import or runtime references remain to removed files, and focused checkpoint tests pass.
- The public checkpoint API and load/save behavior are unchanged.
- Maintained code has explicit ownership without blanket quality-gate exclusions.
- Minimum PyTorch 2.6 compatibility is preserved, with any version-specific shims kept bounded.
- Relocation is split into focused PRs if doing it together would create review noise or conflict with the current stack.
Out of scope
- Replacing the AutoModel reader/writer wholesale with upstream PyTorch.
- Changing performance paths in the current stacked PRs.
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.
Assessment
This issue has not been assessed yet.