squeezellm.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  1. from typing import Any, Dict, List, Optional
  2. from contextlib import suppress
  3. import torch
  4. from torch.nn.parameter import Parameter
  5. from aphrodite.common.utils import is_hip
  6. from aphrodite.modeling.layers.linear import LinearBase
  7. from aphrodite.modeling.utils import set_weight_attrs
  8. from aphrodite.quantization.base_config import (QuantizationConfig,
  9. QuantizeMethodBase)
  10. HAS_QUANTS = False
  11. with suppress(ImportError):
  12. from aphrodite._quant_C import quant_ops as ops
  13. HAS_QUANTS = True
  14. class SqueezeLLMConfig(QuantizationConfig):
  15. """Config class for SqueezeLLM.
  16. Reference: https://arxiv.org/pdf/2306.07629
  17. """
  18. def __init__(
  19. self,
  20. weight_bits: int,
  21. ) -> None:
  22. self.weight_bits = weight_bits
  23. if self.weight_bits != 4:
  24. raise ValueError(
  25. "Currently, only 4-bit weight quantization is supported for "
  26. f"SqueezeLLM, but got {self.weight_bits} bits.")
  27. self.pack_factor = 32 // self.weight_bits
  28. def __repr__(self) -> str:
  29. return f"SqueezeLLMConfig(weight_bits={self.weight_bits})"
  30. def get_name(self) -> str:
  31. return "squeezellm"
  32. def get_supported_act_dtypes(self) -> List[torch.dtype]:
  33. return [torch.half]
  34. def get_min_capability(self) -> int:
  35. return 70
  36. @staticmethod
  37. def get_config_filenames() -> List[str]:
  38. return ["quant_config.json"]
  39. @classmethod
  40. def from_config(cls, config: Dict[str, Any]) -> "SqueezeLLMConfig":
  41. weight_bits = cls.get_from_keys(config, ["wbits"])
  42. return cls(weight_bits)
  43. def get_quant_method(
  44. self,
  45. layer: torch.nn.Module) -> Optional["SqueezeLLMLinearMethod"]:
  46. if isinstance(layer, LinearBase):
  47. return SqueezeLLMLinearMethod(self)
  48. return
  49. def get_scaled_act_names(self) -> List[str]:
  50. return []
  51. class SqueezeLLMLinearMethod(QuantizeMethodBase):
  52. """Linear method for SqueezeLLM.
  53. Args:
  54. quant_config: The SqueezeLLM quantization config.
  55. """
  56. def __init__(self, quant_config: SqueezeLLMConfig):
  57. if not HAS_QUANTS:
  58. raise ImportError("Could not find the quantization kernels.")
  59. self.quant_config = quant_config
  60. def create_weights(self, layer: torch.nn.Module,
  61. input_size_per_partition: int,
  62. output_partition_sizes: List[int], input_size: int,
  63. output_size: int, params_dtype: torch.dtype,
  64. **extra_weight_attrs):
  65. if input_size_per_partition % self.quant_config.pack_factor != 0:
  66. raise ValueError(
  67. "The input size is not aligned with the quantized "
  68. "weight shape. This can be caused by too large "
  69. "tensor parallel size.")
  70. output_size_per_partition = sum(output_partition_sizes)
  71. qweight = Parameter(
  72. torch.empty(
  73. input_size_per_partition // self.quant_config.pack_factor,
  74. output_size_per_partition,
  75. dtype=torch.int32,
  76. ),
  77. requires_grad=False,
  78. )
  79. set_weight_attrs(
  80. qweight, {
  81. "input_dim": 0,
  82. "output_dim": 1,
  83. "packed_dim": 0,
  84. "pack_factor": self.quant_config.pack_factor,
  85. })
  86. lookup_table = Parameter(
  87. torch.empty(
  88. output_size,
  89. self.quant_config.weight_bits**2,
  90. dtype=params_dtype,
  91. ),
  92. requires_grad=False,
  93. )
  94. set_weight_attrs(lookup_table, {
  95. "output_dim": 0,
  96. })
  97. layer.register_parameter("qweight", qweight)
  98. set_weight_attrs(qweight, extra_weight_attrs)
  99. layer.register_parameter("lookup_table", lookup_table)
  100. set_weight_attrs(lookup_table, extra_weight_attrs)
  101. def apply(self,
  102. layer: torch.nn.Module,
  103. x: torch.Tensor,
  104. bias: Optional[torch.Tensor] = None) -> torch.Tensor:
  105. qweight = layer.qweight
  106. lookup_table = layer.lookup_table
  107. out_shape = x.shape[:-1] + (qweight.shape[-1], )
  108. reshaped_x = x.reshape(-1, x.shape[-1])
  109. if is_hip():
  110. out_f = torch.zeros(out_shape, dtype=torch.float)
  111. ops.squeezellm_gemm(reshaped_x, qweight, out_f, lookup_table)
  112. out = out_f.to(dtype=torch.float16)
  113. else:
  114. # NOTE: The output tensor should be zero-initialized.
  115. out = torch.zeros(out_shape, dtype=torch.float16)
  116. ops.squeezellm_gemm(reshaped_x, qweight, out, lookup_table)
  117. if bias is not None:
  118. out.add_(bias)
  119. return out.reshape(out_shape)