| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import abstractmethod |
| 5 | from collections.abc import Callable, Iterable |
| 6 | from typing import Any |
| 7 | |
| 8 | import torch |
| 9 | from torch.nn.parameter import Parameter |
| 10 | from typing_extensions import TypeIs |
| 11 | |
| 12 | import vllm.envs as envs |
| 13 | from vllm.config import get_current_vllm_config, get_current_vllm_config_or_none |
| 14 | from vllm.distributed import ( |
| 15 | divide, |
| 16 | get_tensor_model_parallel_rank, |
| 17 | get_tensor_model_parallel_world_size, |
| 18 | split_tensor_along_last_dim, |
| 19 | tensor_model_parallel_all_gather, |
| 20 | tensor_model_parallel_all_reduce, |
| 21 | ) |
| 22 | from vllm.logger import init_logger |
| 23 | from vllm.model_executor.custom_op import PluggableLayer |
| 24 | from vllm.model_executor.determinism.batch_invariant import ( |
| 25 | linear_batch_invariant, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.quantization import resolve_quant_method |
| 28 | from vllm.model_executor.layers.quantization.base_config import ( |
| 29 | QuantizationConfig, |
| 30 | QuantizeMethodBase, |
| 31 | ) |
| 32 | from vllm.model_executor.layers.utils import ( |
| 33 | dispatch_unquantized_gemm, |
| 34 | ) |
| 35 | from vllm.model_executor.parameter import ( |
| 36 | BasevLLMParameter, |
| 37 | BlockQuantScaleParameter, |
| 38 | ModelWeightParameter, |
| 39 | PackedColumnParameter, |
| 40 | PackedvLLMParameter, |
| 41 | PerTensorScaleParameter, |
| 42 | RowvLLMParameter, |
| 43 | ) |
| 44 | from vllm.model_executor.utils import set_weight_attrs |
| 45 | from vllm.platforms import current_platform |
| 46 | |
| 47 | logger = init_logger(__name__) |
| 48 | |
| 49 | WEIGHT_LOADER_V2_SUPPORTED = [ |
| 50 | "UnquantizedLinearMethod", |
| 51 | "CompressedTensorsLinearMethod", |
| 52 | "CompressedTensorsLinearTransformMethod", |
| 53 | "QutlassNvFP4LinearMethod", |
| 54 | "AutoAWQMarlinLinearMethod", |
| 55 | "AutoAWQLinearMethod", |
| 56 | "AutoGPTQLinearMethod", |
| 57 | "Fp8LinearMethod", |
| 58 | "FBGEMMFp8LinearMethod", |
| 59 | "QuarkLinearMethod", |
| 60 | "HummingLinearMethod", |
| 61 | # ModelOptLinearMethod self-registers via |
| 62 | # register_weight_loader_v2_supported_method (see modelopt.py). |
| 63 | ] |
| 64 | |
| 65 | |
| 66 | def register_weight_loader_v2_supported_method(cls): |
| 67 | """Decorator to register a LinearMethod as supporting weight_loader_v2.""" |
| 68 | WEIGHT_LOADER_V2_SUPPORTED.append(cls.__name__) |
| 69 | return cls |
| 70 | |
| 71 | |
| 72 | def adjust_marlin_shard( |
| 73 | param: Parameter, |
| 74 | shard_size: int, |
| 75 | shard_offset: int, |
| 76 | ) -> tuple[int, int]: |
| 77 | marlin_tile_size: int | None = getattr(param, "marlin_tile_size", None) |
| 78 | if marlin_tile_size is None: |
| 79 | return shard_size, shard_offset |
| 80 | |
| 81 | return shard_size * marlin_tile_size, shard_offset * marlin_tile_size |
| 82 | |
| 83 | |
| 84 | def adjust_block_scale_shard( |
| 85 | weight_block_size: tuple[int, ...] | None, |
| 86 | shard_size: int, |
| 87 | shard_offset: int, |
| 88 | ) -> tuple[int, int]: |
| 89 | assert weight_block_size is not None |
| 90 | block_n = weight_block_size[0] |
| 91 | shard_offset = (shard_offset + block_n - 1) // block_n |
| 92 | shard_size = (shard_size + block_n - 1) // block_n |
| 93 | return shard_size, shard_offset |
| 94 | |
| 95 | |
| 96 | def adjust_scalar_to_fused_array( |
| 97 | param_data: torch.Tensor, |
| 98 | loaded_weight: torch.Tensor, |
| 99 | shard_id: int | str, |
| 100 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 101 | """For fused modules (QKV and MLP) we have an array of length |
| 102 | N that holds 1 scale for each "logical" matrix. So the param |
| 103 | is an array of length N. The loaded_weight corresponds to |
| 104 | one of the shards on disk. Here, we slice the param based on |
| 105 | the shard_id for loading. |
| 106 | """ |
| 107 | qkv_idxs = {"q": 0, "k": 1, "v": 2} |
| 108 | |
| 109 | if isinstance(shard_id, str): |
| 110 | shard_id = qkv_idxs[shard_id] |
| 111 | elif not isinstance(shard_id, int): |
| 112 | raise ValueError(f"Unknown Shard Id {shard_id}") |
| 113 | |
| 114 | # AutoFP8 scales do not have a shape |
| 115 | # compressed-tensors scales do have a shape |
| 116 | if len(loaded_weight.shape) != 0: |
| 117 | assert loaded_weight.shape[0] == 1 |
| 118 | loaded_weight = loaded_weight[0] |
| 119 | |
| 120 | return param_data[shard_id], loaded_weight |
| 121 | |
| 122 | |
| 123 | class LinearMethodBase(QuantizeMethodBase): |
| 124 | """Base class for different (maybe quantized) linear methods.""" |
| 125 | |
| 126 | @abstractmethod |
| 127 | def create_weights( |
| 128 | self, |
| 129 | layer: torch.nn.Module, |
| 130 | input_size_per_partition: int, |
| 131 | output_partition_sizes: list[int], |
| 132 | input_size: int, |
| 133 | output_size: int, |
| 134 | params_dtype: torch.dtype, |
| 135 | **extra_weight_attrs, |
| 136 | ): |
| 137 | """Create weights for a linear layer. |
| 138 | The weights will be set as attributes of the layer. |
| 139 | |
| 140 | Args: |
| 141 | layer: The layer that is using the LinearMethodBase factory. |
| 142 | input_size_per_partition: Size of the weight input dim on rank X. |
| 143 | output_partition_sizes: Sizes of the output dim of each logical |
| 144 | weight on rank X. E.g., output_partition_sizes for QKVLinear |
| 145 | is a list contains the width of Wq, Wk, Wv on rank X. |
| 146 | input_size: Size of the input dim of the weight across all ranks. |
| 147 | output_size: Size of the output dim of the weight across all ranks. |
| 148 | params_dtype: Datatype of the parameters. |
| 149 | |
| 150 | """ |
| 151 | raise NotImplementedError |
| 152 | |
| 153 | @abstractmethod |
| 154 | def apply( |
| 155 | self, |
| 156 | layer: torch.nn.Module, |
| 157 | x: torch.Tensor, |
| 158 | bias: torch.Tensor | None = None, |
| 159 | ) -> torch.Tensor: |
| 160 | """Apply the weights in layer to the input tensor. |
| 161 | Expects create_weights to have been called before on the layer.""" |
| 162 | raise NotImplementedError |
| 163 | |
| 164 | |
| 165 | class UnquantizedLinearMethod(LinearMethodBase): |
| 166 | """Linear method without quantization.""" |
| 167 | |
| 168 | supports_pre_processed_weights = True |
| 169 | |
| 170 | def __init__(self) -> None: |
| 171 | config = get_current_vllm_config_or_none() |
| 172 | linear_backend = ( |
| 173 | config.kernel_config.linear_backend if config is not None else "auto" |
| 174 | ) |
| 175 | self._gemm_impl = dispatch_unquantized_gemm(linear_backend) |
| 176 | |
| 177 | def create_weights( |
| 178 | self, |
| 179 | layer: torch.nn.Module, |
| 180 | input_size_per_partition: int, |
| 181 | output_partition_sizes: list[int], |
| 182 | input_size: int, |
| 183 | output_size: int, |
| 184 | params_dtype: torch.dtype, |
| 185 | **extra_weight_attrs, |
| 186 | ): |
| 187 | # This method creates unquantized linear weights. |
| 188 | # The weights are not quantized, and they are not sharded. |
| 189 | # The amount of memory allocated for the weights is |
| 190 | # sum(output_partition_sizes) * input_size_per_partition. |
| 191 | weight_loader = extra_weight_attrs.pop("weight_loader") |
| 192 | weight = ModelWeightParameter( |
| 193 | data=torch.empty( |
| 194 | sum(output_partition_sizes), |
| 195 | input_size_per_partition, |
| 196 | dtype=params_dtype, |
| 197 | ), |
| 198 | input_dim=1, |
| 199 | output_dim=0, |
| 200 | weight_loader=weight_loader, |
| 201 | ) |
| 202 | |
| 203 | layer.register_parameter("weight", weight) |
| 204 | set_weight_attrs(weight, extra_weight_attrs) |
| 205 | |
| 206 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 207 | if current_platform.is_cpu(): |
| 208 | # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not |
| 209 | # perf-critical, so skip packing and use a plain fallback. |
| 210 | if getattr(layer, "_cpu_skip_gemm_dispatch", False): |
| 211 | layer.cpu_linear = torch.nn.functional.linear |
| 212 | return |
| 213 | |
| 214 | from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm |
| 215 | |
| 216 | dispatch_cpu_unquantized_gemm(layer, remove_weight=True) |
| 217 | elif current_platform.is_xpu(): |
| 218 | # Opt-in: F.linear on XPU is faster with an N-contiguous (N, K) weight |
| 219 | # when K > N, but oneDNN's ab-weights matmul is not run-to-run bitwise |
| 220 | # bitwise reproducible. Off by default. |
| 221 | weight = layer.weight.data |
| 222 | if ( |
| 223 | envs.VLLM_XPU_FORCE_N_CONTIG_WEIGHT |
| 224 | and weight.ndim == 2 |
| 225 | and weight.stride(0) != 1 |
| 226 | ): |
| 227 | layer.weight.data = weight.t().contiguous().t() |
| 228 | |
| 229 | def apply( |
| 230 | self, |
| 231 | layer: torch.nn.Module, |
| 232 | x: torch.Tensor, |
| 233 | bias: torch.Tensor | None = None, |
| 234 | ) -> torch.Tensor: |
| 235 | if envs.VLLM_BATCH_INVARIANT and ( |
| 236 | current_platform.is_cuda_alike() or current_platform.is_xpu() |
| 237 | ): |
| 238 | return linear_batch_invariant(x, layer.weight, bias) |
| 239 | return self._gemm_impl(layer, x, layer.weight, bias) |
| 240 | |
| 241 | |
| 242 | class LinearBase(PluggableLayer): |
| 243 | """Base linear layer. |
| 244 | |
| 245 | Args: |
| 246 | input_size: input dimension of the linear layer. |
| 247 | output_size: output dimension of the linear layer. |
| 248 | skip_bias_add: If true, skip adding bias but instead return it. |
| 249 | params_dtype: Data type for the parameters. |
| 250 | quant_config: Quantization configure. |
| 251 | prefix: Prefix for parameter names. |
| 252 | return_bias: If true, return bias together with outputs in forward pass. |
| 253 | disable_tp: If true, tensor parallelism will be disabled for this layer. |
| 254 | |
| 255 | """ |
| 256 | |
| 257 | requires_unquantized_input = False |
| 258 | |
| 259 | def __init__( |
| 260 | self, |
| 261 | input_size: int, |
| 262 | output_size: int, |
| 263 | bias: bool = False, |
| 264 | skip_bias_add: bool = False, |
| 265 | params_dtype: torch.dtype | None = None, |
| 266 | quant_config: QuantizationConfig | None = None, |
| 267 | prefix: str = "", |
| 268 | *, |
| 269 | return_bias: bool = True, |
| 270 | disable_tp: bool = False, |
| 271 | tp_rank: int | None = None, |
| 272 | tp_size: int | None = None, |
| 273 | ): |
| 274 | super().__init__() |
| 275 | |
| 276 | # Keep input parameters |
| 277 | self.input_size = input_size |
| 278 | self.output_size = output_size |
| 279 | self.has_bias = bias |
| 280 | self.skip_bias_add = skip_bias_add |
| 281 | if params_dtype is None: |
| 282 | params_dtype = torch.get_default_dtype() |
| 283 | self.params_dtype = params_dtype |
| 284 | self.quant_config = quant_config |
| 285 | self.prefix = prefix |
| 286 | self.allow_fp8_block_shape_mismatch = False |
| 287 | self.quant_method: QuantizeMethodBase |
| 288 | if quant_config is None: |
| 289 | self.quant_method = UnquantizedLinearMethod() |
| 290 | elif quant_method := resolve_quant_method(quant_config, self, prefix=prefix): |
| 291 | self.quant_method = quant_method |
| 292 | else: |
| 293 | raise ValueError("All linear layers should support quant method.") |
| 294 | self.return_bias = return_bias |
| 295 | self.disable_tp = disable_tp |
| 296 | if disable_tp: |
| 297 | self.tp_rank, self.tp_size = 0, 1 |
| 298 | else: |
| 299 | self.tp_rank = ( |
| 300 | tp_rank if tp_rank is not None else get_tensor_model_parallel_rank() |
| 301 | ) |
| 302 | self.tp_size = ( |
| 303 | tp_size |
| 304 | if tp_size is not None |
| 305 | else get_tensor_model_parallel_world_size() |
| 306 | ) |
| 307 | |
| 308 | def update_param_tp_status(self): |
| 309 | # Single source of truth for a parameter's TP state. BasevLLMParameter |
| 310 | # stamps self.tp_rank with the *global* rank in __init__; this reconciles |
| 311 | # every child parameter to the *layer's* tp_rank/tp_size (which correctly |
| 312 | # accounts for disable_tp -> replicated weights with tp_rank == 0). |
| 313 | # |
| 314 | # Must be re-run whenever parameters are (re-)created after construction, |
| 315 | # e.g. after quant_method.process_weights_after_loading() swaps in fresh |
| 316 | # Parameters. Otherwise a later load_weights()/weight-refit would narrow a |
| 317 | # replicated weight at global_rank * shard_size and overflow. |
| 318 | for param in self.parameters(): |
| 319 | if isinstance(param, BasevLLMParameter): |
| 320 | param.tp_rank = self.tp_rank |
| 321 | param.tp_size = self.tp_size |
| 322 | |
| 323 | |
| 324 | # --8<-- [start:replicated_linear] |
| 325 | @PluggableLayer.register("replicated_linear") |
| 326 | class ReplicatedLinear(LinearBase): |
| 327 | """Replicated linear layer. |
| 328 | |
| 329 | Args: |
| 330 | input_size: input dimension of the linear layer. |
| 331 | output_size: output dimension of the linear layer. |
| 332 | bias: If true, add bias. |
| 333 | skip_bias_add: If true, skip adding bias but instead return it. |
| 334 | params_dtype: Data type for the parameters. |
| 335 | quant_config: Quantization configure. |
| 336 | prefix: The name of the layer in the state dict, including all parents |
| 337 | (e.g. model.layers.0.qkv_proj) |
| 338 | return_bias: If true, return bias together with outputs in forward pass. |
| 339 | disable_tp: Take no effect for replicated linear layers. |
| 340 | |
| 341 | """ |
| 342 | |
| 343 | # --8<-- [end:replicated_linear] |
| 344 | |
| 345 | def __init__( |
| 346 | self, |
| 347 | input_size: int, |
| 348 | output_size: int, |
| 349 | bias: bool = True, |
| 350 | skip_bias_add: bool = False, |
| 351 | params_dtype: torch.dtype | None = None, |
| 352 | quant_config: QuantizationConfig | None = None, |
| 353 | prefix: str = "", |
| 354 | *, |
| 355 | return_bias: bool = True, |
| 356 | disable_tp: bool = False, |
| 357 | ): |
| 358 | # If MergedReplicatedLinear, use output size of each partition. |
| 359 | if hasattr(self, "output_sizes"): |
| 360 | self.output_partition_sizes = self.output_sizes |
| 361 | else: |
| 362 | self.output_partition_sizes = [output_size] |
| 363 | |
| 364 | super().__init__( |
| 365 | input_size, |
| 366 | output_size, |
| 367 | bias, |
| 368 | skip_bias_add, |
| 369 | params_dtype, |
| 370 | quant_config, |
| 371 | prefix=prefix, |
| 372 | return_bias=return_bias, |
| 373 | disable_tp=disable_tp, |
| 374 | ) |
| 375 | |
| 376 | self.quant_method.create_weights( |
| 377 | self, |
| 378 | self.input_size, |
| 379 | self.output_partition_sizes, |
| 380 | self.input_size, |
| 381 | self.output_size, |
| 382 | self.params_dtype, |
| 383 | weight_loader=self.weight_loader, |
| 384 | ) |
| 385 | |
| 386 | if bias: |
| 387 | self.bias = Parameter( |
| 388 | torch.empty(self.output_size, dtype=self.params_dtype), |
| 389 | requires_grad=False, |
| 390 | ) |
| 391 | set_weight_attrs( |
| 392 | self.bias, |
| 393 | { |
| 394 | "output_dim": 0, |
| 395 | "weight_loader": self.weight_loader, |
| 396 | }, |
| 397 | ) |
| 398 | else: |
| 399 | self.register_parameter("bias", None) |
| 400 | |
| 401 | def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): |
| 402 | if len(loaded_weight.shape) == 0: |
| 403 | loaded_weight = loaded_weight.reshape(1) |
| 404 | |
| 405 | assert param.size() == loaded_weight.size(), ( |
| 406 | f"Tried to load weights of size {loaded_weight.size()}" |
| 407 | f"to a parameter of size {param.size()}" |
| 408 | ) |
| 409 | param.data.copy_(loaded_weight) |
| 410 | |
| 411 | def forward( |
| 412 | self, |
| 413 | x: torch.Tensor, |
| 414 | ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: |
| 415 | bias = self.bias if not self.skip_bias_add else None |
| 416 | |
| 417 | output = self.quant_method.apply(self, x, bias) |
| 418 | |
| 419 | if not self.return_bias: |
| 420 | return output |
| 421 | output_bias = self.bias if self.skip_bias_add else None |
| 422 | return output, output_bias |
| 423 | |
| 424 | def extra_repr(self) -> str: |
| 425 | s = f"in_features={self.input_size}" |
| 426 | s += f", output_features={self.output_size}" |
| 427 | s += f", bias={self.bias is not None}" |
| 428 | return s |
| 429 | |
| 430 | |
| 431 | # --8<-- [start:column_parallel_linear] |
| 432 | @PluggableLayer.register("column_parallel_linear") |
| 433 | class ColumnParallelLinear(LinearBase): |
| 434 | """Linear layer with column parallelism. |
| 435 | |
| 436 | The linear layer is defined as Y = XA + b. A is parallelized along |
| 437 | its second dimension as A = [A_1, ..., A_p]. |
| 438 | |
| 439 | Args: |
| 440 | input_size: first dimension of matrix A. |
| 441 | output_size: second dimension of matrix A. |
| 442 | bias: If true, add bias. |
| 443 | gather_output: If true, call all-gather on output and make Y available |
| 444 | to all GPUs, otherwise, every GPU will have its output |
| 445 | which is Y_i = XA_i |
| 446 | skip_bias_add: This was added to enable performance optimizations where |
| 447 | bias can be fused with other element-wise operations. we |
| 448 | skip adding bias but instead return it. |
| 449 | params_dtype: Data type for the parameters. |
| 450 | quant_config: Quantization configure. |
| 451 | prefix: The name of the layer in the state dict, including all parents |
| 452 | (e.g. model.layers.0.qkv_proj) |
| 453 | return_bias: If true, return bias together with outputs in forward pass. |
| 454 | disable_tp: If true, weights matrix won't be sharded through tp rank. |
| 455 | tp_rank: Override the tensor-parallel rank used for sharding. Defaults to |
| 456 | the global TP rank. Used to shard at a coarser granularity than one |
| 457 | shard per rank (see ``DCPGroupColumnParallelLinear``). |
| 458 | tp_size: Override the tensor-parallel world size used for sharding. |
| 459 | Defaults to the global TP world size. |
| 460 | |
| 461 | """ |
| 462 | |
| 463 | # --8<-- [end:column_parallel_linear] |
| 464 | |
| 465 | def __init__( |
| 466 | self, |
| 467 | input_size: int, |
| 468 | output_size: int, |
| 469 | bias: bool = True, |
| 470 | gather_output: bool = False, |
| 471 | skip_bias_add: bool = False, |
| 472 | params_dtype: torch.dtype | None = None, |
| 473 | quant_config: QuantizationConfig | None = None, |
| 474 | prefix: str = "", |
| 475 | *, |
| 476 | return_bias: bool = True, |
| 477 | disable_tp: bool = False, |
| 478 | tp_rank: int | None = None, |
| 479 | tp_size: int | None = None, |
| 480 | ): |
| 481 | # Divide the weight matrix along the last dimension. |
| 482 | if disable_tp: |
| 483 | self.tp_rank, self.tp_size = 0, 1 |
| 484 | else: |
| 485 | self.tp_rank = ( |
| 486 | tp_rank if tp_rank is not None else get_tensor_model_parallel_rank() |
| 487 | ) |
| 488 | self.tp_size = ( |
| 489 | tp_size |
| 490 | if tp_size is not None |
| 491 | else get_tensor_model_parallel_world_size() |
| 492 | ) |
| 493 | self.input_size_per_partition = input_size |
| 494 | self.output_size_per_partition = divide(output_size, self.tp_size) |
| 495 | self.output_partition_sizes = [self.output_size_per_partition] |
| 496 | # If QKV or MergedColumn, use output size of each partition. |
| 497 | if hasattr(self, "output_sizes"): |
| 498 | self.output_partition_sizes = [ |
| 499 | divide(output_size, self.tp_size) for output_size in self.output_sizes |
| 500 | ] |
| 501 | |
| 502 | super().__init__( |
| 503 | input_size, |
| 504 | output_size, |
| 505 | bias, |
| 506 | skip_bias_add, |
| 507 | params_dtype, |
| 508 | quant_config, |
| 509 | prefix, |
| 510 | return_bias=return_bias, |
| 511 | disable_tp=disable_tp, |
| 512 | tp_rank=self.tp_rank, |
| 513 | tp_size=self.tp_size, |
| 514 | ) |
| 515 | |
| 516 | self._maybe_allow_fp8_block_shape_mismatch() |
| 517 | self.gather_output = gather_output |
| 518 | |
| 519 | self.quant_method.create_weights( |
| 520 | layer=self, |
| 521 | input_size_per_partition=self.input_size_per_partition, |
| 522 | output_partition_sizes=self.output_partition_sizes, |
| 523 | input_size=self.input_size, |
| 524 | output_size=self.output_size, |
| 525 | params_dtype=self.params_dtype, |
| 526 | weight_loader=( |
| 527 | self.weight_loader_v2 |
| 528 | if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED |
| 529 | else self.weight_loader |
| 530 | ), |
| 531 | ) |
| 532 | |
| 533 | if bias: |
| 534 | self.bias = Parameter( |
| 535 | torch.empty(self.output_size_per_partition, dtype=params_dtype), |
| 536 | requires_grad=False, |
| 537 | ) |
| 538 | set_weight_attrs( |
| 539 | self.bias, |
| 540 | { |
| 541 | "output_dim": 0, |
| 542 | "weight_loader": self.weight_loader, |
| 543 | }, |
| 544 | ) |
| 545 | else: |
| 546 | self.register_parameter("bias", None) |
| 547 | self.update_param_tp_status() |
| 548 | |
| 549 | def _maybe_allow_fp8_block_shape_mismatch(self) -> None: |
| 550 | quant_config = getattr(self, "quant_config", None) |
| 551 | weight_block = getattr(quant_config, "weight_block_size", None) |
| 552 | if ( |
| 553 | weight_block is None |
| 554 | or len(weight_block) < 1 |
| 555 | or len(self.output_partition_sizes) <= 1 |
| 556 | ): |
| 557 | return |
| 558 | |
| 559 | try: |
| 560 | block_n = int(weight_block[0]) |
| 561 | except (ValueError, TypeError): |
| 562 | return |
| 563 | |
| 564 | if block_n <= 0: |
| 565 | return |
| 566 | |
| 567 | if any(size % block_n != 0 for size in self.output_partition_sizes): |
| 568 | self.allow_fp8_block_shape_mismatch = True |
| 569 | logger.debug( |
| 570 | "Allowing FP8 block shape mismatch for %s (block_n=%d, partitions=%s)", |
| 571 | getattr(self, "prefix", "<unknown>"), |
| 572 | block_n, |
| 573 | self.output_partition_sizes, |
| 574 | ) |
| 575 | |
| 576 | def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): |
| 577 | output_dim = getattr(param, "output_dim", None) |
| 578 | |
| 579 | is_sharded_weight = getattr(param, "is_sharded_weight", False) |
| 580 | |
| 581 | param_data = param.data |
| 582 | if output_dim is not None and not is_sharded_weight: |
| 583 | shard_size = param_data.shape[output_dim] |
| 584 | start_idx = self.tp_rank * shard_size |
| 585 | loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) |
| 586 | |
| 587 | # Special case for loading scales off disk, which often do not |
| 588 | # have a shape (such as in the case of AutoFP8). |
| 589 | if len(loaded_weight.shape) == 0: |
| 590 | loaded_weight = loaded_weight.reshape(1) |
| 591 | |
| 592 | assert param_data.shape == loaded_weight.shape |
| 593 | param_data.copy_(loaded_weight) |
| 594 | |
| 595 | def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor): |
| 596 | # Special case for loading scales off disk, which often do not |
| 597 | # have a shape (such as in the case of AutoFP8). |
| 598 | if len(loaded_weight.shape) == 0: |
| 599 | assert loaded_weight.numel() == 1 |
| 600 | loaded_weight = loaded_weight.reshape(1) |
| 601 | param.load_column_parallel_weight(loaded_weight=loaded_weight) |
| 602 | |
| 603 | def forward( |
| 604 | self, |
| 605 | input_, |
| 606 | ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: |
| 607 | bias = self.bias if not self.skip_bias_add else None |
| 608 | |
| 609 | # Matrix multiply. |
| 610 | output_parallel = self.quant_method.apply(self, input_, bias) |
| 611 | |
| 612 | if self.gather_output and self.tp_size > 1: |
| 613 | # All-gather across the partitions. |
| 614 | output = tensor_model_parallel_all_gather(output_parallel) |
| 615 | else: |
| 616 | output = output_parallel |
| 617 | |
| 618 | if not self.return_bias: |
| 619 | return output |
| 620 | output_bias = self.bias if self.skip_bias_add else None |
| 621 | return output, output_bias |
| 622 | |
| 623 | def extra_repr(self) -> str: |
| 624 | s = f"in_features={self.input_size}" |
| 625 | s += f", output_features={self.output_size_per_partition}" |
| 626 | s += f", bias={self.bias is not None}" |
| 627 | s += f", tp_size={self.tp_size}" |
| 628 | s += f", gather_output={self.gather_output}" |
| 629 | return s |
| 630 | |
| 631 | |
| 632 | class DCPGroupColumnParallelLinear(ColumnParallelLinear): |
| 633 | """Column-parallel linear whose weight is sharded across DCP groups. |
| 634 | |
| 635 | With Decode Context Parallelism (DCP) the KV cache is sharded across a DCP |
| 636 | group, so MLA decode must attend the group's full head set. This layer shards |
| 637 | its output across DCP *groups* (effective tp size ``tp_size // |
| 638 | dcp_world_size``) rather than across every rank, so each rank in a group |
| 639 | holds the whole group's heads, letting decode skip the query all-gather. |
| 640 | |
| 641 | :meth:`forward` returns the group's full head set. :meth:`_local_view` |
| 642 | extracts this rank's TP shard for prefill. |
| 643 | """ |
| 644 | |
| 645 | def __init__(self, *args, **kwargs): |
| 646 | dcp_world_size = ( |
| 647 | get_current_vllm_config().parallel_config.decode_context_parallel_size |
| 648 | ) |
| 649 | rank = get_tensor_model_parallel_rank() |
| 650 | world_size = get_tensor_model_parallel_world_size() |
| 651 | self.group_size = max(dcp_world_size, 1) |
| 652 | self.qrep_active = self.group_size > 1 |
| 653 | self.rank_in_group = rank % self.group_size |
| 654 | super().__init__( |
| 655 | *args, |
| 656 | **kwargs, |
| 657 | tp_rank=rank // self.group_size, |
| 658 | tp_size=world_size // self.group_size, |
| 659 | ) |
| 660 | |
| 661 | def _local_view(self, out: torch.Tensor) -> torch.Tensor: |
| 662 | """Slice this rank's tp head shard from a group-heads output. |
| 663 | |
| 664 | ``out`` is head-shaped, i.e. ``(..., group_heads, head_dim)``. |
| 665 | """ |
| 666 | if self.group_size == 1: |
| 667 | return out |
| 668 | n = out.shape[-2] // self.group_size |
| 669 | start = self.rank_in_group * n |
| 670 | return out[..., start : start + n, :].contiguous() |
| 671 | |
| 672 | |
| 673 | class MergedColumnParallelLinear(ColumnParallelLinear): |
| 674 | """Packed linear layers with column parallelism. |
| 675 | |
| 676 | Similar to ColumnParallelLinear, but the weight matrix is concatenated |
| 677 | along the output dimension. When the weight matrix is loaded, the |
| 678 | different partitions are sharded separately. |
| 679 | |
| 680 | Args: |
| 681 | input_size: input dimension of the linear layer. |
| 682 | output_sizes: list of output dimensions of the linear layer. |
| 683 | bias: If true, add bias. |
| 684 | gather_output: If true, call all-gather on output and make the output |
| 685 | available to all GPUs, otherwise, every GPU will have |
| 686 | its own output. |
| 687 | skip_bias_add: This was added to enable performance optimizations where |
| 688 | bias can be fused with other element-wise operations. we |
| 689 | skip adding bias but instead return it. |
| 690 | params_dtype: Data type for the parameters. |
| 691 | quant_config: Quantization configure. |
| 692 | prefix: The name of the layer in the state dict, including all parents |
| 693 | (e.g. model.layers.0.qkv_proj) |
| 694 | return_bias: If true, return bias together with outputs in forward pass. |
| 695 | disable_tp: If true, all weights matrix won't be sharded, this layer |
| 696 | will be treated as a "Replicated" MergedLinear. |
| 697 | |
| 698 | """ |
| 699 | |
| 700 | def __init__( |
| 701 | self, |
| 702 | input_size: int, |
| 703 | output_sizes: list[int], |
| 704 | bias: bool = True, |
| 705 | gather_output: bool = False, |
| 706 | skip_bias_add: bool = False, |
| 707 | params_dtype: torch.dtype | None = None, |
| 708 | quant_config: QuantizationConfig | None = None, |
| 709 | prefix: str = "", |
| 710 | *, |
| 711 | return_bias: bool = True, |
| 712 | disable_tp: bool = False, |
| 713 | ): |
| 714 | self.output_sizes = output_sizes |
| 715 | self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1 |
| 716 | self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0 |
| 717 | |
| 718 | assert all(output_size % self.tp_size == 0 for output_size in output_sizes) |
| 719 | super().__init__( |
| 720 | input_size=input_size, |
| 721 | output_size=sum(output_sizes), |
| 722 | bias=bias, |
| 723 | gather_output=gather_output, |
| 724 | skip_bias_add=skip_bias_add, |
| 725 | params_dtype=params_dtype, |
| 726 | quant_config=quant_config, |
| 727 | prefix=prefix, |
| 728 | return_bias=return_bias, |
| 729 | disable_tp=disable_tp, |
| 730 | ) |
| 731 | |
| 732 | def validate_shard_id(self, shard_id: Any) -> TypeIs[int | tuple[int, ...] | None]: |
| 733 | if isinstance(shard_id, int): |
| 734 | if shard_id < 0 or shard_id >= len(self.output_sizes): |
| 735 | raise ValueError( |
| 736 | f"Shard id should be between 0 and {len(self.output_sizes) - 1}. " |
| 737 | f"Got shard id {shard_id}." |
| 738 | ) |
| 739 | return True |
| 740 | if shard_id is None: |
| 741 | return True |
| 742 | if isinstance(shard_id, tuple): |
| 743 | for idx in shard_id: |
| 744 | if not (0 <= idx < len(self.output_sizes)): |
| 745 | raise ValueError( |
| 746 | f"Shard id index {idx} should be between 0 and " |
| 747 | f"{len(self.output_sizes) - 1}. Got shard id {shard_id}." |
| 748 | ) |
| 749 | if len(shard_id) > 1 and any( |
| 750 | b - a != 1 for a, b in zip(shard_id[:-1], shard_id[1:]) |
| 751 | ): |
| 752 | raise ValueError( |
| 753 | "Shard id with multiple indices should be consecutive. " |
| 754 | f"Got shard id {shard_id}." |
| 755 | ) |
| 756 | return True |
| 757 | raise ValueError("This line should not be reached") |
| 758 | |
| 759 | def weight_loader( |
| 760 | self, |
| 761 | param: Parameter, |
| 762 | loaded_weight: torch.Tensor, |
| 763 | loaded_shard_id: tuple[int, ...] | int | None = None, |
| 764 | ): |
| 765 | self.validate_shard_id(loaded_shard_id) |
| 766 | |
| 767 | param_data = param.data |
| 768 | output_dim = getattr(param, "output_dim", None) |
| 769 | # Special case for per-tensor scale to load scalar into fused array. |
| 770 | needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) |
| 771 | |
| 772 | if loaded_shard_id is None or isinstance(loaded_shard_id, tuple): |
| 773 | # Loaded weight is already fused on disk (mlp). |
| 774 | # (e.g., Phi-3's gate_up_proj). |
| 775 | if output_dim is None: |
| 776 | if needs_scalar_to_array: |
| 777 | param_data, loaded_weight = adjust_scalar_to_fused_array( |
| 778 | param_data, loaded_weight, 0 |
| 779 | ) |
| 780 | |
| 781 | assert param_data.shape == loaded_weight.shape |
| 782 | param_data.copy_(loaded_weight) |
| 783 | return |
| 784 | |
| 785 | output_sizes = ( |
| 786 | self.output_sizes[loaded_shard_id[0] : loaded_shard_id[-1] + 1] |
| 787 | if loaded_shard_id is not None |
| 788 | else self.output_sizes |
| 789 | ) |
| 790 | current_shard_offset = 0 |
| 791 | shard_offsets: list[tuple[int, int, int]] = [] |
| 792 | for i, output_size in enumerate(output_sizes): |
| 793 | shard_offsets.append((i, current_shard_offset, output_size)) |
| 794 | current_shard_offset += output_size |
| 795 | packed_dim = getattr(param, "packed_dim", None) |
| 796 | for shard_id, shard_offset, shard_size in shard_offsets: |
| 797 | # Special case for Quantization. |
| 798 | # If quantized, we need to adjust the offset and size to account |
| 799 | # for the packing. |
| 800 | # Add check to adjust the size/offset for FP8 block scales |
| 801 | if isinstance(param, BlockQuantScaleParameter): |
| 802 | weight_block_size = getattr(self, "weight_block_size", None) |
| 803 | shard_size, shard_offset = adjust_block_scale_shard( |
| 804 | weight_block_size, shard_size, shard_offset |
| 805 | ) |
| 806 | |
| 807 | if packed_dim == output_dim: |
| 808 | shard_size = shard_size // param.packed_factor |
| 809 | shard_offset = shard_offset // param.packed_factor |
| 810 | # Special case for Marlin. |
| 811 | shard_size, shard_offset = adjust_marlin_shard( |
| 812 | param, shard_size, shard_offset |
| 813 | ) |
| 814 | |
| 815 | loaded_weight_shard = loaded_weight.narrow( |
| 816 | output_dim, shard_offset, shard_size |
| 817 | ) |
| 818 | self.weight_loader(param, loaded_weight_shard, shard_id) |
| 819 | return |
| 820 | |
| 821 | assert loaded_shard_id < len(self.output_sizes) |
| 822 | if output_dim is not None: |
| 823 | shard_offset = sum(self.output_sizes[:loaded_shard_id]) |
| 824 | shard_size = self.output_sizes[loaded_shard_id] |
| 825 | shard_offset //= self.tp_size |
| 826 | shard_size //= self.tp_size |
| 827 | |
| 828 | if isinstance(param, BlockQuantScaleParameter): |
| 829 | weight_block_size = getattr(self, "weight_block_size", None) |
| 830 | shard_size, shard_offset = adjust_block_scale_shard( |
| 831 | weight_block_size, shard_size, shard_offset |
| 832 | ) |
| 833 | |
| 834 | # Special case for quantization. |
| 835 | # If quantized, we need to adjust the offset and size to account |
| 836 | # for the packing. |
| 837 | packed_dim = getattr(param, "packed_dim", None) |
| 838 | if packed_dim == output_dim: |
| 839 | shard_size = round(shard_size // param.packed_factor) |
| 840 | shard_offset = round(shard_offset // param.packed_factor) |
| 841 | # Special case for Marlin. |
| 842 | shard_size, shard_offset = adjust_marlin_shard( |
| 843 | param, shard_size, shard_offset |
| 844 | ) |
| 845 | |
| 846 | is_sharded_weight = getattr(param, "is_sharded_weight", False) |
| 847 | param_data = param_data.narrow(output_dim, shard_offset, shard_size) |
| 848 | start_idx = self.tp_rank * shard_size |
| 849 | if not is_sharded_weight: |
| 850 | loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) |
| 851 | # Special case for per-tensor scales in fused case. |
| 852 | elif needs_scalar_to_array: |
| 853 | param_data, loaded_weight = adjust_scalar_to_fused_array( |
| 854 | param_data, loaded_weight, loaded_shard_id |
| 855 | ) |
| 856 | |
| 857 | else: |
| 858 | ignore_warning = getattr(param, "ignore_warning", False) |
| 859 | if not ignore_warning: |
| 860 | logger.warning( |
| 861 | "Loading a weight without `output_dim` attribute in " |
| 862 | "MergedColumnParallelLinear, assume the weight is " |
| 863 | "the same for all partitions." |
| 864 | ) |
| 865 | |
| 866 | assert param_data.shape == loaded_weight.shape |
| 867 | param_data.copy_(loaded_weight) |
| 868 | |
| 869 | def _load_fused_module_from_checkpoint( |
| 870 | self, |
| 871 | param: BasevLLMParameter, |
| 872 | loaded_weight: torch.Tensor, |
| 873 | output_sizes: list[int] | None = None, |
| 874 | ): |
| 875 | """Handle special case for models where MLP layers are already |
| 876 | fused on disk. In this case, we have no shard id. This function |
| 877 | determines the shard id by splitting these layers and then calls |
| 878 | the weight loader using the shard id. |
| 879 | |
| 880 | An example of a model with these fused layers: |
| 881 | https://huggingface.co/microsoft/Phi-3-mini-4k-instruct |
| 882 | """ |
| 883 | current_shard_offset = 0 |
| 884 | shard_offsets: list[tuple[int, int, int]] = [] |
| 885 | output_sizes = output_sizes or self.output_sizes |
| 886 | for i, output_size in enumerate(output_sizes): |
| 887 | shard_offsets.append((i, current_shard_offset, output_size)) |
| 888 | current_shard_offset += output_size |
| 889 | |
| 890 | for shard_id, shard_offset, shard_size in shard_offsets: |
| 891 | # Special case for Quantization. |
| 892 | # If quantized, we need to adjust the offset and size to account |
| 893 | # for the packing. |
| 894 | if ( |
| 895 | isinstance(param, (PackedColumnParameter, PackedvLLMParameter)) |
| 896 | and param.packed_dim == param.output_dim |
| 897 | ): |
| 898 | shard_size, shard_offset = param.adjust_shard_indexes_for_packing( |
| 899 | shard_size=shard_size, shard_offset=shard_offset |
| 900 | ) |
| 901 | |
| 902 | loaded_weight_shard = loaded_weight.narrow( |
| 903 | param.output_dim, shard_offset, shard_size |
| 904 | ) |
| 905 | self.weight_loader_v2(param, loaded_weight_shard, shard_id) |
| 906 | |
| 907 | def weight_loader_v2( |
| 908 | self, |
| 909 | param: BasevLLMParameter, |
| 910 | loaded_weight: torch.Tensor, |
| 911 | loaded_shard_id: tuple[int, ...] | int | None = None, |
| 912 | ): |
| 913 | self.validate_shard_id(loaded_shard_id) |
| 914 | if loaded_shard_id is None or isinstance(loaded_shard_id, tuple): |
| 915 | if isinstance(param, PerTensorScaleParameter): |
| 916 | if isinstance(loaded_shard_id, tuple): |
| 917 | for idx in loaded_shard_id: |
| 918 | param.load_merged_column_weight( |
| 919 | loaded_weight=loaded_weight, shard_id=idx |
| 920 | ) |
| 921 | else: |
| 922 | # When weights are already fused on disk (e.g. Phi-3's |
| 923 | # gate_up_proj), there is only a single scale for the |
| 924 | # entire fused matrix. Fill all slots with this scale |
| 925 | # to ensure that any subsequent reduction (like .max()) |
| 926 | # works correctly while preserving the parameter shape. |
| 927 | for idx in range(param.data.shape[0]): |
| 928 | param.load_merged_column_weight( |
| 929 | loaded_weight=loaded_weight, shard_id=idx |
| 930 | ) |
| 931 | return |
| 932 | elif type(param) in (RowvLLMParameter, BasevLLMParameter): |
| 933 | param.load_merged_column_weight(loaded_weight=loaded_weight) |
| 934 | return |
| 935 | output_sizes = ( |
| 936 | [self.output_sizes[idx] for idx in loaded_shard_id] |
| 937 | if loaded_shard_id |
| 938 | else None |
| 939 | ) |
| 940 | if isinstance(param, BlockQuantScaleParameter): |
| 941 | weight_block_size = getattr(self, "weight_block_size", None) |
| 942 | output_sizes = [ |
| 943 | adjust_block_scale_shard(weight_block_size, size, 0)[0] |
| 944 | for size in (output_sizes or self.output_sizes) |
| 945 | ] |
| 946 | # TODO: @dsikka - move to parameter.py |
| 947 | self._load_fused_module_from_checkpoint( |
| 948 | param, loaded_weight, output_sizes=output_sizes |
| 949 | ) |
| 950 | return |
| 951 | |
| 952 | assert loaded_shard_id < len(self.output_sizes) |
| 953 | |
| 954 | shard_offset = sum(self.output_sizes[:loaded_shard_id]) |
| 955 | shard_size = self.output_sizes[loaded_shard_id] |
| 956 | shard_offset //= self.tp_size |
| 957 | shard_size //= self.tp_size |
| 958 | |
| 959 | if isinstance(param, BlockQuantScaleParameter): |
| 960 | weight_block_size = getattr(self, "weight_block_size", None) |
| 961 | shard_size, shard_offset = adjust_block_scale_shard( |
| 962 | weight_block_size, shard_size, shard_offset |
| 963 | ) |
| 964 | |
| 965 | param.load_merged_column_weight( |
| 966 | loaded_weight=loaded_weight, |
| 967 | shard_id=loaded_shard_id, |
| 968 | shard_offset=shard_offset, |
| 969 | shard_size=shard_size, |
| 970 | ) |
| 971 | |
| 972 | def load_weights( |
| 973 | self, weights: Iterable[tuple[str, torch.Tensor]] |
| 974 | ) -> Iterable[str]: |
| 975 | for name, loaded_weight in weights: |
| 976 | shard_id = getattr(loaded_weight, "shard_id", None) |
| 977 | self.validate_shard_id(shard_id) |
| 978 | # Load into self if name is not an attr of self or its submodules |
| 979 | param: Parameter |
| 980 | if "." in name: |
| 981 | submodule, _, attr = name.rpartition(".") |
| 982 | param = getattr(self.get_submodule(submodule), attr, self) |
| 983 | else: |
| 984 | param = getattr(self, name, self) |
| 985 | if param is None and name == "bias": |
| 986 | continue |
| 987 | param.weight_loader(param, loaded_weight, shard_id) |
| 988 | logger.debug( |
| 989 | "Loaded shard %s with shape %s into %s.%s", |
| 990 | shard_id, |
| 991 | loaded_weight.shape, |
| 992 | self.prefix, |
| 993 | name, |
| 994 | ) |
| 995 | yield name |
| 996 | |
| 997 | |
| 998 | class QKVParallelLinear(ColumnParallelLinear): |
| 999 | """Linear layers for the attention's QKV transformation. |
| 1000 | |
| 1001 | Linear layers for the linear transformation of the query, key, and value |
| 1002 | vectors in the attention layer. The weight matrix is concatenated along |
| 1003 | the output dimension. The layer is parallelized along the head dimension. |
| 1004 | When the number of key/value heads is smaller than the number of query |
| 1005 | heads (e.g., multi-query/grouped-query attention), the key/value head may |
| 1006 | be replicated while the query heads are partitioned. |
| 1007 | |
| 1008 | Args: |
| 1009 | hidden_size: input hidden state size of the transformer. |
| 1010 | head_size: size of each attention head. |
| 1011 | total_num_heads: total number of attention query heads. |
| 1012 | total_num_kv_heads: total number of attention key/value heads. If |
| 1013 | None, assume total_num_kv_heads = total_num_heads. |
| 1014 | bias: If true, add bias. |
| 1015 | skip_bias_add: This was added to enable performance optimizations where |
| 1016 | bias can be fused with other element-wise operations. we |
| 1017 | skip adding bias but instead return it. |
| 1018 | params_dtype: Data type for the parameters. |
| 1019 | quant_config: Quantization configure. |
| 1020 | prefix: The name of the layer in the state dict, including all parents |
| 1021 | (e.g. model.layers.0.qkv_proj) |
| 1022 | return_bias: If true, return bias together with outputs in forward pass. |
| 1023 | disable_tp: If true, weights matrix won't be sharded through tp rank. |
| 1024 | |
| 1025 | """ |
| 1026 | |
| 1027 | def __init__( |
| 1028 | self, |
| 1029 | hidden_size: int, |
| 1030 | head_size: int, |
| 1031 | total_num_heads: int, |
| 1032 | total_num_kv_heads: int | None = None, |
| 1033 | bias: bool = True, |
| 1034 | skip_bias_add: bool = False, |
| 1035 | params_dtype: torch.dtype | None = None, |
| 1036 | quant_config: QuantizationConfig | None = None, |
| 1037 | prefix: str = "", |
| 1038 | *, |
| 1039 | return_bias: bool = True, |
| 1040 | disable_tp: bool = False, |
| 1041 | v_head_size: int | None = None, |
| 1042 | ): |
| 1043 | self.hidden_size = hidden_size |
| 1044 | self.head_size = head_size |
| 1045 | self.v_head_size = v_head_size if v_head_size is not None else head_size |
| 1046 | self.total_num_heads = total_num_heads |
| 1047 | if total_num_kv_heads is None: |
| 1048 | total_num_kv_heads = total_num_heads |
| 1049 | self.total_num_kv_heads = total_num_kv_heads |
| 1050 | # Divide the weight matrix along the last dimension. |
| 1051 | tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1 |
| 1052 | self.num_heads = divide(self.total_num_heads, tp_size) |
| 1053 | if tp_size >= self.total_num_kv_heads: |
| 1054 | self.num_kv_heads = 1 |
| 1055 | self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads) |
| 1056 | else: |
| 1057 | self.num_kv_heads = divide(self.total_num_kv_heads, tp_size) |
| 1058 | self.num_kv_head_replicas = 1 |
| 1059 | input_size = self.hidden_size |
| 1060 | self.output_sizes = [ |
| 1061 | self.num_heads * self.head_size * tp_size, # q_proj |
| 1062 | self.num_kv_heads * self.head_size * tp_size, # k_proj |
| 1063 | self.num_kv_heads * self.v_head_size * tp_size, # v_proj |
| 1064 | ] |
| 1065 | output_size = sum(self.output_sizes) |
| 1066 | |
| 1067 | super().__init__( |
| 1068 | input_size=input_size, |
| 1069 | output_size=output_size, |
| 1070 | bias=bias, |
| 1071 | gather_output=False, |
| 1072 | skip_bias_add=skip_bias_add, |
| 1073 | params_dtype=params_dtype, |
| 1074 | quant_config=quant_config, |
| 1075 | prefix=prefix, |
| 1076 | return_bias=return_bias, |
| 1077 | disable_tp=disable_tp, |
| 1078 | ) |
| 1079 | |
| 1080 | def validate_shard_id(self, shard_id: Any) -> TypeIs[str | None]: |
| 1081 | if shard_id in {"q", "k", "v"} or shard_id is None: |
| 1082 | return True |
| 1083 | raise ValueError( |
| 1084 | "Shard id for QKVParallelLinear should be 'q', 'k', or 'v', " |
| 1085 | f"got shard id {shard_id}." |
| 1086 | ) |
| 1087 | |
| 1088 | def _get_shard_offset_mapping(self, loaded_shard_id: str): |
| 1089 | shard_offset_mapping = { |
| 1090 | "q": 0, |
| 1091 | "k": self.num_heads * self.head_size, |
| 1092 | "v": (self.num_heads + self.num_kv_heads) * self.head_size, |
| 1093 | "total": (self.num_heads + self.num_kv_heads) * self.head_size |
| 1094 | + self.num_kv_heads * self.v_head_size, |
| 1095 | } |
| 1096 | return shard_offset_mapping.get(loaded_shard_id) |
| 1097 | |
| 1098 | def _get_shard_size_mapping(self, loaded_shard_id: str): |
| 1099 | shard_size_mapping = { |
| 1100 | "q": self.num_heads * self.head_size, |
| 1101 | "k": self.num_kv_heads * self.head_size, |
| 1102 | "v": self.num_kv_heads * self.v_head_size, |
| 1103 | } |
| 1104 | return shard_size_mapping.get(loaded_shard_id) |
| 1105 | |
| 1106 | def _load_fused_module_from_checkpoint( |
| 1107 | self, param: BasevLLMParameter, loaded_weight: torch.Tensor |
| 1108 | ): |
| 1109 | """Handle special case for models where QKV layers are already |
| 1110 | fused on disk. In this case, we have no shard id. This function |
| 1111 | determines the shard id by splitting these layers and then calls |
| 1112 | the weight loader using the shard id. |
| 1113 | |
| 1114 | An example of a model with these fused layers: |
| 1115 | https://huggingface.co/microsoft/Phi-3-mini-4k-instruct |
| 1116 | """ |
| 1117 | shard_offsets = [ |
| 1118 | # (shard_id, shard_offset, shard_size) |
| 1119 | ("q", 0, self.total_num_heads * self.head_size), |
| 1120 | ( |
| 1121 | "k", |
| 1122 | self.total_num_heads * self.head_size, |
| 1123 | self.total_num_kv_heads * self.head_size, |
| 1124 | ), |
| 1125 | ( |
| 1126 | "v", |
| 1127 | (self.total_num_heads + self.total_num_kv_heads) * self.head_size, |
| 1128 | self.total_num_kv_heads * self.v_head_size, |
| 1129 | ), |
| 1130 | ] |
| 1131 | |
| 1132 | for shard_id, shard_offset, shard_size in shard_offsets: |
| 1133 | # Special case for Quantization. |
| 1134 | # If quantized, we need to adjust the offset and size to account |
| 1135 | # for the packing. |
| 1136 | if isinstance(param, BlockQuantScaleParameter): |
| 1137 | weight_block_size = getattr(self, "weight_block_size", None) |
| 1138 | shard_size, shard_offset = adjust_block_scale_shard( |
| 1139 | weight_block_size, shard_size, shard_offset |
| 1140 | ) |
| 1141 | elif ( |
| 1142 | isinstance(param, (PackedColumnParameter, PackedvLLMParameter)) |
| 1143 | and param.packed_dim == param.output_dim |
| 1144 | ): |
| 1145 | shard_size, shard_offset = param.adjust_shard_indexes_for_packing( |
| 1146 | shard_size=shard_size, shard_offset=shard_offset |
| 1147 | ) |
| 1148 | |
| 1149 | loaded_weight_shard = loaded_weight.narrow( |
| 1150 | param.output_dim, shard_offset, shard_size |
| 1151 | ) |
| 1152 | self.weight_loader_v2(param, loaded_weight_shard, shard_id) |
| 1153 | |
| 1154 | def weight_loader_v2( |
| 1155 | self, |
| 1156 | param: BasevLLMParameter, |
| 1157 | loaded_weight: torch.Tensor, |
| 1158 | loaded_shard_id: str | None = None, |
| 1159 | ): |
| 1160 | self.validate_shard_id(loaded_shard_id) |
| 1161 | if loaded_shard_id is None: # special case for certain models |
| 1162 | if isinstance(param, PerTensorScaleParameter): |
| 1163 | # When weights are already fused on disk (e.g. Phi-3's |
| 1164 | # qkv_proj), there is only a single scale for the entire |
| 1165 | # fused matrix. Fill all slots (q, k, v) with this scale |
| 1166 | # to ensure that any subsequent reduction (like .max()) |
| 1167 | # works correctly while preserving the parameter shape. |
| 1168 | for idx in range(param.data.shape[0]): |
| 1169 | param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=idx) |
| 1170 | return |
| 1171 | elif type(param) in (RowvLLMParameter, BasevLLMParameter): |
| 1172 | param.load_qkv_weight(loaded_weight=loaded_weight) |
| 1173 | return |
| 1174 | # TODO: @dsikka - move to parameter.py |
| 1175 | self._load_fused_module_from_checkpoint(param, loaded_weight) |
| 1176 | return |
| 1177 | |
| 1178 | assert loaded_shard_id in ["q", "k", "v"] |
| 1179 | |
| 1180 | shard_offset = self._get_shard_offset_mapping(loaded_shard_id) |
| 1181 | shard_size = self._get_shard_size_mapping(loaded_shard_id) |
| 1182 | assert shard_offset is not None and shard_size is not None |
| 1183 | |
| 1184 | if isinstance(param, BlockQuantScaleParameter): |
| 1185 | weight_block_size = getattr(self, "weight_block_size", None) |
| 1186 | shard_size, shard_offset = adjust_block_scale_shard( |
| 1187 | weight_block_size, shard_size, shard_offset |
| 1188 | ) |
| 1189 | |
| 1190 | param.load_qkv_weight( |
| 1191 | loaded_weight=loaded_weight, |
| 1192 | num_heads=self.num_kv_head_replicas, |
| 1193 | shard_id=loaded_shard_id, |
| 1194 | shard_offset=shard_offset, |
| 1195 | shard_size=shard_size, |
| 1196 | ) |
| 1197 | |
| 1198 | def weight_loader( |
| 1199 | self, |
| 1200 | param: Parameter, |
| 1201 | loaded_weight: torch.Tensor, |
| 1202 | loaded_shard_id: str | None = None, |
| 1203 | ): |
| 1204 | self.validate_shard_id(loaded_shard_id) |
| 1205 | |
| 1206 | param_data = param.data |
| 1207 | output_dim = getattr(param, "output_dim", None) |
| 1208 | |
| 1209 | # Special case for per-tensor scales in fused case. |
| 1210 | needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) |
| 1211 | |
| 1212 | if loaded_shard_id is None: |
| 1213 | # Loaded weight is already fused on disk (qkv). |
| 1214 | # (e.g., Phi-3's qkv_proj). |
| 1215 | if output_dim is None: |
| 1216 | if needs_scalar_to_array: |
| 1217 | param_data, loaded_weight = adjust_scalar_to_fused_array( |
| 1218 | param_data, loaded_weight, 0 |
| 1219 | ) |
| 1220 | |
| 1221 | assert param_data.shape == loaded_weight.shape |
| 1222 | param_data.copy_(loaded_weight) |
| 1223 | return |
| 1224 | shard_offsets = [ |
| 1225 | # (shard_id, shard_offset, shard_size) |
| 1226 | ("q", 0, self.total_num_heads * self.head_size), |
| 1227 | ( |
| 1228 | "k", |
| 1229 | self.total_num_heads * self.head_size, |
| 1230 | self.total_num_kv_heads * self.head_size, |
| 1231 | ), |
| 1232 | ( |
| 1233 | "v", |
| 1234 | (self.total_num_heads + self.total_num_kv_heads) * self.head_size, |
| 1235 | self.total_num_kv_heads * self.v_head_size, |
| 1236 | ), |
| 1237 | ] |
| 1238 | packed_dim = getattr(param, "packed_dim", None) |
| 1239 | for shard_id, shard_offset, shard_size in shard_offsets: |
| 1240 | # Special case for Quantized Weights. |
| 1241 | # If quantized, we need to adjust the offset and size to account |
| 1242 | # for the packing. |
| 1243 | # Add check to adjust the size/offset for FP8 block scales |
| 1244 | if isinstance(param, BlockQuantScaleParameter): |
| 1245 | weight_block_size = getattr(self, "weight_block_size", None) |
| 1246 | shard_size, shard_offset = adjust_block_scale_shard( |
| 1247 | weight_block_size, shard_size, shard_offset |
| 1248 | ) |
| 1249 | |
| 1250 | if packed_dim == output_dim: |
| 1251 | shard_size = round(shard_size // param.packed_factor) |
| 1252 | shard_offset = round(shard_offset // param.packed_factor) |
| 1253 | |
| 1254 | # Special case for Marlin. |
| 1255 | shard_size, shard_offset = adjust_marlin_shard( |
| 1256 | param, shard_size, shard_offset |
| 1257 | ) |
| 1258 | |
| 1259 | loaded_weight_shard = loaded_weight.narrow( |
| 1260 | output_dim, shard_offset, shard_size |
| 1261 | ) |
| 1262 | self.weight_loader(param, loaded_weight_shard, shard_id) |
| 1263 | return |
| 1264 | |
| 1265 | assert loaded_shard_id in ["q", "k", "v"] |
| 1266 | |
| 1267 | # If output dim is defined, use the default loading process. |
| 1268 | if output_dim is not None: |
| 1269 | if loaded_shard_id == "q": |
| 1270 | shard_offset = 0 |
| 1271 | shard_size = self.num_heads * self.head_size |
| 1272 | elif loaded_shard_id == "k": |
| 1273 | shard_offset = self.num_heads * self.head_size |
| 1274 | shard_size = self.num_kv_heads * self.head_size |
| 1275 | elif loaded_shard_id == "v": |
| 1276 | shard_offset = (self.num_heads + self.num_kv_heads) * self.head_size |
| 1277 | shard_size = self.num_kv_heads * self.v_head_size |
| 1278 | |
| 1279 | if isinstance(param, BlockQuantScaleParameter): |
| 1280 | weight_block_size = getattr(self, "weight_block_size", None) |
| 1281 | shard_size, shard_offset = adjust_block_scale_shard( |
| 1282 | weight_block_size, shard_size, shard_offset |
| 1283 | ) |
| 1284 | |
| 1285 | # Special case for Quantized Weights. |
| 1286 | # If quantized, we need to adjust the offset and size to account |
| 1287 | # for the packing. |
| 1288 | packed_dim = getattr(param, "packed_dim", None) |
| 1289 | if packed_dim == output_dim: |
| 1290 | shard_size = round(shard_size // param.packed_factor) |
| 1291 | shard_offset = round(shard_offset // param.packed_factor) |
| 1292 | |
| 1293 | # Special case for Marlin. |
| 1294 | shard_size, shard_offset = adjust_marlin_shard( |
| 1295 | param, shard_size, shard_offset |
| 1296 | ) |
| 1297 | |
| 1298 | is_sharded_weight = getattr(param, "is_sharded_weight", False) |
| 1299 | param_data = param_data.narrow(output_dim, shard_offset, shard_size) |
| 1300 | if loaded_shard_id == "q": |
| 1301 | shard_rank = self.tp_rank |
| 1302 | else: |
| 1303 | shard_rank = self.tp_rank // self.num_kv_head_replicas |
| 1304 | start_idx = shard_rank * shard_size |
| 1305 | |
| 1306 | if not is_sharded_weight: |
| 1307 | loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size) |
| 1308 | |
| 1309 | # Special case for per-tensor scales in fused case. |
| 1310 | elif needs_scalar_to_array: |
| 1311 | param_data, loaded_weight = adjust_scalar_to_fused_array( |
| 1312 | param_data, loaded_weight, loaded_shard_id |
| 1313 | ) |
| 1314 | else: |
| 1315 | ignore_warning = getattr(param, "ignore_warning", False) |
| 1316 | if not ignore_warning: |
| 1317 | logger.warning( |
| 1318 | "Loading a weight without `output_dim` attribute in " |
| 1319 | "QKVParallelLinear, assume the weight is the same " |
| 1320 | "for all partitions." |
| 1321 | ) |
| 1322 | |
| 1323 | assert param_data.shape == loaded_weight.shape |
| 1324 | param_data.copy_(loaded_weight) |
| 1325 | |
| 1326 | def load_weights( |
| 1327 | self, weights: Iterable[tuple[str, torch.Tensor]] |
| 1328 | ) -> Iterable[str]: |
| 1329 | for name, loaded_weight in weights: |
| 1330 | shard_id = getattr(loaded_weight, "shard_id", None) |
| 1331 | self.validate_shard_id(shard_id) |
| 1332 | # Load into self if name is not an attr of self or its submodules |
| 1333 | param: Parameter |
| 1334 | if "." in name: |
| 1335 | submodule, _, attr = name.rpartition(".") |
| 1336 | param = getattr(self.get_submodule(submodule), attr, self) |
| 1337 | else: |
| 1338 | param = getattr(self, name, self) |
| 1339 | if param is None and name == "bias": |
| 1340 | continue |
| 1341 | param.weight_loader(param, loaded_weight, shard_id) |
| 1342 | logger.debug( |
| 1343 | "Loaded shard %s with shape %s into %s.%s", |
| 1344 | shard_id, |
| 1345 | loaded_weight.shape, |
| 1346 | self.prefix, |
| 1347 | name, |
| 1348 | ) |
| 1349 | yield name |
| 1350 | |
| 1351 | |
| 1352 | class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear): |
| 1353 | """QKV projection fused with a lightning-indexer's index_q/index_k. |
| 1354 | |
| 1355 | NOTE: MiniMax-M3-specific. This is tailored to the M3 sparse-attention |
| 1356 | layers (it assumes the indexer's head count equals the KV head count and |
| 1357 | shares the main head_dim); it is not a general-purpose linear layer. It |
| 1358 | lives here only to sit alongside QKVParallelLinear, whose sharding / |
| 1359 | weight-loading machinery it reuses. |
| 1360 | |
| 1361 | A single column-parallel GEMM emits, per rank:: |
| 1362 | |
| 1363 | [q | k | v | index_q | index_k] |
| 1364 | |
| 1365 | ``index_q`` must have the same head count as the KV heads |
| 1366 | (``total_num_index_heads == total_num_kv_heads``) and ``index_head_size == |
| 1367 | head_size``, so it shards exactly like K/V -- including the KV-head |
| 1368 | *replication* path when ``tp_size > total_num_kv_heads`` (this is what makes |
| 1369 | a TP size greater than the KV-head count work). ``index_k`` is a single |
| 1370 | shared head, replicated to every rank. |
| 1371 | """ |
| 1372 | |
| 1373 | def __init__( |
| 1374 | self, |
| 1375 | hidden_size: int, |
| 1376 | head_size: int, |
| 1377 | total_num_heads: int, |
| 1378 | total_num_kv_heads: int, |
| 1379 | total_num_index_heads: int, |
| 1380 | index_head_size: int, |
| 1381 | bias: bool = False, |
| 1382 | quant_config: QuantizationConfig | None = None, |
| 1383 | prefix: str = "", |
| 1384 | ) -> None: |
| 1385 | # index_q rides the KV-head sharding/replication path, so its head count |
| 1386 | # must match the KV heads. |
| 1387 | assert total_num_index_heads == total_num_kv_heads, ( |
| 1388 | "MinimaxM3QKVParallelLinearWithIndexer requires " |
| 1389 | "total_num_index_heads == total_num_kv_heads" |
| 1390 | ) |
| 1391 | self.hidden_size = hidden_size |
| 1392 | self.head_size = head_size |
| 1393 | self.v_head_size = head_size |
| 1394 | self.total_num_heads = total_num_heads |
| 1395 | self.total_num_kv_heads = total_num_kv_heads |
| 1396 | self.total_num_index_heads = total_num_index_heads |
| 1397 | self.index_head_size = index_head_size |
| 1398 | |
| 1399 | tp_size = get_tensor_model_parallel_world_size() |
| 1400 | self.num_heads = divide(self.total_num_heads, tp_size) |
| 1401 | if tp_size >= self.total_num_kv_heads: |
| 1402 | self.num_kv_heads = 1 |
| 1403 | self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads) |
| 1404 | else: |
| 1405 | self.num_kv_heads = divide(self.total_num_kv_heads, tp_size) |
| 1406 | self.num_kv_head_replicas = 1 |
| 1407 | # index_q shards identically to the KV heads. |
| 1408 | self.num_index_heads = self.num_kv_heads |
| 1409 | |
| 1410 | # Global per-group sizes (replicated groups counted x tp_size, matching |
| 1411 | # the QKVParallelLinear convention). index_k is a single replicated head. |
| 1412 | q = self.num_heads * self.head_size |
| 1413 | kv = self.num_kv_heads * self.head_size |
| 1414 | iq = self.num_index_heads * self.index_head_size |
| 1415 | ik = self.index_head_size |
| 1416 | self.output_sizes = [ |
| 1417 | q * tp_size, # q |
| 1418 | kv * tp_size, # k |
| 1419 | kv * tp_size, # v |
| 1420 | iq * tp_size, # index_q |
| 1421 | ik * tp_size, # index_k (replicated) |
| 1422 | ] |
| 1423 | |
| 1424 | # Skip QKVParallelLinear.__init__ (3-group layout); build the 5-group |
| 1425 | # column-parallel weight directly. |
| 1426 | ColumnParallelLinear.__init__( |
| 1427 | self, |
| 1428 | input_size=self.hidden_size, |
| 1429 | output_size=sum(self.output_sizes), |
| 1430 | bias=bias, |
| 1431 | gather_output=False, |
| 1432 | quant_config=quant_config, |
| 1433 | prefix=prefix, |
| 1434 | ) |
| 1435 | |
| 1436 | def validate_shard_id(self, shard_id: Any) -> TypeIs[str | None]: |
| 1437 | if shard_id in {"q", "k", "v", "index_q", "index_k"} or shard_id is None: |
| 1438 | return True |
| 1439 | raise ValueError( |
| 1440 | "Shard id for MinimaxM3QKVParallelLinearWithIndexer must be one of " |
| 1441 | "'q', 'k', 'v', 'index_q', 'index_k'; got " |
| 1442 | f"{shard_id}." |
| 1443 | ) |
| 1444 | |
| 1445 | def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None: |
| 1446 | h = self.head_size |
| 1447 | nq, nkv, nidx = self.num_heads, self.num_kv_heads, self.num_index_heads |
| 1448 | return { |
| 1449 | "q": 0, |
| 1450 | "k": nq * h, |
| 1451 | "v": (nq + nkv) * h, |
| 1452 | "index_q": (nq + 2 * nkv) * h, |
| 1453 | "index_k": (nq + 2 * nkv + nidx) * h, |
| 1454 | }.get(loaded_shard_id) |
| 1455 | |
| 1456 | def _get_shard_size_mapping(self, loaded_shard_id: str) -> int | None: |
| 1457 | h = self.head_size |
| 1458 | return { |
| 1459 | "q": self.num_heads * h, |
| 1460 | "k": self.num_kv_heads * h, |
| 1461 | "v": self.num_kv_heads * h, |
| 1462 | "index_q": self.num_index_heads * h, |
| 1463 | "index_k": self.index_head_size, |
| 1464 | }.get(loaded_shard_id) |
| 1465 | |
| 1466 | def weight_loader_v2( |
| 1467 | self, |
| 1468 | param: BasevLLMParameter, |
| 1469 | loaded_weight: torch.Tensor, |
| 1470 | loaded_shard_id: str | None = None, |
| 1471 | ) -> None: |
| 1472 | self.validate_shard_id(loaded_shard_id) |
| 1473 | # Index checkpoints are never pre-fused on disk; a shard id is always given. |
| 1474 | assert loaded_shard_id in ("q", "k", "v", "index_q", "index_k") |
| 1475 | |
| 1476 | shard_offset = self._get_shard_offset_mapping(loaded_shard_id) |
| 1477 | shard_size = self._get_shard_size_mapping(loaded_shard_id) |
| 1478 | assert shard_offset is not None and shard_size is not None |
| 1479 | if isinstance(param, BlockQuantScaleParameter): |
| 1480 | weight_block_size = getattr(self, "weight_block_size", None) |
| 1481 | shard_size, shard_offset = adjust_block_scale_shard( |
| 1482 | weight_block_size, shard_size, shard_offset |
| 1483 | ) |
| 1484 | |
| 1485 | # index_k is fully replicated: num_heads == tp_size makes |
| 1486 | # load_qkv_weight pick shard_id_int == 0 on every rank. q/k/v/index_q ride |
| 1487 | # the KV-head replication factor. |
| 1488 | num_heads = ( |
| 1489 | self.tp_size if loaded_shard_id == "index_k" else self.num_kv_head_replicas |
| 1490 | ) |
| 1491 | param.load_qkv_weight( |
| 1492 | loaded_weight=loaded_weight, |
| 1493 | num_heads=num_heads, |
| 1494 | shard_id=loaded_shard_id, |
| 1495 | shard_offset=shard_offset, |
| 1496 | shard_size=shard_size, |
| 1497 | ) |
| 1498 | |
| 1499 | def weight_loader( |
| 1500 | self, |
| 1501 | param: Parameter, |
| 1502 | loaded_weight: torch.Tensor, |
| 1503 | loaded_shard_id: str | None = None, |
| 1504 | ) -> None: |
| 1505 | # Unquantized (bf16) path. MXFP8 checkpoints use weight_loader_v2; this |
| 1506 | # keeps an unquantized load correct too. |
| 1507 | self.validate_shard_id(loaded_shard_id) |
| 1508 | assert loaded_shard_id in ("q", "k", "v", "index_q", "index_k") |
| 1509 | output_dim = getattr(param, "output_dim", None) |
| 1510 | assert output_dim is not None |
| 1511 | |
| 1512 | shard_offset = self._get_shard_offset_mapping(loaded_shard_id) |
| 1513 | shard_size = self._get_shard_size_mapping(loaded_shard_id) |
| 1514 | assert shard_offset is not None and shard_size is not None |
| 1515 | if isinstance(param, BlockQuantScaleParameter): |
| 1516 | weight_block_size = getattr(self, "weight_block_size", None) |
| 1517 | shard_size, shard_offset = adjust_block_scale_shard( |
| 1518 | weight_block_size, shard_size, shard_offset |
| 1519 | ) |
| 1520 | |
| 1521 | param_data = param.data.narrow(output_dim, shard_offset, shard_size) |
| 1522 | if loaded_shard_id == "q": |
| 1523 | shard_rank = self.tp_rank |
| 1524 | elif loaded_shard_id == "index_k": |
| 1525 | shard_rank = 0 # replicated to every rank |
| 1526 | else: |
| 1527 | shard_rank = self.tp_rank // self.num_kv_head_replicas |
| 1528 | loaded_weight = loaded_weight.narrow( |
| 1529 | output_dim, shard_rank * shard_size, shard_size |
| 1530 | ) |
| 1531 | assert param_data.shape == loaded_weight.shape |
| 1532 | param_data.copy_(loaded_weight) |
| 1533 | |
| 1534 | |
| 1535 | class KimiK3MergedQKVGateLinear(MergedColumnParallelLinear): |
| 1536 | """Kimi-K3 QKV-A projection fused with its output-gate projection. |
| 1537 | |
| 1538 | NOTE: Kimi-K3-specific. This is tailored to its replicated Q/KV latent |
| 1539 | projections and tensor-parallel output gate; it is not a general-purpose |
| 1540 | linear layer. It lives here alongside ``MergedColumnParallelLinear``, whose |
| 1541 | sharding and weight-loading machinery it reuses. |
| 1542 | |
| 1543 | The per-rank shard layout is ``[q_a | kv_a | gate]``. The latent shards |
| 1544 | stay replicated across tensor-parallel ranks while the gate shard is |
| 1545 | tensor-parallel, matching the standalone projections they replace. |
| 1546 | """ |
| 1547 | |
| 1548 | def __init__( |
| 1549 | self, |
| 1550 | hidden_size: int, |
| 1551 | q_lora_rank: int, |
| 1552 | kv_lora_rank: int, |
| 1553 | qk_rope_head_dim: int, |
| 1554 | total_num_heads: int, |
| 1555 | v_head_dim: int, |
| 1556 | bias: bool = False, |
| 1557 | quant_config: QuantizationConfig | None = None, |
| 1558 | prefix: str = "", |
| 1559 | ) -> None: |
| 1560 | tp_size = get_tensor_model_parallel_world_size() |
| 1561 | |
| 1562 | # MergedColumnParallelLinear divides every logical shard by TP size. |
| 1563 | # Inflate the replicated latent shards to retain their checkpoint widths. |
| 1564 | output_sizes = [ |
| 1565 | q_lora_rank * tp_size, |
| 1566 | (kv_lora_rank + qk_rope_head_dim) * tp_size, |
| 1567 | total_num_heads * v_head_dim, |
| 1568 | ] |
| 1569 | super().__init__( |
| 1570 | input_size=hidden_size, |
| 1571 | output_sizes=output_sizes, |
| 1572 | bias=bias, |
| 1573 | quant_config=quant_config, |
| 1574 | prefix=prefix, |
| 1575 | ) |
| 1576 | |
| 1577 | def _load_with_replicated_latents( |
| 1578 | self, |
| 1579 | loader: Callable[..., None], |
| 1580 | param: Parameter, |
| 1581 | loaded_weight: torch.Tensor, |
| 1582 | loaded_shard_id: tuple[int, ...] | int | None, |
| 1583 | ) -> None: |
| 1584 | is_replicated = isinstance(loaded_shard_id, int) and loaded_shard_id < 2 |
| 1585 | if not is_replicated: |
| 1586 | loader(param, loaded_weight, loaded_shard_id) |
| 1587 | return |
| 1588 | |
| 1589 | # The base loader narrows every checkpoint shard using its TP rank. |
| 1590 | # Replicated latent shards instead load the rank-zero slice on every rank. |
| 1591 | tp_rank = self.tp_rank |
| 1592 | param_tp_rank = getattr(param, "tp_rank", None) |
| 1593 | self.tp_rank = 0 |
| 1594 | if param_tp_rank is not None: |
| 1595 | param.tp_rank = 0 |
| 1596 | try: |
| 1597 | loader(param, loaded_weight, loaded_shard_id) |
| 1598 | finally: |
| 1599 | self.tp_rank = tp_rank |
| 1600 | if param_tp_rank is not None: |
| 1601 | param.tp_rank = param_tp_rank |
| 1602 | |
| 1603 | def weight_loader( |
| 1604 | self, |
| 1605 | param: Parameter, |
| 1606 | loaded_weight: torch.Tensor, |
| 1607 | loaded_shard_id: tuple[int, ...] | int | None = None, |
| 1608 | ) -> None: |
| 1609 | self._load_with_replicated_latents( |
| 1610 | super().weight_loader, param, loaded_weight, loaded_shard_id |
| 1611 | ) |
| 1612 | |
| 1613 | def weight_loader_v2( |
| 1614 | self, |
| 1615 | param: BasevLLMParameter, |
| 1616 | loaded_weight: torch.Tensor, |
| 1617 | loaded_shard_id: tuple[int, ...] | int | None = None, |
| 1618 | ) -> None: |
| 1619 | self._load_with_replicated_latents( |
| 1620 | super().weight_loader_v2, param, loaded_weight, loaded_shard_id |
| 1621 | ) |
| 1622 | |
| 1623 | |
| 1624 | # --8<-- [start:row_parallel_linear] |
| 1625 | @PluggableLayer.register("row_parallel_linear") |
| 1626 | class RowParallelLinear(LinearBase): |
| 1627 | """Linear layer with row parallelism. |
| 1628 | |
| 1629 | The linear layer is defined as Y = XA + b. A is parallelized along |
| 1630 | its first dimension and X along its second dimension as: |
| 1631 | - - |
| 1632 | | A_1 | |
| 1633 | | . | |
| 1634 | A = | . | X = [X_1, ..., X_p] |
| 1635 | | . | |
| 1636 | | A_p | |
| 1637 | - - |
| 1638 | |
| 1639 | Arguments: |
| 1640 | input_size: first dimension of matrix A. |
| 1641 | output_size: second dimension of matrix A. |
| 1642 | bias: If true, add bias. Note that bias is not parallelized. |
| 1643 | input_is_parallel: If true, we assume that the input is already |
| 1644 | split across the GPUs and we do not split |
| 1645 | again. |
| 1646 | skip_bias_add: This was added to enable performance optimization where |
| 1647 | bias can be fused with other element-wise operations. |
| 1648 | We skip adding bias but instead return it. |
| 1649 | params_dtype: Data type for the parameters. |
| 1650 | reduce_results: If true, call all-reduce on output and make Y available |
| 1651 | to all GPUs, otherwise, every GPU will have its output |
| 1652 | which is Y = X_iA_i |
| 1653 | quant_config: Quantization configure. |
| 1654 | prefix: The name of the layer in the state dict, including all parents |
| 1655 | (e.g. model.layers.0.down_proj) |
| 1656 | return_bias: If true, return bias together with outputs in forward pass. |
| 1657 | disable_tp: If true, weights matrix won't be sharded through tp rank. |
| 1658 | |
| 1659 | """ |
| 1660 | |
| 1661 | # --8<-- [end:row_parallel_linear] |
| 1662 | |
| 1663 | def __init__( |
| 1664 | self, |
| 1665 | input_size: int, |
| 1666 | output_size: int, |
| 1667 | bias: bool = True, |
| 1668 | input_is_parallel: bool = True, |
| 1669 | skip_bias_add: bool = False, |
| 1670 | params_dtype: torch.dtype | None = None, |
| 1671 | reduce_results: bool = True, |
| 1672 | quant_config: QuantizationConfig | None = None, |
| 1673 | prefix: str = "", |
| 1674 | *, |
| 1675 | return_bias: bool = True, |
| 1676 | disable_tp: bool = False, |
| 1677 | ): |
| 1678 | # Divide the weight matrix along the first dimension. |
| 1679 | self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0 |
| 1680 | self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1 |
| 1681 | self.input_size_per_partition = divide(input_size, self.tp_size) |
| 1682 | self.output_size_per_partition = output_size |
| 1683 | self.output_partition_sizes = [output_size] |
| 1684 | |
| 1685 | super().__init__( |
| 1686 | input_size, |
| 1687 | output_size, |
| 1688 | bias, |
| 1689 | skip_bias_add, |
| 1690 | params_dtype, |
| 1691 | quant_config, |
| 1692 | prefix, |
| 1693 | return_bias=return_bias, |
| 1694 | disable_tp=disable_tp, |
| 1695 | ) |
| 1696 | |
| 1697 | self.input_is_parallel = input_is_parallel |
| 1698 | self.reduce_results = reduce_results |
| 1699 | |
| 1700 | self.quant_method.create_weights( |
| 1701 | layer=self, |
| 1702 | input_size_per_partition=self.input_size_per_partition, |
| 1703 | output_partition_sizes=self.output_partition_sizes, |
| 1704 | input_size=self.input_size, |
| 1705 | output_size=self.output_size, |
| 1706 | params_dtype=self.params_dtype, |
| 1707 | weight_loader=( |
| 1708 | self.weight_loader_v2 |
| 1709 | if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED |
| 1710 | else self.weight_loader |
| 1711 | ), |
| 1712 | ) |
| 1713 | if not reduce_results and (bias and not skip_bias_add): |
| 1714 | raise ValueError( |
| 1715 | "When not reduce the results, adding bias to the " |
| 1716 | "results can lead to incorrect results" |
| 1717 | ) |
| 1718 | |
| 1719 | if bias: |
| 1720 | self.bias = Parameter( |
| 1721 | torch.empty(self.output_size, dtype=params_dtype), |
| 1722 | requires_grad=False, |
| 1723 | ) |
| 1724 | set_weight_attrs( |
| 1725 | self.bias, |
| 1726 | { |
| 1727 | "output_dim": 0, |
| 1728 | "weight_loader": self.weight_loader, |
| 1729 | }, |
| 1730 | ) |
| 1731 | else: |
| 1732 | self.register_parameter("bias", None) |
| 1733 | self.update_param_tp_status() |
| 1734 | |
| 1735 | def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): |
| 1736 | input_dim = getattr(param, "input_dim", None) |
| 1737 | is_sharded_weight = getattr(param, "is_sharded_weight", False) |
| 1738 | |
| 1739 | param_data = param.data |
| 1740 | if input_dim is not None and not is_sharded_weight: |
| 1741 | shard_size = param_data.shape[input_dim] |
| 1742 | start_idx = self.tp_rank * shard_size |
| 1743 | loaded_weight = loaded_weight.narrow(input_dim, start_idx, shard_size) |
| 1744 | |
| 1745 | # Special case for loading scales off disk, which often do not |
| 1746 | # have a shape (such as in the case of AutoFP8). |
| 1747 | if len(loaded_weight.shape) == 0: |
| 1748 | loaded_weight = loaded_weight.reshape(1) |
| 1749 | |
| 1750 | assert param_data.shape == loaded_weight.shape |
| 1751 | param_data.copy_(loaded_weight) |
| 1752 | |
| 1753 | def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor): |
| 1754 | # Special case for loading scales off disk, which often do not |
| 1755 | # have a shape (such as in the case of AutoFP8). |
| 1756 | if len(loaded_weight.shape) == 0: |
| 1757 | assert loaded_weight.numel() == 1 |
| 1758 | loaded_weight = loaded_weight.reshape(1) |
| 1759 | |
| 1760 | param.load_row_parallel_weight(loaded_weight=loaded_weight) |
| 1761 | |
| 1762 | def forward( |
| 1763 | self, |
| 1764 | input_, |
| 1765 | ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: |
| 1766 | if self.input_is_parallel: |
| 1767 | input_parallel = input_ |
| 1768 | else: |
| 1769 | split_input = split_tensor_along_last_dim( |
| 1770 | input_, num_partitions=self.tp_size |
| 1771 | ) |
| 1772 | input_parallel = split_input[self.tp_rank].contiguous() |
| 1773 | |
| 1774 | # Matrix multiply. |
| 1775 | # Only fuse bias add into GEMM for rank 0 (this ensures that |
| 1776 | # bias will not get added more than once in TP>1 case) |
| 1777 | bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias |
| 1778 | output_parallel = self.quant_method.apply(self, input_parallel, bias_) |
| 1779 | |
| 1780 | if self.reduce_results and self.tp_size > 1: |
| 1781 | output = tensor_model_parallel_all_reduce(output_parallel) |
| 1782 | else: |
| 1783 | output = output_parallel |
| 1784 | |
| 1785 | if not self.return_bias: |
| 1786 | return output |
| 1787 | output_bias = self.bias if self.skip_bias_add else None |
| 1788 | return output, output_bias |
| 1789 | |
| 1790 | def extra_repr(self) -> str: |
| 1791 | s = f"in_features={self.input_size_per_partition}" |
| 1792 | s += f", output_features={self.output_size}" |
| 1793 | s += f", bias={self.bias is not None}" |
| 1794 | s += f", tp_size={self.tp_size}" |
| 1795 | s += f", reduce_results={self.reduce_results}" |
| 1796 | return s |