[ Misc ] fbgemm checkpoints (#6559)

2024-07-20 12:36:57 -04:00
parent 9042d68362
commit 683e3cb9c4
24 changed files with 234 additions and 47 deletions
--- a/vllm/model_executor/layers/quantization/base_config.py
+++ b/vllm/model_executor/layers/quantization/base_config.py
@@ -97,12 +97,13 @@ class QuantizationConfig(ABC):
            return default

    @abstractmethod
-    def get_quant_method(
-            self, layer: torch.nn.Module) -> Optional[QuantizeMethodBase]:
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional[QuantizeMethodBase]:
        """Get the quantize method to use for the quantized layer.
        
        Args:
            layer: The layer for the quant method.
+            prefix: The full name of the layer in the state dict
        Returns:
            The quantize method. None if the given layer doesn't support quant
            method.