A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 23:30 UTC
[Quant] Aria SupportsQuant (#13416)
Kyle Sayers committed
Feb 18, 2025 at 00:51 UTC
d1b649f1eff6b8b4ce4683b6d11471e79c40ed7b
1 file changed
+10
-3
vllm/model_executor/models/aria.py
+10
-3
@@ -36,7 +36,7 @@ from .idefics2_vision_model import Idefics2VisionConfig
36
from .idefics2_vision_model import (
37
Idefics2VisionTransformer as Idefics3VisionTransformer)
38
# yapf: enable
39
-from .interfaces import SupportsMultiModal
39
+from .interfaces import SupportsMultiModal, SupportsQuant
40
from .llama import LlamaDecoderLayer, LlamaMLP, LlamaModel
41
from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
42
is_pp_missing_parameter, maybe_prefix,
@@ -53,7 +53,8 @@ class AriaImagePixelInputs(TypedDict):
53
"""
54
55
56
-class AriaVisionTransformer(Idefics3VisionTransformer):
56
+class AriaVisionTransformer(Idefics3VisionTransformer, SupportsQuant):
57
+ packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
58
59
def __init__(
60
self,
@@ -304,11 +305,17 @@ class AriaTextDecoderLayer(LlamaDecoderLayer):
305
self.mlp = AriaTextMoELayer(config, quant_config=quant_config)
306
307
307
-class AriaTextModel(LlamaModel):
308
+class AriaTextModel(LlamaModel, SupportsQuant):
309
"""
310
Custom LlamaModel for the AriaMoE model which modifies the standard
311
LlamaModel by replacing the `LlamaDecoderLayer` with `MoEDecoderLayer`.
312
"""
313
+ packed_modules_mapping = {
314
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
315
+ "gate_up_proj": ["gate_proj", "up_proj"],
316
+ "experts.w13_weight": ["experts.fc1.weight"],
317
+ "experts.w2_weight": ["experts.fc2.weight"],
318
+ }
319
320
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
321
super().__init__(vllm_config=vllm_config,