[Model]Add DebertaV2ForSequenceClassification Model Support #20215
Conversation
Signed-off-by: yashaswipiplani <[email protected]>
yashaswipiplani
requested review from
hmellor,
DarkLight1337 and
ywang96
as code owners
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @yashaswipiplani, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request significantly expands vLLM's model compatibility by introducing full support for DeBERTa v2 models, specifically for sequence classification. This integration allows users to leverage vLLM's optimized inference for a broader range of cross-encoder applications, enhancing the framework's utility for various NLP tasks.
Highlights
- New Model Support: I've added comprehensive support for
DebertaV2ForSequenceClassificationmodels, enabling vLLM to run and serve these models for sequence classification tasks. This includes implementing DeBERTa v2's unique architecture components. - Core Architecture Implementation: The core DeBERTa v2 architecture, featuring disentangled self-attention, relative position embeddings, and a convolutional layer for enhanced position encoding, has been implemented from scratch within a new
deberta_v2.pymodule. - Model Integration and Testing: The new DeBERTa v2 model has been fully integrated into vLLM's model registry and existing test suites, ensuring proper functionality and compatibility. Documentation has also been updated to reflect this new capability.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces support for DebertaV2ForSequenceClassification models. I've identified a performance bottleneck in the DebertaV2DisentangledSelfAttention implementation due to nested Python loops, which should be vectorized. There's also a bug in DebertaV2ForSequenceClassification.forward that will cause a runtime error. I've also suggested sorting dictionary entries in registry files for better readability.
| hidden_states = self.deberta( | ||
| input_ids=input_ids, | ||
| positions=positions, | ||
| inputs_embeds=inputs_embeds, | ||
| intermediate_tensors=intermediate_tensors, | ||
| token_type_ids=token_type_ids, | ||
| ) | ||
|
|
||
| # Apply classification head | ||
| logits = self.classifier(hidden_states[:, 0]) # Use [CLS] token | ||
| return logits |
There was a problem hiding this comment.
The self.deberta() call returns a BaseModelOutput object from the encoder, not a raw tensor. Accessing hidden_states[:, 0] will raise an attribute error at runtime. Access the last_hidden_state attribute of the output object before indexing.
| hidden_states = self.deberta( | |
| input_ids=input_ids, | |
| positions=positions, | |
| inputs_embeds=inputs_embeds, | |
| intermediate_tensors=intermediate_tensors, | |
| token_type_ids=token_type_ids, | |
| ) | |
| # Apply classification head | |
| logits = self.classifier(hidden_states[:, 0]) # Use [CLS] token | |
| return logits | |
| deberta_output = self.deberta( | |
| input_ids=input_ids, | |
| positions=positions, | |
| inputs_embeds=inputs_embeds, | |
| intermediate_tensors=intermediate_tensors, | |
| token_type_ids=token_type_ids, | |
| ) | |
| # Apply classification head | |
| logits = self.classifier(deberta_output.last_hidden_state[:, 0]) # Use [CLS] token |
| content_to_content = torch.matmul(query_layer, | ||
| key_layer.transpose(-2, -1)) | ||
|
|
||
| # 2. Content-to-position attention: query content to relative positions | ||
| content_to_position = torch.zeros_like(content_to_content, | ||
| device=query_layer.device) | ||
|
|
||
| # For each query position i, compute attention to all relative positions | ||
| for i in range(seq_len): | ||
| # Q_i^c: [batch_size, num_heads, head_size] | ||
| q_i = query_layer[:, :, i, :] # [batch_size, num_heads, head_size] | ||
|
|
||
| # Compute attention to all relative positions | ||
| # rel_k: [2*max_rel_pos+1, num_heads, head_size] | ||
| # Result: [batch_size, num_heads, 2*max_rel_pos+1] | ||
| c2p_scores = torch.einsum("bnh,rnh->bnr", q_i, rel_k) | ||
|
|
||
| # Extract the correct relative position for each j | ||
| for j in range(seq_len): | ||
| rel_idx = relative_pos[i, j].item() | ||
| content_to_position[:, :, i, j] = c2p_scores[:, :, rel_idx] | ||
|
|
||
| # 3. Position-to-content attention: relative positions to key content | ||
| position_to_content = torch.zeros_like(content_to_content, | ||
| device=query_layer.device) | ||
|
|
||
| # For each key position j, compute attention from relative positions | ||
| for j in range(seq_len): | ||
| # K_j^c: [batch_size, num_heads, head_size] | ||
| k_j = key_layer[:, :, j, :] # [batch_size, num_heads, head_size] | ||
|
|
||
| # Compute attention from all relative positions | ||
| # rel_q: [2*max_rel_pos+1, num_heads, head_size] | ||
| # Result: [batch_size, num_heads, 2*max_rel_pos+1] | ||
| p2c_scores = torch.einsum("bnh,rnh->bnr", k_j, rel_q) | ||
|
|
||
| # Extract the correct relative position for each i (note: j,i not | ||
| # i,j) | ||
| for i in range(seq_len): | ||
| rel_idx = relative_pos[ | ||
| j, i].item() # Reverse direction for position-to-content | ||
| position_to_content[:, :, i, j] = p2c_scores[:, :, rel_idx] | ||
|
|
||
| # Sum all three components | ||
| attention_scores = (content_to_content + content_to_position + | ||
| position_to_content) | ||
|
|
||
| return attention_scores |
There was a problem hiding this comment.
The nested Python loops in disentangled_attention_bias for calculating content_to_position and position_to_content can be a performance bottleneck. These loops should be vectorized using torch.einsum and torch.gather for better performance.
content_to_content = torch.matmul(query_layer,
key_layer.transpose(-2, -1))
# 2. Content-to-position attention: query content to relative positions
c2p_scores = torch.einsum("bnqh,rnh->bnqr", query_layer,
rel_k) # [B, H, Q, R]
rel_pos_expanded = relative_pos.unsqueeze(0).unsqueeze(0).expand(
batch_size, num_heads, -1, -1) # [B, H, Q, K]
content_to_position = torch.gather(c2p_scores, 3,
rel_pos_expanded) # [B, H, Q, K]
# 3. Position-to-content attention: relative positions to key content
p2c_scores = torch.einsum("bnkh,rnh->bnkr", key_layer,
rel_q) # [B, H, K, R]
rel_pos_T = relative_pos.transpose(
0, 1) # Transpose for p2c, shape [K, Q]
rel_pos_T_expanded = rel_pos_T.unsqueeze(0).unsqueeze(0).expand(
batch_size, num_heads, -1, -1) # [B, H, K, Q]
p2c_gathered = torch.gather(p2c_scores, 3,
rel_pos_T_expanded) # [B, H, K, Q]
position_to_content = p2c_gathered.transpose(
-1, -2) # [B, H, Q, K]
# Sum all three components
attention_scores = (content_to_content + content_to_position +
position_to_content)| "BertForSequenceClassification": ("bert", "BertForSequenceClassification"), | ||
| "DebertaV2ForSequenceClassification": ("deberta_v2", "DebertaV2ForSequenceClassification"), # noqa: E501 | ||
| "RobertaForSequenceClassification": ("roberta", | ||
| "RobertaForSequenceClassification"), | ||
| "XLMRobertaForSequenceClassification": ("roberta", | ||
| "RobertaForSequenceClassification"), | ||
| "ModernBertForSequenceClassification": ("modernbert", | ||
| "ModernBertForSequenceClassification"), | ||
| "Qwen3ForSequenceClassification": ("qwen3", "Qwen3ForSequenceClassification"), # noqa: E501 | ||
| "Qwen3ForSequenceClassification": ("qwen3", "Qwen3ForSequenceClassification"), # noqa: E501 |
There was a problem hiding this comment.
For better readability and maintainability, sort the dictionary entries alphabetically by key.
| "BertForSequenceClassification": ("bert", "BertForSequenceClassification"), | |
| "DebertaV2ForSequenceClassification": ("deberta_v2", "DebertaV2ForSequenceClassification"), # noqa: E501 | |
| "RobertaForSequenceClassification": ("roberta", | |
| "RobertaForSequenceClassification"), | |
| "XLMRobertaForSequenceClassification": ("roberta", | |
| "RobertaForSequenceClassification"), | |
| "ModernBertForSequenceClassification": ("modernbert", | |
| "ModernBertForSequenceClassification"), | |
| "Qwen3ForSequenceClassification": ("qwen3", "Qwen3ForSequenceClassification"), # noqa: E501 | |
| "Qwen3ForSequenceClassification": ("qwen3", "Qwen3ForSequenceClassification"), # noqa: E501 | |
| "BertForSequenceClassification": ("bert", "BertForSequenceClassification"), | |
| "DebertaV2ForSequenceClassification": ("deberta_v2", "DebertaV2ForSequenceClassification"), # noqa: E501 | |
| "ModernBertForSequenceClassification": ("modernbert", | |
| "ModernBertForSequenceClassification"), | |
| "Qwen3ForSequenceClassification": ("qwen3", "Qwen3ForSequenceClassification"), # noqa: E501 | |
| "RobertaForSequenceClassification": ("roberta", | |
| "RobertaForSequenceClassification"), | |
| "XLMRobertaForSequenceClassification": ("roberta", | |
| "RobertaForSequenceClassification"), |
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.Purpose
Added support for DebertaV2ForSequenceClassification models like:
meta-llama/Prompt-Guard-86M,mixedbread-ai/mxbai-rerank-xsmall-v1,microsoft/deberta-v2-xlarge-mnli, etcTest Plan
Test Result
(Optional) Documentation Update