Fix checkpoint rank selection for Ulysses sequence parallelism

[sfc-gh-thonguyen]

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.

[stas00]

Thank you for the fix, Thong.

I have a small simplification suggestion, but otherwise looks great.

[stas00]

is there any time where checkpoint_mp_rank arg isn't self.checkpoint_mp_rank?

if so it doesn't need to be passed as an arg and can be pulled from self, no?

[sfc-gh-thonguyen]

In DeepSpeedEngine save paths it is always self.checkpoint_mp_rank. However, _get_expert_ckpt_name is static and is also called by static load_moe_state_dict, including from InferenceEngine, where no DeepSpeedEngine instance exists. The optional argument preserves those legacy/static callers while letting DeepSpeedEngine override the MPU rank for Ulysses. So it cannot generally pull the value from self.