| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Prepare and execute Humming linear layers.""" |
| 4 | |
| 5 | import json |
| 6 | from typing import TYPE_CHECKING |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm import envs |
| 11 | from vllm.model_executor.layers.linear import LinearBase |
| 12 | from vllm.model_executor.layers.quantization.utils.humming.schema import ( |
| 13 | check_and_fallback_input_schema, |
| 14 | ) |
| 15 | |
| 16 | if TYPE_CHECKING: |
| 17 | from vllm.utils.humming import BaseInputSchema, LayerConfig |
| 18 | |
| 19 | |
| 20 | def convert_linear_layer_to_humming_standard( |
| 21 | layer: LinearBase, name_map: dict[str, str] |
| 22 | ): |
| 23 | """Rename/reshape a linear layer's quantized params (the canonical MPLinear |
| 24 | layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names |
| 25 | and layout humming's weight schema expects (``weight`` / ``weight_scale``).""" |
| 26 | for name, checkpoint_name in name_map.items(): |
| 27 | tensor = getattr(layer, checkpoint_name) |
| 28 | delattr(layer, checkpoint_name) |
| 29 | |
| 30 | if name == "weight": |
| 31 | input_dim = getattr(tensor, "input_dim", 1) |
| 32 | output_dim = getattr(tensor, "output_dim", 0) |
| 33 | |
| 34 | if input_dim == 0 and output_dim == 1: |
| 35 | tensor = tensor.transpose(1, 0).contiguous() |
| 36 | else: |
| 37 | assert output_dim == 0 and input_dim == 1 |
| 38 | |
| 39 | tensor = tensor.view(tensor.size(0), -1).view(torch.int32) |
| 40 | elif name in ["weight_scale", "zero_point"]: |
| 41 | if getattr(tensor, "output_dim", 0) == 1: |
| 42 | tensor = tensor.transpose(0, 1).contiguous() |
| 43 | if tensor.ndim == 1: |
| 44 | tensor = tensor.unsqueeze(1) |
| 45 | |
| 46 | tensor = tensor.view(torch.int32) if name == "zero_point" else tensor |
| 47 | |
| 48 | if isinstance(tensor, torch.nn.Parameter): |
| 49 | param = tensor |
| 50 | else: |
| 51 | param = torch.nn.Parameter(tensor, requires_grad=False) |
| 52 | |
| 53 | setattr(layer, name, param) |
| 54 | |
| 55 | |
| 56 | def prepare_humming_linear_layer_config( |
| 57 | layer: LinearBase, |
| 58 | quant_config: dict, |
| 59 | input_quant_config: dict | None = None, |
| 60 | *, |
| 61 | input_schema: "BaseInputSchema | None" = None, |
| 62 | ) -> "LayerConfig": |
| 63 | from vllm.utils.humming import ( |
| 64 | BaseInputSchema, |
| 65 | BaseWeightSchema, |
| 66 | HummingInputSchema, |
| 67 | prepare_layer_config, |
| 68 | transform_humming_tensors, |
| 69 | ) |
| 70 | |
| 71 | weight_schema = BaseWeightSchema.from_config(quant_config) |
| 72 | if input_schema is None: |
| 73 | if input_quant_config is not None: |
| 74 | input_schema = BaseInputSchema.from_config(input_quant_config) |
| 75 | else: |
| 76 | input_schema = HummingInputSchema() |
| 77 | |
| 78 | # ReplicatedLinear has no TP partitioning and so does not set |
| 79 | # input_size_per_partition; for it that is just input_size. Use hasattr |
| 80 | # rather than getattr's default arg, which is evaluated eagerly and would |
| 81 | # raise on layers lacking input_size (e.g. ParallelLMHead). |
| 82 | if hasattr(layer, "input_size_per_partition"): |
| 83 | input_size_per_partition = layer.input_size_per_partition |
| 84 | else: |
| 85 | input_size_per_partition = layer.input_size |
| 86 | shape_k_stacks = [input_size_per_partition] |
| 87 | shape_n_stacks = layer.output_partition_sizes |
| 88 | |
| 89 | # Step 1: convert weight and input schemas to humming standard format |
| 90 | source_tensors = dict(layer.named_parameters()) |
| 91 | weight_schema, tensors = weight_schema.convert_humming( |
| 92 | tensors=source_tensors, |
| 93 | shape_n_stacks=shape_n_stacks, |
| 94 | shape_k_stacks=shape_k_stacks, |
| 95 | param_dtype=layer.params_dtype, |
| 96 | ) |
| 97 | input_schema, input_tensors = input_schema.convert_humming( |
| 98 | tensors=source_tensors, |
| 99 | shape_n_stacks=shape_n_stacks, |
| 100 | shape_k_stacks=shape_k_stacks, |
| 101 | param_dtype=layer.params_dtype, |
| 102 | ) |
| 103 | |
| 104 | if ( |
| 105 | weight_schema.has_zero_point |
| 106 | and input_schema.a_dtype is not None |
| 107 | and weight_schema.b_dtype.num_bits == input_schema.a_dtype.num_bits |
| 108 | ): |
| 109 | raise ValueError( |
| 110 | "Humming does not support asymmetric (zero_point) weights " |
| 111 | "when weight and activation bit widths are equal " |
| 112 | f"(W{weight_schema.b_dtype.num_bits}" |
| 113 | f"A{input_schema.a_dtype.num_bits}). " |
| 114 | "Use symmetric weights or a different activation bit width." |
| 115 | ) |
| 116 | |
| 117 | tensors.update(input_tensors) |
| 118 | layer.weight_schema = weight_schema |
| 119 | input_schema = check_and_fallback_input_schema( |
| 120 | weight_schema=weight_schema, |
| 121 | input_schema=input_schema, |
| 122 | param_dtype=layer.params_dtype, |
| 123 | ) |
| 124 | for name in ("input_scale", "input_scale_2"): |
| 125 | if name != input_schema.static_tensor_scale_name: |
| 126 | tensors.pop(name, None) |
| 127 | input_tensors.pop(name, None) |
| 128 | |
| 129 | # Step 2: transform weight (humming standard format) for forwarding. |
| 130 | config = prepare_layer_config( |
| 131 | shape_n=sum(layer.output_partition_sizes), |
| 132 | shape_k=input_size_per_partition, |
| 133 | weight_schema=weight_schema, |
| 134 | input_schema=input_schema, |
| 135 | pad_n_to_multiple=256, |
| 136 | pad_k_to_multiple=128, |
| 137 | has_bias=layer.has_bias, |
| 138 | torch_dtype=layer.params_dtype, |
| 139 | device=tensors["weight"].device, |
| 140 | ) |
| 141 | tensors = transform_humming_tensors(config, tensors) |
| 142 | tensors.update(input_tensors) |
| 143 | if "bias" in tensors: |
| 144 | zero_bias = torch.zeros_like(tensors["bias"]) |
| 145 | layer.register_buffer("zero_bias", zero_bias) |
| 146 | for name, _ in list(layer.named_parameters()): |
| 147 | delattr(layer, name) |
| 148 | for name, tensor in tensors.items(): |
| 149 | param = torch.nn.Parameter(tensor, requires_grad=False) |
| 150 | setattr(layer, name, param) |
| 151 | |
| 152 | return config |
| 153 | |
| 154 | |
| 155 | def get_humming_linear_compute_config() -> str: |
| 156 | return json.dumps( |
| 157 | { |
| 158 | "use_batch_invariant": envs.VLLM_BATCH_INVARIANT, |
| 159 | "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM, |
| 160 | "gemm_type": "dense", |
| 161 | } |
| 162 | ) |
| 163 | |
| 164 | |
| 165 | def apply_humming_linear( |
| 166 | layer: LinearBase, |
| 167 | x: torch.Tensor, |
| 168 | *, |
| 169 | skip_bias_add: bool = False, |
| 170 | layer_config: "LayerConfig", |
| 171 | compute_config: str, |
| 172 | locks: torch.Tensor, |
| 173 | ) -> torch.Tensor: |
| 174 | from vllm.utils.humming import humming_forward |
| 175 | |
| 176 | mode = layer_config.input_quant_mode |
| 177 | input_scale = getattr(layer, "input_scale", None) if mode.should_quantize else None |
| 178 | input_scale_2 = ( |
| 179 | getattr(layer, "input_scale_2", None) if mode.has_secondary_scale else None |
| 180 | ) |
| 181 | flatten_inputs = x.reshape(-1, x.size(-1)) |
| 182 | output = humming_forward( |
| 183 | layer_config, |
| 184 | inputs=flatten_inputs, |
| 185 | weight=layer.weight, |
| 186 | weight_scale=getattr(layer, "weight_scale", None), |
| 187 | zero_point=getattr(layer, "zero_point", None), |
| 188 | bias=getattr(layer, "zero_bias" if skip_bias_add else "bias", None), |
| 189 | weight_scale_2=getattr(layer, "weight_scale_2", None), |
| 190 | input_scale=input_scale, |
| 191 | input_scale_2=input_scale_2, |
| 192 | hadamard_block_size=layer.weight_schema.hadamard_block_size, |
| 193 | locks=locks, |
| 194 | compute_config=compute_config, |
| 195 | ) |
| 196 | return output.view(*x.shape[:-1], output.size(-1)) |