[WIP][AutoTP] Complete uneven sharding and universal checkpoint support - #8185
[WIP][AutoTP] Complete uneven sharding and universal checkpoint support#8185jinyouzhi wants to merge 12 commits into
Conversation
jinyouzhi
changed the title
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 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".
| sub_dim_sizes = (sub_dim_sizes, ) | ||
|
|
||
| partition_shape = [sum(d) if isinstance(d, tuple) else d for d in matched_sub_params_shape.shape] | ||
| partition_shape = [d // tp_degree if i == partition_dim else d for i, d in enumerate(partition_shape)] |
There was a problem hiding this comment.
would d always divisible by tp_degree here?
|
|
||
| offset = 0 | ||
| for sub_dim_size in sub_dim_sizes: | ||
| part_sub_dim_size = sub_dim_size // tp_degree |
There was a problem hiding this comment.
same question, does sub_dim_size always divisible by tp_degree. If we always intentionally not handle this case (i.e. the case is too complicated), then we should guard it where the case originated.
| merged_chunks = [] | ||
| for sub_dim_size in sub_dim_sizes: | ||
| sub_slice = full_view.narrow(partition_dim, offset, sub_dim_size) \ | ||
| .chunk(tp_world_size, dim=partition_dim)[tp_rank] |
There was a problem hiding this comment.
here also have even assumption.
| # Shards must tile the dimension exactly, otherwise the partitioned weights no longer | ||
| # reconstruct the original tensor. tp_grain_size quantization can violate this when the | ||
| # dimension is not a multiple of the grain size. | ||
| assert sum(shard_sizes) == total_size, ( |
There was a problem hiding this comment.
I have concern that grain is usually set to meet kernel performance constraints and its global. And vocab size may not be divisible by grain, so the final result will be either grain = 1, which is not performance optimial, or vocab cannot divide by grain and get an assertion. A proper solution should be ither handle the remainder gracefully, or use a different grain for vocab. Preferrable the first one if solution is not too complicated.
There was a problem hiding this comment.
Insightful. I prefer to follow your suggestion and keep the existing logic: for total_size >= grain_size, keep the grain-aligned allocation first, even if some ranks end up with zero width. Align as many ranks as possible to grain_size, and let the remainder go to the last rank.
However, for the more fine-grained case—specifically when total_size falls in [grain_size, grain_size * mp_size)—I think it's worth exploring further. For example, given [64,64,64,64,64,64,0,0], should we instead distribute evenly across all ranks so none sit idle: [48,48,48,48,48,48,48,48]. Taking it a step further, we could search for a more balanced allocation via binary search, e.g. [64,64,64,64,32,32,32,32].
This could be left as a open question.
|
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. |
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>
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>
|
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.
Sure, have pushed the merge with latest code. |
…eplicas Signed-off-by: iLeGend <824040212@qq.com>
get_shard_size() quantizes a split to tp_grain_size by flooring the dimension into whole grains, so total_size % tp_grain_size was dropped and the shards no longer tiled the dimension. A GPT-2 vocabulary of 50257 over two ranks yielded 25152 + 25088 = 50240, silently losing the last 17 rows. Give that tail to the last rank instead. Every other rank keeps the kernel alignment tp_grain_size exists for, and the shards reconstruct the dimension exactly, so the sum check in get_shard_size_list() is now an internal invariant rather than a configuration error a user cannot act on. The band where a dimension holds fewer grains than there are ranks still leaves the high ranks with an empty shard. That is pre-existing behaviour, unrelated to the dropped remainder, and is left alone here. With the remainder preserved, a tp_world_size of 1 no longer needs to bypass the shard helper to avoid truncation, so that special case is removed. Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
2 participants
Follow up #8146.
Summary
This PR completes AutoTP uneven sharding support across column-parallel layers, row-parallel layers, and universal checkpoints.
Changes
adb95dd: enable uneven shared & gather by keeping the original and partition sizes)d32bcdf: let the row-parallel use the uneven partition which remove the transposes to faster)d32bcdf: better gather impl based on partition sizes recomputed locally)d32bcdf)6335a5b: deduplicate replicas across pp stages for ckpt)Testing
Added coverage for uneven vocabulary, GQA projections, checkpoint conversion/restore, and PP + TP tied parameters.