Skip to content

[Bugfix][Pooling] Register Sentence-Transformers projector weights - #96

Open
aoshen02 wants to merge 1 commit into
wu/split-01-frameworkfrom
wu/split-02-st-projector
Open

aoshen02 wants to merge 1 commit into
wu/split-01-frameworkfrom
wu/split-02-st-projector

Conversation

@aoshen02

@aoshen02 aoshen02 commented Oct 5, 2026

Copy link
Copy Markdown
Owner

Overview

Sleep mode level 2 dropped the Sentence-Transformers projector, so embeddings changed after wake. Stacked on #95 (the registration helpers and the CI allowlist); review only the top commit.

Claims

  • Fixes embeddings after sleep(2) / wake_up() / reload_weights() for models with an ST projector (e.g. embeddinggemma-300m).
  • Removes \.head\.projector\. from KNOWN_UNREGISTERED.

Validation

GB200, image nightly-dev-arm64-cu130-ac9126e58aa7 overlaid with the framework PR plus this change.

Check framework only + this PR
embeddinggemma-300m, embed, sleep(2) -> wake -> reload vs before max abs diff 0.154 6.0e-8 (fp32 noise; repeated runs on the combined change: <= 1.5e-8)
test_initialization on the 29 architectures in the allowlist pass 29/29 pass

Details

The projector is built from the ST Dense folders and held by DispatchPooler sub-poolers, outside the module tree, so its weights were neither parameters nor buffers. DispatchPooler now holds the projectors in a ModuleList (sub-poolers can hold modules owned elsewhere, so only the projectors join), and the projector weights become constant buffers so strict checkpoint loading does not expect them.

AI assistance. Written with Claude Code; audited by Claude Fable 5.1, Codex (gpt-6-astra) and Kimi K3 as part of the combined change. The submitter reviewed every changed line and ran the checks above.


Pull Request Checklist
  • I used vLLM's /pr-checklist skill. (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.

🤖 Generated with Claude Code

The ST projector is loaded from the Dense folders and held by DispatchPooler
sub-poolers outside the module tree, so sleep mode level 2 dropped it:
embeddinggemma-300m embeddings differed by up to 0.153 after wake. Register
the projector under DispatchPooler and its weights as constant buffers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant