Repository navigation
fix(gemma4): drive the TP ranks concurrently so long prompts do not stall - #1133
Conversation
github-actions
Bot
added
gemma4
af0046d to
20dfa50
Compare
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Thanks; one thread per rank is the right fix for the stall. Two blockers and three asks are inline. Also, the mixed step now takes prompts of any length at two ranks with no TP gate, so could the old bound stay until a gate covers it? The tp.md TL;DR and the chunked walk bullet still describe the single threaded driver.
| handles.push(scope.spawn(move || -> Result<()> { | ||
| activate_rank(&state.ctx)?; | ||
| extra(state, item) | ||
| .with_context(|| format!("tensor-parallel rank {rank} segment"))?; | ||
| state | ||
| .ctx | ||
| .sync() | ||
| .with_context(|| format!("drain rank {rank}")) | ||
| })); |
There was a problem hiding this comment.
Each step spawns a fresh thread, and activate_rank runs cublas_init on it, creating two cuBLAS handles that nothing destroys at thread exit (CublasThreadGuard covers only the scheduler thread). A standalone copy of this pattern on sm_89 leaks 16.25 MiB of device memory per step (3250 MiB over 200 steps, none with cublas_destroy), and decode steps take it too. Could each extra rank run on a long lived thread holding its own guard?
| @@ -2269,97 +2260,28 @@ impl EngineState { | |||
| scores.extend(scored.into_iter().map(Some)); | |||
| Ok(()) | |||
| }; | |||
| activate_rank(&self.ctx)?; | |||
| // Rank 0's tower is launched first, but its per-row readback blocks | |||
| // on its own collective, so it runs only after every rank has | |||
| // launched its tower — the peer's matching call must already be in | |||
| // flight or the one driving thread deadlocks. | |||
| // | |||
| // A rank-0 failure lands in `stepped` rather than returning from | |||
| // here, so the `tp_broken` gate below is what prices it: a broken | |||
| // pair stops the engine, while a single rank — where `abort_comms` | |||
| // has nothing to abort and leaves `tp_broken` clear — keeps the | |||
| // request-scoped failure this path has always had. | |||
| 'scored: { | |||
| let mut tower = match self.serve.launch_prompt_tower(ctx, &mut kv, prompt) { | |||
| Ok(tower) => tower, | |||
| Err(err) => { | |||
| abort_comms( | |||
| &self.tp_broken, | |||
| &self.ctx, | |||
| &mut self.serve, | |||
| &mut self.arena, | |||
| &mut self.more, | |||
| ); | |||
| break 'scored Err(err.context("scored prompt tower")); | |||
| } | |||
| }; | |||
| // The other ranks run the same prefill to write their KV | |||
| // shards; scoring is rank 0's alone. A failure on them is fatal | |||
| // at any world size: the ranks' frontiers must not drift apart. | |||
| if let Err(err) = self.step_extra_ranks(&mut kv, |state, rank_kv| { | |||
| state | |||
| .serve | |||
| .step_scoring(&state.ctx, rank_kv, prompt, None) | |||
| .map(|_| ()) | |||
| }) { | |||
| abort_comms( | |||
| &self.tp_broken, | |||
| &self.ctx, | |||
| &mut self.serve, | |||
| &mut self.arena, | |||
| &mut self.more, | |||
| ); | |||
| break 'scored Err(err); | |||
| } | |||
| // Every rank's tower has now completed and its collectives are | |||
| // paired, so a failure in rank 0's readback costs this request | |||
| // alone at any world size: the next step's collective sequence | |||
| // is intact, and the request's pages return with `kv` on every | |||
| // rank. | |||
| if let Err(err) = self | |||
| .serve | |||
| .score_prompt_tower(ctx, &mut tower, &mut score) | |||
| .context("prompt logprobs") | |||
| { | |||
| break 'scored Err(err); | |||
| } | |||
| echo = Some(PromptEcho { | |||
| ids: prompt.clone(), | |||
| logprobs: scores, | |||
| }); | |||
| Ok(tower.into_logits()) | |||
| } | |||
| self.serve | |||
| .score_prompt_tower(ctx, &mut tower, &mut score) | |||
| .context("prompt logprobs")?; | |||
There was a problem hiding this comment.
This ?, the one on score_prompt_tower, and the solo arm's all return from admit_and_prefill, so a failed prefill at world size 1 now stops the engine; on main it cost only that request. a_failed_prefill_costs_that_request_not_the_engine should fail here. Could both results stay in stepped?
| if ranked.is_err() && extra_ranks > 0 { | ||
| // Rank 0's sequence may have stopped short, so a peer's matching | ||
| // call will never come. Rank 0's own fields are not borrowed by | ||
| // the threads, so its communicator and graphs can be torn down | ||
| // here to unblock them; the extras' follow the join. | ||
| // | ||
| // Only with extra ranks: with none there is no collective to | ||
| // abort and no peer frontier to keep in step, so this stays the | ||
| // ordinary single-rank failure the caller answers by failing that | ||
| // one request. Marking the engine broken there would turn a | ||
| // request-local error — a per-request scratch allocation — into | ||
| // the end of the engine. | ||
| tp_broken.set(true); | ||
| if select_device(ctx).is_ok() { | ||
| arena.release_graphs(); | ||
| } else { | ||
| // A device that will not come current is the one case where | ||
| // the release cannot run — and it is exactly the case a | ||
| // sticky fault produces, so name the order we are about to | ||
| // break: the abort below then runs against a live graph. | ||
| log::error!( | ||
| "rank 0's device {} will not come current, so its graphs stay alive \ | ||
| across the communicator abort; that abort may wedge", | ||
| ctx.device_ordinal | ||
| ); | ||
| } | ||
| serve.detach_tp_comm(); |
There was a problem hiding this comment.
Dropping rank 0's TpComm aborts only rank 0's communicator; a peer already inside the matching collective stays in ctx.sync(), and its abort in break_extras runs after this join. Could every communicator be aborted before the join, with one TP2 run that injects a rank 0 failure to show the engine stops?
| // The line is `serve_oracle`'s calibration, not a tight one: scoring the same | ||
| // context two ways (there a greedy walk against a single prefill, here a | ||
| // 512-row prefill against a 1024-row one) moves the logits by a chaotically | ||
| // amplifying amount — measured in-repo at 0.31..5.75 raw logits — because the | ||
| // prefill GEMM's tiling depends on the row count. Measured here: 1.69 on the | ||
| // four-layer synthetic, 4.05 on the 60-layer 31B. A grosser gap (a page, | ||
| // position or rope error) still fails. | ||
| assert!( | ||
| worst < PREFIX_LINE, | ||
| "the same prefix scored in a longer run differs by {worst} at {worst_at} (line \ | ||
| {PREFIX_LINE}), so the engine is not prefix-consistent" | ||
| ); | ||
| } | ||
|
|
||
| /// The prefix gate's line, on the `serve_oracle` scale (`DRIFT_LINE = 12.0` | ||
| /// there) rather than the one-rank comparison's `LOGBROB_LINE`: the drift it | ||
| /// looks for is the same chaotic shape-dependence, bounded the same way. | ||
| const PREFIX_LINE: f32 = 12.0; |
There was a problem hiding this comment.
serve_oracle's 12.0 is a raw logit line held beside an argmax check on every row; here it bounds logprob gaps with flips only counted, and nothing shows a position error exceeds it. Could this be held against the same gate at world size 1 on the 12B, as #1132 does with its cold leg?
| // Only the nine-token case is gated. The 1024-token "edge" case runs (the | ||
| // concurrent driver carries it), but at window width its drift is the | ||
| // shape-dependent chaos `serve_oracle`'s `neutral_scale` already measures at | ||
| // 0.31..5.75 raw logits — worst shared-token logprob gap 8.37 here — so a | ||
| // strict top-k containment cannot hold and the case needs that gate's | ||
| // tolerance-plus-top-1-share discipline. See docs/models/gemma4/tp.md "Known | ||
| // bounds"; the calibration is the follow-up. | ||
| const CASES: [&str; 1] = ["short"]; |
There was a problem hiding this comment.
In the 12B window fixture, HF's own sdpa and eager rows never put a pick outside the other's top 64 at any length up to 4096, so four such rows is more than chaos. Could the_two_rank_engine_matches_one_rank at PEGAINFER_TP_PROMPT_TOKENS=1024 on the 12B show whether the second rank adds it?
…tall The single-threaded driver ran rank 0's whole segment before any peer's, so a host-blocking operation inside it — rank 0's per-row readback, and a GEMM workspace allocation behind it — waited on a collective whose peer call had not been issued yet, and the scheduler hung. Admission refused a prompt past the length the gates had been verified at (`PEGAINFER_TP_MAX_PROMPT`) to keep that from being reached. `EngineState::drive_ranks` runs rank 0's segment on the scheduler thread and each extra rank's on its own scoped thread, then joins, so every peer's collectives are in flight before rank 0 can block. The payload is transposed once up front — `rank_payloads`/`prefill_payloads`, with `GemmaKv::split` lending each rank one mutable family per request — because the threads need disjoint borrows. The six rank-driving paths (solo, scored and chunked prefill, the mixed step, both decode paths) go through it, and the ceiling, its parser, its knob and its unit assertions are deleted rather than left as scaffolding. A rank-0 failure releases rank 0's graphs and detaches its communicator *before* the join, so a peer already waiting on a call that will never come is unblocked; the extras' teardown follows the join, and the engine is marked broken only when there is a peer to keep in step — so a world-size-1 failure stays request-scoped, which `a_failed_prefill_costs_that_request_not_the_engine` pins. Verified on 2×L20 (`NCCL_PROTO=LL128`): the stall is gone — a scored prefill of 64, 128, 256 and 1024 tokens all pass on the 60-layer 31B, where 128 and up used to hang. The 12B replicate parity keeps `22/24` picks at gap `0.5738`; the 31B-geometry synthetic keeps `48/48` at `0.0000`; the real 31B against its HF dump keeps `7/8` rows, worst shared-token gap `0.6753`; the new prefix gate reads `4.0538` against its `12.0` line. `fmt`, `clippy --all-targets -D warnings` and the single-rank suite (63 passed) are clean. Signed-off-by: bohaotu <bohaotu@pony.ai>
The TP notes called the ceiling's hazard "a GEMM workspace allocation at longer prompts". The capture does not support that: the frozen frame is a lazy cuBLAS/cuBLASLt library load (`cuLibraryGetModule`, reached through `cublasGemmEx`) inside a projection GEMM — a first touch of a shape, not a function of the prompt's length, and the captured prompt was 5 tokens, with rank 0's device 100 % busy and rank 1's idle while the one driving thread waited on a reduction whose peer call it had not issued. With the ceiling lifted the single-threaded driver completes scored prefills of 50/64/100/128/300 tokens and a 401-token server prefill (4-bit 60-layer geometry; the bf16 31B is not re-measured), so the guard was refusing prompts that serve. Comments and the TP notes only; no code change. Signed-off-by: bohaotu <bohaotu@pony.ai>
Review follow-ups on the concurrent driver and its new gate. - `drive_ranks` now checks that the payload covers every rank with an `ensure!` rather than a `debug_assert_eq!`. The world size is inferred from the batch (`rank_payloads`' empty-slice default is one), so in a release build a mismatch would drive fewer ranks than the step has and leave a peer's collective unissued — a hang, not an error. It runs before any launch, so refusing there costs the step nothing. - The rank-0 failure path and `break_extras` name the case where the graph release is skipped (a device that will not come current, i.e. a sticky fault), because that is the one case where the communicator abort then runs against a live graph — the order the teardown forbids. - The new prefix gate asserts its compared rows are finite (the hole round 8 closed in the other two comparisons) and that the two runs shared at least one token, since the gap is only computed where their top-k overlap. - `PEGAINFER_TP_PROMPT_TOKENS=1` no longer wraps `0..short - 1` in release. - `GemmaServe`'s SAFETY note names the per-step scoped threads that drive the extra ranks now, and the TP notes record that the mixed step has no gate at two ranks once the ceiling's bounds went with the ceiling. Signed-off-by: bohaotu <bohaotu@pony.ai>
dbdee97 to
5404aec
Compare
…ilure scope Six review items on pegainfer-project#1133, two of them blockers. - Every step's extra-rank thread created a cuBLAS handle pair in `activate_rank` that nothing destroyed (the scheduler's `CublasThreadGuard` covers only the scheduler thread), leaking device memory per step. The thread now takes the same guard, so the handles go with it. - A failed prefill returned from `admit_and_prefill` through `?` on the scored and solo arms, so at world size 1 a request-local failure stopped the engine. Both arms' failures land in `stepped` again, where the `tp_broken` gate prices them — the shape main had, and what `a_failed_prefill_costs_that_request_not_the_engine` pins. - `TpComm` is a shareable handle now: the engine holds one per extra rank, and a rank-0 failure aborts *every* communicator before the join, so a peer inside the collective is released by its own abort rather than left to notice rank 0's. `an_aborted_communicator_refuses_the_next_reduction` pins the handle's contract (works before, refuses after, aborts once) without a checkpoint. - The prefix gate's 12.0 was borrowed from `serve_oracle`'s raw-logit line with no control arm; `...within_a_single_rank_control` runs the same comparison at world size 1 on a one-card checkpoint and bounds the two-rank arm by a multiple of it. - `tp.md`'s TL;DR and its chunked-walk bullet still described the single-threaded driver, and a follow-up this branch is. Signed-off-by: bohaotu <bohaotu@pony.ai>
`crate::tensor` is not a module of this crate; the comm gate's contexts come from `pegainfer_core::tensor`, as the serving code already does. Signed-off-by: bohaotu <bohaotu@pony.ai>
The control arm and the 1024-token discriminator the review asked for both ran on the 12B (two L20s, NCCL_PROTO=LL128), and the numbers move the notes: - The prefix comparison at window width reads 6.30 nat at one rank and 9.23 at two — a ratio of 1.46, so the two-rank drift is the same phenomenon rather than a multiple of it. The control constants now carry that reading, and the absolute PREFIX_LINE (12.0, from the 31B's 4.05) is recorded as unable to see either. - The 1024-token case is therefore an open defect rather than the chaotic band the notes claimed: the two-rank-vs-one-rank gate keeps 1/4 picks with a 3.81 nat gap against its 1.0 line, and a *one-rank* arm disagrees with itself across the window by 6.30. The HF gate's comment and tp.md's Known bounds now say that, and name the HF anchor at that length as the follow-up. Signed-off-by: bohaotu <bohaotu@pony.ai>
|
Ran the review's asks on the branch head (
The handle leak (#1): the same 12B TP2 server, The 1024-token case is a defect, not the chaotic band — this is the one thing the runs turned up that is not fixed here. At the window width the two engines disagree materially: and the prefix comparison says the length is the variable, not the second rank: scoring the same 512-token prefix inside a 1024-token run reads 6.3047 on one rank and 9.2288 at two (149/115 argmax flips in 511 rows). So a single rank disagrees with itself across the window width, Two caveats on what these runs are: the rank-0 fault in the new TP2 gate is symmetric ( |
The review's discriminator ran, and then the anchors did. - `the_two_rank_engine_matches_one_rank` fails at every length from 128 up (worst shared-token gap 1.87 at 512, 4.52 at 768, 5.50 at 256, 3.81 at 1024 — non-monotone, so the tiling-dependence and not a threshold), against a 1.0 line that is a 64-token calibration. - `serve::oracle::context_waypoints_match_hf` passes, the window width included: 1.15 at w1023, 1.48 at w1024 (tolerance 4.61), 3.15 at w1025, 3.80 at w4096, 1.17 at w16384 and 5.23 at w32768 (7.65). - The 12B golden's `edge` (1024) case is carried but not probed (`probed_cases: ["single", "short"]`); running it fails a strict top-64 containment on ~5-8 rows at the same rows with one rank as with two, so the second rank does not add it. The notes called it an open defect for one commit; they now record the anchors, the fixture's own marking, and that a strict line at that depth is a flake rather than a detector. Signed-off-by: bohaotu <bohaotu@pony.ai>
|
Follow-up on the 1024-token question, now with the anchors. The second rank does not add it. Everything below is the 12B, fixtures from Length bisect ( The reference anchors pass, the window width included ( The So the conclusion for the review's second ask is negative on the hypothesis: the two-rank engine adds nothing at 1024 that one rank does not already show, and the parity gate's 1.0 line is a 64-token calibration that cannot hold at depth (the same-workload reading is 0.57 at 64 tokens). Two notes on method: the one-rank One incidental failure, not from this PR: |
|
Following up on the incidental failure I reported in my previous comment — Its two arms are a one-token decode against a whole-prompt prefill, which the gate itself acknowledges are not bit-equal (it bounds their raw-logit drift at 2.0, and Same five ids, same tail, top two 0.0625 apart on one side and 0.1875 on the other, rows agreeing to 0.3125 everywhere else — a near-tie flip inside the drift the gate allows, not a moved distribution. Under the repo's fleet-wide rule (a flipped pick must be a near tie, each pick inside the other row's top-k — That fix is #1136 (test-only, one function, |
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Thanks for the runs. The leak, the failure scope at one rank and the control arm are fixed, and the one rank control settles the 1024 question. Two items remain on the abort path and two on the gates, inline. Also, the first known bound in tp.md still says rank 0's communicator is detached and that no gate injects a TP2 failure, its per step cost now includes a handle pair created and destroyed (about 0.8 ms in the standalone loop), and the mixed step bound from the last round is still open.
| handles.push(scope.spawn(move || -> Result<()> { | ||
| // This step's thread creates its own cuBLAS handles in | ||
| // `activate_rank`, and the guard is what destroys them when | ||
| // the thread ends: the one the scheduler holds covers only | ||
| // the scheduler's thread, so without it every step would | ||
| // leave a handle pair's device memory behind. | ||
| let _handles = bind_engine_thread(&state.ctx)?; |
There was a problem hiding this comment.
If this thread fails before its first collective (the handle pair created here every step is about 16 MiB, so an allocation failure lands right here), rank 0 keeps launching, and the first host wait in its segment, the stall this PR removes, sits behind a reduction rank 1 never issues. The abort only runs on rank 0's own error or after the join, so the step hangs. Could each thread hold the comm handles and abort all of them on its own failure, as the rank 0 branch does?
| if ranked.is_err() && extra_ranks > 0 { | ||
| // Rank 0's sequence may have stopped short, so a peer's matching | ||
| // call will never come — and a peer inside that collective is | ||
| // unblocked by aborting *its* communicator, not rank 0's. Every | ||
| // communicator is aborted here, before the join that would | ||
| // otherwise be waiting on the peer whose abort is still to come. | ||
| // Rank 0's own fields are not borrowed by the threads, and the | ||
| // peers' communicators are reached through the engine's own | ||
| // handles rather than their serves. | ||
| // | ||
| // Only with extra ranks: with none there is no collective to | ||
| // abort and no peer frontier to keep in step, so this stays the | ||
| // ordinary single-rank failure the caller answers by failing that | ||
| // one request. Marking the engine broken there would turn a | ||
| // request-local error — a per-request scratch allocation — into | ||
| // the end of the engine. | ||
| tp_broken.set(true); | ||
| if select_device(ctx).is_ok() { | ||
| arena.release_graphs(); | ||
| } else { | ||
| // A device that will not come current is the one case where | ||
| // the release cannot run — and it is exactly the case a | ||
| // sticky fault produces, so name the order we are about to | ||
| // break: the abort below then runs against a live graph. | ||
| log::error!( | ||
| "rank 0's device {} will not come current, so its graphs stay alive \ | ||
| across the communicator abort; that abort may wedge", | ||
| ctx.device_ordinal | ||
| ); | ||
| } | ||
| serve.abort_tp_comm(); | ||
| for comm in extra_comms.iter() { | ||
| comm.abort(); | ||
| } |
There was a problem hiding this comment.
The extras' graphs are still alive here, since their arenas stay on the rank threads until the join, so this aborts their communicators in the order the Drop comment on EngineState says wedges. Is that safe with decode graphs on?
| fn an_aborted_communicator_refuses_the_next_reduction() { | ||
| let (device, peer) = devices(); | ||
| assert_ne!(device, peer, "TP2 needs two distinct device ordinals"); | ||
| let contexts: Vec<DeviceContext> = [device, peer] | ||
| .into_iter() | ||
| .map(|ordinal| DeviceContext::new_with_device(ordinal).expect("a device context")) | ||
| .collect(); | ||
| let comms = cudarc::nccl::safe::Comm::from_devices( | ||
| contexts | ||
| .iter() | ||
| .map(|ctx| ctx.stream.clone()) | ||
| .collect::<Vec<_>>(), | ||
| ) | ||
| .expect("a two-rank communicator"); | ||
| let ranks: Vec<crate::layer::TpComm> = | ||
| comms.into_iter().map(crate::layer::TpComm::new).collect(); | ||
| let mut buffers: Vec<_> = contexts | ||
| .iter() | ||
| .map(|ctx| ctx.stream.alloc_zeros::<f32>(8).expect("a buffer")) | ||
| .collect(); | ||
| // The pair reduces together once, before any abort, each rank's enqueue on | ||
| // its own device — the shape the engine runs and the one the wrapper must | ||
| // not have broken. | ||
| for (index, (rank, buffer)) in ranks.iter().zip(buffers.iter_mut()).enumerate() { | ||
| super::select_device(&contexts[index]).expect("the rank's device is current"); | ||
| rank.all_reduce_in_place(buffer, &cudarc::nccl::safe::ReduceOp::Sum) | ||
| .expect("a reduction over a live communicator"); | ||
| } | ||
| for ctx in &contexts { | ||
| ctx.stream | ||
| .synchronize() | ||
| .expect("the pair's reduction completes"); | ||
| } | ||
| ranks[1].abort(); | ||
| ranks[1].abort(); | ||
| assert!( | ||
| ranks[1] | ||
| .all_reduce_in_place(&mut buffers[1], &cudarc::nccl::safe::ReduceOp::Sum) | ||
| .is_err(), | ||
| "a reduction after the abort must refuse instead of reducing comm-less" |
There was a problem hiding this comment.
This holds the reduction after the abort; the engine relies on the abort releasing a reduction already waiting. Could it enqueue a reduction on rank 1 alone, abort rank 1 from the test thread, and require rank 1's stream to synchronize within a deadline? That covers the asymmetric case without an engine hook.
| // Only the nine-token case is asserted: the fixture's own metadata marks | ||
| // `edge` as carried but not probed (`probed_cases: ["single", "short"]`), and | ||
| // running it anyway fails a strict top-64 containment on ~5-8 rows — the same | ||
| // rows with one rank as with two, so it is the depth-dependent drift these | ||
| // fixtures carry tolerances for, not a two-rank difference. See | ||
| // docs/models/gemma4/tp.md "Known bounds". | ||
| const CASES: [&str; 1] = ["short"]; |
There was a problem hiding this comment.
probed_cases lists the cases that carry hidden state probes (hf-golden.md: edge has none because they would dwarf the file), so it does not mark edge as unasserted. The same rows failing at one and two ranks is what shows the second rank adds nothing; could this comment and tp.md rest on that alone?
A rank that stopped issuing collectives left whoever waited on the matching call stranded: only rank 0's own error aborted the communicators, and rank 0 cannot reach that branch while it is blocked in a host wait — so an extra rank's failure hung the step instead of stopping the engine. Each step thread now holds every communicator as a handle and aborts all of them on its own failure, and a failed extra rank marks the engine broken the way a failed rank 0 does. That failure path inverts the release-then-abort order the teardown order documents, and it has to: a peer stopped inside a collective keeps its graphs on its own thread and cannot release them, and the abort is the only thing that gets it out of one. The window is named in tp.md as a bound argued from the abort's own contract rather than measured, and the contract gains the gate it was missing — a reduction enqueued with no peer, aborted from another thread, required to stop waiting and to return, both within a deadline. The mixed step gets back the row bound the deleted TP ceiling also carried (it is the one TP region no gate covers), the long-prompt bound stops reading `probed_cases` as "asserted" and rests on the one-rank/two-rank row identity alone, and tp.md's first bound stops claiming a detached communicator and no TP2 fault gate while also naming the per-step cuBLAS handle pair. Signed-off-by: bohaotu <bohaotu@pony.ai>
|
这轮 7 条都落在 1. extra rank 自己失败时无人 abort（ 2. abort 时 peer 的图还活着（ 3. abort 门禁没覆盖"放掉正在等待的归约"（ 4. 5. 第一条 known bound 双重过期 —— 已重写。 6. per-step 成本漏了句柄对 —— 已补。 写明 7. 混合步上界 —— 按你的建议放回来了。 验证两个 comm 门禁在双 L20 上跑过（stock NCCL 2.18.3 / CUDA 12.9 / sm_89，
第二条第一次实测了引擎失败路径所依赖的两件事： |
Mutual top-8 membership is not a tie: `[6, 0, …]` against `[0, 6, …]` has each pick as the other run's runner-up while both sit six nats from their own runner-up, and LOGBROB_LINE cannot see that either — it bounds a token both runs kept. distribution_gap now also reads the two picks' margin inside each run's own top-k and requires it within NEAR_TIE_WIDTH — the same quantity serve_oracle's near-tie rule reads (a difference inside one row, so the same on log_softmax as on raw logits) and the same 1.0 LOGBROB_LINE uses between the two runs. The rule's halves and the fold are pinned by unit tests that need no device, one of which also holds the walk to stopping at the flip, so the steps after it are not compared as if the two runs still shared a context. Signed-off-by: bohaotu <bohaotu@pony.ai>
It built two requests with one step each, so its `compared == 1` assertion read 2: `distribution_gap` stops the request whose pick moved, and every other request keeps its own context and goes on being compared — which is the behavior the fixture now pins. tp.md's near-tie row gains the margins both real flips read (0.125 / 0 and 0.375 / 0.5), so the 1.0 width has a reading behind it. Signed-off-by: bohaotu <bohaotu@pony.ai>
|
把平局判据的宽度那一半也搬到这边（先在
实测（双 L20，12B， 即 0.125 / 0 与 0.375 / 0.5，都在 1.0 内、余量 2–8×。这条线就是从这个读数定的，已写进 tp.md 的 near-tie 行（连同"两次翻转的 margin"一行）。 规则的两半有不需设备的单测： |
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Thanks for the quick turnaround. Three issues in the new code and two in the evidence, inline. A correction to my last round: rows failing at both one and two ranks do not bound what the second rank adds, so long prompts are still open. One question on the diagnosis: the stall was measured on the 4-bit geometry at TP2, which validate_for refuses (config.rs:228), and #1115 saw 128 tokens stall on the 31B. Could you share how that run was set up, and whether warming the lazily loaded GEMM shape alone removes the 5 token stall? If so, the change could be much smaller.
| tp_broken.set(true); | ||
| if select_device(ctx).is_ok() { | ||
| arena.release_graphs(); | ||
| } else { | ||
| // A device that will not come current is the one case where | ||
| // the release cannot run — and it is exactly the case a | ||
| // sticky fault produces, so name the order we are about to | ||
| // break: the abort below then runs against a live graph. | ||
| log::error!( | ||
| "rank 0's device {} will not come current, so its graphs stay alive \ | ||
| across the communicator abort; that abort may wedge", | ||
| ctx.device_ordinal | ||
| ); | ||
| } | ||
| serve.abort_tp_comm(); | ||
| for comm in extra_comms.iter() { | ||
| comm.abort(); | ||
| } |
There was a problem hiding this comment.
In NCCL 2.18.3 the last abort waits until no CUDA graph references any comm (commDestroySync, src/init.cc L1756). With decode graphs on, the peers' graphs are released only after the join, so the abort, the join and that release wait on each other; the peer path at L1473 has the same cycle with rank 0's graphs, and the new gates capture no graph.
| Ok(()) | ||
| })); | ||
| } | ||
| activate_rank(ctx)?; |
There was a problem hiding this comment.
If this fails, the scope joins peers that may be waiting on rank 0's collective, and the abort is never reached. Could it run before the spawn?
| fn assert_reads_something(&self, what: &str, long: usize) { | ||
| assert!( | ||
| self.compared > 0, | ||
| "{what}: no shared rows were scored at {} vs {long} tokens", | ||
| long / 2 | ||
| ); | ||
| assert!( | ||
| self.overlapped > 0, | ||
| "{what}: no row's top-{TOP_K} was shared between the {}- and {long}-token runs, so \ | ||
| the gap measured nothing", | ||
| long / 2 |
There was a problem hiding this comment.
One shared row is enough to pass: rows with disjoint top 8 add nothing to worst, so a widespread error can read 0, and flips are only printed. Could those count against the gate?
| fn a_flip_six_nats_wide_is_not_a_near_tie() { | ||
| let one = vec![drained_run(&[&[(100, 6.0), (200, 0.0)]])]; | ||
| let two = vec![drained_run(&[&[(200, 6.0), (100, 0.0)]])]; | ||
| distribution_gap(&one, &two, "six nats"); | ||
| } |
There was a problem hiding this comment.
The flip branch already folded both picks into worst_pick under LOGBROB_LINE, so [6, 0] against [0, 6] failed before; this test skips that check, so it cannot show otherwise. Could the explanation be corrected?
| // Only the nine-token case is asserted. Running the fixture's 1024-token | ||
| // `edge` anyway fails a strict top-64 containment on ~5-8 rows — *the same | ||
| // rows with one rank as with two*, which is what shows the second rank adds | ||
| // nothing and leaves this as the depth-dependent shape-dependence these | ||
| // fixtures carry per-case tolerances for. (`probed_cases` is the fixture's | ||
| // list of cases carrying hidden-state probes; `edge` has none because at | ||
| // window width they would dwarf the file, so its absence there says nothing | ||
| // about whether the case is asserted.) See docs/models/gemma4/tp.md | ||
| // "Known bounds". | ||
| const CASES: [&str; 1] = ["short"]; |
There was a problem hiding this comment.
Failures shared by one and two ranks do not bound what the second rank adds, and this still checks nine tokens. Could the notes mark long prompts at TP2 as unverified?
…them NCCL's abort waits for a captured graph that references the communicator being aborted, and a rank stopped inside a collective cannot release its own — its graphs live in an arena its own thread holds, and that release is exactly what the abort is waiting for. With a captured collective in play the failure path has no order that works: the abort waits for the release, the release runs after the join, and the join waits for the rank that abort was to release. The same cycle runs against rank 0's graphs when an extra rank fails first. Graphs are refused at a tensor-parallel launch rather than left to hang there, and capture is a measured wash under TP (146.0 vs 143.7 tok/s at c8), so nothing priced above zero is given up. Two smaller items from the same round. Rank 0's device is bound before any peer is spawned: a failure there would otherwise return from the scope while the peers were already running, and they would wait on a collective rank 0 never issued while the abort sat behind the join. And the prefix gate counts the rows whose top-8 share no token — they contribute nothing to `worst`, so a widespread error could pass it on one shared row elsewhere; a quarter of the rows is the bound. The width half moves here from pegainfer-project#1136: distribution_gap requires the two picks within NEAR_TIE_WIDTH in both runs' own readings, with the rationale corrected — the folded worst_pick against LOGBROB_LINE already catches the six-nat shape, so the width's own catch is the band between them (margins 0.9 and 1.05 with folds of 0.975), which a unit test now fixes. Signed-off-by: bohaotu <bohaotu@pony.ai>
…es at the same length Three runs of the 12B's scored gate at 128 tokens all hang, while its sampled arm passes at 128 and at 401 in the same session on the same tree: the path is what bites, and the hazard is repeatable at a point while absent at another length, so there is no length to bound. Replaces the one-run-per-point caveat with that. Signed-off-by: bohaotu <bohaotu@pony.ai>
Three more runs of the 12B's scored gate at 401 tokens all pass (4/4 with the first), against 3/3 hangs at 128 on the same arm, geometry and session. Both ends are reproducible and they disagree, so the hazard has no length to bound; which shape in the scored path blocks at one length and not the other is left open. Signed-off-by: bohaotu <bohaotu@pony.ai>
随着进一步 debug，PR 正文已过时，现在此进行更新：问题：
解决方案：
|
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Thanks for measuring the warm on main and for correcting the notes; the scored path result answers the question. The round 3 points are addressed. Two follow ups inline, and one suggestion: the 12B scored hang at 128 tokens now reproduces on main (three of three), so a host stack of that run (gdb -p <pid> -batch -ex 'thread apply all bt') would name the call rank 0 blocks in. If it is one call in the scored segment, moving it behind the peers' launch may fix the single threaded driver directly and keep decode graphs at TP, which is worth knowing before settling on the larger change.
| // NCCL's abort waits for the captured graphs that reference the | ||
| // communicator being aborted, and a rank stopped inside a collective | ||
| // cannot release its own — so with decode graphs on there is no order that | ||
| // releases one rank and destroys another's graphs: the abort waits for the | ||
| // release, the release runs after the join, and the join waits for the | ||
| // rank that abort was to release. Refused rather than left to hang there; | ||
| // the repo's own measurement is that graphs are a wash under TP (146.0 vs | ||
| // 143.7 tok/s at 8-way on two L20s), so nothing priced above zero is given | ||
| // up. | ||
| anyhow::ensure!( | ||
| !graph_enabled, | ||
| "CUDA graphs are unsupported under tensor parallelism: a captured collective cannot be \ | ||
| released while the rank that holds it is stopped inside one, and NCCL's abort waits \ | ||
| for exactly that release. Pass --cuda-graph=false." | ||
| ); |
There was a problem hiding this comment.
enable_cuda_graph defaults to true, so every TP2 launch now fails until the flag is flipped. Could TP fall back to eager with a log line instead?
| // A widespread error makes most rows share nothing, and those rows are | ||
| // invisible to `worst`: this is the bound that sees them. Read against | ||
| // the reading, which is 0 of 511 rows at 1024 tokens — a quarter is far | ||
| // above any run this gate has taken and far below a fault that moves | ||
| // every page. | ||
| assert!( | ||
| self.disjoint * 4 <= self.compared, | ||
| "{what}: {} of {} rows share no token between the {}- and {long}-token runs, more \ | ||
| than a quarter of them — a disagreement that widespread is not the row-count \ | ||
| dependence this line is calibrated for, and those rows never reach `worst`", | ||
| self.disjoint, |
There was a problem hiding this comment.
The reading is 0 of 511, and a quarter lets 127 disjoint rows through, so a fault confined to part of the prompt still passes. Could the bound follow the control arm's own count, as the gap does?
… against its control Two review items on pegainfer-project#1133. Decode graphs are now turned **off** at a tensor-parallel launch instead of refused. The flag is on by default, so a refusal meant no TP2 launch started until it was flipped; graphs are a measured wash under TP (146.0 vs 143.7 tok/s at 8-way), so serving eager is strictly better than not serving. The reason graphs cannot be used under TP is unchanged: NCCL's abort waits for a captured collective to be released, and a rank stopped inside one cannot release it. The TP capture sweep is now unreachable and is recorded as a known bound. Every TP gate runs with `enable_cuda_graph: true` (the default), so a reintroduced refusal fails them at launch rather than going unnoticed. The prefix gate's `disjoint` bound now follows the one-rank control's own count — a ratio with a floor, the shape the gap already uses — instead of a flat quarter of the compared rows, which admitted 127 of 511 and so could not see a fault confined to part of the prompt. Measured at the window width on the 12B: both arms read 0 of 511. Also records the call a stalled single-threaded TP2 driver blocks in: a stack of `main` at the 12B scored gate, 128 tokens (hung 3/3 there), puts the scheduler thread in `cublasGemmEx` for the gate|up projection GEMM under `attention_epilogue_into`, spinning inside libcuda, before it reaches `step_extra_ranks` — so rank 1's tower was never launched and rank 0's collective had no peer.
|
(a) Host 栈：卡住的调用已点名 在 pod 里用 gdb 把测试二进制当子进程跑（ptrace_scope=1 只允许 trace 后代，所以 gdb -p 被拒、gdb --args 可以），timeout -s INT 触发中断，12B scored 128-token 在 main 上挂住时的 thread apply all bt： 往上是 libcublas → cublasLtTSTMatmul → … → cuMemsetD8Async → libcuda（在 sched_yield 自旋，符号到此为止）。 三条结论：
(b) disjoint 界改成跟随 control
(c) TP2 图改成回退 eager + 日志
验证（都在 L20、NCCL_PROTO=LL128、CUDA_VISIBLE_DEVICES=1,4） 后四条都在 enable_cuda_graph: true 下跑，即回退路径被真正执行（若还是拒绝，启动就 panic）。 提交：f6134317（代码+文档）→ bad0580（去掉那条过强的措辞）。 |
… against its control Two review items on pegainfer-project#1133. Decode graphs are now turned **off** at a tensor-parallel launch instead of refused. The flag is on by default, so a refusal meant no TP2 launch started until it was flipped; graphs are a measured wash under TP (146.0 vs 143.7 tok/s at 8-way), so serving eager is strictly better than not serving. The reason graphs cannot be used under TP is unchanged: NCCL's abort waits for a captured collective to be released, and a rank stopped inside one cannot release it. The TP capture sweep is now unreachable and is recorded as a known bound. Every TP gate runs with `enable_cuda_graph: true` (the default), so a reintroduced refusal fails them at launch rather than going unnoticed. The prefix gate's `disjoint` bound now follows the one-rank control's own count — a ratio with a floor, the shape the gap already uses — instead of a flat quarter of the compared rows, which admitted 127 of 511 and so could not see a fault confined to part of the prompt. Measured at the window width on the 12B: both arms read 0 of 511. Also records the call a stalled single-threaded TP2 driver blocks in: a stack of `main` at the 12B scored gate, 128 tokens (hung 3/3 there), puts the scheduler thread in `cublasGemmEx` for the gate|up projection GEMM under `attention_epilogue_into`, spinning inside libcuda, before it reaches `step_extra_ranks` — so rank 1's tower was never launched and rank 0's collective had no peer. Signed-off-by: bohaotu <bohaotu@pony.ai>
The mixed step is also a TP path with no gate behind it and no refusal, so the sweep is not the only one. Signed-off-by: bohaotu <bohaotu@pony.ai>
bad0580 to
b6eeb34
Compare
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Thanks, both fixes read right and the stack settles the reordering question. Three points inline; the first is the one before approving.
| // Unreachable while `start_with_knobs` turns graphs off for every | ||
| // `world > 1` launch: the sweep is kept for whoever lifts that, and its | ||
| // having no runtime is a known bound (`docs/models/gemma4/tp.md`). | ||
| if graph_enabled && world > 1 { |
There was a problem hiding this comment.
Graphs are off for every TP launch, so this sweep, reset_min_bucket, and the release_graphs calls in drive_ranks and break_extras are now dead code. Could they go, as PEGAINFER_TP_MAX_PROMPT did? That also drops the comments that still say TP refuses graphs.
| - **The failure path's order is only sound because graphs are off.** The failure path aborts a peer's communicator *while that peer is inside a collective*, and it has to abort it from another thread because the peer cannot release itself. That is free of the reverse dependency only because a TP2 launch turns captured graphs off (`start_with_knobs`; see "Decode graphs are off under TP"): with a captured collective alive, NCCL's abort waits for a graph that only the join can release, and the abort, the join and the release wait on each other. The graph releases the failure path still makes are vacuous while graphs are off — they are there so the order is right if that is ever lifted. What is measured without an engine is the release itself: `engine::lane_gates_tp::an_aborted_communicator_releases_a_waiting_reduction` enqueues a reduction with no peer call, aborts that rank from another thread, and requires its stream to stop waiting — and the abort itself to return — within a deadline. The composition of the two halves (this order, on an engine that captures nothing) is argued rather than exercised: an asymmetric engine-level failure would need a hook the single-rank half also went out of its way not to add. | ||
| - **The TP capture sweep has no runtime.** With `start_with_knobs` turning graphs off for every `world > 1` launch, `start_with_knobs`' phase-interleaved `precapture_bucket` loop is unreachable: no gate exercises it, and the parity reading it was built for is gone (see "Decode graphs are off under TP"). It is kept for whoever lifts the disable; unlike the chunked walk above it is not refused, so nothing anywhere stops a run from reaching it if the disable is dropped. | ||
| - **The `two_rank` gates run with the graph flag at its default.** `tp2_options` passes `enable_cuda_graph: true`, so every TP gate exercises the turn-off above: a reintroduced *refusal* would fail every one of them at launch instead of going unnoticed. | ||
| - **The stall the prompt ceiling guarded was not a length bug — the ceiling is gone with it.** Rank 0's segment reaches a host-side wait inside the tower. An **earlier** capture (a **5-token** prompt) put rank 0's thread in a **lazy cuBLAS/cuBLASLt library load** (`cuLibraryGetModule`, reached through `cublasGemmEx`); the **named** capture — the 12B scored gate at 128 tokens, which hangs 3/3 on `main` — puts it in the **gate|up projection GEMM** (`cublasGemmEx`, reached from `attention_epilogue_into` under `launch_prompt_tower` on the scheduler thread, `gemma4-engine`), spinning inside libcuda with `cuMemsetD8Async` the innermost named frame and nothing below it. In both, rank 0's device reads 100 % busy with rank 1 idle, and the scheduler thread has not reached `step_extra_ranks` — so rank 1's tower was never launched and rank 0's enqueued collectives have no peer, which is the hang itself. Driven **concurrently** (`EngineState::drive_ranks`: rank 0 on the scheduler thread, one scoped thread per extra rank, then join), a host blocking point in one rank's segment cannot stop another rank's launches, so the class is removed rather than bounded and `PEGAINFER_TP_MAX_PROMPT` — with its refusal and its start-up warning — is deleted rather than left as dead scaffolding. Measured on the single-threaded driver with the ceiling lifted: a scored prefill of 50/64/100/128/300 tokens and a 401-token server prefill all complete — on a 4-bit 60-layer geometry, which `validate_for` **refuses at TP2** (`config.rs:228`: TP does not shard the W4A16 GEMMs), so that run predates or bypassed the refusal and **is not reproducible from this tree**; the bf16 31B is not re-measured on the single-threaded driver, and #1115's own record has 128 tokens stalling there. The ceiling was therefore refusing prompts that serve on the geometry it was measured on. Whether the **warm alone** would have been enough for them — which is what decides whether the concurrent driver is needed at all — was then measured on the shipped geometry, on `main`'s single-threaded driver with the warm in place and the ceiling off: a **sampled** prefill of 11 tokens and of **401** tokens both complete on the bf16 12B (35.2 s / 29.7 s), and on the bf16 **31B** the **scored** prompt passes at 64 tokens (472 s) and **hangs at both 128 and 401**, each killed at its 1500 s timeout — with rank 0's device at 100 % and rank 1 idle at the moment the 401-token run was sampled, the same host-wait signature as the capture above. So the warm is not what covers this class, and it is reachable on the geometry this repo serves. A cross-check separates the geometry from the path: the same **scored** gate at the same 128 tokens **also hangs on the 12B** (again rank 0's device at 100 %, rank 1 idle), where the 12B's **sampled** arm passes at 401 tokens — so what differs is the *admission sequence*, not the arm's name: the named block lands in the prompt tower's MLP projection, which **both** arms execute, and which cuBLASLt config that call first touches is whatever the sequence of shapes ahead of it has left alone. It is **not a length bound**: the 31B's scored arm passes at 64 and hangs at 128 and 401, and the 12B's scored arm hangs at 128 in **all three** runs and passes at 401 in **all four**, with the 12B's *sampled* arm passing at both 128 and 401 in the same session. Same arm, same geometry, same session: the outcome is reproducible at each of those lengths and opposite between them, so there is no length to bound — the ceiling was a length bound on a hazard that has none. The **call** rank 0 blocks in is now named (above); why that call spins inside libcuda at one length and not another — the driver-internal reason, which the capture's symbol-less innermost frames do not give — is the open half. The TP gates pass on this build (one-rank comparison `48/48` at gap `0.0000`, the scored probe, the HF reference). The cost is one thread spawn per rank per step, plus the cuBLAS handle pair each step's thread creates in `activate_rank` and destroys when it ends — **~0.8 ms per step** in a standalone copy of that pattern on sm_89. That pair is also what stops the 16.25 MiB per step the first cut leaked (0 MiB growth over a 300- and a 500-step decode run); a thread per rank living for the engine's whole life would keep one pair instead, at the price of that thread. |
There was a problem hiding this comment.
The open half is documented: under lazy loading, loading a kernel can wait for a context synchronization (CUDA Programming Guide, Lazy Loading, "Impact on Concurrent Kernel Execution"), so it waits on rank 0's spinning collective while rank 1 is never launched. If the 12B scored case at 128 tokens passes on main with CUDA_MODULE_LOADING=EAGER, that confirms it, and the proxy thread explanation can cite the guide instead.
| let two = prefix_comparison(&tp2_options(device, peer), long); | ||
| two.report("prefix tp2", long); | ||
| two.assert_reads_something("prefix tp2", long); | ||
| two.assert_disjoint_absolute("prefix tp2", long); |
There was a problem hiding this comment.
Small one: the 31B still uses the quarter, which passes up to 127 of 511 rows. If its run also reads 0, compared / 64 could apply here too.
…all's cause The two-rank capture sweep, `RankState`'s `Drop`, the extra-rank loop in `EngineState::drop`, `drive_ranks`' two `release_graphs` calls, `break_extras`' release, `StepArena::reset_min_bucket` and `StepArena::bucket_ceiling` all go: a tensor-parallel launch turns graphs off, so none of them could run, and every comment that justified itself by "TP refuses graphs" goes with them. Net 101 lines out of `engine.rs`, 22 out of `serve.rs`. `tp.md` now records what the stall is. EAGER does not cure the 12B scored gate at 128 tokens — nor does `CUBLAS_WORKSPACE_CONFIG=:16:8`, nor `CUDA_LAUNCH_BLOCKING=1` — so it is not the lazy module load the previous revision blamed. Two independent gdb captures are frame-for-frame identical and put the scheduler thread in `cublasLtTSTMatmul` → `cuMemsetD8Async`, spinning inside libcuda on the rank's own stream, with every other thread of the process parked: the wait is device-side, and the device is stuck in an all-reduce whose peer rank 1 never issued because the one host thread is stopped here. The cold-start section keeps the warm's explanation and cites the guide's lazy loading note for it, records that the innermost frame is not fixed, and points at the known-bounds bullet for the instance the warm does not answer. Signed-off-by: bohaotu <bohaotu@pony.ai>
Its own count is unmeasured — the 0-of-511 reading is the 12B control's — and the 31B needs two cards that each hold 34 GiB, so the tighter floor waits for a reading from the checkpoint it would be applied to. Signed-off-by: bohaotu <bohaotu@pony.ai>
Measured while re-running the TP gates on the recorded envelope: at PEGAINFER_TP_PROMPT_TOKENS=1024 the one-rank-vs-two-rank gate trips on a flip whose picks sit 0.5 and 1.25 apart in the two readings, with the folded picked-token bound at 0.9307 against its 1.0. So NEAR_TIE_WIDTH carries the same grain as the line beside it, and only the gate's own default length is asserted. Signed-off-by: bohaotu <bohaotu@pony.ai>
补跑时另外得到一条数据，已进 Known bounds（ 验证（记录 envelope，2 * L20）：4 条单测、parity（36.2 s）、scored（16.0 s）、prefix control（27.7 s）、prefix standalone（15.2 s）、rank-0 failure（15.1 s）、两条 abort（1.4 / 3.0 s）全过。 |
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Thanks, the TP graph code is gone and a default launch now serves eager. Three points remain inline, and the first two are the ones before approving.
| .get(row + 1) | ||
| .and_then(|entry| entry.as_ref()), | ||
| ) else { | ||
| continue; |
There was a problem hiding this comment.
A row missing from either echo is skipped here, so a run that scores one of the 511 rows still passes. Could a missing row fail, as it does in the HF gate?
| - **The chunked walk's TP code has no runtime.** `PEGAINFER_MIX_CHUNK_TOKENS` is refused under TP (a chunked walk's rounds gather across prompts, which the TP gates do not cover), so the walker's extra-rank admission and its `tp_broken` check are unreachable with a non-empty `more`. Its segments drive the ranks through the same `drive_ranks` as every other path, so nothing but the refusal stands between that code and the gates — untested until one covers it. | ||
| - **The mixed step has no TP gate, so the row bound the ceiling carried is back.** `serve::oracle::mixed_step_matches_serial` is single-rank, and the deleted ceiling was also what kept a TP mixed step's prompt rows small (`mixed_fits`, and the gather's `gather_budget`), so both are restored for this step alone: a prompt whose rows do not fit beside the live batch — `TP_MIX_PROMPT_ROWS` (64), counting the live decode rows — is prefilled on its own instead of riding the step, exactly as an over-ceiling prompt was before, and the gather leaves the same room (`PEGAINFER_MIX_GATHER_ROWS` still bounds the gathered rows, and the step rides the same shared segment at every rank). At one rank the step is unbounded, as it always was. That region is still untested at two ranks rather than known-good; the bound stands in for the gate, and both the bound and the knob go when a TP gate covers the step. | ||
| - **The vocabulary projection is replicated.** `embed_tokens` and the tied head are `Whole` on each rank (31B: 262144 × 5376 bf16, ~2.8 GiB per rank), and every extra rank computes a whole batch of `vocab × rows` logits per decode step that is then discarded. Dropping it needs a head-free tower for the extra ranks — in graph mode that means a second per-bucket capture, because the head sits inside `decode_gpu_body`. Not implemented, not measured. | ||
| - **The long-prompt comparison needs the fixtures' tolerance discipline, not a tighter line.** At one rank the engine matches HF at every *probed* long-context anchor, including the window width: `serve::oracle::context_waypoints_match_hf` passes, reading 1.15 at w1023, **1.48 at w1024** (tolerance 4.61), 3.15 at w1025, 3.80 at w4096, 1.17 at w16384 and 5.23 at w32768 (tolerance 7.65). **At two ranks that depth is unverified**: the 12B golden's `edge` (1024-token) case is not asserted by the TP gates, and running it anyway fails a strict top-64 containment on ~5-8 rows *at one rank as well as at two* — evidence about neither, since a failure both arms share is explained by neither, so it cannot be read as "the second rank adds nothing". What the TP arms do pin at depth is the prefix comparison (a 1024-token prompt's rows against the same prompt's first half): 6.30 at one rank and 9.23 at two, bound by its own control arm (ratio 1.46) and by a count of the rows that share no token at all — those rows never reach `worst`, so the two-rank arm's count is held against the control's own, a ratio with a floor (both arms read **0** of 511 at 1024 tokens; the floor is a sixty-fourth of the compared rows, replacing a flat quarter that let 127 through). The one-rank-vs-two-rank gate's **1.0 line is a 64-token calibration** and cannot hold at depth either: its drift reads 0.57 at 64 tokens, then 1.87 at 512, 4.52 at 768, 5.50 at 256 and 3.81 at 1024 — non-monotone, i.e. the tiling-dependence rather than a threshold. The **near-tie width has the same grain and the same limit**: run at 1024 tokens that gate trips on a flip whose picks sit 0.5 and 1.25 apart in the two readings, with the folded picked-token bound at 0.9307 (under its 1.0) — so `NEAR_TIE_WIDTH`'s 1.0 is a shallow-length calibration as well, and only the gate's own default length is asserted. Giving the long cases the tolerance discipline is the follow-up; a strict line at depth is not a detector, it is a flake. |
There was a problem hiding this comment.
Long prompts at TP2 still have no correctness check. The HF gate covers only the nine token case, and at 1024 tokens the parity gate fails NEAR_TIE_WIDTH with nothing showing the flip is drift. Could the fixture's 1024 token edge case run at one rank and at two, with the two rank distance from HF held within a ratio of the one rank distance? The 31B quarter bound also needs one reading from the 31B.
| `PEGAINFER_ASYNC_PREFILL` (its lane stream cannot be lock-stepped across ranks), `PEGAINFER_GLOBAL_ATTN=tilelang*` (the generated kernels are compiled for the whole global family), `PEGAINFER_PREFIX_CACHE`, and `PEGAINFER_MIX_CHUNK_TOKENS` (a chunked walk's rounds gather across prompts, which the TP gates do not cover). The prompt ceiling that used to be refused here (`PEGAINFER_TP_MAX_PROMPT`) is gone with the hazard it guarded — see "Known bounds". | ||
|
|
||
| Two environment constraints are baked into the code. The **kernels must be warm before the communicator exists** (see "Cold start"): a first-shape GEMM that materializes after `Comm::from_devices` wedges in the driver's module loader, which is what made the earlier capture attempts look like a capture defect. And **every rank must release its graphs before the communicators drop**: a captured collective bakes in NCCL kernel launches, and `ncclCommAbort` wedges while a graph that references them is still alive, so `EngineState::drop` releases them first, on each rank's own device. | ||
| ## Decode graphs are off under TP |
There was a problem hiding this comment.
The reason graphs are off is repeated in eight places. This section and one line at start_with_knobs would be enough. step_extra_ranks in the stall bound and prefill_extra_ranks in the serve_batch doc name functions this PR removes.
…ctions The failure path's reason for turning decode graphs off under TP was derived in the doc, the log site, two Drop docs, a gate comment and the index row. The derivation lives in "Decode graphs are off under TP" now and everywhere else points at it. The stall bound also named `step_extra_ranks`, which this branch removed — it says "the scheduler thread has not launched rank 1's tower", and the solo admission path no longer names `prefill_extra_ranks` either. Signed-off-by: bohaotu <bohaotu@pony.ai>
Three review points on the TP gates: - The prefix comparison skipped a row absent from either echo, so a run that scored only part of the prompt passed on the rows it did score. It fails now, as the HF gate's own `was not scored` already does. - The HF reference gate asserted only the nine-token case, leaving long prompts unverified. The golden's `edge` (1024 tokens) now runs at one rank and at two, and the two-rank reading is held within 1.5x the one-rank one (floor of two rows, 0.25 logits): a strict top-64 containment fails at both arms (37 rows at one rank, 24 at two; worst gap 6.58 / 8.49), which bounds neither, but a two-rank arm that adds rows the one-rank arm did not is the second-rank fault this gate exists for. - The no-control 31B arm kept a flat quarter on the rows that share no token, which admitted 127 of 511. Its own run on two L20s reads 0 of 511 (worst shared-token gap 4.0538), so it takes the same `compared / 64` the control-relative form uses. Measured on two L20s: `the_two_rank_engine_matches_the_hf_reference` 8/8 rows at gap 0.6548; the new depth gate passes (21.8 s); the controlled prefix gate passes with both arms at 0 of 511. Signed-off-by: bohaotu <bohaotu@pony.ai>
…ership Its doc said the shape was "the membership half's failure", but membership holds there — each pick is the other run's runner-up — and what fires is the margin assert. The line catches the shape too (both picks' gaps fold into `worst_pick`, which reads 6 against `LOGBROB_LINE`); the test below is the one only the width rejects. Signed-off-by: bohaotu <bohaotu@pony.ai>
|
All three points are in on top of the head you reviewed (
On two L20s on that head: units 4/4, |
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
Thanks, a missing row now fails and the depth gate is in. Four points remain inline. The first changes what the new gate accepts, so it is the one before approving.
| } | ||
| } | ||
| if !ours_top.contains_key(&pb) { | ||
| distance.outside += 1; |
There was a problem hiding this comment.
A row that fails both containment checks adds two to outside, and the depth gate reads that count as rows under a 1.5x ceiling. With four such rows at one rank the count is 8, so ten rows at two ranks that each fail one check read 10 and pass the ceiling of 12. Could each row count once? The 37 and 24 you reported still pass after that, though they may not be the row counts tp.md calls them.
| /// The floor, in rows: a one-rank arm that read zero out-of-top-k rows must | ||
| /// not pin the two-rank ceiling to zero. | ||
| const OUTSIDE_FLOOR: usize = 2; | ||
| /// The floor under the gap ratio, in raw logits: two arm's gaps are both |
There was a problem hiding this comment.
Both sides of hf_distance are logprobs, the fixture's *_topk_logprobs and the engine's TokenLogprob, so this floor is in nats. The doc above also says the gate names a two-rank arm that adds rows the one-rank arm did not. It compares two totals and keeps no row positions, so it can only bound how far the total grows. Could both comments say that?
| one.compared > 0, | ||
| "{case}: the one-rank arm scored no rows, so there is no distance to hold against" | ||
| ); | ||
| assert_eq!( |
There was a problem hiding this comment.
Both arms read the same case of the same fixture, and hf_distance panics on a missing row, so the two compared counts are always equal and this check cannot fail. The four line comment above the was not scored panic in prefix_comparison repeats the panic, and the calibration history beside PREFIX_LINE and the disjoint constants is already in tp.md. Could these go?
| // engine's scheduler thread, which owns the whole `EngineState` and touches a | ||
| // rank's serve only while that rank's device is current. Nothing shares a serve | ||
| // between threads, so the raw NCCL communicator it may hold never crosses a | ||
| // time: the thread that loads it (the load-time warm and the decode-graph sweep |
There was a problem hiding this comment.
This still names the decode graph sweep, which the PR removes. The communicator comment in load also still says the all-reduce lands inside the decode graph, while TP now serves eager. Could both drop the graph part?
Round 7 on pegainfer-project#1133: - `hf_distance` counted a containment failure, not a row, so a row failing both checks weighed two and the 1.5x ceiling could be bought: four such rows at one rank read 8, buying 12 for ten single-failure rows at two. A row now counts once, whichever checks it misses. Re-measured on two L20s: one rank **31** rows, two ranks **22** (the 37 / 24 the gate reported were failures), worst gap 6.58 / 8.49. - The gate's doc and `tp.md` claimed it names the rows the second rank added. It compares two totals and keeps no row positions, so both say it bounds the total now. `WORST_FLOOR` is in logprobs (nats), not raw logits. - Dropped the vacuous `assert_eq!` between the arms' `compared` counts (both read the same case and `hf_distance` panics on a missing row), the comment restating that panic, and the calibration prose the constants duplicated from `tp.md`. - `serve.rs`: `attach_tp_comm` no longer says the reduction lands inside a decode graph, and the `unsafe impl Send` SAFETY no longer names the decode-graph sweep. Signed-off-by: bohaotu <bohaotu@pony.ai>
|
All four are in (
Also on that head: units 3/3, |
FeathBow
left a comment
FeathBow
left a comment
There was a problem hiding this comment.
The remaining issues are just minor style/comment fixes and won't block this. Thanks so much for your quick turnaround and edits!
2 participants
Problem
#1115 landed gemma4 TP2 but refused any prompt past the length its driver had been verified at (
PEGAINFER_TP_MAX_PROMPT, default 64). The guard was on a hang, not a length: one host thread drove rank 0's whole tower before the peers', so when rank 0's tower reached a host-side wait — captured at TP2 as a lazy cuBLAS/cuBLASLt library load (cuLibraryGetModule, reached throughcublasGemmEx) inside a projection GEMM — it waited on a collective whose peer call it had not yet issued. That trigger is a first touch of a shape, not a prompt length (the captured prompt was 5 tokens), so the ceiling refused prompts that serve.Fix
One thread per rank.
EngineState::drive_ranksruns rank 0's segment on the scheduler thread and each extra rank's on its own scoped thread, then joins — so by the time rank 0 reaches a host-blocking call, every peer's collectives are already in flight. Every rank-driving path goes through it — the solo, chunked and scored prefills, the mixed step, and both decode paths — and the payload is transposed once up front (rank_payloads/prefill_payloads,GemmaKv::split) because the threads need disjoint borrows. A rank-0 failure releases rank 0's graphs and detaches its communicator before the join, so a peer already waiting on a call that will never arrive is unblocked; the engine is marked broken only when a peer exists, so world size 1 keeps the request-scoped failure its gate pins. The ceiling — env var, parser, knob, unit assertions, and its start-up warning — is deleted rather than left as dead scaffolding. This is a port of the branch written against #1115's earlier rounds onto the squash-mergedmain: the superseded helpers go, while everything #1115's later rounds established stays.Verification
2×L20, sm_89, CUDA 12.9, NCCL 2.18.3,
PEGAINFER_CUDA_SM=89,NCCL_PROTO=LL128.22/24picks kept, worst gap0.5738(line1.0)48/48, gap0.00007/8rows,0out-of-top-k, worst gap0.6753511shared rows, worst gap4.0538(line12.0), 14 argmax flipscargo fmt --check/clippy --all-targets -D warnings/ lib suiteBounds
The 1024-token HF reference comparison is not gated yet (its drift is the chaotic band
serve_oracle'sneutral_scalemeasures, so it needs tolerance-plus-top1_share); a rank-0 failure at TP2 is still argued from the teardown order rather than exercised, since no gate injects one; the cost is one thread spawn per rank per step. The stall's mechanism is measured on the 4-bit 60-layer geometry — the bf16 31B's own boundary is not re-measured — the TP notes' description of the hazard is corrected on the branch, and the review follow-ups (a fail-closed payload check, the prefix gate's finiteness and overlap asserts) are in the commit after it.