DeepSpeed
b39e07a7 - Fix checkpoint rank selection for Ulysses sequence parallelism (#8226)

Commit
5 days ago
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>
Parents
Loading