NVIDIA-NeMo / NVIDIA-NeMo/Automodel
Profile consolidated checkpoint export storage overhead
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 export-storage profiling follow-up identified under NVIDIA-NeMo/Automodel#3576.
Problem
PR #3621 removes repeated full GC scans but does not address possible storage-layer overhead in consolidate_hf_safetensors. Repeated shard lookup, file-open or mmap churn, and buffer copies are hypotheses rather than confirmed bottlenecks.
Scope
- After PR #3621, profile representative consolidated exports.
- Attribute time and memory to shard lookup, open/mmap churn, page faults, tensor-fragment reads and copies, gathering, and conversion.
- If measured, evaluate an FQN-to-source-shard index, bounded file/mmap handle reuse, and zero-copy memory views where lifetime is safe.
- Implement only a measured improvement in a focused follow-up PR.
Acceptance criteria
- Document the representative model, checkpoint layout, storage location, topology, and exact revision.
- Record end-to-end export time, phase timings, host/device memory, page faults, and output size.
- Verify output indexes and tensors against the reference; preserve tied weights, distributed gathering, asynchronous buffer lifetime, bounded resources, and atomic completion.
- Clean up reused handles or mmaps deterministically.
- If the suspected overhead is not material, record the negative result and close without a production change.
Dependency
Use PR #3621 as the baseline and keep this work separate from that PR.
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.