Fix universal checkpoint resume across AutoTP sizes (#8474)
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](https://github.com/0z5a/DeepSpeed/commit/e753a02b72dd18bc220f31ecac5ce5f6c1a89c54),
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](https://github.com/deepspeedai/DeepSpeed/pull/8385#issuecomment-5598904008).
That GPU run used the pre-merge base; the fresh follow-up validation
above is CPU/Gloo.
```bash
pytest -sv tests/unit/checkpoint/test_autotp_uc_checkpoint.py::TestAffineUniversalCheckpointResume
```
Related: #8252, #8230. Follow-up to #8385.
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
