[AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta - #8241
[AutoTP] Replace tp_shard process-wide globals with per-model AutoTPMeta#8241delock wants to merge 19 commits into
Conversation
|
|
||
|
|
||
| def set_ulysses_num_kv_heads(num): | ||
| global _ulysses_num_kv_heads |
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
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>
83b244f to
1c5cfe4
Compare
_bigcode_type_transpose slices the fused projection at meta.n_embd to separate the query block from the replicated kv block. n_embd is Optional, and Python reads input[:None] as "to the end", so a missing hidden size would hand the whole weight to q and leave kv empty -- a silently wrong shard rather than an error. Assert the value is present before slicing. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
470cfd3 to
396f7a0
Compare
fused_LinearLayer defaulted a missing tp_meta to an empty AutoTPMeta before resolving its sub-parameter layout. That default is not a neutral choice here: with num_kv_heads left as None, fused_qkv_subparam_sizes skips the chatglm kv-head branch and falls back to splitting the fused weight into three equal blocks, so the layer would partition to a structurally wrong layout instead of failing. Index kwargs directly so an omitted tp_meta raises. Every AutoTP construction site already passes it; only one test relied on the default and now states the empty meta explicitly, where the codegen layout it exercises does not consult it. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
396f7a0 to
ba8f049
Compare
get_shard_size / get_shard_size_list took a num_kv_heads argument that shadowed meta.num_kv_heads, so the same name meant "this model's kv-head count" in one place and "use this instead" in another. Call it eff_num_kv_heads: the head count the split is actually aligned to, defaulting to the meta value. get_head_shard_sizes was reading meta.num_kv_heads only to pass it straight back as that argument, which is what get_shard_size_list already does when it is omitted. Drop the round trip. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The override existed because get_shard_size read a process-wide num_kv_heads that was only set during AutoTP replacement. The inference engine patches the alibi helpers before that happens, so it probed the head count itself and passed it in to bypass the uninitialized global. That was the parameter's only caller. With the count carried by a per-model AutoTPMeta, the alibi path receives its own model's value like everyone else, and the previous commit removed the round trip that read meta.num_kv_heads only to pass it straight back. Nothing outside tp_shard supplies the argument now, so remove it and the two tests written against it; get_shard_size has a single source for the head count again. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
ad8db49 to
6cef7ca
Compare
Giving Ulysses its own kv-head state stops AutoTP from overwriting it, but the memo itself is set once and never reset, so the first model to take the uneven path decides the split for every later call in the process. Note the limitation and what fixing it would cost, since the gather direction cannot recover the total head count from the tensor shape and would need it threaded through _SeqAllToAll and its backward pass. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
get_shard_size only takes the kv-head-aligned path for attention projections whose dimension divides by the head count; MLP, the last linear and MoE expert layers are excluded so they keep a near-even split. Only the attention case was asserted, so deleting the last_linear or MoE branch left the suite green. Fold the single assertion into a parametrization that pairs each exclusion with the split it produces, and add the non-divisible case. Removing any one of the four rules now fails a specific case rather than none. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
The isolation this PR provides had no end-to-end coverage. The unit test that claimed it built two frozen AutoTPMeta objects and called a pure function with each, which cannot fail regardless of the implementation, so it is replaced. Inject two models in one process with different kv-head counts and assert the first one's split survives. 3 kv heads over 2 ranks splits 192 unevenly as [128, 64] while 2 heads gives [96, 96], and [96, 96] is exactly what the first model would produce once a shared kv-head count had been overwritten, so the two outcomes are distinguishable. Re-deriving from the first model's meta after the second is built covers the value itself rather than the frozen result. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
c071362 to
e79cdf2
Compare
![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: e79cdf24e3
ℹ️ 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".
| keep_module_on_host=tp_config.keep_module_on_host, | ||
| partition_config=partition_config) | ||
| partition_config=partition_config, | ||
| model_config=model_config, |
There was a problem hiding this comment.
When _apply_autotp_partitioning takes the custom partition-config path for a composite Hugging Face model, it correctly resolves head_config = model_config.text_config above but then passes the outer model_config to AutoTP. Multimodal outer configs commonly keep num_key_value_heads, num_attention_heads, and hidden_size only under text_config, so tp_meta loses those values; with GQA and an uneven TP split this falls back to grain-based Q/K/V partitions that can cut through KV heads and fail the attention reshape. Pass head_config here and in the equivalent HF tp_plan constructor at line 799.
Useful? React with 👍 / 👎.
| def set_num_attention_heads(num): | ||
| global num_attention_heads | ||
| num_attention_heads = num | ||
| from dataclasses import dataclass |
There was a problem hiding this comment.
The reviewed commit is a non-merge commit but has no Signed-off-by trailer, so it does not satisfy the repository's mandatory commit requirement and will fail the corresponding contribution/CI check; recreate the commit with --signoff using the configured Git identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
|
Hi @PKUWZP @sfc-gh-truwase @jinyouzhi This PR removes globals from AutoTP processing, which allows AutoTP for both teacher and students in OPD scenario (previously they would conflict with each other and need WA). Can you help review this PR? Thanks! @sfc-gh-truwase Ulysses-SP still have globals left, we may need to address it in a seperate PR. |
|
It is notable that tests under tests/unit/model_parallelism should be run on multi-accelerator environment. This change is better done with a seperate PR. |
AutoTPMeta.from_model_config probed the config it was handed directly, so a multimodal outer config (head counts only under text_config) lost num_kv_heads / num_attention_heads / hidden_size and the sharding fell back to an even-grain split that can cut through KV heads under GQA. Descend into text_config inside from_model_config so every caller (runtime engine, inference engine, direct AutoTP use) shares the fix. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
|
To use Codex here, create a Codex account and connect to github. |
| # 4. Set linear policies | ||
| _autotp.update_linear_policies() | ||
|
|
||
| # 6. Replace modules |
|
|
||
| # A later model carries a different grain; this layer's split is frozen from its own meta. | ||
| other_meta = AutoTPMeta(tp_grain_size=64) | ||
| assert other_meta.tp_grain_size != layer.tp_meta.tp_grain_size |
There was a problem hiding this comment.
I think it would be clearer to directly assert: assert layer.tp_meta.tp_grain_size == 1.
Furthermore, could we construct a second TP layer with tp_grain_size=64 and call _freeze_partition_sizes on it? This would verify that creating and initializing another layer does not modify the first layer’s per-model metadata or frozen partition sizes.
If feasible, it would also be useful to cover GQA explicitly: construct the second TP layer with a different num_kv_heads value, then verify that the first layer still retains its original num_kv_heads and partition layout. This would directly test that the per-model AutoTPMeta state is not overwritten by another model.
There was a problem hiding this comment.
Done in 606ec52. The stray other_meta assert is gone; the test now asserts layer.tp_meta.tp_grain_size == 1 directly, and builds a real second LmHeadLinearAllreduce with tp_grain_size=64, num_kv_heads=1 (GQA), runs _freeze_partition_sizes on it, then re-verifies the first layer's meta (tp_grain_size, num_kv_heads) and frozen (51, 50) split are untouched before the forward call.
Per maintainer feedback the tests-to-preserve tree is tests/unit/v1, which is also the scope the modal GPU workflow selects, so the AutoTP tests run on multi-accelerator CI once they live there. Update the xpu workflow's path accordingly. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
|
Moved |
…ayers Signed-off-by: Guokai Ma <guokai.ma@intel.com>
The GPU CI run surfaced a NameError in register_replicated_grad_hooks: print_dist was called without being imported (log_dist was). Add it to the existing deepspeed.utils.logging import. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
|
The two GPU-CI failures from the previous run are now understood:
|
|
@sfc-gh-truwase Opened #8291 to track the Ulysses |
Moving them into tests/unit/v1 exposed test_tp_plan_real_models to the modal workflow, which tests against transformers main; that main now injects 'embedding_rowwise' into tp_plan for tied-embedding models (#8290), so every full-suite modal run on master would fail until the upstream drift is handled. Move the directory back for now; it will be relocated into tests/unit/v1 once #8290 is fixed. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
|
Reverted the move of Note the run itself proved the coverage value: it caught a real |
The full model_parallelism directory stays out of tests/unit/v1 until the transformers-main 'embedding_rowwise' drift (#8290) is handled, but the two multi-model regressions this PR adds (#8231 reproduction and the multimodal text_config path) are safe there and belong to the GPU workflow's scope. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
|
Follow-up to e9612b3: the two multi-model regression tests this PR adds now live in |
3 participants
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.