DeepSpeed

aa3914df - [AutoTP] Complete uneven sharding and universal checkpoint support (#8185)

[AutoTP] Complete uneven sharding and universal checkpoint support (#8185) Follow up #8146. ## Summary This pull request introduces support for uneven sub-parameter sharding in DeepSpeed's universal checkpoint conversion, updating the universal checkpoint format to version 0.4. The changes ensure that partitioned parameters with sub-parameters of varying sizes are correctly handled during checkpoint conversion, merging, and restoration. Additionally, the PR adds validation to prevent conversion of unsupported checkpoint layouts and improves error handling and metadata validation. Key updates by theme: **Universal Checkpoint Format and Metadata:** - Bumped the universal checkpoint version to 0.4 and introduced the `SUB_PARAM_SHARD_WIDTHS` field to record per-rank widths for each sub-parameter, enabling correct handling of uneven sub-parameter layouts. (`deepspeed/checkpoint/constants.py`, `deepspeed/checkpoint/ds_to_universal.py`) [[1]](diffhunk://#diff-7dfbb96f4f4bdab1e2be9ef97bda5a23e25e32d7e218991c268aea8065aec05eL61-R64) [[2]](diffhunk://#diff-7dfbb96f4f4bdab1e2be9ef97bda5a23e25e32d7e218991c268aea8065aec05eR93-R97) [[3]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R35) [[4]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R302-R303) [[5]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76L335-R441) - Added `AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS` to checkpoint metadata and implemented validation to prevent conversion if unsupported parameter patterns are present. (`deepspeed/checkpoint/constants.py`, `deepspeed/checkpoint/ds_to_universal.py`) [[1]](diffhunk://#diff-7dfbb96f4f4bdab1e2be9ef97bda5a23e25e32d7e218991c268aea8065aec05eL61-R64) [[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R47) [[3]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R950-R957) [[4]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R1134-R1139) **Parameter Merging and Sharding Logic:** - Enhanced the merging logic to correctly handle missing fragments for ranks with no data in uneven parameter sharding, ensuring proper alignment of slices and placeholder insertion. (`deepspeed/checkpoint/ds_to_universal.py`) [[1]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R239-R248) [[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R273-R285) - Refactored the sub-parameter merging code to use the new shard widths metadata, supporting both legacy (even) and new (uneven) layouts. Added logic to reconstruct logical shapes with placeholder dimensions. (`deepspeed/checkpoint/ds_to_universal.py`) [[1]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76L302-R327) [[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76L335-R441) [[3]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R463-R482) **Validation and Error Handling:** - Added early validation for unsupported AutoTP conversions to fail fast before expensive extraction steps. (`deepspeed/checkpoint/ds_to_universal.py`) [[1]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R950-R957) [[2]](diffhunk://#diff-ef90f4743f09a5fcb81bcf3487a1d0f4b638e1191ca2c08c1d1d1f2b2c6b6f76R1134-R1139) - Improved shape consistency checks for pipeline-parallel parameters to ensure all replicas agree on shape. (`deepspeed/checkpoint/ds_to_universal.py`) **Restoration Logic:** - Updated the restoration logic to use the new shard widths metadata, ensuring correct reconstruction of sub-parameters during model loading. (`deepspeed/checkpoint/universal_checkpoint.py`) [[1]](diffhunk://#diff-22d7b9e3b6eac1dc6e989cd8582946321c36f4fc80b527648eeb2d77a8fd3ee2L13-R13) [[2]](diffhunk://#diff-22d7b9e3b6eac1dc6e989cd8582946321c36f4fc80b527648eeb2d77a8fd3ee2R34-R73) These changes collectively improve the robustness and flexibility of DeepSpeed's checkpoint conversion, especially for advanced tensor parallelism scenarios. ## Testing * Added coverage for uneven vocabulary, GQA projections, checkpoint conversion/restore, and PP + TP tied parameters. * loss curve https://github.com/deepspeedai/DeepSpeedExamples/pull/1008 ## Limitations tp_size > num_kv need to further be optimized --------- Signed-off-by: iLeGend <824040212@qq.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Ma,Guokai <guokai.ma@intel.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>