Skip to content

vllm.model_executor.layers.quantization.torchao

TorchAOConfig

Bases: QuantizationConfig

Config class for torchao.

Source code in vllm/model_executor/layers/quantization/torchao.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
class TorchAOConfig(QuantizationConfig):
    """Config class for torchao."""

    def __init__(self, torchao_config) -> None:
        self.torchao_config = torchao_config

    def __repr__(self) -> str:
        return f"TorchAOConfig({self.torchao_config})"

    def get_name(self) -> QuantizationMethods:
        return "torchao"

    def get_supported_act_dtypes(self) -> list[torch.dtype]:
        return [torch.float32, torch.float16, torch.bfloat16]

    @classmethod
    def get_min_capability(cls) -> int:
        return 75

    @staticmethod
    def get_config_filenames() -> list[str]:
        return ["config.json"]

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "TorchAOConfig":
        """Create the quant config from an hf model config"""
        try:
            from torchao.core.config import config_from_dict
        except ImportError as err:
            raise ImportError(
                "Please install torchao>=0.10.0 via "
                "`pip install torchao>=0.10.0` to use torchao quantization."
            ) from err

        hf_config = cls.get_from_keys_or(config, ["quant_type"], None)
        assert hf_config is not None, "quant_type must be specified"
        assert (len(hf_config) == 1 and "default" in hf_config
                ), "Expected only one key 'default' in quant_type dictionary"
        quant_type = hf_config["default"]
        ao_config = config_from_dict(quant_type)
        return cls(ao_config)

    def get_quant_method(self, layer: torch.nn.Module,
                         prefix: str) -> Optional["QuantizeMethodBase"]:
        if not isinstance(layer, LinearBase):
            return None

        from torchao.quantization import AOPerModuleConfig

        module_fqn = prefix
        if isinstance(self.torchao_config, AOPerModuleConfig):
            module_fqn_to_config = self.torchao_config.module_fqn_to_config
            c = module_fqn_to_config.get(
                module_fqn) or module_fqn_to_config.get("_default", None)
            if c is not None:
                current_torchao_config = TorchAOConfig(c)
                return TorchAOLinearMethod(current_torchao_config)
            else:
                return UnquantizedLinearMethod()

        return TorchAOLinearMethod(self)

    def get_scaled_act_names(self) -> list[str]:
        return []

torchao_config instance-attribute 

torchao_config = torchao_config

__init__ 

__init__(torchao_config) -> None
Source code in vllm/model_executor/layers/quantization/torchao.py 
19
20
def __init__(self, torchao_config) -> None:
    self.torchao_config = torchao_config

__repr__ 

__repr__() -> str
Source code in vllm/model_executor/layers/quantization/torchao.py 
22
23
def __repr__(self) -> str:
    return f"TorchAOConfig({self.torchao_config})"

from_config classmethod 

from_config(config: dict[str, Any]) -> TorchAOConfig

Create the quant config from an hf model config

Source code in vllm/model_executor/layers/quantization/torchao.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
@classmethod
def from_config(cls, config: dict[str, Any]) -> "TorchAOConfig":
    """Create the quant config from an hf model config"""
    try:
        from torchao.core.config import config_from_dict
    except ImportError as err:
        raise ImportError(
            "Please install torchao>=0.10.0 via "
            "`pip install torchao>=0.10.0` to use torchao quantization."
        ) from err

    hf_config = cls.get_from_keys_or(config, ["quant_type"], None)
    assert hf_config is not None, "quant_type must be specified"
    assert (len(hf_config) == 1 and "default" in hf_config
            ), "Expected only one key 'default' in quant_type dictionary"
    quant_type = hf_config["default"]
    ao_config = config_from_dict(quant_type)
    return cls(ao_config)

get_config_filenames staticmethod 

get_config_filenames() -> list[str]
Source code in vllm/model_executor/layers/quantization/torchao.py 
35
36
37
@staticmethod
def get_config_filenames() -> list[str]:
    return ["config.json"]

get_min_capability classmethod 

get_min_capability() -> int
Source code in vllm/model_executor/layers/quantization/torchao.py 
31
32
33
@classmethod
def get_min_capability(cls) -> int:
    return 75

get_name 

get_name() -> QuantizationMethods
Source code in vllm/model_executor/layers/quantization/torchao.py 
25
26
def get_name(self) -> QuantizationMethods:
    return "torchao"

get_quant_method 

get_quant_method(
    layer: Module, prefix: str
) -> Optional[QuantizeMethodBase]
Source code in vllm/model_executor/layers/quantization/torchao.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def get_quant_method(self, layer: torch.nn.Module,
                     prefix: str) -> Optional["QuantizeMethodBase"]:
    if not isinstance(layer, LinearBase):
        return None

    from torchao.quantization import AOPerModuleConfig

    module_fqn = prefix
    if isinstance(self.torchao_config, AOPerModuleConfig):
        module_fqn_to_config = self.torchao_config.module_fqn_to_config
        c = module_fqn_to_config.get(
            module_fqn) or module_fqn_to_config.get("_default", None)
        if c is not None:
            current_torchao_config = TorchAOConfig(c)
            return TorchAOLinearMethod(current_torchao_config)
        else:
            return UnquantizedLinearMethod()

    return TorchAOLinearMethod(self)

get_scaled_act_names 

get_scaled_act_names() -> list[str]
Source code in vllm/model_executor/layers/quantization/torchao.py 
78
79
def get_scaled_act_names(self) -> list[str]:
    return []

get_supported_act_dtypes 

get_supported_act_dtypes() -> list[dtype]
Source code in vllm/model_executor/layers/quantization/torchao.py 
28
29
def get_supported_act_dtypes(self) -> list[torch.dtype]:
    return [torch.float32, torch.float16, torch.bfloat16]

TorchAOLinearMethod

Bases: LinearMethodBase

Linear method for torchao.

Parameters:

Name Type Description Default
torchao_config

The torchao quantization config, a string

 required
Source code in vllm/model_executor/layers/quantization/torchao.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
class TorchAOLinearMethod(LinearMethodBase):
    """Linear method for torchao.

    Args:
        torchao_config: The torchao quantization config, a string
        that encodes the type of quantization and all relevant arguments.
    """

    def __init__(self, quant_config: TorchAOConfig):
        self.quant_config = quant_config

    def create_weights(
        self,
        layer: torch.nn.Module,
        input_size_per_partition: int,
        output_partition_sizes: list[int],
        input_size: int,
        output_size: int,
        params_dtype: torch.dtype,
        **extra_weight_attrs,
    ):
        weight = Parameter(
            torch.empty(
                sum(output_partition_sizes),
                input_size_per_partition,
                dtype=params_dtype,
            ),
            requires_grad=False,
        )
        weight = torchao_quantize_param_data(weight,
                                             self.quant_config.torchao_config)

        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})

        layer.register_parameter("weight", weight)
        set_weight_attrs(weight, extra_weight_attrs)

    def apply(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        return F.linear(x, layer.weight, bias)

quant_config instance-attribute 

quant_config = quant_config

__init__ 

__init__(quant_config: TorchAOConfig)
Source code in vllm/model_executor/layers/quantization/torchao.py 
108
109
def __init__(self, quant_config: TorchAOConfig):
    self.quant_config = quant_config

apply 

apply(
    layer: Module, x: Tensor, bias: Optional[Tensor] = None
) -> Tensor
Source code in vllm/model_executor/layers/quantization/torchao.py 
137
138
139
140
141
142
143
def apply(
    self,
    layer: torch.nn.Module,
    x: torch.Tensor,
    bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    return F.linear(x, layer.weight, bias)

create_weights 

create_weights(
    layer: Module,
    input_size_per_partition: int,
    output_partition_sizes: list[int],
    input_size: int,
    output_size: int,
    params_dtype: dtype,
    **extra_weight_attrs,
)
Source code in vllm/model_executor/layers/quantization/torchao.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def create_weights(
    self,
    layer: torch.nn.Module,
    input_size_per_partition: int,
    output_partition_sizes: list[int],
    input_size: int,
    output_size: int,
    params_dtype: torch.dtype,
    **extra_weight_attrs,
):
    weight = Parameter(
        torch.empty(
            sum(output_partition_sizes),
            input_size_per_partition,
            dtype=params_dtype,
        ),
        requires_grad=False,
    )
    weight = torchao_quantize_param_data(weight,
                                         self.quant_config.torchao_config)

    set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})

    layer.register_parameter("weight", weight)
    set_weight_attrs(weight, extra_weight_attrs)

torchao_quantize_param_data 

torchao_quantize_param_data(
    param: Tensor, torchao_config: Any
) -> Parameter

Quantize a Tensor with torchao quantization specified by torchao_config

Parameters:

Name Type Description Default
`param`

weight parameter of the linear module

 required
`torchao_config`

type of quantization and their arguments we want to use to quantize the Tensor

 required
Source code in vllm/model_executor/layers/quantization/torchao.py 
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def torchao_quantize_param_data(param: torch.Tensor,
                                torchao_config: Any) -> torch.nn.Parameter:
    """Quantize a Tensor with torchao quantization specified by torchao_config

    Args:
       `param`: weight parameter of the linear module
       `torchao_config`: type of quantization and their arguments we want to
        use to quantize the Tensor
    """
    from torchao.core.config import AOBaseConfig
    from torchao.quantization import quantize_
    assert isinstance(torchao_config, AOBaseConfig), f"{torchao_config}"
    dummy_linear = torch.nn.Linear(param.shape[1], param.shape[0], bias=False)
    dummy_linear.weight = param
    quantize_(dummy_linear, torchao_config)
    return dummy_linear.weight