Extend ColBERT support to non-standard BERT backbones#34170
Extend ColBERT support to non-standard BERT backbones#34170noooop merged 7 commits intovllm-project:mainfrom
Conversation
Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com>
|
Documentation preview: https://vllm--34170.org.readthedocs.build/en/34170/ |
There was a problem hiding this comment.
Code Review
This pull request is a significant improvement, extending ColBERT support to multiple backbones like ModernBERT and Jina XLM-RoBERTa. The introduction of ColBERTMixin is a great design choice for code reuse and makes adding new backbones much cleaner. The tests have also been nicely refactored to be parametrized across all supported backbones, which improves test coverage and maintainability. I've found one issue related to loading weights for ColBERT models that have a bias in their projection layer, which could affect models like colbert-ir/colbertv2.0.
|
Hi guys @noooop @DarkLight1337 a follow up PR from me to extend ColBERT support beyond standard BERT backbones. Check this out and let me know what you think. Thanks in advance! |
Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com>
66f8e79 to
8dc9ce8
Compare
Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com>
Head branch was pushed to by a user without write access
…4170) Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com> Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
…4170) Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com> Signed-off-by: Eldar Kurtic <research@neuralmagic.com>
…4170) Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
…4170) Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com>
…4170) Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com>
…4170) Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com> Signed-off-by: EricccYang <yangyang4991@gmail.com>
…4170) Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com>
Purpose
Follow-up to the initial ColBERT support (#33686). Extends vLLM's ColBERT late-interaction support beyond standard BERT to multiple encoder backbones, enabling models like
lightonai/GTE-ModernColBERT-v1(ModernBERT) andjinaai/jina-colbert-v2(Jina XLM-RoBERTa).What changed:
ColBERTMixin— new mixin that extracts all shared ColBERT logic (projection layer init, weight loading, pooler building,colbert_dimauto-detection from config or weight shape). Adding a new backbone is now ~40 lines of code.ColBERTModernBertModel— ColBERT with ModernBERT backboneColBERTJinaRobertaModel— ColBERT with Jina XLM-RoBERTa backbone (usesJinaRobertaModel)Refactored the existing
HF_ColBERTclass to useColBERTMixininternally, keeping backward compatibilityMinor config plumbing so Jina XLM-RoBERTa and pylate-trained models (e.g.
lightonai/GTE-ModernColBERT-v1) are recognized correctlyTests are now parametrized across all three backbones (BERT, ModernBERT, Jina) with per-backbone HF comparison tests
Updated docs with architecture table and usage examples for each backbone
HF_ColBERTanswerdotai/answerai-colbert-small-v1ColBERTModernBertModellightonai/GTE-ModernColBERT-v1ColBERTJinaRobertaModeljinaai/jina-colbert-v2Test Plan