Fix checkpoint rank selection for Ulysses sequence parallelism (#8226)
## Summary
- use checkpoint model-parallel rank 0 for Ulysses sequence parallelism
- centralize checkpoint rank selection across model, expert, optimizer,
and ZeRO checkpoint paths
- preserve the real model-parallel rank for tensor-parallel MPUs
## Problem
Ulysses sequence parallelism does not shard model weights, so DeepSpeed
records `mp_world_size == 1`. However, the Ulysses MPU aliases
`get_model_parallel_rank()` to the sequence-parallel rank. Checkpoint
loading combined those values when selecting a model-state file, causing
SP ranks greater than zero to index past a single non-SP checkpoint
shard.
The same rank interpretation was also used for checkpoint filenames, so
this change keeps save and load behavior consistent for replicated SP
weights and optimizer shards.
## Tests
- `pytest tests/unit/sequence_parallelism/test_ulysses.py -k
'CheckpointRank or load_non_sequence_parallel_checkpoint' -q` (3 passed)
- `pytest tests/unit/checkpoint/test_latest_checkpoint.py -q` (2 passed)
- `pre-commit run --files deepspeed/runtime/engine.py
tests/unit/sequence_parallelism/test_ulysses.py`
The distributed regression saves without SP, loads module-only with
SP=2, then saves and fully resumes the SP checkpoint.
Signed-off-by: Thong Nguyen <thong.nguyen@snowflake.com>