Repository navigation
Conversation
0a03e50 to
df831ba
Compare
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![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: df831bac7c
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
dcf58b7 to
3035791
Compare
GPTBigCode's fused c_attn stores the query rows followed by the whole shared K/V block. transformers 5.x exposes num_key_value_heads=1 for multi-query models, so AutoTP split the fused attention by one head and left every other rank with an empty query shard. Use the query-head count for replicated-KV layouts, keep the K/V head attributes whole, recompute num_key_value_groups from the local shard, and reject TP widths larger than the query-head count from the shared AutoTP path. Reject an explicit partition_config for MQA until the heuristic attributes can follow an arbitrary shard. Make the K/V gradient reduction backend-safe: stage a contiguous copy when the view is non-contiguous, then write the reduction back. Register the hook through register_grad_hook so it runs in ZeRO's post-accumulate phase on torch_npu, and accept the two-argument AccumulateGrad signature used before torch 2.1. Signed-off-by: Venus-Tzz <1601677496@qq.com>
Add TP=2/TP=4 parity with ZeRO 0/1/2/3, shard-shape and checkpoint metadata checks, inference logits, TP-width rejection, unit coverage for the replicated-KV attributes and explicit-config guard, and the legacy AccumulateGrad hook signature. Signed-off-by: Venus-Tzz <1601677496@qq.com>
3035791 to
9066a56
Compare
| # group count follows the local query shard. | ||
| replicated_kv_attention = (getattr(self, "_replicates_kv_block", False) | ||
| and getattr(child, "multi_query", False)) | ||
| kv_attrs = ("num_key_value_heads", "num_kv_heads", "kv_n_heads", "num_kv") |
There was a problem hiding this comment.
why we need to recognize for exactly these four names? There are 8 items in _KV_HEAD_ATTRS and 4 are picked.
| setattr(child, param, get_shard_size(param_val, self.mp_size, self.tp_meta, name, rank=tp_index)) | ||
| # Recompute groups from the local query shard. Older transformers versions expose | ||
| # the K/V head count under different names. | ||
| num_heads = getattr(child, "num_heads", None) |
There was a problem hiding this comment.
Similar here, the code is going through a list of names ("kv_heads", "num_key_value_heads"), why this list and why its different from kv_attrs? Also this code should first define a list then loop through this list for better maintainability.
| # shards by its own model's kv-head / grain values and keeps its own training mode. | ||
| # None is kept in the layers so they keep following set_autotp_mode, as before. | ||
| self.tp_meta = AutoTPMeta.from_model_config(model_config, tp_grain_size, training=training_mode) | ||
| self._replicates_kv_block = replicates_kv_block(self.module) |
There was a problem hiding this comment.
This "replicated K/V ⇒ shard by query heads" policy now lives in five places
(engine warning, init override, init config rejection, _validate_replicated_kv_tp_width,
update_mp_params), and only the detection is shared — nothing keeps the five agreeing.
Could the derived values (shard unit, K/V row span, group count) come from one helper, plus an
assert that kv_attrs ⊆ param_list, so a later edit can't silently diverge?
There was a problem hiding this comment.
Many thanks for your suggestions.
_SHARDED_KV_HEAD_ATTRS contains four names because it is exactly the intersection of the module K/V-head aliases (_MODULE_KV_HEAD_ATTRS) and the attributes update_mp_params shards (_MP_PARAM_ATTRS); _KV_HEAD_ATTRS is a config-level probe, so it also includes query-head aliases (the kv == q fallback) and names this loop never shards.
kv_heads is GPTBigCode's module attribute and is not in param_list, so it is lookup-only; num_key_value_heads is the HF GQA name and stays whole for this replicated-K/V layout, while the normal layout still shards it. Group-count lookup now iterates _MODULE_KV_HEAD_ATTRS, and _SHARDED_KV_HEAD_ATTRS is the subset that must not be split.
I centralized the policy in fusedqkv_utils.py (replicated_kv_shard_*, replicated_kv_query_widths, replicated_kv_row_start, replicated_kv_group_count), so all five call sites share the same derivation, with subset assertions and alias coverage in the new tests.
Route the shard unit, K/V row span, and group count through fusedqkv_utils helpers so the engine warning, init path, width validation, and update_mp_params share one derivation. Split the config, module, and sharded K/V attribute lists, and assert the subset relations so later edits cannot silently diverge. Signed-off-by: Venus-Tzz <1601677496@qq.com>
9066a56 to
355e6c0
Compare
Keep the pre-accumulate tensor hook for the normal path and add a post-accumulate fallback for backends that skip it. The fallback reduces only the K/V rows added since its previous call, so gradient accumulation does not repeatedly scale earlier micro-batches. Add a regression test covering both hook paths. Signed-off-by: Venus-Tzz <1601677496@qq.com>
4a4adf4 to
1a19b1a
Compare
|
Hi @VenusTZZ thanks for the quick response. I have two questions:
|
Many thanks for the detailed questions. For questions 1: The starting point was upstream #8928. When it landed as master HEAD, I ran its GPTBigCode AutoTP training test (TestGPTBigCodeMultiQueryTraining) on 4x Ascend 910B4 as part of preparing NPU CI coverage for DeepSpeed (Ascend-CI). The test hung/diverged on NPU: transformers 5.x produced an empty query shard, HCCL rejected the non-contiguous K/V gradient view, and after those were fixed, the ZeRO 0/1/2/3 coverage showed the plain tensor hook was skipped under ZeRO 1/2 (4 K/V all-reduces under ZeRO 0/3 in the tiny 2-layer model, zero under ZeRO 1/2), so the replicated K/V rows were never summed. The same cases pass on A100, where the hook fires under ZeRO 1/2, so the current CUDA CI cannot see this. Fixing that required a post-accumulate fallback. The last commit fixes the bug that fallback introduced: with GAS>1 it rewrote the already accumulated K/V gradient on every backward, so the first micro-batch was all-reduced and rescaled again on the second one. I found it with a temporary ZeRO-0 + GAS=2 end-to-end parity check against an unsharded two-micro-batch reference, and fixed it by keeping the pre-accumulate hook for the normal path and reducing only the K/V rows added since the previous call in the fallback. The committed regression is test_replicated_kv_gradient_reduces_each_micro_batch_once. So there are two validation gaps: (a) no NPU CI, which the Ascend-CI work is addressing, and (b) the committed parity matrix is GAS=1 only. |
5a1c0d3 to
e02f9f3
Compare
For questions 2: I ran three variants: A is the current per-micro-batch K/V reduce; B turns the K/V reduce off completely (not a real configuration, just a reference point for the communication cost); C reduces once at the end of the accumulation window. "Accumulation steps" means gradient accumulation steps, i.e. how many micro-batches go into one optimizer update. I tested the same 24-layer/hidden=2048 GPTBigCode MQA shape on both NPU (4×Ascend 910B4) and GPU (4×A100-SXM4-80GB), with TP=4, bf16, ZeRO-1:
For the 24-layer case, the K/V reduce count per optimizer step drops from 1,536 to 48 at 32 accumulation steps (384 to 48 at 8). The gain is bigger on NPU; on A100 it is about 10% in this run. It is also shape-dependent: on NPU, 8 layers/hidden=2048 at 8 steps goes from 511 ms to 450 ms (-11.9%, +13.5% tokens/s), 4 layers/hidden=4096 shows no measurable gain, and one accumulation step has no gain by construction. ZeRO-2/3, ZeRO-1 with a separate accumulation buffer, offload, BF16_Optimizer, GAS=1, pipeline, and windows starting under You're right about the scope concern. This is a performance optimization, not a correctness fix, and the diff has grown a lot. Review also turned up two deferral bugs — the gate allowed configurations where earlier micro-batches had left Would you be OK with splitting this? I'd keep #8962 focused on the correctness fixes and open a follow-up PR for the deferred reduce with this data and a smaller design: the existing hook, one helper for the deferral condition, per-window |
@VenusTZZ thanks for investigating this. I agree with a separate PR to review the optimization. Thanks! Will approve this PR when the related commits split out, thanks! |
e02f9f3 to
1a19b1a
Compare
gpu-test-labeler
Bot
added
the
skip-gpu-test
Fixes #8961.
Scope
This PR fixes the empty attention shard, the non-contiguous K/V gradient, and the post-accumulate hook rescaling earlier micro-batches under gradient accumulation in AutoTP multi-query GPTBigCode at TP > 1. Parity tests use
gradient_clipping: 0.0; the default clipping path is a separate pre-existing issue.The branch now contains only the four correctness commits (head
1a19b1aab). The deferred-reduce optimization discussed in review was split into follow-up PR #9064 (branchfix-gptbigcode-autotp-defer-kv, 2 commits, stacked on this one); this PR is correctness-only and should merge first.Background
GPTBigCode's fused
c_attnstores the query rows followed by the whole shared K/V block. With transformers 5.x,num_key_value_heads=1makes AutoTP create empty query shards at TP > 1; some backends also return a non-contiguous K/V gradient, and torch_npu skipped the plain hook under ZeRO 1/2. A post-accumulate fallback is needed there, but it must not reduce the already accumulated K/V gradient again on the next micro-batch.Change
num_key_value_groupsfrom the local shard.tp > query headsfrom the shared AutoTP path, and reject an explicitpartition_config/ HFtp_planfor this layout until attribute sharding can follow the actual shard.AccumulateGradsignature). The fallback reduces only the K/V rows added since its previous call, so a second backward under gradient accumulation does not all-reduce and rescale the first micro-batch again.engine.py; the replicated-KV rejection lives in AutoTP.Files changed
deepspeed/module_inject/auto_tp.pydeepspeed/module_inject/fusedqkv_utils.pydeepspeed/module_inject/layers.pydeepspeed/module_inject/tp_shard.pydeepspeed/runtime/engine.pytp > query headsfor this layout; keep the non-replicated warning.tests/unit/v1/autotp/test_autotp_custom_patterns.pytests/unit/v1/autotp/test_autotp_inference.pytests/unit/v1/autotp/test_autotp_training.pyTest
Run from
tests/unit:TestGPTBigCodeMultiQueryTraining(TP=2): runs ZeRO 0/1/2/3 with DP=2 (world_size=4); each case checks four-step loss parity and the K/V hook, while stage 0 also checks shard shapes, checkpoint metadata, and identical K/V rows.TestGPTBigCodeMultiQueryTraining(TP=4): runs stage 0 and checks the four-way query split.TestGPTBigCodeMultiQueryInference: compares TP=2 logits with the unsharded model; illegal-TP rejection is covered by the shared width-check unit test.gradient_clipping: 0.0.Results
NPU: 4×Ascend 910B4 (32 GB each), CANN 8.5.1, torch 2.10.0, torch_npu 2.10.0.post7, transformers 5.19.0. At the split tip
1a19b1aab:test_autotp_custom_patterns.py+test_autotp_inference.py: 45 passed in 417.04 s.TestGPTBigCodeMultiQueryTrainingplus the three gradient-reducer unit cases: 8 passed in 133.07 s.pre-commitpasses on all changed files.CUDA: 4×A100-SXM4-80GB, driver 580.126.16, torch 2.14.1+cu126, NCCL 2.29.3.
218cbf242, which contains this branch) passed 15 tests: 5 training, 8 accumulation, and 2 reducer unit tests. The full target suite (custom patterns + inference + training) has not been rerun on A100 after the split.Validation matrix
Not covered
c_attnslice is 1.414×, the full-model norm is 1.0027× in the tiny model, and a proper clipped-vs-clipped 4-step run differs by 4.72e-4. A clipped-vs-unclipped run is not a valid parity check. Tracked in [BUG] AutoTP counts replicated GPTBigCode key/value rows once per tensor-parallel rank in the global gradient norm #9024.multi_query=Falseis unchanged; its forward failure under AutoTP is pre-existing and noted in [BUG] AutoTP training of multi-query GPTBigCode diverges from single-device training at TP > 1 #8927.tp_metaadjustment assumes one replicated-KV layout per model.generateand a past-KV-cache run are not covered.CI
engine.pyis a shared runtime file, soci/tests_fetcher.pyselects the full suite. GPU workflows run in the merge queue; the on-device results above are the validation for this change.AI assistance
AI assistance was used to reproduce the failure, implement the change, and run the tests; the human submitter reviewed every changed line and is responsible for the change.
Related