@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 186 lines 6.23 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from typing import Any
5
6 import torch
7 from torch.nn import Module
8 from torch.nn.parameter import Parameter
9
10 from vllm.config import get_current_vllm_config
11 from vllm.logger import init_logger
12 from vllm.model_executor.kernels.linear import (
13 init_fp8_linear_kernel,
14 )
15 from vllm.model_executor.layers.linear import (
16 LinearBase,
17 LinearMethodBase,
18 UnquantizedLinearMethod,
19 )
20 from vllm.model_executor.layers.quantization import QuantizationMethods
21 from vllm.model_executor.layers.quantization.base_config import (
22 QuantizationConfig,
23 QuantizeMethodBase,
24 )
25 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
26 prepare_fp8_layer_for_marlin,
27 )
28 from vllm.model_executor.layers.quantization.utils.quant_utils import (
29 is_layer_skipped,
30 kFp8DynamicTokenSym,
31 kFp8StaticTokenSym,
32 )
33 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
34 normalize_e4m3fn_to_e4m3fnuz,
35 )
36 from vllm.model_executor.parameter import (
37 ChannelQuantScaleParameter,
38 ModelWeightParameter,
39 )
40 from vllm.platforms import current_platform
41
42 logger = init_logger(__name__)
43
44
45 class FBGEMMFp8Config(QuantizationConfig):
46 """Config class for FBGEMM Fp8."""
47
48 def __init__(self, ignore_list: list[str], input_scale_ub: float):
49 super().__init__()
50 self.ignore_list = ignore_list if ignore_list else []
51 self.input_scale_ub = input_scale_ub
52
53 # For GPUs that lack FP8 hardware support, we can leverage the Marlin
54 # kernel for fast weight-only FP8 quantization
55 self.use_marlin = not current_platform.has_device_capability(89)
56
57 @classmethod
58 def get_name(cls) -> QuantizationMethods:
59 return "fbgemm_fp8"
60
61 @classmethod
62 def get_supported_act_dtypes(cls) -> list[torch.dtype]:
63 return [torch.bfloat16, torch.float16]
64
65 @classmethod
66 def get_min_capability(cls) -> int:
67 return 80
68
69 @classmethod
70 def get_config_filenames(cls) -> list[str]:
71 return []
72
73 @classmethod
74 def from_config(cls, config: dict[str, Any]) -> "FBGEMMFp8Config":
75 ignore_list = cls.get_from_keys(config, ["modules_to_not_convert"])
76 input_scale_ub = cls.get_from_keys(config, ["activation_scale_ub"])
77 return cls(ignore_list=ignore_list, input_scale_ub=input_scale_ub)
78
79 def get_quant_method(
80 self, layer: torch.nn.Module, prefix: str
81 ) -> "QuantizeMethodBase | None":
82 if isinstance(layer, LinearBase):
83 if is_layer_skipped(
84 prefix=prefix,
85 ignored_layers=self.ignore_list,
86 fused_mapping=self.packed_modules_mapping,
87 ):
88 return UnquantizedLinearMethod()
89 return FBGEMMFp8LinearMethod(self)
90 return None
91
92
93 class FBGEMMFp8LinearMethod(LinearMethodBase):
94 def __init__(self, quant_config: FBGEMMFp8Config):
95 self.quant_config = quant_config
96 self.out_dtype = torch.get_default_dtype()
97 self.input_dtype = get_current_vllm_config().model_config.dtype
98
99 def create_weights(
100 self,
101 layer: torch.nn.Module,
102 input_size_per_partition: int,
103 output_partition_sizes: list[int],
104 input_size: int,
105 output_size: int,
106 params_dtype: torch.dtype,
107 **extra_weight_attrs,
108 ):
109 weight_loader = extra_weight_attrs.get("weight_loader")
110 del input_size, output_size
111 output_size_per_partition = sum(output_partition_sizes)
112
113 layer.logical_widths = output_partition_sizes
114
115 layer.input_size_per_partition = input_size_per_partition
116 layer.output_size_per_partition = output_size_per_partition
117 layer.orig_dtype = params_dtype
118
119 # WEIGHT
120 weight = ModelWeightParameter(
121 data=torch.empty(
122 output_size_per_partition,
123 input_size_per_partition,
124 dtype=torch.float8_e4m3fn,
125 ),
126 input_dim=1,
127 output_dim=0,
128 weight_loader=weight_loader,
129 )
130 layer.register_parameter("weight", weight)
131
132 # WEIGHT SCALE
133 weight_scale = ChannelQuantScaleParameter(
134 data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32),
135 output_dim=0,
136 weight_loader=weight_loader,
137 )
138 weight_scale[:] = torch.finfo(torch.float32).min
139 layer.register_parameter("weight_scale", weight_scale)
140
141 # INPUT SCALE UPPER BOUND
142 input_scale_ub = torch.nn.Parameter(
143 torch.tensor((self.quant_config.input_scale_ub), dtype=torch.float32),
144 requires_grad=False,
145 )
146 layer.input_scale_ub = input_scale_ub
147
148 self.fp8_linear = init_fp8_linear_kernel(
149 activation_quant_key=kFp8DynamicTokenSym,
150 weight_quant_key=kFp8StaticTokenSym,
151 weight_shape=layer.weight.shape,
152 input_dtype=self.input_dtype,
153 out_dtype=self.out_dtype,
154 module_name=self.__class__.__name__,
155 )
156
157 def process_weights_after_loading(self, layer: Module) -> None:
158 # required by torch.compile
159 layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
160 layer.weight = Parameter(layer.weight.data, requires_grad=False)
161
162 weight = layer.weight
163
164 if current_platform.is_fp8_fnuz():
165 weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
166 weight=weight, weight_scale=layer.weight_scale, input_scale=None
167 )
168 if input_scale is not None:
169 layer.input_scale = Parameter(input_scale, requires_grad=False)
170 layer.weight_scale = Parameter(weight_scale, requires_grad=False)
171
172 layer.weight = Parameter(weight.t(), requires_grad=False)
173 if self.quant_config.use_marlin:
174 prepare_fp8_layer_for_marlin(layer)
175 # Activations not quantized for marlin.
176 del layer.input_scale_ub
177
178 self.fp8_linear.process_weights_after_loading(layer)
179
180 def apply(
181 self,
182 layer: torch.nn.Module,
183 x: torch.Tensor,
184 bias: torch.Tensor | None = None,
185 ) -> torch.Tensor:
186 return self.fp8_linear.apply_weights(layer, x, bias)