[AutoTP] Complete uneven sharding and universal checkpoint support by jinyouzhi · Pull Request #8185 · deepspeedai/DeepSpeed

jinyouzhi · 2026-07-28T03:31:18Z

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] [2] [3] [4] [5]
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] [2] [3] [4]

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] [2]
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] [2] [3]

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] [2]
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] [2]

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 Add an AutoTP equivalence check for uneven TP size DeepSpeedExamples#1008

Limitations

tp_size > num_kv need to further be optimized

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6335a5bb11

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

delock · 2026-07-29T07:41:00Z

Hi @jinyouzhi , thanks for your PR! I have left my comments. One thing is print_dist will add output to default run, should use logging instead, if the message is temporary for debugging before merge, they should be removed.

Have you done any end to end run for this feature? This is a big change and I want to know how it works for real model training.

jinyouzhi · 2026-07-29T10:20:45Z

Hi @jinyouzhi , thanks for your PR! I have left my comments. One thing is print_dist will add output to default run, should use logging instead, if the message is temporary for debugging before merge, they should be removed.

Have you done any end to end run for this feature? This is a big change and I want to know how it works for real model training.

Thank you for the careful review. Your suggestions are very helpful and important.
I have not run a full end-to-end training test with this uneven-sharding implementation yet. For the previous PR, I tested the lm_head replacement using the DeepSpeedExamples fine-tuning scripts, but I agree that this change requires validation with a real model training workload. I will extend those tests to cover uneven sharding.

The ZeRO-3 docs update claimed checkpoint conversion "handles uneven ... fused/GQA sub-parameters", but the matched_sub_params_shape branch uses floor division and cannot handle uneven per-rank sub-parameter shards (tohtana's review). Narrow the claim to what conversion actually supports (an uneven partition dimension via per-TP-rank shapes) and state that uneven sharding within a fused/GQA sub-parameter weight is not yet supported. Full sub-param uneven support is tracked by deepspeedai#8185. Signed-off-by: Guokai Ma <guokai.ma@intel.com>

Signed-off-by: iLeGend <824040212@qq.com>

Making column-parallel layers uneven-aware left the row-parallel side on the old even-split assumption. Because a column layer's output dimension and the following row layer's input dimension are the same physical dimension, the two must agree per rank. They no longer did. With num_kv_heads set (the heuristic AutoTP path), hidden=384 and tp=4, q_proj was sharded [128, 128, 64, 64] by get_shard_size_list while o_proj was still sharded [96, 96, 96, 96] by torch.chunk, so the forward pass died with: RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x128 and 96x384) get_shard_size_list is the correct splitter here: update_mp_params derives each rank's num_attention_heads from the same function, so weights must be split the same way to stay consistent with the head metadata. torch.chunk cannot express this (it never pads, front-loads the remainder, and may even return fewer than tp_world_size chunks). This commit: * Makes LinearAllreduce uneven-aware. _tp_partition now always uses uneven_partition, dropping the training-only torch.chunk branch that existed solely because gather_params could not handle uneven shards. _mark_uc_metadata records the true original shape and partition sizes instead of deriving them as shape[1] * tp_world_size. * Adds TensorParallel_Layer._all_gather_shards, shared by both the row and column paths. Partition sizes are recomputed locally from the same deterministic split rather than discovered with an extra collective, and uneven shards are zero padded to a common size so the faster uniform all_gather_into_tensor stays usable. * Teaches ds_to_universal about uneven shards. main() collapsed every tp rank's PARAM_SHAPES into one flat dict, so _merge_zero_shards reshaped every rank's slice to a single shape and conversion failed with: RuntimeError: shape '[50, 12]' is invalid for input of size 612 Shapes are now kept per tp rank. The concatenation itself was already uneven-safe; only the reshape was wrong. * Skips the legacy vocabulary padding in load_hp_checkpoint_state when AutoTP restore metadata is present. That path derives the padded size as shape[0] * tp_world_size, which contradicts an uneven partition that _resolve_autotp_partition already describes exactly. * Asserts in get_shard_size_list that shard sizes sum to the dimension size. tp_grain_size quantization silently violates this today, e.g. get_shard_size_list(1001, 2) returns [512, 448] with tp_grain_size=64. Removing the transposes that row-side gathering previously needed also makes it faster, and the column path returns to its original cost: tp=4, bf16, 16384x16384 before after column, even shards +6% (regr) +0.1% over comm floor row, even shards baseline -8% Tested with 64 AutoTP unit tests plus a non-AutoTP universal checkpoint subset, including new end-to-end save/convert/load coverage for an uneven lm_head (vocab 101, tp=2) and uneven GQA attention (hidden 384, tp=4). Signed-off-by: iLeGend <824040212@qq.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Collect parameter shapes by explicit TP rank and deduplicate replicas across pipeline stages. Validate that replicated shapes agree before keeping one shape per TP rank, preventing tied parameters from exceeding the expected TP degree during universal checkpoint conversion. Signed-off-by: iLeGend <824040212@qq.com>

Signed-off-by: iLeGend <824040212@qq.com>

get_shard_size_list() reads the process-wide tp_shard globals num_kv_heads and tp_grain_size, which a later init_inference call or a second AutoTP model overwrites. Recomputing the split in the forward gather and in gather_params therefore let them disagree with the shards the layer was built with. Resolve it once in _freeze_partition_sizes() and have every consumer read the cached value. A tp_world_size of 1 short-circuits the helper so its grain quantization cannot truncate a replicated parameter. Signed-off-by: iLeGend <824040212@qq.com>

delock · 2026-08-03T03:30:50Z

Hi @jinyouzhi , can you also resolve conflicts with master branch? Thanks!

Upstream deepspeedai#8168 (AutoTP ZeRO-3 checkpoint consolidation) independently introduced per-TP-rank slice shapes, overlapping this branch's uneven sharding work. - ds_to_universal.py: adopt upstream's implementation wholesale, including _group_per_tp_shapes and the merge_tp_slices(uc_info, ...) signature that the new stage-3 (tp, dp) grid path requires. - layers.py: keep this branch's _all_gather_shards based gather_params for uneven row/column shards, and take upstream's removal of the write-only data_partition attribute. - test_autotp_uc_checkpoint.py: keep both test suites. Retain this branch's uneven (4,3)+(4,2) shards in test_merge_tp_slices_uses_row_parallel_cat_dim, since the shard tensors merged to the uneven version and upstream's even [4,4] shapes would fail to reshape.

jinyouzhi · 2026-08-03T06:31:05Z

Hi @jinyouzhi , can you also resolve conflicts with master branch? Thanks!

Sure, have pushed the merge with latest code.

jinyouzhi · 2026-08-07T11:01:23Z

@jinyouzhi I left some comments from manual walk through, I'll finish the rest walk through next week. Thanks!

Thank you very much for patient detailed review. I will proceed one by one.
Very appreciate the insight about UC and multi-model within one process. Looking forward to further details.

…d_size_list/get_shard_sizes Signed-off-by: iLeGend <824040212@qq.com>

Signed-off-by: iLeGend <824040212@qq.com>

jinyouzhi · 2026-08-14T16:33:48Z

Fixed UT failure and resolved conflicts (about autotp compile)

…ce_transformer_layer to set nkv and warn tp>nkv but available Signed-off-by: iLeGend <824040212@qq.com>

…deepspeedai#8308) `tests/unit/module_inject/test_fused_repartition.py::test_gate_up_partition_covers_the_whole_weight` fails on any accelerator-backed machine: ``` assert torch.equal(torch.cat([gate, up], dim=0), full_weight) E RuntimeError: Expected all tensors to be on the same device, but got other is on cpu, different from other tensors on cuda:0 ``` `GateUpPack_LinearLayer._tp_partition()` finishes with `self.move(_partition)`, and `move()` resolves its target as ```python device = 'cpu' if self.__class__.keep_module_on_host else get_accelerator().current_device_name() ``` so each shard comes back on the accelerator. `full_weight` is never moved, and `torch.equal()` refuses the cross-device comparison. The sibling test `test_gate_up_partition_ignores_later_grain_size_changes` compares two post-partition tensors, so both operands share a device and it passes — this is the only case in the file that mixes a partitioned tensor with the original. This compares against `full_weight` on the shards' device instead. It is a no-op on CPU-only runs, which is why the test passes there. Before, on master (8x H20, CUDA): ``` tests/unit/module_inject/ -> 36 passed, 1 failed ``` After: ``` tests/unit/module_inject/test_fused_repartition.py -> 5 passed tests/unit/module_inject/ -> 37 passed ``` `pre-commit run --files tests/unit/module_inject/test_fused_repartition.py` is clean. The test was added in deepspeedai#8185 four days ago. It is not covered by the live CI: `modal-torch-latest` runs `tests/unit/v1/` only, and the self-hosted GPU workflows that would run `unit/` have not produced a run in a long time (`nv-a6000` last ran 2025-08-01, `nv-nightly` 2026-01-15, `nv-torch-latest-v100` and `nv-inference` have no runs listed). Env: torch 2.13.0+cu130, H20 (sm90), single node. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

`AutoTP.register_replicated_grad_hooks()` calls `print_dist()`, but `auto_tp.py` imports only `log_dist`: ```python from deepspeed.utils.logging import log_dist ... if registered: print_dist( f"AutoTP: registered tensor-parallel grad all-reduce for {len(registered)} replicated " f"parameters, e.g. {registered[0]!r}", ranks=[0]) ``` So the call raises as soon as `registered` is non-empty: ``` NameError: name 'print_dist' is not defined ``` `DeepSpeedEngine` calls this during AutoTP setup (`deepspeed/runtime/engine.py:806` and `:845`), so any AutoTP run that registers a hook on a replicated parameter hits it. `print_dist` and `log_dist` are both in `deepspeed/utils/logging.py` and are not interchangeable — `print_dist` exists specifically for messages that should appear regardless of log level — so this adds the import rather than switching the call to `log_dist`. Before, on master: ``` tests/unit/model_parallelism/test_tp_plan_real_models.py -> 2 failed, 4 passed FAILED TestQwen3UnevenTPPlan::test_qwen3_tp3_keeps_attention_heads_aligned FAILED TestQwen3UnevenTPPlan::test_autotp_size_above_kv_head_count_leaves_trailing_ranks_empty ``` After: ``` tests/unit/model_parallelism/test_tp_plan_real_models.py -> 6 passed ``` `pre-commit run --files deepspeed/module_inject/auto_tp.py` is clean. The call came in with deepspeedai#8185 on 2026-08-20. It is not caught by the live CI: `modal-torch-latest` runs `tests/unit/v1/` only, and the self-hosted GPU workflows that would run `unit/` have not produced a run in a long time (`nv-a6000` last ran 2025-08-01, `nv-nightly` 2026-01-15, `nv-torch-latest-v100` and `nv-inference` have no runs listed). Env: torch 2.13.0+cu130, H20, single node. Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>

…eta (deepspeedai#8241) ## What this does Fixes deepspeedai#8231. `tp_shard` kept `num_kv_heads` / `num_attention_heads` / `n_embd` / `tp_grain_size` as **process-wide mutable globals**, written during AutoTP replacement. A second AutoTP model loaded into the same process overwrote them, so the first model's later sharding / gather / checkpoint conversion silently read the wrong values — making it unsafe to run more than one AutoTP model per process (teacher/student, online distillation, RL actor + reference). This moves that state onto a per-model `AutoTPMeta`, computed once from the model config and threaded through every sharding helper and TP layer, so each model carries its own kv-head / grain state. ## Stacking / merge order **Depends on deepspeedai#8185** (AutoTP uneven sharding). This branch is based on deepspeedai#8185's head and is opened as **draft** until deepspeedai#8185 lands — GitHub will drop deepspeedai#8185's commits from this diff automatically once it merges. Please **merge after deepspeedai#8185**. ## Changes - `AutoTPMeta` dataclass + `from_model_config` (single source for kv-head / attn-head / hidden extraction); `get_shard_size(_list)` take it as a required arg; tp_shard globals + `set_*` / `get_*` removed. - AutoTP threads `tp_meta` from `__init__` through every TP layer and fused-QKV helper. - Ulysses sequence parallelism gets its own `_ulysses_num_kv_heads`, decoupled from AutoTP (the AutoTP↔Ulysses coupling is gone; Ulysses's own multi-model case is left for a separate change). - Inference engine builds one `meta` per model (`_autotp_meta`) and threads it through the alibi head-sharding helpers; `_get_model_head_count` / `_get_model_kv_head_count` deleted. - kv-head / attn-head attribute lists unified behind `_kv_head_count_from` / `_attention_head_count_from` (covers chatglm, falcon, llama-class, dbrx, legacy `n_head_kv`). ## Tests Validated on 4×RTX 4080 (nccl) against deepspeedai#8185: full AutoTP / SP / checkpoint suite passes (127 passed); remaining failures are pre-existing env issues (transformers/HF network `client has been closed`, torch 2.12 `ProcessGroupGloo.perform_nocolor_split`, a cuda/cpu device-mismatch), each confirmed failing on the deepspeedai#8185 baseline too. `test_two_models_do_not_clobber_each_others_meta` is the direct regression test for deepspeedai#8231. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com> Signed-off-by: Ma, Guokai <guokai.ma@intel.com>

jinyouzhi marked this pull request as ready for review July 28, 2026 08:02

jinyouzhi requested review from hwchen2017, loadams, tjruwase and tohtana as code owners July 28, 2026 08:02

jinyouzhi changed the title ~~[AutoTP] Complete uneven sharding and universal checkpoint support~~ [WIP][AutoTP] Complete uneven sharding and universal checkpoint support Jul 28, 2026

chatgpt-codex-connector Bot reviewed Jul 28, 2026

View reviewed changes

Comment thread deepspeed/runtime/engine.py

jinyouzhi mentioned this pull request Jul 28, 2026

[RFC] Support sharding LM heads and adopting Online Softmax #8173

Open