huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-27 04:52:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

marlin: support tp>1 when group_size==-1

Browse Source
This commit is contained in:
Daniël de Kok 2024-06-06 11:51:52 +00:00 committed by Daniël de Kok
parent 4594e6faba
commit 0d96468ebb
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
server/text_generation_server/utils/weights.py
View File

@ -513,7 +513,13 @@ class Weights:
"Cannot load `marlin` weight, make sure the model is already quantized, or quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
)
s = self.get_sharded(f"{prefix}.s", dim=0)
num_groups = self._get_slice(f"{prefix}.s").get_shape()[0]
if num_groups == 1:
# The number of groups is 1 when group_size == -1. share
# scales between all shards in this case.
s = self.get_tensor(f"{prefix}.s")
else:
s = self.get_sharded(f"{prefix}.s", dim=0)
weight = MarlinWeight(B=B, s=s)
else: