Prepared-metadata reuse for collective-free async checkpoint saves

[asolergi-nv]

Reuse a prepared .metadata across checkpoint saves (collective-free save)

This is a linked pair of PRs — both must land for the feature to work:

• NVRx (asolergi-nv/nvidia-resiliency-ext, branch feat/ckpt-metadata-reuse): adds the backend reuse path to the async saver.
• Megatron-LM (asolergi-nv/Megatron-LM, branch feat/ckpt-metadata-reuse): adds the --ckpt-metadata / --ckpt-metadata-create knobs that drive it.

The Megatron PR depends on the NVRx PR (it passes the new reuse_metadata_obj / create_metadata_path arguments). Merge/release NVRx first.

---

Problem

On the NVRx async save path, every distributed checkpoint save performs two
cross-rank collectives that scale with world size and are pure coordination
(no model data):

1. Planning — gather_object(local_plan) to the coordinator to build the
   global Metadata.
2. Finalize — gather_object(write_results) to assemble and write .metadata.

For large jobs these gathers are a non-trivial, recurring tax on every save,
even though the checkpoint structure is constant across an entire run
(same model, parallelism, and world size). The global Metadata is therefore
identical save-to-save and across resumes — it is wasteful to rebuild it every
time via collectives.

Solution

Allow a prepared, complete Metadata object to be supplied to the save, so
both gathers are skipped entirely:

NVRx side (async_ckpt/state_dict_saver.py)

• save_state_dict_async_plan(..., reuse_metadata_obj=None): when a prepared
  metadata is supplied, the decentralized planning branch skips both the
  verify_global_md_reuse all-reduce and the gather_object(local_plan) —
  planning becomes purely local.
• save_state_dict_async_finalize(..., reuse_metadata_obj=None, create_metadata_path=None):
  the reuse path skips gather_object(write_results) and instead detects write
  failures with a single all_reduce(MIN) of a 1-element flag, then writes
  the prepared metadata verbatim as .metadata.
• Create mode (create_metadata_path): the first save runs the normal
  (collective) path, then the coordinator reads back the just-written complete
  metadata, broadcasts it, and persists it to a container-local file so later
  saves in the job/run can reuse it. Returns the complete metadata.

A single, tiny all_reduce(MIN) is intentionally retained on the reuse path for
write-failure safety (so a failed rank still aborts the whole job).

Megatron side

• --ckpt-metadata <FILE>: reuse mode — load the prepared .metadata (pickle)
  and hand it to the save strategy as prepared_metadata (skips both gathers).
• --ckpt-metadata-create: first-job mode — the first save creates the
  <FILE> via the collective path, then subsequent saves reuse it.
• TorchDistSaveShardedStrategy gains prepared_metadata / ckpt_metadata_create_path
  and passes reuse_metadata_obj / create_metadata_path through to NVRx; the
  async finalize_fn captures the returned complete metadata in create mode.

Save-side collectives in NVRx (and which disappear)

Each NVRx async save (decentralized plan path) triggers these cross-rank
collectives — none of which move model/optimizer data; they only coordinate
the checkpoint's Metadata:

Phase	Collective	Purpose
Planning (save_state_dict_async_plan)	all_reduce inside verify_global_md_reuse	check whether the previously-loaded plans still match (only when prior plans were loaded)
Planning	gather_object(local_plan) → coordinator	gather every rank's local plan to build the global Metadata
Finalize (save_state_dict_async_finalize)	gather_object(write_results) → coordinator	gather every rank's write results to assemble and write .metadata
Finalize	broadcast(failures_occurred) (1 int)	propagate write-failure status to all ranks

Which disappear after the first save(s), and why

• Planning all_reduce + gather_object(local_plan) — eliminated by NVRx's
  own CheckpointMetadataCache. With --ckpt-assume-constant-structure, the cache
  caches the central plan and, once it has observed an identical structure on
  two consecutive saves, sets validated_cache_reuse=True and reuses the
  cached central plan. So these run on save Add GitHub Actions workflow for building docs #1 and Lint code (black, isort, ruff checks) #2 (the cache needs two
  matching observations to validate) and then vanish from save add linting code actions #3 onward.
• Finalize gather_object(write_results) + broadcast — not cached by
  stock NVRx: the per-rank write results are re-gathered to rebuild .metadata
  on every save, so this collective never disappears on its own.

What --ckpt-metadata changes

Supplying a prepared, complete Metadata lets the save skip both gathers:

• Reuse mode: reuse_metadata_obj is set → the planning branch skips the
  verify_global_md_reuse all_reduce and gather_object(local_plan)
  (planning is purely local), and the finalize skips gather_object(write_results),
  replacing it with a single all_reduce(MIN) of a 1-element flag for
  write-failure detection. This holds from the very first reuse save — no
  warm-up saves needed — and is independent of CheckpointMetadataCache.
• Create mode (first job): runs the full collective path once to produce the
  metadata, plus a one-time broadcast of the complete metadata to all ranks,
  then persists it; every later save reuses it collective-free.

Net steady-state: the only remaining save-side collective is the 1-element
all_reduce(MIN) failure check (kept by design). (The separate mcore-level
collectives — the FullyParallel save-distribution all_gather_object and the
sharding-integrity determine_global_metadata gather — are addressed by the
sibling Megatron PRs M4 and M5 and are out of scope here.)

Correctness constraints

Reuse is only valid when the run that reuses a .metadata has the same
checkpoint structure, world size, and --dist-ckpt-workers as the run that
created it (the caller owns this guarantee). In particular the prepared
metadata must come from a save done with the same --keep-redundant-extra-state
setting, since that toggle changes the saved item set and per-file byte
offsets. A stale/mismatched prepared metadata produces a wrong shard map and a
load-time unpickling error.

Usage

First job (create once):

--async-save --use-persistent-ckpt-worker \
--ckpt-metadata $CONTAINER/prepared.metadata --ckpt-metadata-create

Subsequent jobs (reuse, collective-free saves):

--async-save --use-persistent-ckpt-worker \
--ckpt-metadata $CONTAINER/prepared.metadata

Result

With --ckpt-metadata reuse active (and constant structure), the only remaining
save-side collective is the 1-element all_reduce(MIN) failure check. Both
gather_objects are eliminated.

[asolergi-nv]

🔗 Paired PR (must land together): Megatron --ckpt-metadata knobs that drive this backend → NVIDIA/Megatron-LM#5555

[shurkat-nvidia]

How are you ensuring that the length/offset are correctly collected from WriteResult reported by every rank here?

[asolergi-nv]

Good question — this is the core invariant the reuse path relies on, so let me be explicit. The per-chunk offset/length in storage_data are a function of the checkpoint structure — item shapes/dtypes and the write-bucketing (which depends on world size and --dist-ckpt-workers) — not of the tensor values. For a fixed structure the byte layout is therefore identical save-to-save, which is exactly what makes a prepared metadata reusable.

I've made the unit test demonstrate this directly (also addressing your other comment): it now re-randomizes every parameter between the create save and the reuse save, then asserts the reuse checkpoint — whose .metadata is the verbatim prepared metadata built from the old weights — loads back to the new weights. If any offset/length depended on the data, that round-trip would fail. The validity constraint (same structure / world size / --dist-ckpt-workers / --keep-redundant-extra-state) is documented in the docstring and PR body and is the caller's responsibility.

If you'd like a stronger guarantee we could add a purely-local save-time check that compares this rank's actual WriteResult offsets/lengths against the prepared metadata (no collective needed) and fails fast on a mismatch — happy to add that as a follow-up if you think it's worth it. Addressed in 7ce1029.

[shurkat-nvidia]

One potential idea, instead of a fail-fast would be to hijack the mechanism that records whether ranks are done writing. Currently, that's an all-reduce/all-gather of a single value tensor, which can be 0 or 1. We could add a value 2 that indicates the rank was done writing, but its offsets/lengths/something else about distcp file has changed and the metadata cache can't be reused.

[greptile-apps]

Greptile Summary

This PR adds a prepared-metadata reuse path to NVRx's async checkpoint saver, allowing both gather_object collectives per save to be eliminated once a complete Metadata is available. On the reuse path a single 1-element all_reduce(MIN) replaces the finalize gather, and a bidirectional layout guard (_find_layout_mismatch) protects against stale prepared metadata before any .metadata file is published.

• save_state_dict_async_plan gains reuse_metadata_obj: when set on the decentralized-plan branch, both the verify_global_md_reuse all-reduce and the gather_object(local_plan) are skipped entirely; a warning is logged for non-decentralized planners.
• save_state_dict_async_finalize gains reuse_metadata_obj (skip the write-result gather, validate layout locally, write prepared metadata verbatim) and create_metadata_path (run the full collective path, broadcast the resulting complete metadata to all ranks, persist it for future jobs); the two arguments are mutually exclusive, enforced with a ValueError.
• Three helper functions (_write_metadata_only, _atomic_pickle_write, _read_written_metadata) and comprehensive unit tests covering both the offset-stability rationale and the bidirectional layout guard are added.

Confidence Score: 5/5

Safe to merge; all previous reviewer concerns are addressed and no new correctness bugs were found.

The implementation is carefully constructed: the bidirectional layout guard catches all documented drift scenarios, the mutual exclusion of reuse and create is enforced with an explicit ValueError, the temp+rename write is atomic on the POSIX path, and f.flush() precedes os.fsync. The only finding is that _write_metadata_only writes the prepared metadata verbatim, so the new checkpoint's .metadata carries storage_meta.checkpoint_id from the create-time checkpoint directory rather than the current one — a provenance staleness issue for tooling, not a loading bug.

No files require special attention beyond the _write_metadata_only call site noted in the inline comment.

Important Files Changed

Filename	Overview
src/nvidia_resiliency_ext/checkpointing/async_ckpt/state_dict_saver.py	Core implementation of metadata reuse/create; previous thread issues (flush-before-fsync, rm_file+rename, docstring, mutual-exclusion assertion, bidirectional layout check) are all addressed. One P2 concern: _write_metadata_only writes reuse_metadata_obj verbatim, carrying a stale storage_meta.checkpoint_id from the create-time checkpoint into every reuse checkpoint's .metadata.
tests/checkpointing/unit/test_async_writer.py	Adds test_prepared_metadata_reuse_roundtrip covering create→reuse→load round-trip with a _GatherObjectSpy that verifies exactly one fewer gather_object call on the reuse path; logic and assertions are correct.
tests/checkpointing/unit/test_metadata_reuse_offsets.py	New file; establishes the offset-stability rationale (tensor offsets are value-independent; only storage_meta differs across saves; BYTE_IO bytes objects can drift but uint8 tensors cannot). No issues found.
tests/checkpointing/unit/test_metadata_reuse_validation.py	New file; exhaustively tests _find_layout_mismatch in both directions (forward: length drift, offset shift, absent item; reverse: last-item removal, resharding) and verifies that a stale layout raises and writes no .metadata. No issues found.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant R as All Ranks
    participant C as Coordinator (rank 0)
    participant FS as Filesystem

    rect rgb(200, 230, 200)
        Note over R,FS: CREATE mode (first save — collective path)
        R->>R: save_state_dict_async_plan (gather_object local_plan)
        R->>C: gather_object(write_results) in finalize
        C->>FS: storage_writer.finish() → writes .metadata
        C->>FS: _read_written_metadata() → reads back complete Metadata
        C-->>R: broadcast_object_list(complete_metadata)
        C->>FS: _atomic_pickle_write(complete_metadata, create_metadata_path)
        Note over R: Returns complete_metadata to caller
    end

    rect rgb(200, 200, 240)
        Note over R,FS: REUSE mode (subsequent saves — collective-free)
        R->>R: save_state_dict_async_plan (skips gather_object, purely local)
        R->>R: _find_layout_mismatch(write_results, reuse_metadata_obj)
        R->>R: all_reduce(MIN, local_status) — single 1-element collective
        alt "global_status == 2 (all match)"
            C->>FS: _write_metadata_only(reuse_metadata_obj) — atomic tmp+rename
        else "global_status == 1 (layout drifted)"
            R->>R: raise CheckpointException(metadata_reuse)
        else "global_status == 0 (write failure)"
            R->>R: raise CheckpointException(write)
        end
    end

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant R as All Ranks
    participant C as Coordinator (rank 0)
    participant FS as Filesystem

    rect rgb(200, 230, 200)
        Note over R,FS: CREATE mode (first save — collective path)
        R->>R: save_state_dict_async_plan (gather_object local_plan)
        R->>C: gather_object(write_results) in finalize
        C->>FS: storage_writer.finish() → writes .metadata
        C->>FS: _read_written_metadata() → reads back complete Metadata
        C-->>R: broadcast_object_list(complete_metadata)
        C->>FS: _atomic_pickle_write(complete_metadata, create_metadata_path)
        Note over R: Returns complete_metadata to caller
    end

    rect rgb(200, 200, 240)
        Note over R,FS: REUSE mode (subsequent saves — collective-free)
        R->>R: save_state_dict_async_plan (skips gather_object, purely local)
        R->>R: _find_layout_mismatch(write_results, reuse_metadata_obj)
        R->>R: all_reduce(MIN, local_status) — single 1-element collective
        alt "global_status == 2 (all match)"
            C->>FS: _write_metadata_only(reuse_metadata_obj) — atomic tmp+rename
        else "global_status == 1 (layout drifted)"
            R->>R: raise CheckpointException(metadata_reuse)
        else "global_status == 0 (write failure)"
            R->>R: raise CheckpointException(write)
        end
    end

Loading

_{Reviews (5): Last reviewed commit: "Make reuse layout guard bidirectional" | Re-trigger Greptile}

[shurkat-nvidia]

You should replace model weights by random weights between the metadata reuse. That will check that the metadata reuse holds for multiple training iterations.

[asolergi-nv]

Done in 7ce1029. The test now re-randomizes every parameter between the create and reuse saves (emulating a later training iteration — same structure, new data) and asserts the reuse checkpoint loads back to the new weights. This exercises reuse across changed data and, as noted on your other comment, demonstrates that the reused offsets/lengths depend only on structure, not on the values.

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.