Repository navigation
[Bugfix] Fix Granite 4.1 Vision deepstack handoff between pipeline stages - #65818
Conversation
…ages With pipeline parallelism, Granite 4.1 Vision failed at startup on every non-first stage with KeyError: 'ds_<layer>', and once that was avoided, later stages received the forwarded deepstack features as zeros. - Declare only the deepstack buffers a stage receives: the previous stage forwards the features of layers from this stage's first layer on. - Forward copies of the features: on the first stage they are views of buffers the outer model zeroes right after the forward, before the PP send. Add the model to the multimodal pipeline-parallel test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yifei Wang <yifei.w@inferact.ai>
…equest Replace the text-only entry in the multimodal pipeline-parallel test with a dedicated test that compares an image request between PP2 and no PP. With four text layers, deepstack layers 0 and 3 land on different stages, so it fails if a later stage receives missing or zeroed features. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yifei Wang <yifei.w@inferact.ai>
|
We hit the same startup failure independently on 2026-11-03 while working on #65412: and the engine then hung on Happy to rerun that setup with this branch on our H100s if a second data point is useful. |
|
Hi @artem-spector, thanks for confirming! This PR should fix the |
|
Reran this on our setup (2× H100 80GB,
PP1 output is also identical between the base and this PR (6/6), so single-GPU behavior is unchanged. Thanks again for the fix. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
/ci run |
|
✅ @xsxszab, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #101280 for commit ℹ️ This PR is 337 commits behind upstream |
Overview
Granite 4.1 Vision (#40282) does not work with
--pipeline-parallel-size > 1. Every non-first stage fails at startup withKeyError: 'ds_<layer>', and once that is fixed, later stages receive the forwarded deepstack features as zeros, which silently drops them from image requests. This PR fixes both: each stage declares only the deepstack buffers it receives, and the first stage forwards copies of the features instead of views of buffers it zeroes before the PP send.Claims
Validation
ibm-granite/granite-vision-4.1-4b, bf16, 1–2× B300,--max-model-len 8192. Six requests sent one at a time: three image chats and three text completions, 32 greedy tokens each with top-5 logprobs.KeyError: 'ds_9'on stage 1 during warmupTwo eager PP1 runs are identical (6/6, difference 0). With CUDA graphs, PP2 starts and serves; its outputs differ from PP1 at near-ties by the same amount as PP1 with CUDA graphs differs from PP1 in eager mode, because the two configurations capture different graph sizes.
Test.
test_pp_granite4_vision_deepstackserves a 4-layer dummy-weight variant with and without PP2 and compares a text and an image request (generate_with_image). With four layers, deepstack layers 0 and 3 land on different stages. It fails with only the declaration fix and passes with this PR:pytest -v -s "tests/distributed/test_pipeline_parallel.py::test_pp_granite4_vision_deepstack"On main the same test cannot start the PP2 server (the
KeyErrorabove). It runs in the "(L4) Distributed" CI step, which runs the whole file. It uses dummy weights because the registry still marks this checkpoint as not available online; the real-weight runs above were done outside CI.Details
Startup failure. The deepstack and spatial features target layers 0, 3, 6, 9, 12, 15, 18 and 21. On a non-last stage,
Granite4VisionLLMModel.forwardforwards onlyds_<l>withl >= end_layer, the features later stages inject.Granite4VisionLLMForCausalLM.make_empty_intermediate_tensorsdeclaredds_<l>for every layer regardless of the stage, and Model Runner V2 copies every declared key from the received tensors (#42187), so the keys of earlier layers were missing. The fix declaresds_<l>only forl >= start_layer, which is exactly what the previous stage sends. The runner's strict copy is left as is: a declared key that never arrives should keep failing loudly rather than reuse a stale buffer.Zeroed features. On the first stage, the forwarded
ds_<l>are views of_ds_buffers.Granite4VisionForConditionalGeneration.forwardzeroes those buffers right after the language model returns, so that stale features do not leak into the next request, and the worker sends the intermediate tensors only after the forward returns. Later stages therefore received zeros: with a 20/20 split, layer 21 lost its features, and with PP4, layers 12–21 lost theirs. Text requests are unaffected because their buffers are zero anyway. The fix forwards copies, at the cost of one(num_tokens, hidden_size)copy per forwarded layer.No open PR changes this code path; #65412, the other open Granite 4.1 Vision PR, touches only the image preprocessing.
AI assistance (Claude Code) was used to find the bugs, write the change and the test, and draft this description. The submitter reviewed every changed line and ran the validation above.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.