@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e17e4488bd

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 21:00 UTC

[LoRA] Remove linear hack outside transformers backend (#14177)

Signed-off-by: Isotr0py <2037008807@qq.com>

Isotr0py committed Mar 5, 2025 at 23:06 UTC e17e4488bdfb6aa23a4731287dbe564a8ff62a9d

4 files changed +142 -105

vllm/lora/layers.py

+30 -21

@@ -395,17 +395,20 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

395 x: torch.Tensor,

396 bias: Optional[torch.Tensor] = None) -> torch.Tensor:

397 output = self.base_layer.quant_method.apply(self.base_layer, x, bias)

398 +

399 + # In transformers backend, x and output have extra batch dimension like

400 + # (1, seq_len, hidden_dim), while punica expects (seq_len, hidden_dim),

401 + # therefore we need to flatten the batch dimensions.

402 + if x.ndim == 3 and output.ndim == 3:

403 + output = output.flatten(0, 1)

404 + x = x.flatten(0, 1)

405 +

406 self.punica_wrapper.add_lora_linear(output, x, self.lora_a_stacked,

407 self.lora_b_stacked,

408 self.lora_bias_stacked, 1.0,

409 self.output_slices)

410 return output

411

404 - @classmethod

405 - def get_source_layer(cls, source_layer: nn.Module) -> type:

406 - # Check parent_cls in case source_layer is a HFCompatibleLinear.

407 - return getattr(source_layer, "parent_cls", type(source_layer))

408 -

412

413 class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):

414

@@ -418,7 +421,7 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):

421

422 def forward(

423 self, input_: torch.Tensor

421 - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:

424 + ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:

425 """Forward of ReplicatedLinearWithLoRA

426

427 Args:

@@ -436,6 +439,10 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):

439

440 output_bias = (self.base_layer.bias

441 if self.base_layer.skip_bias_add else None)

442 +

443 + if not self.base_layer.return_bias:

444 + return output

445 +

446 return output, output_bias

447

448 # ReplicatedLinear should always be replaced, regardless of the fully

@@ -448,8 +455,7 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):

455 packed_modules_list: List,

456 model_config: Optional[PretrainedConfig],

457 ) -> bool:

451 - source_layer = cls.get_source_layer(source_layer)

452 - return source_layer is ReplicatedLinear

458 + return type(source_layer) is ReplicatedLinear

459

460

461 class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

@@ -512,7 +518,7 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

518

519 def forward(

520 self, input_: torch.Tensor

515 - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:

521 + ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:

522 """Forward of ColumnParallelLinear

523

524 Args:

@@ -532,6 +538,10 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

538 output = tensor_model_parallel_all_gather(output_parallel)

539 else:

540 output = output_parallel

541 +

542 + if not self.base_layer.return_bias:

543 + return output

544 +

545 output_bias = (self.base_layer.bias

546 if self.base_layer.skip_bias_add else None)

547 return output, output_bias

@@ -545,9 +555,8 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

555 packed_modules_list: List,

556 model_config: Optional[PretrainedConfig],

557 ) -> bool:

548 - source_layer = cls.get_source_layer(source_layer)

549 - return source_layer is ColumnParallelLinear or (

550 - source_layer is MergedColumnParallelLinear

558 + return type(source_layer) is ColumnParallelLinear or (

559 + type(source_layer) is MergedColumnParallelLinear

560 and len(packed_modules_list) == 1)

561

562

@@ -689,8 +698,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

698 packed_modules_list: List,

699 model_config: Optional[PretrainedConfig],

700 ) -> bool:

692 - source_layer = cls.get_source_layer(source_layer)

693 - return (source_layer is MergedColumnParallelLinear

701 + return (type(source_layer) is MergedColumnParallelLinear

702 and len(packed_modules_list) == 2)

703

704

@@ -758,8 +766,7 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

766 def can_replace_layer(cls, source_layer: nn.Module,

767 lora_config: LoRAConfig, packed_modules_list: List,

768 model_config: Optional[PretrainedConfig]) -> bool:

761 - source_layer = cls.get_source_layer(source_layer)

762 - return source_layer is QKVParallelLinear and len(

769 + return type(source_layer) is QKVParallelLinear and len(

770 packed_modules_list) == 1

771

772

@@ -820,8 +827,7 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):

827 packed_modules_list: List,

828 model_config: Optional[PretrainedConfig],

829 ) -> bool:

823 - source_layer = cls.get_source_layer(source_layer)

824 - return (source_layer is QKVParallelLinear

830 + return (type(source_layer) is QKVParallelLinear

831 and len(packed_modules_list) == 3)

832

833

@@ -855,7 +861,7 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

861

862 def forward(

863 self, input_: torch.Tensor

858 - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:

864 + ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[torch.Tensor]]]:

865 """Forward of RowParallelLinear

866

867 Args:

@@ -890,6 +896,10 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

896 else:

897 output = output_

898 output_bias = self.base_layer.bias

899 +

900 + if not self.base_layer.return_bias:

901 + return output

902 +

903 return output, output_bias

904

905 @property

@@ -906,8 +916,7 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

916 packed_modules_list: List,

917 model_config: Optional[PretrainedConfig],

918 ) -> bool:

909 - source_layer = cls.get_source_layer(source_layer)

910 - return source_layer is RowParallelLinear

919 + return type(source_layer) is RowParallelLinear

920

921

922 class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

vllm/lora/utils.py

-10

@@ -67,16 +67,6 @@ def from_layer(layer: nn.Module,

67 packed_modules_list=packed_modules_list,

68 model_config=model_config):

69 instance_layer = lora_cls(layer)

70 - if layer.__class__.__name__ == "HFCompatibleLinear":

71 - # HACK: Make the forward method compatible with the original

72 - # forward method of the instance_layer.

73 - original_forward = instance_layer.forward

74 -

75 - def new_forward(input):

76 - input = input.squeeze(0)

77 - return original_forward(input)[0] # noqa: B023

78 -

79 - instance_layer.forward = new_forward

70 instance_layer.create_lora_weights(max_loras, lora_config,

71 model_config)

72 return instance_layer

vllm/model_executor/layers/linear.py

+110 -61

@@ -2,7 +2,7 @@

3 import itertools

4 from abc import abstractmethod

5 -from typing import Optional

5 +from typing import Optional, Union

7 import torch

8 import torch.nn.functional as F

@@ -152,6 +152,7 @@ class LinearBase(torch.nn.Module):

152 skip_bias_add: If true, skip adding bias but instead return it.

153 params_dtype: Data type for the parameters.

154 quant_config: Quantization configure.

155 + return_bias: If true, return bias together with outputs in forward pass.

156 """

157

158 def __init__(

@@ -162,6 +163,8 @@ class LinearBase(torch.nn.Module):

163 params_dtype: Optional[torch.dtype] = None,

164 quant_config: Optional[QuantizationConfig] = None,

165 prefix: str = "",

166 + *,

167 + return_bias: bool = True,

168 ):

169 super().__init__()

170

@@ -178,9 +181,11 @@ class LinearBase(torch.nn.Module):

181 else:

182 self.quant_method = quant_config.get_quant_method(self,

183 prefix=prefix)

184 + self.return_bias = return_bias

185

182 - def forward(self,

183 - x: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:

186 + def forward(

187 + self, x: torch.Tensor

188 + ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:

189 raise NotImplementedError

190

191

@@ -198,20 +203,25 @@ class ReplicatedLinear(LinearBase):

203 (e.g. model.layers.0.qkv_proj)

204 """

205

201 - def __init__(self,

202 - input_size: int,

203 - output_size: int,

204 - bias: bool = True,

205 - skip_bias_add: bool = False,

206 - params_dtype: Optional[torch.dtype] = None,

207 - quant_config: Optional[QuantizationConfig] = None,

208 - prefix: str = ""):

206 + def __init__(

207 + self,

208 + input_size: int,

209 + output_size: int,

210 + bias: bool = True,

211 + skip_bias_add: bool = False,

212 + params_dtype: Optional[torch.dtype] = None,

213 + quant_config: Optional[QuantizationConfig] = None,

214 + prefix: str = "",

215 + *,

216 + return_bias: bool = True,

217 + ):

218 super().__init__(input_size,

219 output_size,

220 skip_bias_add,

221 params_dtype,

222 quant_config,

214 - prefix=prefix)

223 + prefix=prefix,

224 + return_bias=return_bias)

225

226 # All the linear layer supports quant method.

227 assert self.quant_method is not None

@@ -254,12 +264,15 @@ class ReplicatedLinear(LinearBase):

264 f"to a parameter of size {param.size()}")

265 param.data.copy_(loaded_weight)

266

257 - def forward(self,

258 - x: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:

267 + def forward(

268 + self, x: torch.Tensor

269 + ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:

270 bias = self.bias if not self.skip_bias_add else None

271 assert self.quant_method is not None

272 output = self.quant_method.apply(self, x, bias)

273 output_bias = self.bias if self.skip_bias_add else None

274 + if not self.return_bias:

275 + return output

276 return output, output_bias

277

278 def extra_repr(self) -> str:

@@ -293,16 +306,20 @@ class ColumnParallelLinear(LinearBase):

306 (e.g. model.layers.0.qkv_proj)

307 """

308

296 - def __init__(self,

297 - input_size: int,

298 - output_size: int,

299 - bias: bool = True,

300 - gather_output: bool = False,

301 - skip_bias_add: bool = False,

302 - params_dtype: Optional[torch.dtype] = None,

303 - quant_config: Optional[QuantizationConfig] = None,

304 - output_sizes: Optional[list[int]] = None,

305 - prefix: str = ""):

309 + def __init__(

310 + self,

311 + input_size: int,

312 + output_size: int,

313 + bias: bool = True,

314 + gather_output: bool = False,

315 + skip_bias_add: bool = False,

316 + params_dtype: Optional[torch.dtype] = None,

317 + quant_config: Optional[QuantizationConfig] = None,

318 + output_sizes: Optional[list[int]] = None,

319 + prefix: str = "",

320 + *,

321 + return_bias: bool = True,

322 + ):

323 # Divide the weight matrix along the last dimension.

324 self.tp_size = get_tensor_model_parallel_world_size()

325 self.input_size_per_partition = input_size

@@ -315,8 +332,13 @@ class ColumnParallelLinear(LinearBase):

332 for output_size in self.output_sizes

333 ]

334

318 - super().__init__(input_size, output_size, skip_bias_add, params_dtype,

319 - quant_config, prefix)

335 + super().__init__(input_size,

336 + output_size,

337 + skip_bias_add,

338 + params_dtype,

339 + quant_config,

340 + prefix,

341 + return_bias=return_bias)

342

343 self.gather_output = gather_output

344

@@ -393,7 +415,9 @@ class ColumnParallelLinear(LinearBase):

415 loaded_weight = loaded_weight.reshape(1)

416 param.load_column_parallel_weight(loaded_weight=loaded_weight)

417

396 - def forward(self, input_) -> tuple[torch.Tensor, Optional[Parameter]]:

418 + def forward(

419 + self, input_

420 + ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:

421 bias = self.bias if not self.skip_bias_add else None

422

423 # Matrix multiply.

@@ -405,6 +429,8 @@ class ColumnParallelLinear(LinearBase):

429 else:

430 output = output_parallel

431 output_bias = self.bias if self.skip_bias_add else None

432 + if not self.return_bias:

433 + return output

434 return output, output_bias

435

436 def extra_repr(self) -> str:

@@ -439,15 +465,19 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

465 (e.g. model.layers.0.qkv_proj)

466 """

467

442 - def __init__(self,

443 - input_size: int,

444 - output_sizes: list[int],

445 - bias: bool = True,

446 - gather_output: bool = False,

447 - skip_bias_add: bool = False,

448 - params_dtype: Optional[torch.dtype] = None,

449 - quant_config: Optional[QuantizationConfig] = None,

450 - prefix: str = ""):

468 + def __init__(

469 + self,

470 + input_size: int,

471 + output_sizes: list[int],

472 + bias: bool = True,

473 + gather_output: bool = False,

474 + skip_bias_add: bool = False,

475 + params_dtype: Optional[torch.dtype] = None,

476 + quant_config: Optional[QuantizationConfig] = None,

477 + prefix: str = "",

478 + *,

479 + return_bias: bool = True,

480 + ):

481 self.output_sizes = output_sizes

482 tp_size = get_tensor_model_parallel_world_size()

483 assert all(output_size % tp_size == 0 for output_size in output_sizes)

@@ -458,7 +488,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

488 skip_bias_add=skip_bias_add,

489 params_dtype=params_dtype,

490 quant_config=quant_config,

461 - prefix=prefix)

491 + prefix=prefix,

492 + return_bias=return_bias)

493

494 def weight_loader(self,

495 param: Parameter,

@@ -711,16 +742,20 @@ class QKVParallelLinear(ColumnParallelLinear):

742 (e.g. model.layers.0.qkv_proj)

743 """

744

714 - def __init__(self,

715 - hidden_size: int,

716 - head_size: int,

717 - total_num_heads: int,

718 - total_num_kv_heads: Optional[int] = None,

719 - bias: bool = True,

720 - skip_bias_add: bool = False,

721 - params_dtype: Optional[torch.dtype] = None,

722 - quant_config: Optional[QuantizationConfig] = None,

723 - prefix: str = ""):

745 + def __init__(

746 + self,

747 + hidden_size: int,

748 + head_size: int,

749 + total_num_heads: int,

750 + total_num_kv_heads: Optional[int] = None,

751 + bias: bool = True,

752 + skip_bias_add: bool = False,

753 + params_dtype: Optional[torch.dtype] = None,

754 + quant_config: Optional[QuantizationConfig] = None,

755 + prefix: str = "",

756 + *,

757 + return_bias: bool = True,

758 + ):

759 self.hidden_size = hidden_size

760 self.head_size = head_size

761 self.total_num_heads = total_num_heads

@@ -753,7 +788,8 @@ class QKVParallelLinear(ColumnParallelLinear):

788 skip_bias_add=skip_bias_add,

789 params_dtype=params_dtype,

790 quant_config=quant_config,

756 - prefix=prefix)

791 + prefix=prefix,

792 + return_bias=return_bias)

793

794 def _get_shard_offset_mapping(self, loaded_shard_id: str):

795 shard_offset_mapping = {

@@ -1048,16 +1084,20 @@ class RowParallelLinear(LinearBase):

1084 quant_config: Quantization configure.

1085 """

1086

1051 - def __init__(self,

1052 - input_size: int,

1053 - output_size: int,

1054 - bias: bool = True,

1055 - input_is_parallel: bool = True,

1056 - skip_bias_add: bool = False,

1057 - params_dtype: Optional[torch.dtype] = None,

1058 - reduce_results: bool = True,

1059 - quant_config: Optional[QuantizationConfig] = None,

1060 - prefix: str = ""):

1087 + def __init__(

1088 + self,

1089 + input_size: int,

1090 + output_size: int,

1091 + bias: bool = True,

1092 + input_is_parallel: bool = True,

1093 + skip_bias_add: bool = False,

1094 + params_dtype: Optional[torch.dtype] = None,

1095 + reduce_results: bool = True,

1096 + quant_config: Optional[QuantizationConfig] = None,

1097 + prefix: str = "",

1098 + *,

1099 + return_bias: bool = True,

1100 + ):

1101 # Divide the weight matrix along the first dimension.

1102 self.tp_rank = get_tensor_model_parallel_rank()

1103 self.tp_size = get_tensor_model_parallel_world_size()

@@ -1065,8 +1105,13 @@ class RowParallelLinear(LinearBase):

1105 self.output_size_per_partition = output_size

1106 self.output_partition_sizes = [output_size]

1107

1068 - super().__init__(input_size, output_size, skip_bias_add, params_dtype,

1069 - quant_config, prefix)

1108 + super().__init__(input_size,

1109 + output_size,

1110 + skip_bias_add,

1111 + params_dtype,

1112 + quant_config,

1113 + prefix,

1114 + return_bias=return_bias)

1115

1116 self.input_is_parallel = input_is_parallel

1117 self.reduce_results = reduce_results

@@ -1145,7 +1190,9 @@ class RowParallelLinear(LinearBase):

1190

1191 param.load_row_parallel_weight(loaded_weight=loaded_weight)

1192

1148 - def forward(self, input_) -> tuple[torch.Tensor, Optional[Parameter]]:

1193 + def forward(

1194 + self, input_

1195 + ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]:

1196 if self.input_is_parallel:

1197 input_parallel = input_

1198 else:

@@ -1169,6 +1216,8 @@ class RowParallelLinear(LinearBase):

1216

1217 output_bias = self.bias if self.skip_bias_add else None

1218

1219 + if not self.return_bias:

1220 + return output

1221 return output, output_bias

1222

1223 def extra_repr(self) -> str:

vllm/model_executor/models/transformers.py

+2 -13

@@ -96,23 +96,12 @@ def replace_linear_class(

96 "rowwise": RowParallelLinear,

97 }.get(style, ReplicatedLinear)

99 - class HFCompatibleLinear(vllm_linear_cls):

100 - """

101 - Wrapper class that removes `output_bias` from returned output.

102 - """

103 - # NOTE: The LoRA layer needs to use `parent_cls`.

104 - @property

105 - def parent_cls(self) -> type:

106 - return vllm_linear_cls

107 -

108 - def forward(self, input: torch.Tensor) -> torch.Tensor:

109 - return super().forward(input)[0]

110 -

111 - return HFCompatibleLinear(

99 + return vllm_linear_cls(

100 input_size=linear.in_features,

101 output_size=linear.out_features,

102 bias=linear.bias is not None,

103 quant_config=quant_config,

104 + return_bias=False,

105 )

106

107