Fix universal checkpoint resume across AutoTP sizes

[0z5a]

Universal checkpoints saved by AutoTP can fail or restore corrupted FP32 parameter mappings when training resumes. TP2 -> TP1 enters the Megatron-specific model-state merger even though universal weights are restored separately from zero/. TP2 -> TP2 can overwrite the FP32 parameter mapping while flattening Adam moments. This fixes metadata-rank selection for universal loading and updates each optimizer-state mapping without replacing the FP32 parameter mapping.

This carries forward my original fix and integration tests, which were cherry-picked as dd5c57f into #8385 after its merge-queue snapshot. The merged squash 1190946 contains the affine IR/converter work but omits these three files. The follow-up applies only the two loader fixes and the existing resume tests on that merged base.

The integration test trains a real column/row AutoTP FP32 model with ZeRO-1 and torch Adam for four steps at TP2, saves and converts its checkpoint through legacy or affine metadata, then resumes four steps at TP1 or TP2. It compares logits, loss, gradients, FP32 parameters, both Adam moments, and step counts against uninterrupted training. The affine maps are injected by the test; production metadata emission remains separate work under #8252. Native row-parallel bias uses scale 1 here; this does not enable general non-unit-scale optimizer resume.

Validation on the merged base plus this patch, using Apple M5 CPU/Gloo, Python 3.12.13 and torch 2.14.0:

• All four resume cases fail on unmodified 1190946: TP1 hits Megatron key validation; TP2 has incorrect FP32 mappings. All four pass with this patch.
• 71 affine/AutoTP unit regressions pass.
• Four existing tensor-fragment and universal DP-resize regressions pass.
• Modified-file pre-commit checks pass.

The same original patch also previously passed all four resume cases on two NVIDIA RTX 4000 Ada GPUs (20 GB each), Python 3.12.14, torch 2.13.0+cu130 and NCCL, as recorded in the original validation. That GPU run used the pre-merge base; the fresh follow-up validation above is CPU/Gloo.

pytest -sv tests/unit/checkpoint/test_autotp_uc_checkpoint.py::TestAffineUniversalCheckpointResume

Related: #8252, #8230. Follow-up to #8385.

[0z5a]

@delock Could you review this follow-up? It restores the two AutoTP universal-checkpoint loader fixes and the four save-convert-resume integration cases omitted from the #8385 merge-queue snapshot. On the merged base, all four cases fail before the patch and pass after it; 75 related regressions also pass. The scope remains native scale-1 training, with production affine metadata emission left to the separate step-4 work.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7be32550bd

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required sign-off trailer

This is a non-merge commit, but its message has no Signed-off-by trailer, so it violates the repository's mandatory commit requirement. Recreate the commit with git commit --signoff using the configured author identity.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

[Achyuthan-S]

Reviewed, and reproduced both halves of your claim independently on the merged base (1190946), CPU/gloo with DS_ACCELERATOR=cpu LOCAL_SIZE=4.

Applying only your test file, with the production changes withheld: all four cases fail. The reasons split cleanly along the two bugs, which is a nice confirmation that the fixes are independent rather than one masking the other:

case	failure without the fix
tp1-legacy, tp1-affine	key: attention.dense.weight is not found in the checkpoint
tp2-legacy, tp2-affine	snapshot 0, fc1.weight/fp32: Tensor-likes are not close!

The TP1 pair hits the Megatron merger, the TP2 pair hits the fp32 mapping. Adding both production changes: 4 passed — 43 with the affine suite alongside. I checked the file contents between stages to be sure the production files really were untouched in the first run, since that is the part of the experiment that is easy to get wrong.

Worth spelling out why the tensor_fragment.py line matters, for whoever merges this. The loop runs once per optimizer key, and get_hp_fragment(optim_state_key=key) returns optim_fragment[key]. Assigning to hp_fragment overwrote the FP32 master-weight pointer with a moment buffer slice, so after the loop every update_lp() / update_hp() read and wrote the wrong tensor. One line, and it corrupts a resume with nothing to signal it — which is exactly the Tensor-likes are not close signature above.

The engine.py change is scoped behind load_universal_checkpoint(), so the non-UC path is untouched.

On the Codex comment: Signed-off-by is present on 7be3255 and the DCO check passes — it looks like the diff is being read without the commit trailers. Same false positive it produced on #8385 and #8477.

LGTM.

@delock this fixes a live bug on master rather than adding anything: universal-checkpoint resume across TP sizes currently fails in all four combinations, so it may be worth taking ahead of #8477.

[delock]

Thanks @0z5a for this fix!