NVIDIA-NeMo / NVIDIA-NeMo/Automodel

Clean up and take ownership of HF/DCP checkpoint storage

Open
#3,669 0 comments 0 reactions 1 assignee View on GitHub

@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 _backports into 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.