[WIP][AutoTP] Complete uneven sharding and universal checkpoint support - #8185
[WIP][AutoTP] Complete uneven sharding and universal checkpoint support#8185jinyouzhi wants to merge 14 commits into
Conversation
…iguration(`colwise_gather_output`/`colwise_rep`) - Introduced `gather_output` option in `TPLayerSpec` for column-parallel layers. - Added validation to prevent using gathered output with tied embeddings. - Updated `LinearLayer` to handle gathered output during forward pass. - Enhanced documentation to reflect new gathered column parallelism capabilities. - Added tests for gathered column parallelism scenarios and configurations. Signed-off-by: iLeGend <824040212@qq.com>
…e tp_plan extraction logic Signed-off-by: iLeGend <824040212@qq.com>
… for fallback behavior Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.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>
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".
|
|
||
| local_size = torch.tensor([input.shape[-1]], dtype=torch.long, device=input.device) | ||
| gathered_sizes = [torch.empty_like(local_size) for _ in range(tp_world_size)] | ||
| dist.all_gather(gathered_sizes, local_size, group=group) |
There was a problem hiding this comment.
should call get_shard_size_list to compute size list, rather than use one additional all_gather. Pass total_size and name from parameter list.
There was a problem hiding this comment.
Very good catch! Have pushed fix.
| input_padded[..., :input.shape[-1]].copy_(input) | ||
|
|
||
| gathered = [torch.empty_like(input_padded) for _ in range(tp_world_size)] | ||
| dist.all_gather(gathered, input_padded, group=group) |
There was a problem hiding this comment.
should use all_gather_into_tensor for better collective efficiency.
There was a problem hiding this comment.
Very good catch! Have pushed fix.
Signed-off-by: iLeGend <824040212@qq.com>
| 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.
| ] | ||
| selected_route = "custom partition_config" if partition_config is not None else "HuggingFace tp_plan or AutoTP" | ||
| model_class = f"{type(model).__module__}.{type(model).__qualname__}" | ||
| print_dist( |
There was a problem hiding this comment.
should remove diagnostics message before merge.
| pattern for pattern, style in hf_tp_plan.items() | ||
| if style.lower() in ("colwise_rep", "colwise_gather_output") | ||
| ] | ||
| print_dist( |
There was a problem hiding this comment.
same here, print_dist will generate output regardless of log level, suggest change to logger.info.
|
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. |
2 participants
Pending on #8146 merged first. (Only 3 commits latest)
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.