Roadmap: checkpoint I/O architecture, correctness, performance, and memory

[yuhezhang-ai]

Summary

Track the multi-PR effort to make checkpoint I/O correct, fast, and memory-bounded across model families, checkpoint layouts, distributed configurations, and hardware memory architectures.

This roadmap covers more than base-model load latency. It includes:

• Base-model loading and training resume.
• MoE tensor merging and other layout conversions.
• FP8 dequantization, scale handling, and other precision transforms.
• Discrete CPU/GPU memory and unified-memory systems such as DGX Spark.
• Checkpoint save and consolidated export time and memory.
• End-to-end correctness and observability.

This is broader than the earlier recipe/thread-budget mitigation in #3160 and #3370. The current loading PRs are concrete first steps, not the complete architecture.

Motivation

A representative 58.82 GiB Nemotron Nano V3 checkpoint took about 581 seconds in an early standalone single-GPU profile. A later scheduled job before #3533 loaded the same amount in 165.45 seconds, while the restored direct route loaded it in 30.36 seconds (5.45x faster, an 81.7% reduction). #3533 restored that route on current main; #3574 adds the safety contract needed to use it only when checkpoint destinations truly refer to final model storage.

Profiling also showed that opening a checkpoint is not the whole cost. Physical reads may happen later during model conversion, where page faults, repeated tensor merging, dtype conversion, destination construction, and final model copies can dominate. A recipe-specific CPU thread change helped one case but regressed another, so it is not a general solution.

The broader problem includes several paths with different requirements:

1. Direct loading: checkpoint data can be written into final model parameters or safe views of their storage.
2. Streaming transformation: related tensors are read in a small group, transformed into final model storage, and released before continuing. Examples include MoE merging and FP8 weight-plus-scale conversion.
3. Materialized fallback: the full checkpoint is loaded and converted before installation. This remains the conservative compatibility path when direct or streaming execution is not yet safe.

The goal is not to force every model through the same path. The goal is to select the smallest safe path and make the remaining fallback explicit and measurable. Route selection depends on the adapter, checkpoint metadata, backend, dtype, and topology; the same model family may use direct loading in one configuration and an allocating fallback in another.

Design direction

Keep the existing public checkpointer surface, but separate four internal responsibilities:

• Storage: discover keys, open shards, read and write tensors, and manage output files.
• Model adapters: own model-specific names, shapes, layouts, scale semantics, and conversion math.
• Execution strategy: choose and run direct, streaming, or materialized loading; saving has its own execution flow.
• Validation and measurement: reconcile keys, preserve tied weights, check results, and report time and memory consistently.

Shared checkpoint code should not own model-specific FP8 or MoE policy. Model adapters should not own filesystem or distributed I/O policy.

Do not introduce a general tensor-operation graph up front. Implement the first real hybrid conversion with the smallest model-owned change that works. Extract a reusable execution plan only if the concrete implementation shows at least one of the following:

• The same orchestration is needed by multiple adapters.
• It materially improves load time or peak memory.
• It enables a model that otherwise runs out of memory.
• It removes duplicated correctness-critical conversion logic.

Work plan

• Gate direct loading with an adapter-owned safety contract — refactor(checkpoint): gate direct loads by adapter capability #3574
  • Preserve the direct route for audited models whose destinations refer to final model storage.
  • Return allocating and dequantizing cases to the conservative fallback.
  • Open each indexed fallback shard once per file instead of once per tensor key.
  • Keep the PR focused on safe routing and representative coverage.
• Measure and improve the allocating MoE fallback — perf(checkpoint): reduce allocating grouped MoE load overhead #3580
  • Split fallback timing into disk read, adapter conversion and cleanup, and final installation phases.
  • Replace avoidable per-expert merge temporaries with one final allocation filled directly.
  • Reduce repeated full-heap collection in the hot loop and avoid unrelated CUDA allocator trimming for host-side merges.
  • Benchmark a model and backend that actually enter the allocating path, and attribute direct-fill and cleanup effects separately.
• Standardize correctness and performance evidence
  • Record phase timings, effective throughput, peak memory, page faults, and swap activity where applicable.
  • Establish a small representative matrix covering direct, allocating MoE, precision-transform, and distributed paths.
  • Make correctness evidence a requirement for every optimization PR rather than a final cleanup phase.
• Eliminate avoidable destination construction work
  • Measure destination construction separately from storage read and adapter conversion.
  • For a concrete non-aliasing conversion, avoid copying initialized tensor contents into load buffers that the checkpoint will fully overwrite.
  • Allocate final buffers without initialization only after proving that every element is written and ownership and lifetime are safe.
• Implement the first streaming transformation case
  • Use one concrete case such as the single-device Gemma4 MoE transpose-and-scale conversion or FP8 weight-plus-scale dequantization.
  • Read one logical dependency group, transform it into final model storage, and release source and scratch memory.
  • Keep only a bounded number of logical dependency groups live at once.
  • Group tensors by conversion dependency rather than by physical checkpoint shard.
• Decide whether a deeper execution-plan refactor is justified
  • Prefer a small internal strategy decision and focused executors over a large framework.
  • Unify indexed shard reading for non-direct consumers when the streaming case needs it.
  • Centralize key reconciliation and load-path reporting when doing so removes real duplication.
• Benchmark unified-memory behavior on DGX Spark
  • Compare a direct-load model, an allocating MoE model, and the streaming prototype.
  • Measure system memory availability, process memory, page cache, swap activity, page faults, and time to first successful forward pass. Treat device allocator metrics as supplementary rather than total physical memory.
  • Test cold-cache and warm-cache conditions.
  • Do not add a DGX-specific loader unless measurements show the generic direct or bounded-streaming paths are insufficient.
  • Target peak physical memory close to final model storage plus a small bounded conversion window, with no swapping and identical model results.
• Improve checkpoint save and consolidated export
  • Separate training checkpoint persistence from consolidated model export internally while keeping the public API stable.
  • Avoid retaining a full second converted model when non-contiguous or transformed tensors must be serialized.
  • Stream conversion and writing by logical tensor group or output shard where practical.
  • Preserve distributed gathering, tied weights, output indexes, asynchronous buffer lifetime, and atomic completion.
• Audit and expand model coverage
  • Promote rename/slice/transpose-only adapters to direct loading after proving their destinations are safe.
  • Refactor simple allocating conversions only when the code and measurements justify it.
• Evaluate bounded read-ahead only if a remaining path is I/O-bound
  • Avoid unconditional whole-checkpoint prefetch.
  • Add bounded, per-node read-ahead only when a representative path shows a clear gain without higher peak memory or multi-node contention.

Each concrete implementation should normally have its own public sub-issue and focused PR. This umbrella tracks the shared goals, sequencing, evidence, and architectural decisions.

Correctness requirements

Correctness is a cross-cutting acceptance condition for every work item:

• Preserve exact tensor values for lossless layout and storage changes.
• Compare precision-changing conversions against a trusted reference with justified tolerances.
• Check model outputs after loading, plus one-step training or resume behavior when relevant.
• Verify load-save-reload behavior for supported exports.
• Cover tied weights, missing or extra keys, non-contiguous tensors, MTP/LoRA, and applicable DTensor, tensor-parallel, pipeline-parallel, and expert-parallel layouts.
• Refuse silent partial loads when the built model and checkpoint do not match.

Performance comparisons must use a semantically correct route. An unsafe direct route is not a valid baseline for a fallback or transformation path.

Measurement requirements

Report the smallest useful set for each PR:

• End-to-end load or save time and phase timings.
• Effective throughput.
• Peak host and device memory on discrete-memory systems.
• Peak system/process memory, page cache, swap activity, and page faults on unified-memory systems.
• Cold-cache and warm-cache behavior when storage effects are relevant.
• The exact model, checkpoint layout, backend, dtype, topology, storage location, and code revision.

DGX Spark uses dynamically shared CPU/GPU memory, so conventional device-memory reporting is not directly comparable with discrete GPUs. See the DGX Spark hardware overview and NVIDIA unified-memory reporting guidance.

Related work

• Nemotron Nano V3 single-GPU checkpoint load: tune CPU thread budget #3160 / ci: tune Nemotron single-GPU model load threads #3370 — earlier Nemotron single-GPU profiling and thread-budget mitigation
• fix(checkpoint): restore single-GPU custom-model DCP loading #3533 — restored the single-GPU custom-model direct route on main
• refactor(checkpoint): gate direct loads by adapter capability #3574 — adapter-owned safety gate for direct checkpoint loading
• perf(checkpoint): reduce allocating grouped MoE load overhead #3580 — direct-fill grouped MoE fallback and reduced cleanup overhead