Skip to content

Read only needed checkpoint snapshot components during finalization - #984

Draft
bradhilton wants to merge 2 commits into
hayek/checkpoint-packed-cpu-snapshotfrom
codex/checkpoint-selective-finalization
Draft

bradhilton wants to merge 2 commits into
hayek/checkpoint-packed-cpu-snapshotfrom
codex/checkpoint-selective-finalization

Conversation

@bradhilton

@bradhilton bradhilton commented Sep 26, 2026 •

Copy link
Copy Markdown
Collaborator

Checkpoint finalization reads each per-block snapshot once for LoRA, three times for optimizer components, and once for step scalars. Each pass currently materializes every tensor. Read only the requested component with safe_open/get_tensor, preserving explicit missing-key errors, tensor lifetime, step ordering and the existing merge, manifest and cleanup paths. There is no retained tensor cache or public API change.

Stacked on #979 (a0c6c767); the production diff is limited to the two reader call sites. A follow-up also changes two inherited test-fixture assignments to setattr so CI type-checks their ModuleType stand-ins; the runtime component is unchanged.

Validation:

  • 24 new CPU cases pass on both safetensors 0.7.0 and the exact 0.8.0 wheel used by the frozen native runtime. They compare complete checkpoint bytes across dtypes/ranks/optimizer modes and cover missing/corrupt inputs, error identity, cancellation, abort cleanup, prefix/order behavior and tensor validity after handle closure/file unlink.
  • The existing 33 packed-snapshot tests pass. Ruff, formatting and lock checks pass. Pinned ty 0.0.59 passes the changed files; the full local project has the same eight diagnostics as the unchanged base. The initial full prek attempt used the shared environment's older ty 0.0.14 and failed type checks; canonical CI remains required.
  • On the two-block public fixture, logical tensor objects requested fall from 100 to 20 and logical tensor bytes from 5,160 to 1,032. These are not physical disk reads or copied bytes; mmap and page-cache behavior matter. Both paths still open each block five times.

A bounded warm-cache CPU reader benchmark using safetensors 0.8.0 measured median wall times of 9.14 → 3.02 ms (16 keys/block, 65,536 elements/key) and 55.65 → 16.30 ms (128 keys/block, 2,048 elements/key). It used two blocks, one Torch CPU thread, and eight alternating-order samples of four passes. CPU time agreed with wall time. This measures header/tensor-view reader work only, excluding tensor-data arithmetic, checkpoint merge/write, physical I/O, comparative RSS and GPU performance. Native background stack samples motivated the change but do not establish interference or speedup.

Canonical CI passes on 5fc36d957: lint/format/type/lock, 1,264 Megatron lightweight tests, and 2,083 unit tests (34 skipped, two deselected). TrainerRank GPU validation completed successfully on two H200s: 95 pytest cases passed (two skipped), followed by the CP2 and TP2 public TrainerRank checks. Sky job 1 reported SUCCEEDED and the workflow completed cluster teardown at 2:29:27 AM MDT. These tests establish compatibility, not native workload performance.

The first head's GPU attempt failed before a training job because two free H200s were unavailable; its launch request was cancelled. Its quality run exposed the two inherited fixture type errors corrected by the follow-up above. Those earlier outcomes remain retained. This PR stays draft and unmerged.

@bradhilton
bradhilton had a problem deploying to trainer-rank-gpu-validation September 26, 2026 08:00 — with GitHub Actions Failure

This branch was successfully deployed

1 active deployment
trainer-rank-gpu-validation — 5fc36d95 Deployed Sep 26, 2026 by bradhilton via Run on 2x H200 #744
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant