Unverified Commit 36d1ccb2 authored by Kyle Sayers's avatar Kyle Sayers Committed by GitHub
Browse files

[Quant] BartModel SupportsQuant (#14699)


Signed-off-by: default avatarKyle Sayers <kylesayrs@gmail.com>
parent 1bc3b739
...@@ -44,7 +44,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader ...@@ -44,7 +44,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.sampling_metadata import SamplingMetadata from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.sequence import IntermediateTensors from vllm.sequence import IntermediateTensors
from .interfaces import SupportsV0Only from .interfaces import SupportsQuant, SupportsV0Only
from .utils import maybe_prefix from .utils import maybe_prefix
logger = logging.get_logger(__name__) logger = logging.get_logger(__name__)
...@@ -697,7 +697,7 @@ class BartDecoder(nn.Module): ...@@ -697,7 +697,7 @@ class BartDecoder(nn.Module):
return hidden_states return hidden_states
class BartModel(nn.Module): class BartModel(nn.Module, SupportsQuant):
_tied_weights_keys = [ _tied_weights_keys = [
"encoder.embed_tokens.weight", "decoder.embed_tokens.weight" "encoder.embed_tokens.weight", "decoder.embed_tokens.weight"
] ]
...@@ -763,7 +763,8 @@ class BartModel(nn.Module): ...@@ -763,7 +763,8 @@ class BartModel(nn.Module):
return decoder_outputs return decoder_outputs
class BartForConditionalGeneration(nn.Module, SupportsV0Only): class BartForConditionalGeneration(nn.Module, SupportsV0Only, SupportsQuant):
packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
base_model_prefix = "model" base_model_prefix = "model"
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment