Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: support loading eetq quantized model #393

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

feat: support loading eetq quantized model #393

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
90d5efe
feat: support loading eetq quantized model
Apr 5, 2024
9d91b37
fix: reduce the changes in EETQLinear
Apr 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion server/lorax_server/models/custom_modeling/flash_llama_modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,7 @@ def _load_gqa(config, prefix: str, weights):
dim=0,
)

if config.quantize not in ["gptq", "awq"]:
if isinstance(weight, torch.Tensor):
weight = weight.to(dtype=weights.dtype).to(device=weights.device)

head_size = config.hidden_size // config.num_attention_heads
Expand Down
26 changes: 23 additions & 3 deletions server/lorax_server/utils/layers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from accelerate import init_empty_weights
from torch import nn
from torch.nn import functional as F
from loguru import logger

from lorax_server.adapters.types import LORA, MEDUSA
from lorax_server.utils.gptq.quant_linear import QuantLinear
Expand Down Expand Up @@ -166,8 +167,17 @@ def __init__(
self,
weight,
bias,
scales=None,
quantized=False,
) -> None:
super().__init__()

if quantized:
self.weight = weight
self.scale = scales
self.bias = bias if bias is not None else None
return

# Get the device where the weight tensor is currently stored.
device = weight.device

Expand Down Expand Up @@ -344,10 +354,20 @@ def get_linear(weight, bias, quantize, fan_in_fan_out=False):
quant_type="fp4",
)
elif quantize == "eetq":
if HAS_EETQ:
linear = EETQLinear(weight, bias)
else:
if not HAS_EETQ:
raise ImportError("Please install EETQ from https://github.com/NetEase-FuXi/EETQ")

try:
qweight, scales = weight
linear = EETQLinear(
qweight,
bias,
scales,
True,
)
except Exception:
logger.info("It seems that weight not quantized, make JIT now")
linear = EETQLinear(weight, bias)
elif quantize == "gptq":
try:
qweight, qzeros, scales, g_idx, bits, groupsize, use_exllama = weight
Expand Down
21 changes: 19 additions & 2 deletions server/lorax_server/utils/weights.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ def get_tensor(self, tensor_name: str):
tensor = f.get_tensor(tensor_name)
# Special case for gptq which shouldn't convert
# u4 which are disguised as int32
if tensor.dtype not in [torch.int32, torch.int64]:
if tensor.dtype not in [torch.int8, torch.int32, torch.int64]:
tensor = tensor.to(dtype=self.dtype)
tensor = tensor.to(device=self.device)
return tensor
Expand Down Expand Up @@ -178,7 +178,7 @@ def get_partial_sharded(self, tensor_name: str, dim: int, range: Optional[Tuple[
raise NotImplementedError("Let's make that generic when needed")
# Special case for gptq which shouldn't convert
# u4 which are disguised as int32
if tensor.dtype != torch.int32:
if tensor.dtype not in [torch.int8, torch.int32]:
tensor = tensor.to(dtype=self.dtype)
tensor = tensor.to(device=self.device)
return tensor
Expand Down Expand Up @@ -226,6 +226,15 @@ def get_multi_weights_col(self, prefixes: List[Union[str, Tuple]], quantize: str

bits, groupsize = self._get_gptq_params()
weight = (qweight, qzeros, scales, g_idx, bits, groupsize, False)
elif quantize == "eetq":
try:
qweight = torch.cat(self.get_sharded_list("qweight", prefixes, dim=1), dim=1)
Copy link

@SidaZh SidaZh Apr 15, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it necessary to merge the weight parameters of multiple cards here? The quantization of eetq includes two steps: quantization and cutlass relayout. If the tensor is sliced or concat, the layout will be destroyed and output will be wrong.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are two options: 1. Save and load using the same tp and pp strategies; 2. Load basic per-channel quantization parameters and relayout when initializing EETQLinear which requires some development. The following two interfaces will help:

from EETQ import quant_weights, preprocess_weights

unprocessed_quantized_weight, processed_quantized_weight, scales = quant_weights(unquantized_weight, torch.int8, True)    # quantize and relayout
processed_quantized_weight = preprocess_weights(unprocessed_quantized_weight)   # relayout

scales = torch.cat(self.get_sharded_list("weight_scales", prefixes, dim=0), dim=0)
weight = (qweight, scales)
except RuntimeError:
logger.info("It seems that weight is not quantized, so load it normally then make JIT later")
w = self.get_sharded_list("weight", prefixes, dim=0)
weight = torch.cat(w, dim=dim)
else:
w = self.get_sharded_list("weight", prefixes, dim=0)
weight = torch.cat(w, dim=dim)
Expand Down Expand Up @@ -310,6 +319,14 @@ def get_multi_weights_row(self, prefix: str, quantize: str):
g_idx = None
use_exllama = False
weight = (qweight, qzeros, scales, g_idx, bits, groupsize, use_exllama)
elif quantize == "eetq":
try:
qweight = self.get_sharded(f"{prefix}.qweight", dim=0)
scales = self.get_sharded(f"{prefix}.weight_scales", dim=0)
weight = (qweight, scales)
except RuntimeError:
logger.info("It seems that weight is not quantized, so load it normally then make JIT later")
weight = self.get_sharded(f"{prefix}.weight", dim=1)
else:
weight = self.get_sharded(f"{prefix}.weight", dim=1)
return weight
Expand Down