Aluode/PerceptionLabPortable
0
1# Copyright 2024 The HuggingFace Inc. team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from typing import TYPE_CHECKING, Optional15 16from .base import HfQuantizer17 18 19if TYPE_CHECKING:20 from ..modeling_utils import PreTrainedModel21 22from ..utils import is_accelerate_available, is_fbgemm_gpu_available, is_torch_available, logging23from .quantizers_utils import get_module_from_name24 25 26if is_torch_available():27 import torch28 29 30logger = logging.get_logger(__name__)31 32 33class FbgemmFp8HfQuantizer(HfQuantizer):34 """35 FP8 quantization using fbgemm kernels36 """37 38 requires_parameters_quantization = True39 requires_calibration = False40 41 required_packages = ["fbgemm-gpu", "accelerate"]42 43 def __init__(self, quantization_config, **kwargs):44 super().__init__(quantization_config, **kwargs)45 self.quantization_config = quantization_config46 47 def validate_environment(self, *args, **kwargs):48 if not is_torch_available():49 raise ImportError(50 "Using fbgemm fp8 quantization requires torch >= 2.1.0"51 "Please install the latest version of torch ( pip install --upgrade torch )"52 )53 if not is_fbgemm_gpu_available():54 raise ImportError(55 "Using fbgemm fp8 quantization requires fbgemm-gpu library"56 "Please install the latest version of fbgemm-gpu library by following : https://pytorch.org/FBGEMM/fbgemm_gpu-development/InstallationInstructions.html#fbgemm-gpu-install-libraries"57 )58 59 if not is_accelerate_available("0.32.2"):60 raise ImportError(61 "Loading an FP8 quantized model requires accelerate > 0.32.1 (`pip install --upgrade accelerate`)"62 )63 64 if not torch.cuda.is_available():65 raise RuntimeError("Using FP8 quantized models with fbgemm kernels requires a GPU")66 67 compute_capability = torch.cuda.get_device_capability()68 major, minor = compute_capability69 if major < 9:70 raise ValueError(71 "FP8 quantized models is only supported on GPUs with compute capability >= 9.0 (e.g H100)"72 )73 74 device_map = kwargs.get("device_map")75 if device_map is None:76 logger.warning_once(77 "You have loaded an FP8 model on CPU and have a CUDA device available, make sure to set "78 "your model on a GPU device in order to run your model. To remove this warning, pass device_map = 'cuda'. "79 )80 elif device_map is not None:81 if (82 not self.pre_quantized83 and isinstance(device_map, dict)84 and ("cpu" in device_map.values() or "disk" in device_map.values())85 ):86 raise ValueError(87 "You are attempting to load an FP8 model with a device_map that contains a CPU or disk device."88 "This is not supported when the model is quantized on the fly. "89 "Please use a quantized checkpoint or remove the CPU or disk device from the device_map."90 )91 92 def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype":93 if dtype is None:94 dtype = torch.bfloat1695 logger.info(96 "Overriding dtype=%s with `dtype=torch.bloat16` due to "97 "requirements of `fbgemm-gpu` to enable model loading in fp8. "98 "Pass your own dtype to specify the dtype of the remaining non-linear layers or pass"99 " dtype=torch.bfloat16 to remove this warning.",100 dtype,101 )102 elif dtype == torch.float16:103 raise ValueError(104 "You cannot use FP8 with dtype=torch.float16.We recommend you passing dtype=torch.bfloat16"105 )106 return dtype107 108 def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool:109 from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts110 111 module, tensor_name = get_module_from_name(model, param_name)112 113 if isinstance(module, FbgemmFp8Linear):114 if self.pre_quantized or tensor_name == "bias":115 return False116 else:117 return True118 if isinstance(module, FbgemmFp8Llama4TextExperts):119 if self.pre_quantized or tensor_name == "bias":120 return False121 else:122 return True123 return False124 125 def create_quantized_param(126 self,127 model: "PreTrainedModel",128 param_value: "torch.Tensor",129 param_name: str,130 target_device: "torch.device",131 **kwargs,132 ):133 from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts134 135 module, tensor_name = get_module_from_name(model, param_name)136 137 # Sanity checks138 if isinstance(module, FbgemmFp8Linear):139 if self.pre_quantized or tensor_name == "bias":140 if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn:141 raise ValueError("Expect quantized weights but got an unquantized weight")142 else:143 if tensor_name == "weight_scale":144 raise ValueError("Expect unquantized weights but got a quantized weight_scale")145 if isinstance(module, FbgemmFp8Llama4TextExperts):146 if not (self.pre_quantized or tensor_name == "bias"):147 if tensor_name == "gate_up_proj_scale" or tensor_name == "down_proj_scale":148 raise ValueError("Expect unquantized weights but got a quantized weight_scale")149 150 if isinstance(module, FbgemmFp8Llama4TextExperts):151 if tensor_name == "gate_up_proj":152 # Process each expert separately153 # Transpose the second and third dimension154 transposed_param = param_value.transpose(1, 2)155 156 # Reshape to 2D for quantization157 original_shape = transposed_param.shape158 flattened_param = transposed_param.reshape(-1, original_shape[-1])159 160 # Quantize using per row instead of per column161 new_value_flat, weight_scale_flat = torch.ops.fbgemm.quantize_fp8_per_row(flattened_param)162 163 # Reshape back to original dimensions164 new_value = new_value_flat.reshape(original_shape)165 new_value = new_value.transpose(1, 2)166 weight_scale = weight_scale_flat.reshape(original_shape[0], 1, original_shape[1])167 elif tensor_name == "down_proj":168 # Process each expert separately169 # Transpose the weights for proper quantization170 transposed_param = param_value.transpose(1, 2)171 172 # Reshape to 2D for quantization173 original_shape = transposed_param.shape174 flattened_param = transposed_param.reshape(-1, original_shape[-1])175 176 # Quantize using per column177 new_value_flat, weight_scale_flat = torch.ops.fbgemm.quantize_fp8_per_row(flattened_param)178 179 # Reshape back to original dimensions180 new_value = new_value_flat.reshape(original_shape)181 new_value = new_value.transpose(1, 2)182 weight_scale = weight_scale_flat.reshape(original_shape[0], original_shape[1], 1)183 184 module._parameters[f"{tensor_name}_scale"] = torch.nn.Parameter(weight_scale.to(target_device))185 else:186 new_value, weight_scale = torch.ops.fbgemm.quantize_fp8_per_row(param_value)187 module._parameters[f"{tensor_name}_scale"] = torch.nn.Parameter(188 weight_scale.view(weight_scale.shape[0], 1).to(target_device)189 )190 191 module._parameters[tensor_name] = torch.nn.Parameter(new_value.to(target_device))192 193 del param_name194 195 def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):196 return model197 198 def _process_model_before_weight_loading(199 self,200 model: "PreTrainedModel",201 keep_in_fp32_modules: Optional[list[str]] = None,202 **kwargs,203 ):204 from ..integrations import replace_with_fbgemm_fp8_linear205 206 tp_plan = model._tp_plan207 self.modules_to_not_convert = self.get_modules_to_not_convert(208 model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules209 )210 211 config = model.config212 model = replace_with_fbgemm_fp8_linear(213 model,214 modules_to_not_convert=self.modules_to_not_convert,215 quantization_config=self.quantization_config,216 pre_quantized=self.pre_quantized,217 config=config,218 tp_plan=tp_plan,219 )220 221 model.config.quantization_config = self.quantization_config222 223 def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:224 from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts225 226 not_missing_keys = []227 for name, module in model.named_modules():228 if isinstance(module, (FbgemmFp8Linear, FbgemmFp8Llama4TextExperts)):229 for missing in missing_keys:230 if (231 (name in missing or name in f"{prefix}.{missing}")232 and not missing.endswith(".weight")233 and not missing.endswith(".bias")234 ):235 not_missing_keys.append(missing)236 return [k for k in missing_keys if k not in not_missing_keys]237 238 def update_tp_plan(self, config):239 if "Llama4" in config.__class__.__name__:240 text_plan = {241 # We are using a different tp plan with local_colwise and local_rowwise for the attention because fbgemm operations cannot be parallelized242 # With local_colwise and local_rowwise, all the operations are done locally, and we add a gather operation to gather the results instead of243 # using dtensors244 "layers.*.self_attn.q_proj.weight": "local_colwise",245 "layers.*.self_attn.q_proj.weight_scale": "local_colwise",246 "layers.*.self_attn.k_proj.weight": "local_colwise",247 "layers.*.self_attn.k_proj.weight_scale": "local_colwise",248 "layers.*.self_attn.v_proj.weight": "local_colwise",249 "layers.*.self_attn.v_proj.weight_scale": "local_colwise",250 "layers.*.self_attn.o_proj.weight": "local_rowwise",251 "layers.*.self_attn": "gather",252 # We keep the same sequence_parallel plan for layernorms253 "layers.*.input_layernorm.weight": "sequence_parallel",254 "layers.*.post_attention_layernorm.weight": "sequence_parallel",255 "norm.weight": "sequence_parallel",256 # We keep the same local_colwise and local_rowwise plan for the feed forward shared expert257 # We also add scales for the shared expert, for local_colwise the scale is also local_colwise258 # For local_rowwise the scale is replicated, so we don't need to add it259 "layers.*.feed_forward.shared_expert.gate_proj.weight": "local_colwise",260 "layers.*.feed_forward.shared_expert.gate_proj.weight_scale": "local_colwise",261 "layers.*.feed_forward.shared_expert.up_proj.weight": "local_colwise",262 "layers.*.feed_forward.shared_expert.up_proj.weight_scale": "local_colwise",263 "layers.*.feed_forward.shared_expert.down_proj.weight": "local_rowwise",264 "layers.*.feed_forward.experts": "local",265 "layers.*.feed_forward": "gather",266 "layers.*.feed_forward.experts.*.gate_proj.weight": "local_colwise",267 "layers.*.feed_forward.experts.*.gate_proj.weight_scale": "local_colwise",268 "layers.*.feed_forward.experts.*.up_proj.weight": "local_colwise",269 "layers.*.feed_forward.experts.*.up_proj.weight_scale": "local_colwise",270 "layers.*.feed_forward.experts.*.down_proj.weight": "local_rowwise",271 # For Fused implementation we use local_packed_rowwise for the gate_up_proj, and the same for the packed scales272 # We use local_colwise for the down_proj, and the scales are replicated so we don't add them273 "layers.*.feed_forward.experts.gate_up_proj": "local_packed_rowwise",274 "layers.*.feed_forward.experts.gate_up_proj_scale": "local_packed_rowwise",275 "layers.*.feed_forward.experts.down_proj": "local_colwise",276 }277 if config.get_text_config() is not None:278 config.get_text_config().base_model_tp_plan = text_plan279 else:280 config.base_model_tp_plan = text_plan281 return config282 283 return config284 285 def is_serializable(self, safe_serialization=None):286 return True287 288 @property289 def is_trainable(self) -> bool:290 return False291