[Tracking] checkpoint I/O architecture, correctness, performance, and memory

Status — September 17, 2026

The core checkpoint loading and save/export work is merged. All seven implementation subissues are closed; four follow-up subissues remain open. All PRs in the original implementation stack have landed, including #3619 and #3623.

This tracker now covers maintenance follow-ups. Prioritize reproduced correctness issues and improvements needed by maintained recipes. Broader refactoring, benchmarking, and speculative optimization are deferred unless a concrete workload justifies them.

Completed

• Adapter-owned direct-load safety contract — #3574; subissue #3662.
  Distinguish final model storage from allocating/converting destinations and retain a safe fallback.
• Grouped-MoE loading overhead — #3580; subissue #3663.
  Fill grouped expert tensors directly, avoid repeated full-heap GC, reuse safe checkpoint views, and read rank-local Gemma4 experts.
• Low-memory single-GPU DCP for custom MoE adapters — #3610; subissue #3664.
  Avoid full CPU checkpoints for supported Nemotron/Ling configurations, with constructor-aligned storage capability and load bookkeeping safeguards.
• Direct DCP loading for ordinary HF/external safetensors — #3616; subissue #3665.
  Use final model storage for supported identity-layout adapters and select conversion routes from checkpoint quantization metadata.
• Bounded FP8 conversion — #3619; subissue #3666.
  Load, convert, and release small decoder groups for Mistral3/Devstral checkpoints, supporting BF16 and FP32 destinations.
• Bounded one-GPU GPT-OSS MXFP4 loading — #3623; subissue #3667.
  Convert one decoder layer at a time, avoiding a complete BF16 CPU model; retain the existing distributed route.
• Grouped-MoE save/export GC cleanup — #3621; subissue #3668.
  Remove repeated full-heap GC scans while preserving checkpoint output and memory behavior.

The earlier direct-route restoration (#3533) is also merged. Detailed model-specific correctness and performance evidence remains in the linked PRs.

Recent validation

• Devstral 24B LoRA on one GB10: scoped release job passed, including training, validation, and checkpoint saving for #3619.
• GPT-OSS 20B LoRA on one GB10: scoped release job passed on #3623's final head: 50 training steps, validation, and checkpoint saving; 25 load parts with a largest compressed temporary allocation of 0.39 GiB. PR CI passed, including CPU and both GPU unit-test shards.

These GB10 runs validate the maintained recipes. They do not replace controlled cold/warm-cache benchmarks or measurements of total unified-system memory.

Remaining follow-ups

• #3669 — Clean up and take ownership of HF/DCP checkpoint storage.
  Remove verified-unused backports and clarify ownership/quality checks for maintained storage code. Behavior-neutral cleanup; defer broader relocation unless it has a clear maintenance benefit.
• #3670 — Profile consolidated checkpoint export storage overhead.
  Shard lookup, file-open/mmap churn, and buffer-copy costs remain hypotheses. Profile only when a representative export has a material bottleneck; implement a change only if measurements justify it.
• #3688 — Support grouped expert bias in checkpoint split/merge adapters.
  Add correct bias merge/split and round-trip coverage when a maintained configuration needs it. Unsupported expert-bias conversion currently fails explicitly; router correction bias is separate.
• #3689 — Centralize MoE runtime backend capability truth for checkpoint adapters.
  Keep construction and checkpoint destination selection consistent, validate backend/dispatcher combinations, and audit composite-adapter namespaces and view-loaded keys. Prioritize concrete correctness gaps; keep any shared-helper refactor focused and component boundaries acyclic.

These are follow-ups to the merged work, not outstanding steps in the original PR stack. Acceptance criteria and detailed scope live in each subissue. Keep this umbrella open while these four remain tracked.

Design and acceptance constraints

• Model adapters own model-specific names, layouts, scale semantics, and conversion math; shared checkpoint code owns storage and execution.
• Select the smallest safe route: direct loading into final storage, bounded transformation, or the materialized compatibility fallback.
• Preserve exact values for lossless changes; compare precision conversions against a trusted reference. Validate the affected load/save/reload or training path, key coverage, tied weights, and relevant distributed layouts.
• Measure performance on a semantically correct baseline and record the model, dtype, topology, storage, and revision. On unified-memory systems, CUDA allocator statistics alone do not establish total physical-memory use.

Deferred investigations and retained review notes

The earlier broad roadmap is not an active implementation commitment. Revisit only for a reproduced correctness issue or a measured bottleneck in a maintained workload:

• Broader model/adapter audits, distributed quantized-load grouping, a unified-memory cold/warm-cache benchmark matrix, generalized bounded export, or bounded read-ahead.
• GB10 CI cache/storage throughput improvements belong to a separate CI infrastructure follow-up.
• Earlier #3610 review notes: resolve alternate checkpoint namespaces before DCP constructs destinations, and give _inplace_loaded_native_keys an attempt-scoped lifecycle at the loader boundary. These notes are retained for future targeted investigation; they are not marked resolved by the merged performance work. Composite-adapter capability and namespace auditing is covered by #3689.

Background

Earlier Nemotron profiling and thread-budget mitigation: #3160 / #3370. The merged work replaces recipe-specific tuning with adapter-owned safety and bounded loading where needed, while preserving compatibility fallbacks.