HF PR accepted.

- Removed granite config. - Updated test with offiicial model
vllm-project · Aug 28, 2024 · 45d1e3b · 45d1e3b
1 parent 1dcb38d
commit 45d1e3b
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 202 deletions.
diff --git a/tests/models/test_granite.py b/tests/models/test_granite.py
@@ -7,7 +7,7 @@
 from .utils import check_logprobs_close
 
 MODELS = [
-    "mayank-mishra/granite-3b-mup",
+    "ibm/PowerLM-3b",
 ]
 
 

diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py
@@ -48,8 +48,7 @@
     default_weight_loader, kv_cache_scales_loader, maybe_remap_kv_scale_name)
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.sequence import IntermediateTensors, SamplerOutput
-# from transformers import GraniteConfig
-from vllm.transformers_utils.configs.granite import GraniteConfig
+from transformers import GraniteConfig
 from vllm.utils import is_hip
 
 from .interfaces import SupportsLoRA
@@ -216,6 +215,7 @@ def __init__(
             cache_config=cache_config,
             prefix=f"{prefix}.self_attn",
         )
+
         self.mlp = GraniteMLP(
             hidden_size=self.hidden_size,
             intermediate_size=config.intermediate_size,

diff --git a/vllm/transformers_utils/configs/granite.py b/vllm/transformers_utils/configs/granite.py