[AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta - #8241
Draft
delock wants to merge 32 commits into
Draft
[AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta#8241delock wants to merge 32 commits into
delock wants to merge 32 commits into
Conversation
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>
Upstream #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.
…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>
Signed-off-by: iLeGend <824040212@qq.com>
- Add assertions to validate sub-parameter sizes and shard widths in merge_tp_slices. - Update collect_autotp_universal_checkpoint_info to publish physical sub-parameter sizes instead of counts. - Introduce tests to ensure correct handling of uneven sub-parameter sizes and rejection of invalid configurations. Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
get_shard_size defaulted to the global rank while every caller indexes shards by the rank within the tensor parallel group, so a group smaller than the world sized the wrong shard. Document the contract and default to the group-local rank. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: iLeGend <824040212@qq.com>
The fused layers re-derived their split from the process-wide tp_shard globals, which a second AutoTP model overwrites, so repartitioning after a gather cut the weight differently than the frozen widths recorded for the gather and the checkpoint. Drive both from the frozen widths, carry QWen's attention split size along with them, and refuse to gather the layouts that were never split in rank order. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: iLeGend <824040212@qq.com>
The alibi and attention bias helpers sliced by the global rank and size while the weights are cut per tensor parallel group, so the masks disagreed with the weights whenever the group was smaller than the world. Bind them to the group, and build them after that group exists. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: iLeGend <824040212@qq.com>
Merge every rank at the width it recorded, keep a rank that wrote no fragment aligned with its neighbours, validate conversion support before the extraction pass so a rejected checkpoint leaves no fragments behind, and refuse on restore the layouts conversion already marked unsupported. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: iLeGend <824040212@qq.com>
…lices Signed-off-by: iLeGend <824040212@qq.com>
…idths Signed-off-by: iLeGend <824040212@qq.com>
…s for gather and partition functionality Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
strengthen uneven autotp ucp tests with element-wise checks ref to tp3 cases clean up the rank-0-only operations to ensure failures fail the test instead of leaving other ranks hanging in the next collective. Co-authored-by: Ma,Guokai <guokai.ma@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: iLeGend <824040212@qq.com>
Signed-off-by: iLeGend <824040212@qq.com>
tp_shard kept num_kv_heads / num_attention_heads / n_embd / tp_grain_size as process-wide mutable globals set 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. Move that state onto a frozen AutoTPMeta dataclass computed once from the model config and threaded through every sharding helper and TP layer. Each model now carries its own kv-head / grain state, so multiple AutoTP models (teacher / student, online distillation, RL actor + reference) can coexist in one process. Ulysses sequence parallelism, which repurposed the same global, gets its own private kv-head state so it no longer depends on whichever AutoTP model was loaded last. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
AutoTP (``AutoTPMeta.from_model_config``) and the inference engine (``_get_model_head_count`` / ``_get_model_kv_head_count``) each kept their own attribute-name lists for kv-head and attention-head counts, so the two paths recognized different model families (e.g. chatglm only on the AutoTP side, legacy ``n_head_kv`` / ``kv_n_heads`` only on the inference side) and could even disagree on a plain transformer. Consolidate each count behind one shared list and helper in tp_shard so coverage and probe order live in a single place: - ``_KV_HEAD_ATTRS`` / ``_kv_head_count_from`` for the key/value head count - ``_ATTN_HEAD_ATTRS`` / ``_attention_head_count_from`` for the attention head count Both ``AutoTPMeta.from_model_config`` and the inference engine consume them, so neither count is re-extracted on either side. The union keeps the legacy aliases for older configs/checkpoints, annotated with the transformers version that superseded each (``n_head_kv`` after 4.33, ``kv_n_heads`` superseded at top-level by ``num_key_value_heads`` in 4.40). Signed-off-by: Guokai Ma <guokai.ma@intel.com>
``get_head_shard_sizes`` and ``install_head_sharded_helper`` grew a ``meta`` parameter that no caller ever passed -- it was always ``None``, the ``if meta is not None`` kv-head discovery branch was unreachable, and the ``meta or AutoTPMeta()`` fallback always evaluated to ``AutoTPMeta()``. Drop the parameter and the dead discovery block; the helpers honestly take ``num_heads`` / ``num_kv_heads``, which every caller already supplies. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The alibi helpers (``get_head_shard_sizes``, ``install_head_sharded_helper``) took ``num_heads`` / ``num_kv_heads`` as scalars, so the inference engine extracted them via ``_get_model_head_count`` / ``_get_model_kv_head_count`` -- a second copy of the head-count probe that ``AutoTPMeta.from_model_config`` already does. With AutoTPMeta carrying both counts, the helpers now take a single ``meta`` and the inference engine builds one per model (via the shared ``_attention_head_count_from`` / ``_kv_head_count_from`` probes), deleting the two ``_get_model_*`` methods. The runtime alibi wrappers and ``_head_shard`` are unchanged: they still consume the ``head_shard_sizes`` + ``total_num_heads`` bound at install time, now derived from meta. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
test_gate_up_partition_ignores_later_grain_size_changes existed to check that a layer's frozen shard widths survived a second AutoTP model overwriting the process-wide grain global. With per-model AutoTPMeta the "second model" leg became vacuous -- a layer holding meta A is unaffected by merely constructing a layer with meta B -- so the test no longer tested what its name says. Fold its one piece of real value (the explicit _subparam_shard_widths == [[3,2],[3,2]] assertion) into test_gate_up_partition_covers_the_whole_weight, which already exercises the same layer and partition. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
delock
commented
Aug 10, 2026
|
|
||
|
|
||
| def set_ulysses_num_kv_heads(num): | ||
| global _ulysses_num_kv_heads |
Collaborator
Author
There was a problem hiding this comment.
A note to Ulysses SP owner: It looks like ulysses still rely on global num_kv_heads to work. Do we need to remove global variable on ulysses path as well? Agent suggested there are two path in Ulysses and the legacy path rely on this global variable. @sfc-gh-truwase
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this does
Fixes #8231.
tp_shardkeptnum_kv_heads/num_attention_heads/n_embd/tp_grain_sizeas 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 #8185 (AutoTP uneven sharding). This branch is based on #8185's head and is opened as draft until #8185 lands — GitHub will drop #8185's commits from this diff automatically once it merges. Please merge after #8185.
Changes
AutoTPMetadataclass +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.tp_metafrom__init__through every TP layer and fused-QKV helper._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).metaper model (_autotp_meta) and threads it through the alibi head-sharding helpers;_get_model_head_count/_get_model_kv_head_countdeleted._kv_head_count_from/_attention_head_count_from(covers chatglm, falcon, llama-class, dbrx, legacyn_head_kv).Tests
Validated on 4×RTX 4080 (nccl) against #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.12ProcessGroupGloo.perform_nocolor_split, a cuda/cpu device-mismatch), each confirmed failing on the #8185 baseline too.test_two_models_do_not_clobber_each_others_metais the direct regression test for #8231.